Bump minimum numpyro to >=0.21.0 and remove epsilon padding for zero-mean Poisson/NegBinomial

pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "jax>=0.7.0",
     "numpy>=2.4.4",
     "polars>=1.2.1",
-    "numpyro>=0.19.0",
+    "numpyro>=0.21.0",
 ]
 
 [dependency-groups]

pyrenew/datasets/datagen_he_CA_126.py

@@ -245,7 +245,7 @@ def sample_negbinom(
     Parameters
     ----------
     mu : np.ndarray
-        Mean values (must be positive).
+        Mean values (must be non-negative).
     concentration : float
         Concentration parameter (higher = less overdispersion).
     rng : np.random.Generator
@@ -256,7 +256,6 @@ def sample_negbinom(
     np.ndarray
         Integer counts.
     """
-    mu = np.maximum(mu, 1e-10)
     p = concentration / (concentration + mu)
     return rng.negative_binomial(n=concentration, p=p)
 

pyrenew/observation/negativebinomial.py

@@ -16,7 +16,6 @@ def __init__(
         self,
         name: str,
         concentration_rv: RandomVariable,
-        eps: float = 1e-10,
     ) -> None:
         """
         Default constructor
@@ -32,17 +31,13 @@ def __init__(
             despite the fact that larger values imply that the distribution
             becomes more Poissonian, while smaller ones imply a greater degree
             of dispersion.
-        eps
-            Small value to add to the predicted mean to prevent numerical
-            instability. Defaults to 1e-10.
 
         Returns
         -------
         None
         """
         super().__init__(name=name)
         self.concentration_rv = concentration_rv
-        self.eps = eps
 
     def sample(
         self,
@@ -71,7 +66,7 @@ def sample(
         negative_binomial_sample = numpyro.sample(
             name=self.name,
             fn=dist.NegativeBinomial2(
-                mean=mu + self.eps,
+                mean=mu,
                 concentration=concentration,
             ),
             obs=obs,

pyrenew/observation/noise.py

@@ -28,8 +28,6 @@
 
 from pyrenew.metaclass import RandomVariable
 
-_EPSILON = 1e-10
-
 
 class CountNoise(ABC):
     """
@@ -127,7 +125,7 @@ def sample(
         with numpyro.handlers.mask(mask=True if mask is None else mask):
             return numpyro.sample(
                 name,
-                dist.Poisson(rate=predicted + _EPSILON),
+                dist.Poisson(rate=predicted),
                 obs=obs,
             )
 
@@ -214,7 +212,7 @@ def sample(
             return numpyro.sample(
                 name,
                 dist.NegativeBinomial2(
-                    mean=predicted + _EPSILON,
+                    mean=predicted,
                     concentration=concentration,
                 ),
                 obs=obs,