581 lines
19 KiB
Go
581 lines
19 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package stats
|
||
|
||
import (
|
||
"math"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// Hidden Markov models over discrete observations: the filtered and
|
||
// smoothed state posteriors, the most likely state path and the
|
||
// Baum-Welch fit of the parameters. Every computation runs on the
|
||
// scaled recursions, where each sweep's normaliser absorbs the
|
||
// likelihood of what has been seen, so a long sequence cannot
|
||
// underflow the way the raw forward probabilities would; the log
|
||
// likelihood itself is the sum of the normalisers' logarithms. The
|
||
// decoding walks in log space with the lowest index winning ties, and
|
||
// the fit draws its starting point from the house generator, so the
|
||
// whole surface is deterministic for a given generator state and
|
||
// sequence.
|
||
|
||
// HiddenMarkovModel is a hidden Markov model over a set of hidden
|
||
// states and a set of observation symbols. The fields hold the
|
||
// distribution parameters in row-major order: the rows of Transition
|
||
// are the conditionals P(next = j | current = k) and the rows of
|
||
// Emission the conditionals P(symbol = m | state = k). The
|
||
// constructors validate every row as a distribution, so a model value
|
||
// is always safe to evaluate.
|
||
type HiddenMarkovModel struct {
|
||
// Initial is the distribution over states at the first step.
|
||
Initial []float64
|
||
// Transition is the states×states matrix of one-step
|
||
// probabilities, row-major.
|
||
Transition []float64
|
||
// Emission is the states×symbols matrix of observation
|
||
// probabilities, row-major.
|
||
Emission []float64
|
||
}
|
||
|
||
// HiddenMarkovFitResult carries the Baum-Welch fit of a hidden Markov
|
||
// model.
|
||
type HiddenMarkovFitResult struct {
|
||
// Model is the fitted model.
|
||
Model *HiddenMarkovModel
|
||
// LogLikelihood is the log likelihood of the training sequence
|
||
// under the fitted model.
|
||
LogLikelihood float64
|
||
// Iterations counts the Baum-Welch sweeps taken; Converged reports
|
||
// whether the log likelihood settled under the tolerance before
|
||
// the budget ran out.
|
||
Iterations int
|
||
Converged bool
|
||
}
|
||
|
||
// hmmMaxIterations caps the Baum-Welch sweeps, with the same
|
||
// reporting convention the other fits carry.
|
||
const hmmMaxIterations = 500
|
||
|
||
// hmmTolerance is the convergence tolerance on the log likelihood:
|
||
// the sweeps stop once it moves by less than hmmTolerance scaled by
|
||
// 1 + |log likelihood|. Baum-Welch's ascent crawls along a flat ridge
|
||
// long before its parameters settle to the last digit, and the
|
||
// tolerance stops the sweep where the objective's remaining climb is
|
||
// immaterial rather than where the parameters stop moving.
|
||
const hmmTolerance = 1e-6
|
||
|
||
// hmmFloor is the probability floor the re-estimation applies before
|
||
// renormalising: a parameter driven to exactly zero would silence a
|
||
// state or a symbol forever, and the floor keeps every path alive at
|
||
// a width no plausible parameter sits under.
|
||
const hmmFloor = 1e-12
|
||
|
||
// NewHiddenMarkovModel validates and copies the parameters of a
|
||
// hidden Markov model: Initial of the length of the state set,
|
||
// Transition with one row per state over the states, Emission with
|
||
// one row per state over the symbols. Every entry must be finite and
|
||
// in [0, 1] and every row must sum to 1 within 1e-9; a model that
|
||
// failed any of these would silently weight impossible paths later.
|
||
func NewHiddenMarkovModel(initial, transition, emission []float64) (*HiddenMarkovModel, error) {
|
||
const name = "NewHiddenMarkovModel"
|
||
states := len(initial)
|
||
if states < 1 {
|
||
return nil, base.Errf("%s: the model needs at least one state", name)
|
||
}
|
||
if len(transition) != states*states {
|
||
return nil, base.Errf("%s: the transition matrix holds %d entries, want %d for %d states",
|
||
name, len(transition), states*states, states)
|
||
}
|
||
if len(emission) == 0 || len(emission)%states != 0 {
|
||
return nil, base.Errf("%s: the emission matrix holds %d entries, not a whole number of rows for %d states",
|
||
name, len(emission), states)
|
||
}
|
||
if err := hmmCheckRow(name, "the initial distribution", initial); err != nil {
|
||
return nil, err
|
||
}
|
||
for k := range states {
|
||
if err := hmmCheckRow(name, "the transition rows", transition[k*states:(k+1)*states]); err != nil {
|
||
return nil, err
|
||
}
|
||
}
|
||
symbols := len(emission) / states
|
||
for k := range states {
|
||
if err := hmmCheckRow(name, "the emission rows", emission[k*symbols:(k+1)*symbols]); err != nil {
|
||
return nil, err
|
||
}
|
||
}
|
||
return &HiddenMarkovModel{
|
||
Initial: append([]float64(nil), initial...),
|
||
Transition: append([]float64(nil), transition...),
|
||
Emission: append([]float64(nil), emission...),
|
||
}, nil
|
||
}
|
||
|
||
// hmmCheckRow validates one distribution row: finite, inside [0, 1],
|
||
// summing to 1 within 1e-9.
|
||
func hmmCheckRow(name, what string, row []float64) error {
|
||
total := 0.0
|
||
for _, v := range row {
|
||
if math.IsNaN(v) || math.IsInf(v, 0) {
|
||
return base.Errf("%s: %s hold the non-finite value %g", name, what, v)
|
||
}
|
||
if v < 0 || v > 1 {
|
||
return base.Errf("%s: %s hold the probability %g outside [0, 1]", name, what, v)
|
||
}
|
||
total += v
|
||
}
|
||
if math.Abs(total-1) > 1e-9 {
|
||
return base.Errf("%s: %s sum to %g, not 1", name, what, total)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// hmmStates returns the model's state count.
|
||
func (m *HiddenMarkovModel) hmmStates() int { return len(m.Initial) }
|
||
|
||
// hmmSymbols returns the model's symbol count.
|
||
func (m *HiddenMarkovModel) hmmSymbols() int { return len(m.Emission) / max(1, len(m.Initial)) }
|
||
|
||
// hmmWidth returns the model's symbol count from its own fields, the
|
||
// shape the recursions read the emission rows with.
|
||
func (m *HiddenMarkovModel) hmmWidth() int { return len(m.Emission) / len(m.Initial) }
|
||
|
||
// hmmCheckSequence validates an observation sequence against the
|
||
// model's symbol set.
|
||
func hmmCheckSequence(name string, observations []int, symbols int) error {
|
||
if len(observations) == 0 {
|
||
return base.Errf("%s: the observation sequence is empty", name)
|
||
}
|
||
for t, o := range observations {
|
||
if o < 0 || o >= symbols {
|
||
return base.Errf("%s: observation %d is symbol %d, outside the model's %d symbols", name, t, o, symbols)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// hmmForward runs the scaled forward recursion: alpha[t] is the
|
||
// P(state, observations 0..t) vector normalised by its own total, and
|
||
// the totals accumulate into the log likelihood. The returned rows
|
||
// are fresh. A normaliser of exactly zero means the sequence has
|
||
// probability zero under the model, the structural zeros the
|
||
// parameters carry having silenced every state at some step: the
|
||
// posteriors are undefined there and the recursion refuses rather
|
||
// than divide the zero into NaNs.
|
||
func (m *HiddenMarkovModel) hmmForward(observations []int) ([][]float64, []float64, error) {
|
||
const name = "HiddenMarkovModel.Forward"
|
||
states := m.hmmStates()
|
||
width := m.hmmWidth()
|
||
alpha := make([][]float64, len(observations))
|
||
scales := make([]float64, len(observations))
|
||
row := make([]float64, states)
|
||
for k := range states {
|
||
row[k] = m.Initial[k] * m.Emission[k*width+observations[0]]
|
||
}
|
||
total := 0.0
|
||
for _, v := range row {
|
||
total += v
|
||
}
|
||
if total == 0 {
|
||
return nil, nil, base.Errf("%s: the observation sequence has probability zero under the model", name)
|
||
}
|
||
scales[0] = total
|
||
for k := range states {
|
||
row[k] /= total
|
||
}
|
||
alpha[0] = row
|
||
for t := 1; t < len(observations); t++ {
|
||
next := make([]float64, states)
|
||
sym := observations[t]
|
||
for j := range states {
|
||
total := 0.0
|
||
for k := range states {
|
||
total += row[k] * m.Transition[k*states+j]
|
||
}
|
||
next[j] = total * m.Emission[j*width+sym]
|
||
}
|
||
total := 0.0
|
||
for _, v := range next {
|
||
total += v
|
||
}
|
||
if total == 0 {
|
||
return nil, nil, base.Errf("%s: the observation sequence has probability zero under the model at step %d", name, t)
|
||
}
|
||
scales[t] = total
|
||
for j := range states {
|
||
next[j] /= total
|
||
}
|
||
alpha[t] = next
|
||
row = next
|
||
}
|
||
return alpha, scales, nil
|
||
}
|
||
|
||
// hmmBackward runs the scaled backward recursion against the forward
|
||
// scales, so the product alpha·beta normalises directly into the
|
||
// smoothed posteriors.
|
||
func (m *HiddenMarkovModel) hmmBackward(observations []int, scales []float64) [][]float64 {
|
||
states := m.hmmStates()
|
||
width := m.hmmWidth()
|
||
beta := make([][]float64, len(observations))
|
||
last := make([]float64, states)
|
||
for k := range states {
|
||
last[k] = 1
|
||
}
|
||
beta[len(observations)-1] = last
|
||
for t := len(observations) - 2; t >= 0; t-- {
|
||
row := make([]float64, states)
|
||
sym := observations[t+1]
|
||
for k := range states {
|
||
total := 0.0
|
||
for j := range states {
|
||
total += m.Transition[k*states+j] * m.Emission[j*width+sym] * beta[t+1][j]
|
||
}
|
||
row[k] = total / scales[t+1]
|
||
}
|
||
beta[t] = row
|
||
}
|
||
return beta
|
||
}
|
||
|
||
// Forward runs the scaled forward recursion over the observation
|
||
// sequence and returns the filtered state posteriors, one row of
|
||
// state probabilities per step, with the sequence's log likelihood
|
||
// P(observations | model). Refuses a nil model, an empty sequence,
|
||
// any observation outside the model's symbol set, and a sequence of
|
||
// probability zero under the model, whose structural zero an emission
|
||
// or transition carries leaves the filtered posteriors undefined.
|
||
func (m *HiddenMarkovModel) Forward(observations []int) (filtered [][]float64, logLikelihood float64, err error) {
|
||
const name = "HiddenMarkovModel.Forward"
|
||
if m == nil {
|
||
return nil, 0, base.Errf("%s: the model is nil", name)
|
||
}
|
||
if err := hmmCheckSequence(name, observations, m.hmmSymbols()); err != nil {
|
||
return nil, 0, err
|
||
}
|
||
alpha, scales, ferr := m.hmmForward(observations)
|
||
if ferr != nil {
|
||
return nil, 0, ferr
|
||
}
|
||
return alpha, hmmLogLikelihood(scales), nil
|
||
}
|
||
|
||
// Smooth runs the forward-backward recursions over the observation
|
||
// sequence and returns the smoothed state posteriors P(state at t |
|
||
// all observations), one row per step, with the sequence's log
|
||
// likelihood. The refusals are Forward's own.
|
||
func (m *HiddenMarkovModel) Smooth(observations []int) (smoothed [][]float64, logLikelihood float64, err error) {
|
||
const name = "HiddenMarkovModel.Smooth"
|
||
if m == nil {
|
||
return nil, 0, base.Errf("%s: the model is nil", name)
|
||
}
|
||
if err := hmmCheckSequence(name, observations, m.hmmSymbols()); err != nil {
|
||
return nil, 0, err
|
||
}
|
||
alpha, scales, ferr := m.hmmForward(observations)
|
||
if ferr != nil {
|
||
return nil, 0, ferr
|
||
}
|
||
beta := m.hmmBackward(observations, scales)
|
||
gamma := make([][]float64, len(observations))
|
||
for t := range observations {
|
||
row := make([]float64, len(m.Initial))
|
||
total := 0.0
|
||
for k := range row {
|
||
row[k] = alpha[t][k] * beta[t][k]
|
||
total += row[k]
|
||
}
|
||
for k := range row {
|
||
row[k] /= total
|
||
}
|
||
gamma[t] = row
|
||
}
|
||
return gamma, hmmLogLikelihood(scales), nil
|
||
}
|
||
|
||
// hmmLogLikelihood folds the forward scales into the sequence's log
|
||
// likelihood.
|
||
func hmmLogLikelihood(scales []float64) float64 {
|
||
total := 0.0
|
||
for _, s := range scales {
|
||
total += math.Log(s)
|
||
}
|
||
return total
|
||
}
|
||
|
||
// Viterbi decodes the most likely state path through the observation
|
||
// sequence in log space, breaking ties toward the lowest state index,
|
||
// and returns the path with its log probability
|
||
// log P(path, observations | model). The refusals are Forward's own.
|
||
func (m *HiddenMarkovModel) Viterbi(observations []int) (states []int, logProbability float64, err error) {
|
||
const name = "HiddenMarkovModel.Viterbi"
|
||
if m == nil {
|
||
return nil, 0, base.Errf("%s: the model is nil", name)
|
||
}
|
||
if err := hmmCheckSequence(name, observations, m.hmmSymbols()); err != nil {
|
||
return nil, 0, err
|
||
}
|
||
n := m.hmmStates()
|
||
width := m.hmmWidth()
|
||
scores := make([]float64, n)
|
||
for k := range n {
|
||
scores[k] = math.Log(m.Initial[k]) + math.Log(m.Emission[k*width+observations[0]])
|
||
}
|
||
back := make([][]int, len(observations))
|
||
back[0] = make([]int, n)
|
||
for t := 1; t < len(observations); t++ {
|
||
next := make([]float64, n)
|
||
prev := make([]int, n)
|
||
sym := observations[t]
|
||
for j := range n {
|
||
best := math.Inf(-1)
|
||
arg := 0
|
||
for k := range n {
|
||
candidate := scores[k] + math.Log(m.Transition[k*n+j])
|
||
if candidate > best {
|
||
best = candidate
|
||
arg = k
|
||
}
|
||
}
|
||
next[j] = best + math.Log(m.Emission[j*width+sym])
|
||
prev[j] = arg
|
||
}
|
||
scores = next
|
||
back[t] = prev
|
||
}
|
||
best := math.Inf(-1)
|
||
arg := 0
|
||
for k := range n {
|
||
if scores[k] > best {
|
||
best = scores[k]
|
||
arg = k
|
||
}
|
||
}
|
||
if math.IsInf(best, -1) {
|
||
// Every path has probability zero, the structural zero an
|
||
// emission or transition carries: the same refusal Forward makes,
|
||
// not a meaningless path beside a −Inf log probability.
|
||
return nil, 0, base.Errf("%s: the observation sequence has probability zero under the model", name)
|
||
}
|
||
path := make([]int, len(observations))
|
||
path[len(observations)-1] = arg
|
||
for t := len(observations) - 1; t > 0; t-- {
|
||
path[t-1] = back[t][path[t]]
|
||
}
|
||
return path, best, nil
|
||
}
|
||
|
||
// FitHiddenMarkovModel fits a hidden Markov model of the given state
|
||
// and symbol counts to the observation sequence by Baum-Welch
|
||
// expectation maximisation. The starting point draws every
|
||
// distribution row from the flat Dirichlet through the house
|
||
// generator, so the fit is deterministic for the generator state;
|
||
// expectation maximisation climbs to a local maximum of the
|
||
// likelihood, and a different generator state may land on a different
|
||
// one.
|
||
//
|
||
// The sweeps stop once the log likelihood moves by less than
|
||
// hmmTolerance relative, within hmmMaxIterations sweeps; Converged
|
||
// names which happened. The re-estimation floors every parameter at
|
||
// hmmFloor and renormalises the rows, so no symbol or transition is
|
||
// silenced outright by one sweep. A nil generator, a state or symbol
|
||
// count below one, an empty sequence or an observation outside the
|
||
// symbol set is an error.
|
||
func FitHiddenMarkovModel(g *core.Generator, observations []int, states, symbols int) (*HiddenMarkovFitResult, error) {
|
||
const name = "FitHiddenMarkovModel"
|
||
if g == nil {
|
||
return nil, base.Errf("%s: the generator is nil", name)
|
||
}
|
||
if states < 1 {
|
||
return nil, base.Errf("%s: the model needs at least one state, got %d", name, states)
|
||
}
|
||
if symbols < 1 {
|
||
return nil, base.Errf("%s: the model needs at least one symbol, got %d", name, symbols)
|
||
}
|
||
if err := hmmCheckSequence(name, observations, symbols); err != nil {
|
||
return nil, err
|
||
}
|
||
initial := hmmDrawRow(g, states)
|
||
transition := make([]float64, 0, states*states)
|
||
for range states {
|
||
transition = append(transition, hmmDrawRow(g, states)...)
|
||
}
|
||
emission := make([]float64, 0, states*symbols)
|
||
for range states {
|
||
emission = append(emission, hmmDrawRow(g, symbols)...)
|
||
}
|
||
model, err := NewHiddenMarkovModel(initial, transition, emission)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
prev := math.Inf(-1)
|
||
converged := false
|
||
iterations := hmmMaxIterations
|
||
for iter := 1; iter <= hmmMaxIterations; iter++ {
|
||
logLik, gamma, xi, serr := model.hmmSweep(observations)
|
||
if serr != nil {
|
||
return nil, serr
|
||
}
|
||
move := logLik - prev
|
||
prev = logLik
|
||
if iter > 1 && math.Abs(move) <= hmmTolerance*(1+math.Abs(logLik)) {
|
||
// The model has not moved since the sweep above, so the
|
||
// log likelihood in hand is the fitted model's own.
|
||
converged = true
|
||
iterations = iter
|
||
break
|
||
}
|
||
model = hmmReestimate(model, observations, gamma, xi)
|
||
}
|
||
if !converged {
|
||
// The budget ran out after a re-estimation: the reported
|
||
// likelihood is the final model's own, not its predecessor's.
|
||
logLik, _, _, serr := model.hmmSweep(observations)
|
||
if serr != nil {
|
||
return nil, serr
|
||
}
|
||
prev = logLik
|
||
}
|
||
return &HiddenMarkovFitResult{
|
||
Model: model,
|
||
LogLikelihood: prev,
|
||
Iterations: iterations,
|
||
Converged: converged,
|
||
}, nil
|
||
}
|
||
|
||
// hmmDrawRow draws one distribution row of the given width from the
|
||
// flat Dirichlet through the generator: the normalised exponentials
|
||
// of uniform draws, the classic construction of the simplex's uniform
|
||
// distribution.
|
||
func hmmDrawRow(g *core.Generator, width int) []float64 {
|
||
row := make([]float64, width)
|
||
total := 0.0
|
||
for i := range width {
|
||
weight := -math.Log(max(g.Unit(), 1e-300))
|
||
row[i] = weight
|
||
total += weight
|
||
}
|
||
for i := range row {
|
||
row[i] /= total
|
||
}
|
||
return row
|
||
}
|
||
|
||
// hmmSweep runs the forward-backward pass at the model's current
|
||
// parameters: the log likelihood, the smoothed posteriors gamma and
|
||
// the pairwise posteriors xi the M step reads. xi[t] is the
|
||
// states×states row-major table P(state at t, state at t+1 |
|
||
// observations) for the step from t to t+1.
|
||
func (m *HiddenMarkovModel) hmmSweep(observations []int) (float64, [][]float64, [][]float64, error) {
|
||
states := m.hmmStates()
|
||
width := m.hmmWidth()
|
||
alpha, scales, ferr := m.hmmForward(observations)
|
||
if ferr != nil {
|
||
return 0, nil, nil, ferr
|
||
}
|
||
beta := m.hmmBackward(observations, scales)
|
||
logLik := hmmLogLikelihood(scales)
|
||
gamma := make([][]float64, len(observations))
|
||
for t := range observations {
|
||
row := make([]float64, states)
|
||
total := 0.0
|
||
for k := range states {
|
||
row[k] = alpha[t][k] * beta[t][k]
|
||
total += row[k]
|
||
}
|
||
for k := range states {
|
||
row[k] /= total
|
||
}
|
||
gamma[t] = row
|
||
}
|
||
xi := make([][]float64, len(observations)-1)
|
||
for t := range xi {
|
||
table := make([]float64, states*states)
|
||
total := 0.0
|
||
sym := observations[t+1]
|
||
for k := range states {
|
||
for j := range states {
|
||
v := alpha[t][k] * m.Transition[k*states+j] * m.Emission[j*width+sym] * beta[t+1][j]
|
||
table[k*states+j] = v
|
||
total += v
|
||
}
|
||
}
|
||
for i := range table {
|
||
table[i] /= total
|
||
}
|
||
xi[t] = table
|
||
}
|
||
return logLik, gamma, xi, nil
|
||
}
|
||
|
||
// hmmReestimate applies one Baum-Welch M step: the counts the
|
||
// posteriors carry become the new rows, floored and renormalised.
|
||
func hmmReestimate(model *HiddenMarkovModel, observations []int, gamma, xi [][]float64) *HiddenMarkovModel {
|
||
states := model.hmmStates()
|
||
symbols := model.hmmSymbols()
|
||
initial := make([]float64, states)
|
||
copy(initial, gamma[0])
|
||
transition := make([]float64, states*states)
|
||
for t := range xi {
|
||
for i, v := range xi[t] {
|
||
transition[i] += v
|
||
}
|
||
}
|
||
for k := range states {
|
||
den := 0.0
|
||
for j := range states {
|
||
den += transition[k*states+j]
|
||
}
|
||
if den == 0 {
|
||
// No transition evidence reached this row, the single-
|
||
// observation sequence being the case in point: dividing the
|
||
// zero fills the row with NaN the constructor refuses, and
|
||
// the sweep after it crashed on the nil model that refusal
|
||
// left. The row keeps its previous estimate instead.
|
||
copy(transition[k*states:(k+1)*states], model.Transition[k*states:(k+1)*states])
|
||
continue
|
||
}
|
||
for j := range states {
|
||
transition[k*states+j] = math.Max(transition[k*states+j]/den, hmmFloor)
|
||
}
|
||
hmmRenormalise(transition[k*states : (k+1)*states])
|
||
}
|
||
emission := make([]float64, states*symbols)
|
||
for t, row := range gamma {
|
||
sym := observations[t]
|
||
for k := range states {
|
||
emission[k*symbols+sym] += row[k]
|
||
}
|
||
}
|
||
for k := range states {
|
||
den := 0.0
|
||
for s := range symbols {
|
||
den += emission[k*symbols+s]
|
||
}
|
||
for s := range symbols {
|
||
emission[k*symbols+s] = math.Max(emission[k*symbols+s]/den, hmmFloor)
|
||
}
|
||
hmmRenormalise(emission[k*symbols : (k+1)*symbols])
|
||
}
|
||
// The floors and renormalisations above keep every row a valid
|
||
// distribution, so the constructor's refusals cannot fire here.
|
||
fitted, _ := NewHiddenMarkovModel(initial, transition, emission)
|
||
return fitted
|
||
}
|
||
|
||
// hmmRenormalise scales one distribution row back to sum 1.
|
||
func hmmRenormalise(row []float64) []float64 {
|
||
total := 0.0
|
||
for _, v := range row {
|
||
total += v
|
||
}
|
||
for i := range row {
|
||
row[i] /= total
|
||
}
|
||
return row
|
||
}
|