Linear predict

[hmgaudecker]

Add linear Kalman predict fast path

Fixes #36.

Summary

• Add linear_kalman_predict that uses direct matrix algebra (F @ x + c) instead of the unscented sigma-point transform, for models where all factors use linear or constant transition functions
• maximization_inputs.py auto-selects the fast path via is_all_linear() — no API changes needed
• Refactor likelihood functions to accept a generic predict_func callable instead of hardcoded kalman_predict + transition_func

Benchmark results

Tested on health-cognition (no_feedback_to_investments_linear, 4 latent factors, GPU 8 GiB):

	om.Constraints (unscented)	linear-predict
GPU: per iter (100 iters)	8.87s	8.36s
Speedup	—	~6%
JIT warmup	43.3s	40.4s
GPU memory	Higher (OOMs with ~5 GiB free)	Lower (runs with ~5 GiB free)
CPU: per iter (10 iters)	109.6s	107.9s (~1.6%, within noise)

The main benefit is reduced GPU memory usage — the unscented transform generates 2n+1 sigma points which are expensive to differentiate through, while the linear path uses a single matrix multiply. On a small model (4 factors), the speed gain is modest (~6% GPU), but the memory reduction is the difference between fitting on GPU vs OOMing when memory is constrained.

Test plan

• Added unit tests for linear_kalman_predict and is_all_linear in test_kalman_filters.py
• All 351 existing tests pass
• Benchmarked against om.Constraints on real estimation task

[codecov]

Codecov Report

❌ Patch coverage is 98.67550% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.91%. Comparing base (023775e) to head (5c63203).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/skillmodels/maximization_inputs.py	60.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #87      +/-   ##
==========================================
+ Coverage   96.86%   96.91%   +0.05%     
==========================================
  Files          57       57              
  Lines        4809     4952     +143     
==========================================
+ Hits         4658     4799     +141     
- Misses        151      153       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[janosg]

Maybe we should compare the speed of just the update step to see if the linear one is implemented efficiently. Without a detailed analysis I would expect the linear update step to be at least twice as fast as the unscented one. Of course, in a model with few factors or many measurements per factor, the unscented predict might not have been the bottleneck anyways.

[janosg]

Looks suboptimal but maybe Jax is smart enough at compiling the small array creation away. Have you tried different implementations?

[hmgaudecker]

Confirmed that Jax is smart enough. But I kept a more idiomatic version from the experiments and added a note.

[hmgaudecker]

Re: the question about a linear update step —

The measurement model in skillmodels is always linear (states @ loadings + controls @ control_params), so kalman_update already is the exact linear update. The unscented transform only appears in the predict step (to propagate states through potentially nonlinear transition functions). There is no "unscented update" counterpart — the same kalman_update runs regardless of whether transitions are linear or nonlinear.

The QR decomposition in the update operates on an (n_states+1) × (n_states+1) matrix whose structure depends only on the measurement model (loadings, meas_sd), not on the transition model. So a linear_kalman_update would be identical to the current function — no optimization opportunity here.

(this is Claude, obviously, but it does appear plausible without checking deeply)