"""Hidden Markov model over arbitrary emission segments.
Each hidden state emits one *segment* object. The segment can be any data type
accepted by that state's emission distribution: a scalar, tuple, set, sequence,
or another composable mixle.stats object. To model variable-length emissions,
use a ``SequenceDistribution`` (or any other distribution over sequences) as an
emission distribution.
Unlike ``HiddenMarkovModelDistribution``, each state keeps its own encoder, so
emission distributions may use different distribution classes as long as they
can all score the same raw segment observations.
"""
import math
from collections.abc import Sequence
from numbers import Real
from typing import Any
import numpy as np
from numpy.random import RandomState
from scipy.special import logsumexp
from mixle.capability import Neutral, supports
from mixle.engines.arithmetic import maxrandint
from mixle.stats.combinator.null_dist import (
NullAccumulator,
NullAccumulatorFactory,
NullDataEncoder,
NullDistribution,
NullEstimator,
)
from mixle.stats.compute.mixture_evidence import validated_probability_vector
from mixle.stats.compute.pdist import (
DataSequenceEncoder,
DistributionEnumerator,
DistributionSampler,
ParameterEstimator,
SequenceEncodableProbabilityDistribution,
SequenceEncodableStatisticAccumulator,
StatisticAccumulatorFactory,
)
from mixle.stats.latent._initialization import broken_symmetry_states, derived_rng
from mixle.stats.latent.effective_sample import (
validate_effective_sample_mass,
validated_count_array,
validated_observation_weight,
validated_observation_weights,
validated_positive_integer,
validated_statistic_tuple,
)
# The segmental accumulator shares the chain HMM's exact statistic geometry (initial/state/
# transition counts, one emission-accumulator list, a length child, and the init/trans/state
# key trio), so its mutators delegate to the family-shared transactional implementation --
# the latent-family mutator audit measured every pre-repair defect class alive in the local
# copies, which is precisely the sibling-bug hazard the shared home exists to end.
from mixle.stats.latent.hidden_markov import (
_keyed_statistics_merge,
_keyed_statistics_replace,
_transactional_combine,
_transactional_restore,
_transactional_scale,
_validate_state_mass,
)
from mixle.stats.latent.markov_stopping import (
DEFAULT_TERMINAL_STEP_CAP,
require_terminal_reached,
validate_terminal_reachability,
validated_terminal_states,
validated_terminal_step_cap,
)
from mixle.stats.sequences.markov_chain import MarkovChainDistribution
from mixle.utils.aliasing import MISSING, broadcast_pseudo_count, coalesce_alias, require
from mixle.utils.vector import require_possible_log_evidence
def _owned_segmental_sequence(values: Any, label: str, *, size: int | None = None) -> list[Any]:
"""Return an owned non-empty sequence with optional exact arity."""
try:
owned = list(values)
except TypeError as exc:
raise TypeError(f"{label} must be an iterable.") from exc
if not owned:
raise ValueError(f"{label} must not be empty.")
if size is not None and len(owned) != size:
raise ValueError(f"{label} must contain exactly {size} entries.")
return owned
def _validated_segmental_transitions(values: Any, n_states: int) -> tuple[np.ndarray, tuple[int, ...]]:
"""Return an owned stochastic transition matrix and explicit zero rows."""
try:
transitions = np.asarray(values, dtype=np.float64)
except (TypeError, ValueError, OverflowError) as exc:
raise TypeError("segmental HMM transitions must be a numeric matrix.") from exc
expected = (n_states, n_states)
if transitions.shape != expected:
raise ValueError(f"segmental HMM transitions must have shape {expected}, got {transitions.shape}.")
if np.any(~np.isfinite(transitions)) or np.any(transitions < 0.0):
raise ValueError("segmental HMM transitions must contain finite non-negative probabilities.")
row_sums = transitions.sum(axis=1)
zero_rows = np.flatnonzero(row_sums == 0.0)
nonzero = row_sums != 0.0
if not np.allclose(row_sums[nonzero], 1.0, rtol=1.0e-10, atol=1.0e-12):
raise ValueError("segmental HMM transition rows must sum to one.")
return transitions.copy(), tuple(int(index) for index in zero_rows)
def _reachable_segmental_states(initial: np.ndarray, transitions: np.ndarray) -> np.ndarray:
"""Return states graph-reachable from positive initial mass."""
reachable = initial > 0.0
frontier = list(np.flatnonzero(reachable))
while frontier:
state = frontier.pop()
for child in np.flatnonzero(transitions[state] > 0.0):
if not reachable[child]:
reachable[child] = True
frontier.append(int(child))
return reachable
def _validated_segmental_pseudo_count(value: Any) -> tuple[float | None, float | None]:
"""Return finite non-negative initial/transition pseudo-counts."""
value = broadcast_pseudo_count(value, 2)
if value is None:
return (None, None)
try:
controls = tuple(value)
except TypeError as exc:
raise TypeError("SegmentalHiddenMarkovEstimator pseudo_count must be a scalar or pair.") from exc
if len(controls) != 2:
raise ValueError("SegmentalHiddenMarkovEstimator pseudo_count must contain exactly two entries.")
result: list[float | None] = []
for index, control in enumerate(controls):
if control is None:
result.append(None)
elif isinstance(control, bool) or not isinstance(control, Real):
raise TypeError(f"SegmentalHiddenMarkovEstimator pseudo_count[{index}] must be real or None.")
elif not np.isfinite(float(control)) or float(control) < 0.0:
raise ValueError(f"SegmentalHiddenMarkovEstimator pseudo_count[{index}] must be finite and non-negative.")
else:
result.append(float(control))
return result[0], result[1]
def _validated_segmental_count_array(values: Any, label: str, shape: tuple[int, ...]) -> np.ndarray:
"""Return an owned finite non-negative sufficient-statistic array."""
try:
counts = np.asarray(values, dtype=np.float64)
except (TypeError, ValueError, OverflowError) as exc:
raise TypeError(f"{label} must be a numeric array.") from exc
if counts.shape != shape:
raise ValueError(f"{label} must have shape {shape}, got {counts.shape}.")
if np.any(~np.isfinite(counts)) or np.any(counts < 0.0):
raise ValueError(f"{label} must contain finite non-negative values.")
return counts.copy()
def _validated_segment_count(value: Any) -> int:
"""Return an exact non-negative sampled segment count."""
if isinstance(value, bool) or not isinstance(value, (int, np.integer)):
raise TypeError("segmental HMM sampled length must be an integer.")
result = int(value)
if result < 0:
raise ValueError("segmental HMM sampled length must be non-negative.")
return result
def _forward_log(log_w: np.ndarray, log_a: np.ndarray, log_emit: np.ndarray) -> float:
"""Return log p(emissions) for one segment sequence."""
n = log_emit.shape[0]
if n == 0:
return 0.0
alpha = log_w + log_emit[0]
for t in range(1, n):
alpha = log_emit[t] + logsumexp(alpha[:, None] + log_a, axis=0)
return float(logsumexp(alpha))
def _forward_backward(
log_w: np.ndarray, log_a: np.ndarray, log_emit: np.ndarray
) -> tuple[float, np.ndarray, np.ndarray]:
"""Forward/backward posterior probabilities for one segment sequence.
Returns (log_likelihood, gamma, xi_sum), where gamma[t, k] is p(z_t=k|x)
and xi_sum[i, j] is sum_t p(z_t=i, z_{t+1}=j|x).
"""
n, k = log_emit.shape
gamma = np.zeros((n, k), dtype=np.float64)
xi_sum = np.zeros((k, k), dtype=np.float64)
if n == 0:
return 0.0, gamma, xi_sum
alpha = np.empty((n, k), dtype=np.float64)
beta = np.zeros((n, k), dtype=np.float64)
alpha[0] = log_w + log_emit[0]
for t in range(1, n):
alpha[t] = log_emit[t] + logsumexp(alpha[t - 1][:, None] + log_a, axis=0)
ll = float(logsumexp(alpha[-1]))
if not np.isfinite(ll):
return ll, gamma, xi_sum
for t in range(n - 2, -1, -1):
beta[t] = logsumexp(log_a + log_emit[t + 1][None, :] + beta[t + 1][None, :], axis=1)
gamma = np.exp(alpha + beta - ll)
for t in range(n - 1):
xi = alpha[t][:, None] + log_a + log_emit[t + 1][None, :] + beta[t + 1][None, :] - ll
xi_sum += np.exp(xi)
return ll, gamma, xi_sum
[docs]
class SegmentalHiddenMarkovModelDistribution(SequenceEncodableProbabilityDistribution):
"""HMM whose states emit arbitrary segment-valued distributions.
Observations are lists of segment objects. For example, with
``SequenceDistribution(GaussianDistribution(...), len_dist=...)`` as an
emission, each state emits a variable-length list of real values.
"""
[docs]
def compute_capabilities(self):
"""Describe backend support shared by emissions and optional length model."""
from mixle.stats.compute.capabilities import DistributionCapabilities, intersect_engine_ready
children = tuple(self.emissions)
if not self.null_len_dist:
children = children + (self.len_dist,)
return DistributionCapabilities(engine_ready=intersect_engine_ready(children), kernel_status="generic_latent")
[docs]
def compute_declaration(self):
"""Return a composite compute declaration for the segmental HMM."""
from mixle.stats.compute.declarations import (
DistributionDeclaration,
ParameterSpec,
StatisticSpec,
declaration_for,
)
emission_children = tuple(declaration_for(emission) for emission in self.emissions)
length = None if self.null_len_dist else declaration_for(self.len_dist)
children = tuple(
child for child in emission_children + ((length,) if length is not None else ()) if child is not None
)
roles = tuple("state_%d_emission" % i for i, child in enumerate(emission_children) if child is not None)
if length is not None:
roles += ("length",)
return DistributionDeclaration(
name="segmental_hidden_markov",
distribution_type=type(self),
parameters=(
ParameterSpec("w", constraint="simplex_vector"),
ParameterSpec("transitions", constraint="row_simplex_matrix"),
),
statistics=(
StatisticSpec("num_states", kind="metadata", additive=False, scales=False),
StatisticSpec("initial_counts"),
StatisticSpec("state_counts"),
StatisticSpec("transition_counts"),
StatisticSpec("emissions", kind="tuple"),
StatisticSpec("length", kind="child_stat"),
),
support="segmental_hidden_state_sequence",
children=children,
child_roles=roles,
differentiable=False,
)
def __init__(
self,
emissions: Sequence[SequenceEncodableProbabilityDistribution],
w: Sequence[float] | np.ndarray = MISSING,
transitions: Sequence[Sequence[float]] | np.ndarray = MISSING,
len_dist: SequenceEncodableProbabilityDistribution | None = NullDistribution(),
name: str | None = None,
weights: Sequence[float] | np.ndarray = MISSING,
terminal_states: set[int] | Sequence[int] | None = None,
) -> None:
w = coalesce_alias("w", w, "weights", weights, default=MISSING)
transitions = require("transitions", transitions, default=MISSING)
context = "SegmentalHiddenMarkovModelDistribution"
self.emissions = _owned_segmental_sequence(emissions, f"{context} emissions")
self.n_states = len(self.emissions)
self.w = validated_probability_vector(w, f"{context} initial probabilities", size=self.n_states)
self.terminal_states = validated_terminal_states(
terminal_states,
self.n_states,
context=context,
)
self.transitions, zero_transition_rows = _validated_segmental_transitions(transitions, self.n_states)
reachable = _reachable_segmental_states(self.w, self.transitions)
reachable_zero_rows = tuple(
index
for index in zero_transition_rows
if reachable[index] and (self.terminal_states is None or index not in self.terminal_states)
)
if reachable_zero_rows:
raise ValueError(
f"{context} reachable non-terminal states cannot have zero transition rows: {reachable_zero_rows}."
)
self.unreachable_transition_rows = zero_transition_rows
for index in zero_transition_rows:
self.transitions[index, index] = 1.0
with np.errstate(divide="ignore"):
self.log_w = np.log(self.w)
self.log_transitions = np.log(self.transitions)
self.len_dist = len_dist if len_dist is not None else NullDistribution()
self.null_len_dist = supports(self.len_dist, Neutral)
self.name = name
if self.terminal_states is not None:
self._terminal_mask = np.zeros(self.n_states, dtype=bool)
self._terminal_mask[list(self.terminal_states)] = True
validate_terminal_reachability(
self.w,
self.transitions,
self.terminal_states,
context=context,
)
def __str__(self) -> str:
s1 = ",".join(str(u) for u in self.emissions)
s2 = repr(list(self.w))
s3 = repr([list(u) for u in self.transitions])
return "SegmentalHiddenMarkovModelDistribution([%s], %s, %s, len_dist=%s, name=%s)" % (
s1,
s2,
s3,
str(self.len_dist),
repr(self.name),
)
@property
def topics(self) -> list[SequenceEncodableProbabilityDistribution]:
"""Compatibility alias with HiddenMarkovModelDistribution terminology."""
return self.emissions
[docs]
def density(self, x: Sequence[Any]) -> float:
"""Return the probability density or mass at a single observation."""
return math.exp(self.log_density(x))
[docs]
def log_density(self, x: Sequence[Any]) -> float:
"""Return the log-density or log-mass at a single observation."""
n = len(x)
if n == 0:
return self.len_dist.log_density(0) if not self.null_len_dist else 0.0
log_emit = np.empty((n, self.n_states), dtype=np.float64)
for k, dist in enumerate(self.emissions):
log_emit[:, k] = np.asarray([dist.log_density(xx) for xx in x], dtype=np.float64)
if self.terminal_states is not None:
from mixle.stats.latent.hidden_markov import terminal_forward_loglik
return terminal_forward_loglik(self.log_w, self.log_transitions, log_emit, self._terminal_mask)
rv = _forward_log(self.log_w, self.log_transitions, log_emit)
if not self.null_len_dist:
rv += self.len_dist.log_density(n)
return rv
[docs]
def seq_log_density(self, x: tuple[np.ndarray, np.ndarray, tuple[Any, ...], Any | None]) -> np.ndarray:
"""Return vectorized log-density values for sequence-encoded observations."""
idx, sz, enc_by_state, len_enc = x
nseq = len(sz)
rv = np.zeros(nseq, dtype=np.float64)
total = len(idx)
log_emit = np.empty((total, self.n_states), dtype=np.float64)
for k, dist in enumerate(self.emissions):
log_emit[:, k] = dist.seq_log_density(enc_by_state[k])
offsets = np.concatenate([[0], np.cumsum(sz)]).astype(int)
if self.terminal_states is not None:
from mixle.stats.latent.hidden_markov import terminal_forward_loglik
for i in range(nseq):
rv[i] = terminal_forward_loglik(
self.log_w, self.log_transitions, log_emit[offsets[i] : offsets[i + 1]], self._terminal_mask
)
return rv
for i in range(nseq):
rv[i] = _forward_log(self.log_w, self.log_transitions, log_emit[offsets[i] : offsets[i + 1]])
if not self.null_len_dist and len_enc is not None:
rv += self.len_dist.seq_log_density(len_enc)
return rv
[docs]
def backend_seq_log_density(
self, x: tuple[np.ndarray, np.ndarray, tuple[Any, ...], Any | None], engine: Any
) -> Any:
"""Engine-neutral segmental-HMM forward scoring for encoded segment sequences."""
from mixle.stats.compute.backend import backend_seq_log_density
idx, sz, enc_by_state, len_enc = x
nseq = len(sz)
total = len(idx)
rv = engine.zeros(nseq)
if total > 0:
emission_scores = [
backend_seq_log_density(dist, enc_by_state[k], engine) for k, dist in enumerate(self.emissions)
]
log_emit = engine.stack(emission_scores, axis=1)
log_w = engine.asarray(self.log_w)
log_a = engine.asarray(self.log_transitions)
offsets = np.concatenate([[0], np.cumsum(sz)]).astype(int)
for i, n in enumerate(sz):
if n == 0:
continue
start = offsets[i]
alpha = log_w + log_emit[start]
for t in range(1, int(n)):
alpha = log_emit[start + t] + engine.logsumexp(alpha[:, None] + log_a, axis=0)
rv[i] = engine.logsumexp(alpha, axis=0)
if not self.null_len_dist and len_enc is not None:
rv = rv + backend_seq_log_density(self.len_dist, len_enc, engine)
return rv
[docs]
def sampler(self, seed: int | None = None) -> "SegmentalHiddenMarkovSampler":
"""Return a sampler for drawing observations from this distribution."""
if self.null_len_dist and self.terminal_states is None:
raise ValueError("SegmentalHiddenMarkovSampler requires a non-null len_dist or terminal_states.")
return SegmentalHiddenMarkovSampler(self, seed)
[docs]
def estimator(self, pseudo_count: float | None = None) -> "SegmentalHiddenMarkovEstimator":
"""Return an estimator for fitting this distribution from data."""
len_est = self.len_dist.estimator(pseudo_count=pseudo_count)
ests = [u.estimator(pseudo_count=pseudo_count) for u in self.emissions]
return SegmentalHiddenMarkovEstimator(
ests,
len_estimator=len_est,
pseudo_count=(pseudo_count, pseudo_count),
name=self.name,
terminal_states=self.terminal_states,
)
[docs]
def dist_to_encoder(self) -> "SegmentalHiddenMarkovDataEncoder":
"""Return the data encoder used by this distribution for vectorized methods."""
return SegmentalHiddenMarkovDataEncoder(
[d.dist_to_encoder() for d in self.emissions], self.len_dist.dist_to_encoder()
)
[docs]
def enumerator(self) -> DistributionEnumerator:
"""Enumerate segment sequences in descending marginal probability order.
The segmental HMM has the *standard* HMM forward semantics -- each position emits one segment
from its state's distribution, scored independently -- so it reuses
:class:`HiddenMarkovModelEnumerator` directly via its per-state emission (``topics``),
``log_w``, ``log_transitions``, and ``len_dist``. Each segment is drawn from the union of the
per-state emission supports, so every emission distribution must itself support enumeration
(and a length distribution must be modeled).
"""
from mixle.stats.latent.hidden_markov import HiddenMarkovModelEnumerator
return HiddenMarkovModelEnumerator(
self,
topics=self.emissions,
log_w=self.log_w,
log_transitions=self.log_transitions,
len_dist=self.len_dist,
path_root="SegmentalHiddenMarkovModelDistribution",
)
[docs]
class SegmentalHiddenMarkovSampler(DistributionSampler):
"""Draw iid segmental-HMM observations."""
def __init__(self, dist: SegmentalHiddenMarkovModelDistribution, seed: int | None = None) -> None:
self.dist = dist
self.rng = RandomState(seed)
self.obs_samplers = [d.sampler(seed=self.rng.randint(0, maxrandint)) for d in dist.emissions]
self.len_sampler = None if dist.null_len_dist else dist.len_dist.sampler(seed=self.rng.randint(0, maxrandint))
p_map = {i: dist.w[i] for i in range(dist.n_states)}
t_map = {i: {j: dist.transitions[i, j] for j in range(dist.n_states)} for i in range(dist.n_states)}
self.state_sampler = MarkovChainDistribution(p_map, t_map).path_sampler(seed=self.rng.randint(0, maxrandint))
[docs]
def sample(self, size: int | None = None, *, batched: bool = True) -> list[Any] | list[list[Any]]:
"""Draw one segment sequence, or ``size`` iid segment sequences."""
if size is not None:
return [self.sample() for _ in range(size)]
if self.dist.terminal_states is not None:
max_steps = validated_terminal_step_cap(DEFAULT_TERMINAL_STEP_CAP)
z = int(self.state_sampler.sample_seq())
states = [z]
while z not in self.dist.terminal_states and len(states) < max_steps:
z = int(self.state_sampler.sample_seq(v0=z))
states.append(z)
require_terminal_reached(
z in self.dist.terminal_states,
mode="segmental terminal-state",
max_steps=max_steps,
last_state=z,
)
return [self.obs_samplers[s].sample() for s in states]
n = _validated_segment_count(self.len_sampler.sample())
states = self.state_sampler.sample_seq(n)
return [self.obs_samplers[s].sample() for s in states]
[docs]
class SegmentalHiddenMarkovAccumulator(SequenceEncodableStatisticAccumulator):
"""Baum-Welch accumulator for segmental HMMs."""
def __init__(
self,
accumulators: Sequence[SequenceEncodableStatisticAccumulator],
len_accumulator: SequenceEncodableStatisticAccumulator | None = NullAccumulator(),
keys: tuple[str | None, str | None, str | None] | None = (None, None, None),
name: str | None = None,
) -> None:
self.accumulators = _owned_segmental_sequence(
accumulators,
"SegmentalHiddenMarkovAccumulator state accumulators",
)
self.num_states = len(self.accumulators)
self.init_counts = np.zeros(self.num_states, dtype=np.float64)
self.state_counts = np.zeros(self.num_states, dtype=np.float64)
self.trans_counts = np.zeros((self.num_states, self.num_states), dtype=np.float64)
self.len_accumulator = len_accumulator if len_accumulator is not None else NullAccumulator()
keys = (None, None, None) if keys is None else tuple(keys)
if len(keys) != 3:
raise ValueError("SegmentalHiddenMarkovAccumulator keys must contain exactly three entries.")
self.init_key, self.trans_key, self.state_key = keys
self.name = name
# When _track_ll is enabled, seq_update accumulates the per-sequence data
# log-likelihood into _seq_ll. Used by the fused-EM fast path in
# optimize(reuse_estep_ll=True); default path is unchanged and zero-cost.
self._track_ll = False
self._seq_ll = 0.0
self._init_rng = False
self._state_rng: RandomState | None = None
self._len_rng: RandomState | None = None
self._acc_rng: list[RandomState] | None = None
def _rng_initialize(self, rng: RandomState) -> None:
seeds = rng.randint(maxrandint, size=2 + self.num_states)
self._state_rng = RandomState(seeds[0])
self._len_rng = RandomState(seeds[1])
self._acc_rng = [RandomState(seeds[i + 2]) for i in range(self.num_states)]
self._init_rng = True
[docs]
def update(self, x: Sequence[Any], weight: float, estimate: SegmentalHiddenMarkovModelDistribution) -> None:
"""Update sufficient statistics for one observed segment sequence."""
weight = validated_observation_weight(weight, "segmental-HMM observation weight")
self.seq_update(estimate.dist_to_encoder().seq_encode([x]), np.asarray([weight]), estimate)
[docs]
def initialize(self, x: Sequence[Any], weight: float, rng: RandomState) -> None:
"""Randomly initialize sufficient statistics for one segment sequence."""
weight = validated_observation_weight(weight, "segmental-HMM initialization weight")
if not self._init_rng:
self._rng_initialize(rng)
enc = self.acc_to_encoder().seq_encode([x])
self.seq_initialize(enc, np.asarray([weight]), rng)
[docs]
def seq_initialize(
self, x: tuple[np.ndarray, np.ndarray, tuple[Any, ...], Any | None], weights: np.ndarray, rng: RandomState
) -> None:
"""Randomly initialize state and emission statistics for encoded sequences."""
weights = validated_observation_weights(weights, len(x[1]), "segmental-HMM initialization weights")
if not self._init_rng:
self._rng_initialize(rng)
idx, sz, enc_by_state, len_enc = x
total = len(idx)
# A uniform draw gives every state the same random 1/K subsample, so every emission starts at
# the marginal and EM opens on the symmetric fixed point (A-01); the gallery's segmental fit
# was 3.4 nats/observation short of the converged one (P09-F11). A k-means++ start over the
# observations breaks that; a non-numeric emission has nothing to cluster and keeps the draw.
states = self._state_rng.choice(self.num_states, size=total)
if self.num_states > 1 and len(enc_by_state):
broken = broken_symmetry_states(enc_by_state[0], self.num_states, derived_rng(states), rows=total)
if broken is not None:
states = broken
offsets = np.concatenate([[0], np.cumsum(sz)]).astype(int)
weighted_state = np.zeros((total, self.num_states), dtype=np.float64)
for i, n in enumerate(sz):
if n == 0:
continue
start, stop = offsets[i], offsets[i + 1]
weighted_state[start:stop, :] = 0.0
weighted_state[np.arange(start, stop), states[start:stop]] = weights[i]
self.init_counts[states[start]] += weights[i]
self.state_counts += np.bincount(
states[start:stop], weights=np.full(n, weights[i]), minlength=self.num_states
)
for t in range(start, stop - 1):
self.trans_counts[states[t], states[t + 1]] += weights[i]
for k in range(self.num_states):
self.accumulators[k].seq_initialize(enc_by_state[k], weighted_state[:, k], self._acc_rng[k])
self.len_accumulator.seq_initialize(len_enc, weights, self._len_rng)
[docs]
def seq_update(
self,
x: tuple[np.ndarray, np.ndarray, tuple[Any, ...], Any | None],
weights: np.ndarray,
estimate: SegmentalHiddenMarkovModelDistribution,
) -> None:
"""Update encoded-sequence statistics with Baum-Welch posteriors."""
weights = validated_observation_weights(weights, len(x[1]), "segmental-HMM observation weights")
if estimate.n_states != self.num_states:
raise ValueError("segmental-HMM estimate and accumulator state counts must match")
require_possible_log_evidence(
estimate.seq_log_density(x),
context="SegmentalHiddenMarkovAccumulator.seq_update",
)
idx, sz, enc_by_state, len_enc = x
total = len(idx)
log_emit = np.empty((total, self.num_states), dtype=np.float64)
for k, dist in enumerate(estimate.emissions):
log_emit[:, k] = dist.seq_log_density(enc_by_state[k])
# When the fused-EM fast path requests it, accumulate the per-sequence data
# log-likelihood from the existing forward pass (ll returned by _forward_backward,
# which equals seq_log_density's _forward_log term exactly). The standard path skips it.
track_ll = self._track_ll
ll_ret = np.zeros(len(sz), dtype=np.float64) if track_ll else None
terminal = estimate.terminal_states is not None
if terminal:
from mixle.stats.latent.hidden_markov import terminal_forward_backward
offsets = np.concatenate([[0], np.cumsum(sz)]).astype(int)
gamma_all = np.zeros((total, self.num_states), dtype=np.float64)
for i, n in enumerate(sz):
if n == 0:
continue
start, stop = offsets[i], offsets[i + 1]
if terminal:
ll_i, gamma, xi = terminal_forward_backward(
estimate.log_w, estimate.log_transitions, log_emit[start:stop], estimate._terminal_mask
)
if gamma is None:
continue
xi_sum = xi.sum(axis=0)
else:
ll_i, gamma, xi_sum = _forward_backward(estimate.log_w, estimate.log_transitions, log_emit[start:stop])
if track_ll:
ll_ret[i] = ll_i
w = weights[i]
gamma_w = gamma * w
gamma_all[start:stop, :] = gamma_w
self.init_counts += gamma_w[0]
self.state_counts += gamma_w.sum(axis=0)
self.trans_counts += xi_sum * w
if track_ll:
if not estimate.null_len_dist and len_enc is not None:
ll_ret = ll_ret + estimate.len_dist.seq_log_density(len_enc)
self._seq_ll += float(np.dot(weights, ll_ret))
for k in range(self.num_states):
self.accumulators[k].seq_update(enc_by_state[k], gamma_all[:, k], estimate.emissions[k])
self.len_accumulator.seq_update(len_enc, weights, estimate.len_dist)
[docs]
def seq_update_engine(self, x, weights, estimate, engine):
"""Engine-resident Baum-Welch E-step (numpy or torch).
Reuses hmm_engine_forward_backward over the sequence-contiguous segment encoding, producing
the same initial/state/transition/emission statistics as the host seq_update.
"""
from mixle.stats.latent.hidden_markov import hmm_engine_forward_backward, hmm_pad_log_emissions
weights_np = validated_observation_weights(
engine.to_numpy(weights) if hasattr(engine, "to_numpy") else weights,
len(x[1]),
"segmental-HMM observation weights",
)
if estimate.n_states != self.num_states:
raise ValueError("segmental-HMM estimate and accumulator state counts must match")
require_possible_log_evidence(
estimate.seq_log_density(x),
context="SegmentalHiddenMarkovAccumulator.seq_update_engine",
)
idx, sz, enc_by_state, len_enc = x
sz = np.asarray(sz)
total = len(idx)
num_states = self.num_states
log_emit_flat = np.empty((total, num_states), dtype=np.float64)
for k, dist in enumerate(estimate.emissions):
log_emit_flat[:, k] = dist.seq_log_density(enc_by_state[k])
padded, mask, offsets = hmm_pad_log_emissions(log_emit_flat, sz)
with np.errstate(divide="ignore"):
log_w = estimate.log_w
log_a = estimate.log_transitions
_, gamma, xi_sum, pi = hmm_engine_forward_backward(engine, padded, log_w, log_a, mask, weights=weights_np)
gamma = np.asarray(engine.to_numpy(gamma))
xi_sum = np.asarray(engine.to_numpy(xi_sum))
pi = np.asarray(engine.to_numpy(pi))
gamma_flat = np.zeros((total, num_states), dtype=np.float64)
for i in range(len(sz)):
n = int(sz[i])
if n > 0:
gamma_flat[offsets[i] : offsets[i + 1], :] = gamma[i, :n, :]
self.init_counts += pi.sum(axis=0)
self.state_counts += gamma_flat.sum(axis=0)
self.trans_counts += xi_sum
for k in range(num_states):
self.accumulators[k].seq_update(enc_by_state[k], gamma_flat[:, k], estimate.emissions[k])
self.len_accumulator.seq_update(len_enc, weights_np, estimate.len_dist)
[docs]
def combine(
self, suff_stat: tuple[int, np.ndarray, np.ndarray, np.ndarray, Sequence[Any], Any | None]
) -> "SegmentalHiddenMarkovAccumulator":
"""Merge another segmental-HMM sufficient-statistic value."""
num_states, init_counts, state_counts, trans_counts, acc_values, len_value = validated_statistic_tuple(
suff_stat,
6,
"segmental-HMM sufficient statistics",
)
num_states = validated_positive_integer(num_states, "segmental-HMM statistic state count")
if num_states != self.num_states:
raise ValueError("segmental-HMM statistic state count must match the accumulator")
init_counts = validated_count_array(init_counts, (num_states,), "segmental-HMM initial counts")
state_counts = validated_count_array(state_counts, (num_states,), "segmental-HMM state counts")
trans_counts = validated_count_array(
trans_counts,
(num_states, num_states),
"segmental-HMM transition counts",
)
# combine receives one accumulator's own statistics -- pooling happens later, at
# key_merge -- so the unconditional equality applies, with the family's sqrt-mass
# tolerance instead of the previous flat 1e-9 (which was tighter than float32 roundoff).
_validate_state_mass(
init_counts,
state_counts,
trans_counts,
init_key=None,
trans_key=None,
label="segmental-HMM",
)
if not isinstance(acc_values, (tuple, list)) or len(acc_values) != num_states:
raise ValueError("segmental-HMM emission statistics must match the state count")
# The ENTIRE combine is transactional in the family-shared core (STAT-RR8-1/RR9-1
# classes, measured on this accumulator by the latent-family mutator audit).
_transactional_combine(
self,
init_counts,
state_counts,
trans_counts,
acc_values,
len_value,
family_label="segmental-HMM",
)
return self
[docs]
def value(self) -> tuple[int, np.ndarray, np.ndarray, np.ndarray, tuple[Any, ...], Any | None]:
"""Return transition, emission, and sequence-length sufficient statistics."""
return (
self.num_states,
self.init_counts,
self.state_counts,
self.trans_counts,
tuple(a.value() for a in self.accumulators),
self.len_accumulator.value(),
)
[docs]
def from_value(
self, x: tuple[int, np.ndarray, np.ndarray, np.ndarray, Sequence[Any], Any | None]
) -> "SegmentalHiddenMarkovAccumulator":
"""Replace this accumulator from serialized sufficient statistics."""
num_states, init_counts, state_counts, trans_counts, acc_values, len_value = validated_statistic_tuple(
x,
6,
"segmental-HMM sufficient statistics",
)
num_states = validated_positive_integer(num_states, "segmental-HMM statistic state count")
if num_states != self.num_states:
raise ValueError("segmental-HMM statistic state count must match the accumulator")
# Validate EVERYTHING before mutating ANYTHING (the previous order assigned each count
# array as it validated), enforce the mode-appropriate mass relation -- a KEYED
# accumulator's own value() carries pooled initial/transition parts, and the previous
# unconditional equality rejected its own round-trip (measured) -- then assign and
# restore children as one transaction in the family-shared core.
candidate_init = validated_count_array(init_counts, (num_states,), "segmental-HMM initial counts")
candidate_state = validated_count_array(state_counts, (num_states,), "segmental-HMM state counts")
candidate_trans = validated_count_array(
trans_counts,
(num_states, num_states),
"segmental-HMM transition counts",
)
_validate_state_mass(
candidate_init,
candidate_state,
candidate_trans,
init_key=self.init_key,
trans_key=self.trans_key,
label="segmental-HMM",
)
if not isinstance(acc_values, (tuple, list)) or len(acc_values) != num_states:
raise ValueError("segmental-HMM emission statistics must match the state count")
_transactional_restore(self, candidate_init, candidate_state, candidate_trans, acc_values, len_value)
return self
[docs]
def scale(self, c: float) -> "SegmentalHiddenMarkovAccumulator":
"""Scale all weight-linear sufficient statistics by ``c``."""
c = validated_observation_weight(c, "segmental-HMM statistic scale")
# Parent counts, children, and the length child scale as ONE transaction with the
# scaled result validated as a postcondition, in the family-shared core (measured;
# STAT-RR8-1/RR10-1 classes).
_transactional_scale(self, c, family_label="segmental-HMM")
return self
[docs]
def key_merge(self, stats_dict: dict[str, Any]) -> None:
"""Merge keyed initial, transition, emission, and length statistics."""
# One shared implementation for the whole HMM family: transactional with in-place-
# healing rollback, pooled-result finiteness checks, and copy-on-adoption aliasing
# discipline -- the local copy left earlier pools merged when a later pool failed
# (measured; STAT-RR8-1/RR9-1/RR10-1 classes).
_keyed_statistics_merge(self, stats_dict, family_label="segmental-HMM")
[docs]
def key_replace(self, stats_dict: dict[str, Any]) -> None:
"""Replace keyed initial, transition, emission, and length statistics."""
# One shared implementation for the whole HMM family: both candidates validated before
# either is assigned, copy-on-replace aliasing discipline, and full accumulator-plus-
# mapping rollback -- the local copy accepted [inf, 0] outright and adopted a poisoned
# pooled child list before anything validated it (measured; STAT-RR8-1/RR9-1/RR10-1).
_keyed_statistics_replace(self, stats_dict, family_label="segmental-HMM")
[docs]
def acc_to_encoder(self) -> "SegmentalHiddenMarkovDataEncoder":
"""Return an encoder compatible with the emission and length accumulators."""
return SegmentalHiddenMarkovDataEncoder(
[a.acc_to_encoder() for a in self.accumulators], self.len_accumulator.acc_to_encoder()
)
[docs]
class SegmentalHiddenMarkovAccumulatorFactory(StatisticAccumulatorFactory):
"""Factory for SegmentalHiddenMarkovAccumulator."""
def __init__(
self,
factories: Sequence[StatisticAccumulatorFactory],
len_factory: StatisticAccumulatorFactory | None = NullAccumulatorFactory(),
keys: tuple[str | None, str | None, str | None] | None = (None, None, None),
name: str | None = None,
) -> None:
self.factories = _owned_segmental_sequence(
factories,
"SegmentalHiddenMarkovAccumulatorFactory state factories",
)
self.len_factory = len_factory if len_factory is not None else NullAccumulatorFactory()
self.keys = (None, None, None) if keys is None else tuple(keys)
if len(self.keys) != 3:
raise ValueError("SegmentalHiddenMarkovAccumulatorFactory keys must contain exactly three entries.")
self.name = name
[docs]
def make(self) -> SegmentalHiddenMarkovAccumulator:
"""Create a fresh segmental-HMM accumulator."""
return SegmentalHiddenMarkovAccumulator(
[f.make() for f in self.factories], self.len_factory.make(), keys=self.keys, name=self.name
)
[docs]
class SegmentalHiddenMarkovEstimator(ParameterEstimator):
"""Baum-Welch estimator for SegmentalHiddenMarkovModelDistribution."""
def __init__(
self,
estimators: Sequence[ParameterEstimator],
len_estimator: ParameterEstimator | None = NullEstimator(),
pseudo_count: float | tuple[float | None, float | None] | None = (None, None),
name: str | None = None,
keys: tuple[str | None, str | None, str | None] | None = (None, None, None),
terminal_states: set[int] | Sequence[int] | None = None,
) -> None:
self.estimators = _owned_segmental_sequence(
estimators,
"SegmentalHiddenMarkovEstimator estimators",
)
self.num_states = len(self.estimators)
self.len_estimator = len_estimator if len_estimator is not None else NullEstimator()
self.pseudo_count = _validated_segmental_pseudo_count(pseudo_count)
self.keys = (None, None, None) if keys is None else tuple(keys)
if len(self.keys) != 3:
raise ValueError("SegmentalHiddenMarkovEstimator keys must contain exactly three entries.")
self.name = name
self.terminal_states = validated_terminal_states(
terminal_states,
self.num_states,
context="SegmentalHiddenMarkovEstimator",
)
[docs]
def accumulator_factory(self) -> SegmentalHiddenMarkovAccumulatorFactory:
"""Return an accumulator factory for Baum-Welch sufficient statistics."""
return SegmentalHiddenMarkovAccumulatorFactory(
[e.accumulator_factory() for e in self.estimators],
self.len_estimator.accumulator_factory(),
keys=self.keys,
name=self.name,
)
[docs]
def estimate(
self,
nobs: float | None,
suff_stat: tuple[int, np.ndarray, np.ndarray, np.ndarray, Sequence[Any], Any | None],
) -> SegmentalHiddenMarkovModelDistribution:
"""Estimate initial, transition, emission, and length distributions."""
values = validated_statistic_tuple(suff_stat, 6, "segmental-HMM sufficient statistics")
num_states, init_counts, state_counts, trans_counts, emission_ss, len_ss = values
if isinstance(num_states, bool) or not isinstance(num_states, (int, np.integer)):
raise TypeError("segmental HMM sufficient-statistic state count must be an integer.")
if int(num_states) != self.num_states:
raise ValueError(
f"segmental HMM sufficient-statistic state count {num_states} does not match {self.num_states}."
)
num_states = self.num_states
init_counts = _validated_segmental_count_array(
init_counts,
"segmental HMM initial counts",
(num_states,),
)
state_counts = _validated_segmental_count_array(
state_counts,
"segmental HMM state counts",
(num_states,),
)
trans_counts = _validated_segmental_count_array(
trans_counts,
"segmental HMM transition counts",
(num_states, num_states),
)
emission_ss = _owned_segmental_sequence(
emission_ss,
"segmental HMM emission sufficient statistics",
size=num_states,
)
_est_keys = self.keys if isinstance(self.keys, (tuple, list)) and len(self.keys) == 3 else (None, None, None)
# the M-step receives post-key_merge statistics, so this is where pooling actually
# lands: equality when initial/transition are unkeyed, the pooled upper bound otherwise
# (the previous unconditional equality rejected every keyed fit's M-step -- measured)
_validate_state_mass(
init_counts,
state_counts,
trans_counts,
init_key=_est_keys[0],
trans_key=_est_keys[1],
label="segmental-HMM",
)
if _est_keys[0] is None:
# a POOLED initial-count vector carries the mass of every site sharing the key, so
# comparing it against THIS site's observation count stops being a corruption
# check; an unkeyed one is still bounded by the observations that produced it
validate_effective_sample_mass(
nobs,
float(init_counts.sum()),
label="segmental-HMM effective sample",
allow_unassigned=True,
)
emissions = [self.estimators[k].estimate(state_counts[k], emission_ss[k]) for k in range(num_states)]
len_dist = self.len_estimator.estimate(nobs, len_ss)
if self.pseudo_count[0] is not None:
w = init_counts + self.pseudo_count[0] / float(num_states)
else:
w = init_counts.copy()
w = np.ones(num_states) / float(num_states) if w.sum() <= 0.0 else w / w.sum()
if self.pseudo_count[1] is not None:
transitions = trans_counts + self.pseudo_count[1] / float(num_states * num_states)
else:
transitions = trans_counts.copy()
row_sum = transitions.sum(axis=1, keepdims=True)
empty_rows = row_sum[:, 0] == 0.0
transitions = np.divide(
transitions,
row_sum,
out=np.zeros_like(transitions),
where=row_sum > 0.0,
)
empty_indices = np.flatnonzero(empty_rows)
transitions[empty_indices, empty_indices] = 1.0
return SegmentalHiddenMarkovModelDistribution(
emissions, w, transitions, len_dist=len_dist, name=self.name, terminal_states=self.terminal_states
)
[docs]
class SegmentalHiddenMarkovDataEncoder(DataSequenceEncoder):
"""Encode a batch of segment sequences for a segmental HMM."""
def __init__(
self,
emission_encoders: Sequence[DataSequenceEncoder],
len_encoder: DataSequenceEncoder | None = NullDataEncoder(),
) -> None:
self.emission_encoders = _owned_segmental_sequence(
emission_encoders,
"SegmentalHiddenMarkovDataEncoder emission encoders",
)
self.len_encoder = len_encoder if len_encoder is not None else NullDataEncoder()
def __str__(self) -> str:
return "SegmentalHiddenMarkovDataEncoder([%s], len_encoder=%s)" % (
",".join(str(e) for e in self.emission_encoders),
str(self.len_encoder),
)
def __eq__(self, other: object) -> bool:
return (
isinstance(other, SegmentalHiddenMarkovDataEncoder)
and self.emission_encoders == other.emission_encoders
and self.len_encoder == other.len_encoder
)
[docs]
def seq_encode(self, x: Sequence[Sequence[Any]]) -> tuple[np.ndarray, np.ndarray, tuple[Any, ...], Any | None]:
"""Flatten segment sequences and encode them for every state emission."""
lengths = np.asarray([len(seq) for seq in x], dtype=np.int32)
idx = np.repeat(np.arange(len(x), dtype=np.int32), lengths)
flat: list[Any] = []
for seq in x:
flat.extend(seq)
enc_by_state = tuple(enc.seq_encode(flat) for enc in self.emission_encoders)
len_enc = None if supports(self.len_encoder, Neutral) else self.len_encoder.seq_encode(lengths)
return idx, lengths, enc_by_state, len_enc
[docs]
def row_count(self, x: Any) -> int:
"""Return and validate the number of encoded segment sequences."""
if not isinstance(x, tuple) or len(x) != 4:
raise ValueError("segmental HMM encoding must be a four-slot payload.")
idx, lengths, enc_by_state, len_enc = x
lengths = np.asarray(lengths)
if lengths.ndim != 1 or np.any(lengths < 0):
raise ValueError("segmental HMM encoded lengths must be one-dimensional and non-negative.")
total = int(lengths.sum())
if np.asarray(idx).shape != (total,):
raise ValueError("segmental HMM node-to-sequence indices do not match encoded lengths.")
if len(enc_by_state) != len(self.emission_encoders):
raise ValueError("segmental HMM encoded emission-state arity does not match the encoder.")
for encoder, encoded in zip(self.emission_encoders, enc_by_state):
if encoder.row_count(encoded) != total:
raise ValueError("segmental HMM encoded emission rows do not match the flattened segment count.")
count = int(lengths.size)
if len_enc is not None and self.len_encoder.row_count(len_enc) != count:
raise ValueError("segmental HMM encoded length rows do not match the sequence count.")
return count
SegmentalHiddenMarkovDistribution = SegmentalHiddenMarkovModelDistribution
# --- Backward-compatible API naming aliases ---
SegmentalHiddenMarkovModelAccumulator = SegmentalHiddenMarkovAccumulator
SegmentalHiddenMarkovModelAccumulatorFactory = SegmentalHiddenMarkovAccumulatorFactory
SegmentalHiddenMarkovModelDataEncoder = SegmentalHiddenMarkovDataEncoder
SegmentalHiddenMarkovModelEstimator = SegmentalHiddenMarkovEstimator
SegmentalHiddenMarkovModelSampler = SegmentalHiddenMarkovSampler
def _register_segmental_engine_kernel():
"""Register the engine-resident segmental-HMM kernel (idempotent; called at import)."""
from mixle.stats.compute.kernel import GenericKernel, GenericKernelFactory, KernelFactory, register_kernel_factory
class SegmentalHiddenMarkovModelKernel(GenericKernel):
"""Segmental-HMM kernel whose E-step runs the forward-backward on the active engine."""
def accumulate(self, enc, weights):
if self.estimator is None:
raise ValueError("SegmentalHiddenMarkovModelKernel.accumulate requires an estimator.")
if not getattr(self.engine, "resident_estep", True):
return super().accumulate(enc, weights)
host_enc = getattr(enc, "host_payload", enc)
accumulator = self.estimator.accumulator_factory().make()
accumulator.seq_update_engine(host_enc, weights, self.dist, self.engine)
return accumulator.value()
class SegmentalHiddenMarkovModelKernelFactory(KernelFactory):
def build(self, dist, engine, estimator=None):
if not dist.supports_engine(engine):
return GenericKernelFactory().build(dist, engine, estimator=estimator)
return SegmentalHiddenMarkovModelKernel(dist, engine=engine, estimator=estimator)
register_kernel_factory(SegmentalHiddenMarkovModelDistribution, SegmentalHiddenMarkovModelKernelFactory())
_register_segmental_engine_kernel()