sc_neurocore_engine/bindings/
fixed_point_lif.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 — fixed-point LIF PyO3 bindings
8
9//! Python bindings for sequential and parallel fixed-point LIF batch kernels.
10
11use numpy::{PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray1};
12use pyo3::exceptions::PyValueError;
13use pyo3::prelude::*;
14use pyo3::types::PyDict;
15
16use crate::neuron;
17
18/// Register fixed-point LIF batch kernels with the extension module.
19pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
20    module.add_class::<FixedPointLif>()?;
21    module.add_function(wrap_pyfunction!(batch_lif_run, module)?)?;
22    module.add_function(wrap_pyfunction!(batch_lif_run_multi, module)?)?;
23    module.add_function(wrap_pyfunction!(batch_lif_run_varying, module)?)?;
24    Ok(())
25}
26
27#[pyclass(module = "sc_neurocore_engine.sc_neurocore_engine")]
28pub struct FixedPointLif {
29    inner: neuron::FixedPointLif,
30}
31
32#[pymethods]
33impl FixedPointLif {
34    #[new]
35    #[pyo3(signature = (
36        data_width=16,
37        fraction=8,
38        v_rest=0,
39        v_reset=0,
40        v_threshold=256,
41        refractory_period=2
42    ))]
43    fn new(
44        data_width: u32,
45        fraction: u32,
46        v_rest: i16,
47        v_reset: i16,
48        v_threshold: i16,
49        refractory_period: i32,
50    ) -> Self {
51        Self {
52            inner: neuron::FixedPointLif::new(
53                data_width,
54                fraction,
55                v_rest,
56                v_reset,
57                v_threshold,
58                refractory_period,
59            ),
60        }
61    }
62
63    #[pyo3(signature = (leak_k, gain_k, i_t, noise_in=0))]
64    fn step(&mut self, leak_k: i16, gain_k: i16, i_t: i16, noise_in: i16) -> (i32, i16) {
65        self.inner.step(leak_k, gain_k, i_t, noise_in)
66    }
67
68    fn reset(&mut self) {
69        self.inner.reset();
70    }
71
72    fn reset_state(&mut self) {
73        self.reset();
74    }
75
76    fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
77        let dict = PyDict::new(py);
78        dict.set_item("v", self.inner.v)?;
79        dict.set_item("refractory_counter", self.inner.refractory_counter)?;
80        Ok(dict.into_any().unbind())
81    }
82}
83
84/// Run a LIF neuron for N steps with constant inputs.
85///
86/// Returns (spikes: ndarray[i32], voltages: ndarray[i16]).
87#[pyfunction]
88#[pyo3(signature = (
89    n_steps,
90    leak_k,
91    gain_k,
92    i_t,
93    noise_in=0,
94    data_width=16,
95    fraction=8,
96    v_rest=0,
97    v_reset=0,
98    v_threshold=256,
99    refractory_period=2
100))]
101#[allow(clippy::too_many_arguments)]
102fn batch_lif_run<'py>(
103    py: Python<'py>,
104    n_steps: usize,
105    leak_k: i16,
106    gain_k: i16,
107    i_t: i16,
108    noise_in: i16,
109    data_width: u32,
110    fraction: u32,
111    v_rest: i16,
112    v_reset: i16,
113    v_threshold: i16,
114    refractory_period: i32,
115) -> (Bound<'py, PyArray1<i32>>, Bound<'py, PyArray1<i16>>) {
116    let mut lif = neuron::FixedPointLif::new(
117        data_width,
118        fraction,
119        v_rest,
120        v_reset,
121        v_threshold,
122        refractory_period,
123    );
124    let spikes_arr = PyArray1::<i32>::zeros(py, n_steps, false);
125    let voltages_arr = PyArray1::<i16>::zeros(py, n_steps, false);
126
127    // SAFETY: Arrays are newly allocated and contiguous.
128    let spikes_slice = unsafe {
129        spikes_arr
130            .as_slice_mut()
131            .expect("newly allocated spikes array must be contiguous")
132    };
133    // SAFETY: Arrays are newly allocated and contiguous.
134    let voltages_slice = unsafe {
135        voltages_arr
136            .as_slice_mut()
137            .expect("newly allocated voltages array must be contiguous")
138    };
139
140    for i in 0..n_steps {
141        let (s, v) = lif.step(leak_k, gain_k, i_t, noise_in);
142        spikes_slice[i] = s;
143        voltages_slice[i] = v;
144    }
145
146    (spikes_arr, voltages_arr)
147}
148
149/// Run N independent LIF neurons in parallel, each with its own constant input.
150///
151/// Returns (spikes: ndarray[i32, (n_neurons, n_steps)],
152///          voltages: ndarray[i16, (n_neurons, n_steps)]).
153#[pyfunction]
154#[pyo3(signature = (
155    n_neurons,
156    n_steps,
157    leak_k,
158    gain_k,
159    currents,
160    data_width=16,
161    fraction=8,
162    v_rest=0,
163    v_reset=0,
164    v_threshold=256,
165    refractory_period=2
166))]
167#[allow(clippy::too_many_arguments)]
168#[allow(clippy::type_complexity)]
169fn batch_lif_run_multi<'py>(
170    py: Python<'py>,
171    n_neurons: usize,
172    n_steps: usize,
173    leak_k: i16,
174    gain_k: i16,
175    currents: PyReadonlyArray1<'py, i16>,
176    data_width: u32,
177    fraction: u32,
178    v_rest: i16,
179    v_reset: i16,
180    v_threshold: i16,
181    refractory_period: i32,
182) -> PyResult<(Bound<'py, PyArray2<i32>>, Bound<'py, PyArray2<i16>>)> {
183    use rayon::prelude::*;
184
185    let curr_slice = currents
186        .as_slice()
187        .map_err(|e| PyValueError::new_err(format!("Cannot read currents: {e}")))?;
188    if curr_slice.len() != n_neurons {
189        return Err(PyValueError::new_err(format!(
190            "currents length {} does not match n_neurons {}.",
191            curr_slice.len(),
192            n_neurons
193        )));
194    }
195
196    let spikes_arr = PyArray2::<i32>::zeros(py, [n_neurons, n_steps], false);
197    let voltages_arr = PyArray2::<i16>::zeros(py, [n_neurons, n_steps], false);
198
199    if n_neurons == 0 || n_steps == 0 {
200        return Ok((spikes_arr, voltages_arr));
201    }
202
203    // SAFETY: Arrays are newly allocated and contiguous.
204    let spikes_flat = unsafe {
205        spikes_arr
206            .as_slice_mut()
207            .expect("newly allocated spikes array must be contiguous")
208    };
209    // SAFETY: Arrays are newly allocated and contiguous.
210    let voltages_flat = unsafe {
211        voltages_arr
212            .as_slice_mut()
213            .expect("newly allocated voltages array must be contiguous")
214    };
215
216    spikes_flat
217        .par_chunks_mut(n_steps)
218        .zip(voltages_flat.par_chunks_mut(n_steps))
219        .zip(curr_slice.par_iter().copied())
220        .for_each(|((spike_row, voltage_row), i_t)| {
221            let mut lif = neuron::FixedPointLif::new(
222                data_width,
223                fraction,
224                v_rest,
225                v_reset,
226                v_threshold,
227                refractory_period,
228            );
229            for step in 0..n_steps {
230                let (s, v) = lif.step(leak_k, gain_k, i_t, 0);
231                spike_row[step] = s;
232                voltage_row[step] = v;
233            }
234        });
235
236    Ok((spikes_arr, voltages_arr))
237}
238
239/// Run a LIF neuron for N steps with per-step current and optional noise arrays.
240#[pyfunction]
241#[pyo3(signature = (
242    leak_k,
243    gain_k,
244    currents,
245    noises=None,
246    data_width=16,
247    fraction=8,
248    v_rest=0,
249    v_reset=0,
250    v_threshold=256,
251    refractory_period=2
252))]
253#[allow(clippy::too_many_arguments)]
254#[allow(clippy::type_complexity)]
255fn batch_lif_run_varying<'py>(
256    py: Python<'py>,
257    leak_k: i16,
258    gain_k: i16,
259    currents: PyReadonlyArray1<'py, i16>,
260    noises: Option<PyReadonlyArray1<'py, i16>>,
261    data_width: u32,
262    fraction: u32,
263    v_rest: i16,
264    v_reset: i16,
265    v_threshold: i16,
266    refractory_period: i32,
267) -> PyResult<(Bound<'py, PyArray1<i32>>, Bound<'py, PyArray1<i16>>)> {
268    let curr_slice = currents
269        .as_slice()
270        .map_err(|e| PyValueError::new_err(format!("Cannot read currents: {e}")))?;
271    let noise_slice: Option<&[i16]> = match noises.as_ref() {
272        Some(n) => Some(
273            n.as_slice()
274                .map_err(|e| PyValueError::new_err(format!("Cannot read noises: {e}")))?,
275        ),
276        None => None,
277    };
278
279    let n_steps = curr_slice.len();
280    if let Some(ns) = noise_slice {
281        if ns.len() != n_steps {
282            return Err(PyValueError::new_err(format!(
283                "noises length {} does not match currents length {}.",
284                ns.len(),
285                n_steps
286            )));
287        }
288    }
289
290    let mut lif = neuron::FixedPointLif::new(
291        data_width,
292        fraction,
293        v_rest,
294        v_reset,
295        v_threshold,
296        refractory_period,
297    );
298    let spikes_arr = PyArray1::<i32>::zeros(py, n_steps, false);
299    let voltages_arr = PyArray1::<i16>::zeros(py, n_steps, false);
300
301    // SAFETY: Arrays are newly allocated and contiguous.
302    let spikes_slice = unsafe {
303        spikes_arr
304            .as_slice_mut()
305            .expect("newly allocated spikes array must be contiguous")
306    };
307    // SAFETY: Arrays are newly allocated and contiguous.
308    let voltages_slice = unsafe {
309        voltages_arr
310            .as_slice_mut()
311            .expect("newly allocated voltages array must be contiguous")
312    };
313
314    for i in 0..n_steps {
315        let noise_in = noise_slice.map_or(0, |ns| ns[i]);
316        let (s, v) = lif.step(leak_k, gain_k, curr_slice[i], noise_in);
317        spikes_slice[i] = s;
318        voltages_slice[i] = v;
319    }
320
321    Ok((spikes_arr, voltages_arr))
322}
sc_neurocore_engine/bindings/fixed_point_lif.rs

sc_neurocore_engine/bindings/
fixed_point_lif.rs
