mixle.stats.multivariate.diagonal_gaussian module

Diagonal Gaussian distributions, estimators, accumulators, and encoders.

The log-density of an n-dimensional diagonal Gaussian observation x = (x_1, x_2, ..., x_n) with mean mu and diagonal covariance covar = diag(s2_1, s2_2, ..., s2_n) is:

log(p_mat(x)) = -0.5*sum_{i=1}^{n} (x_i-m_i)^2 / s2_i - 0.5*log(s2_i) - (n/2)*log(2*pi).

Reference: Mardia, Kent & Bibby, Multivariate Analysis (Academic Press, 1979).

class DiagonalGaussianSuffStat(sum_, sum2_, count_, anchored=None)

Bases: tuple

A (sum, sum_squares, count) sufficient statistic that also carries shift-anchored moments.

Behaves exactly like the plain 3-tuple everywhere it is indexed, unpacked, or iterated (it is one), so generic consumers – scale_suff_stat, the declared StatisticSpec reader, engine kernels – see nothing new. anchored is extra payload: (anchor, sum_i w_i*(x_i - anchor), sum_i w_i*(x_i - anchor)^2) per coordinate, which DiagonalGaussianAccumulator.combine() folds in and DiagonalGaussianEstimator.estimate() uses to compute shift-invariant variances. Mirrors mixle.stats.univariate.continuous.gaussian.GaussianSuffStat.

Parameters:

• sum_ (ndarray)

• sum2_ (ndarray)

• count_ (float)

• anchored (tuple | None)

class DiagonalGaussianFisherView(dist)

Bases: FixedFisherView

Fisher view over per-dimension first and second moments for a diagonal Gaussian.

Parameters:

dist (Any)

class DiagonalGaussianDistribution(mu, covar=MISSING, name=None, keys=None, covariance=MISSING, prior=None)

Bases: SequenceEncodableProbabilityDistribution

Multivariate Gaussian distribution with independent components (diagonal covariance matrix).

Parameters:

• mu (Sequence[float] | ndarray)

• covar (Sequence[float] | ndarray)

• name (str | None)

• keys (str | None)

• covariance (Sequence[float] | ndarray)

• prior (SequenceEncodableProbabilityDistribution | None)

classmethod compute_capabilities()

Declare backend support for diagonal Gaussian generated kernels.

classmethod compute_declaration()

Return the generated-compute declaration for the diagonal Gaussian.

static exp_family_sufficient_statistics(x, engine)

Return vector sufficient statistics for generated diagonal-Gaussian scoring.

Parameters:

• x (Any)

• engine (Any)

Return type:

tuple[Any, …]

static exp_family_natural_parameters(params, engine)

Return vector natural parameters for generated diagonal-Gaussian scoring.

Parameters:

• params (dict[str, Any])

• engine (Any)

Return type:

tuple[Any, …]

static exp_family_log_partition(params, engine)

Return the diagonal-Gaussian log partition for generated scoring.

Parameters:

• params (dict[str, Any])

• engine (Any)

Return type:

Any

static backend_legacy_sufficient_statistics(x, params, engine)

Return row-wise legacy accumulator statistics for generated resident reductions.

Parameters:

• x (Any)

• params (dict[str, Any])

• engine (Any)

Return type:

tuple[Any, …]

set_prior(prior)

Attach a parameter prior and precompute conjugate-prior expectations.

With a MultivariateNormalGamma(mu0, lam, a, b) prior over (mu, tau=1/covar) this caches the expected natural parameters [ea, eb, e1, e2] with e1 = E[mu*tau] and e2 = -0.5*E[tau] per component (ea, eb scalars summed over components), so that expected_log_density(x) = x.e1 + (x*x).e2 - ea + eb. Any other prior (including None) leaves the distribution a plain point model.

Parameters:

prior (SequenceEncodableProbabilityDistribution | None)

Return type:

None

expected_log_density(x)

Variational expectation E_q[log p(x | mu, tau)] under the prior.

Falls back to the plug-in log_density(x) when no conjugate prior is attached.

Return type:

float

seq_expected_log_density(x)

Vectorized expected_log_density over sequence-encoded observations.

Parameters:

x (ndarray)

Return type:

ndarray

density(x)

Evaluate the density at observation x.

See log_density() for details.

Parameters:

x (Union[Sequence[float], np.ndarray]) – Length-dim observation vector.

Returns:

Density at x.

log_density(x)

Evaluate the log-density at observation x.

The log-density is given by

log(p(x)) = -0.5*sum_{i=1}^{n} (x_i-m_i)^2 / s2_i - 0.5*log(s2_i) - (n/2)*log(2*pi).

Computed in the CENTERED form the formula above states, -0.5*sum (x-mu)^2/covar, not in the algebraically equal expanded form x^2*ca + x*cb + cc. The expanded form is three terms of size O(|x|^2/covar) summing to an O(1) answer, so at |x| ~ 1e9 with unit variance it cancels away every significant digit: the scorer returned round-number garbage (measured: log-likelihoods of exactly 66560.0 and -135168.0 on unit-variance data at offset 1.7e9), and because EM selects the model with the best objective, an optimize/fit run at that offset returned the WORSE of its own iterates. x - mu is exact for operands of the same magnitude, so the centered form is accurate at any offset; it is also one gemv instead of two. The full-covariance estimator’s scorer always used the centered form – this is the diagonal one catching up, and it makes the documented formula true (T1-F1). The expanded coefficients remain available as ca/cb/cc for the fused kernels.

Parameters:

x (Union[Sequence[float], np.ndarray]) – Length-dim observation vector.

Returns:

Log-density at x.

condition(observed)

Return the conditional over the unobserved dimensions given observed.

A diagonal Gaussian has independent coordinates, so conditioning on some of them leaves the rest unchanged: the result is just DiagonalGaussian(mu[unobserved], covar[unobserved]) (the observed values do not shift the unobserved mean or variance). Provided so diagonal-covariance Gaussian mixtures support MixtureDistribution.conditional() – there the responsibilities still update from how well each component explains the observed coordinates, even though the within-component coordinates are independent. Raises if no dimension is left unobserved.

Parameters:

observed (dict[int, float])

Return type:

DiagonalGaussianDistribution

marginal(keep)

Return the marginal over the dimensions keep: DiagonalGaussian(mu[keep], covar[keep]).

Marginalizing a diagonal Gaussian simply drops the other independent coordinates (order kept).

Parameters:

keep (Sequence[int])

Return type:

DiagonalGaussianDistribution

density_cumulative(x)

Exact probability-ordered cumulative G(x) = P(p(Y) >= p(x)) – the highest-density-region mass through x (multivariate analogue of a CDF). For a diagonal Gaussian the squared Mahalanobis distance sum_i (x_i-mu_i)^2/var_i is chi-square(dim), so G = chi2.cdf(maha2, dim). Used by mixle.enumeration.density_rank.density_rank() to return an EXACT cumulative.

Parameters:

x (Sequence[float] | ndarray)

Return type:

float

density_quantile(q)

Inverse of density_cumulative(): a representative point at cumulative-density index q.

q is the highest-density-region mass, whose boundary is the squared-Mahalanobis level chi2.ppf(q, dim); a representative point on that contour offsets the first coordinate by sqrt(level * var_0) (Mahalanobis distance exactly the level). Sweeping q enumerates the support in descending density.

Parameters:

q (float)

Return type:

ndarray

seq_log_density(x)

Vectorized evaluation of the log-density at a sequence-encoded input x.

Parameters:

x (np.ndarray) – Encoded data matrix with shape (sz, dim) from DiagonalGaussianDataEncoder.seq_encode().

Returns:

Numpy array of length sz containing the log-density of each encoded observation.

Return type:

ndarray

static backend_log_density_from_params(x, mu, covar, engine)

Engine-neutral diagonal Gaussian log-density from explicit parameters.

Parameters:

• x (Any)

• mu (Any)

• covar (Any)

• engine (Any)

Return type:

Any

backend_seq_log_density(x, engine)

Engine-neutral vectorized log-density for encoded data.

Parameters:

• x (Any)

• engine (Any)

Return type:

Any

classmethod backend_stacked_params(dists, engine)

Return stacked diagonal-Gaussian parameters for a homogeneous mixture kernel.

Parameters:

• dists (Sequence[DiagonalGaussianDistribution])

• engine (Any)

Return type:

dict[str, Any]

classmethod backend_stacked_log_density(x, params, engine)

Return an (n, k) matrix of diagonal-Gaussian log densities.

Parameters:

• x (Any)

• params (dict[str, Any])

• engine (Any)

Return type:

Any

to_fisher(**kwargs)

Return this distribution’s own Fisher view.

sampler(seed=None)

Return a sampler for iid draws from this distribution.

Parameters:

seed (int | None) – Optional seed for the sampler’s random state.

Returns:

A configured DiagonalGaussianSampler.

Return type:

DiagonalGaussianSampler

estimator(pseudo_count=None)

Return an estimator initialized from this distribution’s shape.

Parameters:

pseudo_count (float | None) – Optional smoothing count applied to mean and variance estimates.

Returns:

A DiagonalGaussianEstimator.

Return type:

DiagonalGaussianEstimator

dist_to_encoder()

Return an encoder for iid diagonal Gaussian observations.

Return type:

DiagonalGaussianDataEncoder

class DiagonalGaussianSampler(dist, seed=None)

Bases: DistributionSampler

Sampler for iid diagonal Gaussian observations.

Parameters:

• dist (DiagonalGaussianDistribution)

• seed (int | None)

sample(size=None, *, batched=True)

Draw iid samples from the diagonal Gaussian distribution.

Parameters:

• size (Optional[int]) – Number of iid samples to draw. If None, a single sample is drawn.

• batched (bool)

Returns:

Numpy array with shape (dim,) if size is None, else a list of ‘size’ such arrays.

Return type:

Sequence[ndarray] | ndarray

class DiagonalGaussianAccumulator(dim=None, keys=None)

Bases: SequenceEncodableStatisticAccumulator

Accumulator for diagonal Gaussian sufficient statistics.

Alongside the declared raw moments (sum, sum2, count) this keeps a SHIFT-ANCHORED moment track, for the reason the univariate GaussianAccumulator keeps one: the variance computed from raw reduced moments is the classic cancellation-prone E[x^2] - mu^2 form, which loses ~2*log2(abs(mean)/sd) bits, so unit-spread data at offset 1e7 fits variances tens of percent wrong and data at offset 1.7e9 fits variances thousands of times too large (or collapsed onto the floor – the sign of the garbage is data-dependent), on data whose variance a two-pass computation returns exactly. Epoch seconds are ~1.7e9. Anchoring at the first value seen keeps every term of the scatter O(count * spread^2), making the M-step variances shift-invariant. The track is CONDITIONING-GATED (see _needs_anchor()): a chunk the raw form handles to ~1e-9 relative error accumulates exactly the historical single-pass way – bit-identical statistics, no second pass – and the anchor activates only when a chunk (or a scalar update) would corrupt a variance.

Parameters:

• dim (int | None)

• keys (str | None)

update(x, weight, estimate)

Update sufficient statistics with a single weighted observation.

Parameters:

• x (Union[Sequence[float], np.ndarray]) – Length-dim observation vector.

• weight (float) – Weight for the observation.

• estimate (Optional[DiagonalGaussianDistribution]) – Kept for consistency with SequenceEncodableStatisticAccumulator (not used).

Returns:

None.

Return type:

None

initialize(x, weight, rng)

Initialize the accumulator with a weighted observation. Calls update().

Parameters:

• x (Union[Sequence[float], np.ndarray]) – Length-dim observation vector.

• weight (float) – Weight for the observation.

• rng (RandomState) – Kept for consistency with SequenceEncodableStatisticAccumulator.

Returns:

None.

Return type:

None

seq_update(x, weights, estimate)

Vectorized update of sufficient statistics with an encoded sequence of observations.

Parameters:

• x (np.ndarray) – Encoded data matrix with shape (sz, dim).

• weights (np.ndarray) – Numpy array of sz observation weights.

• estimate (Optional[DiagonalGaussianDistribution]) – Kept for consistency (not used).

Returns:

None.

Return type:

None

seq_initialize(x, weights, rng)

Vectorized initialization of the accumulator. Calls seq_update().

Parameters:

• x (np.ndarray) – Encoded data matrix with shape (sz, dim).

• weights (np.ndarray) – Numpy array of sz observation weights.

• rng (RandomState) – Kept for consistency with SequenceEncodableStatisticAccumulator.

Returns:

None.

Return type:

None

combine(suff_stat)

Merge sufficient statistics into this accumulator.

Parameters:

suff_stat (Tuple[np.ndarray, np.ndarray, float]) – Tuple of (weighted sum of observations, weighted sum of squared observations, sum of weights).

Returns:

This accumulator.

Return type:

DiagonalGaussianAccumulator

value()

Return (sum, sum_squares, count) sufficient statistics.

Once the shift-anchored moment track is live the returned value is a DiagonalGaussianSuffStat – a drop-in 3-tuple (indexing/unpacking/iteration all behave identically) that additionally carries those moments in its anchored attribute, so combine() can fold them in and DiagonalGaussianEstimator.estimate() can compute shift-invariant variances. Well-conditioned data never activates the track and gets the historical plain tuple.

Return type:

tuple[ndarray, ndarray, float]

from_value(x)

Replace this accumulator’s sufficient statistics.

Parameters:

x (Tuple[np.ndarray, np.ndarray, float]) – Tuple of (weighted sum of observations, weighted sum of squared observations, sum of weights).

Returns:

This accumulator.

Return type:

DiagonalGaussianAccumulator

scale(c)

Scale the accumulated statistics in-place by c, anchored track included.

The structural default routes through from_value(scale_suff_stat(self.value(), c)), which rebuilds from a PLAIN tuple and would therefore drop the anchored payload – turning a scaled large-offset accumulator back into a cancellation-prone one. Scaling every weight by c scales both anchored moments by c and leaves the anchor (a data value) alone.

Parameters:

c (float)

Return type:

DiagonalGaussianAccumulator

acc_to_encoder()

Return an encoder compatible with this accumulator’s dimension.

Return type:

DiagonalGaussianDataEncoder

class DiagonalGaussianAccumulatorFactory(dim=None, keys=None)

Bases: StatisticAccumulatorFactory

Factory for diagonal Gaussian accumulators.

Parameters:

• dim (int | None)

• keys (str | None)

make()

Return a fresh accumulator with the factory configuration.

Return type:

DiagonalGaussianAccumulator

class DiagonalGaussianEstimator(dim=None, pseudo_count=(None, None), suff_stat=(None, None), name=None, keys=None, prior=None, min_covar=None, ridge=None)

Bases: ParameterEstimator

Estimator for diagonal Gaussian distributions.

The fit is SHIFT-EQUIVARIANT: fit(x + c) returns the variances of x + c to within about 1e-7 relative for any constant c the data can carry, epoch seconds (~1.7e9) included. That is not free from the declared (sum, sum2, count) statistics – their variance is the cancellation-prone E[x^2] - mu^2, which at offset 1e7 was tens of percent wrong and at 1.7e9 thousands-fold wrong or collapsed onto the floor – so DiagonalGaussianAccumulator carries a conditioning-gated shift-anchored moment track and this estimator reads it. Statistics that arrive already reduced and WITHOUT that track (an engine kernel’s stacked moments, a hand-built tuple) cannot be corrected here and take the historical path; when they are too ill-conditioned for it, estimate warns rather than returning variances it cannot stand behind.

The 1e-7 is a measured relative bound, not a figure of speech, and the error is NOT uniform across the offset range. Once the anchored track engages, the variances are exact to a few ulps (~1e-14 relative, and unchanged from offset 1.7e9 to 1e10). The bound is set by the band just BELOW the conditioning gate, where the raw E[x^2] - mu^2 form is still in use because it is accurate enough to be worth its single pass: at mean^2/variance just under the gate’s 4e6 the worst case measured over seeds, sample sizes and scales is 3.1e-8 relative. An earlier revision of this docstring promised “a few ulps” everywhere, which overstated that band by roughly eight orders of magnitude (T1-07).

VARIANCE FLOOR. The fitted variances are the exact per-coordinate maximum-likelihood ones whenever every coordinate has a positive variance of its own. The M-step floors each coordinate at ridge * var_i (default ridge=1e-6) – RELATIVE to that coordinate’s own variance, so the safeguard is equivariant under any per-coordinate change of units and, sitting strictly below the variance it guards, never binds on such a coordinate. It is not priced off mean(var) across coordinates: that coupled columns a diagonal covariance exists to keep independent, and inflated the smallest variance of a plain probability/dollars/count table 3.03x (T1-03).

A coordinate with no positive variance of its own – a constant column, an EM component holding a single point – still has to be lifted, because a zero variance is an infinite density. It gets ridge * mean(var over the coordinates that have one), the same manufactured jitter the full covariance estimator uses for a degenerate direction. The absolute min_covar (default 1e-8) is the last resort when nothing carries a scale at all (every coordinate non-positive, or ridge=0.0), and becomes a hard per-coordinate lower bound as soon as it is passed explicitly – which is how mixle.task regularizes its diagonal Gaussians.

A floor that actually lifted a coordinate is recorded on the fitted distribution as a variance-floored(...) entry in numerical_repairs() – and therefore in fit_provenance().repairs – naming the coordinate and how far it moved. Manufactured jitter is always reported; an explicit min_covar dominating a positive variance is reported once it exceeds _FLOOR_MATERIAL_RATIO of it. An ordinary fit lifts nothing and records nothing. Because the default floor never binds on a coordinate with its own variance, ridge=0.0 is not the route to the maximum-likelihood fit here – it removes the only scale the degenerate-column jitter can be priced from and hands those coordinates to the absolute min_covar instead.

Parameters:

• dim (int | None)

• pseudo_count (float | tuple[float | None, float | None])

• suff_stat (tuple[ndarray | None, ndarray | None])

• name (str | None)

• keys (str | None)

• prior (SequenceEncodableProbabilityDistribution | None)

• min_covar (float | None)

• ridge (float | None)

accumulator_factory()

Return an accumulator factory matching this estimator.

Return type:

DiagonalGaussianAccumulatorFactory

model_log_density(model)

Log-density of the model parameters under the MultivariateNormalGamma prior (ELBO global term).

The prior is over (mu, tau=1/covar), so the model’s covariance is inverted before scoring.

Parameters:

model (DiagonalGaussianDistribution)

Return type:

float

estimate(nobs, suff_stat)

Estimate a diagonal Gaussian distribution from aggregated sufficient statistics.

Suff_stat is a Tuple of size 3 containing:

suff_stat[0] (np.ndarray): Component-wise sum of weighted observation values. suff_stat[1] (np.ndarray): Component-wise sum of weighted squared observation values. suff_stat[2] (float): Sum of weights for each observation.

Parameters:

• nobs (Optional[float]) – Weighted number of observations used in aggregation of suff stats.

• suff_stat (Tuple[np.ndarray, np.ndarray, float]) – See above for details.

Returns:

DiagonalGaussianDistribution object.

Return type:

DiagonalGaussianDistribution

class DiagonalGaussianDataEncoder(dim=None)

Bases: DataSequenceEncoder

Encoder for iid diagonal Gaussian observations.

Parameters:

dim (int | None)

seq_encode(x)

Encode a sequence of iid length-dim observations for vectorized seq_ calls.

Parameters:

x (Sequence[Union[List[float], np.ndarray]]) – Sequence of length-dim observation vectors.

Returns:

Encoded data matrix with shape (len(x), dim).

Return type:

ndarray