mixle.stats.latent.structured_hmm module

Structured HMMs: a composable transition layer (dense / low-rank / combinators) + forward-backward.

A standard HMM stores a dense K x K transition matrix and the forward-backward does O(K^2) work per step. Rich structure – a low-rank transition, a block of independent chains, a factorial (Kronecker) product – is hard to express that way. This module factors the transition behind a small TransitionOperator interface so the forward-backward only needs two primitives:

forward(alpha) = alpha @ A (push a state-belief forward one step) backward(v) = A @ v (pull an emission-weighted belief back one step)

and an M-step that re-estimates the operator from expected transition mass. Any operator that implements those plugs into the SAME forward-backward / EM. Implementations:

• DenseTransition – the usual K x K matrix (O(K^2)).

• LowRankTransition – A = G @ Phi with an inner rank r (K x r, r x K row-stochastic): each state mixes over r shared “transition profiles”. Forward/backward and the M-step are O(K r), and the parameter count drops from K^2 to 2 K r. (Combinators – block-diagonal, Kronecker/factorial – are the same interface; see TransitionOperator subclasses.)

The forgetting / mixing property of an ergodic chain (beliefs forget the distant past) is what lets the forward-backward be split into chunks and run in parallel; see parallel in this package’s estimation.

class TransitionOperator

Bases: object

A row-stochastic state-transition operator behind the HMM forward-backward.

Subclasses provide the two linear maps the recursions need plus an M-step from expected transition mass. forward/backward must be consistent with as_matrix (forward(a) == a @ A, backward(v) == A @ v); the cheap operators never materialize A.

n_states: int

forward(alpha)

Parameters:

alpha (ndarray)

Return type:

ndarray

backward(v)

Parameters:

v (ndarray)

Return type:

ndarray

as_matrix()

Return type:

ndarray

new_accumulator()

Return type:

Any

accumulate(acc, alpha_t, w_next, scale)

Add one transition’s expected mass. alpha_t is the (normalized) forward belief at t, w_next = b_{t+1} * beta_{t+1} the emission-weighted backward belief at t+1, scale the forward normalizer c_{t+1} (so the per-step posterior transition mass is exact).

Parameters:

• acc (Any)

• alpha_t (ndarray)

• w_next (ndarray)

• scale (float)

Return type:

None

estimate(acc)

Parameters:

acc (Any)

Return type:

TransitionOperator

random_accumulator(rng)

A randomly-filled accumulator whose estimate yields a random (structured) transition – used to seed EM when there is no warm start. Fills new_accumulator shapes (nested) with noise.

Return type:

Any

class DenseTransition(a, prior=None)

Bases: TransitionOperator

The usual dense K x K row-stochastic transition (O(K^2) forward-backward).

prior (a K x K pseudocount matrix) is added to the expected counts before each M-step re-normalization – a Dirichlet/MAP transition. A diagonal prior is a sticky self-transition bias (see sticky_transition()); a flat prior is symmetric-Dirichlet smoothing.

Parameters:

• a (np.ndarray)

• prior (np.ndarray | None)

forward(alpha)

backward(v)

as_matrix()

new_accumulator()

accumulate(acc, alpha_t, w_next, scale)

Add one transition’s expected mass. alpha_t is the (normalized) forward belief at t, w_next = b_{t+1} * beta_{t+1} the emission-weighted backward belief at t+1, scale the forward normalizer c_{t+1} (so the per-step posterior transition mass is exact).

estimate(acc)

sticky_transition(a, kappa)

A dense transition with a STICKY self-transition prior: kappa pseudocounts on the diagonal favor staying in a state (longer dwell times, cleaner segmentation – the sticky-HMM idea).

Parameters:

kappa (float)

Return type:

DenseTransition

dirichlet_transition(a, alpha)

A dense transition with a symmetric Dirichlet(alpha) smoothing prior on every row (MAP).

Parameters:

alpha (float)

Return type:

DenseTransition

kron_initial(pi1, pi2)

Factorized initial distribution pi1 (x) pi2 for a factorial (Kronecker) HMM – the two chains start independently. Matches a KroneckerTransition so the joint initial respects the factors.

Return type:

ndarray

class LowRankTransition(g, phi)

Bases: TransitionOperator

A = G @ Phi: each state’s next-state distribution is a mix of r shared transition profiles.

G is K x r row-stochastic (state -> profile mixing), Phi is r x K row-stochastic (profile -> next-state). A = G @ Phi is K x K row-stochastic with rank <= r. Forward ((alpha @ G) @ Phi), backward (G @ (Phi @ v)) and the M-step are all O(K r) – never forming A – and the parameter count is 2 K r instead of K^2.

Parameters:

• g (np.ndarray)

• phi (np.ndarray)

forward(alpha)

backward(v)

as_matrix()

new_accumulator()

accumulate(acc, alpha_t, w_next, scale)

Add one transition’s expected mass. alpha_t is the (normalized) forward belief at t, w_next = b_{t+1} * beta_{t+1} the emission-weighted backward belief at t+1, scale the forward normalizer c_{t+1} (so the per-step posterior transition mass is exact).

estimate(acc)

class SparseTransition(n_states, edges, values=None)

Bases: TransitionOperator

Only the given (from, to) edges are allowed (left-to-right / banded HMMs). Forward, backward and the M-step are O(#edges) – transitions outside the edge set stay exactly zero through EM, so the structure is preserved. Build edges yourself or with left_to_right_edges() / banded_edges().

Parameters:

n_states (int)

forward(alpha)

backward(v)

as_matrix()

new_accumulator()

accumulate(acc, alpha_t, w_next, scale)

Add one transition’s expected mass. alpha_t is the (normalized) forward belief at t, w_next = b_{t+1} * beta_{t+1} the emission-weighted backward belief at t+1, scale the forward normalizer c_{t+1} (so the per-step posterior transition mass is exact).

estimate(acc)

left_to_right_edges(n_states, skip=1)

Edges for a left-to-right (Bakis) HMM: each state may stay or advance up to skip states.

Parameters:

• n_states (int)

• skip (int)

banded_edges(n_states, bandwidth=1)

Edges for a banded transition: state i connects to i-bandwidth .. i+bandwidth (local time-series).

Parameters:

• n_states (int)

• bandwidth (int)

class FinalStateEnumeration(hmm, len_dist)

Bases: object

Result of _final_state_enumerate(): top_k(k) -> [(sequence, log_prob), …] descending.

top_k(k)

class StructuredHMM(emissions, pi, transition, emission_estimators=None, keys=(None, None), name=None, len_dist=None, terminal_states=None, final_states=None)

Bases: object

An HMM whose transition is a TransitionOperator (dense / low-rank / a combinator).

emissions is one observation distribution per state; pi the initial-state distribution; transition any TransitionOperator. The scaled forward-backward and EM call the operator’s forward/backward/accumulate/estimate, so a low-rank or factorial transition runs the SAME inference at its own cost (O(K r) for low-rank). emission_estimators (one per state) drives the emission M-step; default reuses emissions[k].estimator().

Parameters:

transition (TransitionOperator)

viterbi(seq)

Most-likely state path (Viterbi / max-product). Uses the transition matrix, so it works for any operator; O(T K^2) – a read-out, not the EM hot loop.

posterior_decode(seq)

Per-position MAP state argmax_k P(z_t = k | x) from the forward-backward posteriors gamma.

enumerator()

Enumerate observation sequences in descending marginal probability (top_k / rank / seek / nucleus / certified estimates). Enumeration depends only on pi, the transition MATRIX, the emissions and a length distribution – not on the operator’s internal structure – so it reuses the built-in HMM enumerator (an A*-style best-first search over the trellis) on the dense matrix. Requires len_dist (a distribution over sequence length) and enumerable (discrete) emissions.

dist_to_enumerator()

state_posteriors(seq)

The full smoothing posteriors gamma[t,k] = P(z_t = k | x).

seq_log_density(seqs)

Return type:

ndarray

sampler(seed=None)

fit(seqs, *, max_its=50, tol=1e-6, fast=True)

EM (Baum-Welch) through the transition operator. Returns (fitted_hmm, loglik_trace).

fast=True uses the numba-jitted dense forward-backward (~30x over the numpy Python loop) when the transition is a plain DenseTransition with no terminal states; structured operators (low-rank / sparse / combinator) use the operator’s per-step accumulate as before.

Parameters:

• max_its (int)

• tol (float)

• fast (bool)

dist_to_encoder()

estimator(pseudo_count=None)

log_density(x)

class BlockDiagonalTransition(blocks)

Bases: TransitionOperator

Independent sub-chains: the states partition into blocks and transitions stay within a block.

A model whose initial state picks a block and then evolves inside it – a mixture of regimes that do not switch. Build it from any sub-operators (each block can itself be dense or low-rank). Exact, block-local forward-backward and M-step.

forward(alpha)

backward(v)

as_matrix()

new_accumulator()

accumulate(acc, alpha_t, w_next, scale)

Add one transition’s expected mass. alpha_t is the (normalized) forward belief at t, w_next = b_{t+1} * beta_{t+1} the emission-weighted backward belief at t+1, scale the forward normalizer c_{t+1} (so the per-step posterior transition mass is exact).

estimate(acc)

class KroneckerTransition(op1, op2)

Bases: TransitionOperator

Factorial HMM: the state is the pair (s1, s2) of two chains evolving in parallel, with A = A1 (x) A2 (Kronecker). State index is i1 * K2 + i2.

Forward-backward uses the reshape identity (alpha @ (A1 (x) A2) reshapes to A1^T @ M @ A2), so a step is O(K1 K2 (K1 + K2)) instead of O((K1 K2)^2) – the whole point of a factorial HMM. The E-step is exact over the joint state; the M-step is the standard factorial marginal update (each factor re-estimated from the marginalized joint transition mass), verified to keep EM monotone.

Parameters:

• op1 (TransitionOperator)

• op2 (TransitionOperator)

forward(alpha)

backward(v)

as_matrix()

new_accumulator()

accumulate(acc, alpha_t, w_next, scale)

Add one transition’s expected mass. alpha_t is the (normalized) forward belief at t, w_next = b_{t+1} * beta_{t+1} the emission-weighted backward belief at t+1, scale the forward normalizer c_{t+1} (so the per-step posterior transition mass is exact).

estimate(acc)

chunked_state_posteriors(hmm, seq, *, chunk, overlap)

State posteriors gamma for one long sequence via overlapping chunks, each run INDEPENDENTLY (embarrassingly parallel). The first chunk uses the model’s pi; interior chunks start from the uniform belief and the overlap context lets the chain forget that wrong boundary – so the kept interior matches the exact forward-backward up to an error that decays at the mixing rate in overlap.

Parameters:

• hmm (StructuredHMM)

• chunk (int)

• overlap (int)

Return type:

ndarray

fit_chunked(hmm, seqs, *, chunk, overlap, max_its=50, workers=0, tol=1e-6)

Baum-Welch where each long sequence’s forward-backward is split into overlapping chunks run in PARALLEL (the forgetting property bounds the boundary error). workers>0 runs the per-chunk E-steps on a thread pool (NumPy releases the GIL in its array kernels); workers=0 runs them serially. The interior suff-statistics are accumulated exactly as in StructuredHMM.fit(); this only changes how the E-step is computed, trading a small, overlap-controlled approximation for intra-sequence parallelism. Returns (fitted_hmm, loglik_trace) (LL is the chunk-summed approximation).

Parameters:

• hmm (StructuredHMM)

• chunk (int)

• overlap (int)

• max_its (int)

• workers (int)

• tol (float)

class StructuredHMMDataEncoder

Bases: DataSequenceEncoder

Sequences pass through as lists – the structured forward-backward scores raw observations through the per-state emission log_density (no flattened columnar encoding; composability over raw speed).

seq_encode(x)

Encode the iid observation sequence x for vectorized evaluation.

class StructuredHMMAccumulator(emission_accumulators, transition_proto, keys=(None, None))

Bases: SequenceEncodableStatisticAccumulator

Baum-Welch E-step accumulator: per-sequence forward-backward, accumulating initial-state mass, transition-operator mass, and per-state weighted emission statistics.

update(x, weight, estimate)

seq_update(x, weights, estimate)

seq_initialize(x, weights, rng)

combine(suff_stat)

value()

from_value(x)

scale(factor)

Multiply the running statistics by factor – the decay primitive online/streaming Baum-Welch (StreamingEstimator) uses to fold a new batch into a forgetting running estimate.

acc_to_encoder()

key_merge(store)

Pool this accumulator’s statistics into stats_dict under its merge key.

The structural default implements the common single-key pattern: store the accumulator under self.keys the first time the key is seen, else combine into the one already there. Accumulators with several named keys (e.g. an HMM’s init/trans/state keys) or a non-accumulator stats payload override this. A keys of None (the default) is a no-op.

key_replace(store)

Replace this accumulator’s statistics from the pooled stats_dict entry (see key_merge).

class StructuredHMMAccumulatorFactory(emission_estimators, transition_proto, keys)

Bases: StatisticAccumulatorFactory

make()

class StructuredHMMEstimator(emission_estimators, transition_proto, keys=(None, None), name=None, len_dist=None, terminal_states=None)

Bases: ParameterEstimator

Estimator (M-step) for a StructuredHMM: re-estimates pi, the transition OPERATOR (any structure – dense/low-rank/combinator), and each state’s emission from the Baum-Welch statistics. keys=(init_key, trans_key) tie the initial / transition parameters across HMMs that share them.

accumulator_factory()

estimate(nobs, suff_stat)

stationary_initial(op, *, iters=2000, tol=1e-13)

The transition’s stationary distribution (pi @ A == pi), by power iteration through op.forward – so it is O(K r) for a low-rank op, never forming A. Use it to COUPLE a StructuredHMM’s initial state to its transition (pi = stationary_initial(transition)): the chain starts in its long-run distribution instead of a free, separately-estimated pi. Answers “do the initial states match the transition?” – they can, by construction.

Parameters:

• op (TransitionOperator)

• iters (int)

• tol (float)

Return type:

ndarray

class InputOutputHMM(emissions, pi, transitions, emission_estimators=None, name=None, terminal_states=None)

Bases: object

Input-output HMM (IOHMM): an exogenous discrete input u_t selects which transition governs each step. Holds one TransitionOperator per input symbol; the emission is per-state. Data is (obs_seq, input_seq) pairs where input_seq[t] in {0..M-1} drives the transition from t to t+1.

Lets a covariate steer the dynamics – regime switching driven by an observed control, the difference between a plain HMM and a controlled Markov model. (Input-dependent emissions are a natural extension; here emissions depend on state only.)

seq_log_density(x, input_seqs=None)

Per-sequence forward log-likelihood. Two call forms: - seq_log_density(obs_seqs, input_seqs) – the explicit two-list API; or - seq_log_density(records) – one list of (obs, input)-pair sequences (the 5-part contract).

fit(obs_seqs, input_seqs, *, max_its=50, tol=1e-6)

Parameters:

• max_its (int)

• tol (float)

log_density(seq)

dist_to_encoder()

estimator(pseudo_count=None)

class IOHMMDataEncoder

Bases: DataSequenceEncoder

An IOHMM record is one (obs, input) sequence – a list of (observation, input_symbol) pairs.

seq_encode(x)

Encode the iid observation sequence x for vectorized evaluation.

class IOHMMAccumulator(emission_accumulators, transition_protos)

Bases: SequenceEncodableStatisticAccumulator

update(x, weight, estimate)

seq_update(x, weights, estimate)

seq_initialize(x, weights, rng)

combine(suff_stat)

value()

from_value(x)

acc_to_encoder()

class IOHMMAccumulatorFactory(emission_estimators, transition_protos)

Bases: StatisticAccumulatorFactory

make()

class IOHMMEstimator(emission_estimators, transition_protos, name=None)

Bases: ParameterEstimator

Estimator (M-step) for an InputOutputHMM: re-estimates pi, one transition operator per input symbol (from the per-input expected counts), and each state’s emission.

accumulator_factory()

estimate(nobs, suff_stat)

class ExplicitDurationHMM(emissions, pi, transition_matrix, durations, max_duration, name=None)

Bases: object

Hidden semi-Markov model (explicit-duration HMM): each state emits for a random duration drawn from a per-state duration distribution, then switches state (the transition matrix has a zero diagonal – dwell time is modeled explicitly, not as a self-loop). This captures non-geometric state durations a plain HMM cannot.

durations is one length-max_duration probability vector per state (over d = 1..max_duration). The forward variable alpha_t(j) = P(obs_1:t, a segment ends at t in state j); the likelihood is sum_j alpha_T(j). Forward/EM are O(T * K * max_duration). Verified against brute-force segmentation.

forward_loglik(seq)

Total log-likelihood log sum_j alpha_T(j) via the scaled explicit-duration forward.

fit(seqs, *, max_its=50, tol=1e-6)

Baum-Welch (EM) for the explicit-duration HMM: re-estimates emissions, the per-state duration distributions, the (zero-diagonal) transition, and pi. Returns (fitted_hmm, loglik_trace).

Parameters:

• max_its (int)

• tol (float)

log_density(seq)

seq_log_density(x)

dist_to_encoder()

estimator(pseudo_count=None)

to_structured_hmm(len_dist=None)

The HSMM as an EQUIVALENT StructuredHMM via the remaining-duration expansion: K*D sub-states (k, r) = “state k with r steps left in the segment”. The expanded chain emits from state k at every sub-state, decrements deterministically (k,r)->(k,r-1), and at (k,1) switches segment with A[k,k’]*dur[k’](d’). final_states = the (k,1) sub-states require the last segment to COMPLETE, so the expanded forward log-likelihood EQUALS this EDHMM’s exactly. This hands the HSMM the full StructuredHMM read-out API – Viterbi (recover state+remaining-duration), posterior decoding, the standard forward – and, with len_dist, enumeration. O(K*D) states.

enumerator(len_dist)

Enumerate observation sequences in descending marginal probability under this HSMM (complete final segment), given a len_dist over total sequence length. Built on the exact HMM expansion + the final-state best-first enumerator; .top_k(k) -> [(sequence, log_prob), …]. Needs discrete (Categorical) emissions and a Categorical-like len_dist.

state_posteriors(seq)

Per-position smoothing posteriors gamma[t, j] = P(z_t = j | obs), marginalizing the durations (sum the posterior of every segment that covers position t). Rows sum to 1.

posterior_decode(seq)

Per-position MAP state argmax_j P(z_t = j | obs).

viterbi_segments(seq)

Most-likely segmentation (max-product over the segment lattice): a list of (state, start, duration) segments covering the sequence, O(T K D). The HSMM analog of Viterbi decoding.

sampler(seed=None)

class EDHMMDataEncoder

Bases: DataSequenceEncoder

An ExplicitDurationHMM record is one observation sequence (the durations are latent).

seq_encode(x)

Encode the iid observation sequence x for vectorized evaluation.

class EDHMMAccumulator(emission_accumulators, k, d)

Bases: SequenceEncodableStatisticAccumulator

E-step accumulator for an explicit-duration HMM: per-sequence segment posteriors -> initial / transition / per-state DURATION counts + emission occupancy statistics.

update(x, weight, estimate)

seq_update(x, weights, estimate)

seq_initialize(x, weights, rng)

combine(suff_stat)

value()

from_value(x)

acc_to_encoder()

class EDHMMAccumulatorFactory(emission_estimators, k, d)

Bases: StatisticAccumulatorFactory

make()

class EDHMMEstimator(emission_estimators, k, d, name=None)

Bases: ParameterEstimator

Estimator (M-step) for an ExplicitDurationHMM: re-estimates pi, the zero-diagonal transition, the per-state DURATION distributions, and each state’s emission from the segment-posterior statistics.

accumulator_factory()

estimate(nobs, suff_stat)

jit_forward_loglik(hmm)

Compile the scaled forward log-likelihood recursion to a single jax.jit XLA program (lax.scan over time). Returns a callable score(seq) -> float: emission log-densities are evaluated on the host (arbitrary emissions), then the forward scan runs jitted on the transition matrix. Works for any operator (uses as_matrix()); the win is large T / K. Requires the JAX optional extra.

Parameters:

hmm (StructuredHMM)