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, and results agree with the legacy seq path to floating-point tolerance.

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 wins decisively: ~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 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]
seq_component_log_density(enc, model=None, n_threads=None)[source]
Parameters:

n_threads (int | None)

Return type:

ndarray

seq_log_density(enc, model=None)[source]
Return type:

ndarray

posteriors(enc, model=None)[source]
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])