mixle.stats.latent.tree_hidden_markov_model module

Create, estimate, and sample from a hidden Markov model on a rooted tree.

Defines the TreeHiddenMarkovModelDistribution, TreeHiddenMarkovSampler, TreeHiddenMarkovAccumulatorFactory, TreeHiddenMarkovAccumulator, TreeHiddenMarkovEstimator, and the TreeHiddenMarkovDataEncoder classes for use with mixle.

Data type: Sequence[Tuple[Tuple[int, int], T]]. A single observation is a rooted tree given as a sequence of ((node_id, parent_id), emission) tuples, where node_id is an integer in 0,1,…,n-1, the root node has parent_id = -1, and the emission has the data type T of the emission (topic) distributions.

Each node u carries a hidden state Z_u in {0,…,K-1}. The root state is drawn from the initial state weights w, the state of every child is drawn from the transition probability matrix A conditioned on its parent’s state (children of a node are conditionally independent given the parent state), and the observed emission at node u is drawn from topics[Z_u]. The number of children of each node is modeled by an optional length distribution (len_dist) with support on the non-negative integers, which is required for sampling.

Inference uses an upward-downward (beta/eta) message passing recursion over tree levels. Two equivalent vectorized implementations are provided and selected with the use_numba flag: numba-compiled kernels that parallelize over trees, and a pure-numpy implementation that batches nodes level by level.

find_level(parents)

Find the level in the tree for nodes, given an array of parents.

Parameters:

parents (np.ndarray) – Numpy array of integers with first entry -1.

Returns:

Level of each node in the free excluding the first entry which is the root (level = 0).

Return type:

list[int]

class TreeHiddenMarkovModelDistribution(topics, w=MISSING, transitions=MISSING, len_dist=NullDistribution(), terminal_level=10, name=None, use_numba=False, weights=MISSING)

Bases: SequenceEncodableProbabilityDistribution

Hidden Markov model on a rooted tree with emission distributions of type T.

Data type: Sequence[Tuple[Tuple[int, int], T]] (((node_id, parent_id), emission) per node, root parent -1).

Parameters:

• topics (Sequence[SequenceEncodableProbabilityDistribution])

• w (Sequence[float] | ndarray)

• transitions (list[list[float]] | ndarray)

• len_dist (SequenceEncodableProbabilityDistribution | None)

• terminal_level (int)

• name (str | None)

• use_numba (bool)

• weights (Sequence[float] | ndarray)

compute_capabilities()

compute_declaration()

density(x)

Density of tree HMM at a single observed tree x.

See log_density() for details.

Parameters:

x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

Returns:

Density at observation x.

Return type:

float

log_density(x)

Log-density of tree HMM at a single observed tree x.

The hidden states are marginalized out with an upward (beta) message passing recursion over the tree. When a non-null length distribution (len_dist) is set, its contribution to the likelihood (the sum over nodes of len_dist.log_density(num_children)) is included; a NullDistribution length contributes nothing.

Parameters:

x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

Returns:

Log-density at observation x.

Return type:

float

seq_log_density(x)

Vectorized evaluation of log-density at sequence encoded input x.

Dispatches to numba kernels or to the pure-numpy level-by-level recursion depending on which encoding (use_numba) the input was created with.

Parameters:

x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

Returns:

Numpy array of log-density (float), one entry per encoded tree.

Return type:

ndarray

backend_seq_log_density(x, engine)

Engine-neutral tree-HMM scoring for the pure non-numba encoded layout.

Parameters:

• x (tuple[tuple[ndarray, tuple[int, ndarray, ndarray, ndarray], tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray, ndarray], E3, tuple[ndarray, E4] | None] | None, tuple[int, ndarray, tuple[ndarray, ndarray, ndarray], tuple[ndarray, ndarray, ndarray, ndarray, list[ndarray], list[ndarray], list[ndarray], list[ndarray], ndarray], E3, tuple[ndarray, E4] | None] | None])

• engine (Any)

Return type:

Any

seq_posterior(x)

Posterior state membership probabilities for each node of each encoded tree.

Parameters:

x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

Returns:

List with one numpy array per tree, each of shape (num_nodes, num_states), containing the posterior probability of each hidden state at each node.

Return type:

list[ndarray] | None

viterbi(x)

Most likely hidden state assignment for each node of a single observed tree.

Parameters:

x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

Returns:

Numpy array of int states, one per node of the tree.

Return type:

ndarray

seq_viterbi(x)

Vectorized Viterbi state assignments for sequence encoded trees.

Parameters:

x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

Returns:

List with one numpy array of int states per tree (one state per node).

Return type:

list[ndarray]

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.

sampler(seed=None)

Create a TreeHiddenMarkovSampler object from parameters of distribution instance.

Parameters:

seed (Optional[int]) – Used to set seed in random sampler.

Returns:

TreeHiddenMarkovSampler object.

Raises:

Exception – If len_dist is a NullDistribution (a length distribution with support on the non-negative integers is required for sampling).

Return type:

TreeHiddenMarkovSampler

enumerator()

Not supported: observations are rooted trees, not chains.

The chain best-first enumerator (HiddenMarkovModelEnumerator) does not apply – the marginal sums over hidden states on a branching tree whose shape is itself governed by the per-node child-count len_dist, so the support is a set of trees rather than sequences. Use sampler() or the exact log_density instead.

Return type:

DistributionEnumerator

estimator(pseudo_count=None)

Create a TreeHiddenMarkovEstimator with estimators for the topics and length distribution.

Parameters:

pseudo_count (Optional[float]) – Used to inflate sufficient statistics of the initial state weights, transition matrix, topics, and length distribution.

Returns:

TreeHiddenMarkovEstimator object.

Return type:

TreeHiddenMarkovEstimator

dist_to_encoder()

Returns TreeHiddenMarkovDataEncoder object for encoding sequences of iid Tree HMM observations.

Return type:

TreeHiddenMarkovDataEncoder

class TreeHiddenMarkovSampler(dist, seed=None)

Bases: DistributionSampler

Sampler for the TreeHiddenMarkovModelDistribution. Draws rooted trees of state-emitted observations.

Parameters:

• dist (TreeHiddenMarkovModelDistribution)

• seed (int | None)

sample_state(given_state, size=None)

Draw child state(s) from the transition matrix row of a given parent state.

Parameters:

• given_state (int) – State of the parent node.

• size (Optional[int]) – Number of child states to draw. If None, a single int is returned.

Returns:

Single state (int) if size is None, else numpy array of size states.

Return type:

int | ndarray

sample_tree(size=None)

Sample rooted tree(s) of ((node_id, parent_id), emission) tuples.

The root state is drawn from the initial weights, each child state from the transition matrix, and the number of children of each node from the length sampler. Sampling stops once the terminal_level of the distribution is reached.

Parameters:

size (Optional[int]) – Number of trees to draw. If None, a single tree is returned.

Returns:

A single tree (list of ((node_id, parent_id), emission) tuples with root parent -1) if size is None, else a list of size trees.

sample(size=None)

Draw iid tree observations from the tree HMM.

Parameters:

size (Optional[int]) – Number of trees to draw. If None, a single tree is returned.

Returns:

A single tree if size is None, else a list of size trees. See sample_tree().

Raises:

RuntimeError – If no length sampler is available for the number of children.

class TreeHiddenMarkovAccumulator(accumulators, len_accumulator=NullAccumulator(), keys=(None, None, None), name=None, use_numba=True)

Bases: SequenceEncodableStatisticAccumulator

Accumulator for the tree HMM. Tracks initial-state, state, and transition counts plus emission stats.

Parameters:

• accumulators (Sequence[SequenceEncodableStatisticAccumulator])

• len_accumulator (SequenceEncodableStatisticAccumulator | None)

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

• name (str | None)

• use_numba (bool)

update(x, weight, estimate)

Update sufficient statistics with a single weighted tree observation.

Encodes the tree and delegates to seq_update().

Parameters:

• x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

• weight (float) – Weight for observation.

• estimate (TreeHiddenMarkovModelDistribution) – Previous estimate used for the E-step.

Return type:

None

initialize(x, weight, rng)

Initialize sufficient statistics with a single weighted tree observation.

Parameters:

• x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

• weight (float) – Weight for observation.

• rng (RandomState) – Random number generator used to draw random state assignments.

Return type:

None

seq_initialize(x, weights, rng)

Vectorized initialization of sufficient statistics from sequence encoded trees.

Hidden states are assigned uniformly at random and the initial-state, state, transition, emission, and length statistics are accumulated under that random assignment.

Parameters:

• x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

• weights (np.ndarray) – Weights for each encoded tree.

• rng (np.random.RandomState) – Random number generator used to draw state assignments.

Return type:

None

seq_update(x, weights, estimate)

Vectorized E-step update of sufficient statistics from sequence encoded trees.

Runs the upward-downward (Baum-Welch style) recursion under the previous estimate to obtain expected initial-state, state membership, and transition counts, and passes the node posteriors as weights to the emission accumulators (and tree weights to the length accumulator).

Parameters:

• x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

• weights (np.ndarray) – Weights for each encoded tree.

• estimate (TreeHiddenMarkovModelDistribution) – Previous estimate used for the E-step.

Return type:

None

seq_update_engine(x, weights, estimate, engine)

Engine-resident E-step for the pure (non-numba) tree encoding.

Runs the upward-downward recursion entirely with engine ops (numpy or torch): emission scoring, the level-by-level beta/eta upward pass, and the alpha/xi downward pass all live on the active engine. Node posteriors and aggregated initial/state/transition counts are produced on the engine and converted to host arrays only to feed the child accumulators. Mirrors the numpy branch of seq_update.

Parameters:

• x (tuple[tuple[ndarray, tuple[int, ndarray, ndarray, ndarray], tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray, ndarray], E3, tuple[ndarray, E4] | None] | None, tuple[int, ndarray, tuple[ndarray, ndarray, ndarray], tuple[ndarray, ndarray, ndarray, ndarray, list[ndarray], list[ndarray], list[ndarray], list[ndarray], ndarray], E3, tuple[ndarray, E4] | None] | None])

• weights (ndarray)

• estimate (TreeHiddenMarkovModelDistribution)

• engine (Any)

Return type:

None

combine(suff_stat)

Combine sufficient statistics from another accumulator into this one.

Parameters:

suff_stat (Tuple) – Tuple of number of states, initial-state counts, state counts, transition counts, emission sufficient statistics per state, and length sufficient statistics.

Returns:

Self, with aggregated sufficient statistics.

Return type:

TreeHiddenMarkovAccumulator

value()

Returns sufficient statistics as a Tuple of number of states, initial-state counts, state counts, transition counts, emission sufficient statistics per state, and length statistics.

Return type:

tuple[int, ndarray, ndarray, ndarray, Sequence[Any], Any | None]

from_value(x)

Set sufficient statistics of accumulator from value x.

Parameters:

x (Tuple) – Tuple of number of states, initial-state counts, state counts, transition counts, emission sufficient statistics per state, and length sufficient statistics.

Returns:

Self, with sufficient statistics set to x.

Return type:

TreeHiddenMarkovAccumulator

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:

TreeHiddenMarkovAccumulator

key_merge(stats_dict)

Merge keyed sufficient statistics into stats_dict (and recurse into member accumulators).

Parameters:

stats_dict (Dict[str, Any]) – Dictionary mapping keys to shared sufficient statistics.

Return type:

None

key_replace(stats_dict)

Replace keyed sufficient statistics with values from stats_dict (and recurse into members).

Parameters:

stats_dict (Dict[str, Any]) – Dictionary mapping keys to shared sufficient statistics.

Return type:

None

acc_to_encoder()

Returns a TreeHiddenMarkovDataEncoder object for encoding sequences of tree HMM observations.

Return type:

TreeHiddenMarkovDataEncoder

class TreeHiddenMarkovAccumulatorFactory(factories, len_factory=NullAccumulatorFactory(), keys=(None, None, None), name=None, use_numba=True)

Bases: StatisticAccumulatorFactory

Factory for creating TreeHiddenMarkovAccumulator objects.

Parameters:

• factories (Sequence[StatisticAccumulatorFactory])

• len_factory (StatisticAccumulatorFactory)

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

• name (str | None)

• use_numba (bool)

make()

Returns a new TreeHiddenMarkovAccumulator object.

Return type:

TreeHiddenMarkovAccumulator

class TreeHiddenMarkovEstimator(estimators, len_estimator=NullEstimator(), pseudo_count=(None, None), name=None, keys=(None, None, None), use_numba=True)

Bases: ParameterEstimator

Estimator for the TreeHiddenMarkovModelDistribution from aggregated sufficient statistics.

Parameters:

• estimators (list[ParameterEstimator])

• len_estimator (ParameterEstimator | None)

• pseudo_count (tuple[float | None, float | None] | None)

• name (str | None)

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

• use_numba (bool)

accumulator_factory()

Returns a TreeHiddenMarkovAccumulatorFactory for creating TreeHiddenMarkovAccumulator objects.

Return type:

TreeHiddenMarkovAccumulatorFactory

estimate(nobs, suff_stat)

Estimate a TreeHiddenMarkovModelDistribution from sufficient statistics (M-step).

Initial state weights and transition rows are normalized counts, optionally smoothed with the pseudo_count pair. Rows of the transition matrix with no observed transitions are left as zeros when no pseudo-count is given. Topics and the length distribution are estimated from their respective sufficient statistics.

Parameters:

• nobs (Optional[float]) – Number of observations.

• suff_stat (Tuple) – Tuple of number of states, initial-state counts, state counts, transition counts, emission sufficient statistics per state, and length sufficient statistics.

Returns:

TreeHiddenMarkovModelDistribution object.

Return type:

TreeHiddenMarkovModelDistribution

class TreeHiddenMarkovDataEncoder(emission_encoder, len_encoder=NullDataEncoder(), use_numba=True)

Bases: DataSequenceEncoder

Data encoder for sequences of tree HMM observations (flattens trees into level-indexed arrays).

Parameters:

• emission_encoder (DataSequenceEncoder)

• len_encoder (DataSequenceEncoder | None)

• use_numba (bool)

seq_encode(x)

Encode a sequence of tree observations for vectorized functions.

Returns a pair (numba_encoding, numpy_encoding) with exactly one entry set, depending on use_numba: the numba encoding stores per-tree slice offsets consumed by the numba kernels, while the numpy encoding (see _seq_encode()) batches nodes across trees by level.

Parameters:

x (Sequence[Sequence[Tuple[D, T]]]) – Sequence of trees of ((node_id, parent_id), emission) tuples (root parent -1).

Returns:

Tuple (E) with the numba encoding in slot 0 (numpy slot None), or vice versa.

Return type:

tuple[tuple[ndarray, tuple[int, ndarray, ndarray, ndarray], tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray, ndarray], Any, tuple[ndarray, Any] | None] | None, tuple[int, ndarray, tuple[ndarray, ndarray, ndarray], tuple[ndarray, ndarray, ndarray, ndarray, list[ndarray], list[ndarray], list[ndarray], list[ndarray], ndarray], Any, tuple[ndarray, Any] | None] | None]

numba_seq_log_density(num_states, tz, txz, tp, tpz, tlnz, xp, xc, xl, xbi, xln, xlnl, pr_obs, p_level, tr_mat, pr_max0, betas, etas, out)

Numba kernel: per-tree log-likelihood via the upward (beta) recursion, written to out.

numba_baum_welch(num_states, tz, txz, tp, tpz, tlnz, xp, xc, xl, xbi, xln, xlnl, pr_obs, p_level, tr_mat, weights, betas, etas, alphas, xi_acc, pi_acc)

Numba kernel: upward-downward E-step writing node posteriors (alphas), per-tree expected transition counts (xi_acc), and root-state counts (pi_acc).

numba_posteriors(num_states, tz, txz, tp, tpz, tlnz, xp, xc, xl, xbi, xln, xlnl, pr_obs, p_level, tr_mat, betas, etas)

Numba kernel: upward (beta) recursion writing normalized per-node state posteriors into betas.

numba_initialize(tz, txz, tp, tpz, xp, xc, states, weights, init_counts, state_counts, trans_counts)

Numba kernel: accumulate initial-state, state, and transition counts for random state assignments.

numba_viterbi(num_states, tz, txz, tp, tpz, tlnz, xp, xc, xl, xbi, xln, log_pr_obs, log_init_p, log_tr_mat, betas, etas, out)

Numba kernel: max-product upward recursion writing the most likely state per node into out.

vec_bincount(idx, ll, out)

Numba kernel: scatter-add ll into out at positions idx (weighted bincount).

level_state_prob(levels, num_states, tr_mat, init_prob, out)

Numba kernel: marginal state probabilities per tree level, out[k] = init_prob @ tr_mat^k.

TreeHiddenMarkovModelAccumulator

alias of TreeHiddenMarkovAccumulator

TreeHiddenMarkovModelAccumulatorFactory

alias of TreeHiddenMarkovAccumulatorFactory

TreeHiddenMarkovModelDataEncoder

alias of TreeHiddenMarkovDataEncoder

TreeHiddenMarkovModelEstimator

alias of TreeHiddenMarkovEstimator

TreeHiddenMarkovModelSampler

alias of TreeHiddenMarkovSampler