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)[source]

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()[source]

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

Return type:

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

set_parameters(params)[source]

Set the parameters and refresh the cached log-weights.

Parameters:

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

Return type:

None

density(x)[source]

Density of a group x; see log_density().

Parameters:

x (Any)

Return type:

float

density_semantics()[source]

Return density semantics for the expected-weight HDP mixture approximation.

log_density(x)[source]

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

Parameters:

x (Any)

Return type:

float

seq_encode(x)[source]

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) consumed by vectorized seq_* methods.

Return type:

Any

seq_log_density(x)[source]

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)[source]

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)[source]

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)[source]

Create a HierarchicalDirichletProcessMixtureSampler for this distribution.

Parameters:

seed (int | None)

Return type:

HierarchicalDirichletProcessMixtureSampler

estimator(pseudo_count=None)[source]

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

Parameters:

pseudo_count (float | None)

Return type:

HierarchicalDirichletProcessMixtureEstimator

dist_to_encoder()[source]

Returns a HierarchicalDirichletProcessMixtureDataEncoder for this distribution.

Return type:

HierarchicalDirichletProcessMixtureDataEncoder

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

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)[source]

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)[source]

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)[source]

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)[source]

Initialize with random Dirichlet assignments for group x.

Parameters:
Return type:

None

seq_initialize(x, weights, rng)[source]

Vectorized initialize() with random Dirichlet assignments.

Parameters:
Return type:

None

update(x, weight, estimate)[source]

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)[source]

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)[source]

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

Parameters:

suff_stat (tuple)

Return type:

HierarchicalDirichletProcessMixtureAccumulator

scale(c)[source]

Scale linear HDP mixture sufficient statistics while preserving metadata.

Parameters:

c (float)

Return type:

HierarchicalDirichletProcessMixtureAccumulator

value()[source]

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

Return type:

tuple

from_value(x)[source]

Set the sufficient statistics from a value() tuple.

Parameters:

x (tuple)

Return type:

HierarchicalDirichletProcessMixtureAccumulator

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 HierarchicalDirichletProcessMixtureDataEncoder for this accumulator.

Return type:

HierarchicalDirichletProcessMixtureDataEncoder

class HierarchicalDirichletProcessMixtureAccumulatorFactory(factories, len_factory, name, keys)[source]

Bases: StatisticAccumulatorFactory

Factory for HDP-mixture sufficient-statistic accumulators.

Parameters:

  • factories (Sequence[StatisticAccumulatorFactory])

  • len_factory (StatisticAccumulatorFactory | None)

  • name (str | None)

  • keys (str | None)

make()[source]

Returns a new HierarchicalDirichletProcessMixtureAccumulator.

Return type:

HierarchicalDirichletProcessMixtureAccumulator

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

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()[source]

Returns a HierarchicalDirichletProcessMixtureAccumulatorFactory for this estimator.

Return type:

HierarchicalDirichletProcessMixtureAccumulatorFactory

model_log_density(model)[source]

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)[source]

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:

Fitted hierarchical Dirichlet-process mixture approximation.

Return type:

HierarchicalDirichletProcessMixtureDistribution

class HierarchicalDirichletProcessMixtureDataEncoder(encoder, len_encoder=None)[source]

Bases: DataSequenceEncoder

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

Parameters:

  • encoder (DataSequenceEncoder)

  • len_encoder (DataSequenceEncoder | None)

seq_encode(x)[source]

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

Parameters:

x (Sequence[Sequence])

Return type:

Any