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]

Return exact-or-approximate density semantics joined from component models.

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:
Return type:

None

seq_initialize(x, weights, rng)[source]

Vectorized initialize() with random Dirichlet assignments.

Parameters:
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 DP mixture sufficient statistics while preserving metadata.

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 for DP-mixture sufficient-statistic accumulators.

Parameters:
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:

Fitted Dirichlet-process mixture approximation.

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