Skip to main content

sc_neurocore_engine/bindings/
surrogate.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 — Surrogate-gradient LIF PyO3 binding
8
9//! Python binding for surrogate-gradient LIF training and shared surrogate parsing.
10
11use pyo3::exceptions::PyValueError;
12use pyo3::prelude::*;
13
14use crate::grad;
15
16/// Register surrogate-gradient LIF training with the extension module.
17pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
18    module.add_class::<PySurrogateLif>()?;
19    Ok(())
20}
21
22pub(crate) fn parse_surrogate(name: &str, k: Option<f32>) -> PyResult<grad::SurrogateType> {
23    let normalized = name.to_ascii_lowercase().replace('-', "_");
24    match normalized.as_str() {
25        "fast_sigmoid" => Ok(grad::SurrogateType::FastSigmoid {
26            k: k.unwrap_or(25.0),
27        }),
28        "superspike" | "super_spike" => Ok(grad::SurrogateType::SuperSpike {
29            k: k.unwrap_or(100.0),
30        }),
31        "arctan" | "arc_tan" => Ok(grad::SurrogateType::ArcTan { k: k.unwrap_or(10.0) }),
32        "straightthrough" | "straight_through" | "ste" => Ok(grad::SurrogateType::StraightThrough),
33        _ => Err(PyValueError::new_err(format!(
34            "Unknown surrogate '{}'. Use one of: fast_sigmoid, superspike, arctan, straight_through.",
35            name
36        ))),
37    }
38}
39
40#[pyclass(
41    name = "SurrogateLif",
42    module = "sc_neurocore_engine.sc_neurocore_engine"
43)]
44pub struct PySurrogateLif {
45    inner: grad::SurrogateLif,
46}
47
48#[pymethods]
49impl PySurrogateLif {
50    #[new]
51    #[pyo3(signature = (
52        data_width=16,
53        fraction=8,
54        v_rest=0,
55        v_reset=0,
56        v_threshold=256,
57        refractory_period=2,
58        surrogate="fast_sigmoid",
59        k=None
60    ))]
61    #[allow(clippy::too_many_arguments)]
62    fn new(
63        data_width: u32,
64        fraction: u32,
65        v_rest: i16,
66        v_reset: i16,
67        v_threshold: i16,
68        refractory_period: i32,
69        surrogate: &str,
70        k: Option<f32>,
71    ) -> PyResult<Self> {
72        let surrogate = parse_surrogate(surrogate, k)?;
73        Ok(Self {
74            inner: grad::SurrogateLif::new(
75                data_width,
76                fraction,
77                v_rest,
78                v_reset,
79                v_threshold,
80                refractory_period,
81                surrogate,
82            ),
83        })
84    }
85
86    #[pyo3(signature = (leak_k, gain_k, i_t, noise_in=0))]
87    fn forward(&mut self, leak_k: i16, gain_k: i16, i_t: i16, noise_in: i16) -> (i32, i16) {
88        self.inner.forward(leak_k, gain_k, i_t, noise_in)
89    }
90
91    fn backward(&mut self, grad_output: f32) -> PyResult<f32> {
92        self.inner
93            .backward(grad_output)
94            .map_err(PyValueError::new_err)
95    }
96
97    fn clear_trace(&mut self) {
98        self.inner.clear_trace();
99    }
100
101    fn reset(&mut self) {
102        self.inner.reset();
103    }
104
105    fn trace_len(&self) -> usize {
106        self.inner.trace_len()
107    }
108}