fix continous poisson log prob

pfackeldey · May 10, 2024 · 23d4d72 · 23d4d72
1 parent b950146
commit 23d4d72
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/evermore/pdf.py b/src/evermore/pdf.py
@@ -5,6 +5,7 @@
 import equinox as eqx
 import jax
 import jax.numpy as jnp
+from jax.scipy.special import gammaln, xlogy
 from jaxtyping import Array, PRNGKeyArray
 
 __all__ = [
@@ -46,8 +47,11 @@ class Poisson(PDF):
     lamb: Array = eqx.field(converter=jnp.atleast_1d)
 
     def log_prob(self, x: Array) -> Array:
-        logpdf_max = jax.scipy.stats.poisson.logpmf(self.lamb, mu=self.lamb)
-        unnormalized = jax.scipy.stats.poisson.logpmf((x + 1) * self.lamb, mu=self.lamb)
+        def _continous_poisson_log_prob(x, lamb):
+            return xlogy(x, lamb) - lamb - gammaln(x + 1)
+
+        logpdf_max = _continous_poisson_log_prob(self.lamb, self.lamb)
+        unnormalized = _continous_poisson_log_prob((x + 1) * self.lamb, self.lamb)
         return unnormalized - logpdf_max
 
     def sample(self, key: PRNGKeyArray) -> Array: