`sc_neurocore.world_model.predictive_model` — Linear Gaussian SSM¶

1. Scope¶

This module implements a Linear Gaussian State-Space Model (LGSSM) with the standard inference + learning trinity: Kalman filter (forward), Rauch-Tung-Striebel (RTS) smoother (backward), and Expectation-Maximisation (EM) parameter learner. It serves as the probabilistic predictive substrate for sc_neurocore.world_model.planner.SCPlanner and for any downstream consumer that needs to predict, smooth, or fit linear-Gaussian dynamics over noisy observations.

Implementation references:

Kalman, R.E. (1960). A New Approach to Linear Filtering and Prediction Problems. J. Basic Engineering 82(1): 35-45.
Rauch, H.E., Tung, F. & Striebel, C.T. (1965). Maximum likelihood estimates of linear dynamic systems. AIAA J 3(8): 1445-1450. (the RTS smoother)
Shumway, R.H. & Stoffer, D.S. (1982). An approach to time series smoothing and forecasting using the EM algorithm. J Time Series Analysis 3(4): 253-264.
Bishop, C.M. (2006). Pattern Recognition and Machine Learning, Springer. §13.3 (linear dynamical systems).
Murphy, K.P. (2023). Probabilistic Machine Learning: Advanced Topics, MIT Press. §29 (state-space models).

The previous implementation was a deterministic linear matmul on a randomly-initialised matrix that called itself "stochastic" and admitted "(simplified)" in a comment. It was replaced 2026-04-17 per feedback_sophisticated_from_start.md — a placeholder masquerading as a world model is unacceptable.

2. Model¶

Latent state x_t ∈ R^d, observation y_t ∈ R^p, control u_t ∈ R^m:

Text Only

x_{t+1} = A · x_t + B · u_t + w_t,    w_t ~ N(0, Q)
y_t     = C · x_t + D · u_t + v_t,    v_t ~ N(0, R)

with prior x_0 ~ N(μ_0, Σ_0). All parameters {A, B, C, D, Q, R, μ_0, Σ_0} are estimable from data via the EM algorithm (Shumway & Stoffer 1982).

3. Public API¶

Python

from sc_neurocore.world_model.predictive_model import (
    LinearGaussianSSM,    # the model parameters (dataclass)
    KalmanFilter,         # forward inference
    RTSSmoother,          # backward smoothing
    EMLearner,            # parameter estimation via EM
    PredictiveWorldModel, # legacy wrapper (still re-exported)
    FilterResult,         # KalmanFilter output
    SmoothResult,         # RTSSmoother output
)

LinearGaussianSSM.random(state_dim, obs_dim, control_dim) constructs a random stable LGSSM (spectral radius < 1) for smoke tests and EM initialisation.

4. Algorithm details¶

4.1 Kalman filter (forward pass)¶

Per step t:

Predict: x_pred = A x_filt + B u, P_pred = A P_filt A^T + Q.
Update (Joseph form for numerical stability):
innovation e_t = y_t - C x_pred - D u_t,
innovation covariance S = C P_pred C^T + R,
Kalman gain K = P_pred C^T S^{-1},
filtered mean x_filt = x_pred + K e_t,
filtered covariance P_filt = (I - K C) P_pred (I - K C)^T + K R K^T.
Log-likelihood contribution: -½ (p log 2π + log |S| + e_t^T S^{-1} e_t).

4.2 RTS smoother (backward pass)¶

Per step t = T-2 ... 0:

RTS gain: J_t = P_filt(t) A^T P_pred(t+1)^{-1}.
Smoothed mean: x_smooth(t) = x_filt(t) + J_t (x_smooth(t+1) - x_pred(t+1)).
Smoothed covariance: P_smooth(t) = P_filt(t) + J_t (P_smooth(t+1) - P_pred(t+1)) J_t^T.
Lag-1 cross-covariance: C_smooth(t,t+1) = J_t P_smooth(t+1). (Required by EM.)

The smoother is invariant at t = T-1 (no future to incorporate) and reduces uncertainty for all earlier steps (verified by test_rts_smoother_reduces_uncertainty).

4.3 EM learner¶

Per iteration:

E-step: Kalman filter + RTS smoother → posterior moments.
M-step (closed-form, A and C only — B/D held fixed):
A_new = (Σ E[x_{t+1} x_t^T])(Σ E[x_t x_t^T])^{-1} (sum over t = 0..T-2).
Q_new = (1/(T-1)) Σ (E[x_{t+1} x_{t+1}^T] - A E[x_{t+1} x_t^T]^T - ...) (collapsed RHS).
C_new = (Σ y_t E[x_t]^T)(Σ E[x_t x_t^T])^{-1} (sum over t = 0..T-1).
R_new = (1/T) Σ ((y_t - C E[x_t])(y_t - C E[x_t])^T + C P_t C^T).
μ_0, Σ_0 ← smoothed first state.

EM theory (Dempster et al. 1977) guarantees the log-likelihood is monotone non-decreasing across iterations under exact arithmetic; the Python implementation respects this to within a few units of float64 round-off (test_em_log_likelihood_monotone_non_decreasing).

5. Identifiability¶

Linear Gaussian SSMs have a well-known sign + scale ambiguity (Bishop 2006 §13.3.4): the pair (A, C) and (αA, C/α) are observationally equivalent for any α > 0. Direct parameter recovery is therefore brittle — what IS identifiable is the observation likelihood. Tests verify recovery via held-out log-likelihood, not raw parameter agreement.

6. Pipeline wiring¶

PredictiveWorldModel is consumed by:

sc_neurocore.world_model.planner.SCPlanner — calls predict_next_state / forecast to evaluate candidate action sequences.
tests/test_interfaces_generative_worldmodel.py — exercises the legacy API with state_dim=4, action_dim=2.

Re-exported via sc_neurocore.world_model.__init__:

Python

from sc_neurocore.world_model import PredictiveWorldModel

7. Multi-language acceleration chain¶

Per feedback_multi_language_accel.md (Rust + Julia + Go + Mojo + Python fallback). Current state:

Backend	Status	Source
python (NumPy `linalg.solve`)	✅ implemented	this module
rust (PyO3 + `ndarray` Cholesky)	✅ implemented	`engine/src/lgssm.rs`
julia (juliacall + `LinearAlgebra` LAPACK)	✅ implemented	`src/sc_neurocore/accel/julia/world_model/predictive_model.jl`
go (cgo + ctypes shared library)	✅ implemented	`src/sc_neurocore/accel/go/lgssm/lgssm.go`
mojo (Mojo `--emit shared-lib` + ctypes)	✅ implemented	`src/sc_neurocore/accel/mojo/world_model/lgssm.mojo`

All 5 backends are implemented. The dispatcher (KalmanFilter.filter(backend='auto'|'rust'|'julia'|'go'|'mojo'|'python')) keeps Rust as the 'auto' default for backwards compatibility; explicit backend='mojo' is the fastest measured choice (see §8). All five backends return identical (means, covariances, log-likelihood) results to atol=1e-9 on the parity tests (test_four_backend_parity_when_all_available covers the four non-Mojo backends; Mojo parity verified ad-hoc at ≤1e-15 absolute tolerance on means and covariances, ≤1e-7 on the log-likelihood scalar — see commit message for the parity probe).

The benchmark benchmarks/bench_predictive_model.py runs the workload on every available backend and records unavailable_reason for the ones not yet wired. The accel/julia/world_model/predictive_model.jl file that previously existed was a non-functional placeholder (Python syntax inside a Julia module) and was deleted in the

68 commit; #69 (Mojo LGSSM) is now closed.¶

Mojo @export FFI pattern (per feedback_mojo_026_ffi_pattern): the @export decorator forbids parametric signatures, so the kalman_filter_c function accepts every matrix as a raw Int address (numpy arr.ctypes.data) and the Mojo body reconstructs UnsafePointer[Float64, MutAnyOrigin](unsafe_from_address=addr) inside. Working buffers are heap-allocated via std.memory.alloc and re-cast to the same origin for uniform helper signatures. All matrix algebra (matmul, transpose-matmul, Cholesky, triangular solve) is hand-rolled in 100 lines of Mojo to keep the @export boundary minimal — no parametric helpers leak into the C ABI.

8. Performance¶

Reproducible via:

Bash

python benchmarks/bench_predictive_model.py \
    --json benchmarks/results/bench_predictive_model.json

Workload: 4-D state, 3-D obs, T=200 sequence sampled from the true model. Median + min over 5 repeats. Hardware: Linux 6.17 x86_64, NumPy 2.2.0, Python 3.12.3.

Workload	Backend	Median	Min	Speedup vs Python
Forward Kalman filter	python	10.18 ms	8.15 ms	1.0×
Forward Kalman filter	rust	1.77 ms	1.09 ms	5.7×
Forward Kalman filter	julia	1.99 ms	0.98 ms	5.1×
Forward Kalman filter	go	3.14 ms	2.63 ms	3.2×
Forward Kalman filter	mojo	0.22 ms	0.20 ms	46×
RTS smoother	python	~11 ms	—	1.0×
EM (10 iters)	python	~145 ms	—	1.0×

All five implemented backends produce identical log-likelihood (-288.0601) on this workload (parity verified at ≤ 1e-15 abs-tol on means/covariances, ≤ 1e-7 on the scalar log-lik).

Mojo is the fastest on this (T=200, d=4, p=3) workload by a wide margin — 8× over Rust, 46× over NumPy. The reason: the Mojo kernel does ALL matrix algebra in one straight-line function (no PyO3/cgo per-call marshalling, no LAPACK setup overhead, no dynamic dispatch), and the LLVM backend produces tight SIMD-friendly inner loops with the constant 4×4 matrix shapes inferred at JIT time. Rust is second (PyO3 marshalling adds ~0.5 ms), Julia third (juliacall has comparable overhead to PyO3), Go fourth (ctypes call overhead is the highest of the four FFI paths).

The auto dispatcher still prefers Rust under 'auto' for backwards-compatibility — change to backend='mojo' explicitly to take the speedup. A future cleanup will switch the auto priority to fastest-first per the feedback_fallback_chain_ordering rule.

The RTS smoother and EM learner currently dispatch only to the Python path. Extending all 5 accel backends to RTS + EM is deferred — the marginal value is low because RTS smoothing is already sub-15 ms even in pure Python.

Captured run in benchmarks/results/bench_predictive_model.json.

9. Tests¶

tests/test_world_model/test_predictive_model.py — 20 cases: shape validation, PSD invariance, log-likelihood monotonicity, RTS smoother covariance reduction, low-noise tracking, high-noise prior reliance, EM held-out log-likelihood improvement, legacy wrapper compatibility.
tests/test_planner.py — updated: dropped the 3 tests that enforced the placeholder transition_matrix design; added test_predict_next_state_obeys_ssm_dynamics asserting output == A·x + B·u.
tests/test_interfaces_generative_worldmodel.py — pre-existing 4 tests for the legacy API; pass unchanged after the rewrite (API preserved for backwards compatibility).
src/sc_neurocore/accel/rust/safety/predictive_model.rs — Rust safety mirror for the LGSSM Kalman contract, including shape validation, positive-definite covariance validation, Cholesky solve, Joseph-form covariance update, and log-likelihood.

Focused verification: PYTHONPATH=bridge:src .venv/bin/python -m pytest -q tests/test_world_model.py tests/test_world_model/test_predictive_model.py tests/test_world_model/test_predictive_model_backends.py --no-cov → 77 passed, 3 skipped.

10. Audit completeness — 7-point rule¶

#	Criterion	Status	Notes
1	Pipeline wiring	✅ PASS	`world_model/__init__` re-exports preserved; SCPlanner consumer still passes
2	Multi-angle tests	✅ PASS	77 focused predictive-model tests passed, 3 skipped for unavailable optional paths; PSD invariance, EM monotonicity, identifiability caveat
3	Acceleration path	✅ PASS	python + rust + julia + go + mojo forward Kalman paths documented above; Rust safety mirror validates the LGSSM contract independently
4	Benchmarks	✅ PASS	`benchmarks/bench_predictive_model.py` committed; multi-backend harness handles unavailable backends gracefully
5	Performance docs	✅ PASS	§8 with measured numbers
6	Documentation page	✅ PASS	This page
7	Rules followed	✅ PASS	SPDX/copyright header present. No `# mypy: ignore-errors`. Citation list cites 5 published references.

Net: 0 WARN, 0 FAIL for the documented forward Kalman filter surface.

11. Known issues / followups¶

11.1 Forward Kalman acceleration scope¶

The forward Kalman filter is the implemented multi-language acceleration surface. RTS smoothing and EM learning stay on the Python path; the benchmark harness records this explicitly instead of reporting unavailable backends as failures.

11.2 EM does not estimate B and D¶

The M-step in EMLearner updates only A, C, Q, R, μ_0, Σ_0. Joint estimation of B and D requires augmenting the sufficient statistics with control terms; this is documented in Shumway & Stoffer (1982) Appendix A but not implemented here. Open follow-up: extend EMLearner to optimise B and D when controls are present.

11.3 Identifiability test is held-out-LL not parameter recovery¶

By design — see §5. Direct parameter recovery is brittle due to the LGSSM sign+scale ambiguity. The held-out log-likelihood test is the proper identifiability check.

12. Audit batch identification¶

This page was produced as part of the Antigravity audit (#66 / #62) — third complete audit cycle (after chiplet_gen.simulate_thermal and physics/heat.py). One commit per task per feedback_per_task_full_workflow.md.

7. Performance benchmarks¶

Output from `bench_predictive_model.py`¶