Sequences
Hidden Markov models (forward, backward, posteriors, Viterbi), the scalar Kalman filter, Kaplan–Meier survival and Cox proportional hazards.
This module holds the models that read time. A hidden Markov model is a loop over time whose body is one contraction, and so is a Kalman filter; each definition that names itself compiles to a loop. The forward pass, the backward pass and Viterbi are the same recursion over different semirings: the forward pass sums products and answers how likely the sequence is, Viterbi takes the maximum of products and answers how likely the best path is, and the backward pass is the forward one read from the other end. The module also covers survival analysis — the Kaplan–Meier estimator and Cox’s proportional hazards model — where the “time” is how long each subject lasted.
Hidden Markov model conventions. The emissions arrive already evaluated: em is a T×S matrix whose entry \((t, s)\) is the likelihood of the observation at time t under state s. Build it from a symbol table for discrete observations, or from a density in Distributions for continuous ones. The transition matrix tr is S×S with \(\mathrm{tr}_{ij} = P(\text{next} = j \mid \text{now} = i)\), each row summing to one, and p0 is the initial state distribution, a vector of length S. The sequence needs at least one step.
Limitations. The plain forward and backward passes multiply probabilities at every step, so after a few hundred steps they underflow; use seqLogLikelihood, which rescales each step, for long sequences. alphas, betas and the posteriors built from them are unscaled. Of Baum–Welch, only the emission re-estimation is provided, not the transition re-estimation. viterbi returns the scores of the best path ending in each state, not the path itself. The Kalman filter is the scalar one-state, one-measurement filter.
Notation. Time steps are counted from zero, \(t = 0, 1, \dots, T-1\), and so are states, \(i, j, s = 0, 1, \dots, S-1\). \(E \in \mathbb{R}^{T \times S}\) (the argument em) has entries \(E_{ts}\), and \(e_t \in \mathbb{R}^{S}\) is its row \(t\); \(A \in \mathbb{R}^{S \times S}\) (the argument tr) has entries \(A_{ij}\); \(p_0 \in \Delta^{S-1}\) (the argument p0). \(\odot\) is the entrywise product and \([\,\cdot\,]\) the Iverson bracket.
open Prelude.Sequences
Reading one step of a sequence
A single time step is read by contracting against a one-hot vector over time, and written back by the reverse contraction. The vector of step numbers is built once and carried through each loop.
Building blocks
stepIds(em)— Let \(E \in \mathbb{R}^{T \times S}\). Then \(\operatorname{stepIds}(E) = (0, 1, \dots, T-1) \in \mathbb{R}^{T}\). The step numbers.atStep(em, ids, t)— Let \(E \in \mathbb{R}^{T \times S}\), let \(\iota \in \mathbb{R}^{T}\) (the argumentids) be step numbers and \(t \in \{0, \dots, T-1\}\). Then \(\operatorname{atStep}(E, \iota, t)_s = \sum_{t'} E_{t's} [\,\iota_{t'} = t\,]\), which is \((e_t)_s = E_{ts}\) when \(\iota = (0, 1, \dots, T-1)\). Row t of a T×S matrix, as a vector of length S.scatterAt(v, ids, t)— Let \(v \in \mathbb{R}^{S}\), \(\iota \in \mathbb{R}^{T}\) and \(t \in \{0, \dots, T-1\}\). Then \(\operatorname{scatterAt}(v, \iota, t)_{t's} = [\,\iota_{t'} = t\,]\, v_s\). A T×S matrix holding v in row t and zeros elsewhere.
The forward pass
The forward variable \(\alpha_t(j)\) is the joint probability of the first t+1 observations and being in state j at time t:
\[\alpha_0(j) = p_0(j)\, e_0(j), \qquad \alpha_{t+1}(j) = \Bigl[\sum_i \alpha_t(i)\, A_{ij}\Bigr] e_{t+1}(j), \qquad P(\text{observations}) = \sum_j \alpha_{T-1}(j)\]The scaled pass normalises \(\alpha\) to a distribution at every step and adds up the logarithms of the normalising constants \(c_t\), whose sum is the log-likelihood: \(\ln P = \sum_t \ln c_t\).
forward
forward(em, tr, p0) · em emission likelihoods T×S · tr transition matrix S×S · p0 initial distribution vector S → final forward variable vector S
Let \(E \in \mathbb{R}^{T \times S}\) (the argument em) be the emission likelihoods with rows \(e_t\), \(A \in \mathbb{R}^{S \times S}\) (the argument tr) a row-stochastic transition matrix, and \(p_0 \in \Delta^{S-1}\) (the argument p0) the initial distribution; let \(T \ge 1\). Let \(\alpha_0 = p_0 \odot e_0\) and \(\alpha_{t+1} = (A^{\top} \alpha_t) \odot e_{t+1}\) for \(t = 0, 1, \dots, T-2\). Then
The unscaled forward variable at the last time step.
Intuition Entry \(j\) is the probability of everything observed so far and of ending in state \(j\). It sums over every hidden path, but each step folds all paths into one vector, so the cost is \(T \cdot S^2\) rather than \(S^T\). Normalised to sum to one, it answers “which state are we in now, given the history?”
Use it when you want the current hidden state after a sequence, such as a machine’s health mode or a market regime (normalise the result yourself), or the building block for seqLikelihood. It is unscaled and underflows after a few hundred steps. For a long sequence’s score, use seqLogLikelihood.
state-estimationsequence-modelling
def forward(em, tr, p0) ≔ forwardLoop(rows(em) − 1, em, stepIds(em), tr, p0 · atStep(em, stepIds(em), 0))
seqLikelihood
seqLikelihood(em, tr, p0) · arguments as forward → likelihood scalar
Let \(E, A, p_0\) and \(\alpha_{T-1}\) be as in forward. Then
The probability of the whole observation sequence. Suitable for short sequences only; it underflows on long ones.
Intuition How probable the observations are under this model, counting every hidden path that could have produced them. A model that explains the data well assigns it a high probability.
Use it when you are comparing models on a short sequence: score it under one HMM per class (a word, a gesture, a fault type) and pick the highest, or flag sequences the model finds unlikely. Probabilities shrink geometrically with length, so beyond a few dozen steps switch to seqLogLikelihood, which gives the same answer as a logarithm without underflow.
sequence-classificationmodel-selectionanomaly-detection
def seqLikelihood(em, tr, p0) ≔ einsum("s->", forward(em, tr, p0))
seqLogLikelihood
seqLogLikelihood(em, tr, p0) · arguments as forward → log-likelihood scalar
Let \(E, A, p_0\) be as in forward. Let \(c_0 = \max\bigl(\sum_j (p_0 \odot e_0)_j,\ 10^{-300}\bigr)\) and \(\hat\alpha_0 = (p_0 \odot e_0)/c_0\), and for \(t = 1, 2, \dots, T-1\) let
Then
\[\operatorname{seqLogLikelihood}(E, A, p_0) = \sum_{t=0}^{T-1} \ln c_t .\]The natural logarithm of the sequence’s probability, computed with a rescaled forward pass so it stays accurate on long sequences. Each normalising constant is floored at \(10^{-300}\) before its logarithm is taken.
Intuition Rescale the forward vector to a distribution at every step and record what was divided out. The likelihood is the product of those factors, so its logarithm is their sum, and no intermediate number ever gets tiny.
Use it when you score or compare HMMs on sequences of any length: model selection, classifying a sequence by the best-scoring model, or tracking progress across re-estimation rounds. Sequences of different lengths are not directly comparable, so divide by \(T\) for a per-step score. A step the model deems impossible adds \(\ln 10^{-300} \approx -691\) rather than \(-\infty\).
model-selectionsequence-classificationnumerical-stability
def seqLogLikelihood(em, tr, p0) ≔ logLoop(rows(em) − 1, em, stepIds(em), tr, p0 · atStep(em, stepIds(em), 0) / firstScale(em, p0), ln(firstScale(em, p0)))
Building blocks
ahead(a, tr)— Let \(a \in \mathbb{R}^{S}\) and \(A \in \mathbb{R}^{S \times S}\) (the argumenttr). Then \(\operatorname{ahead}(a, A)_j = \sum_{i} a_i A_{ij}\), that is \(\operatorname{ahead}(a, A) = A^{\top} a\). Carries a state vector one step through the transition matrix.forwardStep(a, at, tr)— Let \(a \in \mathbb{R}^{S}\), let \(u \in \mathbb{R}^{S}\) (the argumentat) be one step’s emission likelihoods and \(A \in \mathbb{R}^{S \times S}\). Then \(\operatorname{forwardStep}(a, u, A) = (A^{\top} a) \odot u\). One forward step: carry through the transitions, then weight by this step’s emissions.forwardLoop(k, em, ids, tr, a)— Let \(k \in \mathbb{N}\), \(E \in \mathbb{R}^{T \times S}\), \(\iota \in \mathbb{R}^{T}\) (the argumentids), \(A \in \mathbb{R}^{S \times S}\) and \(a \in \mathbb{R}^{S}\). Then \(\operatorname{forwardLoop}(0, E, \iota, A, a) = a\) and \(\operatorname{forwardLoop}(k, E, \iota, A, a) = \operatorname{forwardLoop}\bigl(k - 1, E, \iota, A, \operatorname{forwardStep}(a, \operatorname{atStep}(E, \iota, T - k), A)\bigr)\). Applies the remaining k forward steps.scaleOf(a, at, tr)— Let \(a, u \in \mathbb{R}^{S}\) (the argumentsa,at) and \(A \in \mathbb{R}^{S \times S}\). Then \(\operatorname{scaleOf}(a, u, A) = \max\bigl(\sum_j \operatorname{forwardStep}(a, u, A)_j,\ 10^{-300}\bigr)\). The total of one forward step, floored at 1e-300.rescaled(a, at, tr)— Let \(a, u, A\) be as inscaleOf. Then \(\operatorname{rescaled}(a, u, A) = \operatorname{forwardStep}(a, u, A) / \operatorname{scaleOf}(a, u, A)\). One forward step divided by its total.logLoop(k, em, ids, tr, a, run)— Let \(k \in \mathbb{N}\), \(E, \iota, A, a\) be as inforwardLoop, and \(\rho \in \mathbb{R}\) (the argumentrun). Then \(\operatorname{logLoop}(0, E, \iota, A, a, \rho) = \rho\) and, writing \(u = \operatorname{atStep}(E, \iota, T - k)\), \(\operatorname{logLoop}(k, E, \iota, A, a, \rho) = \operatorname{logLoop}\bigl(k - 1, E, \iota, A, \operatorname{rescaled}(a, u, A), \rho + \ln \operatorname{scaleOf}(a, u, A)\bigr)\). Runs the rescaled pass for k steps, adding the log of each scale to the running total.firstScale(em, p0)— Let \(E \in \mathbb{R}^{T \times S}\) and \(p_0 \in \Delta^{S-1}\). Then \(\operatorname{firstScale}(E, p_0) = \max\bigl(\sum_j p_0(j) E_{0j},\ 10^{-300}\bigr)\). The total of the first step’s forward variable, floored at 1e-300.
Every forward variable, kept
The posteriors need \(\alpha_t\) at every time step, not only the last. The loop accumulates a T×S matrix by adding each step’s vector scattered into its own row.
alphas
alphas(em, tr, p0) · em emission likelihoods T×S · tr transition matrix S×S · p0 initial distribution vector S → forward variables T×S
Let \(E, A, p_0\) and the forward variables \(\alpha_0, \dots, \alpha_{T-1}\) be as in forward. Then
Row t is \(\alpha_t\). Unscaled.
Intuition The whole history of the forward pass: row \(t\) is what the model believes about the state at \(t\) from observations up to \(t\) only. Normalised, the rows are the filtered state distribution as it would have looked in real time.
Use it when you want to replay how the belief about the hidden state evolved, for example how early a regime change became visible, or when you need the forward half of posteriors. It is unscaled, so keep sequences short. The rows shrink toward zero as \(t\) grows, so normalise each row before reading it.
state-estimationtime-series
def alphas(em, tr, p0) ≔ alphaLoop(rows(em) − 1, em, stepIds(em), tr, p0 · atStep(em, stepIds(em), 0), scatterAt(p0 · atStep(em, stepIds(em), 0), stepIds(em), 0))
Building blocks
alphaLoop(k, em, ids, tr, a, acc)— Let \(k \in \mathbb{N}\), \(E, \iota, A, a\) be as inforwardLoop, and \(P \in \mathbb{R}^{T \times S}\) (the argumentacc). Then \(\operatorname{alphaLoop}(0, E, \iota, A, a, P) = P\) and, writing \(a' = \operatorname{forwardStep}(a, \operatorname{atStep}(E, \iota, T - k), A)\), \(\operatorname{alphaLoop}(k, E, \iota, A, a, P) = \operatorname{alphaLoop}\bigl(k - 1, E, \iota, A, a', P + \operatorname{scatterAt}(a', \iota, T - k)\bigr)\). Runs the remaining k forward steps, adding each result into its row of acc.
The backward pass
The backward variable \(\beta_t(i)\) is the probability of the observations after time t, given state i at time t. It is the forward recursion with the transition read the other way and the emission taken from the next step:
\[\beta_{T-1}(i) = 1, \qquad \beta_t(i) = \sum_j A_{ij}\, e_{t+1}(j)\, \beta_{t+1}(j)\]betas
betas(em, tr) · em emission likelihoods T×S · tr transition matrix S×S → backward variables T×S
Let \(E \in \mathbb{R}^{T \times S}\) (the argument em) with rows \(e_t\) and \(A \in \mathbb{R}^{S \times S}\) (the argument tr). Let \(\beta_{T-1} = \mathbf{1}_S\) and \(\beta_t = A\,(e_{t+1} \odot \beta_{t+1})\) for \(t = T-2, T-3, \dots, 0\). Then
Row t is \(\beta_t\); the last row is all ones. Unscaled.
Intuition Entry \((t, i)\) says how well being in state \(i\) at time \(t\) explains everything observed afterwards. It is the future’s evidence about the present, which the forward pass cannot see. It does not depend on the initial distribution.
Use it when you combine it with alphas to use both past and future evidence, which posteriors does for you. On its own it is mainly useful for checking or extending the E-step. It is unscaled, so keep sequences short.
state-estimationsequence-modelling
def betas(em, tr) ≔ backwardLoop(rows(em) − 1, em, stepIds(em), tr, rowOnes(transposed(em)), scatterAt(rowOnes(transposed(em)), stepIds(em), rows(em) − 1))
Building blocks
transposed(em)— Let \(E \in \mathbb{R}^{T \times S}\). Then \(\operatorname{transposed}(E) = E^{\top} \in \mathbb{R}^{S \times T}\). The transpose of a T×S matrix.behind(b, at, tr)— Let \(b \in \mathbb{R}^{S}\) be the next step’s backward variable, \(u \in \mathbb{R}^{S}\) (the argumentat) the next step’s emission likelihoods and \(A \in \mathbb{R}^{S \times S}\). Then \(\operatorname{behind}(b, u, A)_i = \sum_{j} A_{ij}\, u_j\, b_j\), that is \(\operatorname{behind}(b, u, A) = A\,(u \odot b)\). One backward step from the next step’s backward variable and emissions.backwardLoop(k, em, ids, tr, b, acc)— Let \(k \in \mathbb{N}\), \(E \in \mathbb{R}^{T \times S}\), \(\iota \in \mathbb{R}^{T}\), \(A \in \mathbb{R}^{S \times S}\), \(b \in \mathbb{R}^{S}\) and \(P \in \mathbb{R}^{T \times S}\) (the argumentacc). Then \(\operatorname{backwardLoop}(0, E, \iota, A, b, P) = P\) and, writing \(b' = \operatorname{behind}(b, \operatorname{atStep}(E, \iota, k), A)\), \(\operatorname{backwardLoop}(k, E, \iota, A, b, P) = \operatorname{backwardLoop}\bigl(k - 1, E, \iota, A, b', P + \operatorname{scatterAt}(b', \iota, k - 1)\bigr)\). Runs the remaining k backward steps, adding each result into its row of acc.
The posteriors and emission re-estimation
The posterior probability of each state at each time is the product of the forward and backward variables, normalised per time step — the expectation step of Baum–Welch:
\[\gamma_t(i) = \frac{\alpha_t(i)\, \beta_t(i)}{\sum_j \alpha_t(j)\, \beta_t(j)}\]With \(\gamma\) in hand, the emission table re-estimates as the share of each state’s time spent emitting each symbol: \(\hat B_{sk} \propto \sum_t \gamma_t(s)\, o_{tk}\).
posteriors
posteriors(em, tr, p0) · em emission likelihoods T×S · tr transition matrix S×S · p0 initial distribution vector S → state posteriors T×S
Let \(E, A, p_0\) be as in forward, with forward variables \(\alpha_t\) (alphas) and backward variables \(\beta_t\) (betas). Then
Row t sums to one. The row totals are floored at \(10^{-300}\) before dividing.
Intuition For each time step, the probability of each hidden state given the whole sequence: the past through \(\alpha\) and the future through \(\beta\). Looking both ways makes it sharper than filtering, because a later observation can revise what an earlier moment probably was.
Use it when you label a recorded sequence after the fact: regime detection, segmenting a signal into phases, soft state labels for another model, or the E-step of Baum–Welch. Taking the most probable state per row can produce a sequence of states the transitions forbid. For one consistent path use Viterbi (viterbi). The internals are unscaled, so keep sequences short.
state-estimationregime-modellingdata-segmentation
def posteriors(em, tr, p0) ≔ joint(em, tr, p0) / lifted(joint(em, tr, p0), max(einsum("ts->t", joint(em, tr, p0)), 1e-300))
reEmission
reEmission(em, tr, p0, o) · arguments as posteriors · o one-hot observed symbols T×K → re-estimated emission table S×K
Let \(E, A, p_0\) and \(\gamma\) be as in posteriors, and let \(O \in \{0, 1\}^{T \times K}\) (the argument o) hold in row \(t\) the one-hot code of the symbol observed at time \(t\). Then
The Baum–Welch update of a discrete emission table: row s is the distribution of symbols observed while in state s, weighted by the posteriors. Each row sums to one.
Intuition If we knew the hidden states, the emission table would just be counts of symbols per state. We do not, so each time step splits its vote among the states in proportion to the posteriors, and the expected counts are normalised per state.
Use it when you learn a discrete emission table from unlabelled sequences. Run this update, rebuild em from the new table, and repeat. The likelihood (seqLogLikelihood) should not fall between rounds. Transitions are not re-estimated here, so they stay as given. Rows of o must be one-hot symbols, and sequences must be short enough for the unscaled posteriors.
parameter-estimationclusteringsequence-modelling
def reEmission(em, tr, p0, o) ≔ emissionCounts(em, tr, p0, o) / lifted(emissionCounts(em, tr, p0, o), max(einsum("sk->s", emissionCounts(em, tr, p0, o)), 1e-300))
Building blocks
joint(em, tr, p0)— Let \(E, A, p_0\) be as inforward. Then \(\operatorname{joint}(E, A, p_0) = \operatorname{alphas}(E, A, p_0) \odot \operatorname{betas}(E, A)\), that is \(\operatorname{joint}(E, A, p_0)_{ti} = \alpha_t(i)\, \beta_t(i)\). The unnormalised posteriors.emissionCounts(em, tr, p0, o)— Let \(E, A, p_0\) be as inforward, \(O \in \{0, 1\}^{T \times K}\) and \(\Gamma = \operatorname{posteriors}(E, A, p_0)\). Then \(\operatorname{emissionCounts}(E, A, p_0, O) = \Gamma^{\top} O\), that is \(C_{sk} = \sum_{t} \Gamma_{ts} O_{tk}\). The expected count of each symbol emitted from each state, S×K.
Viterbi
Viterbi is the forward pass with the sum over previous states replaced by a maximum:
\[\delta_0(j) = p_0(j)\, e_0(j), \qquad \delta_{t+1}(j) = \Bigl[\max_i \delta_t(i)\, A_{ij}\Bigr] e_{t+1}(j)\]The largest entry of the final \(\delta\) is the probability of the single most likely state path. No back-pointers are kept, so the path itself is not returned.
viterbi
viterbi(em, tr, p0) · em emission likelihoods T×S · tr transition matrix S×S · p0 initial distribution vector S → best-path scores vector S
Let \(E, A, p_0\) be as in forward. Let \(\delta_0 = p_0 \odot e_0\) and, for \(t = 0, 1, \dots, T-2\),
Then
\[\operatorname{viterbi}(E, A, p_0) = \delta_{T-1} .\]For each state, the probability of the most likely path that ends there.
Intuition The forward pass asks “how likely, counting all paths?”. Viterbi asks “how likely is the single best path?”, so the sum over previous states becomes a maximum. Entry \(j\) is the probability of the best explanation that ends in state \(j\).
Use it when you decode: you want the one most plausible story, not averaged beliefs. Typical uses are part-of-speech tagging, gene finding and fault diagnosis. This returns the scores only, with no back-pointers, so the path itself is not recovered. Use lastState for its final state and posteriors for per-step beliefs. It is unscaled, so keep sequences short.
sequence-decodingstate-estimation
def viterbi(em, tr, p0) ≔ viterbiLoop(rows(em) − 1, em, stepIds(em), tr, p0 · atStep(em, stepIds(em), 0))
bestPathOdds
bestPathOdds(em, tr, p0) · arguments as viterbi → probability scalar
Let \(E, A, p_0\) and \(\delta_{T-1}\) be as in viterbi. Then
The joint probability of the observations and the single most likely state path.
Intuition The joint probability of the data and its single best explanation. Divided by seqLikelihood, which sums over all paths, it gives the share of the total probability that one path carries: near 1 means the decoding is unambiguous, near 0 means many paths compete.
Use it when you want to judge how much to trust a Viterbi decoding, or to score a sequence by its best explanation rather than by all of them (Viterbi scoring). It is unscaled, so on long sequences it underflows to 0.
sequence-decodinguncertainty-quantification
def bestPathOdds(em, tr, p0) ≔ maximum(viterbi(em, tr, p0))
lastState
lastState(em, tr, p0) · arguments as viterbi → 0/1 mask vector S
Let \(E, A, p_0\) and \(\delta_{T-1}\) be as in viterbi. Then
A mask, not an index: one at the final state of the most likely path. If several states tie, all of them are marked.
Intuition Where the most likely story ends. It is marked as a mask rather than an index so it stays a tensor, and ties mark every tied state.
Use it when you need a hard answer to “what state is the system in at the end?” under the best single explanation, such as the current regime or the final tag. For the probability of each current state instead, normalise forward. If several entries are 1, the best paths tie and the data cannot choose between them.
sequence-decodingstate-estimation
def lastState(em, tr, p0) ≔ viterbi(em, tr, p0) = bestPathOdds(em, tr, p0)
Building blocks
bestFrom(d, tr)— Let \(\delta \in \mathbb{R}^{S}\) (the argumentd) and \(A \in \mathbb{R}^{S \times S}\). Then \(\operatorname{bestFrom}(\delta, A)_j = \max_{i} \delta_i A_{ij}\). For each next state, the largest product of a current score and its transition.viterbiStep(d, at, tr)— Let \(\delta \in \mathbb{R}^{S}\), \(u \in \mathbb{R}^{S}\) (the argumentat) and \(A \in \mathbb{R}^{S \times S}\). Then \(\operatorname{viterbiStep}(\delta, u, A) = \operatorname{bestFrom}(\delta, A) \odot u\). One Viterbi step: best transition, then this step’s emissions.viterbiLoop(k, em, ids, tr, d)— Let \(k \in \mathbb{N}\), \(E \in \mathbb{R}^{T \times S}\), \(\iota \in \mathbb{R}^{T}\), \(A \in \mathbb{R}^{S \times S}\) and \(\delta \in \mathbb{R}^{S}\). Then \(\operatorname{viterbiLoop}(0, E, \iota, A, \delta) = \delta\) and \(\operatorname{viterbiLoop}(k, E, \iota, A, \delta) = \operatorname{viterbiLoop}\bigl(k - 1, E, \iota, A, \operatorname{viterbiStep}(\delta, \operatorname{atStep}(E, \iota, T - k), A)\bigr)\). Applies the remaining k Viterbi steps.
The Kalman filter
The scalar Kalman filter tracks one hidden value \(x_{t+1} = f x_t + \text{noise}(q)\) observed as \(z_t = x_t + \text{noise}(r)\). Each step predicts, then corrects by a gain that weighs the measurement against the prediction by their variances:
\[x^- = f x, \quad P^- = f^2 P + q, \quad K = \frac{P^-}{P^- + r}, \quad x' = x^- + K (z - x^-), \quad P' = (1 - K) P^-\]Here q is the process noise variance, r the measurement noise variance, and p the current variance of the estimate.
gainOf
gainOf(p, f, q, r) · p estimate variance scalar · f transition factor scalar · q process variance scalar · r measurement variance scalar → Kalman gain scalar
Let \(P \ge 0\) (the argument p) be the estimate’s variance, \(f \in \mathbb{R}\) (the argument f) the transition factor, \(q \ge 0\) (the argument q) the process variance and \(r > 0\) (the argument r) the measurement variance. Then
The gain lies between 0 (ignore the measurement) and 1 (trust it fully).
Intuition The gain is a trust ratio: the prediction’s uncertainty divided by the total uncertainty. When the prediction is vague compared with the sensor, \(K\) is near 1 and the measurement wins. When the sensor is noisy compared with the prediction, \(K\) is near 0 and the estimate barely moves.
Use it when you tune a filter. The ratio of q to r sets how responsive the filter is, and the gain shows the effect directly. With constant f, q and r the gain settles to a steady value after a few steps, which tells you the effective smoothing.
state-estimationdenoisingsensor-fusion
def gainOf(p, f, q, r) ≔ spreadAhead(p, f, q) / (spreadAhead(p, f, q) + r)
corrected
corrected(x, p, f, q, r, z) · x current estimate scalar · p its variance · f, q, r as gainOf · z new measurement scalar → updated estimate scalar
Let \(x \in \mathbb{R}\) (the argument x) be the current estimate, \(P, f, q, r\) be as in gainOf, and \(z \in \mathbb{R}\) (the argument z) the new measurement. Then
One predict-and-correct step of the estimate.
Intuition Predict where the value should be, see how far the measurement lands from that prediction (the surprise), and move a fraction \(K\) of the way toward it. It is a weighted average of prediction and measurement, weighted by how much each is trusted.
Use it when measurements arrive one at a time, as in streaming sensors, tracking or live dashboards, and you update the estimate on each arrival. Carry the variance alongside it with narrowed, using the same p for both calls. For a whole recorded series, kalman runs the loop.
state-estimationdenoisingstreaming-data
def corrected(x, p, f, q, r, z) ≔ predicted(x, f) + gainOf(p, f, q, r) · (z − predicted(x, f))
narrowed
narrowed(p, f, q, r) · arguments as gainOf → updated variance scalar
Let \(P, f, q, r\) be as in gainOf. Then
The estimate’s variance after the same step. It does not depend on the measurement.
Intuition Each measurement shrinks the uncertainty: the new variance is the predicted variance times \(1 - K\), so it is always at most the prediction’s. It ignores the measured value, so how confident the filter becomes is fixed in advance by f, q and r.
Use it when you need the estimate’s uncertainty, for error bars (about \(\pm 2\sqrt{P}\) for roughly 95% under Gaussian noise) or to decide when enough measurements have arrived. Because it does not depend on data, the whole variance sequence can be computed before any measurement arrives.
uncertainty-quantificationstate-estimation
def narrowed(p, f, q, r) ≔ (1 − gainOf(p, f, q, r)) · spreadAhead(p, f, q)
kalman
kalman(zs, f, q, r, x0, p0) · zs measurements vector T · f transition factor scalar · q process variance scalar · r measurement variance scalar · x0 initial estimate scalar · p0 initial variance scalar → final estimate scalar
Let \(z = (z_0, \dots, z_{T-1}) \in \mathbb{R}^{T}\) (the argument zs) be the measurements, \(f, q, r\) be as in gainOf, \(x_0 \in \mathbb{R}\) (the argument x0) the initial estimate and \(P_0 \ge 0\) (the argument p0) its variance. For \(t = 0, 1, \dots, T-1\) let
Then
\[\operatorname{kalman}(z, f, q, r, x_0, P_0) = x_T .\]The estimate after the last measurement, the measurements taken in order. A large p0 expresses a vague starting guess.
Intuition Feed a noisy series through predict-and-correct and read off the best current estimate. With \(f = 1\) and \(q = 0\) the level is assumed constant, and the filter converges to the running mean. A positive q lets the level drift, so old measurements are gradually forgotten.
Use it when you want the current value of a slowly varying quantity from noisy readings: a sensor level, a position, or a smoothed price. It returns only the final estimate. For the estimate at every step, loop corrected and narrowed yourself. It handles one scalar state, not vectors or matrices.
denoisingstate-estimationtime-series
def kalman(zs, f, q, r, x0, p0) ≔ filterLoop(count(zs), zs, positions(zs), f, q, r, x0, p0)
Building blocks
predicted(x, f)— Let \(x, f \in \mathbb{R}\). Then \(\operatorname{predicted}(x, f) = f x\). The predicted value.spreadAhead(p, f, q)— Let \(P \ge 0\) (the argumentp), \(f \in \mathbb{R}\) and \(q \ge 0\). Then \(\operatorname{spreadAhead}(P, f, q) = f^2 P + q\). The predicted variance.pickOne(v, ids, t)— Let \(v \in \mathbb{R}^{T}\), \(\iota \in \mathbb{R}^{T}\) (the argumentids) and \(t \in \{0, \dots, T-1\}\). Then \(\operatorname{pickOne}(v, \iota, t) = \sum_{i} v_i [\,\iota_i = t\,]\), which is \(v_t\) when \(\iota = (0, 1, \dots, T-1)\). Entry t of a vector, read by contraction against a one-hot.filterLoop(k, zs, ids, f, q, r, x, p)— Let \(k \in \mathbb{N}\), \(z \in \mathbb{R}^{T}\) (the argumentzs), \(\iota \in \mathbb{R}^{T}\), \(f, q, r\) be as ingainOf, and \(x \in \mathbb{R}\), \(P \ge 0\) (the argumentsx,p). Then \(\operatorname{filterLoop}(0, z, \iota, f, q, r, x, P) = x\) and \(\operatorname{filterLoop}(k, z, \iota, f, q, r, x, P) = \operatorname{filterLoop}\bigl(k - 1, z, \iota, f, q, r, \operatorname{corrected}(x, P, f, q, r, \operatorname{pickOne}(z, \iota, T - k)), \operatorname{narrowed}(P, f, q, r)\bigr)\). Applies the remaining k filter steps and answers the estimate.
Survival
Censoring is why survival analysis exists: a subject who left the study still alive lasted at least that long, which a plain average would throw away. The Kaplan–Meier estimator keeps that information by counting who was still at risk at each event time:
\[\hat S(t) = \prod_{t_j \le t} \left(1 - \frac{d_j}{n_j}\right)\]where the product runs over distinct event times, \(d_j\) is the number of failures at \(t_j\) and \(n_j\) the number still at risk. The product is taken as a sum of logarithms. Subjects sharing a time contribute one factor, not one each. The input need not be sorted.
Conventions: t is a vector of n observed times, and d a vector of n event indicators — 1 if the subject failed at that time, 0 if it was censored. The estimator builds n×n comparison matrices. In the statements below subjects are numbered \(i, j = 0, 1, \dots, n-1\) in the order given, \(t \in \mathbb{R}^{n}\) and \(d \in \{0, 1\}^{n}\) are the arguments t and d, and \(\operatorname{safe}(v) = v + [v = 0]\) entrywise.
hazardAt
hazardAt(t, d) · t times vector n · d event indicators vector n → hazard vector n
Let \(t \in \mathbb{R}^{n}\) (the argument t) be the observed times and \(d \in \{0, 1\}^{n}\) (the argument d) the event indicators. Then
For each subject, the number of failures at its time divided by the number of subjects whose time is at least that time (which always includes the subject itself, so the guard never acts).
Intuition The risk of failing at a moment, given that you made it that far. Only subjects still under observation count in the denominator, so a subject censored earlier stops counting, and one censored later still counts as at risk. Nobody is assumed dead or alive beyond what was seen.
Use it when you want the discrete hazard at each observed time, to spot periods of high risk (such as early failures or wear-out), or as the input to survival. Subjects sharing a time get the same value. At a time with only censorings the hazard is 0.
survival-analysiscensored-data
def hazardAt(t, d) ≔ failuresAt(t, d) / safe(atRisk(t))
survival
survival(t, d) · t times vector n · d event indicators vector n → survival probability vector n
Let \(t \in \mathbb{R}^{n}\) and \(d \in \{0, 1\}^{n}\) be as in hazardAt, with hazards \(h_j = \operatorname{hazardAt}(t, d)_j\), and let \(\phi_j = \operatorname{isFirst}(t)_j\) mark the first subject listed at each distinct time. Then
The Kaplan–Meier estimate \(\hat S(t_i)\) at each subject’s own time, in the order the subjects were given. Factors are floored at \(10^{-300}\) before their logarithm, so a curve that reaches zero reads as a very small positive number.
Intuition To survive past time \(t\) you must survive every event time up to \(t\), so multiply the chances of getting through each one. Censored subjects help by staying in the at-risk counts for as long as they were watched, so their partial information is not thrown away. The curve only steps down at failure times.
Use it when you estimate a survival or retention curve: patient survival, customer churn, or time to machine failure. You can read median survival where the curve first reaches 0.5, and compare groups by computing a curve for each. Values come back in input order, so sort by t to plot. Confidence intervals are not provided.
survival-analysiscensored-data
def survival(t, d) ≔ exp(einsum("ij,j->i", upTo(t), ln(max(1 − hazardAt(t, d), 1e-300))))
Building blocks
sameTime(t)— Let \(t \in \mathbb{R}^{n}\). Then \(\operatorname{sameTime}(t)_{ij} = [\,t_j = t_i\,]\). One where two subjects share a time.idxHeld(t)— Let \(t \in \mathbb{R}^{n}\). Then \(\operatorname{idxHeld}(t)_{ij} = i\). An n×n matrix holding the row index in every row.idxLaid(t)— Let \(t \in \mathbb{R}^{n}\). Then \(\operatorname{idxLaid}(t)_{ij} = j\). An n×n matrix holding the column index in every column.isFirst(t)— Let \(t \in \mathbb{R}^{n}\). Then \(\operatorname{isFirst}(t)_i = \bigl[\, \sum_{j} [\,t_j = t_i\,]\,[\,j < i\,] = 0 \,\bigr]\). One for the first subject listed at each distinct time.atRisk(t)— Let \(t \in \mathbb{R}^{n}\). Then \(\operatorname{atRisk}(t)_i = \sum_{j} [\,t_j \ge t_i\,]\). For each subject, how many subjects have a time at least as large.failuresAt(t, d)— Let \(t \in \mathbb{R}^{n}\) and \(d \in \{0, 1\}^{n}\). Then \(\operatorname{failuresAt}(t, d)_i = \sum_{j} [\,t_j = t_i\,]\, d_j\). For each subject, how many failures share its time.upTo(t)— Let \(t \in \mathbb{R}^{n}\). Then \(\operatorname{upTo}(t)_{ij} = [\,t_j \le t_i\,]\, \operatorname{isFirst}(t)_j\). One where subject j’s time is at most subject i’s and j is the first at its time.
Cox’s proportional hazards
Cox’s model says each subject’s hazard is an unknown baseline multiplied by \(\exp(x_i \cdot \beta)\). The partial likelihood asks only which subject failed first among those still at risk, and the baseline cancels from that comparison, so the hazard ratios \(\exp(\beta)\) are estimated without ever specifying the baseline:
\[\ell(\beta) = \sum_i d_i \Bigl[ x_i \cdot \beta - \ln \sum_{j : t_j \ge t_i} \exp(x_j \cdot \beta) \Bigr], \qquad \nabla \ell = \sum_i d_i \bigl( x_i - \bar x_i(\beta) \bigr)\]where \(\bar x_i(\beta)\) is the hazard-weighted mean of the covariates over the risk set of subject i. Tied times all stay in each other’s risk sets (the Breslow approach). The linear predictor is clamped to \([-30, 30]\) inside the exponential.
Conventions: t and d as for survival; x is the n×p covariate matrix with no intercept column (the baseline absorbs it); b is a coefficient vector of length p. In the statements below \(X \in \mathbb{R}^{n \times p}\) (the argument x) has rows \(x_i^{\top}\), \(\beta \in \mathbb{R}^{p}\) is the argument b, and \(\eta_i = \min\bigl(\max(x_i^{\top} \beta, -30), 30\bigr)\) is the clamped linear predictor.
hazardOf
hazardOf(x, b) · x covariates n×p · b coefficients vector p → relative hazards vector n
Let \(X \in \mathbb{R}^{n \times p}\) (the argument x) be the covariates and \(\beta \in \mathbb{R}^{p}\) (the argument b) the coefficients. Then
Each subject’s relative hazard, with the exponent clamped to \([-30, 30]\).
Intuition How many times riskier a subject is than the unknown baseline. The effects are multiplicative, so a coefficient \(\beta_c\) means each unit of covariate \(c\) multiplies the hazard by \(e^{\beta_c}\) at every moment.
Use it when you score subjects by relative risk with fitted coefficients (from coxFit): rank patients, customers or machines by how soon they are likely to fail. These are ratios, not probabilities or times, because the baseline is never estimated. The exponent is clamped to \([-30, 30]\), so extreme scores saturate.
risk-scoringsurvival-analysis
def hazardOf(x, b) ≔ exp(min(max(apply(x, b), 0 − 30), 30))
coxLogPartial
coxLogPartial(t, d, x, b) · t times vector n · d event indicators vector n · x covariates n×p · b coefficients vector p → log partial likelihood scalar
Let \(t \in \mathbb{R}^{n}\) (the argument t) be the times, \(d \in \{0, 1\}^{n}\) (the argument d) the event indicators, and \(X, \beta\) be as in hazardOf. Then
The objective \(\ell(\beta)\) above. The first term uses the unclamped predictor \(x_i^{\top} \beta\); only the one inside the exponential is clamped, so the two agree with \(\ell(\beta)\) exactly while every \(\lvert x_j^{\top} \beta \rvert \le 30\).
Intuition At each failure, ask: given that someone at risk failed now, how probable is it that it was this subject? That probability is their relative hazard over the total of those at risk. The unknown baseline multiplies both, so it cancels. The objective adds up the logs of those probabilities over all failures.
Use it when you compare coefficient vectors or nested models. Twice the difference in log partial likelihood between nested models gives a likelihood-ratio test. It is also how you watch coxFit improve. Censored subjects contribute only by being in the risk sets.
survival-analysismodel-selection
def coxLogPartial(t, d, x, b) ≔ einsum("i,i->", d, apply(x, b) − ln(max(riskTotal(t, x, b), 1e-300)))
coxScore
coxScore(t, d, x, b) · arguments as coxLogPartial → gradient vector p
Let \(t, d, X, \beta\) be as in coxLogPartial. Then
The score \(\nabla \ell(\beta)\) (exact wherever no predictor is clamped); it is zero at the maximum.
Intuition For each failure, compare the failing subject’s covariates with the risk-weighted average of everyone still at risk. If those who fail tend to have larger values than the model expects, the gradient pushes \(\beta\) up. At the best fit the surprises balance and the score is zero.
Use it when you check whether coxFit has converged (every entry near zero), or drive your own optimiser. Its size depends on the number of events and the covariate scales, so judge “near zero” relative to those.
survival-analysistraining-diagnostics
def coxScore(t, d, x, b) ≔ einsum("i,ip->p", d, x − riskMean(t, x, b))
coxFit
coxFit(k, r, t, d, x) · k number of steps scalar · r step size scalar · t times vector n · d event indicators vector n · x covariates n×p → coefficients vector p
Let \(k \in \mathbb{N}\) (the argument k) be the number of steps, \(r > 0\) (the argument r) the step size, and \(t, d, X\) be as in coxLogPartial. Let \(\beta_0 = 0 \in \mathbb{R}^{p}\) and \(\beta_{s+1} = \beta_s + r\, \operatorname{coxScore}(t, d, X, \beta_s)\) for \(s = 0, 1, \dots, k-1\). Then
Maximises the partial likelihood by k steps of plain gradient ascent, starting from zero. There is no line search and no Newton step, and no standard errors are returned; check convergence by calling coxScore at the result.
Intuition Climb the partial likelihood from “no effect” (\(\beta = 0\)) until the score vanishes. The result says how each covariate multiplies the hazard, \(e^{\beta}\) per unit, while holding the others fixed, without ever modelling how risk changes with time.
Use it when you estimate hazard ratios for risk factors: a treatment’s effect in a trial, drivers of churn, or predictors of failure. The score sums over events, so use a smaller r as the number of events grows, and standardise covariates so one step size suits all. There are no standard errors. If a covariate perfectly orders the failures, \(\beta\) keeps growing, so check coxScore at the result.
survival-analysisrisk-factor-analysis
def coxFit(k, r, t, d, x) ≔ coxLoop(k, r, t, d, x, zeroed(columnTotals(x)))
Building blocks
riskSet(t)— Let \(t \in \mathbb{R}^{n}\). Then \(\operatorname{riskSet}(t)_{ij} = [\,t_j \ge t_i\,]\). One where subject j is still at risk at subject i’s time.riskTotal(t, x, b)— Let \(t \in \mathbb{R}^{n}\), \(X \in \mathbb{R}^{n \times p}\) and \(\beta \in \mathbb{R}^{p}\). Then \(\operatorname{riskTotal}(t, X, \beta)_i = \sum_{j} [\,t_j \ge t_i\,]\, e^{\eta_j}\). The sum of relative hazards over each subject’s risk set.riskMean(t, x, b)— Let \(t, X, \beta\) be as inriskTotal. Then \(\operatorname{riskMean}(t, X, \beta)_{i\cdot} = \bar x_i^{\top}\) with \(\bar x_i = \sum_{j} [\,t_j \ge t_i\,]\, e^{\eta_j}\, x_j \big/ \operatorname{safe}\bigl(\operatorname{riskTotal}(t, X, \beta)_i\bigr)\). The hazard-weighted mean covariate row over each subject’s risk set, n×p.columnTotals(x)— Let \(X \in \mathbb{R}^{n \times p}\). Then \(\operatorname{columnTotals}(X)_c = \sum_{i} X_{ic}\) for each column \(c\). The sum of each column; used to shape the zero starting coefficients.coxLoop(k, r, t, d, x, b)— Let \(k \in \mathbb{N}\), \(r > 0\), \(t, d, X\) be as incoxLogPartialand \(\beta \in \mathbb{R}^{p}\). Then \(\operatorname{coxLoop}(0, r, t, d, X, \beta) = \beta\) and \(\operatorname{coxLoop}(k, r, t, d, X, \beta) = \operatorname{coxLoop}\bigl(k - 1, r, t, d, X, \beta + r \operatorname{coxScore}(t, d, X, \beta)\bigr)\). Applies k gradient-ascent steps to b.
Examples
The examples below use two files, Hmm.hb and Lifetimes.hb:
open Prelude.Sequences
fn likelihoodOf(em, tr, p0) ↦ seqLikelihood(em, tr, p0)
fn logLikelihoodOf(em, tr, p0) ↦ seqLogLikelihood(em, tr, p0)
fn forwardVector(em, tr, p0) ↦ forward(em, tr, p0)
fn statePosteriors(em, tr, p0) ↦ posteriors(em, tr, p0)
fn emissionsRefit(em, tr, p0, o) ↦ reEmission(em, tr, p0, o)
fn viterbiScores(em, tr, p0) ↦ viterbi(em, tr, p0)
fn bestFinal(em, tr, p0) ↦ lastState(em, tr, p0)
fn bestOdds(em, tr, p0) ↦ bestPathOdds(em, tr, p0)
open Prelude.Sequences
fn filteredLevel(zs, f, q, r, x0, p0) ↦ kalman(zs, f, q, r, x0, p0)
fn survivalCurve(t, d) ↦ survival(t, d)
fn coxCoefficients(k, r, t, d, x) ↦ coxFit(k, r, t, d, x)
fn coxGradient(t, d, x, b) ↦ coxScore(t, d, x, b)
fn coxObjective(t, d, x, b) ↦ coxLogPartial(t, d, x, b)
The hidden Markov examples share one model: two states, transitions [[0.7, 0.3], [0.4, 0.6]], a start of [0.6, 0.4], and two symbols emitted with probabilities [0.9, 0.1] from state 0 and [0.2, 0.8] from state 1. The observed symbols are 0, 1, 0, so the emission matrix takes rows from that table.
Score an observation sequence
By hand: \(\alpha_0 = [0.54, 0.08]\), \(\alpha_1 = [0.041, 0.168]\), \(\alpha_2 = [0.08631, 0.02262]\), so the likelihood is 0.10893 and its logarithm −2.21705. The scaled pass gives the same logarithm.
using MLambda.Hilbert.Runtime;
using Examples;
var em = (Tensor)new double[,] { { 0.9, 0.2 }, { 0.1, 0.8 }, { 0.9, 0.2 } };
var tr = (Tensor)new double[,] { { 0.7, 0.3 }, { 0.4, 0.6 } };
var p0 = (Tensor)new[] { 0.6, 0.4 };
var alpha = Hmm.ForwardVector(em, tr, p0);
var like = Hmm.LikelihoodOf(em, tr, p0);
var logL = Hmm.LogLikelihoodOf(em, tr, p0);
Output (verified, six significant digits):
forward = [0.08631, 0.02262]
likelihood = 0.10893
logLikelihood = -2.21705
Recover the hidden states and re-estimate the emissions
The posteriors put the middle observation (symbol 1) mostly in state 1 and the outer ones in state 0. One Baum–Welch emission update sharpens state 0 towards symbol 0.
var o = (Tensor)new double[,] { { 1, 0 }, { 0, 1 }, { 1, 0 } };
var gamma = Hmm.StatePosteriors(em, tr, p0);
var table = Hmm.EmissionsRefit(em, tr, p0, o);
Output (verified, six decimals; matrices row by row):
posteriors = [[0.810521, 0.189479], [0.259708, 0.740292], [0.792344, 0.207656]]
reEmission = [[0.860565, 0.139435], [0.349153, 0.650847]]
Find the most likely path’s probability and final state
By hand: \(\delta_1 = [0.0378, 0.1296]\) and \(\delta_2 = [0.1296 \cdot 0.4 \cdot 0.9,\ 0.1296 \cdot 0.6 \cdot 0.2] = [0.046656, 0.015552]\). The best path (0, 1, 0) has probability 0.046656 and ends in state 0.
var delta = Hmm.ViterbiScores(em, tr, p0);
var odds = Hmm.BestOdds(em, tr, p0);
var last = Hmm.BestFinal(em, tr, p0);
Output (verified):
viterbi = [0.046656, 0.015552]
bestOdds = 0.046656
lastState = [1, 0]
Filter noisy measurements of a constant
With f = 1 and no process noise, one measurement of 2 against a prior of 0 with equal variances gives a gain of one half and an estimate of 1. From a vague prior (variance \(10^6\)), three measurements 1, 2, 3 give very nearly their mean.
var one = Lifetimes.FilteredLevel((Tensor)new[] { 2d }, 1.0, 0.0, 1.0, 0.0, 1.0);
var three = Lifetimes.FilteredLevel((Tensor)new[] { 1d, 2, 3 }, 1.0, 0.0, 1.0, 0.0, 1000000.0);
Output (verified, six decimals):
one measurement = 1
three measurements = 1.999999
Estimate a survival curve with censoring
Six subjects, given in no particular order, with times [3, 1, 4, 2, 5, 2]; the subjects at times 4 and 5 were censored. At time 1, one of six fails: \(5/6\). At time 2, two of the five at risk fail: \(5/6 \cdot 3/5 = 0.5\). At time 3, one of three fails: \(0.5 \cdot 2/3 = 0.3333\), which then holds. Each subject receives the curve at its own time.
var times = (Tensor)new[] { 3d, 1, 4, 2, 5, 2 };
var events = (Tensor)new[] { 1d, 1, 0, 1, 0, 1 };
var curve = Lifetimes.SurvivalCurve(times, events);
Output (verified, six decimals):
survival = [0.333333, 0.833333, 0.333333, 0.5, 0.333333, 0.5]
Fit a Cox model and check the score
Six subjects, all failing, with one binary covariate; the three with covariate 1 tend to fail earlier, so the coefficient is positive. After 200 gradient steps of size 0.1 the score is zero to printing precision, and the log partial likelihood has risen from its value at \(\beta = 0\).
var ct = (Tensor)new[] { 1d, 2, 3, 4, 5, 6 };
var cd = (Tensor)new[] { 1d, 1, 1, 1, 1, 1 };
var cx = (Tensor)new double[,] { { 1 }, { 1 }, { 0 }, { 1 }, { 0 }, { 0 } };
var beta = Lifetimes.CoxCoefficients(200.0, 0.1, ct, cd, cx);
var score = Lifetimes.CoxGradient(ct, cd, cx, beta);
var atFit = Lifetimes.CoxObjective(ct, cd, cx, beta);
var atNil = Lifetimes.CoxObjective(ct, cd, cx, (Tensor)new[] { 0d });
Output (verified, six decimals):
coefficient = [1.687846]
score = [0]
logPartial = -5.34324
logPartial b=0 = -6.579251