sc_neurocore_engine/bindings/
surrogate.rs1use pyo3::exceptions::PyValueError;
12use pyo3::prelude::*;
13
14use crate::grad;
15
16pub(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}