mixle.stats.bayes.dirichlet_process_mixture module¶

Dirichlet process mixture (truncated stick-breaking) with variational Bayes estimation, adapted onto the mixle.stats base-class protocol.

A truncated stick-breaking representation with K components:

alpha ~ Gamma(s1, 1/s2) (concentration hyper-prior) v_k | alpha ~ Beta(1, alpha) (stick fractions, k < K) w_k = v_k * prod_{j<k} (1 - v_j) (mixture weights) z_i | w ~ Categorical(w) x_i | z_i = k ~ components[k]

Data type: whatever the component distributions accept; a datum is a single observation scored under the mixture log sum_k w_k p(x | theta_k).

Estimation is mean-field variational Bayes: accumulators collect the optimal local assignments phi_ik (computed from the components’ expected_log_density, i.e. the VB E-step), and the estimator updates the variational Beta posteriors gamma_k on the stick fractions, the Gamma hyper-posterior on alpha, and each component’s conjugate update. Components are re-sorted by expected count each iteration, and each component’s posterior (carried as its prior, i.e. component.get_prior()) serves as the variational factor q(theta_k). seq_local_elbo provides the per-observation data terms of the ELBO; the data-independent terms live in DirichletProcessMixtureEstimator.model_log_density.

This is a port of mixle.bstats.dpm. The variational math is preserved exactly; only the surrounding object protocol is adapted to mixle.stats: SequenceEncodableProbabilityDistribution / SequenceEncodableStatisticAccumulator / StatisticAccumulatorFactory / ParameterEstimator, a DataSequenceEncoder (delegated to components[0]), the two-argument estimate(nobs, suff_stat) signature, and seq_initialize on the accumulator.

cbg(x, s1, s2)[source]

Log-density of a compound Beta-Gamma stick fraction: x = 1 - exp(-y) with y ~ Exponential(alpha) and alpha ~ Gamma(s1, 1/s2), marginalized over alpha.

Parameters:

x (float) – Stick fraction in (0, 1).
s1 (float) – Gamma shape of the concentration hyper-prior.
s2 (float) – Gamma rate of the concentration hyper-prior.

Returns:

Log-density at x.

Return type:

float

class DirichletProcessMixtureDistribution(components, w, a, g, component_priors, name=None, prior=default_prior)[source]

Bases: SequenceEncodableProbabilityDistribution

Truncated Dirichlet process mixture with stick-breaking weights w over K component distributions, carrying the variational Beta posteriors.

Parameters:

components (Sequence[SequenceEncodableProbabilityDistribution])
w (ndarray | list[float])
a (float)
g (ndarray)
component_priors (Sequence[SequenceEncodableProbabilityDistribution])
name (str | None)
prior (SequenceEncodableProbabilityDistribution | None)

get_prior()[source]

Return the Gamma hyper-posterior on the concentration alpha.

Return type:: SequenceEncodableProbabilityDistribution | None

set_prior(prior)[source]

Set the Gamma hyper-prior (or hyper-posterior) on alpha.

Parameters:: prior (SequenceEncodableProbabilityDistribution | None)
Return type:: None

get_parameters()[source]

Returns the parameter tuple (alpha, weights, component parameters).

Return type:: tuple[float, ndarray, list[Any]]

set_parameters(params)[source]

Set the parameters and refresh the cached log-weights.

Parameters:: params (tuple[float, ndarray, Sequence[Any]]) – Tuple (alpha, weights, components).
Return type:: None

density(x)[source]

Density of the mixture at observation x; see log_density().

Parameters:: x (Any)
Return type:: float

log_density(x)[source]

Mixture log-density log sum_k w_k p(x | theta_k) at observation x.

Parameters:: x (Any)
Return type:: float

expected_log_density(x)[source]

Mixture log-density with each component’s plug-in log-density replaced by its variational expectation E_q[log p(x | theta_k)].

Parameters:: x (Any)
Return type:: float

seq_log_density(x)[source]

Vectorized log-density at sequence-encoded input x.

Parameters:: x (Any)
Return type:: ndarray

posterior(x)[source]

Return the component posterior p(z = k | x) at a single observation.

This is the plug-in mixture posterior consistent with log_density(): softmax_k( log p(x | theta_k) + log w_k ). An observation with no support under any component falls back to the mixture weights w. Returns a length-K array summing to 1.

Parameters:: x (Any)
Return type:: ndarray

seq_posterior(x)[source]

Vectorized component posterior over a sequence-encoded input.

Returns an (sz, K) array whose row i is the plug-in posterior p(z = k | x_i) (see posterior()); rows for observations with no support under any component fall back to the mixture weights. A row-wise log-sum-exp keeps the softmax numerically stable.

Parameters:: x (Any)
Return type:: ndarray

seq_local_elbo(x)[source]

Per-observation local ELBO contributions.

For each observation i this returns

sum_k phi_ik * ( E_q[log p(z_i = k | v)] + E_q[log p(x_i | theta_k)] - log phi_ik )

where phi_i is the optimal variational assignment for x_i. The global (data-independent) ELBO terms are returned by DirichletProcessMixtureEstimator.model_log_density.

Parameters:: x (Any)
Return type:: ndarray

seq_expected_log_density(x)[source]

Vectorized expected_log_density() at sequence-encoded input x.

Parameters:: x (Any)
Return type:: ndarray

density_semantics()[source]

What log_density returns relative to the true log-density (default: exact).

Override to declare that this distribution’s log_density is a variational lower bound (ELBO), an upper bound, or an approximation rather than the exact log p(x). This is surfaced as the ExactDensity capability and noted in mixle.describe(), so code that needs an exact likelihood can require(x, ExactDensity) instead of silently trusting a bound.

sampler(seed=None)[source]

Create a DirichletProcessMixtureSampler for this distribution.

Parameters:: seed (int | None)
Return type:: DirichletProcessMixtureSampler

estimator(pseudo_count=None)[source]

Create a DirichletProcessMixtureEstimator from this distribution’s components.

Parameters:: pseudo_count (float | None)
Return type:: DirichletProcessMixtureEstimator

dist_to_encoder()[source]

Returns a DirichletProcessMixtureDataEncoder delegating to the components.

Return type:: DirichletProcessMixtureDataEncoder

class DirichletProcessMixtureSampler(dist, seed=None)[source]

Bases: DistributionSampler

Draws samples from a DirichletProcessMixtureDistribution.

Parameters:

dist (DirichletProcessMixtureDistribution)
seed (int | None)

sample(size=None, *, batched=True)[source]

Draw size samples (a single observation when size is None).

A component is chosen with probability w_k and an observation is drawn from that component. With batched=True (default) component draws are grouped and scattered – bit-identical to the per-draw loop (batched=False) but far faster, since each component sampler owns an independent RNG.

Parameters:

size (Optional[int]) – Number of samples to draw.
batched (bool) – Vectorize component draws (default); set False for the per-draw loop.

Returns:

A single observation if size is None, else a list of size observations.

Return type:

Any

class DirichletProcessMixtureAccumulator(accumulators, keys=(None, None))[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulates DPM sufficient statistics: expected component counts, Beta stick-fraction counts, and each component’s weighted statistics.

Parameters:

accumulators (Sequence[SequenceEncodableStatisticAccumulator])
keys (tuple[str | None, str | None])

update(x, weight, estimate)[source]

Accumulate the VB E-step statistics for observation x.

Computes the optimal variational assignment phi from the current estimate’s expected log-densities and weights, then adds the weighted phi to the component and stick-fraction counts and pushes phi-weighted updates into the component accumulators.

Parameters:

x (Any)
weight (float)
estimate (DirichletProcessMixtureDistribution)

Return type:

None

seq_update(x, weights, estimate)[source]

Vectorized update() on sequence-encoded data.

Parameters:

x (Any)
weights (ndarray)
estimate (DirichletProcessMixtureDistribution)

Return type:

None

initialize(x, weight, rng)[source]

Initialize with a random Dirichlet assignment of observation x.

Parameters:

x (Any)
weight (float)
rng (RandomState)

Return type:

None

seq_initialize(x, weights, rng)[source]

Vectorized initialize() with random Dirichlet assignments.

Parameters:

x (Any)
weights (ndarray)
rng (RandomState)

Return type:

None

combine(suff_stat)[source]

Add another accumulator’s sufficient-statistic value into this one.

Parameters:: suff_stat (tuple)
Return type:: DirichletProcessMixtureAccumulator

scale(c)[source]

Scale linear sufficient statistics in-place by c.

The structural default is correct for ordinary weighted sums, nested tuples/lists/dicts, and numeric arrays. Families whose value() payload includes non-linear metadata such as support bounds must override this method and leave that metadata unscaled.

Parameters:: c (float)
Return type:: DirichletProcessMixtureAccumulator

value()[source]

Returns (comp_counts, beta_counts, alpha, prev_nw, component values).

Return type:: tuple

from_value(x)[source]

Set the sufficient statistics from a value() tuple.

Parameters:: x (tuple)
Return type:: DirichletProcessMixtureAccumulator

key_merge(stats_dict)[source]

Merge this accumulator’s keyed statistics into a shared dict.

Parameters:: stats_dict (dict[str, Any])
Return type:: None

key_replace(stats_dict)[source]

Replace this accumulator’s statistics with the pooled keyed values.

Parameters:: stats_dict (dict[str, Any])
Return type:: None

acc_to_encoder()[source]

Returns a DirichletProcessMixtureDataEncoder delegating to the components.

Return type:: DirichletProcessMixtureDataEncoder

class DirichletProcessMixtureAccumulatorFactory(factories, dim, keys)[source]

Bases: StatisticAccumulatorFactory

Factory that creates DirichletProcessMixtureAccumulator objects.

Parameters:

factories (Sequence[StatisticAccumulatorFactory])
dim (int)
keys (tuple[str | None, str | None])

make()[source]

Returns a new DirichletProcessMixtureAccumulator.

Return type:: DirichletProcessMixtureAccumulator

class DirichletProcessMixtureEstimator(estimators, name=None, prior=default_prior, pseudo_count=None, keys=(None, None))[source]

Bases: ParameterEstimator

Estimates a DirichletProcessMixtureDistribution by mean-field variational Bayes from accumulated assignment statistics.

Parameters:

estimators (Sequence[ParameterEstimator])
name (str | None)
prior (SequenceEncodableProbabilityDistribution | None)
pseudo_count (float | None)
keys (tuple[str | None, str | None])

accumulator_factory()[source]

Returns a DirichletProcessMixtureAccumulatorFactory for this estimator.

Return type:: DirichletProcessMixtureAccumulatorFactory

get_prior()[source]

Return the Gamma hyper-prior on the concentration alpha.

Return type:: SequenceEncodableProbabilityDistribution | None

set_prior(prior)[source]

Set the Gamma hyper-prior on the concentration alpha.

Parameters:: prior (SequenceEncodableProbabilityDistribution | None)
Return type:: None

model_log_density(model)[source]

Data-independent ELBO terms of the variational approximation.

Combines the cross-entropies of the stick-fraction prior and the component priors against their variational posteriors with the entropies of those posteriors. Together with DirichletProcessMixtureDistribution.seq_local_elbo this forms the full ELBO maximized by the fit driver.

Parameters:: model (DirichletProcessMixtureDistribution)
Return type:: float

estimate(nobs, suff_stat)[source]

Estimate a DirichletProcessMixtureDistribution by one VB M-step.

Re-estimates each component (whose conjugate update carries its posterior forward as its prior), re-sorts components by expected count, updates the variational Beta posteriors gamma_k on the stick fractions, updates the Gamma hyper-posterior on the concentration alpha (carried as the returned distribution’s prior), and converts the expected log stick fractions into the mixture weights w.

Parameters:

nobs (Optional[float]) – Not used. Kept for the stats ParameterEstimator.estimate(nobs, suff_stat) signature.
suff_stat (tuple) – Tuple (comp_counts, beta_counts, alpha, prev_nw, component suff stats) as returned by DirichletProcessMixtureAccumulator.value().

Returns:

DirichletProcessMixtureDistribution object.

Return type:

DirichletProcessMixtureDistribution

class DirichletProcessMixtureDataEncoder(encoder)[source]

Bases: DataSequenceEncoder

Encodes observations with the shared component encoding.

Parameters:: encoder (DataSequenceEncoder)

seq_encode(x)[source]

Encode a sequence of observations with the component encoder.

Parameters:: x (Sequence[Any])
Return type:: Any