fix: beta-binomial numerical instability

storopoli · Dec 3, 2023 · 188894e · 188894e
1 parent e23bf1b
commit 188894e
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/turing/07-robust_beta_binomial_regression-wells.jl b/turing/07-robust_beta_binomial_regression-wells.jl
@@ -24,7 +24,14 @@ X = standardize(ZScoreTransform, X; dims=1)
 y = wells[:, :switch]
 
 # define alternate parameterizations
-BetaBinomial2(n, μ, ϕ) = BetaBinomial(n, μ * ϕ, (1 - μ) * ϕ)
+function BetaBinomial2(n, μ, ϕ)
+    α = μ * ϕ
+    β = (1 - μ) * ϕ
+    α = α > 0 ? α : 1e-4 # numerical stability
+    β = β > 0 ? β : 1e-4 # numerical stability
+
+    return BetaBinomial(n, α, β)
+end
 
 # define the model
 @model function beta_binomial_regression(X, y; predictors=size(X, 2))