From d8597802b15e84634c1af823a27b3ad63fa2996a Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sun, 19 Jul 2026 13:31:56 +0200 Subject: [PATCH 1/2] add halfcauchy distribution --- .../stats/IndependentDistributions.scala | 26 +++++++++++++ .../dimwit/stats/DistributionSuite.scala | 39 +++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala index 39d1b55..1009e64 100644 --- a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala +++ b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala @@ -116,6 +116,32 @@ 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]: diff --git a/core/src/test/scala/dimwit/stats/DistributionSuite.scala b/core/src/test/scala/dimwit/stats/DistributionSuite.scala index adb8402..ab0de81 100644 --- a/core/src/test/scala/dimwit/stats/DistributionSuite.scala +++ b/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)) From 623b1078647f645133a1b42cd1a9266621dd4caf Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Mon, 20 Jul 2026 10:26:33 +0200 Subject: [PATCH 2/2] fix formatting --- core/src/main/scala/dimwit/stats/IndependentDistributions.scala | 2 -- 1 file changed, 2 deletions(-) diff --git a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala index 1009e64..2132c49 100644 --- a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala +++ b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala @@ -116,7 +116,6 @@ 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]: @@ -141,7 +140,6 @@ object HalfCauchy: 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]: