diff --git a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala index 975ed4a7..39d1b558 100644 --- a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala +++ b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala @@ -175,3 +175,29 @@ object Beta: def apply[T <: Tuple: Labels, V: IsFloating](alpha: Tensor[T, V], beta: Tensor[T, V]): Beta[T, V] = require(alpha.shape.dimensions == beta.shape.dimensions, "alpha and beta must have the same dimensions") new Beta(alpha, beta) + +/** Exponential distribution */ +class Exponential[T <: Tuple: Labels, V: IsFloating](val rate: Tensor[T, V]) extends IndependentDistribution[T, V]: + + lazy val scale = (Tensor0(1f).asType(VType[V]) /! rate).jaxValue + + override def elementWiseLogProb(x: Tensor[T, V]): Tensor[T, LogProb] = + liftPyTensor(jstats.expon.logpdf(x.jaxValue, scale = scale)) + + override def sample(k: Random.Key): Tensor[T, V] = + // samples from jrandom.exponential are always with scale=1, + // so we need to divide by the rate to get the correct distribution + val stndExp = liftPyTensor[T, V]( + Jax.jrandom.exponential(k.jaxKey, shape = rate.shape.dimensions.toPythonProxy) + ) + stndExp / rate + +class Poisson[T <: Tuple: Labels, V: IsInteger](val rate: Tensor[T, V]) extends IndependentDistribution[T, V]: + + override def elementWiseLogProb(x: Tensor[T, V]): Tensor[T, LogProb] = + liftPyTensor(jstats.poisson.logpmf(x.jaxValue, mu = rate.jaxValue)) + + override def sample(k: Random.Key): Tensor[T, V] = + liftPyTensor( + Jax.jrandom.poisson(k.jaxKey, lam = rate.jaxValue, shape = rate.shape.dimensions.toPythonProxy) + ) diff --git a/core/src/test/scala/dimwit/stats/DistributionSuite.scala b/core/src/test/scala/dimwit/stats/DistributionSuite.scala index 66095632..adb84026 100644 --- a/core/src/test/scala/dimwit/stats/DistributionSuite.scala +++ b/core/src/test/scala/dimwit/stats/DistributionSuite.scala @@ -335,3 +335,50 @@ class DistributionSuite extends DimwitTest: // Mean of Beta distribution is alpha / (alpha + beta) val expectedMeans = betaDist.alpha / (betaDist.alpha + betaDist.beta) sampleMeans should approxEqual(expectedMeans, 0.02f) + + describe("Exponential"): + it("logProbs matches JAX"): + val rate = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 1.0f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 1.0f, 0.3f)) + + val dist = Exponential(rate) + val scalaLogProbs = dist.elementWiseLogProb(x) + val scale = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(2.0f, 1.0f, 0.5f)) + val jaxLogProbs = liftPyTensor1(Axis[A], VType[Float32])( + jstats.expon.logpdf(x.jaxValue, scale = scale.jaxValue) + ) + scalaLogProbs.asFloat should approxEqual(jaxLogProbs) + + it("sample means approximates 1/rate"): + val exponential = Exponential( + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.5f, 2.0f)) + ) + val key = Random.Key(42) + val samples = key.splitvmap(Axis[Samples] -> 10000)(k => exponential.sample(k)) + val sampleMeans = samples.mean(Axis[Samples]) + // Mean of Exponential(rate) is 1/rate + val expectedMeans = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(2.0f, 0.5f)) + sampleMeans should approxEqual(expectedMeans, 0.1f) + + describe("Poisson"): + it("logProbs matches JAX"): + val rate = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1, 3, 10)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1, 2, 8)) + + val dist = Poisson(rate) + val scalaLogProbs = dist.elementWiseLogProb(x) + val jaxLogProbs = liftPyTensor1(Axis[A], VType[Float32])( + jstats.poisson.logpmf(x.jaxValue, mu = rate.jaxValue) + ) + scalaLogProbs.asFloat should approxEqual(jaxLogProbs) + + it("sample means approximates rate"): + val poisson = Poisson( + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1, 5)) + ) + val key = Random.Key(42) + val samples = key.splitvmap(Axis[Samples] -> 10000)(k => poisson.sample(k)) + val sampleMeans = samples.asFloat32.mean(Axis[Samples]) + // Mean of Poisson(lambda) is lambda + val expectedMeans = poisson.rate.asFloat32 + sampleMeans should approxEqual(expectedMeans, 0.2f)