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 low-overhead operators never materialize A.

to_dict()

Return a safe JSON-compatible representation of this operator.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct an operator from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

TransitionOperator

forward(alpha)

Push a state-belief row vector one transition forward.

Parameters:

alpha (ndarray)

Return type:

ndarray

backward(v)

Pull an emission-weighted belief vector one transition backward.

Parameters:

v (ndarray)

Return type:

ndarray

as_matrix()

Materialize the transition matrix when an explicit matrix is needed.

Return type:

ndarray

new_accumulator()

Create an empty transition sufficient-statistic 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)

Estimate a transition operator of the same structural family from accumulated statistics.

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 HMMFitDiagnostics(algorithm, converged, iterations, termination_reason, log_likelihood_trace, initial_log_likelihood, final_log_likelihood, final_absolute_delta, final_relative_delta, monotone, n_sequences, total_weight, approximate)

Bases: object

Machine-readable receipt for a structured-family EM fit.

Parameters:

• algorithm (str)

• converged (bool)

• iterations (int)

• termination_reason (str)

• log_likelihood_trace (tuple[float, ...])

• initial_log_likelihood (float)

• final_log_likelihood (float)

• final_absolute_delta (float | None)

• final_relative_delta (float | None)

• monotone (bool)

• n_sequences (int)

• total_weight (float)

• approximate (bool)

class HMMFitResult(model, trace, diagnostics)

Bases: tuple

Backward-compatible (model, trace) result with a fit diagnostics receipt.

Parameters:

• model (Any)

• trace (list[float])

• diagnostics (HMMFitDiagnostics)

property model

Return the fitted model.

property log_likelihood_trace

Return the accepted model log-likelihood trajectory.

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)

Push a state-belief row vector forward with the dense matrix.

backward(v)

Pull a vector backward with the dense matrix.

as_matrix()

Return the dense row-stochastic transition matrix.

new_accumulator()

Create dense expected-transition-count storage.

accumulate(acc, alpha_t, w_next, scale)

Accumulate one dense expected-transition-count contribution.

estimate(acc)

Estimate a row-normalized dense transition from expected counts.

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)

Push a state-belief row vector through the low-rank transition.

backward(v)

Pull a vector backward through the low-rank transition.

as_matrix()

Materialize the implied dense transition matrix.

new_accumulator()

Create state-profile and profile-next sufficient-statistic storage.

accumulate(acc, alpha_t, w_next, scale)

Accumulate one exact low-rank transition contribution.

estimate(acc)

Estimate low-rank transition factors from accumulated statistics.

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)

Push a state-belief row vector through the sparse transition.

backward(v)

Pull a vector backward through the sparse transition.

as_matrix()

Materialize the sparse transition as a dense matrix.

new_accumulator()

Create one expected-count slot per allowed edge.

accumulate(acc, alpha_t, w_next, scale)

Accumulate one sparse expected-transition contribution.

estimate(acc)

Estimate a sparse transition over the same allowed edge set.

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)

Return the top k final-state-constrained sequences.

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()

Return the sequence enumerator for this HMM.

state_posteriors(seq)

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

seq_log_density(seqs)

Score a batch of observation sequences.

Return type:

ndarray

sampler(seed=None)

Return a sampler for observation sequences.

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

Fit by Baum-Welch through the transition operator.

fast=True uses the numba-jitted dense forward-backward when possible. The returned HMMFitResult remains unpackable as (fitted_hmm, loglik_trace) and also exposes a validated diagnostics receipt.

Parameters:

• max_its (int)

• tol (float)

• fast (bool)

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)

Push each block’s belief mass forward independently.

backward(v)

Pull each block’s vector backward independently.

as_matrix()

Materialize the block-diagonal transition matrix.

new_accumulator()

Create one transition accumulator per block.

accumulate(acc, alpha_t, w_next, scale)

Accumulate block-local expected transition mass.

estimate(acc)

Estimate each block and return a new block-diagonal transition.

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)

Push a joint belief through the Kronecker-factorized transition.

backward(v)

Pull a joint vector backward through the Kronecker transition.

as_matrix()

Materialize the dense Kronecker transition matrix.

new_accumulator()

Create marginal transition-count accumulators for both factors.

accumulate(acc, alpha_t, w_next, scale)

Accumulate exact marginalized factor transition statistics.

estimate(acc)

Estimate both Kronecker factors from marginalized expected counts.

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, weights=None)

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. The returned HMMFitResult remains unpackable as (fitted_hmm, loglik_trace) and marks the receipt’s objective as approximate.

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 sequences as lists without changing their observations.

row_count(x)

Return the number of pass-through structured-HMM records.

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)

Accumulate sufficient statistics from one weighted sequence.

seq_update(x, weights, estimate)

Run forward-backward and accumulate weighted sufficient statistics for a batch.

seq_initialize(x, weights, rng)

Initialize sufficient statistics with random soft state responsibilities.

combine(suff_stat)

Merge serialized HMM sufficient statistics.

value()

Return serialized HMM sufficient statistics.

from_value(x)

Restore accumulator state from serialized sufficient statistics.

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()

Return the encoder associated with this accumulator.

key_merge(store)

Merge tied initial or transition sufficient statistics into store.

key_replace(store)

Replace tied initial or transition statistics from store.

class StructuredHMMAccumulatorFactory(emission_estimators, transition_proto, keys)

Bases: StatisticAccumulatorFactory

Factory for structured HMM accumulators.

make()

Create a fresh structured HMM accumulator.

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()

Return the accumulator factory used by this estimator.

estimate(nobs, suff_stat)

Estimate an HMM from Baum-Welch sufficient statistics.

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).

viterbi(seq, inputs=None)

Most-likely state path (Viterbi / max-product), conditioned on the input/control sequence: step t -> t+1 maximizes over the transition inputs[t] selects. Call as viterbi(record) on one (observation, input)-pair sequence or as viterbi(obs_seq, input_seq). Uses the per-input transition matrices, so it works for any operator; O(T K^2) – a read-out, not the EM hot loop.

posterior_decode(seq, inputs=None)

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

state_posteriors(seq, inputs=None)

The full smoothing posteriors gamma[t,k] = P(z_t = k | x, u), conditioned on the inputs.

sampler(seed=None)

Return a sampler for (observation, input)-pair records along a given control sequence.

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

Fit the IOHMM by weighted Baum-Welch and return a diagnostics-bearing result.

Parameters:

• max_its (int)

• tol (float)

log_density(seq)

Return the log likelihood of one (observation, input) sequence.

dist_to_encoder()

Return the pass-through IOHMM sequence encoder.

estimator(pseudo_count=None)

Return the estimator for this IOHMM structure.

to_dict()

Return a safe JSON-compatible representation of this IOHMM (decodes via load_models).

Return type:

dict[str, Any]

class IOHMMDataEncoder

Bases: DataSequenceEncoder

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

seq_encode(x)

Encode IOHMM records as lists of (observation, input) pairs.

row_count(x)

Return the number of pass-through input-output HMM records.

class IOHMMAccumulator(emission_accumulators, transition_protos)

Bases: SequenceEncodableStatisticAccumulator

Accumulator for IOHMM Baum-Welch sufficient statistics.

update(x, weight, estimate)

Accumulate sufficient statistics from one weighted IOHMM record.

seq_update(x, weights, estimate)

Accumulate weighted sufficient statistics from a batch of IOHMM records.

seq_initialize(x, weights, rng)

Initialize IOHMM sufficient statistics with random soft responsibilities.

combine(suff_stat)

Merge serialized IOHMM sufficient statistics.

value()

Return serialized IOHMM sufficient statistics.

from_value(x)

Restore accumulator state from serialized IOHMM statistics.

scale(factor)

Scale IOHMM transition, emission, and initial-state statistics.

acc_to_encoder()

Return the encoder associated with this accumulator.

class IOHMMAccumulatorFactory(emission_estimators, transition_protos)

Bases: StatisticAccumulatorFactory

Factory for IOHMM accumulators.

make()

Create a fresh IOHMM accumulator.

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()

Return the accumulator factory used by this estimator.

estimate(nobs, suff_stat)

Estimate an IOHMM from Baum-Welch sufficient statistics.

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). A sequence has an exogenous fixed observation horizon. Its final latent duration is right-censored when it extends beyond that horizon, exactly matching sample(length). Forward/EM are O(T * K * max_duration).

forward_loglik(seq)

Fixed-horizon log likelihood, marginalizing a possibly right-censored final duration.

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

Fit the explicit-duration HMM by weighted Baum-Welch with a diagnostics receipt.

Parameters:

• max_its (int)

• tol (float)

log_density(seq)

Return the explicit-duration HMM log likelihood for one sequence.

seq_log_density(x)

Score a batch of observation sequences.

dist_to_encoder()

Return the pass-through explicit-duration sequence encoder.

estimator(pseudo_count=None)

Return the estimator for this explicit-duration HMM structure.

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’). Reading the first T expanded-chain emissions exactly represents the right-censored fixed horizon, so no final-state constraint is applied. This hands the HSMM the full StructuredHMM read-out API and, with len_dist, enumeration. O(K*D) states.

enumerator(len_dist)

Enumerate fixed-horizon observation sequences in descending marginal probability, given a len_dist over total sequence length. Built on the exact right-censored HMM expansion; .top_k(k) -> [(sequence, log_prob), …].

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, observed_duration) segments covering the sequence, O(T K D). The last tuple describes the observed portion of a right-censored latent duration.

sampler(seed=None)

Return a sampler for explicit-duration HMM sequences.

class EDHMMDataEncoder

Bases: DataSequenceEncoder

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

seq_encode(x)

Encode EDHMM records as observation-sequence lists.

row_count(x)

Return the number of pass-through explicit-duration HMM records.

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)

Accumulate sufficient statistics from one weighted sequence.

seq_update(x, weights, estimate)

Accumulate weighted explicit-duration HMM statistics from a batch.

seq_initialize(x, weights, rng)

Initialize explicit-duration HMM statistics with random soft responsibilities.

combine(suff_stat)

Merge serialized explicit-duration HMM sufficient statistics.

value()

Return serialized explicit-duration HMM sufficient statistics.

from_value(x)

Restore accumulator state from serialized EDHMM statistics.

scale(factor)

Scale explicit-duration HMM statistics.

acc_to_encoder()

Return the encoder associated with this accumulator.

class EDHMMAccumulatorFactory(emission_estimators, k, d)

Bases: StatisticAccumulatorFactory

Factory for explicit-duration HMM accumulators.

make()

Create a fresh explicit-duration HMM accumulator.

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()

Return the accumulator factory used by this estimator.

estimate(nobs, suff_stat)

Estimate an explicit-duration HMM from segment posterior statistics.

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.

Precision follows the caller’s ambient JAX policy rather than mixle’s: with the default jax_enable_x64=False the scan runs in float32, so the returned score agrees with seq_log_density to roughly float32 relative precision, not to float64. Enable x64 in your own process if you need more – mixle does not set it for you (MXR-080-0147 removed the import-time global mutation that used to make this path silently float64).

Parameters:

hmm (StructuredHMM)