mixle.stats.bayes.hierarchical_dirichlet_process_mixture module

Hierarchical Dirichlet process mixture (truncated) for grouped data, adapted onto the mixle.stats base-class protocol.

Observations arrive in groups (a datum is a sequence of observations). All groups share K global atoms; each group mixes them with its own weights. This implements the finite “direct-assignment” truncation of the HDP (Teh et al. 2006):

beta ~ Dirichlet(gamma/K, …, gamma/K) (global weights) pi_j | beta ~ Dirichlet(alpha * beta) (group j weights) z_ji | pi_j ~ pi_j x_ji | z_ji = k ~ components[k]

Estimation alternates:

• E-step at point estimates: responsibilities phi_jik from the group’s current weights and the atom densities,

• posterior-mean update for each group’s weights under Dirichlet(alpha*beta + expected_counts), deliberately using the mean rather than the boundary-degenerate MAP when alpha*beta_k < 1, together with the atoms’ estimator updates,

• global-weight update via the standard expected-table-count approximation m_jk = alpha*beta_k*(psi(alpha*beta_k + n_jk) - psi(alpha*beta_k)), with beta set to the Dirichlet(gamma/K + m_.k) posterior mean. Applying this table-count formula to fractional responsibility counts is a deterministic approximation, not an exact collapsed-HDP CAVI step.

seq_local_elbo scores training groups with their fitted weights (this is what the fit driver maximizes); seq_log_density scores a (possibly new) group with the global weights beta, i.e. the expected weights of an unseen group. For multi-observation new groups this is a beta plug-in score, not the integrated finite-HDP predictive density obtained by integrating over a new group row pi ~ Dirichlet(alpha*beta).

Group sizes are exogenous unless len_dist is supplied (used for sampling and added to the per-group score). The length model uses the mixle.stats NullDistribution/NullEstimator/NullAccumulator family.

This is a port of mixle.bstats.hdpm onto the mixle.stats protocol. The object should be read as the finite direct-assignment approximation described above, with posterior-mean rows and an expected-table global-row heuristic.

class HierarchicalDirichletProcessMixtureDistribution(components, beta, alpha, gamma, group_weights=None, name=None, len_dist=None)

Bases: SequenceEncodableProbabilityDistribution

Truncated hierarchical DP mixture over K shared atoms with global weights beta and (optionally) fitted per-group weights.

Parameters:

• components (Sequence[SequenceEncodableProbabilityDistribution])

• beta (ndarray | list[float])

• alpha (float)

• gamma (float)

• group_weights (ndarray | None)

• name (str | None)

• len_dist (SequenceEncodableProbabilityDistribution | None)

get_parameters()

Returns the parameter tuple (beta, alpha, gamma, component parameters).

Return type:

tuple[ndarray, float, float, list[Any]]

set_parameters(params)

Set the parameters and refresh the cached log-weights.

Parameters:

params (tuple[ndarray, float, float, Sequence[Any]])

Return type:

None

density(x)

Density of a group x; see log_density().

Parameters:

x (Any)

Return type:

float

density_semantics()

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.

log_density(x)

Score a group with the global weights beta (expected weights of a new group).

Parameters:

x (Any)

Return type:

float

seq_encode(x)

Encode groups into a flat component encoding with offsets.

Parameters:

x (Sequence[Sequence]) – Iterable of groups (sequences of observations).

Returns:

Tuple (lengths, offsets, flat_enc, len_enc) for use with seq_ methods.

Return type:

Any

seq_log_density(x)

Vectorized log_density() at sequence-encoded input x (each group scored with the global weights beta).

Parameters:

x (Any)

Return type:

ndarray

seq_local_elbo(x)

Per-group data term of the penalized objective: training groups are scored with their own fitted weights. Falls back to beta when the group count does not match the fit.

Parameters:

x (Any)

Return type:

ndarray

group_posteriors(x)

Posterior atom-usage (mean responsibility) per group.

Groups are scored with their fitted weights when available (else with the global weights beta).

Parameters:

x (Sequence[Sequence])

Return type:

ndarray

sampler(seed=None)

Create a HierarchicalDirichletProcessMixtureSampler for this distribution.

Parameters:

seed (int | None)

Return type:

HierarchicalDirichletProcessMixtureSampler

estimator(pseudo_count=None)

Create a HierarchicalDirichletProcessMixtureEstimator from this distribution’s components, concentrations, and length estimator.

Parameters:

pseudo_count (float | None)

Return type:

HierarchicalDirichletProcessMixtureEstimator

dist_to_encoder()

Returns a HierarchicalDirichletProcessMixtureDataEncoder for this distribution.

Return type:

HierarchicalDirichletProcessMixtureDataEncoder

class HierarchicalDirichletProcessMixtureSampler(dist, seed=None)

Bases: DistributionSampler

Draws groups from a HierarchicalDirichletProcessMixtureDistribution (per-group weights drawn from Dirichlet(alpha*beta)).

Parameters:

• dist (HierarchicalDirichletProcessMixtureDistribution)

• seed (int | None)

sample_group(n=None)

Draw a single group of n observations.

Group weights pi ~ Dirichlet(alpha*beta) are drawn once for the group, then each observation draws an atom from pi.

Parameters:

n (int | None)

Return type:

list[Any]

sample(size=None)

Draw size groups (a single group when size is None).

Parameters:

size (int | None)

Return type:

Any

class HierarchicalDirichletProcessMixtureAccumulator(accumulators, len_accumulator=None, name=None, keys=None)

Bases: SequenceEncodableStatisticAccumulator

Accumulates HDP mixture sufficient statistics: per-group expected atom counts (in data order) plus each atom’s weighted statistics.

Parameters:

• accumulators (Sequence[SequenceEncodableStatisticAccumulator])

• len_accumulator (SequenceEncodableStatisticAccumulator | None)

• name (str | None)

• keys (str | None)

initialize(x, weight, rng)

Initialize with random Dirichlet assignments for group x.

Parameters:

• x (Any)

• weight (float)

• rng (RandomState)

Return type:

None

seq_initialize(x, weights, rng)

Vectorized initialize() with random Dirichlet assignments.

Parameters:

• x (Any)

• weights (ndarray)

• rng (RandomState)

Return type:

None

update(x, weight, estimate)

Accumulate the E-step statistics for one group (delegates to seq_update on a singleton encoding).

Parameters:

• x (Any)

• weight (float)

• estimate (HierarchicalDirichletProcessMixtureDistribution)

Return type:

None

seq_update(x, weights, estimate)

E-step on sequence-encoded data at the current point estimates.

Computes responsibilities phi from each group’s current weights (the fitted group weights, or beta for new/unmatched groups) and the atom densities, recording per-group expected counts and pushing phi-weighted updates into the atom accumulators. Also records the estimate’s beta and alpha for the estimator’s global-weight update.

Parameters:

• x (Any)

• weights (ndarray)

• estimate (HierarchicalDirichletProcessMixtureDistribution)

Return type:

None

combine(suff_stat)

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

Parameters:

suff_stat (tuple)

Return type:

HierarchicalDirichletProcessMixtureAccumulator

scale(c)

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:

HierarchicalDirichletProcessMixtureAccumulator

value()

Returns (group_counts, prev_beta, prev_alpha, atom values, len_value).

Return type:

tuple

from_value(x)

Set the sufficient statistics from a value() tuple.

Parameters:

x (tuple)

Return type:

HierarchicalDirichletProcessMixtureAccumulator

key_merge(stats_dict)

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

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)

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

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()

Returns a HierarchicalDirichletProcessMixtureDataEncoder for this accumulator.

Return type:

HierarchicalDirichletProcessMixtureDataEncoder

class HierarchicalDirichletProcessMixtureAccumulatorFactory(factories, len_factory, name, keys)

Bases: StatisticAccumulatorFactory

Factory that creates HierarchicalDirichletProcessMixtureAccumulator objects.

Parameters:

• factories (Sequence[StatisticAccumulatorFactory])

• len_factory (StatisticAccumulatorFactory | None)

• name (str | None)

• keys (str | None)

make()

Returns a new HierarchicalDirichletProcessMixtureAccumulator.

Return type:

HierarchicalDirichletProcessMixtureAccumulator

class HierarchicalDirichletProcessMixtureEstimator(estimators, gamma=1.0, alpha=1.0, name=None, keys=None, len_estimator=None)

Bases: ParameterEstimator

Estimates a HierarchicalDirichletProcessMixtureDistribution from accumulated group counts via the direct-assignment truncation updates.

Parameters:

• estimators (Sequence[ParameterEstimator])

• gamma (float)

• alpha (float)

• name (str | None)

• keys (str | None)

• len_estimator (ParameterEstimator | None)

accumulator_factory()

Returns a HierarchicalDirichletProcessMixtureAccumulatorFactory for this estimator.

Return type:

HierarchicalDirichletProcessMixtureAccumulatorFactory

model_log_density(model)

Log-density of the model parameters under the HDP priors.

Sums the Dirichlet(gamma/K) log-density of the global weights beta, the Dirichlet(alpha*beta) log-density of each fitted group’s weights (all floored at a tiny constant for boundary estimates), and each atom estimator’s model_log_density of its atom. Together with seq_local_elbo this forms the penalized objective maximized by the fit driver.

Parameters:

model (HierarchicalDirichletProcessMixtureDistribution)

Return type:

float

estimate(nobs, suff_stat)

Estimate a HierarchicalDirichletProcessMixtureDistribution.

Re-estimates each atom (whose conjugate update carries its posterior forward as its prior), updates the global weights beta via the expected-table-count approximation followed by the Dirichlet(gamma/K + m_.k) posterior mean, and sets each group’s weights to the Dirichlet(alpha*beta) posterior mean (deliberately the mean, not the MAP, which degenerates when alpha*beta_k < 1).

Parameters:

• nobs (Optional[float]) – Not used. Kept for the stats ParameterEstimator.estimate(nobs, suff_stat) signature.

• suff_stat (tuple) – Tuple (group_counts, prev_beta, prev_alpha, atom stats, len_value) as returned by HierarchicalDirichletProcessMixtureAccumulator.value().

Returns:

HierarchicalDirichletProcessMixtureDistribution object.

Return type:

HierarchicalDirichletProcessMixtureDistribution

class HierarchicalDirichletProcessMixtureDataEncoder(encoder, len_encoder=None)

Bases: DataSequenceEncoder

Encodes groups into a flat component encoding with per-group offsets.

Parameters:

• encoder (DataSequenceEncoder)

• len_encoder (DataSequenceEncoder | None)

seq_encode(x)

Encode groups into (lengths, offsets, flat_enc, len_enc).

Parameters:

x (Sequence[Sequence])

Return type:

Any