211 lines
7.0 KiB
Go
211 lines
7.0 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package stats
|
||
|
||
import (
|
||
"math"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// The multivariate normal: densities through a Cholesky solve and
|
||
// draws through the same triangular factor, the two operations any
|
||
// Bayesian or Monte Carlo workflow reaches for.
|
||
|
||
// mvnSymmetryEps is the relative tolerance the covariance's mirror
|
||
// check allows. A covariance assembled as A·Aᵀ can differ from its
|
||
// mirror by an ulp of rounding; a matrix that is genuinely asymmetric
|
||
// differs by far more than this, and reading only its lower triangle
|
||
// would describe a different distribution from the one handed in.
|
||
const mvnSymmetryEps = 1e-12
|
||
|
||
// mvnCholesky factors a symmetric positive-definite matrix into the
|
||
// lower triangular L with A = L·Lᵀ; a non-positive pivot names the row
|
||
// in the error, and a mirror pair that disagrees beyond a relative
|
||
// mvnSymmetryEps names the entry, because only the lower triangle is
|
||
// read.
|
||
func mvnCholesky(name string, cov *core.Array, d int) ([][]float64, error) {
|
||
// The covariance is read once into a flat local copy: the symmetry
|
||
// sweep and the O(d³) factorisation then index a plain slice. The
|
||
// copy carries the values FloatAt returns, so every difference,
|
||
// product and pivot below keeps its exact bits.
|
||
covVals := make([]float64, d*d)
|
||
if fs := rawFloats(cov); fs != nil {
|
||
copy(covVals, fs)
|
||
} else {
|
||
for i := range covVals {
|
||
covVals[i] = cov.FloatAt(i)
|
||
}
|
||
}
|
||
return mvnCholeskyFlat(name, covVals, d)
|
||
}
|
||
|
||
// mvnCholeskyFlat factors a matrix already held flat row-major, the
|
||
// shared body of the array entry point and of a caller that assembles
|
||
// the matrix in place. The slice is read, never written.
|
||
func mvnCholeskyFlat(name string, covVals []float64, d int) ([][]float64, error) {
|
||
l := make([][]float64, d)
|
||
for i := range d {
|
||
l[i] = make([]float64, d)
|
||
}
|
||
if err := mvnCholeskyFlatInto(name, covVals, l, d); err != nil {
|
||
return nil, err
|
||
}
|
||
return l, nil
|
||
}
|
||
|
||
// mvnCholeskyFlatInto is mvnCholeskyFlat on a caller-owned d×d
|
||
// destination, for a sweep that factors one covariance after another:
|
||
// the rows are cleared and refilled, and every entry the factorisation
|
||
// reads is one it has already written in the same pass. The arithmetic
|
||
// is mvnCholeskyFlat's own, unchanged.
|
||
func mvnCholeskyFlatInto(name string, covVals []float64, l [][]float64, d int) error {
|
||
for i := range d {
|
||
for j := range i {
|
||
lo, hi := covVals[i*d+j], covVals[j*d+i]
|
||
if math.Abs(lo-hi) > mvnSymmetryEps*math.Max(math.Abs(lo), math.Abs(hi)) {
|
||
return base.Errf("%s: the covariance is not symmetric at (%d, %d): %g against %g",
|
||
name, i+1, j+1, lo, hi)
|
||
}
|
||
}
|
||
}
|
||
for i := range d {
|
||
row := l[i]
|
||
clear(row)
|
||
for j := range i + 1 {
|
||
total := covVals[i*d+j]
|
||
for k := range j {
|
||
total -= l[i][k] * l[j][k]
|
||
}
|
||
if i == j {
|
||
if total <= 0 {
|
||
return base.Errf("%s: the covariance is not positive definite at row %d", name, i+1)
|
||
}
|
||
row[j] = math.Sqrt(total)
|
||
} else {
|
||
row[j] = total / l[j][j]
|
||
}
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// MultivariateNormalLogDensity evaluates the density of the
|
||
// d-dimensional normal N(mean, cov) at one point x, all three rank-1
|
||
// and cov symmetric positive definite, through the Cholesky factor:
|
||
// the log determinant is twice the sum of the factor's diagonal and
|
||
// the quadratic form the squared norm of the forward solve. Every
|
||
// entry of all three inputs must be finite, and the covariance must
|
||
// mirror itself: a non-finite value or an asymmetric pair is refused
|
||
// by name rather than answered with a NaN or with the density of a
|
||
// different distribution.
|
||
func MultivariateNormalLogDensity(mean, cov, x *core.Array) (float64, error) {
|
||
const name = "MultivariateNormalLogDensity"
|
||
if mean.NDim() != 1 || x.NDim() != 1 || cov.NDim() != 2 {
|
||
return 0, base.Errf("%s: mean and x must be rank 1 and cov rank 2", name)
|
||
}
|
||
if mean.Dtype() == core.Complex || cov.Dtype() == core.Complex || x.Dtype() == core.Complex {
|
||
return 0, base.Errf("%s: complex inputs are not supported", name)
|
||
}
|
||
d := mean.Len()
|
||
if x.Len() != d || cov.Shape()[0] != d || cov.Shape()[1] != d {
|
||
return 0, base.Errf("%s: mean holds %d entries, x %d, cov %s", name, d, x.Len(), base.ShapeText(cov.Shape()))
|
||
}
|
||
if err := checkFinite(name, "the mean", mean); err != nil {
|
||
return 0, err
|
||
}
|
||
if err := checkFinite(name, "the covariance", cov); err != nil {
|
||
return 0, err
|
||
}
|
||
if err := checkFinite(name, "the point", x); err != nil {
|
||
return 0, err
|
||
}
|
||
l, err := mvnCholesky(name, cov, d)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
diff := make([]float64, d)
|
||
solve := make([]float64, d)
|
||
logDet := 0.0
|
||
for i := range d {
|
||
diff[i] = x.FloatAt(i) - mean.FloatAt(i)
|
||
logDet += math.Log(l[i][i])
|
||
total := diff[i]
|
||
for j := range i {
|
||
total -= l[i][j] * solve[j]
|
||
}
|
||
solve[i] = total / l[i][i]
|
||
}
|
||
quad := 0.0
|
||
for i := range d {
|
||
quad += solve[i] * solve[i]
|
||
}
|
||
return -0.5*float64(d)*math.Log(2*math.Pi) - logDet - 0.5*quad, nil
|
||
}
|
||
|
||
// MultivariateNormalDraws returns n draws from N(mean, cov) as an
|
||
// (n × d) array, one draw per row: standard normals through the
|
||
// generator, coloured by the Cholesky factor of the covariance. The
|
||
// draws are deterministic for a given generator state. The mean and
|
||
// the covariance must be finite and the covariance symmetric, the
|
||
// same contract MultivariateNormalLogDensity enforces.
|
||
func MultivariateNormalDraws(g *core.Generator, n int, mean, cov *core.Array) (*core.Array, error) {
|
||
const name = "MultivariateNormalDraws"
|
||
if g == nil {
|
||
return nil, base.Errf("%s: the generator is nil", name)
|
||
}
|
||
if n < 1 {
|
||
return nil, base.Errf("%s: the count must be at least 1, got %d", name, n)
|
||
}
|
||
if mean.NDim() != 1 || cov.NDim() != 2 {
|
||
return nil, base.Errf("%s: mean must be rank 1 and cov rank 2", name)
|
||
}
|
||
if mean.Dtype() == core.Complex || cov.Dtype() == core.Complex {
|
||
return nil, base.Errf("%s: complex inputs are not supported", name)
|
||
}
|
||
d := mean.Len()
|
||
if cov.Shape()[0] != d || cov.Shape()[1] != d {
|
||
return nil, base.Errf("%s: mean holds %d entries but cov is %s", name, d, base.ShapeText(cov.Shape()))
|
||
}
|
||
if err := checkFinite(name, "the mean", mean); err != nil {
|
||
return nil, err
|
||
}
|
||
if err := checkFinite(name, "the covariance", cov); err != nil {
|
||
return nil, err
|
||
}
|
||
l, err := mvnCholesky(name, cov, d)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
out := core.New(core.Float, n, d)
|
||
vals := out.RawFloats()
|
||
z := make([]float64, d)
|
||
// The mean is read once: every draw reuses the same values the
|
||
// accessor walk returned.
|
||
meanVals := make([]float64, d)
|
||
if fs := rawFloats(mean); fs != nil {
|
||
copy(meanVals, fs)
|
||
} else {
|
||
for i := range meanVals {
|
||
meanVals[i] = mean.FloatAt(i)
|
||
}
|
||
}
|
||
for r := range n {
|
||
// One standard normal vector per draw, coloured by L: the
|
||
// shared z is what makes the off-diagonal covariance appear.
|
||
for j := range d {
|
||
z[j] = g.NormalUnit()
|
||
}
|
||
for i := range d {
|
||
total := meanVals[i]
|
||
for j := range i + 1 {
|
||
total += l[i][j] * z[j]
|
||
}
|
||
vals[r*d+i] = total
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|