Kalman Filter From Scratch in Python: the Local Level Model

A level x_t follows a random walk: x_t = x_t−1 + w_t with w_t ~ N(0, q). You observe y_t = x_t + v_t with v_t ~ N(0, r), all noises independent. Before the first observation your belief about x₁ is N(m0, p0).

Process ys in order. After each observation, record [mean, variance] of the level given everything seen so far. There is no prediction step before the first observation. Return the records.

Examples

Input ys = [10, 12, 11], q = 1, r = 4, m0 = 8, p0 = 9

Output [[9.3846, 2.7692], [10.6535, 1.9406], [10.8003, 1.6947]]

The gain starts at 9/13 and settles as the variance approaches its steady state.

Input ys = [5, 5, 5, 5], q = 0, r = 1, m0 = 0, p0 = 100

Output [[4.9505, 0.9901], [4.9751, 0.4975], [4.9834, 0.3322], [4.9875, 0.2494]]

q = 0: the level never moves, and the filter becomes a running average that forgets the poor prior.

Input ys = [3], q = 2, r = 1, m0 = 0, p0 = 1

Output [[1.5, 0.5]]

One observation: the posterior is halfway between prior and data, since p0 = r.

Constraints

1 ≤ len(ys) ≤ 10⁵
q ≥ 0, r > 0, p0 > 0
Checked to a relative tolerance of 10⁻⁹

Hints

Hint 1

Keep two numbers, the mean m and the variance P of your belief about the current level.

Hint 2

Before each observation after the first, the level may have moved: add q to P.

Hint 3

The gain is K = P / (P + r). Then m ← m + K(y − m) and P ← (1 − K)P. Record [m, P].

Where candidates go wrong

Adding q before the first observation as well: the prior N(m0, p0) already describes x₁.
Updating P before computing the new mean with the old gain, or computing K after shrinking P.
Writing the update as P ← P − K, which mixes a variance with a dimensionless gain.

The idea

Both the prior and the measurement are Gaussian, so the posterior is Gaussian and two numbers describe it. Combining a belief N(m, P) with a measurement y of variance r gives precision 1/P + 1/r and a mean weighted by precision. Written with the gain K = P/(P + r), that is m + K(y − m) and (1 − K)P. Between measurements the random walk adds variance q and leaves the mean alone.

The result at time t is exactly the conditional distribution of x_t given y₁, …, y_t, the same answer you would get by writing down the joint Gaussian of all levels and observations and conditioning on all of them at once. The filter gets there in O(1) work per observation.

Time O(n) · Space O(1) beyond the output

Reference solution

def kalman_filter(ys, q, r, m0, p0):
    m, p = float(m0), float(p0)
    out = []
    for t, y in enumerate(ys):
        if t > 0:
            p += q                     # predict: the level may have moved
        k = p / (p + r)                # gain: how far to trust the new measurement
        m += k * (y - m)               # update the mean towards y
        p *= 1 - k                     # and shrink the uncertainty
        out.append([m, p])
    return out

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

Trading firms use filters like this to estimate fair prices, hedge ratios and spreads that drift. The interview version is one-dimensional so that the reasoning, not the matrix algebra, is on show.

The theory behind it is in the lesson Statistical inference in code: standard errors, permutation tests, the Kalman filter.

Keep going

Easy10 min

A one-dimensional Kalman filter

Examples

Constraints

Hints

Where candidates go wrong

The idea

Reference solution

Follow-ups an interviewer asks

Why interviewers ask it

Keep going

The autocorrelation function of a series

Regression coefficients with standard errors and t-statistics

An exact permutation test for a difference in means

A one-dimensional Kalman filter