feat: initial release

Assisted-by: GLM 5.3 Flash

stats/distrib.go

@@ -0,0 +1,182 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package stats
+
+import (
+	"sourcedock.dev/petrbalvin/tensor/internal/base"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+)
+
+import (
+	"math"
+)
+
+// ExponentialDraws returns n draws from the exponential distribution
+// with the given rate (mean 1/rate), by inverse CDF.
+func ExponentialDraws(g *core.Generator, n int, rate float64) (*core.Array, error) {
+	if n < 1 {
+		return nil, base.Errf("ExponentialDraws: n must be ≥ 1")
+	}
+	if !(rate > 0) {
+		return nil, base.Errf("ExponentialDraws: rate must be positive, got %v", rate)
+	}
+	out := core.New(core.Float, []int{n}...)
+	for i := range n {
+		u := 1 - g.Unit()
+		out.RawFloats()[i] = -math.Log(u) / rate
+	}
+	return out, nil
+}
+
+// GammaDraws returns n draws from the gamma distribution with shape
+// α > 0 and rate β > 0, by Marsaglia-Tsang for α ≥ 1 and by boosting
+// with the exponential for α < 1.
+func GammaDraws(g *core.Generator, n int, alpha, beta float64) (*core.Array, error) {
+	if n < 1 {
+		return nil, base.Errf("GammaDraws: n must be ≥ 1")
+	}
+	if !(alpha > 0 && beta > 0) {
+		return nil, base.Errf("GammaDraws: shape and rate must be positive")
+	}
+	out := core.New(core.Float, []int{n}...)
+	d := alpha - 1.0/3
+	c := 1 / math.Sqrt(9*d)
+	for i := range n {
+		if alpha >= 1 {
+			out.RawFloats()[i] = gammaMarsagliaTsang(g, alpha, d, c) / beta
+			continue
+		}
+		// α < 1: boost to α + 1 and scale by a uniform^(1/α).
+		boost := alpha + 1
+		dd := boost - 1.0/3
+		cc := 1 / math.Sqrt(9*dd)
+		v := gammaMarsagliaTsang(g, boost, dd, cc)
+		out.RawFloats()[i] = v * math.Pow(g.Unit(), 1/alpha) / beta
+	}
+	return out, nil
+}
+
+// gammaMarsagliaTsang draws one gamma(α, 1) for α ≥ 1 by the
+// Marsaglia-Tsang squeeze: a normal draw shapes the cube root of the
+// scale, an exponential-tilted accept/reject polishes the tail.
+func gammaMarsagliaTsang(g *core.Generator, alpha, d, c float64) float64 {
+	for {
+		x := g.NormalUnit()
+		v := 1 + c*x
+		if v <= 0 {
+			continue
+		}
+		vvv := v * v * v
+		u := g.Unit()
+		if u < 1-0.0331*x*x*x*x {
+			return d * vvv
+		}
+		if math.Log(u) < 0.5*x*x+d*(1-vvv+math.Log(vvv)) {
+			return d * vvv
+		}
+	}
+}
+
+// ChiSquareDraws returns n draws from the chi-squared distribution
+// with df degrees of freedom, the gamma(df/2, 2) distribution.
+func ChiSquareDraws(g *core.Generator, n int, df int) (*core.Array, error) {
+	if df < 1 {
+		return nil, base.Errf("ChiSquareDraws: df must be ≥ 1")
+	}
+	return GammaDraws(g, n, float64(df)/2, 0.5)
+}
+
+// StudentTDraws returns n draws from Student's t with df degrees of
+// freedom, as N(0,1)/√(χ²_df/df).
+func StudentTDraws(g *core.Generator, n int, df int) (*core.Array, error) {
+	if df < 1 {
+		return nil, base.Errf("StudentTDraws: df must be ≥ 1")
+	}
+	chi, err := ChiSquareDraws(g, n, df)
+	if err != nil {
+		return nil, err
+	}
+	out := core.New(core.Float, []int{n}...)
+	for i := range n {
+		z := g.NormalUnit()
+		den := math.Sqrt(chi.FloatAt(i) / float64(df))
+		if den == 0 {
+			// The χ² draw underflowed to exactly zero: the ratio is an
+			// infinity carrying the numerator's sign, not an unsigned
+			// one (and not the NaN a 0/0 numerator of zero would make).
+			out.RawFloats()[i] = math.Copysign(math.Inf(1), z)
+		} else {
+			out.RawFloats()[i] = z / den
+		}
+	}
+	return out, nil
+}
+
+// PoissonDraws returns n draws from the Poisson distribution with the
+// given λ, by the Knuth multiplication method for λ < 30 and the
+// normal approximation above.
+func PoissonDraws(g *core.Generator, n int, lambda float64) (*core.Array, error) {
+	if n < 1 {
+		return nil, base.Errf("PoissonDraws: n must be ≥ 1")
+	}
+	if !(lambda >= 0) {
+		return nil, base.Errf("PoissonDraws: λ must be ≥ 0")
+	}
+	out := core.New(core.Float, []int{n}...)
+	L := math.Exp(-lambda)
+	for i := range n {
+		if lambda < 30 {
+			// λ = 0 (or tiny enough that exp(−λ) rounds to 1) is the
+			// degenerate distribution at 0: the multiplication loop
+			// would never run and hand back k−1 = −1.
+			if L >= 1 {
+				out.RawFloats()[i] = 0
+				continue
+			}
+			k := 0.0
+			p := 1.0
+			for p > L {
+				k++
+				p *= g.Unit()
+			}
+			out.RawFloats()[i] = k - 1
+		} else {
+			// Normal approximation for large λ.
+			z := g.NormalUnit()
+			out.RawFloats()[i] = max(0, math.Floor(lambda+math.Sqrt(lambda)*z+0.5))
+		}
+	}
+	return out, nil
+}
+
+// BinomialDraws returns n draws from the binomial distribution with
+// the given number of trials and success probability: each draw runs
+// trials uniform comparisons against p and counts the successes, the
+// exact per-trial loop. It is exact but costs O(trials) uniforms per
+// draw, so keep trials modest.
+func BinomialDraws(g *core.Generator, n int, trials int, p float64) (*core.Array, error) {
+	if n < 1 {
+		// The sibling draws all refuse n < 1; without the guard this one
+		// returned a nil array with a nil error, which a caller that only
+		// checks the error then dereferenced.
+		return nil, base.Errf("BinomialDraws: n must be ≥ 1")
+	}
+	if trials < 1 {
+		return nil, base.Errf("BinomialDraws: trials must be ≥ 1")
+	}
+	if !(p >= 0 && p <= 1) {
+		return nil, base.Errf("BinomialDraws: p must be in [0, 1]")
+	}
+	out := core.New(core.Float, []int{n}...)
+	for i := range n {
+		count := 0.0
+		for range trials {
+			if g.Unit() < p {
+				count++
+			}
+		}
+		out.RawFloats()[i] = count
+	}
+	return out, nil
+}