Allow customizing truncation max_n_steps in Hurdle Mixtures

bwengals · May 31, 2024 · 508a134 · 508a134
1 parent 7d15175
commit 508a134
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pymc/distributions/mixture.py b/pymc/distributions/mixture.py
@@ -798,7 +798,7 @@ def dist(cls, psi, mu=None, alpha=None, p=None, n=None, **kwargs):
         )
 
 
-def _hurdle_mixture(*, name, nonzero_p, nonzero_dist, dtype, **kwargs):
+def _hurdle_mixture(*, name, nonzero_p, nonzero_dist, dtype, max_n_steps=10_000, **kwargs):
     """Helper function to create a hurdle mixtures
 
     If name is `None`, this function returns an unregistered variable
@@ -819,7 +819,7 @@ def _hurdle_mixture(*, name, nonzero_p, nonzero_dist, dtype, **kwargs):
     weights = pt.stack([1 - nonzero_p, nonzero_p], axis=-1)
     comp_dists = [
         DiracDelta.dist(zero),
-        Truncated.dist(nonzero_dist, lower=lower),
+        Truncated.dist(nonzero_dist, lower=lower, max_n_steps=max_n_steps),
     ]
 
     if name is not None: