From 110e34e1b17fea5ae97625cccc6b14db94144871 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:29:58 +0100
Subject: [PATCH 01/18] Make RealUnivariateDistribution type

---
 EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl |  2 ++
 EpiAware/src/EpiAwareUtils/RealValued.jl    | 11 +++++++++++
 2 files changed, 13 insertions(+)
 create mode 100644 EpiAware/src/EpiAwareUtils/RealValued.jl

diff --git a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
index bcabad2d0..be7ae8b3c 100644
--- a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
+++ b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
@@ -10,6 +10,7 @@ using DynamicPPL: Model, fix, condition, @submodel, @model
 using MCMCChains: Chains
 using Random: AbstractRNG, randexp
 using Tables: rowtable
+import Base: eltype
 
 using Distributions, DocStringExtensions, QuadGK, Statistics, Turing
 
@@ -32,6 +33,7 @@ include("turing-methods.jl")
 include("DirectSample.jl")
 include("post-inference.jl")
 include("get_param_array.jl")
+include("RealValued.jl")
 include("SafePoisson.jl")
 include("SafeNegativeBinomial.jl")
 
diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
new file mode 100644
index 000000000..2358df546
--- /dev/null
+++ b/EpiAware/src/EpiAwareUtils/RealValued.jl
@@ -0,0 +1,11 @@
+"""
+A type to represent real-valued distributions, the purpose of this type is to avoid problems
+with the `eltype` function when having `rand` calls in the model.
+"""
+struct RealValued <: Distributions.ValueSupport end
+Base.eltype(::Type{<:Sampleable{F,RealValued}}) where {F} = Real
+
+"""
+A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
+"""
+const RealUnivariateDistribution = Distribution{Univariate, RealValued}

From 9cc415a79c8980c9137849fa3da4eb673a7f95ea Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:30:23 +0100
Subject: [PATCH 02/18] test RealUnivariateDistribution

---
 EpiAware/test/EpiAwareUtils/RealValued.jl | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 EpiAware/test/EpiAwareUtils/RealValued.jl

diff --git a/EpiAware/test/EpiAwareUtils/RealValued.jl b/EpiAware/test/EpiAwareUtils/RealValued.jl
new file mode 100644
index 000000000..589830ae8
--- /dev/null
+++ b/EpiAware/test/EpiAwareUtils/RealValued.jl
@@ -0,0 +1,14 @@
+@testitem "RealValued Type Tests" begin
+    @test RealValued <: Distributions.ValueSupport
+end
+
+# Test for eltype function
+@testitem "eltype Function Tests" begin
+    struct DummySampleable <: Sampleable{Univariate, RealValued} end
+    @test eltype(DummySampleable) == Real
+end
+
+# Test for RealUnivariateDistribution alias
+@testitem "RealUnivariateDistribution Alias Tests" begin
+    @test RealUnivariateDistribution == Distribution{Univariate, RealValued}
+end

From 678f807d21975f7092d49c07706564ea05dbf9c5 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:30:45 +0100
Subject: [PATCH 03/18] update method

---
 EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl |  2 +-
 EpiAware/src/EpiAwareUtils/SafePoisson.jl          | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
index 4341bcc28..010cdb0ac 100644
--- a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
+++ b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
@@ -65,7 +65,7 @@ var(d)
 2.4617291430060293e40
 ```
 "
-struct SafeNegativeBinomial{T <: Real} <: DiscreteUnivariateDistribution
+struct SafeNegativeBinomial{T <: Real} <: RealUnivariateDistribution
     r::T
     p::T
 
diff --git a/EpiAware/src/EpiAwareUtils/SafePoisson.jl b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
index 46e5201ad..d234dded0 100644
--- a/EpiAware/src/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
@@ -45,7 +45,7 @@ var(d)
 7.016735912097631e20
 ```
 "
-struct SafePoisson{T <: Real} <: DiscreteUnivariateDistribution
+struct SafePoisson{T <: Real} <: RealUnivariateDistribution
     λ::T
 
     SafePoisson{T}(λ::Real) where {T <: Real} = new{T}(λ)
@@ -86,7 +86,7 @@ Distributions.rate(d::SafePoisson) = d.λ
 ### Statistics
 
 Distributions.mean(d::SafePoisson) = d.λ
-Distributions.mode(d::SafePoisson) = _safe_int_floor(d.λ)
+Distributions.mode(d::SafePoisson) = floor(d.λ)
 Distributions.var(d::SafePoisson) = d.λ
 Distributions.skewness(d::SafePoisson) = one(typeof(d.λ)) / sqrt(d.λ)
 Distributions.kurtosis(d::SafePoisson) = one(typeof(d.λ)) / d.λ
@@ -142,12 +142,12 @@ ad_rand(λ) = ad_rand(Random.GLOBAL_RNG, λ)
 function ad_rand(rng::AbstractRNG, λ)
     s = sqrt(λ)
     d = 6.0 * λ^2
-    L = _safe_int_floor(λ - 1.1484)
+    L = floor(λ - 1.1484)
     # Step N
     G = λ + s * randn(rng)
 
     if G >= 0.0
-        K = _safe_int_floor(G)
+        K = floor(G)
         # Step I
         if K >= L
             return K
@@ -177,7 +177,7 @@ function ad_rand(rng::AbstractRNG, λ)
             continue
         end
 
-        K = _safe_int_floor(λ + s * T)
+        K = floor(λ + s * T)
         px, py, fx, fy = procf(λ, K, s)
         c = 0.1069 / λ
 

From 21683e629a031262046fc7514a560cc75d4597bc Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:40:01 +0100
Subject: [PATCH 04/18] reformat

---
 EpiAware/src/EpiAwareUtils/RealValued.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
index 2358df546..d4c27ca88 100644
--- a/EpiAware/src/EpiAwareUtils/RealValued.jl
+++ b/EpiAware/src/EpiAwareUtils/RealValued.jl
@@ -3,7 +3,7 @@ A type to represent real-valued distributions, the purpose of this type is to av
 with the `eltype` function when having `rand` calls in the model.
 """
 struct RealValued <: Distributions.ValueSupport end
-Base.eltype(::Type{<:Sampleable{F,RealValued}}) where {F} = Real
+Base.eltype(::Type{<:Sampleable{F, RealValued}}) where {F} = Real
 
 """
 A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.

From 554ff37cb89cde8bd5b12fd7c88bf47dab3acbcc Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:47:17 +0100
Subject: [PATCH 05/18] modify procf

---
 EpiAware/src/EpiAwareUtils/SafePoisson.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/EpiAware/src/EpiAwareUtils/SafePoisson.jl b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
index d234dded0..357b22fd8 100644
--- a/EpiAware/src/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
@@ -229,7 +229,7 @@ function log1pmx(x::Float64)
 end
 
 # Procedure F
-function procf(λ, K::Int, s::Float64)
+function procf(λ, K, s::Float64)
     # can be pre-computed, but does not seem to affect performance
     ω = 0.3989422804014327 / s
     b1 = 0.041666666666666664 / λ

From c93cbc37b5a2eca06d5c282ab3e65b55a180314e Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:52:20 +0100
Subject: [PATCH 06/18] Explicit use of Distributions

---
 EpiAware/src/EpiAwareUtils/RealValued.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
index d4c27ca88..44ebaa3ce 100644
--- a/EpiAware/src/EpiAwareUtils/RealValued.jl
+++ b/EpiAware/src/EpiAwareUtils/RealValued.jl
@@ -3,9 +3,9 @@ A type to represent real-valued distributions, the purpose of this type is to av
 with the `eltype` function when having `rand` calls in the model.
 """
 struct RealValued <: Distributions.ValueSupport end
-Base.eltype(::Type{<:Sampleable{F, RealValued}}) where {F} = Real
+Base.eltype(::Type{<:Distributions.Sampleable{F, RealValued}}) where {F} = Real
 
 """
 A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
 """
-const RealUnivariateDistribution = Distribution{Univariate, RealValued}
+const RealUnivariateDistribution = Distributions.Distribution{Distributions.Univariate, RealValued}

From 655709629807e81b068ac90a61eb10540f055d72 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 16:06:23 +0100
Subject: [PATCH 07/18] fix unit test fails

---
 EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl          |  3 ++-
 EpiAware/src/EpiAwareUtils/RealValued.jl             |  3 ++-
 EpiAware/src/EpiAwareUtils/SafePoisson.jl            |  2 +-
 EpiAware/test/EpiAwareUtils/RealValued.jl            | 12 +++---------
 EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl  |  2 +-
 EpiAware/test/EpiAwareUtils/SafePoisson.jl           |  8 ++++----
 .../EpiObsModels/modifiers/PrefixObservationModel.jl |  2 +-
 7 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
index be7ae8b3c..af0bab2f4 100644
--- a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
+++ b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
@@ -15,7 +15,8 @@ import Base: eltype
 using Distributions, DocStringExtensions, QuadGK, Statistics, Turing
 
 #Export Structures
-export HalfNormal, DirectSample, SafePoisson, SafeNegativeBinomial
+export HalfNormal, DirectSample, SafePoisson, SafeNegativeBinomial, RealValued,
+       RealUnivariateDistribution
 
 #Export functions
 export scan, spread_draws, censored_cdf, censored_pmf, get_param_array, prefix_submodel, ∫F
diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
index 44ebaa3ce..9b5c4c462 100644
--- a/EpiAware/src/EpiAwareUtils/RealValued.jl
+++ b/EpiAware/src/EpiAwareUtils/RealValued.jl
@@ -8,4 +8,5 @@ Base.eltype(::Type{<:Distributions.Sampleable{F, RealValued}}) where {F} = Real
 """
 A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
 """
-const RealUnivariateDistribution = Distributions.Distribution{Distributions.Univariate, RealValued}
+const RealUnivariateDistribution = Distributions.Distribution{
+    Distributions.Univariate, RealValued}
diff --git a/EpiAware/src/EpiAwareUtils/SafePoisson.jl b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
index 357b22fd8..0ec1122b9 100644
--- a/EpiAware/src/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
@@ -241,7 +241,7 @@ function procf(λ, K, s::Float64)
 
     if K < 10
         px = -float(λ)
-        py = λ^K / factorial(K)
+        py = λ^K / factorial(floor(Int, K))
     else
         δ = 0.08333333333333333 / K
         δ -= 4.8 * δ^3
diff --git a/EpiAware/test/EpiAwareUtils/RealValued.jl b/EpiAware/test/EpiAwareUtils/RealValued.jl
index 589830ae8..13205084d 100644
--- a/EpiAware/test/EpiAwareUtils/RealValued.jl
+++ b/EpiAware/test/EpiAwareUtils/RealValued.jl
@@ -1,14 +1,8 @@
 @testitem "RealValued Type Tests" begin
-    @test RealValued <: Distributions.ValueSupport
-end
-
-# Test for eltype function
-@testitem "eltype Function Tests" begin
+    using Distributions
     struct DummySampleable <: Sampleable{Univariate, RealValued} end
-    @test eltype(DummySampleable) == Real
-end
 
-# Test for RealUnivariateDistribution alias
-@testitem "RealUnivariateDistribution Alias Tests" begin
+    @test RealValued <: Distributions.ValueSupport
+    @test eltype(DummySampleable) == Real
     @test RealUnivariateDistribution == Distribution{Univariate, RealValued}
 end
diff --git a/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl b/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
index 8c828e49d..7e5bb98cd 100644
--- a/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
+++ b/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
@@ -70,7 +70,7 @@ end
 
     dist = SafeNegativeBinomial(r, p)
     @testset "Large value of mean samples a BigInt with SafePoisson" begin
-        @test rand(dist) isa BigInt
+        @test rand(dist) isa Real
     end
     @testset "Large value of mean sample failure with Poisson" begin
         _dist = EpiAware.EpiAwareUtils._negbin(dist)
diff --git a/EpiAware/test/EpiAwareUtils/SafePoisson.jl b/EpiAware/test/EpiAwareUtils/SafePoisson.jl
index 6630e066b..e4f889254 100644
--- a/EpiAware/test/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/test/EpiAwareUtils/SafePoisson.jl
@@ -2,9 +2,9 @@
     λ = 10.0
     dist = SafePoisson(λ)
     @test typeof(dist) <: SafePoisson
-    @test rand(dist) isa Int
-    @test rand(dist, 10) isa Vector{Int}
-    @test rand(dist, 10, 10) isa Array{Int}
+    @test rand(dist) isa Real
+    @test rand(dist, 10) isa Vector{Real}
+    @test rand(dist, 10, 10) isa Array{Real}
 end
 
 @testitem "Check distribution properties of SafePoisson" begin
@@ -54,7 +54,7 @@ end
     bigλ = exp(48.0) #Large value of λ
     dist = SafePoisson(bigλ)
     @testset "Large value of mean samples a BigInt with SafePoisson" begin
-        @test rand(dist) isa BigInt
+        @test rand(dist) isa Real
     end
     @testset "Large value of mean sample failure with Poisson" begin
         _dist = Poisson(dist.λ)
diff --git a/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl b/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl
index b46ba54c4..346d2b9ba 100644
--- a/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl
+++ b/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl
@@ -10,5 +10,5 @@ end
 
     mdl = generate_observations(model, missing, 10)
     draw = rand(mdl)
-    @test typeof(draw[:var"Test.y_t[1]"]) <: Int
+    @test typeof(draw[:var"Test.y_t[1]"]) <: Real
 end

From 986ef7e48d2e635c4390514e85991b81c8ef9674 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 22:29:05 +0100
Subject: [PATCH 08/18] Discrete valued dists expecting rand return of
 Union{int,BigInt}

---
 EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl      |  6 +++---
 EpiAware/src/EpiAwareUtils/RealValued.jl         | 12 ------------
 EpiAware/src/EpiAwareUtils/SafeInt.jl            | 16 ++++++++++++++++
 .../src/EpiAwareUtils/SafeNegativeBinomial.jl    |  2 +-
 EpiAware/src/EpiAwareUtils/SafePoisson.jl        | 12 ++++++------
 EpiAware/test/EpiAwareUtils/RealValued.jl        |  8 --------
 EpiAware/test/EpiAwareUtils/SafeInt.jl           |  8 ++++++++
 .../test/EpiAwareUtils/SafeNegativeBinomial.jl   |  2 +-
 EpiAware/test/EpiAwareUtils/SafePoisson.jl       |  8 ++++----
 9 files changed, 39 insertions(+), 35 deletions(-)
 delete mode 100644 EpiAware/src/EpiAwareUtils/RealValued.jl
 create mode 100644 EpiAware/src/EpiAwareUtils/SafeInt.jl
 delete mode 100644 EpiAware/test/EpiAwareUtils/RealValued.jl
 create mode 100644 EpiAware/test/EpiAwareUtils/SafeInt.jl

diff --git a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
index af0bab2f4..22aceaa91 100644
--- a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
+++ b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
@@ -15,8 +15,8 @@ import Base: eltype
 using Distributions, DocStringExtensions, QuadGK, Statistics, Turing
 
 #Export Structures
-export HalfNormal, DirectSample, SafePoisson, SafeNegativeBinomial, RealValued,
-       RealUnivariateDistribution
+export HalfNormal, DirectSample, SafePoisson, SafeNegativeBinomial, SafeIntValued, SafeInt,
+       SafeDiscreteUnivariateDistribution
 
 #Export functions
 export scan, spread_draws, censored_cdf, censored_pmf, get_param_array, prefix_submodel, ∫F
@@ -34,7 +34,7 @@ include("turing-methods.jl")
 include("DirectSample.jl")
 include("post-inference.jl")
 include("get_param_array.jl")
-include("RealValued.jl")
+include("SafeInt.jl")
 include("SafePoisson.jl")
 include("SafeNegativeBinomial.jl")
 
diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
deleted file mode 100644
index 9b5c4c462..000000000
--- a/EpiAware/src/EpiAwareUtils/RealValued.jl
+++ /dev/null
@@ -1,12 +0,0 @@
-"""
-A type to represent real-valued distributions, the purpose of this type is to avoid problems
-with the `eltype` function when having `rand` calls in the model.
-"""
-struct RealValued <: Distributions.ValueSupport end
-Base.eltype(::Type{<:Distributions.Sampleable{F, RealValued}}) where {F} = Real
-
-"""
-A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
-"""
-const RealUnivariateDistribution = Distributions.Distribution{
-    Distributions.Univariate, RealValued}
diff --git a/EpiAware/src/EpiAwareUtils/SafeInt.jl b/EpiAware/src/EpiAwareUtils/SafeInt.jl
new file mode 100644
index 000000000..a4d26fad4
--- /dev/null
+++ b/EpiAware/src/EpiAwareUtils/SafeInt.jl
@@ -0,0 +1,16 @@
+const SafeInt = Union{Int, BigInt}
+
+"""
+A type to represent real-valued distributions, the purpose of this type is to avoid problems
+with the `eltype` function when having `rand` calls in the model.
+"""
+struct SafeIntValued <: Distributions.ValueSupport end
+function Base.eltype(::Type{<:Distributions.Sampleable{F, SafeIntValued}}) where {F}
+    Union{Int, BigInt}
+end
+
+"""
+A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
+"""
+const SafeDiscreteUnivariateDistribution = Distributions.Distribution{
+    Distributions.Univariate, SafeIntValued}
diff --git a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
index 010cdb0ac..e57714b80 100644
--- a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
+++ b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
@@ -65,7 +65,7 @@ var(d)
 2.4617291430060293e40
 ```
 "
-struct SafeNegativeBinomial{T <: Real} <: RealUnivariateDistribution
+struct SafeNegativeBinomial{T <: Real} <: SafeDiscreteUnivariateDistribution
     r::T
     p::T
 
diff --git a/EpiAware/src/EpiAwareUtils/SafePoisson.jl b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
index 0ec1122b9..c85243ebc 100644
--- a/EpiAware/src/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
@@ -45,7 +45,7 @@ var(d)
 7.016735912097631e20
 ```
 "
-struct SafePoisson{T <: Real} <: RealUnivariateDistribution
+struct SafePoisson{T <: Real} <: SafeDiscreteUnivariateDistribution
     λ::T
 
     SafePoisson{T}(λ::Real) where {T <: Real} = new{T}(λ)
@@ -142,12 +142,12 @@ ad_rand(λ) = ad_rand(Random.GLOBAL_RNG, λ)
 function ad_rand(rng::AbstractRNG, λ)
     s = sqrt(λ)
     d = 6.0 * λ^2
-    L = floor(λ - 1.1484)
+    L = _safe_int_floor(λ - 1.1484)
     # Step N
     G = λ + s * randn(rng)
 
     if G >= 0.0
-        K = floor(G)
+        K = _safe_int_floor(G)
         # Step I
         if K >= L
             return K
@@ -177,7 +177,7 @@ function ad_rand(rng::AbstractRNG, λ)
             continue
         end
 
-        K = floor(λ + s * T)
+        K = _safe_int_floor(λ + s * T)
         px, py, fx, fy = procf(λ, K, s)
         c = 0.1069 / λ
 
@@ -229,7 +229,7 @@ function log1pmx(x::Float64)
 end
 
 # Procedure F
-function procf(λ, K, s::Float64)
+function procf(λ, K::SafeInt, s::Float64)
     # can be pre-computed, but does not seem to affect performance
     ω = 0.3989422804014327 / s
     b1 = 0.041666666666666664 / λ
@@ -241,7 +241,7 @@ function procf(λ, K, s::Float64)
 
     if K < 10
         px = -float(λ)
-        py = λ^K / factorial(floor(Int, K))
+        py = λ^K / factorial(K)
     else
         δ = 0.08333333333333333 / K
         δ -= 4.8 * δ^3
diff --git a/EpiAware/test/EpiAwareUtils/RealValued.jl b/EpiAware/test/EpiAwareUtils/RealValued.jl
deleted file mode 100644
index 13205084d..000000000
--- a/EpiAware/test/EpiAwareUtils/RealValued.jl
+++ /dev/null
@@ -1,8 +0,0 @@
-@testitem "RealValued Type Tests" begin
-    using Distributions
-    struct DummySampleable <: Sampleable{Univariate, RealValued} end
-
-    @test RealValued <: Distributions.ValueSupport
-    @test eltype(DummySampleable) == Real
-    @test RealUnivariateDistribution == Distribution{Univariate, RealValued}
-end
diff --git a/EpiAware/test/EpiAwareUtils/SafeInt.jl b/EpiAware/test/EpiAwareUtils/SafeInt.jl
new file mode 100644
index 000000000..d1a59dfab
--- /dev/null
+++ b/EpiAware/test/EpiAwareUtils/SafeInt.jl
@@ -0,0 +1,8 @@
+@testitem "SafeInt Type Tests" begin
+    using Distributions
+    struct DummySampleable <: Sampleable{Univariate, SafeIntValued} end
+
+    @test SafeIntValued <: Distributions.ValueSupport
+    @test eltype(DummySampleable) <: Union{Int, BigInt}
+    @test SafeDiscreteUnivariateDistribution == Distribution{Univariate, SafeIntValued}
+end
diff --git a/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl b/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
index 7e5bb98cd..4e1ee09db 100644
--- a/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
+++ b/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
@@ -70,7 +70,7 @@ end
 
     dist = SafeNegativeBinomial(r, p)
     @testset "Large value of mean samples a BigInt with SafePoisson" begin
-        @test rand(dist) isa Real
+        @test rand(dist) isa Union{Int, BigInt}
     end
     @testset "Large value of mean sample failure with Poisson" begin
         _dist = EpiAware.EpiAwareUtils._negbin(dist)
diff --git a/EpiAware/test/EpiAwareUtils/SafePoisson.jl b/EpiAware/test/EpiAwareUtils/SafePoisson.jl
index e4f889254..874cf3423 100644
--- a/EpiAware/test/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/test/EpiAwareUtils/SafePoisson.jl
@@ -2,9 +2,9 @@
     λ = 10.0
     dist = SafePoisson(λ)
     @test typeof(dist) <: SafePoisson
-    @test rand(dist) isa Real
-    @test rand(dist, 10) isa Vector{Real}
-    @test rand(dist, 10, 10) isa Array{Real}
+    @test rand(dist) isa SafeInt
+    @test rand(dist, 10) isa Vector{SafeInt}
+    @test rand(dist, 10, 10) isa Array{SafeInt}
 end
 
 @testitem "Check distribution properties of SafePoisson" begin
@@ -54,7 +54,7 @@ end
     bigλ = exp(48.0) #Large value of λ
     dist = SafePoisson(bigλ)
     @testset "Large value of mean samples a BigInt with SafePoisson" begin
-        @test rand(dist) isa Real
+        @test rand(dist) isa SafeInt
     end
     @testset "Large value of mean sample failure with Poisson" begin
         _dist = Poisson(dist.λ)

From 4e2f876f0ee8cb0cd56cc67c9ee70e0151917657 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:29:58 +0100
Subject: [PATCH 09/18] Make RealUnivariateDistribution type

---
 EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl |  2 ++
 EpiAware/src/EpiAwareUtils/RealValued.jl    | 11 +++++++++++
 2 files changed, 13 insertions(+)
 create mode 100644 EpiAware/src/EpiAwareUtils/RealValued.jl

diff --git a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
index bcabad2d0..be7ae8b3c 100644
--- a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
+++ b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
@@ -10,6 +10,7 @@ using DynamicPPL: Model, fix, condition, @submodel, @model
 using MCMCChains: Chains
 using Random: AbstractRNG, randexp
 using Tables: rowtable
+import Base: eltype
 
 using Distributions, DocStringExtensions, QuadGK, Statistics, Turing
 
@@ -32,6 +33,7 @@ include("turing-methods.jl")
 include("DirectSample.jl")
 include("post-inference.jl")
 include("get_param_array.jl")
+include("RealValued.jl")
 include("SafePoisson.jl")
 include("SafeNegativeBinomial.jl")
 
diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
new file mode 100644
index 000000000..2358df546
--- /dev/null
+++ b/EpiAware/src/EpiAwareUtils/RealValued.jl
@@ -0,0 +1,11 @@
+"""
+A type to represent real-valued distributions, the purpose of this type is to avoid problems
+with the `eltype` function when having `rand` calls in the model.
+"""
+struct RealValued <: Distributions.ValueSupport end
+Base.eltype(::Type{<:Sampleable{F,RealValued}}) where {F} = Real
+
+"""
+A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
+"""
+const RealUnivariateDistribution = Distribution{Univariate, RealValued}

From 66b9b7cf1a51c231fdc3c89bacfaba6c9815d96e Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:30:23 +0100
Subject: [PATCH 10/18] test RealUnivariateDistribution

---
 EpiAware/test/EpiAwareUtils/RealValued.jl | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 EpiAware/test/EpiAwareUtils/RealValued.jl

diff --git a/EpiAware/test/EpiAwareUtils/RealValued.jl b/EpiAware/test/EpiAwareUtils/RealValued.jl
new file mode 100644
index 000000000..589830ae8
--- /dev/null
+++ b/EpiAware/test/EpiAwareUtils/RealValued.jl
@@ -0,0 +1,14 @@
+@testitem "RealValued Type Tests" begin
+    @test RealValued <: Distributions.ValueSupport
+end
+
+# Test for eltype function
+@testitem "eltype Function Tests" begin
+    struct DummySampleable <: Sampleable{Univariate, RealValued} end
+    @test eltype(DummySampleable) == Real
+end
+
+# Test for RealUnivariateDistribution alias
+@testitem "RealUnivariateDistribution Alias Tests" begin
+    @test RealUnivariateDistribution == Distribution{Univariate, RealValued}
+end

From a5814c46a520efb3d79185c2d107c80fe908814d Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:30:45 +0100
Subject: [PATCH 11/18] update method

---
 EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl |  2 +-
 EpiAware/src/EpiAwareUtils/SafePoisson.jl          | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
index 4341bcc28..010cdb0ac 100644
--- a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
+++ b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
@@ -65,7 +65,7 @@ var(d)
 2.4617291430060293e40
 ```
 "
-struct SafeNegativeBinomial{T <: Real} <: DiscreteUnivariateDistribution
+struct SafeNegativeBinomial{T <: Real} <: RealUnivariateDistribution
     r::T
     p::T
 
diff --git a/EpiAware/src/EpiAwareUtils/SafePoisson.jl b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
index 46e5201ad..d234dded0 100644
--- a/EpiAware/src/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
@@ -45,7 +45,7 @@ var(d)
 7.016735912097631e20
 ```
 "
-struct SafePoisson{T <: Real} <: DiscreteUnivariateDistribution
+struct SafePoisson{T <: Real} <: RealUnivariateDistribution
     λ::T
 
     SafePoisson{T}(λ::Real) where {T <: Real} = new{T}(λ)
@@ -86,7 +86,7 @@ Distributions.rate(d::SafePoisson) = d.λ
 ### Statistics
 
 Distributions.mean(d::SafePoisson) = d.λ
-Distributions.mode(d::SafePoisson) = _safe_int_floor(d.λ)
+Distributions.mode(d::SafePoisson) = floor(d.λ)
 Distributions.var(d::SafePoisson) = d.λ
 Distributions.skewness(d::SafePoisson) = one(typeof(d.λ)) / sqrt(d.λ)
 Distributions.kurtosis(d::SafePoisson) = one(typeof(d.λ)) / d.λ
@@ -142,12 +142,12 @@ ad_rand(λ) = ad_rand(Random.GLOBAL_RNG, λ)
 function ad_rand(rng::AbstractRNG, λ)
     s = sqrt(λ)
     d = 6.0 * λ^2
-    L = _safe_int_floor(λ - 1.1484)
+    L = floor(λ - 1.1484)
     # Step N
     G = λ + s * randn(rng)
 
     if G >= 0.0
-        K = _safe_int_floor(G)
+        K = floor(G)
         # Step I
         if K >= L
             return K
@@ -177,7 +177,7 @@ function ad_rand(rng::AbstractRNG, λ)
             continue
         end
 
-        K = _safe_int_floor(λ + s * T)
+        K = floor(λ + s * T)
         px, py, fx, fy = procf(λ, K, s)
         c = 0.1069 / λ
 

From 4f22138fa1c673a4b6838ddd6d6358105c340a12 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:40:01 +0100
Subject: [PATCH 12/18] reformat

---
 EpiAware/src/EpiAwareUtils/RealValued.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
index 2358df546..d4c27ca88 100644
--- a/EpiAware/src/EpiAwareUtils/RealValued.jl
+++ b/EpiAware/src/EpiAwareUtils/RealValued.jl
@@ -3,7 +3,7 @@ A type to represent real-valued distributions, the purpose of this type is to av
 with the `eltype` function when having `rand` calls in the model.
 """
 struct RealValued <: Distributions.ValueSupport end
-Base.eltype(::Type{<:Sampleable{F,RealValued}}) where {F} = Real
+Base.eltype(::Type{<:Sampleable{F, RealValued}}) where {F} = Real
 
 """
 A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.

From 2576c20b2b55aef4bd9c2b41c49cf93434baae11 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:47:17 +0100
Subject: [PATCH 13/18] modify procf

---
 EpiAware/src/EpiAwareUtils/SafePoisson.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/EpiAware/src/EpiAwareUtils/SafePoisson.jl b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
index d234dded0..357b22fd8 100644
--- a/EpiAware/src/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
@@ -229,7 +229,7 @@ function log1pmx(x::Float64)
 end
 
 # Procedure F
-function procf(λ, K::Int, s::Float64)
+function procf(λ, K, s::Float64)
     # can be pre-computed, but does not seem to affect performance
     ω = 0.3989422804014327 / s
     b1 = 0.041666666666666664 / λ

From 83b59d7d8d1fe5159fc77a04c0978a05d565499f Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 15:52:20 +0100
Subject: [PATCH 14/18] Explicit use of Distributions

---
 EpiAware/src/EpiAwareUtils/RealValued.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
index d4c27ca88..44ebaa3ce 100644
--- a/EpiAware/src/EpiAwareUtils/RealValued.jl
+++ b/EpiAware/src/EpiAwareUtils/RealValued.jl
@@ -3,9 +3,9 @@ A type to represent real-valued distributions, the purpose of this type is to av
 with the `eltype` function when having `rand` calls in the model.
 """
 struct RealValued <: Distributions.ValueSupport end
-Base.eltype(::Type{<:Sampleable{F, RealValued}}) where {F} = Real
+Base.eltype(::Type{<:Distributions.Sampleable{F, RealValued}}) where {F} = Real
 
 """
 A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
 """
-const RealUnivariateDistribution = Distribution{Univariate, RealValued}
+const RealUnivariateDistribution = Distributions.Distribution{Distributions.Univariate, RealValued}

From 1ff41c2be4ac8ee583f4bd2b6b8e55ca0fea88d1 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 16:06:23 +0100
Subject: [PATCH 15/18] fix unit test fails

---
 EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl          |  3 ++-
 EpiAware/src/EpiAwareUtils/RealValued.jl             |  3 ++-
 EpiAware/src/EpiAwareUtils/SafePoisson.jl            |  2 +-
 EpiAware/test/EpiAwareUtils/RealValued.jl            | 12 +++---------
 EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl  |  2 +-
 EpiAware/test/EpiAwareUtils/SafePoisson.jl           |  8 ++++----
 .../EpiObsModels/modifiers/PrefixObservationModel.jl |  2 +-
 7 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
index be7ae8b3c..af0bab2f4 100644
--- a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
+++ b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
@@ -15,7 +15,8 @@ import Base: eltype
 using Distributions, DocStringExtensions, QuadGK, Statistics, Turing
 
 #Export Structures
-export HalfNormal, DirectSample, SafePoisson, SafeNegativeBinomial
+export HalfNormal, DirectSample, SafePoisson, SafeNegativeBinomial, RealValued,
+       RealUnivariateDistribution
 
 #Export functions
 export scan, spread_draws, censored_cdf, censored_pmf, get_param_array, prefix_submodel, ∫F
diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
index 44ebaa3ce..9b5c4c462 100644
--- a/EpiAware/src/EpiAwareUtils/RealValued.jl
+++ b/EpiAware/src/EpiAwareUtils/RealValued.jl
@@ -8,4 +8,5 @@ Base.eltype(::Type{<:Distributions.Sampleable{F, RealValued}}) where {F} = Real
 """
 A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
 """
-const RealUnivariateDistribution = Distributions.Distribution{Distributions.Univariate, RealValued}
+const RealUnivariateDistribution = Distributions.Distribution{
+    Distributions.Univariate, RealValued}
diff --git a/EpiAware/src/EpiAwareUtils/SafePoisson.jl b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
index 357b22fd8..0ec1122b9 100644
--- a/EpiAware/src/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
@@ -241,7 +241,7 @@ function procf(λ, K, s::Float64)
 
     if K < 10
         px = -float(λ)
-        py = λ^K / factorial(K)
+        py = λ^K / factorial(floor(Int, K))
     else
         δ = 0.08333333333333333 / K
         δ -= 4.8 * δ^3
diff --git a/EpiAware/test/EpiAwareUtils/RealValued.jl b/EpiAware/test/EpiAwareUtils/RealValued.jl
index 589830ae8..13205084d 100644
--- a/EpiAware/test/EpiAwareUtils/RealValued.jl
+++ b/EpiAware/test/EpiAwareUtils/RealValued.jl
@@ -1,14 +1,8 @@
 @testitem "RealValued Type Tests" begin
-    @test RealValued <: Distributions.ValueSupport
-end
-
-# Test for eltype function
-@testitem "eltype Function Tests" begin
+    using Distributions
     struct DummySampleable <: Sampleable{Univariate, RealValued} end
-    @test eltype(DummySampleable) == Real
-end
 
-# Test for RealUnivariateDistribution alias
-@testitem "RealUnivariateDistribution Alias Tests" begin
+    @test RealValued <: Distributions.ValueSupport
+    @test eltype(DummySampleable) == Real
     @test RealUnivariateDistribution == Distribution{Univariate, RealValued}
 end
diff --git a/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl b/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
index 8c828e49d..7e5bb98cd 100644
--- a/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
+++ b/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
@@ -70,7 +70,7 @@ end
 
     dist = SafeNegativeBinomial(r, p)
     @testset "Large value of mean samples a BigInt with SafePoisson" begin
-        @test rand(dist) isa BigInt
+        @test rand(dist) isa Real
     end
     @testset "Large value of mean sample failure with Poisson" begin
         _dist = EpiAware.EpiAwareUtils._negbin(dist)
diff --git a/EpiAware/test/EpiAwareUtils/SafePoisson.jl b/EpiAware/test/EpiAwareUtils/SafePoisson.jl
index 6630e066b..e4f889254 100644
--- a/EpiAware/test/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/test/EpiAwareUtils/SafePoisson.jl
@@ -2,9 +2,9 @@
     λ = 10.0
     dist = SafePoisson(λ)
     @test typeof(dist) <: SafePoisson
-    @test rand(dist) isa Int
-    @test rand(dist, 10) isa Vector{Int}
-    @test rand(dist, 10, 10) isa Array{Int}
+    @test rand(dist) isa Real
+    @test rand(dist, 10) isa Vector{Real}
+    @test rand(dist, 10, 10) isa Array{Real}
 end
 
 @testitem "Check distribution properties of SafePoisson" begin
@@ -54,7 +54,7 @@ end
     bigλ = exp(48.0) #Large value of λ
     dist = SafePoisson(bigλ)
     @testset "Large value of mean samples a BigInt with SafePoisson" begin
-        @test rand(dist) isa BigInt
+        @test rand(dist) isa Real
     end
     @testset "Large value of mean sample failure with Poisson" begin
         _dist = Poisson(dist.λ)
diff --git a/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl b/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl
index b46ba54c4..346d2b9ba 100644
--- a/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl
+++ b/EpiAware/test/EpiObsModels/modifiers/PrefixObservationModel.jl
@@ -10,5 +10,5 @@ end
 
     mdl = generate_observations(model, missing, 10)
     draw = rand(mdl)
-    @test typeof(draw[:var"Test.y_t[1]"]) <: Int
+    @test typeof(draw[:var"Test.y_t[1]"]) <: Real
 end

From 7c51d1f535babcc6470f23d766237bf24942aa02 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Thu, 10 Oct 2024 22:29:05 +0100
Subject: [PATCH 16/18] Discrete valued dists expecting rand return of
 Union{int,BigInt}

---
 EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl      |  6 +++---
 EpiAware/src/EpiAwareUtils/RealValued.jl         | 12 ------------
 EpiAware/src/EpiAwareUtils/SafeInt.jl            | 16 ++++++++++++++++
 .../src/EpiAwareUtils/SafeNegativeBinomial.jl    |  2 +-
 EpiAware/src/EpiAwareUtils/SafePoisson.jl        | 12 ++++++------
 EpiAware/test/EpiAwareUtils/RealValued.jl        |  8 --------
 EpiAware/test/EpiAwareUtils/SafeInt.jl           |  8 ++++++++
 .../test/EpiAwareUtils/SafeNegativeBinomial.jl   |  2 +-
 EpiAware/test/EpiAwareUtils/SafePoisson.jl       |  8 ++++----
 9 files changed, 39 insertions(+), 35 deletions(-)
 delete mode 100644 EpiAware/src/EpiAwareUtils/RealValued.jl
 create mode 100644 EpiAware/src/EpiAwareUtils/SafeInt.jl
 delete mode 100644 EpiAware/test/EpiAwareUtils/RealValued.jl
 create mode 100644 EpiAware/test/EpiAwareUtils/SafeInt.jl

diff --git a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
index af0bab2f4..22aceaa91 100644
--- a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
+++ b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl
@@ -15,8 +15,8 @@ import Base: eltype
 using Distributions, DocStringExtensions, QuadGK, Statistics, Turing
 
 #Export Structures
-export HalfNormal, DirectSample, SafePoisson, SafeNegativeBinomial, RealValued,
-       RealUnivariateDistribution
+export HalfNormal, DirectSample, SafePoisson, SafeNegativeBinomial, SafeIntValued, SafeInt,
+       SafeDiscreteUnivariateDistribution
 
 #Export functions
 export scan, spread_draws, censored_cdf, censored_pmf, get_param_array, prefix_submodel, ∫F
@@ -34,7 +34,7 @@ include("turing-methods.jl")
 include("DirectSample.jl")
 include("post-inference.jl")
 include("get_param_array.jl")
-include("RealValued.jl")
+include("SafeInt.jl")
 include("SafePoisson.jl")
 include("SafeNegativeBinomial.jl")
 
diff --git a/EpiAware/src/EpiAwareUtils/RealValued.jl b/EpiAware/src/EpiAwareUtils/RealValued.jl
deleted file mode 100644
index 9b5c4c462..000000000
--- a/EpiAware/src/EpiAwareUtils/RealValued.jl
+++ /dev/null
@@ -1,12 +0,0 @@
-"""
-A type to represent real-valued distributions, the purpose of this type is to avoid problems
-with the `eltype` function when having `rand` calls in the model.
-"""
-struct RealValued <: Distributions.ValueSupport end
-Base.eltype(::Type{<:Distributions.Sampleable{F, RealValued}}) where {F} = Real
-
-"""
-A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
-"""
-const RealUnivariateDistribution = Distributions.Distribution{
-    Distributions.Univariate, RealValued}
diff --git a/EpiAware/src/EpiAwareUtils/SafeInt.jl b/EpiAware/src/EpiAwareUtils/SafeInt.jl
new file mode 100644
index 000000000..a4d26fad4
--- /dev/null
+++ b/EpiAware/src/EpiAwareUtils/SafeInt.jl
@@ -0,0 +1,16 @@
+const SafeInt = Union{Int, BigInt}
+
+"""
+A type to represent real-valued distributions, the purpose of this type is to avoid problems
+with the `eltype` function when having `rand` calls in the model.
+"""
+struct SafeIntValued <: Distributions.ValueSupport end
+function Base.eltype(::Type{<:Distributions.Sampleable{F, SafeIntValued}}) where {F}
+    Union{Int, BigInt}
+end
+
+"""
+A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
+"""
+const SafeDiscreteUnivariateDistribution = Distributions.Distribution{
+    Distributions.Univariate, SafeIntValued}
diff --git a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
index 010cdb0ac..e57714b80 100644
--- a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
+++ b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
@@ -65,7 +65,7 @@ var(d)
 2.4617291430060293e40
 ```
 "
-struct SafeNegativeBinomial{T <: Real} <: RealUnivariateDistribution
+struct SafeNegativeBinomial{T <: Real} <: SafeDiscreteUnivariateDistribution
     r::T
     p::T
 
diff --git a/EpiAware/src/EpiAwareUtils/SafePoisson.jl b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
index 0ec1122b9..c85243ebc 100644
--- a/EpiAware/src/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/src/EpiAwareUtils/SafePoisson.jl
@@ -45,7 +45,7 @@ var(d)
 7.016735912097631e20
 ```
 "
-struct SafePoisson{T <: Real} <: RealUnivariateDistribution
+struct SafePoisson{T <: Real} <: SafeDiscreteUnivariateDistribution
     λ::T
 
     SafePoisson{T}(λ::Real) where {T <: Real} = new{T}(λ)
@@ -142,12 +142,12 @@ ad_rand(λ) = ad_rand(Random.GLOBAL_RNG, λ)
 function ad_rand(rng::AbstractRNG, λ)
     s = sqrt(λ)
     d = 6.0 * λ^2
-    L = floor(λ - 1.1484)
+    L = _safe_int_floor(λ - 1.1484)
     # Step N
     G = λ + s * randn(rng)
 
     if G >= 0.0
-        K = floor(G)
+        K = _safe_int_floor(G)
         # Step I
         if K >= L
             return K
@@ -177,7 +177,7 @@ function ad_rand(rng::AbstractRNG, λ)
             continue
         end
 
-        K = floor(λ + s * T)
+        K = _safe_int_floor(λ + s * T)
         px, py, fx, fy = procf(λ, K, s)
         c = 0.1069 / λ
 
@@ -229,7 +229,7 @@ function log1pmx(x::Float64)
 end
 
 # Procedure F
-function procf(λ, K, s::Float64)
+function procf(λ, K::SafeInt, s::Float64)
     # can be pre-computed, but does not seem to affect performance
     ω = 0.3989422804014327 / s
     b1 = 0.041666666666666664 / λ
@@ -241,7 +241,7 @@ function procf(λ, K, s::Float64)
 
     if K < 10
         px = -float(λ)
-        py = λ^K / factorial(floor(Int, K))
+        py = λ^K / factorial(K)
     else
         δ = 0.08333333333333333 / K
         δ -= 4.8 * δ^3
diff --git a/EpiAware/test/EpiAwareUtils/RealValued.jl b/EpiAware/test/EpiAwareUtils/RealValued.jl
deleted file mode 100644
index 13205084d..000000000
--- a/EpiAware/test/EpiAwareUtils/RealValued.jl
+++ /dev/null
@@ -1,8 +0,0 @@
-@testitem "RealValued Type Tests" begin
-    using Distributions
-    struct DummySampleable <: Sampleable{Univariate, RealValued} end
-
-    @test RealValued <: Distributions.ValueSupport
-    @test eltype(DummySampleable) == Real
-    @test RealUnivariateDistribution == Distribution{Univariate, RealValued}
-end
diff --git a/EpiAware/test/EpiAwareUtils/SafeInt.jl b/EpiAware/test/EpiAwareUtils/SafeInt.jl
new file mode 100644
index 000000000..d1a59dfab
--- /dev/null
+++ b/EpiAware/test/EpiAwareUtils/SafeInt.jl
@@ -0,0 +1,8 @@
+@testitem "SafeInt Type Tests" begin
+    using Distributions
+    struct DummySampleable <: Sampleable{Univariate, SafeIntValued} end
+
+    @test SafeIntValued <: Distributions.ValueSupport
+    @test eltype(DummySampleable) <: Union{Int, BigInt}
+    @test SafeDiscreteUnivariateDistribution == Distribution{Univariate, SafeIntValued}
+end
diff --git a/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl b/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
index 7e5bb98cd..4e1ee09db 100644
--- a/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
+++ b/EpiAware/test/EpiAwareUtils/SafeNegativeBinomial.jl
@@ -70,7 +70,7 @@ end
 
     dist = SafeNegativeBinomial(r, p)
     @testset "Large value of mean samples a BigInt with SafePoisson" begin
-        @test rand(dist) isa Real
+        @test rand(dist) isa Union{Int, BigInt}
     end
     @testset "Large value of mean sample failure with Poisson" begin
         _dist = EpiAware.EpiAwareUtils._negbin(dist)
diff --git a/EpiAware/test/EpiAwareUtils/SafePoisson.jl b/EpiAware/test/EpiAwareUtils/SafePoisson.jl
index e4f889254..874cf3423 100644
--- a/EpiAware/test/EpiAwareUtils/SafePoisson.jl
+++ b/EpiAware/test/EpiAwareUtils/SafePoisson.jl
@@ -2,9 +2,9 @@
     λ = 10.0
     dist = SafePoisson(λ)
     @test typeof(dist) <: SafePoisson
-    @test rand(dist) isa Real
-    @test rand(dist, 10) isa Vector{Real}
-    @test rand(dist, 10, 10) isa Array{Real}
+    @test rand(dist) isa SafeInt
+    @test rand(dist, 10) isa Vector{SafeInt}
+    @test rand(dist, 10, 10) isa Array{SafeInt}
 end
 
 @testitem "Check distribution properties of SafePoisson" begin
@@ -54,7 +54,7 @@ end
     bigλ = exp(48.0) #Large value of λ
     dist = SafePoisson(bigλ)
     @testset "Large value of mean samples a BigInt with SafePoisson" begin
-        @test rand(dist) isa Real
+        @test rand(dist) isa SafeInt
     end
     @testset "Large value of mean sample failure with Poisson" begin
         _dist = Poisson(dist.λ)

From fa2b48c48da7251ae2d505834bca9ae3e6714f8b Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Fri, 11 Oct 2024 11:10:46 +0100
Subject: [PATCH 17/18] Typo

---
 EpiAware/src/EpiAwareUtils/SafeInt.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/EpiAware/src/EpiAwareUtils/SafeInt.jl b/EpiAware/src/EpiAwareUtils/SafeInt.jl
index a4d26fad4..565327eed 100644
--- a/EpiAware/src/EpiAwareUtils/SafeInt.jl
+++ b/EpiAware/src/EpiAwareUtils/SafeInt.jl
@@ -6,11 +6,11 @@ with the `eltype` function when having `rand` calls in the model.
 """
 struct SafeIntValued <: Distributions.ValueSupport end
 function Base.eltype(::Type{<:Distributions.Sampleable{F, SafeIntValued}}) where {F}
-    Union{Int, BigInt}
+    SafeInt
 end
 
 """
-A constant alias for `Distribution{Univariate, RealValued}`. This type represents a univariate distribution with real-valued outcomes.
+A constant alias for `Distribution{Univariate, SafeIntValued}`. This type represents a univariate distribution with real-valued outcomes.
 """
 const SafeDiscreteUnivariateDistribution = Distributions.Distribution{
     Distributions.Univariate, SafeIntValued}

From dfe658a0260241768f1259958e814f0e54a6d264 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Fri, 11 Oct 2024 11:10:54 +0100
Subject: [PATCH 18/18] Revert CI test

---
 EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl b/EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl
index 2b4b52102..69d4ec93d 100644
--- a/EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl
+++ b/EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl
@@ -175,9 +175,7 @@ end
         ExpGrowthRate,
         Renewal] .|>
                 em_type -> em_type(
-        data = EpiData([0.2, 0.5, 0.3],
-            em_type == Renewal ? softplus : exp
-        ),
+        data = EpiData([0.2, 0.5, 0.3], exp),
         initialisation_prior = Normal(log(100.0), 0.25)
     )