404 lines
13 KiB
Go
404 lines
13 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package stats
|
||
|
||
import (
|
||
"math"
|
||
"math/big"
|
||
"testing"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// The hidden Markov machinery against exact referents. The model
|
||
// probabilities are rational, so a short sequence's likelihood and
|
||
// its most likely path are computed exactly by enumerating every
|
||
// path, and the recursions are compared against that enumeration, not
|
||
// against quoted figures.
|
||
|
||
// exactModel holds one small model's parameters as rationals.
|
||
type exactModel struct {
|
||
initial []float64
|
||
transition []float64 // 2×2
|
||
emission []float64 // 2×2
|
||
}
|
||
|
||
// ratRow converts a float row to rationals.
|
||
func ratRow(vals []float64) []*big.Rat {
|
||
out := make([]*big.Rat, len(vals))
|
||
for i, v := range vals {
|
||
out[i] = big.NewRat(int64(v*1e15), 1e15)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// exactPathProbability evaluates one path's joint probability with
|
||
// the observations it explains.
|
||
func exactPathProbability(m *exactModel, observations, path []int) *big.Rat {
|
||
initial := ratRow(m.initial)
|
||
trans := make([][]*big.Rat, 2)
|
||
for k := range 2 {
|
||
trans[k] = ratRow(m.transition[k*2 : (k+1)*2])
|
||
}
|
||
emis := make([][]*big.Rat, 2)
|
||
for k := range 2 {
|
||
emis[k] = ratRow(m.emission[k*2 : (k+1)*2])
|
||
}
|
||
total := new(big.Rat).Set(initial[path[0]])
|
||
total.Mul(total, emis[path[0]][observations[0]])
|
||
for t := 1; t < len(path); t++ {
|
||
total.Mul(total, trans[path[t-1]][path[t]])
|
||
total.Mul(total, emis[path[t]][observations[t]])
|
||
}
|
||
return total
|
||
}
|
||
|
||
func TestHiddenMarkovAgainstExactEnumeration(t *testing.T) {
|
||
m := &exactModel{
|
||
initial: []float64{0.6, 0.4},
|
||
transition: []float64{0.7, 0.3, 0.2, 0.8},
|
||
emission: []float64{0.9, 0.1, 0.25, 0.75},
|
||
}
|
||
model, err := NewHiddenMarkovModel(m.initial, m.transition, m.emission)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
observations := []int{0, 0, 1, 0, 1, 1, 0, 1}
|
||
// Every path over the sequence, exactly.
|
||
bestProb := new(big.Rat)
|
||
totalProb := new(big.Rat)
|
||
var bestPath []int
|
||
path := make([]int, len(observations))
|
||
var walk func(t int)
|
||
walk = func(t int) {
|
||
if t == len(path) {
|
||
p := exactPathProbability(m, observations, path)
|
||
totalProb.Add(totalProb, p)
|
||
if p.Cmp(bestProb) > 0 {
|
||
bestProb.Set(p)
|
||
bestPath = append([]int(nil), path...)
|
||
}
|
||
return
|
||
}
|
||
for _, s := range []int{0, 1} {
|
||
path[t] = s
|
||
walk(t + 1)
|
||
}
|
||
}
|
||
walk(0)
|
||
// The forward recursion's likelihood against the exact total.
|
||
_, ll, err := model.Forward(observations)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
wantLL, _ := new(big.Float).SetRat(totalProb).Float64()
|
||
if math.Abs(ll-math.Log(wantLL)) > 1e-12 {
|
||
t.Fatalf("the forward likelihood = %.16f, want the exact %.16f", ll, math.Log(wantLL))
|
||
}
|
||
// Viterbi against the exact best path.
|
||
states, pathProb, err := model.Viterbi(observations)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
wantBest, _ := new(big.Float).SetRat(bestProb).Float64()
|
||
if math.Abs(pathProb-math.Log(wantBest)) > 1e-12 {
|
||
t.Fatalf("Viterbi's log probability = %.16f, want the exact %.16f", pathProb, math.Log(wantBest))
|
||
}
|
||
for step := range states {
|
||
if states[step] != bestPath[step] {
|
||
t.Fatalf("Viterbi's path %v leaves the exact best path %v at step %d", states, bestPath, step)
|
||
}
|
||
}
|
||
// The smoothing's last row must coincide with filtering's: both
|
||
// see the whole sequence by then.
|
||
filtered, _, err := model.Forward(observations)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
smoothed, _, err := model.Smooth(observations)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
for k := range filtered[len(observations)-1] {
|
||
if filtered[len(observations)-1][k] != smoothed[len(observations)-1][k] {
|
||
t.Fatalf("the last filtered posterior %v disagrees with the smoothed one %v",
|
||
filtered[len(observations)-1], smoothed[len(observations)-1])
|
||
}
|
||
if smoothed[len(observations)-1][k] < 0 || smoothed[len(observations)-1][k] > 1 {
|
||
t.Fatalf("the smoothed posterior %g is not a probability", smoothed[len(observations)-1][k])
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestHiddenMarkovSingleStateLikelihood(t *testing.T) {
|
||
// One state: the model is an independent-symbol law and the
|
||
// likelihood is the product of the emissions, by hand.
|
||
model, err := NewHiddenMarkovModel([]float64{1}, []float64{1}, []float64{0.25, 0.75})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
observations := []int{1, 1, 0, 1, 1, 1, 0, 1, 1}
|
||
_, ll, err := model.Forward(observations)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
want := 0.0
|
||
for _, o := range observations {
|
||
if o == 0 {
|
||
want += math.Log(0.25)
|
||
} else {
|
||
want += math.Log(0.75)
|
||
}
|
||
}
|
||
if math.Abs(ll-want) > 1e-12 {
|
||
t.Fatalf("the single-state likelihood = %.16f, want %.16f", ll, want)
|
||
}
|
||
path, _, err := model.Viterbi(observations)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
for _, s := range path {
|
||
if s != 0 {
|
||
t.Fatalf("the single-state decode left the only state: %v", path)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestHiddenMarkovScalingSurvivesLongSequences(t *testing.T) {
|
||
// A thousand-step sequence under a sticky model: the raw forward
|
||
// probabilities underflow long before the end, and the scaled
|
||
// recursions must answer a finite likelihood and rows that stay
|
||
// distributions.
|
||
model, err := NewHiddenMarkovModel(
|
||
[]float64{0.5, 0.5},
|
||
[]float64{0.99, 0.01, 0.02, 0.98},
|
||
[]float64{0.8, 0.2, 0.3, 0.7})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
observations := make([]int, 1000)
|
||
for t := range observations {
|
||
observations[t] = (t * 7 / 3) % 2
|
||
}
|
||
filtered, ll, err := model.Forward(observations)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if math.IsInf(ll, 0) || math.IsNaN(ll) {
|
||
t.Fatalf("the long sequence's likelihood = %g", ll)
|
||
}
|
||
for step, row := range filtered {
|
||
total := 0.0
|
||
for _, v := range row {
|
||
total += v
|
||
}
|
||
if math.Abs(total-1) > 1e-9 {
|
||
t.Fatalf("the filtered row at %d sums to %g", step, total)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestHiddenMarkovFitRecovery(t *testing.T) {
|
||
// Baum-Welch on a sequence drawn from a known model: the fitted
|
||
// likelihood must clear the known model's own, the fit must
|
||
// converge, and a repeat run must be bit-identical.
|
||
truth, err := NewHiddenMarkovModel(
|
||
[]float64{0.65, 0.35},
|
||
[]float64{0.8, 0.2, 0.15, 0.85},
|
||
[]float64{0.85, 0.15, 0.3, 0.7})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
observations := hmmSimulate(t, truth, 400, 20260925)
|
||
_, truthLL, err := truth.Forward(observations)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
fit, err := FitHiddenMarkovModel(core.NewGenerator(7), observations, 2, 2)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if !fit.Converged {
|
||
t.Fatalf("the fit did not converge (%d iterations)", fit.Iterations)
|
||
}
|
||
if fit.LogLikelihood < truthLL {
|
||
t.Fatalf("the fitted likelihood %.6f sits below the truth's %.6f, expectation maximisation failed to climb",
|
||
fit.LogLikelihood, truthLL)
|
||
}
|
||
again, err := FitHiddenMarkovModel(core.NewGenerator(7), observations, 2, 2)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if again.LogLikelihood != fit.LogLikelihood || again.Iterations != fit.Iterations {
|
||
t.Fatal("two identical fits disagreed")
|
||
}
|
||
for i := range fit.Model.Initial {
|
||
if fit.Model.Initial[i] != again.Model.Initial[i] ||
|
||
fit.Model.Transition[i] != again.Model.Transition[i] ||
|
||
fit.Model.Emission[i] != again.Model.Emission[i] {
|
||
t.Fatal("two identical fits returned different parameters")
|
||
}
|
||
}
|
||
// The fitted model answers the same likelihood through Forward as
|
||
// the fit reported.
|
||
_, ll, err := fit.Model.Forward(observations)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if ll != fit.LogLikelihood {
|
||
t.Fatalf("the fitted model's likelihood %.16f differs from the fit's report %.16f", ll, fit.LogLikelihood)
|
||
}
|
||
}
|
||
|
||
// hmmSimulate draws an observation sequence from a model through the
|
||
// house generator, by inverse transform on the cumulative rows.
|
||
func hmmSimulate(t *testing.T, model *HiddenMarkovModel, length int, seed int64) []int {
|
||
t.Helper()
|
||
g := core.NewGenerator(seed)
|
||
states := len(model.Initial)
|
||
symbols := len(model.Emission) / states
|
||
draw := func(row []float64) int {
|
||
u := g.Unit()
|
||
total := 0.0
|
||
for i, v := range row {
|
||
total += v
|
||
if u < total {
|
||
return i
|
||
}
|
||
}
|
||
return len(row) - 1
|
||
}
|
||
state := draw(model.Initial)
|
||
out := make([]int, length)
|
||
for t := range length {
|
||
out[t] = draw(model.Emission[state*symbols : (state+1)*symbols])
|
||
state = draw(model.Transition[state*states : (state+1)*states])
|
||
}
|
||
return out
|
||
}
|
||
|
||
func TestHiddenMarkovRefusals(t *testing.T) {
|
||
model, err := NewHiddenMarkovModel([]float64{0.6, 0.4}, []float64{0.7, 0.3, 0.2, 0.8}, []float64{0.9, 0.1, 0.25, 0.75})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, _, err := model.Forward(nil); err == nil {
|
||
t.Fatal("an empty sequence was accepted")
|
||
}
|
||
if _, _, err := model.Forward([]int{0, 2, 1}); err == nil {
|
||
t.Fatal("an out-of-range symbol was accepted")
|
||
}
|
||
if _, err := NewHiddenMarkovModel([]float64{0.6, 0.5}, []float64{0.7, 0.3, 0.2, 0.8}, []float64{0.9, 0.1, 0.25, 0.75}); err == nil {
|
||
t.Fatal("an initial distribution summing over 1 was accepted")
|
||
}
|
||
if _, err := NewHiddenMarkovModel([]float64{0.6, 0.4}, []float64{0.7, 0.3, 0.2}, []float64{0.9, 0.1, 0.25, 0.75}); err == nil {
|
||
t.Fatal("a short transition matrix was accepted")
|
||
}
|
||
if _, err := NewHiddenMarkovModel([]float64{0.6, 0.4}, []float64{0.7, 0.3, 0.2, 0.8}, []float64{0.9, 0.1, 0.25}); err == nil {
|
||
t.Fatal("a ragged emission matrix was accepted")
|
||
}
|
||
if _, err := NewHiddenMarkovModel([]float64{1.2, -0.2}, []float64{0.7, 0.3, 0.2, 0.8}, []float64{0.9, 0.1, 0.25, 0.75}); err == nil {
|
||
t.Fatal("a probability outside [0, 1] was accepted")
|
||
}
|
||
if _, err := FitHiddenMarkovModel(nil, []int{0, 1}, 2, 2); err == nil {
|
||
t.Fatal("a nil generator was accepted")
|
||
}
|
||
if _, err := FitHiddenMarkovModel(core.NewGenerator(1), []int{0, 1, 5}, 2, 2); err == nil {
|
||
t.Fatal("an out-of-range training symbol was accepted")
|
||
}
|
||
var nilModel *HiddenMarkovModel
|
||
if _, _, err := nilModel.Forward([]int{0}); err == nil {
|
||
t.Fatal("a nil model was accepted")
|
||
}
|
||
}
|
||
|
||
func TestHiddenMarkovZeroProbabilitySequence(t *testing.T) {
|
||
// A structural zero in the emissions: the only state never emits
|
||
// symbol 1, so a sequence holding it has probability zero and no
|
||
// posterior exists. The recursions must refuse rather than divide
|
||
// the zero normaliser into NaN posteriors.
|
||
model, err := NewHiddenMarkovModel([]float64{1}, []float64{1}, []float64{1, 0})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, _, err := model.Forward([]int{1}); err == nil {
|
||
t.Fatal("a zero-probability sequence was accepted by Forward")
|
||
}
|
||
if _, _, err := model.Smooth([]int{1}); err == nil {
|
||
t.Fatal("a zero-probability sequence was accepted by Smooth")
|
||
}
|
||
// A sequence that dies one step in: the first symbol is live, the
|
||
// second is not, and the refusal is the same.
|
||
if _, _, err := model.Forward([]int{0, 1}); err == nil {
|
||
t.Fatal("a sequence that dies at step 1 was accepted by Forward")
|
||
}
|
||
// The live half of the same model still answers: the zero emission
|
||
// silences nothing a legal sequence needs.
|
||
filtered, ll, err := model.Forward([]int{0, 0})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if ll != 0 || filtered[0][0] != 1 {
|
||
t.Fatalf("the certain sequence answered (%g, %v), want (0, [1])", ll, filtered[0])
|
||
}
|
||
}
|
||
|
||
func TestHiddenMarkovFitSingleObservation(t *testing.T) {
|
||
// One observation carries emission and initial evidence but no
|
||
// transition evidence: the re-estimation divided the zero count into
|
||
// NaN rows, the constructor refused them, and the sweep then called
|
||
// a method on the nil model and crashed the process. The fit must
|
||
// answer with a valid model whose transitions keep their starting
|
||
// estimate.
|
||
g := core.NewGenerator(11)
|
||
res, err := FitHiddenMarkovModel(g, []int{0}, 2, 2)
|
||
if err != nil {
|
||
t.Fatalf("FitHiddenMarkovModel on one observation: %v", err)
|
||
}
|
||
if res.Model == nil {
|
||
t.Fatal("FitHiddenMarkovModel on one observation returned no model")
|
||
}
|
||
if math.IsNaN(res.LogLikelihood) || math.IsInf(res.LogLikelihood, 0) {
|
||
t.Fatalf("log likelihood = %g, want a finite value", res.LogLikelihood)
|
||
}
|
||
for k := range 2 {
|
||
sum := 0.0
|
||
for _, v := range res.Model.Transition[k*2 : k*2+2] {
|
||
if math.IsNaN(v) || v < 0 {
|
||
t.Fatalf("transition row %d holds %g, want probabilities", k, v)
|
||
}
|
||
sum += v
|
||
}
|
||
if math.Abs(sum-1) > 1e-9 {
|
||
t.Fatalf("transition row %d sums to %g, want 1", k, sum)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestViterbiZeroProbabilitySequence(t *testing.T) {
|
||
// Symbol 2 is emitted by no state, so every path through the sequence
|
||
// has probability zero. Forward and Smooth refuse such a sequence, and
|
||
// Viterbi documents the same refusals: a meaningless path beside a
|
||
// −Inf log probability is not an answer.
|
||
model, err := NewHiddenMarkovModel(
|
||
[]float64{0.5, 0.5},
|
||
[]float64{0.5, 0.5, 0.5, 0.5},
|
||
[]float64{0.5, 0.5, 0, 0.5, 0.5, 0},
|
||
)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, _, err := model.Viterbi([]int{0, 2, 1}); err == nil {
|
||
t.Fatal("a zero-probability sequence was accepted by Viterbi")
|
||
}
|
||
// The live symbols of the same model still decode.
|
||
path, lp, err := model.Viterbi([]int{0, 1})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if path[0] != path[1] || lp >= 0 {
|
||
t.Fatalf("Viterbi([0, 1]) = (%v, %g), want a coherent path with a negative log probability", path, lp)
|
||
}
|
||
}
|