|
| 1 | +package vecxt_re |
| 2 | + |
| 3 | +import org.apache.commons.rng.simple.RandomSource |
| 4 | +import io.github.quafadas.plots.SetupVega.{*, given} |
| 5 | +import io.circe.syntax.* |
| 6 | + |
| 7 | +/** Empirical distribution (JVM only). |
| 8 | + * |
| 9 | + * This is a nonparametric distribution built directly from observed samples. |
| 10 | + * It supports positive weights $w_i$. |
| 11 | + * |
| 12 | + * The distribution is represented as a discrete measure on the (possibly repeated) sample values: |
| 13 | + * $$\mathbb{P}(X = x) = \sum_{i: x_i = x} \frac{w_i}{\sum_k w_k}.$$ |
| 14 | + * |
| 15 | + * Consequently, the CDF is a right-continuous step function |
| 16 | + * $$F(t) = \mathbb{P}(X \le t) = \sum_{x \le t} \mathbb{P}(X=x).$$ |
| 17 | + * |
| 18 | + * Sampling is performed by inverse-transform sampling on the cumulative weights. |
| 19 | + * |
| 20 | + * Note: Since this is an empirical (atomic) distribution, a density/PDF in the usual continuous sense is not defined. |
| 21 | + * The `plot` method instead displays a (weighted) histogram density estimate. |
| 22 | + */ |
| 23 | +case class Empirical(values: IArray[Double], weights: IArray[Double]) |
| 24 | + extends DiscreteDistr[Double] |
| 25 | + with HasMean[Double] |
| 26 | + with HasVariance[Double] |
| 27 | + with HasCdf |
| 28 | + with HasInverseCdf: |
| 29 | + |
| 30 | + require(values.nonEmpty, "values must not be empty") |
| 31 | + require(values.forall(_.isFinite), "all values must be finite") |
| 32 | + require(weights.length == values.length, "weights must be the same length as values") |
| 33 | + require(weights.forall(w => w > 0 && w.isFinite), "weights must be positive and finite") |
| 34 | + |
| 35 | + private val rng = RandomSource.XO_RO_SHI_RO_128_PP.create() |
| 36 | + |
| 37 | + private val n = values.length |
| 38 | + |
| 39 | + private val pairs: Array[(Double, Double)] = |
| 40 | + val out = new Array[(Double, Double)](n) |
| 41 | + var j = 0 |
| 42 | + while j < n do |
| 43 | + out(j) = (values(j), weights(j)) |
| 44 | + j += 1 |
| 45 | + end while |
| 46 | + out |
| 47 | + |
| 48 | + scala.util.Sorting.stableSort(pairs, (a: (Double, Double), b: (Double, Double)) => a._1 < b._1) |
| 49 | + |
| 50 | + // Compress duplicates so we have unique support points. |
| 51 | + private val xsBuf = scala.collection.mutable.ArrayBuffer.empty[Double] |
| 52 | + private val wBuf = scala.collection.mutable.ArrayBuffer.empty[Double] |
| 53 | + |
| 54 | + private var totalWeight = 0.0 |
| 55 | + private var i = 0 |
| 56 | + while i < pairs.length do |
| 57 | + val x = pairs(i)._1 |
| 58 | + var wSum = 0.0 |
| 59 | + while i < pairs.length && pairs(i)._1 == x do |
| 60 | + val w = pairs(i)._2 |
| 61 | + wSum += w |
| 62 | + totalWeight += w |
| 63 | + i += 1 |
| 64 | + end while |
| 65 | + xsBuf += x |
| 66 | + wBuf += wSum |
| 67 | + end while |
| 68 | + |
| 69 | + private val xs: Array[Double] = xsBuf.toArray |
| 70 | + private val probs: Array[Double] = wBuf.toArray.map(_ / totalWeight) |
| 71 | + |
| 72 | + private val cdfVals: Array[Double] = |
| 73 | + val out = new Array[Double](probs.length) |
| 74 | + var acc = 0.0 |
| 75 | + var j = 0 |
| 76 | + while j < probs.length do |
| 77 | + acc += probs(j) |
| 78 | + out(j) = acc |
| 79 | + j += 1 |
| 80 | + end while |
| 81 | + out |
| 82 | + |
| 83 | + private val meanVal: Double = |
| 84 | + var s = 0.0 |
| 85 | + var j = 0 |
| 86 | + while j < xs.length do |
| 87 | + s += xs(j) * probs(j) |
| 88 | + j += 1 |
| 89 | + end while |
| 90 | + s |
| 91 | + |
| 92 | + private val varVal: Double = |
| 93 | + var s2 = 0.0 |
| 94 | + var j = 0 |
| 95 | + while j < xs.length do |
| 96 | + val d = xs(j) - meanVal |
| 97 | + s2 += probs(j) * d * d |
| 98 | + j += 1 |
| 99 | + end while |
| 100 | + s2 |
| 101 | + |
| 102 | + def mean: Double = meanVal |
| 103 | + |
| 104 | + def variance: Double = varVal |
| 105 | + |
| 106 | + /** Probability mass at exactly `x` (sums weights for duplicates). */ |
| 107 | + def probabilityOf(x: Double): Double = |
| 108 | + val idx = java.util.Arrays.binarySearch(xs, x) |
| 109 | + if idx >= 0 then probs(idx) else 0.0 |
| 110 | + |
| 111 | + /** Draw a sample using inverse CDF sampling over the atomic masses. */ |
| 112 | + def draw: Double = |
| 113 | + val u = rng.nextDouble() |
| 114 | + inverseCdf(u) |
| 115 | + |
| 116 | + /** CDF $F(t)=P(X\le t)$ (right-continuous). */ |
| 117 | + def cdf(x: Double): Double = |
| 118 | + if x < xs(0) then 0.0 |
| 119 | + else if x >= xs(xs.length - 1) then 1.0 |
| 120 | + else |
| 121 | + // Find the last index with xs(idx) <= x |
| 122 | + val ip = java.util.Arrays.binarySearch(xs, x) |
| 123 | + val idx = if ip >= 0 then ip else -ip - 2 |
| 124 | + cdfVals(idx) |
| 125 | + |
| 126 | + /** Probability that $x < X \le y$. */ |
| 127 | + def probability(x: Double, y: Double): Double = |
| 128 | + if y <= x then 0.0 |
| 129 | + else cdf(y) - cdf(x) |
| 130 | + |
| 131 | + /** Inverse CDF (quantile function): returns the smallest `x` with `F(x) >= p`. */ |
| 132 | + def inverseCdf(p: Double): Double = |
| 133 | + require(p >= 0.0 && p <= 1.0, "p must be in [0,1]") |
| 134 | + if p <= 0.0 then xs(0) |
| 135 | + else |
| 136 | + val ip = java.util.Arrays.binarySearch(cdfVals, p) |
| 137 | + val idx = if ip >= 0 then ip else -ip - 1 |
| 138 | + xs(math.min(idx, xs.length - 1)) |
| 139 | + |
| 140 | + /** Plot a (weighted) histogram density estimate. */ |
| 141 | + def plot(using viz.LowPriorityPlotTarget) = |
| 142 | + val plot = VegaPlot.fromResource("empiricalPdf.vl.json") |
| 143 | + val data = (0 until n).map { i => (x = values(i), w = weights(i)) } |
| 144 | + plot.plot( |
| 145 | + _.data.values := data.asJson, |
| 146 | + _ += (title = s"Empirical Distribution (n=$n)").asJson |
| 147 | + ) |
| 148 | + |
| 149 | + /** Plot the empirical CDF (step function). */ |
| 150 | + def plotCdf(using viz.LowPriorityPlotTarget) = |
| 151 | + val plot = VegaPlot.fromResource("empiricalCdf.vl.json") |
| 152 | + |
| 153 | + // Add an initial point at (min, 0) so the step is visible from the left. |
| 154 | + val points = |
| 155 | + val pts = scala.collection.mutable.ArrayBuffer.empty[(x: Double, cdf: Double)] |
| 156 | + pts += ((x = xs(0), cdf = 0.0)) |
| 157 | + var j = 0 |
| 158 | + while j < xs.length do |
| 159 | + pts += ((x = xs(j), cdf = cdfVals(j))) |
| 160 | + j += 1 |
| 161 | + end while |
| 162 | + pts.toVector |
| 163 | + |
| 164 | + plot.plot( |
| 165 | + _.data.values := points.asJson, |
| 166 | + _ += (title = s"Empirical CDF (n=$n)").asJson |
| 167 | + ) |
| 168 | + |
| 169 | +end Empirical |
| 170 | + |
| 171 | +object Empirical: |
| 172 | + /** Construct an unweighted empirical distribution (all weights equal to $1$). |
| 173 | + * |
| 174 | + * Note: We intentionally avoid `apply` overloads here because `IArray[Double]` erases to `Array[Double]` |
| 175 | + * on the JVM, which can create signature collisions with the case class companion methods. |
| 176 | + */ |
| 177 | + inline def equalWeights(values: Array[Double]): Empirical = |
| 178 | + Empirical( |
| 179 | + IArray.from(values), |
| 180 | + IArray.from(Array.fill(values.length)(1.0)) |
| 181 | + ) |
| 182 | + |
| 183 | + /** Construct a weighted empirical distribution from arrays. */ |
| 184 | + inline def weighted(values: Array[Double], weights: Array[Double]): Empirical = |
| 185 | + Empirical(IArray.from(values), IArray.from(weights)) |
| 186 | + |
0 commit comments