Restore NB2 zero-mean padding in NegativeBinomial observation paths

pyrenew/observation/negativebinomial.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import jax.numpy as jnp
 import numpyro
 import numpyro.distributions as dist
 from jax.typing import ArrayLike
@@ -61,12 +62,14 @@ def sample(
         -------
         ArrayLike
         """
+        # NB2 log_prob can be NaN at exact zero mean; pad by epsilon for stability.
+        padded_mean = jnp.asarray(mu) + jnp.finfo(float).eps
         concentration = self.concentration_rv.sample()
 
         negative_binomial_sample = numpyro.sample(
             name=self.name,
             fn=dist.NegativeBinomial2(
-                mean=mu,
+                mean=padded_mean,
                 concentration=concentration,
             ),
             obs=obs,

pyrenew/observation/noise.py

@@ -207,12 +207,14 @@ def sample(
         ArrayLike
             Negative Binomial-distributed counts.
         """
+        # NB2 log_prob can be NaN at exact zero mean; pad by epsilon for stability.
+        padded_mean = jnp.asarray(predicted) + jnp.finfo(float).eps
         concentration = self.concentration_rv()
         with numpyro.handlers.mask(mask=True if mask is None else mask):
             return numpyro.sample(
                 name,
                 dist.NegativeBinomial2(
-                    mean=predicted,
+                    mean=padded_mean,
                     concentration=concentration,
                 ),
                 obs=obs,

test/test_observation_counts.py

@@ -382,6 +382,21 @@ def test_negative_binomial_noise_validate_negative_concentration(self):
         with pytest.raises(ValueError, match="concentration must be positive"):
             noise.validate_concentration_rv()
 
+    def test_negative_binomial_noise_zero_mean_has_finite_log_prob(self):
+        """Test NegativeBinomialNoise yields finite log-probability at zero mean."""
+        noise = NegativeBinomialNoise(DeterministicVariable("conc", 10.0))
+        with numpyro.handlers.seed(rng_seed=223):
+            tr = numpyro.handlers.trace(
+                lambda: noise.sample(
+                    "noise_obs",
+                    predicted=jnp.array([0.0, 0.0]),
+                    obs=jnp.array([0, 0]),
+                )
+            ).get_trace()
+
+        log_prob = tr["noise_obs"]["fn"].log_prob(tr["noise_obs"]["value"])
+        assert jnp.all(jnp.isfinite(log_prob))
+
 
 class TestBaseObservationProcessValidation:
     """Test base observation process PMF validation."""

test/test_observation_negativebinom.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import numpy.testing as testing
+import jax.numpy as jnp
 import numpyro
 from jax.typing import ArrayLike
 
@@ -60,3 +61,19 @@ def test_negativebinom_random_obs():
     # Sample mean should be close to the expected rate (5.0)
     testing.assert_almost_equal(np.mean(sim_nb1), 5.0, decimal=0)
     testing.assert_almost_equal(np.mean(sim_nb2), 5.0, decimal=0)
+
+
+def test_negativebinom_zero_mean_has_finite_log_prob():
+    """Check that zero means do not produce NaN log-probability."""
+    negb = NegativeBinomialObservation(
+        "negbinom_rv",
+        concentration_rv=DeterministicVariable(name="concentration", value=10),
+    )
+
+    with numpyro.handlers.seed(rng_seed=223):
+        tr = numpyro.handlers.trace(
+            lambda: negb(mu=jnp.array([0.0, 0.0]), obs=jnp.array([0, 0]))
+        ).get_trace()
+
+    log_prob = tr["negbinom_rv"]["fn"].log_prob(tr["negbinom_rv"]["value"])
+    assert jnp.all(jnp.isfinite(log_prob))