Add truncated normal distribution

core/src/main/scala/dimwit/stats/IndependentDistributions.scala

@@ -34,6 +34,32 @@ object Normal:
   def standardSample(key: Random.Key): Tensor0[Float32] = new Normal(Tensor0(0f), Tensor0(1f)).sample(key)
   def standardNormal[T <: Tuple: Labels](shape: Shape[T]): Normal[T, Float32] = Normal.standardIsotropic(shape, scale = Tensor0(VType[Float32])(1f))
 
+/** Truncated Normal (Gaussian) distribution */
+class TruncatedNormal[T <: Tuple: Labels, V: IsFloating](val loc: Tensor[T, V], val scale: Tensor[T, V], val low: Tensor[T, V], val high: Tensor[T, V]) extends IndependentDistribution[T, V]:
+
+  private lazy val a = (low - loc) / scale
+  private lazy val b = (high - loc) / scale
+
+  override def elementWiseLogProb(x: Tensor[T, V]): Tensor[T, LogProb] =
+    liftPyTensor(jstats.truncnorm.logpdf(x.jaxValue, a = a.jaxValue, b = b.jaxValue, loc = loc.jaxValue, scale = scale.jaxValue))
+
+  override def sample(key: Random.Key): Tensor[T, V] =
+    val stdTruncated = liftPyTensor(loc.shape, VType[V])(
+      Jax.jrandom.truncated_normal(key.jaxKey, lower = a.jaxValue, upper = b.jaxValue, shape = loc.shape.dimensions.toPythonProxy)
+    )
+    stdTruncated * scale + loc
+
+object TruncatedNormal:
+
+  def apply[T <: Tuple: Labels, V: IsFloating](loc: Tensor[T, V], scale: Tensor[T, V], low: Tensor[T, V], high: Tensor[T, V]): TruncatedNormal[T, V] =
+    new TruncatedNormal(loc, scale, low, high)
+
+  def symmetric[T <: Tuple: Labels, V: IsFloating](loc: Tensor[T, V], scale: Tensor[T, V], numStd: Float = 2.0f): TruncatedNormal[T, V] =
+    val boundOffset = scale *! Tensor0(VType[V])(numStd)
+    val low = loc - boundOffset
+    val high = loc + boundOffset
+    new TruncatedNormal(loc, scale, low, high)
+
 /** Uniform distribution */
 class Uniform[T <: Tuple: Labels, V: IsFloating](val low: Tensor[T, V], val high: Tensor[T, V]) extends IndependentDistribution[T, V]:
 

core/src/test/scala/dimwit/stats/DistributionSuite.scala

@@ -40,6 +40,53 @@ class DistributionSuite extends DimwitTest:
       val expectedMeans = normal.loc
       sampleMeans should approxEqual(expectedMeans, 0.2f)
 
+  describe("TruncatedNormal Distribution"):
+    it("logProbs matches JAX"):
+      val loc = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 1.0f, -0.5f))
+      val scale = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 0.5f, 2.0f))
+      val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 1.5f, -1.0f))
+
+      val numStd = 2.0f
+      val dist = TruncatedNormal.symmetric(loc, scale, numStd = numStd)
+      val scalaLogProbs = dist.elementWiseLogProb(x)
+      val low = (dist.loc - numStd *! scale)
+      val high = (dist.loc + numStd *! scale)
+      val a = (low - loc) / scale
+      val b = (high - loc) / scale
+      val jaxLogProbs = liftPyTensor1(Axis[A], VType[Float32])(
+        jstats.truncnorm.logpdf(
+          x.jaxValue,
+          a = a.jaxValue,
+          b = b.jaxValue,
+          loc = loc.jaxValue,
+          scale = scale.jaxValue
+        )
+      )
+      scalaLogProbs.asFloat should approxEqual(jaxLogProbs)
+
+    it("sample means approximates means"):
+      val loc = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 1.0f))
+      val scale = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 0.5f))
+      val dist = TruncatedNormal.symmetric(loc, scale, numStd = 2.0f)
+
+      val key = Random.Key(42)
+      val samples = key.splitvmap(Axis[Samples] -> 10000)(k => dist.sample(k))
+      val sampleMeans = samples.mean(Axis[Samples])
+      val expectedMeans = dist.loc
+      sampleMeans should approxEqual(expectedMeans, 0.2f)
+
+    it("samples strictly respect bounds"):
+      val loc = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 1.0f))
+      val scale = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 0.5f))
+      // Use a tight boundary (e.g. 1 standard deviation) to heavily force truncation
+      val dist = TruncatedNormal.symmetric(loc, scale, numStd = 1.0f)
+
+      val key = Random.Key(42)
+      val samples = key.splitvmap(Axis[Samples] -> 10000)(k => dist.sample(k))
+
+      (samples.min(Axis[Samples]) >= dist.low).all.item shouldBe true
+      (samples.max(Axis[Samples]) <= dist.high).all.item shouldBe true
+
   describe("Uniform Distribution"):
     it("logProbs matches JAX"):
       val low = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, -1.0f, 2.0f))