mixle.stats.compute.fused_kernels module

Fused numba kernels for mixture estimation over heterogeneous data.

Instead of one vectorized numpy pass per distribution per component (the seq_* path, which makes K * n_fields full passes over the data with temporaries), each supported distribution contributes a scalar row kernel

log_density(i, k, params, cols) -> float accumulate(i, k, w, cols, stats) -> None

where params/stats are tuples of arrays with a leading component axis (so EM iterations never trigger recompilation: new parameters are passed as data) and cols is a tuple of flat columnar arrays produced by one generic encoding pass. Kernels compose structurally - composites add child kernels, sequences sum a child kernel over an offset range - and numba inlines the composition into one fused loop: scoring all K components of an arbitrarily nested model touches each observation once.

This replaces the per-distribution vectorized encoders and seq_* math for supported models: the encoder is a single generic pass that flattens leaves into typed columns, and estimation is

enc = CompiledMixture(model).encode(data) model = compiled.fit(enc, estimator) # fused EM to convergence

The kernel E-step produces the same sufficient-statistic structures as the legacy accumulators, so the existing estimator M-steps (pseudo-counts included) are reused unchanged.

Agreement with the legacy seq path, precisely:

  • SCORING agrees to a few ulp at any data magnitude. Every Gaussian-family row kernel evaluates the CENTERED quadratic form -(x-mu)^2/(2*covar), the same form the seq_log_density implementations use, rather than the natural- parameter expansion whose three O(|x|^2/covar) terms cancel away the whole answer on offset data.

  • The E-step emits the HISTORICAL RAW moments (sum w*x, sum w*x^2, sum w). It does not reproduce the shift-anchored moment track the Python accumulators (GaussianAccumulator, DiagonalGaussianAccumulator, …) now carry, because the anchor has to be a DATA value chosen at encode time and the columnar encoding does not yet carry one. On data whose mean^2/variance is large the fitted VARIANCE therefore loses digits to cancellation exactly as the pre-anchor accumulators did (the fitted MEAN is a same-sign sum and is unaffected). Measured at offset 1.7e9 on unit-variance data: the diagonal M-step returns covar ~ 4.8e3 and the univariate M-step sigma2 ~ 2.9e10, against a true variance of ~1.0. DiagonalGaussianEstimator detects the missing anchored payload and emits a RuntimeWarning naming the conditioning and the remedy first; GaussianEstimator currently does NOT, so the scalar case is silent. Until an encode-time anchor lands, center the data (or accumulate through the Python accumulators) when fitting at large offsets on this path. Scoring is unaffected and stays exact at every magnitude.

Supported distributions: Gaussian, LogGaussian, Gamma, Categorical, IntegerCategorical, Bernoulli, StudentT, Logistic, Weibull, Rayleigh, Pareto, Uniform, Poisson, Exponential, Geometric, Binomial, NegativeBinomial, DiagonalGaussian (vector leaf), Optional (missing-data wrapper around any supported child), Ignored (fixed wrapped dist, scored but never estimated), Composite (any nesting), Sequence (variable-length, with optional length model), and a Mixture over same-structure components at the top level. Unsupported models should continue to use the seq_* path.

Still excluded: latent-variable families (HMMs, LDA/PLSI, tree models, hidden association) require per-row dynamic programs that do not fit the stateless scalar row kernel; heterogeneous mixtures break the identical-structure requirement; full-covariance MVN and von Mises-Fisher need linear algebra / Bessel-function evaluations outside the scalar kernels’ scope for now.

Performance characteristics (Apple silicon, float64): on flat scalar fields the fused scalar loops run at rough parity with the legacy vectorized numpy path (numpy’s per-field passes are SIMD-optimized; the kernels make a single threaded pass). On variable-length sequence data - the heterogeneous case this library targets - fusion is substantially faster: ~8-9x faster scoring and EM on topic-model-like mixtures of token sequences, because the legacy path re-walks the flattened token arrays once per component per field while the kernel path streams each observation exactly once.

class CompiledEncoding(n, columns, structure)[source]

Bases: Sequence[Any]

Columnar data tied to one compiled distribution structure.

The object intentionally preserves the historical two-item sequence interface (n, columns = encoding) while carrying the structure and array-geometry fingerprints needed to reject incompatible dispatch.

Parameters:
CompiledEncoding.count -> integer -- return number of occurrences of value
class CompiledMixture(model)[source]

Bases: object

Fused-kernel estimation engine for a mixture (or single distribution).

Compile once per model structure; parameters are re-read from the current model on every call, so EM iterations never recompile.

encode(data)[source]

Encode raw observations into the compiled mixture’s columnar buffers.

close()[source]

Release worker threads; safe to call repeatedly.

Return type:

None

seq_component_log_density(enc, model=None, n_threads=None)[source]

Return per-component log-density matrix for encoded observations.

Parameters:

n_threads (int | None)

Return type:

ndarray

seq_log_density(enc, model=None)[source]

Return mixture or single-model log densities for encoded observations.

Return type:

ndarray

posteriors(enc, model=None)[source]

Return posterior component weights for encoded observations.

Return type:

ndarray

weighted_suff_stats(enc, gamma, model=None, n_threads=None)[source]

Legacy-format sufficient statistics from per-row component weights.

Parameters:
em_step(enc, estimator, model=None, weights=None)[source]

One fused EM step; returns the re-estimated model via the legacy M-step.

Parameters:

weights (ndarray | None)

initialize(enc, estimator, rng, p=0.1)[source]

Random sparse-responsibility initialization (replaces seq_initialize).

Parameters:
fit(enc, estimator, max_its=100, delta=1.0e-8, rng=None, init_p=0.1, model=None, out=None, max_iter=None)[source]

EM to convergence with fused kernels. Returns (model, log_likelihood).

max_iter is the preferred spelling of max_its; when given it overrides max_its.

Parameters:
build_kernel(dists)[source]

Build the fused kernel node for K same-structure distributions.

Parameters:

dists (Sequence[Any])