sc_neurocore_engine/bindings/
lgssm.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Commercial license available
3// © Concepts 1996–2026 Miroslav Šotek. All rights reserved.
4// © Code 2020–2026 Miroslav Šotek. All rights reserved.
5// ORCID: 0009-0009-3560-0851
6// Contact: www.anulum.li | protoscience@anulum.li
7// SC-NeuroCore — LGSSM Kalman filter PyO3 binding
8
9//! Python binding for the Linear Gaussian state-space model Kalman filter.
10
11use pyo3::prelude::*;
12use pyo3::types::PyDict;
13
14use crate::lgssm;
15
16/// Register the LGSSM Kalman filter with the extension module.
17pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
18    module.add_function(wrap_pyfunction!(py_lgssm_kalman_filter, module)?)?;
19    Ok(())
20}
21
22/// Forward Kalman filter for a Linear Gaussian State-Space Model.
23///
24/// Parity contract with `sc_neurocore.world_model.predictive_model.KalmanFilter`:
25/// for the same model parameters and observation sequence, the
26/// returned (means, covariances, log_likelihood) must agree with
27/// the Python implementation to within float64 round-off.
28///
29/// All matrices are passed as flat row-major Vec<f64>; the caller
30/// supplies their shapes explicitly. Returns a dict with keys:
31///   - "means": Vec<Vec<f64>> shape (T, d)
32///   - "covariances": Vec<Vec<Vec<f64>>> shape (T, d, d)
33///   - "pred_means": Vec<Vec<f64>> shape (T, d)
34///   - "pred_covariances": Vec<Vec<Vec<f64>>> shape (T, d, d)
35///   - "log_likelihood": f64
36///   - "backend": "rust"
37#[pyfunction]
38#[pyo3(signature = (
39    obs_flat, controls_flat, t_len, p_dim, m_dim,
40    a_flat, b_flat, c_flat, d_flat, q_flat, r_flat,
41    mu_0, sigma_0_flat, d_dim,
42))]
43#[allow(clippy::too_many_arguments)]
44fn py_lgssm_kalman_filter<'py>(
45    py: Python<'py>,
46    obs_flat: Vec<f64>,
47    controls_flat: Vec<f64>,
48    t_len: usize,
49    p_dim: usize,
50    m_dim: usize,
51    a_flat: Vec<f64>,
52    b_flat: Vec<f64>,
53    c_flat: Vec<f64>,
54    d_flat: Vec<f64>,
55    q_flat: Vec<f64>,
56    r_flat: Vec<f64>,
57    mu_0: Vec<f64>,
58    sigma_0_flat: Vec<f64>,
59    d_dim: usize,
60) -> PyResult<Py<PyAny>> {
61    use ndarray::Array1;
62    use ndarray::Array2;
63
64    let to_2d = |flat: &[f64], rows: usize, cols: usize| -> Array2<f64> {
65        Array2::from_shape_vec((rows, cols), flat.to_vec()).expect("shape")
66    };
67    let obs = to_2d(&obs_flat, t_len, p_dim);
68    let controls = to_2d(&controls_flat, t_len, m_dim);
69    let a = to_2d(&a_flat, d_dim, d_dim);
70    let b = to_2d(&b_flat, d_dim, m_dim);
71    let c = to_2d(&c_flat, p_dim, d_dim);
72    let d = to_2d(&d_flat, p_dim, m_dim);
73    let q = to_2d(&q_flat, d_dim, d_dim);
74    let r = to_2d(&r_flat, p_dim, p_dim);
75    let mu_0_arr = Array1::from(mu_0);
76    let sigma_0 = to_2d(&sigma_0_flat, d_dim, d_dim);
77
78    let result = lgssm::kalman_filter(
79        obs.view(),
80        controls.view(),
81        a.view(),
82        b.view(),
83        c.view(),
84        d.view(),
85        q.view(),
86        r.view(),
87        mu_0_arr.view(),
88        sigma_0.view(),
89    );
90
91    // Convert to Python-friendly nested Vec
92    let means: Vec<Vec<f64>> = (0..t_len)
93        .map(|t| (0..d_dim).map(|i| result.means[(t, i)]).collect())
94        .collect();
95    let covs: Vec<Vec<Vec<f64>>> = (0..t_len)
96        .map(|t| {
97            (0..d_dim)
98                .map(|i| (0..d_dim).map(|j| result.covariances[(t, i, j)]).collect())
99                .collect()
100        })
101        .collect();
102    let pred_means: Vec<Vec<f64>> = (0..t_len)
103        .map(|t| (0..d_dim).map(|i| result.pred_means[(t, i)]).collect())
104        .collect();
105    let pred_covs: Vec<Vec<Vec<f64>>> = (0..t_len)
106        .map(|t| {
107            (0..d_dim)
108                .map(|i| {
109                    (0..d_dim)
110                        .map(|j| result.pred_covariances[(t, i, j)])
111                        .collect()
112                })
113                .collect()
114        })
115        .collect();
116
117    let dict = PyDict::new(py);
118    dict.set_item("means", means)?;
119    dict.set_item("covariances", covs)?;
120    dict.set_item("pred_means", pred_means)?;
121    dict.set_item("pred_covariances", pred_covs)?;
122    dict.set_item("log_likelihood", result.log_likelihood)?;
123    dict.set_item("backend", "rust")?;
124    Ok(dict.into_any().unbind())
125}
sc_neurocore_engine/bindings/lgssm.rs

sc_neurocore_engine/bindings/
lgssm.rs
