mixle.stats.latent.mixture module

Homogeneous finite mixtures with stable scoring and EM accumulation.

This module defines MixtureDistribution, MixtureEstimator and the sampler, accumulator, factory, and encoder types used by the standard Mixle estimation loop.

A MixtureDistribution has density

p(y) = sum_k p(y | z=k) p(z=k).

All components are expected to model the same observation type. Scoring uses log-sum-exp over component log densities and log weights; impossible rows are represented as -inf scores rather than NaN.

mixture_prior(weight_prior, component_priors)

Build the joint mixture prior: a weight prior plus one prior per component.

Parameters:

• weight_prior (SequenceEncodableProbabilityDistribution) – Prior on the mixture weights (a DirichletDistribution or SymmetricDirichletDistribution).

• component_priors (Sequence[SequenceEncodableProbabilityDistribution]) – Sequence of one conjugate prior per component.

Returns:

A (weight_prior, tuple(component_priors)) pair consumed by MixtureDistribution/MixtureEstimator set_prior.

Return type:

tuple[SequenceEncodableProbabilityDistribution, tuple[SequenceEncodableProbabilityDistribution, …]]

class MixtureDistribution(components, w=MISSING, name=None, weights=MISSING, prior=None)

Bases: SequenceEncodableProbabilityDistribution

Finite mixture over homogeneous component distributions.

components define both the conditional families p(x | z=k) and the observation type accepted by the mixture. w contains the component probabilities and is cached as log_w for stable scalar and vectorized scoring. Zero-weight components are retained for structural compatibility but contribute -inf to log-density calculations.

Parameters:

• components (Sequence[SequenceEncodableProbabilityDistribution]) – Component distributions. Each component should support the same raw observation shape and sequence-encoding contract.

• w (np.ndarray | list[float]) – Component weights. The values are interpreted as simplex weights and should sum to one.

• name (str | None) – Optional display name for diagnostics and generated artifacts.

• weights (np.ndarray | list[float]) – Alias for w.

• prior (SequenceEncodableProbabilityDistribution | None) – Optional joint mixture prior or weight prior.

components

Component distribution objects.

w

Component weights as a NumPy array.

zw

Boolean mask for zero-weight components.

log_w

Log weights, with zero-weight entries represented as -inf.

num_components

Number of mixture components.

compute_capabilities()

Return compute-backend metadata shared by all mixture components.

compute_declaration()

Return the symbolic declaration for mixture weights and component statistics.

get_prior()

Return the joint mixture prior, or None for a plain point model.

When a weight prior is attached the joint prior is the (weight_prior, tuple(component priors)) pair produced by mixture_prior(); otherwise None.

Return type:

SequenceEncodableProbabilityDistribution | None

set_prior(prior)

Attach a weight prior (and optional per-component priors), caching weight expectations.

With a (symmetric) Dirichlet weight prior this caches the variational weight expectations E[log w_k] = digamma(alpha_k) - digamma(sum_j alpha_j) used by expected_log_density. Component priors, when supplied, are delegated to each component via component.set_prior. prior=None (the default) leaves the mixture a plain point model (byte-identical MLE behaviour).

Parameters:

prior (SequenceEncodableProbabilityDistribution | None)

Return type:

None

expected_log_density(x)

Variational expected log-density at observation x.

Uses E[log w_k] under the (symmetric) Dirichlet weight prior together with each component’s expected_log_density. Falls back to the plug-in log_density(x) when no conjugate weight prior is attached.

Parameters:

x (T)

Return type:

float

seq_expected_log_density(x)

Vectorized variational expected log-density at sequence-encoded input x.

Falls back to seq_log_density(x) when no conjugate weight prior is attached.

Parameters:

x (T1)

Return type:

ndarray

density(x)

Return the mixture density at one raw observation.

Parameters:

x (T) – Observation accepted by every component family.

Returns:

exp(log_density(x)).

Return type:

float

density_semantics()

Return joined density semantics over all mixture components.

log_density(x)

Return the mixture log-density at one raw observation.

The calculation is logsumexp_k(log p_k(x) + log w_k). Component families are responsible for returning -inf for observations outside support; the mixture combines those values without converting them to NaN.

Parameters:

x (T) – Observation accepted by every component family.

Returns:

Finite log-density when at least one positive-weight component can score the observation, otherwise -inf.

Return type:

float

conditional(observed)

Return the conditional mixture over the unobserved coordinates given observed.

The conditional of a mixture is itself a mixture: for sum_k w_k f_k observing x_o,

P(x_u | x_o) = sum_k w’_k f_k(x_u | x_o), w’_k proportional to w_k f_k.marginal(x_o)(x_o),

i.e. the component responsibilities are updated by how well each component explains the observed coordinates and each component is replaced by its own conditional. Because the result is a full MixtureDistribution you can both score it and .sampler(seed).sample() from it – the latter is given=-style conditional sampling that first draws a component from the posterior responsibilities, then draws the unobserved coordinates from that component’s conditional.

Requires each component to support marginal(indices) and condition(observed) (e.g. the multivariate Gaussian / Student-t). observed maps coordinate index to its fixed value.

Parameters:

observed (dict[int, float])

Return type:

MixtureDistribution

component_log_density(x)

Return component-wise log densities for one raw observation.

Parameters:

x (T) – Observation accepted by every component family.

Returns:

One log-density per component, before mixture weights are applied.

Return type:

ndarray

posterior(x)

Return component responsibilities for one raw observation.

Responsibilities are proportional to w[k] * p_k(x). If every positive-weight component reports an impossible observation, the method returns a copy of the prior mixture weights so callers receive a finite responsibility vector rather than NaN.

Parameters:

x (T) – Observation accepted by every component family.

Returns:

Probability vector over component labels.

Return type:

ndarray

seq_component_log_density(x)

Return vectorized component log densities for encoded observations.

x must be produced by MixtureDataEncoder.seq_encode or by an equivalent component encoder. The output has shape (n, k) where n is the number of encoded observations and k is the number of mixture components.

Parameters:

x (T1) – Encoded observation batch.

Returns:

Component log-density matrix before mixture weights are applied.

Return type:

ndarray

seq_log_density(x)

Return vectorized mixture log densities for encoded observations.

Each row is evaluated with a row-wise log-sum-exp over component scores plus log weights. Rows for which every positive-weight component is impossible return -inf.

Parameters:

x (T1) – Encoded observation batch.

Returns:

One log-density per encoded observation.

Return type:

ndarray

backend_seq_component_log_density(x, engine)

Engine-neutral component log densities for encoded data.

Parameters:

• x (T1)

• engine (Any)

Return type:

Any

backend_seq_log_density(x, engine)

Engine-neutral mixture log-density for encoded data.

Parameters:

• x (T1)

• engine (Any)

Return type:

Any

gradient_fit_state(engine, torch, leaves, recurse, tensor_param)

Return distribution-owned state for autograd fitting.

Parameters:

• engine (Any)

• torch (Any)

• leaves (list[Any])

• recurse (Any)

• tensor_param (Any)

Return type:

Any

seq_posterior(x)

Return vectorized component responsibilities for encoded observations.

Each row is proportional to w[k] * p_k(x_i). Rows where all positive-weight components are impossible fall back to the prior mixture weights, matching posterior() and avoiding NaN responsibility rows during EM accumulation.

Parameters:

x (T1) – Encoded observation batch.

Returns:

(n, k) probability matrix whose rows sum to one.

Return type:

ndarray

latent_posterior(x)

Return the latent posterior q(z | x) over component labels for raw observations x.

q(z) is the exact independent-categorical posterior whose marginals are the EM responsibilities. The returned CategoricalLatentPosterior can .marginals() (the responsibilities), .sample(rng) component labels, .mode() (the MAP labels), or .entropy().

Parameters:

x (Sequence[T])

Return type:

CategoricalLatentPosterior

posterior_predictive(x, seed=None)

Draw posterior-predictive observations conditioned on x.

For each observed x_i the component is sampled from the latent posterior q(z_i | x_i) and a fresh observation is emitted from that component – i.e. “given I saw x_i, draw a new point from the same mixture component it likely came from”. Returns a list the length of x. Draws are grouped by component and scattered (vectorized) via the shared sampling helper.

Parameters:

• x (Sequence[T])

• seed (int | None)

Return type:

list[Any]

support_size()

Upper bound on distinct support points: the sum over components (union <= sum).

Return type:

int | None

tropical_displacement_bits()

log2(#positive-weight components) – the tropical-vs-marginal cost gap (in bits).

The marginal log p(x) = logsumexp_k (log w_k + log p_k(x)) is bounded by its largest term M(x) = max_k (log w_k + log p_k(x)) via M(x) <= log p(x) <= M(x) + log K, where K is the number of components that can contribute (positive weight). The structural seek bins by the tropical cost M(x); mixle.enumeration.density_rank.marginal_seek() widens its smear window by this many bits so the reported rank bracket provably contains the TRUE marginal rank. K <= 1 means the marginal is a single term -> 0.0 (the seek is then exact). When the component supports are provably disjoint every value lands in one component, so M(x) equals the marginal and there is likewise no displacement -> 0.0 (the seek is exact and tight).

Return type:

float

to_fisher(**kwargs)

Structural Fisher view for the mixture.

sampler(seed=None)

Return a sampler that draws from the mixture distribution.

Parameters:

seed (int | None) – Optional RandomState seed for reproducible draws.

Returns:

MixtureSampler bound to this distribution.

Return type:

MixtureSampler

estimator(pseudo_count=None)

Return an estimator with matching component structure.

Parameters:

pseudo_count (float | None) – Optional smoothing mass applied through the estimator path.

Returns:

MixtureEstimator suitable for fitting observations of the same type as this distribution.

Return type:

MixtureEstimator

decomposition()

Mixture components split along the component axis. Responsibilities (logsumexp) are computed INSIDE a shard; across shards the per-component sufficient stats SUM-reduce plus one scalar total-count all-reduce – the homogeneous stacked-kernel + DTensor path (engine_axis=0).

dist_to_encoder()

Return an encoder that delegates observation encoding to components.

Return type:

MixtureDataEncoder

enumerator()

Return an enumerator over the union of component supports.

Return type:

MixtureEnumerator

quantized_index(max_bits, bin_width_bits=1.0)

Build a bounded bit-quantized index from a global mixture frontier.

The primary path pulls candidates from weighted component enumerator heads. The log-sum of those heads bounds every unseen value, so construction stops when the live global frontier falls below 2**(-max_bits). This avoids the looser per-component log2(K) candidate expansion. If a component cannot enumerate, the method falls back to the structured cross-index path.

Parameters:

• max_bits (float)

• bin_width_bits (float)

Return type:

QuantizedEnumerationIndex

quantized_count_index(quantizer, max_fine_bucket)

BoundedCount for the MARGINAL mixture law: pool weight-scaled component count indices.

log p(x) = logsumexp_k (log w_k + log p_k(x)) has no exact structural count – overlapping component supports would need value-level deduplication. This builds the count semiring’s plus-fold over scale(component_index, log w_k) instead, which:

• reaches a 2**M budget structurally (no enumeration), and

• is a conservative UPPER bound – a value shared by several components is counted once per component, and each value is binned by its dominant weighted component (the tropical cost, within log2(K) bits of the exact logsumexp).

Every unranked value still carries its exact mixture log_density (re-evaluated by the budget builder). For an exact small-budget index (best-first union with dedup), use quantized_index. Components that cannot count structurally raise EnumerationError.

Parameters:

max_fine_bucket (int)

structural_fine_bucket(value, quantizer)

Dominant weighted-component structural bucket (mirrors the plus-of-scaled-children index).

Return type:

int

is_canonical_copy(value, coarse_bin, quantizer)

Stateless dedup: keep value only at its dominant (best-weighted) component’s bin.

The canonical bin is the coarse bin of the minimum, over components, of the component’s structural fine bucket shifted by the weight term. O(K) model evaluations, no state.

Parameters:

coarse_bin (int)

Return type:

bool

class MixtureEnumerator(dist)

Bases: DistributionEnumerator

Enumerator over the deduplicated union of weighted component supports.

Parameters:

dist (MixtureDistribution)

class MixtureSampler(dist, seed=None)

Bases: DistributionSampler

Sampler that draws a latent component and then samples from that component.

Parameters:

• dist (MixtureDistribution)

• seed (int | None)

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

Draw iid samples from a mixture distribution.

The data type drawn from ‘comp_samplers’ is type T, corresponding to the data type of the mixture components.

If size is None, a single sample (of data type T) is drawn and returned. If size is not None, ‘size’-iid mixture samples are drawn and returned as a List with data type List[T].

With batched=True (default) each component sampler is invoked once with the number of draws assigned to it and the results are scattered back into draw order. Because every component sampler owns an independent RandomState, this yields the same draws as the legacy per-draw loop (batched=False) but far faster.

Parameters:

• size (Optional[int]) – Number of iid samples to draw.

• batched (bool) – Vectorize component draws (default); set False for the legacy per-draw loop.

Returns:

Data type T or List[T].

Return type:

list[Any] | Any

class MixtureAccumulator(accumulators, keys=(None, None), name=None, init='dirichlet')

Bases: SequenceEncodableStatisticAccumulator

EM accumulator for mixture weights and component sufficient statistics.

Parameters:

• accumulators (Sequence[SequenceEncodableStatisticAccumulator])

• keys (tuple[str | None, str | None])

• name (str | None)

• init (str)

seq_update(x, weights, estimate)

Accumulate a vectorized EM E-step from encoded observations.

Responsibilities are computed from estimate using the same log-sum-exp normalization as MixtureDistribution.seq_posterior. Rows where every component is impossible fall back to the estimate’s mixture weights, so the accumulator receives finite responsibility weights rather than NaN.

Parameters:

• x (T1) – Encoded observation batch.

• weights (ndarray) – Non-negative observation weights.

• estimate (MixtureDistribution) – Previous EM iterate used to compute responsibilities.

Return type:

None

update(x, weight, estimate)

Accumulate one weighted raw observation under an EM estimate.

The observation is routed to each component accumulator with weight * estimate.posterior(x)[k].

Parameters:

• x (T) – Raw observation.

• weight (float) – Observation weight.

• estimate (MixtureDistribution) – Previous EM iterate used to compute responsibilities.

Return type:

None

initialize(x, weight, rng)

Initialize component sufficient statistics from one observation.

The default initialization draws a responsibility vector from a Dirichlet distribution and delegates responsibility-weighted initialization to every component accumulator.

Parameters:

• x (T) – Raw observation.

• weight (float) – Observation weight.

• rng (RandomState) – Random state used to seed component initializers.

Return type:

None

seq_initialize(x, weights, rng)

Initialize component sufficient statistics from encoded observations.

With init="kmeans++" the method uses a numeric feature matrix when one can be extracted from the encoded data. Ragged, object, hetero, or non-finite encodings fall back to Dirichlet responsibilities rather than mutating input data or forcing an invalid numeric representation.

Parameters:

• x (T1) – Encoded observation batch.

• weights (ndarray) – Non-negative observation weights.

• rng (RandomState) – Random state used to seed component initializers.

Return type:

None

combine(suff_stat)

Merge serialized mixture sufficient statistics into this accumulator.

Parameters:

suff_stat (tuple[ndarray, tuple[T2, ...]]) – (component_counts, component_suff_stats) tuple.

Returns:

self for accumulator chaining.

Return type:

MixtureAccumulator

value()

Return serialized mixture sufficient statistics.

Returns:

(component_counts, component_suff_stats) where the second item contains one serialized child accumulator value per component.

Return type:

tuple[ndarray, tuple[Any, …]]

from_value(x)

Restore this accumulator from serialized sufficient statistics.

Parameters:

x (tuple[ndarray, tuple[T2, ...]]) – (component_counts, component_suff_stats) tuple.

Returns:

self after restoring child accumulator state.

Return type:

MixtureAccumulator

scale(c)

Scale component counts and delegate child sufficient statistics.

Parameters:

c (float)

Return type:

MixtureAccumulator

key_merge(stats_dict)

Merge keyed mixture statistics into a shared statistics dictionary.

Parameters:

stats_dict (dict[str, Any]) – Mutable shared sufficient-statistics mapping keyed by estimator key names.

Return type:

None

key_replace(stats_dict)

Replace local keyed statistics from a shared statistics dictionary.

Parameters:

stats_dict (dict[str, Any]) – Shared sufficient-statistics mapping keyed by estimator key names.

Return type:

None

acc_to_encoder()

Return an encoder assembled from the component accumulators.

Return type:

MixtureDataEncoder

class MixtureAccumulatorFactory(factories, keys=(None, None), name=None, init='dirichlet')

Bases: StatisticAccumulatorFactory

Factory for mixture accumulators built from component accumulator factories.

Parameters:

• factories (Sequence[StatisticAccumulatorFactory])

• keys (tuple[str | None, str | None])

• name (str | None)

• init (str)

make()

Return a fresh mixture accumulator with fresh component accumulators.

Return type:

MixtureAccumulator

class MixtureEstimator(estimators, fixed_weights=None, suff_stat=None, pseudo_count=None, name=None, keys=(None, None), prior=None, w_min=0.0, robust=False, init=None)

Bases: ParameterEstimator

Estimator for mixture weights and component distributions from EM sufficient statistics.

Parameters:

• estimators (Sequence[ParameterEstimator])

• fixed_weights (list[float] | np.ndarray | None)

• suff_stat (np.ndarray | None)

• pseudo_count (float | None)

• name (str | None)

• keys (tuple[str | None, str | None])

• prior (SequenceEncodableProbabilityDistribution | None)

• w_min (float)

• robust (bool)

• init (str | None)

accumulator_factory()

Return a mixture accumulator factory matching the component estimators.

Return type:

MixtureAccumulatorFactory

get_prior()

Return the joint mixture prior, or None for a plain MLE estimator.

When a weight prior is attached the joint prior is the (weight_prior, tuple(component priors)) pair produced by mixture_prior().

Return type:

SequenceEncodableProbabilityDistribution | None

set_prior(prior)

Attach a weight prior (and optional per-component priors).

With a (symmetric) Dirichlet weight prior the estimator switches to the conjugate MAP weight update; component priors, when supplied, are delegated to each component estimator via estimator.set_prior (those carry out their own conjugate updates). prior=None leaves the estimator a plain MLE estimator (byte-identical behaviour).

Parameters:

prior (SequenceEncodableProbabilityDistribution | None)

Return type:

None

model_log_density(model)

Log density of the model parameters under this estimator’s prior (ELBO global term).

Returns the Dirichlet weight-prior log-density evaluated at model.w plus the sum of each component estimator’s model_log_density at the corresponding component model. Returns 0.0 for a plain MLE estimator with no priors anywhere.

Parameters:

model (MixtureDistribution)

Return type:

float

estimate(nobs, suff_stat)

Estimate a mixture distribution from EM sufficient statistics.

suff_stat is (component_counts, component_suff_stats). Component parameters are delegated to the child estimators. Mixture weights follow the fixed-weight, conjugate-prior, pseudo-count, or plain-MLE path selected by the estimator configuration. Plain-MLE weights may be floored by w_min and are always renormalized.

Parameters:

• nobs (float | None) – Unused compatibility argument from ParameterEstimator.

• suff_stat (tuple[ndarray, tuple[Any, ...]]) – Serialized mixture sufficient statistics.

Returns:

Fitted MixtureDistribution.

Return type:

MixtureDistribution

class MixtureDataEncoder(encoder)

Bases: DataSequenceEncoder

Encoder for homogeneous or heterogeneous mixture component encodings.

Parameters:

encoder (DataSequenceEncoder | Sequence[DataSequenceEncoder])

seq_encode(x)

Sequence encode a sequence of iid observations drawn from the mixture distribution.

For a homogeneous mixture this delegates to the single shared component encoder. For a heterogeneous mixture each component encoder encodes the data separately and the encodings are bundled in a _HeteroMixtureEncoded wrapper.

Parameters:

x (Sequence[T]) – A Sequence of iid observations drawn from a mixture distribution with component distributions consistent with the per-component encoders.

Returns:

Encoded sequence (single shared encoding, or a per-component wrapper).

Return type:

Any

class MixtureFisherView(dist)

Bases: FixedFisherView

Complete-data Fisher view for finite mixture distributions.

Coordinates are component assignment indicators followed by each component’s sufficient statistics gated by that assignment. Observed data map to posterior-expected complete-data statistics.

Parameters:

dist (Any)

structured_statistics(x, estimate=None, weight=1.0)

Return mixture responsibility statistics and weighted component statistics for one observation.

Parameters:

• x (Any)

• estimate (Any | None)

• weight (float)

Return type:

Any