add halfcauchy distribution

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

@@ -116,6 +116,30 @@ object Cauchy:
     require(loc.shape.dimensions == scale.shape.dimensions, "Location and scale must have the same dimensions")
     new Cauchy(loc, scale)
 
+/** Half-Cauchy distribution - absolute value of a Cauchy-distributed random variable */
+class HalfCauchy[T <: Tuple: Labels, V: IsFloating](val loc: Tensor[T, V], val scale: Tensor[T, V]) extends IndependentDistribution[T, V]:
+
+  override def elementWiseLogProb(x: Tensor[T, V]): Tensor[T, LogProb] =
+    // Half-Cauchy logpdf = log(2) + cauchy.logpdf for x >= loc, -inf otherwise
+    val rawLogProb = liftPyTensor(x.shape, VType[LogProb])(
+      Jax.jnp.log(2.0) + jstats.cauchy.logpdf(x.jaxValue, loc = loc.jaxValue, scale = scale.jaxValue)
+    )
+    val valid = x >= loc
+    val negInf = LogProb(Tensor.like(x).fill(Float.NegativeInfinity).asFloat32)
+    where(valid, rawLogProb, negInf)
+
+  override def sample(k: Random.Key): Tensor[T, V] =
+    // Half-Cauchy: |Cauchy(0, 1)| * scale + loc
+    val cauchy = liftPyTensor(loc.shape, VType[V])(Jax.jrandom.cauchy(k.jaxKey, shape = loc.shape.dimensions.toPythonProxy))
+    cauchy.abs * scale + loc
+
+object HalfCauchy:
+
+  /** Create a HalfCauchy distribution from location and scale tensors */
+  def apply[T <: Tuple: Labels, V: IsFloating](loc: Tensor[T, V], scale: Tensor[T, V]): HalfCauchy[T, V] =
+    require(loc.shape.dimensions == scale.shape.dimensions, "Location and scale must have the same dimensions")
+    new HalfCauchy(loc, scale)
+
 /** Half-normal distribution */
 class HalfNormal[T <: Tuple: Labels, V: IsFloating](val loc: Tensor[T, V], val scale: Tensor[T, V]) extends IndependentDistribution[T, V]:
 

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

@@ -156,6 +156,45 @@ class DistributionSuite extends DimwitTest:
       val expectedMedian = cauchy.loc
       sampleMedian should approxEqual(expectedMedian, 0.5f)
 
+  describe("HalfCauchy"):
+    it("logProbs matches scipy"):
+      val scipy_stats = py.module("scipy.stats")
+      val loc = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 0.0f, 0.0f))
+      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.0f, 0.8f))
+
+      val dist = HalfCauchy(loc, scale)
+      val scalaLogProbs = dist.elementWiseLogProb(x)
+      val scipyLogProbs = liftPyTensor1(Axis[A], VType[Float32])(
+        scipy_stats.halfcauchy.logpdf(x.jaxValue, loc = loc.jaxValue, scale = scale.jaxValue)
+      )
+      scalaLogProbs.asFloat should approxEqual(scipyLogProbs)
+
+    it("logProb is -inf for x < loc"):
+      val loc = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 1.0f, 1.0f))
+      val scale = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 1.0f, 1.0f))
+      val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 0.5f, 0.9f))
+
+      val dist = HalfCauchy(loc, scale)
+      val logProbs = dist.elementWiseLogProb(x)
+      logProbs.asFloat.toArray.foreach(v => v should be(Float.NegativeInfinity))
+
+    it("sample medians approximates location"):
+      val halfCauchy = HalfCauchy(
+        Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 0.0f)),
+        Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 2.0f))
+      )
+      val key = Random.Key(42)
+      val samples = key.splitvmap(Axis[Samples] -> 50000)(k => halfCauchy.sample(k))
+      // All samples should be >= loc
+      val minSample = samples.min(Axis[Samples])
+      minSample.toArray.zip(halfCauchy.loc.toArray).foreach { (s, l) =>
+        s should be >= l
+      }
+      // Median of HalfCauchy(0, scale) is scale
+      val sampleMedian = samples.median(Axis[Samples])
+      sampleMedian should approxEqual(halfCauchy.scale, 0.2f)
+
   describe("HalfNormal"):
     it("logProbs computed correctly"):
       val loc = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 0.0f, 0.0f))