1#![allow(
10 clippy::useless_conversion,
11 clippy::needless_range_loop,
12 clippy::too_many_arguments,
13 deprecated
14)]
15
16use numpy::{
17 IntoPyArray, PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2,
18 PyReadwriteArray1, PyUntypedArrayMethods,
19};
20use pyo3::exceptions::PyValueError;
21use pyo3::prelude::*;
22use pyo3::types::PyDict;
23use pyo3::IntoPyObject;
24
25pub mod adc_to_spike;
26pub mod analysis;
27pub mod attention;
28pub mod bitstream;
29pub mod brunel;
30pub mod connectome;
31pub mod conv;
32pub mod cordiv;
33pub mod cortical_column;
34pub mod cortical_inject;
35pub mod dna;
36pub mod ei_network;
37pub mod encoder;
38pub mod evo;
39pub mod fault;
40pub mod fusion;
41#[cfg(feature = "gpu")]
42pub mod gpu;
43pub mod grad;
44pub mod graph;
45pub mod ir;
46pub mod layer;
47pub mod lgssm;
48pub mod network_runner;
49pub mod neuron;
50pub mod neurons;
51pub mod optimizer;
52pub mod partition;
53pub mod phi;
54pub mod photonic;
55pub mod ping;
56pub mod predictive_coding;
57pub mod pyo3_neurons;
58pub mod quantum;
59pub mod rall_dendrite;
60pub mod recorder;
61pub mod recurrent;
62pub mod rk4_neurons;
63pub mod sc_inference;
64pub mod scpn;
65pub mod simd;
66pub mod sobol;
67pub mod supervisor;
68pub mod synapses;
69pub mod topology;
70pub mod wilson_cowan;
71pub mod wong_wang;
72
73#[pyclass(
76 name = "BitStreamTensor",
77 module = "sc_neurocore_engine.sc_neurocore_engine"
78)]
79pub struct PyBitStreamTensor {
80 inner: bitstream::BitStreamTensor,
81}
82
83#[pymethods]
84impl PyBitStreamTensor {
85 #[new]
87 #[pyo3(signature = (dimension=10000, seed=0xACE1))]
88 fn new(dimension: usize, seed: u64) -> Self {
89 use rand::SeedableRng;
90 let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(seed);
91 let data = bitstream::bernoulli_packed(0.5, dimension, &mut rng);
92 Self {
93 inner: bitstream::BitStreamTensor::from_words(data, dimension),
94 }
95 }
96
97 #[staticmethod]
99 fn from_packed(data: Vec<u64>, length: usize) -> PyResult<Self> {
100 if length == 0 {
101 return Err(pyo3::exceptions::PyValueError::new_err(
102 "bitstream length must be > 0",
103 ));
104 }
105 Ok(Self {
106 inner: bitstream::BitStreamTensor::from_words(data, length),
107 })
108 }
109
110 fn xor_inplace(&mut self, other: &PyBitStreamTensor) {
112 self.inner.xor_inplace(&other.inner);
113 }
114
115 fn xor(&self, other: &PyBitStreamTensor) -> PyBitStreamTensor {
117 PyBitStreamTensor {
118 inner: self.inner.xor(&other.inner),
119 }
120 }
121
122 fn rotate_right(&mut self, shift: usize) {
124 self.inner.rotate_right(shift);
125 }
126
127 fn hamming_distance(&self, other: &PyBitStreamTensor) -> f32 {
129 self.inner.hamming_distance(&other.inner)
130 }
131
132 #[staticmethod]
134 fn bundle(vectors: Vec<PyRef<'_, PyBitStreamTensor>>) -> PyBitStreamTensor {
135 let refs: Vec<&bitstream::BitStreamTensor> = vectors.iter().map(|v| &v.inner).collect();
136 PyBitStreamTensor {
137 inner: bitstream::BitStreamTensor::bundle(&refs),
138 }
139 }
140
141 fn popcount(&self) -> u64 {
143 bitstream::popcount(&self.inner)
144 }
145
146 #[getter]
148 fn data(&self) -> Vec<u64> {
149 self.inner.data.clone()
150 }
151
152 #[getter]
154 fn length(&self) -> usize {
155 self.inner.length
156 }
157
158 fn __len__(&self) -> usize {
159 self.inner.length
160 }
161
162 fn __repr__(&self) -> String {
163 format!(
164 "BitStreamTensor(length={}, popcount={})",
165 self.inner.length,
166 bitstream::popcount(&self.inner)
167 )
168 }
169}
170
171#[pyclass(
172 name = "StdpSynapse",
173 module = "sc_neurocore_engine.sc_neurocore_engine"
174)]
175pub struct StdpSynapse {
176 inner: synapses::StdpSynapse,
177}
178
179#[pymethods]
180impl StdpSynapse {
181 #[new]
182 #[pyo3(signature = (initial_weight, data_width=16, fraction=8))]
183 fn new(initial_weight: i16, data_width: u32, fraction: u32) -> Self {
184 Self {
185 inner: synapses::StdpSynapse::new(initial_weight, data_width, fraction),
186 }
187 }
188
189 #[allow(clippy::too_many_arguments)]
190 #[pyo3(signature = (pre_spike, post_spike, a_plus=16, a_minus=-16, decay=250, w_min=0, w_max=32767))]
191 fn step(
192 &mut self,
193 pre_spike: bool,
194 post_spike: bool,
195 a_plus: i16,
196 a_minus: i16,
197 decay: i16,
198 w_min: i16,
199 w_max: i16,
200 ) {
201 let params = synapses::StdpParams {
202 a_plus,
203 a_minus,
204 decay,
205 w_min,
206 w_max,
207 };
208 self.inner.step(pre_spike, post_spike, ¶ms);
209 }
210
211 #[getter]
212 fn weight(&self) -> i16 {
213 self.inner.weight
214 }
215
216 #[setter]
217 fn set_weight(&mut self, value: i16) {
218 self.inner.weight = value;
219 }
220}
221
222#[pyclass(
225 name = "BrunelNetwork",
226 module = "sc_neurocore_engine.sc_neurocore_engine"
227)]
228pub struct PyBrunelNetwork {
229 inner: brunel::BrunelNetwork,
230}
231
232#[pymethods]
233impl PyBrunelNetwork {
234 #[new]
235 #[pyo3(signature = (
236 n_neurons,
237 w_indptr,
238 w_indices,
239 w_data,
240 leak_k,
241 gain_k,
242 ext_lambda,
243 ext_weight_fp,
244 data_width=16,
245 fraction=8,
246 v_rest=0,
247 v_reset=0,
248 v_threshold=256,
249 refractory_period=2,
250 seed=42
251 ))]
252 #[allow(clippy::too_many_arguments)]
253 fn new(
254 n_neurons: usize,
255 w_indptr: PyReadonlyArray1<'_, i64>,
256 w_indices: PyReadonlyArray1<'_, i64>,
257 w_data: PyReadonlyArray1<'_, i16>,
258 leak_k: i16,
259 gain_k: i16,
260 ext_lambda: f64,
261 ext_weight_fp: i16,
262 data_width: u32,
263 fraction: u32,
264 v_rest: i16,
265 v_reset: i16,
266 v_threshold: i16,
267 refractory_period: i32,
268 seed: u64,
269 ) -> PyResult<Self> {
270 let indptr = w_indptr
271 .as_slice()
272 .map_err(|e| PyValueError::new_err(format!("Cannot read w_indptr: {e}")))?;
273 let indices = w_indices
274 .as_slice()
275 .map_err(|e| PyValueError::new_err(format!("Cannot read w_indices: {e}")))?;
276 let data = w_data
277 .as_slice()
278 .map_err(|e| PyValueError::new_err(format!("Cannot read w_data: {e}")))?;
279
280 let row_offsets: Vec<usize> = indptr.iter().map(|&v| v as usize).collect();
281 let col_indices: Vec<usize> = indices.iter().map(|&v| v as usize).collect();
282 let values: Vec<i16> = data.to_vec();
283
284 let inner = brunel::BrunelNetwork::new(
285 n_neurons,
286 row_offsets,
287 col_indices,
288 values,
289 data_width,
290 fraction,
291 v_rest,
292 v_reset,
293 v_threshold,
294 refractory_period,
295 leak_k,
296 gain_k,
297 ext_lambda,
298 ext_weight_fp,
299 seed,
300 )
301 .map_err(PyValueError::new_err)?;
302
303 Ok(Self { inner })
304 }
305
306 fn run<'py>(&mut self, py: Python<'py>, n_steps: usize) -> Bound<'py, PyArray1<u32>> {
307 let counts = self.inner.run(n_steps);
308 counts.into_pyarray(py)
309 }
310}
311
312#[pyclass(
315 name = "NetworkRunner",
316 module = "sc_neurocore_engine.sc_neurocore_engine"
317)]
318pub struct PyNetworkRunner {
319 inner: network_runner::NetworkRunner,
320}
321
322#[pymethods]
323impl PyNetworkRunner {
324 #[new]
325 fn new() -> Self {
326 Self {
327 inner: network_runner::NetworkRunner::new(),
328 }
329 }
330
331 fn add_population(&mut self, model: &str, n: usize) -> PyResult<usize> {
332 let pop = network_runner::create_population(model, n).map_err(PyValueError::new_err)?;
333 Ok(self.inner.add_population(pop))
334 }
335
336 #[pyo3(signature = (src, tgt, row_offsets, col_indices, values, delay=0))]
337 fn add_projection(
338 &mut self,
339 src: usize,
340 tgt: usize,
341 row_offsets: Vec<i64>,
342 col_indices: Vec<i64>,
343 values: Vec<f64>,
344 delay: usize,
345 ) {
346 let ro: Vec<usize> = row_offsets.iter().map(|&x| x as usize).collect();
347 let ci: Vec<usize> = col_indices.iter().map(|&x| x as usize).collect();
348 let proj = network_runner::ProjectionRunner::new(src, tgt, ro, ci, values, delay);
349 self.inner.add_projection(proj);
350 }
351
352 fn step_population<'py>(
353 &mut self,
354 py: Python<'py>,
355 pop_index: usize,
356 currents: PyReadonlyArray1<'py, f64>,
357 ) -> PyResult<Py<PyAny>> {
358 let currents = currents.as_slice()?;
359 let (spikes, voltages) = self
360 .inner
361 .step_population_with_currents(pop_index, currents)
362 .map_err(PyValueError::new_err)?;
363 let dict = PyDict::new(py);
364 dict.set_item("spikes", spikes.into_pyarray(py))?;
365 dict.set_item("voltages", voltages.into_pyarray(py))?;
366 Ok(dict.into_any().unbind())
367 }
368
369 fn run<'py>(&mut self, py: Python<'py>, n_steps: usize) -> PyResult<Py<PyAny>> {
370 let results = self.inner.run(n_steps);
371 let dict = PyDict::new(py);
372 let spike_counts: Vec<u64> = results.spike_counts.iter().map(|&c| c as u64).collect();
373 dict.set_item("spike_counts", spike_counts.into_pyarray(py))?;
374 let spike_data: Vec<Py<PyArray1<u64>>> = results
375 .spike_data
376 .into_iter()
377 .map(|v: Vec<u64>| v.into_pyarray(py).unbind())
378 .collect();
379 dict.set_item("spike_data", spike_data)?;
380 let voltages: Vec<Py<PyArray1<f64>>> = results
381 .voltages
382 .into_iter()
383 .map(|v: Vec<f64>| v.into_pyarray(py).unbind())
384 .collect();
385 dict.set_item("voltages", voltages)?;
386 Ok(dict.into_any().unbind())
387 }
388
389 #[staticmethod]
390 fn supported_models() -> Vec<&'static str> {
391 network_runner::supported_models()
392 }
393}
394
395#[pyfunction]
400fn py_batch_simulate<'py>(
401 py: Python<'py>,
402 model_name: &str,
403 n_steps: usize,
404 current_trace: PyReadonlyArray1<'py, f64>,
405) -> PyResult<Py<PyAny>> {
406 let mut neuron = network_runner::create_neuron(model_name).map_err(PyValueError::new_err)?;
407 let currents = current_trace.as_slice()?;
408 let steps = n_steps.min(currents.len());
409
410 let mut voltages = vec![0.0f64; steps];
411 let mut spikes: Vec<u64> = Vec::new();
412
413 for t in 0..steps {
414 let fired = neuron.step(currents[t]);
415 voltages[t] = neuron.soma_voltage();
416 if fired != 0 {
417 spikes.push(t as u64);
418 }
419 }
420
421 let d = PyDict::new(py);
422 d.set_item("voltages", voltages.into_pyarray(py))?;
423 d.set_item("spikes", spikes.into_pyarray(py))?;
424 d.set_item("n_steps", steps)?;
425 Ok(d.into_any().unbind())
426}
427
428#[pyfunction]
431#[pyo3(signature = (
432 n_exc=80, n_inh=20,
433 w_ee=0.1, w_ei=0.4, w_ie=0.1, w_ii=0.4,
434 p_conn=0.2, ext_rate=5.0,
435 duration=200.0, dt=0.1, seed=42
436))]
437fn py_simulate_ei_network<'py>(
438 py: Python<'py>,
439 n_exc: usize,
440 n_inh: usize,
441 w_ee: f64,
442 w_ei: f64,
443 w_ie: f64,
444 w_ii: f64,
445 p_conn: f64,
446 ext_rate: f64,
447 duration: f64,
448 dt: f64,
449 seed: u64,
450) -> PyResult<Py<PyAny>> {
451 let r = ei_network::simulate_ei(
452 n_exc, n_inh, w_ee, w_ei, w_ie, w_ii, p_conn, ext_rate, duration, dt, seed,
453 );
454 let n_spikes = r.spike_times.len();
455 let d = PyDict::new(py);
456 d.set_item("spike_times", r.spike_times.into_pyarray(py))?;
457 d.set_item(
458 "spike_neurons",
459 r.spike_neurons
460 .iter()
461 .map(|&x| x as i64)
462 .collect::<Vec<_>>()
463 .into_pyarray(py),
464 )?;
465 d.set_item("n_exc", r.n_exc)?;
466 d.set_item("n_inh", r.n_inh)?;
467 d.set_item("n_total", r.n_exc + r.n_inh)?;
468 d.set_item("n_spikes", n_spikes)?;
469 d.set_item("rate_time", r.rate_time.into_pyarray(py))?;
470 d.set_item("exc_rates", r.exc_rates.into_pyarray(py))?;
471 d.set_item("inh_rates", r.inh_rates.into_pyarray(py))?;
472 d.set_item("duration", duration)?;
473 d.set_item("dt", dt)?;
474 d.set_item("mean_exc_rate", r.mean_exc_rate)?;
475 d.set_item("mean_inh_rate", r.mean_inh_rate)?;
476 Ok(d.into_any().unbind())
477}
478
479#[pyfunction]
492#[pyo3(signature = (
493 s1_init, s2_init,
494 tau_s, gamma, j_n, j_cross, i_0, sigma, dt,
495 stim1, stim2, xi,
496))]
497#[allow(clippy::too_many_arguments)]
498fn py_wong_wang_simulate<'py>(
499 py: Python<'py>,
500 s1_init: f64,
501 s2_init: f64,
502 tau_s: f64,
503 gamma: f64,
504 j_n: f64,
505 j_cross: f64,
506 i_0: f64,
507 sigma: f64,
508 dt: f64,
509 stim1: PyReadonlyArray1<'py, f64>,
510 stim2: PyReadonlyArray1<'py, f64>,
511 xi: PyReadonlyArray1<'py, f64>,
512) -> PyResult<Py<PyAny>> {
513 let stim1 = stim1.as_slice()?;
514 let stim2 = stim2.as_slice()?;
515 let xi = xi.as_slice()?;
516 let n = stim1.len();
517 if stim2.len() != n {
518 return Err(PyValueError::new_err(format!(
519 "stim1 and stim2 length mismatch: {n} vs {}",
520 stim2.len()
521 )));
522 }
523 if xi.len() != 2 * n {
524 return Err(PyValueError::new_err(format!(
525 "xi length must be 2 * n_steps ({}): got {}",
526 2 * n,
527 xi.len()
528 )));
529 }
530 let mut s1o = vec![0.0_f64; n];
531 let mut s2o = vec![0.0_f64; n];
532 let mut r1o = vec![0.0_f64; n];
533 let mut r2o = vec![0.0_f64; n];
534 let (s1_final, s2_final) = wong_wang::simulate(
535 s1_init, s2_init, tau_s, gamma, j_n, j_cross, i_0, sigma, dt, stim1, stim2, xi, &mut s1o,
536 &mut s2o, &mut r1o, &mut r2o,
537 );
538 let d = PyDict::new(py);
539 d.set_item("s1", s1o.into_pyarray(py))?;
540 d.set_item("s2", s2o.into_pyarray(py))?;
541 d.set_item("r1", r1o.into_pyarray(py))?;
542 d.set_item("r2", r2o.into_pyarray(py))?;
543 d.set_item("s1_final", s1_final)?;
544 d.set_item("s2_final", s2_final)?;
545 Ok(d.into_any().unbind())
546}
547
548#[pyfunction]
563#[pyo3(signature = (spikes, weights_q88, centres_q88, sigmas_q88, n_taps))]
564fn py_dcls_max_forward_batch_q88<'py>(
565 py: Python<'py>,
566 spikes: PyReadonlyArray1<'py, u8>,
567 weights_q88: PyReadonlyArray1<'py, i16>,
568 centres_q88: PyReadonlyArray1<'py, i16>,
569 sigmas_q88: PyReadonlyArray1<'py, i16>,
570 n_taps: usize,
571) -> PyResult<Py<PyAny>> {
572 let result = scpn::dcls_max_forward_batch_q88(
573 spikes.as_slice()?,
574 weights_q88.as_slice()?,
575 centres_q88.as_slice()?,
576 sigmas_q88.as_slice()?,
577 n_taps,
578 )
579 .map_err(|e| PyValueError::new_err(e.to_string()))?;
580 let active_tap_counts: Vec<i64> = result.active_tap_counts.iter().map(|&c| c as i64).collect();
581 let d = PyDict::new(py);
582 d.set_item("outputs_q88", result.outputs_q88.into_pyarray(py))?;
583 d.set_item(
584 "accumulators_q16_16",
585 result.accumulators_q16_16.into_pyarray(py),
586 )?;
587 d.set_item("overflow", result.overflow.into_pyarray(py))?;
588 d.set_item("active_tap_counts", active_tap_counts.into_pyarray(py))?;
589 d.set_item("max_gates_q88", result.max_gates_q88.into_pyarray(py))?;
590 Ok(d.into_any().unbind())
591}
592
593#[pyfunction]
606#[pyo3(signature = (weights_packed, n_out, n_in, n_words, input_probs, length, seed))]
607#[allow(clippy::too_many_arguments)]
608fn py_sc_forward_packed<'py>(
609 py: Python<'py>,
610 weights_packed: PyReadonlyArray1<'py, u64>,
611 n_out: usize,
612 n_in: usize,
613 n_words: usize,
614 input_probs: PyReadonlyArray1<'py, f64>,
615 length: usize,
616 seed: u64,
617) -> PyResult<Bound<'py, PyArray1<f64>>> {
618 let outputs = sc_inference::sc_forward_packed(
619 weights_packed.as_slice()?,
620 n_out,
621 n_in,
622 n_words,
623 input_probs.as_slice()?,
624 length,
625 seed,
626 )
627 .map_err(|e| PyValueError::new_err(e.to_string()))?;
628 Ok(outputs.into_pyarray(py))
629}
630
631#[pyfunction]
643#[pyo3(signature = (
644 samples, adc_width, q_int, q_frac, decimation, signed_input, threshold_q,
645))]
646#[allow(clippy::too_many_arguments)]
647fn py_adc_to_spike_windows<'py>(
648 py: Python<'py>,
649 samples: PyReadonlyArray1<'py, i64>,
650 adc_width: u32,
651 q_int: u32,
652 q_frac: u32,
653 decimation: u32,
654 signed_input: i64,
655 threshold_q: i64,
656) -> PyResult<Py<PyAny>> {
657 let result = adc_to_spike::adc_to_spike_windows(
658 samples.as_slice()?,
659 adc_width,
660 q_int,
661 q_frac,
662 decimation,
663 signed_input != 0,
664 threshold_q,
665 )
666 .map_err(|e| PyValueError::new_err(e.to_string()))?;
667 let d = PyDict::new(py);
668 d.set_item("window_values_q", result.window_values_q.into_pyarray(py))?;
669 d.set_item("spike_counts", result.spike_counts.into_pyarray(py))?;
670 d.set_item("polarities", result.polarities.into_pyarray(py))?;
671 Ok(d.into_any().unbind())
672}
673
674#[pyfunction]
686#[pyo3(signature = (weights_q88, inputs_q1616, n_outputs, n_inputs))]
687fn py_mixed_dense_forward_batch_q88_q1616<'py>(
688 py: Python<'py>,
689 weights_q88: PyReadonlyArray1<'py, i16>,
690 inputs_q1616: PyReadonlyArray1<'py, i32>,
691 n_outputs: usize,
692 n_inputs: usize,
693) -> PyResult<Py<PyAny>> {
694 let result = crate::ir::qformat::mixed_dense_forward_batch_q88_q1616(
695 weights_q88.as_slice()?,
696 inputs_q1616.as_slice()?,
697 n_outputs,
698 n_inputs,
699 )
700 .map_err(|e| PyValueError::new_err(e.to_string()))?;
701 let d = PyDict::new(py);
702 d.set_item("outputs_q1616", result.outputs_q1616.into_pyarray(py))?;
703 d.set_item("overflow", result.overflow.into_pyarray(py))?;
704 d.set_item("underflow", result.underflow.into_pyarray(py))?;
705 Ok(d.into_any().unbind())
706}
707
708#[pyfunction]
716#[pyo3(signature = (
717 e_init, i_init,
718 w_ee, w_ei, w_ie, w_ii,
719 tau_e, tau_i,
720 a, theta, dt,
721 ext_input,
722))]
723#[allow(clippy::too_many_arguments)]
724fn py_wilson_cowan_simulate<'py>(
725 py: Python<'py>,
726 e_init: f64,
727 i_init: f64,
728 w_ee: f64,
729 w_ei: f64,
730 w_ie: f64,
731 w_ii: f64,
732 tau_e: f64,
733 tau_i: f64,
734 a: f64,
735 theta: f64,
736 dt: f64,
737 ext_input: PyReadonlyArray1<'py, f64>,
738) -> PyResult<Py<PyAny>> {
739 let ext = ext_input.as_slice()?;
740 let n = ext.len();
741 let mut e_out = vec![0.0_f64; n];
742 let mut i_out = vec![0.0_f64; n];
743 let (e_final, i_final) = wilson_cowan::simulate(
744 e_init, i_init, w_ee, w_ei, w_ie, w_ii, tau_e, tau_i, a, theta, dt, ext, &mut e_out,
745 &mut i_out,
746 );
747 let d = PyDict::new(py);
748 d.set_item("e", e_out.into_pyarray(py))?;
749 d.set_item("i", i_out.into_pyarray(py))?;
750 d.set_item("e_final", e_final)?;
751 d.set_item("i_final", i_final)?;
752 Ok(d.into_any().unbind())
753}
754
755#[pymodule]
758fn sc_neurocore_engine(m: &Bound<'_, PyModule>) -> PyResult<()> {
759 m.add("__version__", env!("CARGO_PKG_VERSION"))?;
760 m.add_function(wrap_pyfunction!(simd_tier, m)?)?;
761 m.add_function(wrap_pyfunction!(set_num_threads, m)?)?;
762 m.add_function(wrap_pyfunction!(pack_bitstream, m)?)?;
763 m.add_function(wrap_pyfunction!(unpack_bitstream, m)?)?;
764 m.add_function(wrap_pyfunction!(popcount, m)?)?;
765 m.add_function(wrap_pyfunction!(pack_bitstream_numpy, m)?)?;
766 m.add_function(wrap_pyfunction!(popcount_numpy, m)?)?;
767 m.add_function(wrap_pyfunction!(unpack_bitstream_numpy, m)?)?;
768 m.add_function(wrap_pyfunction!(batch_lif_run, m)?)?;
769 m.add_function(wrap_pyfunction!(batch_lif_run_multi, m)?)?;
770 m.add_function(wrap_pyfunction!(batch_lif_run_varying, m)?)?;
771 m.add_function(wrap_pyfunction!(batch_encode, m)?)?;
772 m.add_function(wrap_pyfunction!(batch_encode_numpy, m)?)?;
773 m.add_function(wrap_pyfunction!(py_wong_wang_simulate, m)?)?;
774 m.add_function(wrap_pyfunction!(py_dcls_max_forward_batch_q88, m)?)?;
775 m.add_function(wrap_pyfunction!(py_mixed_dense_forward_batch_q88_q1616, m)?)?;
776 m.add_function(wrap_pyfunction!(py_adc_to_spike_windows, m)?)?;
777 m.add_function(wrap_pyfunction!(py_sc_forward_packed, m)?)?;
778 m.add_function(wrap_pyfunction!(py_wilson_cowan_simulate, m)?)?;
779 m.add_class::<Lfsr16>()?;
780 m.add_class::<BitstreamEncoder>()?;
781 m.add_class::<FixedPointLif>()?;
782 m.add_class::<DenseLayer>()?;
783 #[cfg(feature = "gpu")]
784 m.add_class::<gpu::PyGpuDenseLayer>()?;
785 m.add_class::<StdpSynapse>()?;
786 m.add_class::<PySurrogateLif>()?;
787 m.add_class::<PyDifferentiableDenseLayer>()?;
788 m.add_class::<PyStochasticAttention>()?;
789 m.add_class::<PyStochasticGraphLayer>()?;
790 m.add_class::<PyKuramotoSolver>()?;
791 m.add_class::<PySCPNMetrics>()?;
792 m.add_class::<PyBitStreamTensor>()?;
793 m.add_class::<PyBrunelNetwork>()?;
794 m.add_class::<PyIzhikevich>()?;
795 m.add_class::<PyBitstreamAverager>()?;
796 m.add_class::<PyScGraph>()?;
797 m.add_class::<PyScGraphBuilder>()?;
798 m.add_function(wrap_pyfunction!(ir_verify, m)?)?;
799 m.add_function(wrap_pyfunction!(ir_print, m)?)?;
800 m.add_function(wrap_pyfunction!(ir_parse, m)?)?;
801 m.add_function(wrap_pyfunction!(ir_emit_sv, m)?)?;
802 m.add_class::<PyAdExNeuron>()?;
803 m.add_class::<PyExpIFNeuron>()?;
804 m.add_class::<PyLapicqueNeuron>()?;
805 pyo3_neurons::register_neuron_classes(m)?;
806 m.add_class::<PyNetworkRunner>()?;
807 m.add_class::<supervisor::PySpikingControllerPool>()?;
808 m.add_function(wrap_pyfunction!(py_simulate_ei_network, m)?)?;
809 m.add_function(wrap_pyfunction!(py_batch_simulate, m)?)?;
810 m.add_function(wrap_pyfunction!(rk4_neurons::py_rk4_neuron_simulate, m)?)?;
811 m.add_function(wrap_pyfunction!(py_cordiv, m)?)?;
812 m.add_function(wrap_pyfunction!(py_adaptive_length, m)?)?;
813 m.add_function(wrap_pyfunction!(py_prediction_error, m)?)?;
814 m.add_function(wrap_pyfunction!(py_predict_xor_ema, m)?)?;
815 m.add_function(wrap_pyfunction!(py_predict_xor_lfsr, m)?)?;
816 m.add_function(wrap_pyfunction!(py_recover_xor_ema, m)?)?;
817 m.add_function(wrap_pyfunction!(py_recover_xor_lfsr, m)?)?;
818 m.add_function(wrap_pyfunction!(py_phi_star, m)?)?;
819 m.add_class::<PyCorticalColumn>()?;
820 m.add_class::<PyRallDendrite>()?;
821 m.add_function(wrap_pyfunction!(py_spike_times, m)?)?;
823 m.add_function(wrap_pyfunction!(py_isi, m)?)?;
824 m.add_function(wrap_pyfunction!(py_firing_rate, m)?)?;
825 m.add_function(wrap_pyfunction!(py_spike_count, m)?)?;
826 m.add_function(wrap_pyfunction!(py_bin_spike_train, m)?)?;
827 m.add_function(wrap_pyfunction!(py_instantaneous_rate, m)?)?;
828 m.add_function(wrap_pyfunction!(py_psth, m)?)?;
829 m.add_function(wrap_pyfunction!(py_cv_isi, m)?)?;
830 m.add_function(wrap_pyfunction!(py_cv2, m)?)?;
831 m.add_function(wrap_pyfunction!(py_local_variation, m)?)?;
832 m.add_function(wrap_pyfunction!(py_fano_factor, m)?)?;
833 m.add_function(wrap_pyfunction!(py_lempel_ziv_complexity, m)?)?;
834 m.add_function(wrap_pyfunction!(py_permutation_entropy, m)?)?;
835 m.add_function(wrap_pyfunction!(py_hurst_exponent, m)?)?;
836 m.add_function(wrap_pyfunction!(py_approximate_entropy, m)?)?;
837 m.add_function(wrap_pyfunction!(py_sample_entropy, m)?)?;
838 m.add_function(wrap_pyfunction!(py_cross_correlation, m)?)?;
840 m.add_function(wrap_pyfunction!(py_pairwise_correlation, m)?)?;
841 m.add_function(wrap_pyfunction!(py_event_synchronization, m)?)?;
842 m.add_function(wrap_pyfunction!(py_spike_train_coherence, m)?)?;
843 m.add_function(wrap_pyfunction!(py_spike_time_tiling_coefficient, m)?)?;
844 m.add_function(wrap_pyfunction!(py_covariance_matrix, m)?)?;
845 m.add_function(wrap_pyfunction!(py_autocorrelation_time, m)?)?;
846 m.add_function(wrap_pyfunction!(py_noise_correlation, m)?)?;
847 m.add_function(wrap_pyfunction!(py_signal_correlation, m)?)?;
848 m.add_function(wrap_pyfunction!(py_spike_count_covariance, m)?)?;
849 m.add_function(wrap_pyfunction!(py_joint_psth, m)?)?;
850 m.add_function(wrap_pyfunction!(py_coincidence_index, m)?)?;
851 m.add_function(wrap_pyfunction!(py_van_rossum_distance, m)?)?;
853 m.add_function(wrap_pyfunction!(py_victor_purpura_distance, m)?)?;
854 m.add_function(wrap_pyfunction!(py_isi_distance, m)?)?;
855 m.add_function(wrap_pyfunction!(py_spike_distance, m)?)?;
856 m.add_function(wrap_pyfunction!(py_spike_sync, m)?)?;
857 m.add_function(wrap_pyfunction!(py_spike_sync_profile, m)?)?;
858 m.add_function(wrap_pyfunction!(py_spike_profile, m)?)?;
859 m.add_function(wrap_pyfunction!(py_isi_profile, m)?)?;
860 m.add_function(wrap_pyfunction!(py_adaptive_spike_distance, m)?)?;
861 m.add_function(wrap_pyfunction!(py_schreiber_similarity, m)?)?;
862 m.add_function(wrap_pyfunction!(py_hunter_milton_similarity, m)?)?;
863 m.add_function(wrap_pyfunction!(py_earth_movers_distance, m)?)?;
864 m.add_function(wrap_pyfunction!(py_multi_neuron_victor_purpura, m)?)?;
865 m.add_function(wrap_pyfunction!(py_spike_distance_matrix, m)?)?;
866 m.add_function(wrap_pyfunction!(py_mutual_information, m)?)?;
868 m.add_function(wrap_pyfunction!(py_transfer_entropy, m)?)?;
869 m.add_function(wrap_pyfunction!(py_spike_train_entropy, m)?)?;
870 m.add_function(wrap_pyfunction!(py_noise_entropy, m)?)?;
871 m.add_function(wrap_pyfunction!(py_stimulus_specific_information, m)?)?;
872 m.add_function(wrap_pyfunction!(py_kozachenko_leonenko_mi, m)?)?;
873 m.add_function(wrap_pyfunction!(py_pairwise_granger_causality, m)?)?;
875 m.add_function(wrap_pyfunction!(py_conditional_granger_causality, m)?)?;
876 m.add_function(wrap_pyfunction!(py_spectral_granger_causality, m)?)?;
877 m.add_function(wrap_pyfunction!(py_partial_directed_coherence, m)?)?;
878 m.add_function(wrap_pyfunction!(py_directed_transfer_function, m)?)?;
879 m.add_function(wrap_pyfunction!(py_population_vector_decode, m)?)?;
881 m.add_function(wrap_pyfunction!(py_bayesian_decode, m)?)?;
882 m.add_function(wrap_pyfunction!(py_maximum_likelihood_decode, m)?)?;
883 m.add_function(wrap_pyfunction!(py_linear_discriminant_decode, m)?)?;
884 m.add_function(wrap_pyfunction!(py_naive_bayes_decode, m)?)?;
885 m.add_function(wrap_pyfunction!(py_tokenise_spikes, m)?)?;
887 m.add_function(wrap_pyfunction!(py_sinusoidal_position_encode, m)?)?;
888 m.add_function(wrap_pyfunction!(py_scaled_dot_product_attention, m)?)?;
889 m.add_function(wrap_pyfunction!(py_gaussian_attention, m)?)?;
890 m.add_function(wrap_pyfunction!(py_infonce_loss, m)?)?;
891 m.add_function(wrap_pyfunction!(py_functional_connectivity, m)?)?;
893 m.add_function(wrap_pyfunction!(py_unitary_events, m)?)?;
894 m.add_function(wrap_pyfunction!(py_cell_assembly_detection, m)?)?;
895 m.add_function(wrap_pyfunction!(py_synfire_chain_detection, m)?)?;
896 m.add_function(wrap_pyfunction!(py_surrogate_isi_shuffle, m)?)?;
898 m.add_function(wrap_pyfunction!(py_surrogate_dither, m)?)?;
899 m.add_function(wrap_pyfunction!(py_homogeneous_poisson, m)?)?;
900 m.add_function(wrap_pyfunction!(py_gamma_process, m)?)?;
901 m.add_function(wrap_pyfunction!(py_compound_poisson_process, m)?)?;
902 m.add_function(wrap_pyfunction!(py_surrogate_joint_isi, m)?)?;
903 m.add_function(wrap_pyfunction!(py_surrogate_bin_shuffling, m)?)?;
904 m.add_function(wrap_pyfunction!(py_surrogate_spike_train_shifting, m)?)?;
905 m.add_function(wrap_pyfunction!(py_burst_detection, m)?)?;
907 m.add_function(wrap_pyfunction!(py_first_spike_latency, m)?)?;
908 m.add_function(wrap_pyfunction!(py_response_onset, m)?)?;
909 m.add_function(wrap_pyfunction!(py_change_point_detection, m)?)?;
910 m.add_function(wrap_pyfunction!(py_spike_directionality, m)?)?;
912 m.add_function(wrap_pyfunction!(py_spike_train_order, m)?)?;
913 m.add_function(wrap_pyfunction!(py_cubic_higher_order, m)?)?;
914 m.add_function(wrap_pyfunction!(py_power_spectrum, m)?)?;
916 m.add_function(wrap_pyfunction!(py_waveform_width, m)?)?;
918 m.add_function(wrap_pyfunction!(py_waveform_amplitude, m)?)?;
919 m.add_function(wrap_pyfunction!(py_waveform_repolarization_slope, m)?)?;
920 m.add_function(wrap_pyfunction!(py_waveform_recovery_slope, m)?)?;
921 m.add_function(wrap_pyfunction!(py_waveform_halfwidth, m)?)?;
922 m.add_function(wrap_pyfunction!(py_waveform_pt_ratio, m)?)?;
923 m.add_function(wrap_pyfunction!(py_conditional_intensity, m)?)?;
925 m.add_function(wrap_pyfunction!(py_isi_hazard_function, m)?)?;
926 m.add_function(wrap_pyfunction!(py_isi_survivor_function, m)?)?;
927 m.add_function(wrap_pyfunction!(py_renewal_density, m)?)?;
928 m.add_function(wrap_pyfunction!(py_spike_triggered_average, m)?)?;
930 m.add_function(wrap_pyfunction!(py_spike_triggered_covariance, m)?)?;
931 m.add_function(wrap_pyfunction!(py_spatial_information, m)?)?;
932 m.add_function(wrap_pyfunction!(py_place_field_detection, m)?)?;
933 m.add_function(wrap_pyfunction!(py_tuning_curve, m)?)?;
934 m.add_function(wrap_pyfunction!(py_phase_locking_value, m)?)?;
936 m.add_function(wrap_pyfunction!(py_spike_field_coherence, m)?)?;
937 m.add_function(wrap_pyfunction!(py_spike_phase_histogram, m)?)?;
938 m.add_function(wrap_pyfunction!(py_isolation_distance, m)?)?;
940 m.add_function(wrap_pyfunction!(py_l_ratio, m)?)?;
941 m.add_function(wrap_pyfunction!(py_silhouette_score, m)?)?;
942 m.add_function(wrap_pyfunction!(py_d_prime, m)?)?;
943 m.add_function(wrap_pyfunction!(py_isi_violation_rate, m)?)?;
944 m.add_function(wrap_pyfunction!(py_presence_ratio, m)?)?;
945 m.add_function(wrap_pyfunction!(py_amplitude_cutoff, m)?)?;
946 m.add_function(wrap_pyfunction!(py_snr, m)?)?;
947 m.add_function(wrap_pyfunction!(py_nn_hit_rate, m)?)?;
948 m.add_function(wrap_pyfunction!(py_drift_metric, m)?)?;
949 m.add_function(wrap_pyfunction!(py_spike_train_pca, m)?)?;
951 m.add_function(wrap_pyfunction!(py_demixed_pca, m)?)?;
952 m.add_function(wrap_pyfunction!(py_factor_analysis, m)?)?;
953 m.add_function(wrap_pyfunction!(py_pca_components, m)?)?;
954 m.add_function(wrap_pyfunction!(py_demixed_components, m)?)?;
955 m.add_function(wrap_pyfunction!(py_factor_loadings, m)?)?;
956 m.add_function(wrap_pyfunction!(py_gpfa, m)?)?;
958 m.add_function(wrap_pyfunction!(py_gpfa_em, m)?)?;
959 m.add_function(wrap_pyfunction!(py_gpfa_transform, m)?)?;
960 m.add_function(wrap_pyfunction!(py_spade_detect, m)?)?;
962 m.add_function(wrap_pyfunction!(py_dna_design_sequence, m)?)?;
964 m.add_function(wrap_pyfunction!(py_dna_design_orthogonal_set, m)?)?;
965 m.add_function(wrap_pyfunction!(py_dna_check_cross_hybridization, m)?)?;
966 m.add_function(wrap_pyfunction!(py_dna_simulate_kinetics, m)?)?;
967 m.add_function(wrap_pyfunction!(py_dna_detect_hairpins, m)?)?;
968 m.add_function(wrap_pyfunction!(py_qa_ising_energy, m)?)?;
970 m.add_function(wrap_pyfunction!(py_qa_batch_ising_energy, m)?)?;
971 m.add_function(wrap_pyfunction!(py_qa_simulated_annealing, m)?)?;
972 m.add_function(wrap_pyfunction!(py_qa_gauge_transform, m)?)?;
973 m.add_function(wrap_pyfunction!(py_qa_generate_gauges, m)?)?;
974 m.add_function(wrap_pyfunction!(py_qa_greedy_partition, m)?)?;
975 m.add_function(wrap_pyfunction!(py_ph_route_waveguides, m)?)?;
977 m.add_function(wrap_pyfunction!(py_ph_mzi_transfer_matrix, m)?)?;
978 m.add_function(wrap_pyfunction!(py_ph_cascade_mzi, m)?)?;
979 m.add_function(wrap_pyfunction!(py_ph_analyze_crosstalk, m)?)?;
980 m.add_function(wrap_pyfunction!(py_ph_analyze_power_budget, m)?)?;
981 m.add_function(wrap_pyfunction!(py_ph_analyze_crosstalk_bank, m)?)?;
982 m.add_function(wrap_pyfunction!(py_ph_analyze_crosstalk_pairs, m)?)?;
983 m.add_function(wrap_pyfunction!(py_opt_sa_search, m)?)?;
985 m.add_function(wrap_pyfunction!(py_opt_extract_pareto, m)?)?;
986 m.add_function(wrap_pyfunction!(py_evo_batch_mutate, m)?)?;
988 m.add_function(wrap_pyfunction!(py_evo_batch_fitness, m)?)?;
989 m.add_function(wrap_pyfunction!(py_evo_batch_crossover, m)?)?;
990 m.add_function(wrap_pyfunction!(py_evo_diversity, m)?)?;
991 m.add_function(wrap_pyfunction!(py_evo_novelty, m)?)?;
992 m.add_function(wrap_pyfunction!(py_evo_tournament, m)?)?;
993 m.add_function(wrap_pyfunction!(py_lgssm_kalman_filter, m)?)?;
995 m.add_function(wrap_pyfunction!(py_ollivier_ricci_curvature, m)?)?;
996 m.add_function(wrap_pyfunction!(py_cazelles_map_simulate, m)?)?;
997 m.add_function(wrap_pyfunction!(py_courage_nekorkin_map_simulate, m)?)?;
998 m.add_function(wrap_pyfunction!(py_mckean_simulate, m)?)?;
999 m.add_function(wrap_pyfunction!(py_wilson_hr_simulate, m)?)?;
1000 m.add_function(wrap_pyfunction!(py_pernarowski_simulate, m)?)?;
1001 m.add_function(wrap_pyfunction!(py_terman_wang_simulate, m)?)?;
1002 m.add_function(wrap_pyfunction!(py_mihalas_niebur_simulate, m)?)?;
1003 m.add_function(wrap_pyfunction!(py_glif_simulate, m)?)?;
1004 m.add_function(wrap_pyfunction!(py_rulkov_map_simulate, m)?)?;
1005 m.add_function(wrap_pyfunction!(py_ibarz_tanaka_map_simulate, m)?)?;
1006 m.add_function(wrap_pyfunction!(py_medvedev_map_simulate, m)?)?;
1007 m.add_function(wrap_pyfunction!(py_ermentrout_kopell_map_simulate, m)?)?;
1008 m.add_function(wrap_pyfunction!(py_fitzhugh_nagumo_simulate, m)?)?;
1009 m.add_function(wrap_pyfunction!(py_hindmarsh_rose_simulate, m)?)?;
1010 m.add_function(wrap_pyfunction!(py_fitzhugh_rinzel_simulate, m)?)?;
1011 m.add_function(wrap_pyfunction!(py_izhikevich2007_simulate, m)?)?;
1012 m.add_function(wrap_pyfunction!(py_inject_bitflip_u8, m)?)?;
1014 m.add_function(wrap_pyfunction!(py_inject_stuck_at_0_u8, m)?)?;
1015 m.add_function(wrap_pyfunction!(py_inject_stuck_at_1_u8, m)?)?;
1016 m.add_function(wrap_pyfunction!(py_inject_dropout_u8, m)?)?;
1017 m.add_function(wrap_pyfunction!(py_inject_gaussian_u8, m)?)?;
1018 m.add_function(wrap_pyfunction!(py_kl_refine, m)?)?;
1020 m.add_function(wrap_pyfunction!(py_ping_step, m)?)?;
1022 m.add_function(wrap_pyfunction!(py_parallel_csr_spmv_add, m)?)?;
1024 m.add_function(wrap_pyfunction!(py_parallel_csr_multi_spmv_add, m)?)?;
1025 Ok(())
1026}
1027
1028#[pyfunction]
1031fn py_cordiv(
1032 py: Python<'_>,
1033 numerator: PyReadonlyArray1<'_, u8>,
1034 denominator: PyReadonlyArray1<'_, u8>,
1035) -> PyResult<Py<PyArray1<u8>>> {
1036 let num = numerator.as_slice()?;
1037 let den = denominator.as_slice()?;
1038 let result = cordiv::cordiv(num, den);
1039 Ok(result.into_pyarray(py).into())
1040}
1041
1042#[pyfunction]
1043fn py_adaptive_length(epsilon: f64, confidence: f64) -> usize {
1044 cordiv::adaptive_length_hoeffding(epsilon, confidence)
1045}
1046
1047#[pyfunction]
1050fn py_prediction_error(
1051 _py: Python<'_>,
1052 predicted: PyReadonlyArray1<'_, u64>,
1053 actual: PyReadonlyArray1<'_, u64>,
1054 length: usize,
1055) -> PyResult<f64> {
1056 let pred = predicted
1057 .as_slice()
1058 .map_err(|e| PyValueError::new_err(format!("predicted array must be contiguous: {e}")))?;
1059 let act = actual
1060 .as_slice()
1061 .map_err(|e| PyValueError::new_err(format!("actual array must be contiguous: {e}")))?;
1062 Ok(predictive_coding::prediction_error_packed(
1063 pred, act, length,
1064 ))
1065}
1066
1067#[pyfunction]
1070fn py_predict_xor_ema(
1071 _py: Python<'_>,
1072 spikes: PyReadonlyArray1<'_, i8>,
1073 n_channels: usize,
1074 alpha: f64,
1075 threshold: f64,
1076) -> PyResult<(Py<PyArray1<i8>>, usize)> {
1077 let data = spikes
1078 .as_slice()
1079 .map_err(|e| PyValueError::new_err(format!("spikes must be contiguous: {e}")))?;
1080 let (errors, correct) =
1081 predictive_coding::predict_and_xor_ema(data, n_channels, alpha, threshold);
1082 Ok((PyArray1::from_vec(_py, errors).into(), correct))
1083}
1084
1085#[pyfunction]
1086fn py_recover_xor_ema(
1087 _py: Python<'_>,
1088 errors: PyReadonlyArray1<'_, i8>,
1089 n_channels: usize,
1090 alpha: f64,
1091 threshold: f64,
1092) -> PyResult<Py<PyArray1<i8>>> {
1093 let data = errors
1094 .as_slice()
1095 .map_err(|e| PyValueError::new_err(format!("errors must be contiguous: {e}")))?;
1096 let spikes = predictive_coding::xor_and_recover_ema(data, n_channels, alpha, threshold);
1097 Ok(PyArray1::from_vec(_py, spikes).into())
1098}
1099
1100#[pyfunction]
1101fn py_predict_xor_lfsr(
1102 _py: Python<'_>,
1103 spikes: PyReadonlyArray1<'_, i8>,
1104 n_channels: usize,
1105 alpha_q8: i32,
1106 seed: u16,
1107) -> PyResult<(Py<PyArray1<i8>>, usize)> {
1108 let data = spikes
1109 .as_slice()
1110 .map_err(|e| PyValueError::new_err(format!("spikes must be contiguous: {e}")))?;
1111 let (errors, correct) =
1112 predictive_coding::predict_and_xor_lfsr(data, n_channels, alpha_q8, seed);
1113 Ok((PyArray1::from_vec(_py, errors).into(), correct))
1114}
1115
1116#[pyfunction]
1117fn py_recover_xor_lfsr(
1118 _py: Python<'_>,
1119 errors: PyReadonlyArray1<'_, i8>,
1120 n_channels: usize,
1121 alpha_q8: i32,
1122 seed: u16,
1123) -> PyResult<Py<PyArray1<i8>>> {
1124 let data = errors
1125 .as_slice()
1126 .map_err(|e| PyValueError::new_err(format!("errors must be contiguous: {e}")))?;
1127 let spikes = predictive_coding::xor_and_recover_lfsr(data, n_channels, alpha_q8, seed);
1128 Ok(PyArray1::from_vec(_py, spikes).into())
1129}
1130
1131#[pyfunction]
1134fn py_phi_star(_py: Python<'_>, data: PyReadonlyArray2<'_, f64>, tau: usize) -> PyResult<f64> {
1135 if !data.is_c_contiguous() {
1136 return Err(PyValueError::new_err(
1137 "py_phi_star requires C-contiguous array input",
1138 ));
1139 }
1140 let shape = data.shape();
1141 let n_channels = shape[0];
1142 let n_timesteps = shape[1];
1143 let flat = data.as_slice().map_err(|e| {
1144 PyValueError::new_err(format!("py_phi_star requires C-contiguous array: {e}"))
1145 })?;
1146 let channels: Vec<Vec<f64>> = (0..n_channels)
1147 .map(|i| flat[i * n_timesteps..(i + 1) * n_timesteps].to_vec())
1148 .collect();
1149 Ok(phi::phi_star(&channels, tau))
1150}
1151
1152#[pyclass(
1155 name = "CorticalColumnRust",
1156 module = "sc_neurocore_engine.sc_neurocore_engine"
1157)]
1158pub struct PyCorticalColumn {
1159 inner: cortical_column::CorticalColumnRust,
1160}
1161
1162#[pymethods]
1163impl PyCorticalColumn {
1164 #[new]
1165 fn new(n: usize, tau: f64, dt: f64, threshold: f64, w_exc: f64, w_inh: f64, seed: u64) -> Self {
1166 Self {
1167 inner: cortical_column::CorticalColumnRust::new(
1168 n, tau, dt, threshold, w_exc, w_inh, seed,
1169 ),
1170 }
1171 }
1172
1173 fn step<'py>(
1174 &mut self,
1175 py: Python<'py>,
1176 thalamic_input: PyReadonlyArray1<'py, f64>,
1177 ) -> PyResult<Py<PyDict>> {
1178 let input = thalamic_input.as_slice()?;
1179 let spikes = self.inner.step(input);
1180 let dict = PyDict::new(py);
1181 let names = ["l4", "l23_exc", "l23_inh", "l5", "l6"];
1182 for (i, name) in names.iter().enumerate() {
1183 dict.set_item(*name, spikes[i].clone().into_pyarray(py))?;
1184 }
1185 Ok(dict.into())
1186 }
1187
1188 fn reset(&mut self) {
1189 self.inner.reset();
1190 }
1191}
1192
1193#[pyclass(
1196 name = "RallDendriteRust",
1197 module = "sc_neurocore_engine.sc_neurocore_engine"
1198)]
1199pub struct PyRallDendrite {
1200 inner: rall_dendrite::RallDendriteRust,
1201}
1202
1203#[pymethods]
1204impl PyRallDendrite {
1205 #[new]
1206 fn new(n_branches: usize, branch_length: usize, tau: f64, coupling: f64, dt: f64) -> Self {
1207 Self {
1208 inner: rall_dendrite::RallDendriteRust::new(
1209 n_branches,
1210 branch_length,
1211 tau,
1212 coupling,
1213 dt,
1214 ),
1215 }
1216 }
1217
1218 fn step(&mut self, branch_inputs: PyReadonlyArray1<'_, f64>) -> PyResult<f64> {
1219 let inputs = branch_inputs.as_slice()?;
1220 Ok(self.inner.step(inputs))
1221 }
1222
1223 fn reset(&mut self) {
1224 self.inner.reset();
1225 }
1226
1227 #[getter]
1228 fn soma_v(&self) -> f64 {
1229 self.inner.soma_v
1230 }
1231}
1232
1233#[pyfunction]
1235fn simd_tier() -> &'static str {
1236 #[cfg(target_arch = "x86_64")]
1237 {
1238 if is_x86_feature_detected!("avx512vpopcntdq") {
1239 return "avx512-vpopcntdq";
1240 }
1241 if is_x86_feature_detected!("avx512bw") {
1242 return "avx512bw";
1243 }
1244 if is_x86_feature_detected!("avx512f") {
1245 return "avx512f";
1246 }
1247 if is_x86_feature_detected!("avx2") {
1248 return "avx2";
1249 }
1250 if is_x86_feature_detected!("popcnt") {
1251 return "popcnt";
1252 }
1253 }
1254 #[cfg(target_arch = "aarch64")]
1255 {
1256 return "neon";
1257 }
1258 "portable"
1259}
1260
1261#[pyfunction]
1266fn set_num_threads(n: usize) -> PyResult<()> {
1267 if n == 0 {
1268 return Ok(());
1269 }
1270 rayon::ThreadPoolBuilder::new()
1271 .num_threads(n)
1272 .build_global()
1273 .map_err(|e| PyValueError::new_err(format!("Cannot set thread pool: {e}")))
1274}
1275
1276#[pyfunction]
1277fn pack_bitstream(py: Python<'_>, bits: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
1278 if let Ok(rows) = bits.extract::<Vec<Vec<u8>>>() {
1279 let packed_rows: Vec<Vec<u64>> = rows.iter().map(|row| bitstream::pack(row).data).collect();
1280 return Ok(packed_rows
1281 .into_pyobject(py)
1282 .map_err(|e| PyValueError::new_err(e.to_string()))?
1283 .into_any()
1284 .unbind());
1285 }
1286
1287 let flat = bits
1288 .extract::<Vec<u8>>()
1289 .map_err(|_| PyValueError::new_err("Expected a 1-D or 2-D array of uint8 bits."))?;
1290 Ok(bitstream::pack(&flat)
1291 .data
1292 .into_pyobject(py)
1293 .map_err(|e| PyValueError::new_err(e.to_string()))?
1294 .into_any()
1295 .unbind())
1296}
1297
1298#[pyfunction]
1299#[pyo3(signature = (packed, original_length, original_shape=None))]
1300fn unpack_bitstream(
1301 py: Python<'_>,
1302 packed: &Bound<'_, PyAny>,
1303 original_length: usize,
1304 original_shape: Option<(usize, usize)>,
1305) -> PyResult<Py<PyAny>> {
1306 if let Ok(rows) = packed.extract::<Vec<Vec<u64>>>() {
1307 let batch = rows.len();
1308 let per_batch_len = if let Some((expected_batch, length)) = original_shape {
1309 if expected_batch != batch {
1310 return Err(PyValueError::new_err(format!(
1311 "original_shape batch {} does not match packed batch {}.",
1312 expected_batch, batch
1313 )));
1314 }
1315 length
1316 } else {
1317 original_length.checked_div(batch).unwrap_or(0)
1318 };
1319
1320 let unpacked_rows: Vec<Vec<u8>> = rows
1321 .into_iter()
1322 .map(|row| {
1323 bitstream::unpack(&bitstream::BitStreamTensor::from_words(row, per_batch_len))
1324 })
1325 .collect();
1326 return Ok(unpacked_rows
1327 .into_pyobject(py)
1328 .map_err(|e| PyValueError::new_err(e.to_string()))?
1329 .into_any()
1330 .unbind());
1331 }
1332
1333 let words = packed.extract::<Vec<u64>>().map_err(|_| {
1334 PyValueError::new_err("Expected packed uint64 words as 1-D or 2-D sequence.")
1335 })?;
1336 let tensor = bitstream::BitStreamTensor::from_words(words, original_length);
1337 Ok(bitstream::unpack(&tensor)
1338 .into_pyobject(py)
1339 .map_err(|e| PyValueError::new_err(e.to_string()))?
1340 .into_any()
1341 .unbind())
1342}
1343
1344#[pyfunction]
1345fn popcount(packed: &Bound<'_, PyAny>) -> PyResult<u64> {
1346 if let Ok(rows) = packed.extract::<Vec<Vec<u64>>>() {
1347 return Ok(rows
1348 .iter()
1349 .map(|row| simd::popcount_dispatch(row))
1350 .sum::<u64>());
1351 }
1352
1353 let words = packed.extract::<Vec<u64>>().map_err(|_| {
1354 PyValueError::new_err("Expected packed uint64 words as 1-D or 2-D sequence.")
1355 })?;
1356 Ok(simd::popcount_dispatch(&words))
1357}
1358
1359#[pyfunction]
1362fn pack_bitstream_numpy<'py>(
1363 py: Python<'py>,
1364 bits: PyReadonlyArray1<'py, u8>,
1365) -> PyResult<Bound<'py, PyArray1<u64>>> {
1366 let slice = bits
1367 .as_slice()
1368 .map_err(|e| PyValueError::new_err(format!("Cannot read numpy array: {e}")))?;
1369 let tensor = simd::pack_dispatch(slice);
1370 Ok(tensor.data.into_pyarray(py))
1371}
1372
1373#[pyfunction]
1375fn popcount_numpy(packed: PyReadonlyArray1<'_, u64>) -> PyResult<u64> {
1376 let words = packed
1377 .as_slice()
1378 .map_err(|e| PyValueError::new_err(format!("Cannot read numpy array: {e}")))?;
1379 Ok(simd::popcount_dispatch(words))
1380}
1381
1382#[pyfunction]
1384fn unpack_bitstream_numpy<'py>(
1385 py: Python<'py>,
1386 packed: PyReadonlyArray1<'py, u64>,
1387 original_length: usize,
1388) -> PyResult<Bound<'py, PyArray1<u8>>> {
1389 let words = packed
1390 .as_slice()
1391 .map_err(|e| PyValueError::new_err(format!("Cannot read numpy array: {e}")))?;
1392 let tensor = bitstream::BitStreamTensor::from_words(words.to_vec(), original_length);
1393 let bits = bitstream::unpack(&tensor);
1394 Ok(bits.into_pyarray(py))
1395}
1396
1397#[pyfunction]
1401#[pyo3(signature = (
1402 n_steps,
1403 leak_k,
1404 gain_k,
1405 i_t,
1406 noise_in=0,
1407 data_width=16,
1408 fraction=8,
1409 v_rest=0,
1410 v_reset=0,
1411 v_threshold=256,
1412 refractory_period=2
1413))]
1414#[allow(clippy::too_many_arguments)]
1415fn batch_lif_run<'py>(
1416 py: Python<'py>,
1417 n_steps: usize,
1418 leak_k: i16,
1419 gain_k: i16,
1420 i_t: i16,
1421 noise_in: i16,
1422 data_width: u32,
1423 fraction: u32,
1424 v_rest: i16,
1425 v_reset: i16,
1426 v_threshold: i16,
1427 refractory_period: i32,
1428) -> (Bound<'py, PyArray1<i32>>, Bound<'py, PyArray1<i16>>) {
1429 let mut lif = neuron::FixedPointLif::new(
1430 data_width,
1431 fraction,
1432 v_rest,
1433 v_reset,
1434 v_threshold,
1435 refractory_period,
1436 );
1437 let spikes_arr = PyArray1::<i32>::zeros(py, n_steps, false);
1438 let voltages_arr = PyArray1::<i16>::zeros(py, n_steps, false);
1439
1440 let spikes_slice = unsafe {
1442 spikes_arr
1443 .as_slice_mut()
1444 .expect("newly allocated spikes array must be contiguous")
1445 };
1446 let voltages_slice = unsafe {
1448 voltages_arr
1449 .as_slice_mut()
1450 .expect("newly allocated voltages array must be contiguous")
1451 };
1452
1453 for i in 0..n_steps {
1454 let (s, v) = lif.step(leak_k, gain_k, i_t, noise_in);
1455 spikes_slice[i] = s;
1456 voltages_slice[i] = v;
1457 }
1458
1459 (spikes_arr, voltages_arr)
1460}
1461
1462#[pyfunction]
1467#[pyo3(signature = (
1468 n_neurons,
1469 n_steps,
1470 leak_k,
1471 gain_k,
1472 currents,
1473 data_width=16,
1474 fraction=8,
1475 v_rest=0,
1476 v_reset=0,
1477 v_threshold=256,
1478 refractory_period=2
1479))]
1480#[allow(clippy::too_many_arguments)]
1481#[allow(clippy::type_complexity)]
1482fn batch_lif_run_multi<'py>(
1483 py: Python<'py>,
1484 n_neurons: usize,
1485 n_steps: usize,
1486 leak_k: i16,
1487 gain_k: i16,
1488 currents: PyReadonlyArray1<'py, i16>,
1489 data_width: u32,
1490 fraction: u32,
1491 v_rest: i16,
1492 v_reset: i16,
1493 v_threshold: i16,
1494 refractory_period: i32,
1495) -> PyResult<(Bound<'py, PyArray2<i32>>, Bound<'py, PyArray2<i16>>)> {
1496 use rayon::prelude::*;
1497
1498 let curr_slice = currents
1499 .as_slice()
1500 .map_err(|e| PyValueError::new_err(format!("Cannot read currents: {e}")))?;
1501 if curr_slice.len() != n_neurons {
1502 return Err(PyValueError::new_err(format!(
1503 "currents length {} does not match n_neurons {}.",
1504 curr_slice.len(),
1505 n_neurons
1506 )));
1507 }
1508
1509 let spikes_arr = PyArray2::<i32>::zeros(py, [n_neurons, n_steps], false);
1510 let voltages_arr = PyArray2::<i16>::zeros(py, [n_neurons, n_steps], false);
1511
1512 if n_neurons == 0 || n_steps == 0 {
1513 return Ok((spikes_arr, voltages_arr));
1514 }
1515
1516 let spikes_flat = unsafe {
1518 spikes_arr
1519 .as_slice_mut()
1520 .expect("newly allocated spikes array must be contiguous")
1521 };
1522 let voltages_flat = unsafe {
1524 voltages_arr
1525 .as_slice_mut()
1526 .expect("newly allocated voltages array must be contiguous")
1527 };
1528
1529 spikes_flat
1530 .par_chunks_mut(n_steps)
1531 .zip(voltages_flat.par_chunks_mut(n_steps))
1532 .zip(curr_slice.par_iter().copied())
1533 .for_each(|((spike_row, voltage_row), i_t)| {
1534 let mut lif = neuron::FixedPointLif::new(
1535 data_width,
1536 fraction,
1537 v_rest,
1538 v_reset,
1539 v_threshold,
1540 refractory_period,
1541 );
1542 for step in 0..n_steps {
1543 let (s, v) = lif.step(leak_k, gain_k, i_t, 0);
1544 spike_row[step] = s;
1545 voltage_row[step] = v;
1546 }
1547 });
1548
1549 Ok((spikes_arr, voltages_arr))
1550}
1551
1552#[pyfunction]
1554#[pyo3(signature = (
1555 leak_k,
1556 gain_k,
1557 currents,
1558 noises=None,
1559 data_width=16,
1560 fraction=8,
1561 v_rest=0,
1562 v_reset=0,
1563 v_threshold=256,
1564 refractory_period=2
1565))]
1566#[allow(clippy::too_many_arguments)]
1567#[allow(clippy::type_complexity)]
1568fn batch_lif_run_varying<'py>(
1569 py: Python<'py>,
1570 leak_k: i16,
1571 gain_k: i16,
1572 currents: PyReadonlyArray1<'py, i16>,
1573 noises: Option<PyReadonlyArray1<'py, i16>>,
1574 data_width: u32,
1575 fraction: u32,
1576 v_rest: i16,
1577 v_reset: i16,
1578 v_threshold: i16,
1579 refractory_period: i32,
1580) -> PyResult<(Bound<'py, PyArray1<i32>>, Bound<'py, PyArray1<i16>>)> {
1581 let curr_slice = currents
1582 .as_slice()
1583 .map_err(|e| PyValueError::new_err(format!("Cannot read currents: {e}")))?;
1584 let noise_slice: Option<&[i16]> = match noises.as_ref() {
1585 Some(n) => Some(
1586 n.as_slice()
1587 .map_err(|e| PyValueError::new_err(format!("Cannot read noises: {e}")))?,
1588 ),
1589 None => None,
1590 };
1591
1592 let n_steps = curr_slice.len();
1593 if let Some(ns) = noise_slice {
1594 if ns.len() != n_steps {
1595 return Err(PyValueError::new_err(format!(
1596 "noises length {} does not match currents length {}.",
1597 ns.len(),
1598 n_steps
1599 )));
1600 }
1601 }
1602
1603 let mut lif = neuron::FixedPointLif::new(
1604 data_width,
1605 fraction,
1606 v_rest,
1607 v_reset,
1608 v_threshold,
1609 refractory_period,
1610 );
1611 let spikes_arr = PyArray1::<i32>::zeros(py, n_steps, false);
1612 let voltages_arr = PyArray1::<i16>::zeros(py, n_steps, false);
1613
1614 let spikes_slice = unsafe {
1616 spikes_arr
1617 .as_slice_mut()
1618 .expect("newly allocated spikes array must be contiguous")
1619 };
1620 let voltages_slice = unsafe {
1622 voltages_arr
1623 .as_slice_mut()
1624 .expect("newly allocated voltages array must be contiguous")
1625 };
1626
1627 for i in 0..n_steps {
1628 let noise_in = noise_slice.map_or(0, |ns| ns[i]);
1629 let (s, v) = lif.step(leak_k, gain_k, curr_slice[i], noise_in);
1630 spikes_slice[i] = s;
1631 voltages_slice[i] = v;
1632 }
1633
1634 Ok((spikes_arr, voltages_arr))
1635}
1636
1637#[pyfunction]
1641#[pyo3(signature = (probs, length=1024, seed=0xACE1))]
1642fn batch_encode<'py>(
1643 _py: Python<'py>,
1644 probs: PyReadonlyArray1<'py, f64>,
1645 length: usize,
1646 seed: u64,
1647) -> PyResult<Vec<Vec<u64>>> {
1648 let prob_slice = probs
1649 .as_slice()
1650 .map_err(|e| PyValueError::new_err(format!("Cannot read probs: {e}")))?;
1651 let words = length.div_ceil(64);
1652
1653 use rand::SeedableRng;
1654 let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(seed);
1655
1656 let packed: Vec<Vec<u64>> = prob_slice
1657 .iter()
1658 .map(|&p| {
1659 let mut data = bitstream::bernoulli_packed(p, length, &mut rng);
1660 data.resize(words, 0);
1661 data
1662 })
1663 .collect();
1664
1665 Ok(packed)
1666}
1667
1668#[pyfunction]
1672#[pyo3(signature = (probs, length=1024, seed=0xACE1))]
1673fn batch_encode_numpy<'py>(
1674 py: Python<'py>,
1675 probs: PyReadonlyArray1<'py, f64>,
1676 length: usize,
1677 seed: u64,
1678) -> PyResult<Bound<'py, PyArray2<u64>>> {
1679 use rayon::prelude::*;
1680
1681 let prob_slice = probs
1682 .as_slice()
1683 .map_err(|e| PyValueError::new_err(format!("Cannot read probs: {e}")))?;
1684 let words = length.div_ceil(64);
1685 let n_probs = prob_slice.len();
1686
1687 let rows: Vec<Vec<u64>> = prob_slice
1688 .par_iter()
1689 .enumerate()
1690 .map(|(idx, &p)| {
1691 use rand::SeedableRng;
1692
1693 let prob_seed = seed.wrapping_add(idx as u64);
1694 let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(prob_seed);
1695 let mut row = bitstream::bernoulli_packed_simd(p, length, &mut rng);
1696 row.resize(words, 0);
1697 row
1698 })
1699 .collect();
1700
1701 let mut flat = Vec::with_capacity(n_probs * words);
1702 for row in &rows {
1703 flat.extend_from_slice(row);
1704 }
1705
1706 let arr = ndarray::Array2::from_shape_vec((n_probs, words), flat)
1707 .map_err(|e| PyValueError::new_err(format!("Shape construction failed: {e}")))?;
1708 Ok(arr.into_pyarray(py))
1709}
1710
1711#[pyclass(module = "sc_neurocore_engine.sc_neurocore_engine")]
1712pub struct Lfsr16 {
1713 inner: encoder::Lfsr16,
1714 seed_init: u16,
1715}
1716
1717#[pymethods]
1718impl Lfsr16 {
1719 #[new]
1720 #[pyo3(signature = (seed=0xACE1))]
1721 fn new(seed: u16) -> PyResult<Self> {
1722 if seed == 0 {
1723 return Err(PyValueError::new_err("LFSR seed must be non-zero."));
1724 }
1725 Ok(Self {
1726 inner: encoder::Lfsr16::new(seed),
1727 seed_init: seed,
1728 })
1729 }
1730
1731 fn step(&mut self) -> u16 {
1732 self.inner.step()
1733 }
1734
1735 #[getter]
1736 fn reg(&self) -> u16 {
1737 self.inner.reg
1738 }
1739
1740 #[getter]
1741 fn width(&self) -> u32 {
1742 self.inner.width
1743 }
1744
1745 #[pyo3(signature = (seed=None))]
1746 fn reset(&mut self, seed: Option<u16>) -> PyResult<()> {
1747 let next = seed.unwrap_or(self.seed_init);
1748 if next == 0 {
1749 return Err(PyValueError::new_err("LFSR seed must be non-zero."));
1750 }
1751 self.inner = encoder::Lfsr16::new(next);
1752 self.seed_init = next;
1753 Ok(())
1754 }
1755}
1756
1757#[pyclass(module = "sc_neurocore_engine.sc_neurocore_engine")]
1758pub struct BitstreamEncoder {
1759 inner: encoder::BitstreamEncoder,
1760 seed_init: u16,
1761}
1762
1763#[pymethods]
1764impl BitstreamEncoder {
1765 #[new]
1766 #[pyo3(signature = (data_width=16, seed=0xACE1))]
1767 fn new(data_width: u32, seed: u16) -> PyResult<Self> {
1768 if seed == 0 {
1769 return Err(PyValueError::new_err("LFSR seed must be non-zero."));
1770 }
1771 Ok(Self {
1772 inner: encoder::BitstreamEncoder::new(data_width, seed),
1773 seed_init: seed,
1774 })
1775 }
1776
1777 fn step(&mut self, x_value: u16) -> u8 {
1778 self.inner.step(x_value)
1779 }
1780
1781 #[getter]
1782 fn data_width(&self) -> u32 {
1783 self.inner.data_width
1784 }
1785
1786 #[getter]
1787 fn reg(&self) -> u16 {
1788 self.inner.lfsr.reg
1789 }
1790
1791 #[pyo3(signature = (seed=None))]
1792 fn reset(&mut self, seed: Option<u16>) -> PyResult<()> {
1793 let next = seed.unwrap_or(self.seed_init);
1794 if next == 0 {
1795 return Err(PyValueError::new_err("LFSR seed must be non-zero."));
1796 }
1797 self.inner.reset(Some(next));
1798 self.seed_init = next;
1799 Ok(())
1800 }
1801}
1802
1803#[pyclass(module = "sc_neurocore_engine.sc_neurocore_engine")]
1804pub struct FixedPointLif {
1805 inner: neuron::FixedPointLif,
1806}
1807
1808#[pymethods]
1809impl FixedPointLif {
1810 #[new]
1811 #[pyo3(signature = (
1812 data_width=16,
1813 fraction=8,
1814 v_rest=0,
1815 v_reset=0,
1816 v_threshold=256,
1817 refractory_period=2
1818 ))]
1819 fn new(
1820 data_width: u32,
1821 fraction: u32,
1822 v_rest: i16,
1823 v_reset: i16,
1824 v_threshold: i16,
1825 refractory_period: i32,
1826 ) -> Self {
1827 Self {
1828 inner: neuron::FixedPointLif::new(
1829 data_width,
1830 fraction,
1831 v_rest,
1832 v_reset,
1833 v_threshold,
1834 refractory_period,
1835 ),
1836 }
1837 }
1838
1839 #[pyo3(signature = (leak_k, gain_k, i_t, noise_in=0))]
1840 fn step(&mut self, leak_k: i16, gain_k: i16, i_t: i16, noise_in: i16) -> (i32, i16) {
1841 self.inner.step(leak_k, gain_k, i_t, noise_in)
1842 }
1843
1844 fn reset(&mut self) {
1845 self.inner.reset();
1846 }
1847
1848 fn reset_state(&mut self) {
1849 self.reset();
1850 }
1851
1852 fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
1853 let dict = PyDict::new(py);
1854 dict.set_item("v", self.inner.v)?;
1855 dict.set_item("refractory_counter", self.inner.refractory_counter)?;
1856 Ok(dict.into_any().unbind())
1857 }
1858}
1859
1860#[pyclass(
1863 name = "Izhikevich",
1864 module = "sc_neurocore_engine.sc_neurocore_engine"
1865)]
1866pub struct PyIzhikevich {
1867 inner: neuron::Izhikevich,
1868}
1869
1870#[pymethods]
1871impl PyIzhikevich {
1872 #[new]
1873 #[pyo3(signature = (a=0.02, b=0.2, c=-65.0, d=8.0, dt=1.0))]
1874 fn new(a: f64, b: f64, c: f64, d: f64, dt: f64) -> Self {
1875 Self {
1876 inner: neuron::Izhikevich::new(a, b, c, d, dt),
1877 }
1878 }
1879
1880 fn step(&mut self, current: f64) -> i32 {
1881 self.inner.step(current)
1882 }
1883
1884 fn reset(&mut self) {
1885 self.inner.reset();
1886 }
1887
1888 fn reset_state(&mut self) {
1889 self.reset();
1890 }
1891
1892 fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
1893 let dict = PyDict::new(py);
1894 dict.set_item("v", self.inner.v)?;
1895 dict.set_item("u", self.inner.u)?;
1896 Ok(dict.into_any().unbind())
1897 }
1898}
1899
1900#[pyclass(
1903 name = "BitstreamAverager",
1904 module = "sc_neurocore_engine.sc_neurocore_engine"
1905)]
1906pub struct PyBitstreamAverager {
1907 inner: neuron::BitstreamAverager,
1908}
1909
1910#[pymethods]
1911impl PyBitstreamAverager {
1912 #[new]
1913 #[pyo3(signature = (window=1024))]
1914 fn new(window: usize) -> Self {
1915 Self {
1916 inner: neuron::BitstreamAverager::new(window),
1917 }
1918 }
1919
1920 fn push(&mut self, bit: u8) {
1921 self.inner.push(bit);
1922 }
1923
1924 fn estimate(&self) -> f64 {
1925 self.inner.estimate()
1926 }
1927
1928 fn reset(&mut self) {
1929 self.inner.reset();
1930 }
1931
1932 #[getter]
1933 fn window(&self) -> usize {
1934 self.inner.window()
1935 }
1936}
1937
1938#[pyclass(
1941 name = "AdExNeuron",
1942 module = "sc_neurocore_engine.sc_neurocore_engine"
1943)]
1944#[derive(Clone)]
1945pub struct PyAdExNeuron {
1946 inner: neuron::AdExNeuron,
1947}
1948
1949#[pymethods]
1950impl PyAdExNeuron {
1951 #[new]
1952 fn new() -> Self {
1953 Self {
1954 inner: neuron::AdExNeuron::new(),
1955 }
1956 }
1957 fn step(&mut self, current: f64) -> i32 {
1958 self.inner.step(current)
1959 }
1960 fn reset(&mut self) {
1961 self.inner.reset();
1962 }
1963 fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
1964 let d = PyDict::new(py);
1965 d.set_item("v", self.inner.v)?;
1966 d.set_item("w", self.inner.w)?;
1967 Ok(d.into_any().unbind())
1968 }
1969}
1970
1971#[pyclass(
1972 name = "ExpIFNeuron",
1973 module = "sc_neurocore_engine.sc_neurocore_engine"
1974)]
1975#[derive(Clone)]
1976pub struct PyExpIFNeuron {
1977 inner: neuron::ExpIfNeuron,
1978}
1979
1980#[pymethods]
1981impl PyExpIFNeuron {
1982 #[new]
1983 fn new() -> Self {
1984 Self {
1985 inner: neuron::ExpIfNeuron::new(),
1986 }
1987 }
1988 fn step(&mut self, current: f64) -> i32 {
1989 self.inner.step(current)
1990 }
1991 fn reset(&mut self) {
1992 self.inner.reset();
1993 }
1994 fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
1995 let d = PyDict::new(py);
1996 d.set_item("v", self.inner.v)?;
1997 Ok(d.into_any().unbind())
1998 }
1999}
2000
2001#[pyclass(
2002 name = "LapicqueNeuron",
2003 module = "sc_neurocore_engine.sc_neurocore_engine"
2004)]
2005#[derive(Clone)]
2006pub struct PyLapicqueNeuron {
2007 inner: neuron::LapicqueNeuron,
2008}
2009
2010#[pymethods]
2011impl PyLapicqueNeuron {
2012 #[new]
2013 #[pyo3(signature = (tau=20.0, resistance=1.0, threshold=1.0, dt=1.0))]
2014 fn new(tau: f64, resistance: f64, threshold: f64, dt: f64) -> Self {
2015 Self {
2016 inner: neuron::LapicqueNeuron::new(tau, resistance, threshold, dt),
2017 }
2018 }
2019 fn step(&mut self, current: f64) -> i32 {
2020 self.inner.step(current)
2021 }
2022 fn reset(&mut self) {
2023 self.inner.reset();
2024 }
2025 fn get_state(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
2026 let d = PyDict::new(py);
2027 d.set_item("v", self.inner.v)?;
2028 Ok(d.into_any().unbind())
2029 }
2030}
2031
2032#[pyclass(module = "sc_neurocore_engine.sc_neurocore_engine")]
2035pub struct DenseLayer {
2036 inner: layer::DenseLayer,
2037}
2038
2039#[pymethods]
2040impl DenseLayer {
2041 #[new]
2042 #[pyo3(signature = (n_inputs, n_neurons, length=1024, seed=24301))]
2043 fn new(n_inputs: usize, n_neurons: usize, length: usize, seed: u64) -> Self {
2044 Self {
2045 inner: layer::DenseLayer::new(n_inputs, n_neurons, length, seed),
2046 }
2047 }
2048
2049 fn get_weights(&self) -> Vec<Vec<f64>> {
2050 self.inner.get_weights()
2051 }
2052
2053 fn set_weights(&mut self, weights: Vec<Vec<f64>>) -> PyResult<()> {
2054 self.inner
2055 .set_weights(weights)
2056 .map_err(PyValueError::new_err)
2057 }
2058
2059 fn refresh_packed_weights(&mut self) {
2060 self.inner.refresh_packed_weights();
2061 }
2062
2063 #[pyo3(signature = (input_values, seed=44257))]
2064 fn forward(&self, input_values: Vec<f64>, seed: u64) -> PyResult<Vec<f64>> {
2065 self.inner
2066 .forward(&input_values, seed)
2067 .map_err(PyValueError::new_err)
2068 }
2069
2070 #[pyo3(signature = (input_values, seed=44257))]
2071 fn forward_fast(&self, input_values: Vec<f64>, seed: u64) -> PyResult<Vec<f64>> {
2072 self.inner
2073 .forward_fast(&input_values, seed)
2074 .map_err(PyValueError::new_err)
2075 }
2076
2077 #[pyo3(signature = (input_values, seed=44257))]
2081 fn forward_numpy<'py>(
2082 &self,
2083 py: Python<'py>,
2084 input_values: PyReadonlyArray1<'py, f64>,
2085 seed: u64,
2086 ) -> PyResult<Bound<'py, PyArray1<f64>>> {
2087 let slice = input_values
2088 .as_slice()
2089 .map_err(|e| PyValueError::new_err(format!("Cannot read input array: {e}")))?;
2090 let out = self
2091 .inner
2092 .forward_numpy_inner(slice, seed)
2093 .map_err(PyValueError::new_err)?;
2094 Ok(out.into_pyarray(py))
2095 }
2096
2097 #[pyo3(signature = (inputs, seed=44257))]
2102 fn forward_batch_numpy<'py>(
2103 &self,
2104 py: Python<'py>,
2105 inputs: PyReadonlyArray2<'py, f64>,
2106 seed: u64,
2107 ) -> PyResult<Bound<'py, PyArray2<f64>>> {
2108 let shape = inputs.shape();
2109 let n_samples = shape[0];
2110 let n_inputs = shape[1];
2111 if n_inputs != self.inner.n_inputs {
2112 return Err(PyValueError::new_err(format!(
2113 "Expected {} input features, got {}.",
2114 self.inner.n_inputs, n_inputs
2115 )));
2116 }
2117
2118 let flat_inputs = inputs
2119 .as_slice()
2120 .map_err(|e| PyValueError::new_err(format!("Array not contiguous: {e}")))?;
2121 let out = PyArray2::<f64>::zeros(py, [n_samples, self.inner.n_neurons], false);
2122 let out_slice = unsafe {
2124 out.as_slice_mut()
2125 .expect("newly allocated output array must be contiguous")
2126 };
2127
2128 self.inner
2129 .forward_batch_into(flat_inputs, n_samples, seed, out_slice)
2130 .map_err(PyValueError::new_err)?;
2131 Ok(out)
2132 }
2133
2134 fn forward_prepacked(&self, packed_inputs: &Bound<'_, PyAny>) -> PyResult<Vec<f64>> {
2140 if let Ok(arr) = packed_inputs.extract::<PyReadonlyArray2<u64>>() {
2141 let view = arr.as_array();
2142 let rows: Vec<Vec<u64>> = (0..view.nrows()).map(|i| view.row(i).to_vec()).collect();
2143 return self
2144 .inner
2145 .forward_prepacked(&rows)
2146 .map_err(PyValueError::new_err);
2147 }
2148
2149 let rows = packed_inputs.extract::<Vec<Vec<u64>>>().map_err(|_| {
2150 PyValueError::new_err(
2151 "packed_inputs must be a 2-D numpy uint64 array or list[list[int]].",
2152 )
2153 })?;
2154 self.inner
2155 .forward_prepacked(&rows)
2156 .map_err(PyValueError::new_err)
2157 }
2158
2159 #[pyo3(signature = (packed_inputs,))]
2164 fn forward_prepacked_numpy<'py>(
2165 &self,
2166 py: Python<'py>,
2167 packed_inputs: PyReadonlyArray2<'py, u64>,
2168 ) -> PyResult<Bound<'py, PyArray1<f64>>> {
2169 let shape = packed_inputs.shape();
2170 let n_inputs = shape[0];
2171 let words = shape[1];
2172 let flat = packed_inputs
2173 .as_slice()
2174 .map_err(|e| PyValueError::new_err(format!("Array not contiguous: {e}")))?;
2175 let out = self
2176 .inner
2177 .forward_prepacked_2d(flat, n_inputs, words)
2178 .map_err(PyValueError::new_err)?;
2179 Ok(out.into_pyarray(py))
2180 }
2181}
2182
2183fn parse_surrogate(name: &str, k: Option<f32>) -> PyResult<grad::SurrogateType> {
2184 let normalized = name.to_ascii_lowercase().replace('-', "_");
2185 match normalized.as_str() {
2186 "fast_sigmoid" => Ok(grad::SurrogateType::FastSigmoid {
2187 k: k.unwrap_or(25.0),
2188 }),
2189 "superspike" | "super_spike" => Ok(grad::SurrogateType::SuperSpike {
2190 k: k.unwrap_or(100.0),
2191 }),
2192 "arctan" | "arc_tan" => Ok(grad::SurrogateType::ArcTan { k: k.unwrap_or(10.0) }),
2193 "straightthrough" | "straight_through" | "ste" => Ok(grad::SurrogateType::StraightThrough),
2194 _ => Err(PyValueError::new_err(format!(
2195 "Unknown surrogate '{}'. Use one of: fast_sigmoid, superspike, arctan, straight_through.",
2196 name
2197 ))),
2198 }
2199}
2200
2201fn extract_matrix_f64(data: &Bound<'_, PyAny>, name: &str) -> PyResult<(Vec<f64>, usize, usize)> {
2202 if let Ok(rows) = data.extract::<Vec<Vec<f64>>>() {
2203 if rows.is_empty() {
2204 return Err(PyValueError::new_err(format!(
2205 "{} must not be an empty matrix.",
2206 name
2207 )));
2208 }
2209 let row_count = rows.len();
2210 let cols = rows[0].len();
2211 if cols == 0 {
2212 return Err(PyValueError::new_err(format!(
2213 "{} must not have zero columns.",
2214 name
2215 )));
2216 }
2217 if rows.iter().any(|r| r.len() != cols) {
2218 return Err(PyValueError::new_err(format!(
2219 "{} must be a rectangular matrix.",
2220 name
2221 )));
2222 }
2223 let out = rows.into_iter().flatten().collect::<Vec<f64>>();
2224 return Ok((out, row_count, cols));
2225 }
2226
2227 if let Ok(flat) = data.extract::<Vec<f64>>() {
2228 if flat.is_empty() {
2229 return Err(PyValueError::new_err(format!(
2230 "{} must not be an empty vector.",
2231 name
2232 )));
2233 }
2234 let cols = flat.len();
2235 return Ok((flat, 1, cols));
2236 }
2237
2238 Err(PyValueError::new_err(format!(
2239 "{} must be a 1-D or 2-D float array.",
2240 name
2241 )))
2242}
2243
2244fn reshape_flat_to_rows(flat: Vec<f64>, rows: usize, cols: usize) -> Vec<Vec<f64>> {
2245 let mut out = Vec::with_capacity(rows);
2246 for i in 0..rows {
2247 out.push(flat[i * cols..(i + 1) * cols].to_vec());
2248 }
2249 out
2250}
2251
2252#[pyclass(
2253 name = "SurrogateLif",
2254 module = "sc_neurocore_engine.sc_neurocore_engine"
2255)]
2256pub struct PySurrogateLif {
2257 inner: grad::SurrogateLif,
2258}
2259
2260#[pymethods]
2261impl PySurrogateLif {
2262 #[new]
2263 #[pyo3(signature = (
2264 data_width=16,
2265 fraction=8,
2266 v_rest=0,
2267 v_reset=0,
2268 v_threshold=256,
2269 refractory_period=2,
2270 surrogate="fast_sigmoid",
2271 k=None
2272 ))]
2273 #[allow(clippy::too_many_arguments)]
2274 fn new(
2275 data_width: u32,
2276 fraction: u32,
2277 v_rest: i16,
2278 v_reset: i16,
2279 v_threshold: i16,
2280 refractory_period: i32,
2281 surrogate: &str,
2282 k: Option<f32>,
2283 ) -> PyResult<Self> {
2284 let surrogate = parse_surrogate(surrogate, k)?;
2285 Ok(Self {
2286 inner: grad::SurrogateLif::new(
2287 data_width,
2288 fraction,
2289 v_rest,
2290 v_reset,
2291 v_threshold,
2292 refractory_period,
2293 surrogate,
2294 ),
2295 })
2296 }
2297
2298 #[pyo3(signature = (leak_k, gain_k, i_t, noise_in=0))]
2299 fn forward(&mut self, leak_k: i16, gain_k: i16, i_t: i16, noise_in: i16) -> (i32, i16) {
2300 self.inner.forward(leak_k, gain_k, i_t, noise_in)
2301 }
2302
2303 fn backward(&mut self, grad_output: f32) -> PyResult<f32> {
2304 self.inner
2305 .backward(grad_output)
2306 .map_err(PyValueError::new_err)
2307 }
2308
2309 fn clear_trace(&mut self) {
2310 self.inner.clear_trace();
2311 }
2312
2313 fn reset(&mut self) {
2314 self.inner.reset();
2315 }
2316
2317 fn trace_len(&self) -> usize {
2318 self.inner.trace_len()
2319 }
2320}
2321
2322#[pyclass(
2323 name = "DifferentiableDenseLayer",
2324 module = "sc_neurocore_engine.sc_neurocore_engine"
2325)]
2326pub struct PyDifferentiableDenseLayer {
2327 inner: grad::DifferentiableDenseLayer,
2328}
2329
2330#[pymethods]
2331impl PyDifferentiableDenseLayer {
2332 #[new]
2333 #[pyo3(signature = (
2334 n_inputs,
2335 n_neurons,
2336 length=1024,
2337 seed=24301,
2338 surrogate="fast_sigmoid",
2339 k=None
2340 ))]
2341 fn new(
2342 n_inputs: usize,
2343 n_neurons: usize,
2344 length: usize,
2345 seed: u64,
2346 surrogate: &str,
2347 k: Option<f32>,
2348 ) -> PyResult<Self> {
2349 let surrogate = parse_surrogate(surrogate, k)?;
2350 Ok(Self {
2351 inner: grad::DifferentiableDenseLayer::new(
2352 n_inputs, n_neurons, length, seed, surrogate,
2353 ),
2354 })
2355 }
2356
2357 fn get_weights(&self) -> Vec<Vec<f64>> {
2358 self.inner.layer.get_weights()
2359 }
2360
2361 #[pyo3(signature = (input_values, seed=44257))]
2362 fn forward(&mut self, input_values: Vec<f64>, seed: u64) -> PyResult<Vec<f64>> {
2363 self.inner
2364 .forward(&input_values, seed)
2365 .map_err(PyValueError::new_err)
2366 }
2367
2368 fn backward(&self, grad_output: Vec<f64>) -> PyResult<(Vec<f64>, Vec<Vec<f64>>)> {
2369 self.inner
2370 .backward(&grad_output)
2371 .map_err(PyValueError::new_err)
2372 }
2373
2374 fn update_weights(&mut self, weight_grads: Vec<Vec<f64>>, lr: f64) -> PyResult<()> {
2375 if weight_grads.len() != self.inner.layer.n_neurons {
2376 return Err(PyValueError::new_err(format!(
2377 "Expected {} grad rows, got {}.",
2378 self.inner.layer.n_neurons,
2379 weight_grads.len()
2380 )));
2381 }
2382 if weight_grads
2383 .iter()
2384 .any(|row| row.len() != self.inner.layer.n_inputs)
2385 {
2386 return Err(PyValueError::new_err(format!(
2387 "Expected each grad row to have length {}.",
2388 self.inner.layer.n_inputs
2389 )));
2390 }
2391 self.inner.update_weights(&weight_grads, lr);
2392 Ok(())
2393 }
2394
2395 fn clear_cache(&mut self) {
2396 self.inner.clear_cache();
2397 }
2398}
2399
2400#[pyclass(
2401 name = "StochasticAttention",
2402 module = "sc_neurocore_engine.sc_neurocore_engine"
2403)]
2404pub struct PyStochasticAttention {
2405 inner: attention::StochasticAttention,
2406}
2407
2408#[pymethods]
2409impl PyStochasticAttention {
2410 #[new]
2411 #[pyo3(signature = (dim_k, temperature=None))]
2412 fn new(dim_k: usize, temperature: Option<f64>) -> Self {
2413 Self {
2414 inner: match temperature {
2415 Some(t) => attention::StochasticAttention::with_temperature(dim_k, t),
2416 None => attention::StochasticAttention::new(dim_k),
2417 },
2418 }
2419 }
2420
2421 fn forward_softmax(
2422 &self,
2423 q: &Bound<'_, PyAny>,
2424 k: &Bound<'_, PyAny>,
2425 v: &Bound<'_, PyAny>,
2426 ) -> PyResult<Vec<Vec<f64>>> {
2427 let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2428 let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2429 let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2430
2431 let out = self
2432 .inner
2433 .forward_softmax(
2434 &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols,
2435 )
2436 .map_err(PyValueError::new_err)?;
2437
2438 Ok(reshape_flat_to_rows(out, q_rows, v_cols))
2439 }
2440
2441 #[pyo3(signature = (q, k, v, n_heads))]
2442 fn forward_multihead_softmax(
2443 &self,
2444 q: &Bound<'_, PyAny>,
2445 k: &Bound<'_, PyAny>,
2446 v: &Bound<'_, PyAny>,
2447 n_heads: usize,
2448 ) -> PyResult<Vec<Vec<f64>>> {
2449 let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2450 let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2451 let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2452
2453 let out = self
2454 .inner
2455 .forward_multihead_softmax(
2456 &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols, n_heads,
2457 )
2458 .map_err(PyValueError::new_err)?;
2459
2460 let out_cols = v_cols;
2461 Ok(reshape_flat_to_rows(out, q_rows, out_cols))
2462 }
2463
2464 fn forward(
2465 &self,
2466 q: &Bound<'_, PyAny>,
2467 k: &Bound<'_, PyAny>,
2468 v: &Bound<'_, PyAny>,
2469 ) -> PyResult<Vec<Vec<f64>>> {
2470 let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2471 let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2472 let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2473
2474 let out = self
2475 .inner
2476 .forward(
2477 &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols,
2478 )
2479 .map_err(PyValueError::new_err)?;
2480
2481 Ok(reshape_flat_to_rows(out, q_rows, v_cols))
2482 }
2483
2484 #[pyo3(signature = (q, k, v, length=1024, seed=44257))]
2485 fn forward_sc(
2486 &self,
2487 q: &Bound<'_, PyAny>,
2488 k: &Bound<'_, PyAny>,
2489 v: &Bound<'_, PyAny>,
2490 length: usize,
2491 seed: u64,
2492 ) -> PyResult<Vec<Vec<f64>>> {
2493 let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2494 let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2495 let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2496
2497 let out = self
2498 .inner
2499 .forward_sc(
2500 &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols, length,
2501 seed,
2502 )
2503 .map_err(PyValueError::new_err)?;
2504
2505 Ok(reshape_flat_to_rows(out, q_rows, v_cols))
2506 }
2507
2508 #[pyo3(signature = (q, k, v, n_heads))]
2509 fn forward_multihead(
2510 &self,
2511 q: &Bound<'_, PyAny>,
2512 k: &Bound<'_, PyAny>,
2513 v: &Bound<'_, PyAny>,
2514 n_heads: usize,
2515 ) -> PyResult<Vec<Vec<f64>>> {
2516 let (q_data, q_rows, q_cols) = extract_matrix_f64(q, "Q")?;
2517 let (k_data, k_rows, k_cols) = extract_matrix_f64(k, "K")?;
2518 let (v_data, v_rows, v_cols) = extract_matrix_f64(v, "V")?;
2519
2520 let out = self
2521 .inner
2522 .forward_multihead(
2523 &q_data, q_rows, q_cols, &k_data, k_rows, k_cols, &v_data, v_rows, v_cols, n_heads,
2524 )
2525 .map_err(PyValueError::new_err)?;
2526
2527 let out_cols = v_cols;
2528 Ok(reshape_flat_to_rows(out, q_rows, out_cols))
2529 }
2530}
2531
2532#[pyclass(
2533 name = "StochasticGraphLayer",
2534 module = "sc_neurocore_engine.sc_neurocore_engine"
2535)]
2536pub struct PyStochasticGraphLayer {
2537 inner: graph::StochasticGraphLayer,
2538}
2539
2540#[pymethods]
2541impl PyStochasticGraphLayer {
2542 #[new]
2543 #[pyo3(signature = (adj_matrix, n_features, seed=42))]
2544 fn new(adj_matrix: &Bound<'_, PyAny>, n_features: usize, seed: u64) -> PyResult<Self> {
2545 let (adj_flat, n_rows, n_cols) = extract_matrix_f64(adj_matrix, "adj_matrix")?;
2546 if n_rows != n_cols {
2547 return Err(PyValueError::new_err(format!(
2548 "adj_matrix must be square, got {}x{}.",
2549 n_rows, n_cols
2550 )));
2551 }
2552 Ok(Self {
2553 inner: graph::StochasticGraphLayer::new(adj_flat, n_rows, n_features, seed),
2554 })
2555 }
2556
2557 #[staticmethod]
2559 #[pyo3(signature = (row_offsets, col_indices, values, n_nodes, n_features, seed=42))]
2560 fn from_sparse(
2561 row_offsets: Vec<usize>,
2562 col_indices: Vec<usize>,
2563 values: Vec<f64>,
2564 n_nodes: usize,
2565 n_features: usize,
2566 seed: u64,
2567 ) -> PyResult<Self> {
2568 let csr = graph::CsrMatrix::new(row_offsets, col_indices, values, n_nodes, n_nodes)
2569 .map_err(PyValueError::new_err)?;
2570 let inner = graph::StochasticGraphLayer::new_sparse(csr, n_features, seed)
2571 .map_err(PyValueError::new_err)?;
2572 Ok(Self { inner })
2573 }
2574
2575 #[staticmethod]
2577 #[pyo3(signature = (adj_matrix, n_features, seed=42, density_threshold=0.3))]
2578 fn from_dense_auto(
2579 adj_matrix: &Bound<'_, PyAny>,
2580 n_features: usize,
2581 seed: u64,
2582 density_threshold: f64,
2583 ) -> PyResult<Self> {
2584 let (adj_flat, n_rows, n_cols) = extract_matrix_f64(adj_matrix, "adj_matrix")?;
2585 if n_rows != n_cols {
2586 return Err(PyValueError::new_err(format!(
2587 "adj_matrix must be square, got {}x{}.",
2588 n_rows, n_cols
2589 )));
2590 }
2591 Ok(Self {
2592 inner: graph::StochasticGraphLayer::from_dense_auto(
2593 adj_flat,
2594 n_rows,
2595 n_features,
2596 seed,
2597 density_threshold,
2598 ),
2599 })
2600 }
2601
2602 fn is_sparse(&self) -> bool {
2603 self.inner.is_sparse()
2604 }
2605
2606 fn forward(&self, node_features: &Bound<'_, PyAny>) -> PyResult<Vec<Vec<f64>>> {
2607 let (x_flat, x_rows, x_cols) = extract_matrix_f64(node_features, "node_features")?;
2608 if x_rows != self.inner.n_nodes || x_cols != self.inner.n_features {
2609 return Err(PyValueError::new_err(format!(
2610 "Expected node_features shape ({}, {}), got ({}, {}).",
2611 self.inner.n_nodes, self.inner.n_features, x_rows, x_cols
2612 )));
2613 }
2614 let out = self.inner.forward(&x_flat).map_err(PyValueError::new_err)?;
2615 Ok(reshape_flat_to_rows(
2616 out,
2617 self.inner.n_nodes,
2618 self.inner.n_features,
2619 ))
2620 }
2621
2622 #[pyo3(signature = (node_features, length=1024, seed=44257))]
2623 fn forward_sc(
2624 &self,
2625 node_features: &Bound<'_, PyAny>,
2626 length: usize,
2627 seed: u64,
2628 ) -> PyResult<Vec<Vec<f64>>> {
2629 let (x_flat, x_rows, x_cols) = extract_matrix_f64(node_features, "node_features")?;
2630 if x_rows != self.inner.n_nodes || x_cols != self.inner.n_features {
2631 return Err(PyValueError::new_err(format!(
2632 "Expected node_features shape ({}, {}), got ({}, {}).",
2633 self.inner.n_nodes, self.inner.n_features, x_rows, x_cols
2634 )));
2635 }
2636 let out = self
2637 .inner
2638 .forward_sc(&x_flat, length, seed)
2639 .map_err(PyValueError::new_err)?;
2640 Ok(reshape_flat_to_rows(
2641 out,
2642 self.inner.n_nodes,
2643 self.inner.n_features,
2644 ))
2645 }
2646
2647 fn get_weights(&self) -> Vec<f64> {
2648 self.inner.get_weights()
2649 }
2650
2651 fn set_weights(&mut self, weights: Vec<f64>) -> PyResult<()> {
2652 self.inner
2653 .set_weights(weights)
2654 .map_err(PyValueError::new_err)
2655 }
2656}
2657
2658#[pyclass(
2659 name = "KuramotoSolver",
2660 module = "sc_neurocore_engine.sc_neurocore_engine"
2661)]
2662pub struct PyKuramotoSolver {
2663 inner: scpn::KuramotoSolver,
2664}
2665
2666fn validate_kuramoto_finite(name: &str, values: &[f64]) -> PyResult<()> {
2667 if values.iter().all(|value| value.is_finite()) {
2668 Ok(())
2669 } else {
2670 Err(PyValueError::new_err(format!(
2671 "{name} values must be finite"
2672 )))
2673 }
2674}
2675
2676fn validate_kuramoto_dt(dt: f64) -> PyResult<()> {
2677 if dt.is_finite() && dt > 0.0 {
2678 Ok(())
2679 } else {
2680 Err(PyValueError::new_err("dt must be finite and positive"))
2681 }
2682}
2683
2684fn validate_kuramoto_matrix_shape(
2685 name: &str,
2686 values_len: usize,
2687 rows: usize,
2688 cols: usize,
2689 n: usize,
2690) -> PyResult<()> {
2691 let is_absent = rows == 0 && cols == 0 && values_len == 0;
2692 let is_flat = rows == 1 && values_len == n * n;
2693 let is_square = rows == n && cols == n;
2694 if is_absent || is_flat || is_square {
2695 Ok(())
2696 } else {
2697 Err(PyValueError::new_err(format!(
2698 "{name} must be shape ({n}, {n}) or flat length {}",
2699 n * n
2700 )))
2701 }
2702}
2703
2704#[pymethods]
2705impl PyKuramotoSolver {
2706 #[getter]
2707 fn phases(&self) -> Vec<f64> {
2708 self.inner.phases.clone()
2709 }
2710 #[new]
2711 #[pyo3(signature = (omega, coupling, phases, noise_amp=0.1))]
2712 fn new(
2713 omega: Vec<f64>,
2714 coupling: &Bound<'_, PyAny>,
2715 phases: Vec<f64>,
2716 noise_amp: f64,
2717 ) -> PyResult<Self> {
2718 let n = omega.len();
2719 if n == 0 {
2720 return Err(PyValueError::new_err("omega must not be empty."));
2721 }
2722 if phases.len() != n {
2723 return Err(PyValueError::new_err(format!(
2724 "phases length mismatch: got {}, expected {}.",
2725 phases.len(),
2726 n
2727 )));
2728 }
2729 validate_kuramoto_finite("omega", &omega)?;
2730 validate_kuramoto_finite("initial_phases", &phases)?;
2731 if !(noise_amp.is_finite() && noise_amp >= 0.0) {
2732 return Err(PyValueError::new_err(
2733 "noise_amp must be finite and non-negative",
2734 ));
2735 }
2736
2737 let (coupling_flat, rows, cols) = extract_matrix_f64(coupling, "coupling")?;
2738 if rows == 1 {
2739 if coupling_flat.len() != n * n {
2740 return Err(PyValueError::new_err(format!(
2741 "Flat coupling length mismatch: got {}, expected {}.",
2742 coupling_flat.len(),
2743 n * n
2744 )));
2745 }
2746 } else if rows != n || cols != n {
2747 return Err(PyValueError::new_err(format!(
2748 "coupling must be shape ({}, {}) or flat length {}, got ({}, {}).",
2749 n,
2750 n,
2751 n * n,
2752 rows,
2753 cols
2754 )));
2755 }
2756 validate_kuramoto_finite("coupling", &coupling_flat)?;
2757
2758 Ok(Self {
2759 inner: scpn::KuramotoSolver::new(omega, coupling_flat, phases, noise_amp),
2760 })
2761 }
2762
2763 #[pyo3(signature = (dt, seed=0))]
2764 fn step(&mut self, dt: f64, seed: u64) -> PyResult<f64> {
2765 validate_kuramoto_dt(dt)?;
2766 Ok(self.inner.step(dt, seed))
2767 }
2768
2769 #[pyo3(signature = (n_steps, dt, seed=0))]
2770 fn run(&mut self, n_steps: usize, dt: f64, seed: u64) -> PyResult<Vec<f64>> {
2771 validate_kuramoto_dt(dt)?;
2772 Ok(self.inner.run(n_steps, dt, seed))
2773 }
2774
2775 fn set_field_pressure(&mut self, f: f64) -> PyResult<()> {
2776 if !f.is_finite() {
2777 return Err(PyValueError::new_err("field_pressure must be finite"));
2778 }
2779 self.inner.set_field_pressure(f);
2780 Ok(())
2781 }
2782
2783 #[pyo3(signature = (
2784 dt,
2785 seed=0,
2786 W=None,
2787 sigma_g=0.0,
2788 h_munu=None,
2789 pgbo_weight=0.0,
2790 ))]
2791 #[allow(non_snake_case)]
2792 fn step_ssgf(
2793 &mut self,
2794 dt: f64,
2795 seed: u64,
2796 W: Option<&Bound<'_, PyAny>>,
2797 sigma_g: f64,
2798 h_munu: Option<&Bound<'_, PyAny>>,
2799 pgbo_weight: f64,
2800 ) -> PyResult<f64> {
2801 validate_kuramoto_dt(dt)?;
2802 if !sigma_g.is_finite() {
2803 return Err(PyValueError::new_err("sigma_g must be finite"));
2804 }
2805 if !pgbo_weight.is_finite() {
2806 return Err(PyValueError::new_err("pgbo_weight must be finite"));
2807 }
2808 let (w_flat, w_rows, w_cols) = match W {
2809 Some(w) => extract_matrix_f64(w, "W")?,
2810 None => (vec![], 0, 0),
2811 };
2812 let (h_flat, h_rows, h_cols) = match h_munu {
2813 Some(h) => extract_matrix_f64(h, "h_munu")?,
2814 None => (vec![], 0, 0),
2815 };
2816 validate_kuramoto_matrix_shape("W", w_flat.len(), w_rows, w_cols, self.inner.n)?;
2817 validate_kuramoto_matrix_shape("h_munu", h_flat.len(), h_rows, h_cols, self.inner.n)?;
2818 validate_kuramoto_finite("w_flat", &w_flat)?;
2819 validate_kuramoto_finite("h_flat", &h_flat)?;
2820 Ok(self
2821 .inner
2822 .step_ssgf(dt, seed, &w_flat, sigma_g, &h_flat, pgbo_weight))
2823 }
2824
2825 #[pyo3(signature = (
2826 n_steps,
2827 dt,
2828 seed=0,
2829 W=None,
2830 sigma_g=0.0,
2831 h_munu=None,
2832 pgbo_weight=0.0,
2833 ))]
2834 #[allow(clippy::too_many_arguments, non_snake_case)]
2835 fn run_ssgf(
2836 &mut self,
2837 n_steps: usize,
2838 dt: f64,
2839 seed: u64,
2840 W: Option<&Bound<'_, PyAny>>,
2841 sigma_g: f64,
2842 h_munu: Option<&Bound<'_, PyAny>>,
2843 pgbo_weight: f64,
2844 ) -> PyResult<Vec<f64>> {
2845 validate_kuramoto_dt(dt)?;
2846 if !sigma_g.is_finite() {
2847 return Err(PyValueError::new_err("sigma_g must be finite"));
2848 }
2849 if !pgbo_weight.is_finite() {
2850 return Err(PyValueError::new_err("pgbo_weight must be finite"));
2851 }
2852 let (w_flat, w_rows, w_cols) = match W {
2853 Some(w) => extract_matrix_f64(w, "W")?,
2854 None => (vec![], 0, 0),
2855 };
2856 let (h_flat, h_rows, h_cols) = match h_munu {
2857 Some(h) => extract_matrix_f64(h, "h_munu")?,
2858 None => (vec![], 0, 0),
2859 };
2860 validate_kuramoto_matrix_shape("W", w_flat.len(), w_rows, w_cols, self.inner.n)?;
2861 validate_kuramoto_matrix_shape("h_munu", h_flat.len(), h_rows, h_cols, self.inner.n)?;
2862 validate_kuramoto_finite("w_flat", &w_flat)?;
2863 validate_kuramoto_finite("h_flat", &h_flat)?;
2864 Ok(self
2865 .inner
2866 .run_ssgf(n_steps, dt, seed, &w_flat, sigma_g, &h_flat, pgbo_weight))
2867 }
2868
2869 fn order_parameter(&self) -> f64 {
2870 self.inner.order_parameter()
2871 }
2872
2873 fn apply_phases(&mut self, phases: Vec<f64>) -> PyResult<()> {
2874 if phases.len() != self.inner.n {
2875 return Err(PyValueError::new_err(format!(
2876 "phases length mismatch: got {}, expected {}.",
2877 phases.len(),
2878 self.inner.n
2879 )));
2880 }
2881 validate_kuramoto_finite("phases", &phases)?;
2882 self.inner.set_phases(phases);
2883 Ok(())
2884 }
2885
2886 fn set_phases(&mut self, phases: Vec<f64>) -> PyResult<()> {
2887 self.apply_phases(phases)
2888 }
2889
2890 #[setter(phases)]
2891 fn set_phases_attr(&mut self, phases: Vec<f64>) -> PyResult<()> {
2892 self.apply_phases(phases)
2893 }
2894
2895 fn set_coupling(&mut self, coupling: &Bound<'_, PyAny>) -> PyResult<()> {
2896 let n = self.inner.n;
2897 let (coupling_flat, rows, cols) = extract_matrix_f64(coupling, "coupling")?;
2898 if rows == 1 {
2899 if coupling_flat.len() != n * n {
2900 return Err(PyValueError::new_err(format!(
2901 "Flat coupling length mismatch: got {}, expected {}.",
2902 coupling_flat.len(),
2903 n * n
2904 )));
2905 }
2906 } else if rows != n || cols != n {
2907 return Err(PyValueError::new_err(format!(
2908 "coupling must be shape ({}, {}) or flat length {}, got ({}, {}).",
2909 n,
2910 n,
2911 n * n,
2912 rows,
2913 cols
2914 )));
2915 }
2916 validate_kuramoto_finite("coupling", &coupling_flat)?;
2917 self.inner.set_coupling(coupling_flat);
2918 Ok(())
2919 }
2920}
2921
2922#[pyclass(
2923 name = "SCPNMetrics",
2924 module = "sc_neurocore_engine.sc_neurocore_engine"
2925)]
2926pub struct PySCPNMetrics;
2927
2928#[pymethods]
2929impl PySCPNMetrics {
2930 #[new]
2931 fn new() -> Self {
2932 Self
2933 }
2934
2935 #[staticmethod]
2936 fn global_coherence(weights: [f64; 7], metrics: [f64; 7]) -> f64 {
2937 scpn::SCPNMetrics::global_coherence(&weights, &metrics)
2938 }
2939
2940 #[staticmethod]
2941 fn consciousness_index(phases_l4: Vec<f64>, glyph_l7: [f64; 6]) -> f64 {
2942 scpn::SCPNMetrics::consciousness_index(&phases_l4, &glyph_l7)
2943 }
2944}
2945
2946#[pyclass(name = "ScGraph", module = "sc_neurocore_engine.sc_neurocore_engine")]
2949pub struct PyScGraph {
2950 inner: ir::graph::ScGraph,
2951}
2952
2953#[pymethods]
2954impl PyScGraph {
2955 fn len(&self) -> usize {
2957 self.inner.len()
2958 }
2959
2960 fn __len__(&self) -> usize {
2961 self.inner.len()
2962 }
2963
2964 fn is_empty(&self) -> bool {
2966 self.inner.is_empty()
2967 }
2968
2969 #[getter]
2971 fn name(&self) -> &str {
2972 &self.inner.name
2973 }
2974
2975 fn num_inputs(&self) -> usize {
2977 self.inner.inputs().len()
2978 }
2979
2980 fn num_outputs(&self) -> usize {
2982 self.inner.outputs().len()
2983 }
2984
2985 fn __repr__(&self) -> String {
2986 format!("ScGraph('{}', ops={})", self.inner.name, self.inner.len())
2987 }
2988}
2989
2990#[pyclass(
2991 name = "ScGraphBuilder",
2992 module = "sc_neurocore_engine.sc_neurocore_engine"
2993)]
2994pub struct PyScGraphBuilder {
2995 inner: Option<ir::builder::ScGraphBuilder>,
2996}
2997
2998impl PyScGraphBuilder {
2999 fn builder_mut(&mut self) -> PyResult<&mut ir::builder::ScGraphBuilder> {
3000 self.inner
3001 .as_mut()
3002 .ok_or_else(|| PyValueError::new_err("Builder already consumed by build()."))
3003 }
3004}
3005
3006#[pymethods]
3007impl PyScGraphBuilder {
3008 #[new]
3009 fn new(name: String) -> Self {
3010 Self {
3011 inner: Some(ir::builder::ScGraphBuilder::new(name)),
3012 }
3013 }
3014
3015 fn input(&mut self, name: &str, ty: &str) -> PyResult<u32> {
3017 let sc_type = parse_sc_type(ty)?;
3018 Ok(self.builder_mut()?.input(name, sc_type).0)
3019 }
3020
3021 fn output(&mut self, name: &str, source_id: u32) -> PyResult<u32> {
3023 Ok(self
3024 .builder_mut()?
3025 .output(name, ir::graph::ValueId(source_id))
3026 .0)
3027 }
3028
3029 fn constant_f64(&mut self, value: f64, ty: &str) -> PyResult<u32> {
3031 let sc_type = parse_sc_type(ty)?;
3032 Ok(self
3033 .builder_mut()?
3034 .constant(ir::graph::ScConst::F64(value), sc_type)
3035 .0)
3036 }
3037
3038 fn constant_i64(&mut self, value: i64, ty: &str) -> PyResult<u32> {
3040 let sc_type = parse_sc_type(ty)?;
3041 Ok(self
3042 .builder_mut()?
3043 .constant(ir::graph::ScConst::I64(value), sc_type)
3044 .0)
3045 }
3046
3047 fn encode(&mut self, prob_id: u32, length: usize, seed: u64) -> PyResult<u32> {
3049 let seed = u16::try_from(seed)
3050 .map_err(|_| PyValueError::new_err(format!("Seed out of range for u16: {seed}")))?;
3051 Ok(self
3052 .builder_mut()?
3053 .encode(ir::graph::ValueId(prob_id), length, seed)
3054 .0)
3055 }
3056
3057 fn bitwise_and(&mut self, lhs_id: u32, rhs_id: u32) -> PyResult<u32> {
3059 Ok(self
3060 .builder_mut()?
3061 .bitwise_and(ir::graph::ValueId(lhs_id), ir::graph::ValueId(rhs_id))
3062 .0)
3063 }
3064
3065 fn popcount(&mut self, input_id: u32) -> PyResult<u32> {
3067 Ok(self.builder_mut()?.popcount(ir::graph::ValueId(input_id)).0)
3068 }
3069
3070 #[pyo3(signature = (
3072 current_id,
3073 leak_id,
3074 gain_id,
3075 noise_id,
3076 data_width=16,
3077 fraction=8,
3078 v_rest=0,
3079 v_reset=0,
3080 v_threshold=256,
3081 refractory_period=2
3082 ))]
3083 #[allow(clippy::too_many_arguments)]
3084 fn lif_step(
3085 &mut self,
3086 current_id: u32,
3087 leak_id: u32,
3088 gain_id: u32,
3089 noise_id: u32,
3090 data_width: u32,
3091 fraction: u32,
3092 v_rest: i64,
3093 v_reset: i64,
3094 v_threshold: i64,
3095 refractory_period: u32,
3096 ) -> PyResult<u32> {
3097 let params = ir::graph::LifParams {
3098 data_width,
3099 fraction,
3100 v_rest,
3101 v_reset,
3102 v_threshold,
3103 refractory_period,
3104 };
3105 Ok(self
3106 .builder_mut()?
3107 .lif_step(
3108 ir::graph::ValueId(current_id),
3109 ir::graph::ValueId(leak_id),
3110 ir::graph::ValueId(gain_id),
3111 ir::graph::ValueId(noise_id),
3112 params,
3113 )
3114 .0)
3115 }
3116
3117 #[pyo3(signature = (
3119 inputs_id,
3120 weights_id,
3121 leak_id,
3122 gain_id,
3123 n_inputs=3,
3124 n_neurons=7,
3125 data_width=16,
3126 stream_length=1024,
3127 seed_base=0xACE1u64,
3128 y_min=0,
3129 y_max=65535
3130 ))]
3131 #[allow(clippy::too_many_arguments)]
3132 fn dense_forward(
3133 &mut self,
3134 inputs_id: u32,
3135 weights_id: u32,
3136 leak_id: u32,
3137 gain_id: u32,
3138 n_inputs: usize,
3139 n_neurons: usize,
3140 data_width: u32,
3141 stream_length: usize,
3142 seed_base: u64,
3143 y_min: i64,
3144 y_max: i64,
3145 ) -> PyResult<u32> {
3146 let input_seed_base = u16::try_from(seed_base).map_err(|_| {
3147 PyValueError::new_err(format!("seed_base out of range for u16: {seed_base}"))
3148 })?;
3149 let params = ir::graph::DenseParams {
3150 n_inputs,
3151 n_neurons,
3152 data_width,
3153 stream_length,
3154 input_seed_base,
3155 weight_seed_base: input_seed_base.wrapping_add(1),
3156 y_min,
3157 y_max,
3158 };
3159 Ok(self
3160 .builder_mut()?
3161 .dense_forward(
3162 ir::graph::ValueId(inputs_id),
3163 ir::graph::ValueId(weights_id),
3164 ir::graph::ValueId(leak_id),
3165 ir::graph::ValueId(gain_id),
3166 params,
3167 )
3168 .0)
3169 }
3170
3171 fn scale(&mut self, input_id: u32, factor: f64) -> PyResult<u32> {
3173 Ok(self
3174 .builder_mut()?
3175 .scale(ir::graph::ValueId(input_id), factor)
3176 .0)
3177 }
3178
3179 fn offset(&mut self, input_id: u32, offset_val: f64) -> PyResult<u32> {
3181 Ok(self
3182 .builder_mut()?
3183 .offset(ir::graph::ValueId(input_id), offset_val)
3184 .0)
3185 }
3186
3187 fn div_const(&mut self, input_id: u32, divisor: u64) -> PyResult<u32> {
3189 Ok(self
3190 .builder_mut()?
3191 .div_const(ir::graph::ValueId(input_id), divisor)
3192 .0)
3193 }
3194
3195 fn build(&mut self) -> PyResult<PyScGraph> {
3197 let builder = self
3198 .inner
3199 .take()
3200 .ok_or_else(|| PyValueError::new_err("Builder already consumed by build()."))?;
3201 Ok(PyScGraph {
3202 inner: builder.build(),
3203 })
3204 }
3205}
3206
3207#[pyfunction]
3209fn ir_verify(graph: PyRef<'_, PyScGraph>) -> Option<Vec<String>> {
3210 match ir::verify::verify(&graph.inner) {
3211 Ok(()) => None,
3212 Err(errors) => Some(errors.iter().map(|e| e.to_string()).collect()),
3213 }
3214}
3215
3216#[pyfunction]
3218fn ir_print(graph: PyRef<'_, PyScGraph>) -> String {
3219 ir::printer::print(&graph.inner)
3220}
3221
3222#[pyfunction]
3224fn ir_parse(text: &str) -> PyResult<PyScGraph> {
3225 ir::parser::parse(text)
3226 .map(|graph| PyScGraph { inner: graph })
3227 .map_err(|e| PyValueError::new_err(e.to_string()))
3228}
3229
3230#[pyfunction]
3232fn ir_emit_sv(graph: PyRef<'_, PyScGraph>) -> PyResult<String> {
3233 ir::emit_sv::emit(&graph.inner).map_err(PyValueError::new_err)
3234}
3235
3236fn parse_sc_type(s: &str) -> PyResult<ir::graph::ScType> {
3241 let s = s.trim();
3242 let lower = s.to_ascii_lowercase();
3243 match lower.as_str() {
3244 "bool" => Ok(ir::graph::ScType::Bool),
3245 "rate" => Ok(ir::graph::ScType::Rate),
3246 "u32" => Ok(ir::graph::ScType::UInt { width: 32 }),
3247 "u64" => Ok(ir::graph::ScType::UInt { width: 64 }),
3248 "i16" => Ok(ir::graph::ScType::SInt { width: 16 }),
3249 "i32" => Ok(ir::graph::ScType::SInt { width: 32 }),
3250 "bitstream" => Ok(ir::graph::ScType::Bitstream { length: 0 }),
3251 _ => {
3252 if let Some(width) = lower.strip_prefix('u') {
3253 if let Ok(width) = width.parse::<u32>() {
3254 return Ok(ir::graph::ScType::UInt { width });
3255 }
3256 }
3257 if let Some(width) = lower.strip_prefix('i') {
3258 if let Ok(width) = width.parse::<u32>() {
3259 return Ok(ir::graph::ScType::SInt { width });
3260 }
3261 }
3262 if let Some(inner) = lower
3263 .strip_prefix("bitstream<")
3264 .and_then(|r| r.strip_suffix('>'))
3265 {
3266 let length = inner.parse::<usize>().map_err(|_| {
3267 PyValueError::new_err(format!("Invalid bitstream length: '{inner}'"))
3268 })?;
3269 return Ok(ir::graph::ScType::Bitstream { length });
3270 }
3271 if let Some(inner) = lower
3272 .strip_prefix("fixed<")
3273 .and_then(|r| r.strip_suffix('>'))
3274 {
3275 let parts: Vec<&str> = inner.split(',').collect();
3276 if parts.len() != 2 {
3277 return Err(PyValueError::new_err(format!(
3278 "fixed type needs 2 params: '{s}'"
3279 )));
3280 }
3281 let width = parts[0].trim().parse::<u32>().map_err(|_| {
3282 PyValueError::new_err(format!("Invalid fixed width: '{}'", parts[0]))
3283 })?;
3284 let frac = parts[1].trim().parse::<u32>().map_err(|_| {
3285 PyValueError::new_err(format!("Invalid fixed frac: '{}'", parts[1]))
3286 })?;
3287 return Ok(ir::graph::ScType::FixedPoint { width, frac });
3288 }
3289 if let Some(inner) = lower.strip_prefix("vec<").and_then(|r| r.strip_suffix('>')) {
3290 if let Some(comma_pos) = inner.rfind(',') {
3291 let inner_ty_str = &inner[..comma_pos];
3292 let count_str = inner[comma_pos + 1..].trim();
3293 let inner_ty = parse_sc_type(inner_ty_str)?;
3294 let count = count_str.parse::<usize>().map_err(|_| {
3295 PyValueError::new_err(format!("Invalid vec count: '{count_str}'"))
3296 })?;
3297 return Ok(ir::graph::ScType::Vec {
3298 element: Box::new(inner_ty),
3299 count,
3300 });
3301 }
3302 }
3303 Err(PyValueError::new_err(format!("Unknown IR type: '{s}'")))
3304 }
3305 }
3306}
3307
3308#[pyfunction]
3311#[pyo3(signature = (binary_train, dt=0.001))]
3312fn py_spike_times(
3313 py: Python<'_>,
3314 binary_train: PyReadonlyArray1<'_, i32>,
3315 dt: f64,
3316) -> Py<PyArray1<f64>> {
3317 let data = binary_train.as_slice().unwrap();
3318 analysis::basic::spike_times(data, dt)
3319 .into_pyarray(py)
3320 .into()
3321}
3322
3323#[pyfunction]
3324#[pyo3(signature = (binary_train, dt=0.001))]
3325fn py_isi(py: Python<'_>, binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> Py<PyArray1<f64>> {
3326 let data = binary_train.as_slice().unwrap();
3327 analysis::basic::isi(data, dt).into_pyarray(py).into()
3328}
3329
3330#[pyfunction]
3331#[pyo3(signature = (binary_train, dt=0.001))]
3332fn py_firing_rate(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
3333 let data = binary_train.as_slice().unwrap();
3334 analysis::basic::firing_rate(data, dt)
3335}
3336
3337#[pyfunction]
3338fn py_spike_count(binary_train: PyReadonlyArray1<'_, i32>) -> i64 {
3339 let data = binary_train.as_slice().unwrap();
3340 analysis::basic::spike_count(data)
3341}
3342
3343#[pyfunction]
3344#[pyo3(signature = (binary_train, bin_size=10))]
3345fn py_bin_spike_train(
3346 py: Python<'_>,
3347 binary_train: PyReadonlyArray1<'_, i32>,
3348 bin_size: usize,
3349) -> Py<PyArray1<i64>> {
3350 let data = binary_train.as_slice().unwrap();
3351 analysis::basic::bin_spike_train(data, bin_size)
3352 .into_pyarray(py)
3353 .into()
3354}
3355
3356#[pyfunction]
3357#[pyo3(signature = (binary_train, dt=0.001, kernel="gaussian", sigma_ms=10.0))]
3358fn py_instantaneous_rate(
3359 py: Python<'_>,
3360 binary_train: PyReadonlyArray1<'_, f64>,
3361 dt: f64,
3362 kernel: &str,
3363 sigma_ms: f64,
3364) -> Py<PyArray1<f64>> {
3365 let data = binary_train.as_slice().unwrap();
3366 analysis::rate::instantaneous_rate(data, dt, kernel, sigma_ms)
3367 .into_pyarray(py)
3368 .into()
3369}
3370
3371#[pyfunction]
3372#[pyo3(signature = (trials, bin_ms=10.0, dt=0.001))]
3373fn py_psth(
3374 py: Python<'_>,
3375 trials: Vec<PyReadonlyArray1<'_, f64>>,
3376 bin_ms: f64,
3377 dt: f64,
3378) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
3379 let vecs: Vec<Vec<f64>> = trials
3380 .iter()
3381 .map(|t| t.as_slice().unwrap().to_vec())
3382 .collect();
3383 let (rates, centers) = analysis::rate::psth(&vecs, bin_ms, dt);
3384 (
3385 rates.into_pyarray(py).into(),
3386 centers.into_pyarray(py).into(),
3387 )
3388}
3389
3390#[pyfunction]
3391#[pyo3(signature = (binary_train, dt=0.001))]
3392fn py_cv_isi(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
3393 let data = binary_train.as_slice().unwrap();
3394 analysis::variability::cv_isi(data, dt)
3395}
3396
3397#[pyfunction]
3398#[pyo3(signature = (binary_train, dt=0.001))]
3399fn py_cv2(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
3400 let data = binary_train.as_slice().unwrap();
3401 analysis::variability::cv2(data, dt)
3402}
3403
3404#[pyfunction]
3405#[pyo3(signature = (binary_train, dt=0.001))]
3406fn py_local_variation(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
3407 let data = binary_train.as_slice().unwrap();
3408 analysis::variability::local_variation(data, dt)
3409}
3410
3411#[pyfunction]
3412#[pyo3(signature = (binary_train, window_ms=50.0, dt=0.001))]
3413fn py_fano_factor(binary_train: PyReadonlyArray1<'_, i32>, window_ms: f64, dt: f64) -> f64 {
3414 let data = binary_train.as_slice().unwrap();
3415 analysis::variability::fano_factor(data, window_ms, dt)
3416}
3417
3418#[pyfunction]
3419fn py_lempel_ziv_complexity(binary_train: PyReadonlyArray1<'_, i32>) -> f64 {
3420 let data = binary_train.as_slice().unwrap();
3421 analysis::variability::lempel_ziv_complexity(data)
3422}
3423
3424#[pyfunction]
3425#[pyo3(signature = (binary_train, order=3, delay=1))]
3426fn py_permutation_entropy(
3427 binary_train: PyReadonlyArray1<'_, i32>,
3428 order: usize,
3429 delay: usize,
3430) -> f64 {
3431 let data = binary_train.as_slice().unwrap();
3432 analysis::variability::permutation_entropy(data, order, delay)
3433}
3434
3435#[pyfunction]
3436#[pyo3(signature = (binary_train, min_window=10))]
3437fn py_hurst_exponent(binary_train: PyReadonlyArray1<'_, i32>, min_window: usize) -> f64 {
3438 let data = binary_train.as_slice().unwrap();
3439 analysis::variability::hurst_exponent(data, min_window)
3440}
3441
3442#[pyfunction]
3443#[pyo3(signature = (binary_train, m=2, r_factor=0.2))]
3444fn py_approximate_entropy(binary_train: PyReadonlyArray1<'_, i32>, m: usize, r_factor: f64) -> f64 {
3445 let data = binary_train.as_slice().unwrap();
3446 analysis::variability::approximate_entropy(data, m, r_factor)
3447}
3448
3449#[pyfunction]
3450#[pyo3(signature = (binary_train, m=2, r_factor=0.2))]
3451fn py_sample_entropy(binary_train: PyReadonlyArray1<'_, i32>, m: usize, r_factor: f64) -> f64 {
3452 let data = binary_train.as_slice().unwrap();
3453 analysis::variability::sample_entropy(data, m, r_factor)
3454}
3455
3456#[pyfunction]
3459#[pyo3(signature = (train_a, train_b, max_lag_ms=50.0, dt=0.001))]
3460fn py_cross_correlation(
3461 py: Python<'_>,
3462 train_a: PyReadonlyArray1<'_, i32>,
3463 train_b: PyReadonlyArray1<'_, i32>,
3464 max_lag_ms: f64,
3465 dt: f64,
3466) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
3467 let a = train_a.as_slice().unwrap();
3468 let b = train_b.as_slice().unwrap();
3469 let (cc, lags) = analysis::correlation::cross_correlation(a, b, max_lag_ms, dt);
3470 (cc.into_pyarray(py).into(), lags.into_pyarray(py).into())
3471}
3472
3473#[pyfunction]
3474#[pyo3(signature = (trains, dt=0.001))]
3475fn py_pairwise_correlation(
3476 py: Python<'_>,
3477 trains: Vec<PyReadonlyArray1<'_, i32>>,
3478 dt: f64,
3479) -> Py<PyArray2<f64>> {
3480 let vecs: Vec<Vec<i32>> = trains
3481 .iter()
3482 .map(|t| t.as_slice().unwrap().to_vec())
3483 .collect();
3484 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3485 let mat = analysis::correlation::pairwise_correlation(&refs, dt);
3486 let n = mat.len();
3487 let flat: Vec<f64> = mat.into_iter().flatten().collect();
3488 numpy::PyArray2::from_vec2(py, &flat.chunks(n).map(|c| c.to_vec()).collect::<Vec<_>>())
3489 .unwrap()
3490 .into()
3491}
3492
3493#[pyfunction]
3494#[pyo3(signature = (train_a, train_b, dt=0.001, tau_ms=5.0))]
3495fn py_event_synchronization(
3496 train_a: PyReadonlyArray1<'_, i32>,
3497 train_b: PyReadonlyArray1<'_, i32>,
3498 dt: f64,
3499 tau_ms: f64,
3500) -> f64 {
3501 let a = train_a.as_slice().unwrap();
3502 let b = train_b.as_slice().unwrap();
3503 analysis::correlation::event_synchronization(a, b, dt, tau_ms)
3504}
3505
3506#[pyfunction]
3507#[pyo3(signature = (train_a, train_b, dt=0.001))]
3508fn py_spike_train_coherence(
3509 py: Python<'_>,
3510 train_a: PyReadonlyArray1<'_, i32>,
3511 train_b: PyReadonlyArray1<'_, i32>,
3512 dt: f64,
3513) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
3514 let a = train_a.as_slice().unwrap();
3515 let b = train_b.as_slice().unwrap();
3516 let (coh, freqs) = analysis::correlation::spike_train_coherence(a, b, dt);
3517 (coh.into_pyarray(py).into(), freqs.into_pyarray(py).into())
3518}
3519
3520#[pyfunction]
3521#[pyo3(signature = (train_a, train_b, dt=0.001, delta_ms=5.0))]
3522fn py_spike_time_tiling_coefficient(
3523 train_a: PyReadonlyArray1<'_, i32>,
3524 train_b: PyReadonlyArray1<'_, i32>,
3525 dt: f64,
3526 delta_ms: f64,
3527) -> f64 {
3528 let a = train_a.as_slice().unwrap();
3529 let b = train_b.as_slice().unwrap();
3530 analysis::correlation::spike_time_tiling_coefficient(a, b, dt, delta_ms)
3531}
3532
3533#[pyfunction]
3534#[pyo3(signature = (trains, bin_size=10))]
3535fn py_covariance_matrix(
3536 py: Python<'_>,
3537 trains: Vec<PyReadonlyArray1<'_, i32>>,
3538 bin_size: usize,
3539) -> Py<PyArray2<f64>> {
3540 let vecs: Vec<Vec<i32>> = trains
3541 .iter()
3542 .map(|t| t.as_slice().unwrap().to_vec())
3543 .collect();
3544 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3545 let mat = analysis::correlation::covariance_matrix(&refs, bin_size);
3546 let n = mat.len();
3547 let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3548 numpy::PyArray2::from_vec2(py, &rows)
3549 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3550 .into()
3551}
3552
3553#[pyfunction]
3554#[pyo3(signature = (binary_train, dt=0.001, max_lag_ms=100.0))]
3555fn py_autocorrelation_time(
3556 binary_train: PyReadonlyArray1<'_, i32>,
3557 dt: f64,
3558 max_lag_ms: f64,
3559) -> f64 {
3560 let data = binary_train.as_slice().unwrap();
3561 analysis::correlation::autocorrelation_time(data, dt, max_lag_ms)
3562}
3563
3564#[pyfunction]
3565#[pyo3(signature = (trains, bin_size=50))]
3566fn py_noise_correlation(
3567 py: Python<'_>,
3568 trains: Vec<PyReadonlyArray1<'_, i32>>,
3569 bin_size: usize,
3570) -> Py<PyArray2<f64>> {
3571 let vecs: Vec<Vec<i32>> = trains
3572 .iter()
3573 .map(|t| t.as_slice().unwrap().to_vec())
3574 .collect();
3575 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3576 let mat = analysis::correlation::noise_correlation(&refs, bin_size);
3577 let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3578 let n = rows.len();
3579 numpy::PyArray2::from_vec2(py, &rows)
3580 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3581 .into()
3582}
3583
3584#[pyfunction]
3585#[pyo3(signature = (trains, bin_size=50))]
3586fn py_signal_correlation(
3587 py: Python<'_>,
3588 trains: Vec<PyReadonlyArray1<'_, i32>>,
3589 bin_size: usize,
3590) -> Py<PyArray2<f64>> {
3591 let vecs: Vec<Vec<i32>> = trains
3592 .iter()
3593 .map(|t| t.as_slice().unwrap().to_vec())
3594 .collect();
3595 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3596 let mat = analysis::correlation::signal_correlation(&refs, bin_size);
3597 let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3598 let n = rows.len();
3599 numpy::PyArray2::from_vec2(py, &rows)
3600 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3601 .into()
3602}
3603
3604#[pyfunction]
3605#[pyo3(signature = (trains, window=50))]
3606fn py_spike_count_covariance(
3607 py: Python<'_>,
3608 trains: Vec<PyReadonlyArray1<'_, i32>>,
3609 window: usize,
3610) -> Py<PyArray2<f64>> {
3611 let vecs: Vec<Vec<i32>> = trains
3612 .iter()
3613 .map(|t| t.as_slice().unwrap().to_vec())
3614 .collect();
3615 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3616 let mat = analysis::correlation::spike_count_covariance(&refs, window);
3617 let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3618 let n = rows.len();
3619 numpy::PyArray2::from_vec2(py, &rows)
3620 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3621 .into()
3622}
3623
3624#[pyfunction]
3625#[pyo3(signature = (train_a, train_b, bin_size=10))]
3626fn py_joint_psth(
3627 py: Python<'_>,
3628 train_a: PyReadonlyArray1<'_, i32>,
3629 train_b: PyReadonlyArray1<'_, i32>,
3630 bin_size: usize,
3631) -> Py<PyArray2<f64>> {
3632 let a = train_a.as_slice().unwrap();
3633 let b = train_b.as_slice().unwrap();
3634 let (flat, n) = analysis::correlation::joint_psth(a, b, bin_size);
3635 if n == 0 {
3636 return numpy::PyArray2::zeros(py, [0, 0], false).into();
3637 }
3638 let rows: Vec<Vec<f64>> = flat.chunks(n).map(|c| c.to_vec()).collect();
3639 numpy::PyArray2::from_vec2(py, &rows).unwrap().into()
3640}
3641
3642#[pyfunction]
3643#[pyo3(signature = (train_a, train_b, dt=0.001, delta_ms=2.0))]
3644fn py_coincidence_index(
3645 train_a: PyReadonlyArray1<'_, i32>,
3646 train_b: PyReadonlyArray1<'_, i32>,
3647 dt: f64,
3648 delta_ms: f64,
3649) -> f64 {
3650 let a = train_a.as_slice().unwrap();
3651 let b = train_b.as_slice().unwrap();
3652 analysis::correlation::coincidence_index(a, b, dt, delta_ms)
3653}
3654
3655#[pyfunction]
3658#[pyo3(signature = (train_a, train_b, dt=0.001, tau_ms=10.0))]
3659fn py_van_rossum_distance(
3660 train_a: PyReadonlyArray1<'_, i32>,
3661 train_b: PyReadonlyArray1<'_, i32>,
3662 dt: f64,
3663 tau_ms: f64,
3664) -> f64 {
3665 let a = train_a.as_slice().unwrap();
3666 let b = train_b.as_slice().unwrap();
3667 analysis::distance::van_rossum_distance(a, b, dt, tau_ms)
3668}
3669
3670#[pyfunction]
3671#[pyo3(signature = (times_a, times_b, cost_per_s=1000.0))]
3672fn py_victor_purpura_distance(
3673 times_a: PyReadonlyArray1<'_, f64>,
3674 times_b: PyReadonlyArray1<'_, f64>,
3675 cost_per_s: f64,
3676) -> f64 {
3677 let a = times_a.as_slice().unwrap();
3678 let b = times_b.as_slice().unwrap();
3679 analysis::distance::victor_purpura_distance(a, b, cost_per_s)
3680}
3681
3682#[pyfunction]
3683#[pyo3(signature = (train_a, train_b, dt=0.001))]
3684fn py_isi_distance(
3685 train_a: PyReadonlyArray1<'_, i32>,
3686 train_b: PyReadonlyArray1<'_, i32>,
3687 dt: f64,
3688) -> f64 {
3689 let a = train_a.as_slice().unwrap();
3690 let b = train_b.as_slice().unwrap();
3691 analysis::distance::isi_distance(a, b, dt)
3692}
3693
3694#[pyfunction]
3695#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0))]
3696fn py_spike_distance(
3697 times_a: PyReadonlyArray1<'_, f64>,
3698 times_b: PyReadonlyArray1<'_, f64>,
3699 t_start: f64,
3700 t_end: f64,
3701) -> f64 {
3702 let a = times_a.as_slice().unwrap();
3703 let b = times_b.as_slice().unwrap();
3704 analysis::distance::spike_distance(a, b, t_start, t_end)
3705}
3706
3707#[pyfunction]
3708#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0))]
3709fn py_spike_sync(
3710 times_a: PyReadonlyArray1<'_, f64>,
3711 times_b: PyReadonlyArray1<'_, f64>,
3712 t_start: f64,
3713 t_end: f64,
3714) -> f64 {
3715 let a = times_a.as_slice().unwrap();
3716 let b = times_b.as_slice().unwrap();
3717 analysis::distance::spike_sync(a, b, t_start, t_end)
3718}
3719
3720#[pyfunction]
3721#[pyo3(signature = (times_a, times_b, n_bins=50, t_start=0.0, t_end=1.0))]
3722fn py_spike_sync_profile(
3723 py: Python<'_>,
3724 times_a: PyReadonlyArray1<'_, f64>,
3725 times_b: PyReadonlyArray1<'_, f64>,
3726 n_bins: usize,
3727 t_start: f64,
3728 t_end: f64,
3729) -> Py<PyArray1<f64>> {
3730 let a = times_a.as_slice().unwrap();
3731 let b = times_b.as_slice().unwrap();
3732 analysis::distance::spike_sync_profile(a, b, n_bins, t_start, t_end)
3733 .into_pyarray(py)
3734 .into()
3735}
3736
3737#[pyfunction]
3738#[pyo3(signature = (times_a, times_b, n_bins=50, t_start=0.0, t_end=1.0))]
3739fn py_spike_profile(
3740 py: Python<'_>,
3741 times_a: PyReadonlyArray1<'_, f64>,
3742 times_b: PyReadonlyArray1<'_, f64>,
3743 n_bins: usize,
3744 t_start: f64,
3745 t_end: f64,
3746) -> Py<PyArray1<f64>> {
3747 let a = times_a.as_slice().unwrap();
3748 let b = times_b.as_slice().unwrap();
3749 analysis::distance::spike_profile(a, b, n_bins, t_start, t_end)
3750 .into_pyarray(py)
3751 .into()
3752}
3753
3754#[pyfunction]
3755#[pyo3(signature = (binary_train_a, binary_train_b, dt=0.001, n_bins=50))]
3756fn py_isi_profile(
3757 py: Python<'_>,
3758 binary_train_a: PyReadonlyArray1<'_, i32>,
3759 binary_train_b: PyReadonlyArray1<'_, i32>,
3760 dt: f64,
3761 n_bins: usize,
3762) -> Py<PyArray1<f64>> {
3763 let a = binary_train_a.as_slice().unwrap();
3764 let b = binary_train_b.as_slice().unwrap();
3765 analysis::distance::isi_profile(a, b, dt, n_bins)
3766 .into_pyarray(py)
3767 .into()
3768}
3769
3770#[pyfunction]
3771#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0, cost=0.0))]
3772fn py_adaptive_spike_distance(
3773 times_a: PyReadonlyArray1<'_, f64>,
3774 times_b: PyReadonlyArray1<'_, f64>,
3775 t_start: f64,
3776 t_end: f64,
3777 cost: f64,
3778) -> f64 {
3779 let a = times_a.as_slice().unwrap();
3780 let b = times_b.as_slice().unwrap();
3781 analysis::distance::adaptive_spike_distance(a, b, t_start, t_end, cost)
3782}
3783
3784#[pyfunction]
3785#[pyo3(signature = (train_a, train_b, dt=0.001, sigma_ms=5.0))]
3786fn py_schreiber_similarity(
3787 train_a: PyReadonlyArray1<'_, i32>,
3788 train_b: PyReadonlyArray1<'_, i32>,
3789 dt: f64,
3790 sigma_ms: f64,
3791) -> f64 {
3792 let a = train_a.as_slice().unwrap();
3793 let b = train_b.as_slice().unwrap();
3794 analysis::distance::schreiber_similarity(a, b, dt, sigma_ms)
3795}
3796
3797#[pyfunction]
3798#[pyo3(signature = (times_a, times_b, dt_max=0.01))]
3799fn py_hunter_milton_similarity(
3800 times_a: PyReadonlyArray1<'_, f64>,
3801 times_b: PyReadonlyArray1<'_, f64>,
3802 dt_max: f64,
3803) -> f64 {
3804 let a = times_a.as_slice().unwrap();
3805 let b = times_b.as_slice().unwrap();
3806 analysis::distance::hunter_milton_similarity(a, b, dt_max)
3807}
3808
3809#[pyfunction]
3810#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0, n_bins=100))]
3811fn py_earth_movers_distance(
3812 times_a: PyReadonlyArray1<'_, f64>,
3813 times_b: PyReadonlyArray1<'_, f64>,
3814 t_start: f64,
3815 t_end: f64,
3816 n_bins: usize,
3817) -> f64 {
3818 let a = times_a.as_slice().unwrap();
3819 let b = times_b.as_slice().unwrap();
3820 analysis::distance::earth_movers_distance(a, b, t_start, t_end, n_bins)
3821}
3822
3823#[pyfunction]
3824#[pyo3(signature = (spike_times_list, cost_per_s=1000.0))]
3825fn py_multi_neuron_victor_purpura(
3826 py: Python<'_>,
3827 spike_times_list: Vec<PyReadonlyArray1<'_, f64>>,
3828 cost_per_s: f64,
3829) -> Py<PyArray2<f64>> {
3830 let vecs: Vec<Vec<f64>> = spike_times_list
3831 .iter()
3832 .map(|t| t.as_slice().unwrap().to_vec())
3833 .collect();
3834 let refs: Vec<&[f64]> = vecs.iter().map(|v| v.as_slice()).collect();
3835 let mat = analysis::distance::multi_neuron_victor_purpura(&refs, cost_per_s);
3836 let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3837 let n = rows.len();
3838 numpy::PyArray2::from_vec2(py, &rows)
3839 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3840 .into()
3841}
3842
3843#[pyfunction]
3844#[pyo3(signature = (spike_times_list, metric="spike_distance", t_start=0.0, t_end=1.0))]
3845fn py_spike_distance_matrix(
3846 py: Python<'_>,
3847 spike_times_list: Vec<PyReadonlyArray1<'_, f64>>,
3848 metric: &str,
3849 t_start: f64,
3850 t_end: f64,
3851) -> Py<PyArray2<f64>> {
3852 let vecs: Vec<Vec<f64>> = spike_times_list
3853 .iter()
3854 .map(|t| t.as_slice().unwrap().to_vec())
3855 .collect();
3856 let refs: Vec<&[f64]> = vecs.iter().map(|v| v.as_slice()).collect();
3857 let mat = analysis::distance::spike_distance_matrix(&refs, metric, t_start, t_end);
3858 let rows: Vec<Vec<f64>> = mat.into_iter().collect();
3859 let n = rows.len();
3860 numpy::PyArray2::from_vec2(py, &rows)
3861 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
3862 .into()
3863}
3864
3865#[pyfunction]
3868#[pyo3(signature = (train_a, train_b, bin_size=10))]
3869fn py_mutual_information(
3870 train_a: PyReadonlyArray1<'_, i32>,
3871 train_b: PyReadonlyArray1<'_, i32>,
3872 bin_size: usize,
3873) -> f64 {
3874 let a = train_a.as_slice().unwrap();
3875 let b = train_b.as_slice().unwrap();
3876 analysis::information::mutual_information(a, b, bin_size)
3877}
3878
3879#[pyfunction]
3880#[pyo3(signature = (source, target, bin_size=10, lag=1))]
3881fn py_transfer_entropy(
3882 source: PyReadonlyArray1<'_, i32>,
3883 target: PyReadonlyArray1<'_, i32>,
3884 bin_size: usize,
3885 lag: usize,
3886) -> f64 {
3887 let s = source.as_slice().unwrap();
3888 let t = target.as_slice().unwrap();
3889 analysis::information::transfer_entropy(s, t, bin_size, lag)
3890}
3891
3892#[pyfunction]
3893#[pyo3(signature = (binary_train, bin_size=10, word_length=4))]
3894fn py_spike_train_entropy(
3895 binary_train: PyReadonlyArray1<'_, i32>,
3896 bin_size: usize,
3897 word_length: usize,
3898) -> f64 {
3899 let data = binary_train.as_slice().unwrap();
3900 analysis::information::spike_train_entropy(data, bin_size, word_length)
3901}
3902
3903#[pyfunction]
3904#[pyo3(signature = (binary_train, n_trials=10, bin_size=10, word_length=4))]
3905fn py_noise_entropy(
3906 binary_train: PyReadonlyArray1<'_, i32>,
3907 n_trials: usize,
3908 bin_size: usize,
3909 word_length: usize,
3910) -> f64 {
3911 let data = binary_train.as_slice().unwrap();
3912 analysis::information::noise_entropy(data, n_trials, bin_size, word_length)
3913}
3914
3915#[pyfunction]
3916fn py_stimulus_specific_information(
3917 spike_counts: PyReadonlyArray1<'_, f64>,
3918 stimulus_ids: PyReadonlyArray1<'_, i64>,
3919) -> f64 {
3920 let counts = spike_counts.as_slice().unwrap();
3921 let ids = stimulus_ids.as_slice().unwrap();
3922 analysis::information::stimulus_specific_information(counts, ids)
3923}
3924
3925#[pyfunction]
3926#[pyo3(signature = (x, y, k=3))]
3927fn py_kozachenko_leonenko_mi(
3928 x: PyReadonlyArray1<'_, f64>,
3929 y: PyReadonlyArray1<'_, f64>,
3930 k: usize,
3931) -> f64 {
3932 let xd = x.as_slice().unwrap();
3933 let yd = y.as_slice().unwrap();
3934 analysis::information::kozachenko_leonenko_mi(xd, yd, k)
3935}
3936
3937#[pyfunction]
3940#[pyo3(signature = (source, target, bin_size=10, order=5))]
3941fn py_pairwise_granger_causality(
3942 source: PyReadonlyArray1<'_, i32>,
3943 target: PyReadonlyArray1<'_, i32>,
3944 bin_size: usize,
3945 order: usize,
3946) -> f64 {
3947 let s = source.as_slice().unwrap();
3948 let t = target.as_slice().unwrap();
3949 analysis::causality::pairwise_granger_causality(s, t, bin_size, order)
3950}
3951
3952#[pyfunction]
3953#[pyo3(signature = (source, target, condition, bin_size=10, order=5))]
3954fn py_conditional_granger_causality(
3955 source: PyReadonlyArray1<'_, i32>,
3956 target: PyReadonlyArray1<'_, i32>,
3957 condition: PyReadonlyArray1<'_, i32>,
3958 bin_size: usize,
3959 order: usize,
3960) -> f64 {
3961 let s = source.as_slice().unwrap();
3962 let t = target.as_slice().unwrap();
3963 let c = condition.as_slice().unwrap();
3964 analysis::causality::conditional_granger_causality(s, t, c, bin_size, order)
3965}
3966
3967#[pyfunction]
3968#[pyo3(signature = (trains, bin_size=10, order=5, n_freqs=64))]
3969fn py_spectral_granger_causality(
3970 py: Python<'_>,
3971 trains: Vec<PyReadonlyArray1<'_, i32>>,
3972 bin_size: usize,
3973 order: usize,
3974 n_freqs: usize,
3975) -> (Py<PyArray1<f64>>, usize, usize) {
3976 let vecs: Vec<Vec<i32>> = trains
3977 .iter()
3978 .map(|t| t.as_slice().unwrap().to_vec())
3979 .collect();
3980 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3981 let (gc, d) = analysis::causality::spectral_granger_causality(&refs, bin_size, order, n_freqs);
3982 (gc.into_pyarray(py).into(), d, n_freqs)
3983}
3984
3985#[pyfunction]
3986#[pyo3(signature = (trains, bin_size=10, order=5, n_freqs=64))]
3987fn py_partial_directed_coherence(
3988 py: Python<'_>,
3989 trains: Vec<PyReadonlyArray1<'_, i32>>,
3990 bin_size: usize,
3991 order: usize,
3992 n_freqs: usize,
3993) -> (Py<PyArray1<f64>>, usize, usize) {
3994 let vecs: Vec<Vec<i32>> = trains
3995 .iter()
3996 .map(|t| t.as_slice().unwrap().to_vec())
3997 .collect();
3998 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
3999 let (pdc, d) = analysis::causality::partial_directed_coherence(&refs, bin_size, order, n_freqs);
4000 (pdc.into_pyarray(py).into(), d, n_freqs)
4001}
4002
4003#[pyfunction]
4004#[pyo3(signature = (trains, bin_size=10, order=5, n_freqs=64))]
4005fn py_directed_transfer_function(
4006 py: Python<'_>,
4007 trains: Vec<PyReadonlyArray1<'_, i32>>,
4008 bin_size: usize,
4009 order: usize,
4010 n_freqs: usize,
4011) -> (Py<PyArray1<f64>>, usize, usize) {
4012 let vecs: Vec<Vec<i32>> = trains
4013 .iter()
4014 .map(|t| t.as_slice().unwrap().to_vec())
4015 .collect();
4016 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4017 let (dtf, d) = analysis::causality::directed_transfer_function(&refs, bin_size, order, n_freqs);
4018 (dtf.into_pyarray(py).into(), d, n_freqs)
4019}
4020
4021#[pyfunction]
4024#[pyo3(signature = (trains, preferred_directions, window=50))]
4025fn py_population_vector_decode(
4026 py: Python<'_>,
4027 trains: Vec<PyReadonlyArray1<'_, i32>>,
4028 preferred_directions: PyReadonlyArray1<'_, f64>,
4029 window: usize,
4030) -> Py<PyArray1<f64>> {
4031 let vecs: Vec<Vec<i32>> = trains
4032 .iter()
4033 .map(|t| t.as_slice().unwrap().to_vec())
4034 .collect();
4035 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4036 let dirs = preferred_directions.as_slice().unwrap();
4037 analysis::decoding::population_vector_decode(&refs, dirs, window)
4038 .into_pyarray(py)
4039 .into()
4040}
4041
4042#[pyfunction]
4043#[pyo3(signature = (spike_counts, tuning_rates, n_stimuli, n_neurons, prior=None))]
4044fn py_bayesian_decode(
4045 spike_counts: PyReadonlyArray1<'_, f64>,
4046 tuning_rates: PyReadonlyArray1<'_, f64>,
4047 n_stimuli: usize,
4048 n_neurons: usize,
4049 prior: Option<PyReadonlyArray1<'_, f64>>,
4050) -> usize {
4051 let counts = spike_counts.as_slice().unwrap();
4052 let rates = tuning_rates.as_slice().unwrap();
4053 let p: Vec<f64> = prior
4054 .map(|p| p.as_slice().unwrap().to_vec())
4055 .unwrap_or_default();
4056 analysis::decoding::bayesian_decode(counts, rates, n_stimuli, n_neurons, &p)
4057}
4058
4059#[pyfunction]
4060fn py_maximum_likelihood_decode(
4061 spike_counts: PyReadonlyArray1<'_, f64>,
4062 tuning_rates: PyReadonlyArray1<'_, f64>,
4063 n_stimuli: usize,
4064 n_neurons: usize,
4065) -> usize {
4066 let counts = spike_counts.as_slice().unwrap();
4067 let rates = tuning_rates.as_slice().unwrap();
4068 analysis::decoding::maximum_likelihood_decode(counts, rates, n_stimuli, n_neurons)
4069}
4070
4071#[pyfunction]
4072fn py_linear_discriminant_decode(
4073 train_data: PyReadonlyArray1<'_, f64>,
4074 n_samples: usize,
4075 n_features: usize,
4076 labels: PyReadonlyArray1<'_, i64>,
4077 test_point: PyReadonlyArray1<'_, f64>,
4078) -> i64 {
4079 let data = train_data.as_slice().unwrap();
4080 let lbl = labels.as_slice().unwrap();
4081 let tp = test_point.as_slice().unwrap();
4082 analysis::decoding::linear_discriminant_decode(data, n_samples, n_features, lbl, tp)
4083}
4084
4085#[pyfunction]
4086fn py_naive_bayes_decode(
4087 train_data: PyReadonlyArray1<'_, f64>,
4088 n_samples: usize,
4089 n_features: usize,
4090 labels: PyReadonlyArray1<'_, i64>,
4091 test_point: PyReadonlyArray1<'_, f64>,
4092) -> i64 {
4093 let data = train_data.as_slice().unwrap();
4094 let lbl = labels.as_slice().unwrap();
4095 let tp = test_point.as_slice().unwrap();
4096 analysis::decoding::naive_bayes_decode(data, n_samples, n_features, lbl, tp)
4097}
4098
4099#[pyfunction]
4102#[pyo3(signature = (trains, dt=1.0))]
4103fn py_tokenise_spikes(
4104 py: Python<'_>,
4105 trains: Vec<PyReadonlyArray1<'_, i32>>,
4106 dt: f64,
4107) -> (Py<PyArray1<i64>>, Py<PyArray1<f64>>) {
4108 let vecs: Vec<Vec<i32>> = trains
4109 .iter()
4110 .map(|t| t.as_slice().unwrap().to_vec())
4111 .collect();
4112 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4113 let tokens = analysis::neural_decoders::tokenise_spikes(&refs, dt);
4114 let uids: Vec<i64> = tokens.iter().map(|t| t.0 as i64).collect();
4115 let times: Vec<f64> = tokens.iter().map(|t| t.1).collect();
4116 (uids.into_pyarray(py).into(), times.into_pyarray(py).into())
4117}
4118
4119#[pyfunction]
4120fn py_sinusoidal_position_encode(
4121 py: Python<'_>,
4122 timestamps: PyReadonlyArray1<'_, f64>,
4123 d_model: usize,
4124) -> Py<PyArray1<f64>> {
4125 let ts = timestamps.as_slice().unwrap();
4126 analysis::neural_decoders::sinusoidal_position_encode(ts, d_model)
4127 .into_pyarray(py)
4128 .into()
4129}
4130
4131#[pyfunction]
4132fn py_scaled_dot_product_attention(
4133 py: Python<'_>,
4134 queries: PyReadonlyArray1<'_, f64>,
4135 keys: PyReadonlyArray1<'_, f64>,
4136 values: PyReadonlyArray1<'_, f64>,
4137 nq: usize,
4138 nk: usize,
4139 d: usize,
4140) -> Py<PyArray1<f64>> {
4141 let q = queries.as_slice().unwrap();
4142 let k = keys.as_slice().unwrap();
4143 let v = values.as_slice().unwrap();
4144 analysis::neural_decoders::scaled_dot_product_attention(q, k, v, nq, nk, d)
4145 .into_pyarray(py)
4146 .into()
4147}
4148
4149#[pyfunction]
4150fn py_gaussian_attention(
4151 py: Python<'_>,
4152 queries: PyReadonlyArray1<'_, f64>,
4153 keys: PyReadonlyArray1<'_, f64>,
4154 values: PyReadonlyArray1<'_, f64>,
4155 nq: usize,
4156 nk: usize,
4157 d: usize,
4158 sigma: f64,
4159) -> Py<PyArray1<f64>> {
4160 let q = queries.as_slice().unwrap();
4161 let k = keys.as_slice().unwrap();
4162 let v = values.as_slice().unwrap();
4163 analysis::neural_decoders::gaussian_attention(q, k, v, nq, nk, d, sigma)
4164 .into_pyarray(py)
4165 .into()
4166}
4167
4168#[pyfunction]
4169fn py_infonce_loss(
4170 anchors: PyReadonlyArray1<'_, f64>,
4171 positives: PyReadonlyArray1<'_, f64>,
4172 n: usize,
4173 d: usize,
4174 temperature: f64,
4175) -> f64 {
4176 let a = anchors.as_slice().unwrap();
4177 let p = positives.as_slice().unwrap();
4178 analysis::neural_decoders::infonce_loss(a, p, n, d, temperature)
4179}
4180
4181#[pyfunction]
4184#[pyo3(signature = (trains, max_lag_ms=20.0, dt=0.001))]
4185fn py_functional_connectivity(
4186 py: Python<'_>,
4187 trains: Vec<PyReadonlyArray1<'_, i32>>,
4188 max_lag_ms: f64,
4189 dt: f64,
4190) -> Py<PyArray2<f64>> {
4191 let vecs: Vec<Vec<i32>> = trains
4192 .iter()
4193 .map(|t| t.as_slice().unwrap().to_vec())
4194 .collect();
4195 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4196 let mat = analysis::network::functional_connectivity(&refs, max_lag_ms, dt);
4197 let n = refs.len();
4198 let rows: Vec<Vec<f64>> = mat.chunks(n).map(|c| c.to_vec()).collect();
4199 numpy::PyArray2::from_vec2(py, &rows)
4200 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
4201 .into()
4202}
4203
4204#[pyfunction]
4205#[pyo3(signature = (trains, bin_size=5, alpha=0.05))]
4206fn py_unitary_events(
4207 py: Python<'_>,
4208 trains: Vec<PyReadonlyArray1<'_, i32>>,
4209 bin_size: usize,
4210 alpha: f64,
4211) -> Py<PyArray1<i64>> {
4212 let vecs: Vec<Vec<i32>> = trains
4213 .iter()
4214 .map(|t| t.as_slice().unwrap().to_vec())
4215 .collect();
4216 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4217 let result = analysis::network::unitary_events(&refs, bin_size, alpha);
4218 let as_i64: Vec<i64> = result.into_iter().map(|v| v as i64).collect();
4219 as_i64.into_pyarray(py).into()
4220}
4221
4222#[pyfunction]
4223#[pyo3(signature = (trains, bin_size=5, threshold=2.0))]
4224fn py_cell_assembly_detection(
4225 trains: Vec<PyReadonlyArray1<'_, i32>>,
4226 bin_size: usize,
4227 threshold: f64,
4228) -> Vec<Vec<usize>> {
4229 let vecs: Vec<Vec<i32>> = trains
4230 .iter()
4231 .map(|t| t.as_slice().unwrap().to_vec())
4232 .collect();
4233 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4234 analysis::network::cell_assembly_detection(&refs, bin_size, threshold)
4235}
4236
4237#[pyfunction]
4238#[pyo3(signature = (trains, dt=0.001, max_delay_ms=20.0, min_chain_length=3))]
4239fn py_synfire_chain_detection(
4240 trains: Vec<PyReadonlyArray1<'_, i32>>,
4241 dt: f64,
4242 max_delay_ms: f64,
4243 min_chain_length: usize,
4244) -> Vec<Vec<usize>> {
4245 let vecs: Vec<Vec<i32>> = trains
4246 .iter()
4247 .map(|t| t.as_slice().unwrap().to_vec())
4248 .collect();
4249 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4250 analysis::network::synfire_chain_detection(&refs, dt, max_delay_ms, min_chain_length)
4251}
4252
4253#[pyfunction]
4256#[pyo3(signature = (binary_train, seed=0))]
4257fn py_surrogate_isi_shuffle(
4258 py: Python<'_>,
4259 binary_train: PyReadonlyArray1<'_, i32>,
4260 seed: u64,
4261) -> Py<PyArray1<i32>> {
4262 let data = binary_train.as_slice().unwrap();
4263 analysis::surrogates::surrogate_isi_shuffle(data, seed)
4264 .into_pyarray(py)
4265 .into()
4266}
4267
4268#[pyfunction]
4269#[pyo3(signature = (binary_train, dither_ms=5.0, dt=0.001, seed=0))]
4270fn py_surrogate_dither(
4271 py: Python<'_>,
4272 binary_train: PyReadonlyArray1<'_, i32>,
4273 dither_ms: f64,
4274 dt: f64,
4275 seed: u64,
4276) -> Py<PyArray1<i32>> {
4277 let data = binary_train.as_slice().unwrap();
4278 analysis::surrogates::surrogate_dither(data, dither_ms, dt, seed)
4279 .into_pyarray(py)
4280 .into()
4281}
4282
4283#[pyfunction]
4284#[pyo3(signature = (rate_hz, duration_s, dt=0.001, seed=0))]
4285fn py_homogeneous_poisson(
4286 py: Python<'_>,
4287 rate_hz: f64,
4288 duration_s: f64,
4289 dt: f64,
4290 seed: u64,
4291) -> Py<PyArray1<f64>> {
4292 analysis::surrogates::homogeneous_poisson(rate_hz, duration_s, dt, seed)
4293 .into_pyarray(py)
4294 .into()
4295}
4296
4297#[pyfunction]
4298#[pyo3(signature = (rate_hz, shape, duration_s, dt=0.001, seed=0))]
4299fn py_gamma_process(
4300 py: Python<'_>,
4301 rate_hz: f64,
4302 shape: f64,
4303 duration_s: f64,
4304 dt: f64,
4305 seed: u64,
4306) -> Py<PyArray1<f64>> {
4307 analysis::surrogates::gamma_process(rate_hz, shape, duration_s, dt, seed)
4308 .into_pyarray(py)
4309 .into()
4310}
4311
4312#[pyfunction]
4313#[pyo3(signature = (rate_hz, burst_mean, duration_s, dt=0.001, seed=0))]
4314fn py_compound_poisson_process(
4315 py: Python<'_>,
4316 rate_hz: f64,
4317 burst_mean: f64,
4318 duration_s: f64,
4319 dt: f64,
4320 seed: u64,
4321) -> Py<PyArray1<f64>> {
4322 analysis::surrogates::compound_poisson_process(rate_hz, burst_mean, duration_s, dt, seed)
4323 .into_pyarray(py)
4324 .into()
4325}
4326
4327#[pyfunction]
4328#[pyo3(signature = (binary_train, seed=0))]
4329fn py_surrogate_joint_isi(
4330 py: Python<'_>,
4331 binary_train: PyReadonlyArray1<'_, i32>,
4332 seed: u64,
4333) -> Py<PyArray1<i32>> {
4334 let data = binary_train.as_slice().unwrap();
4335 analysis::surrogates::surrogate_joint_isi(data, seed)
4336 .into_pyarray(py)
4337 .into()
4338}
4339
4340#[pyfunction]
4341#[pyo3(signature = (binary_train, bin_size=10, seed=0))]
4342fn py_surrogate_bin_shuffling(
4343 py: Python<'_>,
4344 binary_train: PyReadonlyArray1<'_, i32>,
4345 bin_size: usize,
4346 seed: u64,
4347) -> Py<PyArray1<i32>> {
4348 let data = binary_train.as_slice().unwrap();
4349 analysis::surrogates::surrogate_bin_shuffling(data, bin_size, seed)
4350 .into_pyarray(py)
4351 .into()
4352}
4353
4354#[pyfunction]
4355#[pyo3(signature = (binary_train, max_shift=50, seed=0))]
4356fn py_surrogate_spike_train_shifting(
4357 py: Python<'_>,
4358 binary_train: PyReadonlyArray1<'_, i32>,
4359 max_shift: usize,
4360 seed: u64,
4361) -> Py<PyArray1<i32>> {
4362 let data = binary_train.as_slice().unwrap();
4363 analysis::surrogates::surrogate_spike_train_shifting(data, max_shift, seed)
4364 .into_pyarray(py)
4365 .into()
4366}
4367
4368#[pyfunction]
4371#[pyo3(signature = (binary_train, dt=0.001, max_isi_ms=10.0, min_spikes=3))]
4372fn py_burst_detection(
4373 binary_train: PyReadonlyArray1<'_, i32>,
4374 dt: f64,
4375 max_isi_ms: f64,
4376 min_spikes: usize,
4377) -> Vec<(f64, f64, usize)> {
4378 let data = binary_train.as_slice().unwrap();
4379 analysis::temporal::burst_detection(data, dt, max_isi_ms, min_spikes)
4380}
4381
4382#[pyfunction]
4383#[pyo3(signature = (binary_train, dt=0.001))]
4384fn py_first_spike_latency(binary_train: PyReadonlyArray1<'_, i32>, dt: f64) -> f64 {
4385 let data = binary_train.as_slice().unwrap();
4386 analysis::temporal::first_spike_latency(data, dt)
4387}
4388
4389#[pyfunction]
4390#[pyo3(signature = (binary_train, baseline_steps=100, dt=0.001, threshold_sigma=3.0))]
4391fn py_response_onset(
4392 binary_train: PyReadonlyArray1<'_, i32>,
4393 baseline_steps: usize,
4394 dt: f64,
4395 threshold_sigma: f64,
4396) -> f64 {
4397 let data = binary_train.as_slice().unwrap();
4398 analysis::temporal::response_onset(data, baseline_steps, dt, threshold_sigma)
4399}
4400
4401#[pyfunction]
4402#[pyo3(signature = (binary_train, bin_size=50, threshold=3.0))]
4403fn py_change_point_detection(
4404 py: Python<'_>,
4405 binary_train: PyReadonlyArray1<'_, i32>,
4406 bin_size: usize,
4407 threshold: f64,
4408) -> Py<PyArray1<i64>> {
4409 let data = binary_train.as_slice().unwrap();
4410 let cps = analysis::temporal::change_point_detection(data, bin_size, threshold);
4411 let as_i64: Vec<i64> = cps.into_iter().map(|v| v as i64).collect();
4412 as_i64.into_pyarray(py).into()
4413}
4414
4415#[pyfunction]
4418#[pyo3(signature = (times_a, times_b, t_start=0.0, t_end=1.0))]
4419fn py_spike_directionality(
4420 times_a: PyReadonlyArray1<'_, f64>,
4421 times_b: PyReadonlyArray1<'_, f64>,
4422 t_start: f64,
4423 t_end: f64,
4424) -> f64 {
4425 let a = times_a.as_slice().unwrap();
4426 let b = times_b.as_slice().unwrap();
4427 analysis::patterns::spike_directionality(a, b, t_start, t_end)
4428}
4429
4430#[pyfunction]
4431#[pyo3(signature = (times_list, t_start=0.0, t_end=1.0))]
4432fn py_spike_train_order(
4433 py: Python<'_>,
4434 times_list: Vec<PyReadonlyArray1<'_, f64>>,
4435 t_start: f64,
4436 t_end: f64,
4437) -> Py<PyArray2<f64>> {
4438 let vecs: Vec<Vec<f64>> = times_list
4439 .iter()
4440 .map(|t| t.as_slice().unwrap().to_vec())
4441 .collect();
4442 let refs: Vec<&[f64]> = vecs.iter().map(|v| v.as_slice()).collect();
4443 let mat = analysis::patterns::spike_train_order(&refs, t_start, t_end);
4444 let n = refs.len();
4445 let rows: Vec<Vec<f64>> = mat.chunks(n).map(|c| c.to_vec()).collect();
4446 numpy::PyArray2::from_vec2(py, &rows)
4447 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [n, n], false))
4448 .into()
4449}
4450
4451#[pyfunction]
4452#[pyo3(signature = (binary_train, dt=0.001, max_lag=20))]
4453fn py_cubic_higher_order(
4454 py: Python<'_>,
4455 binary_train: PyReadonlyArray1<'_, i32>,
4456 dt: f64,
4457 max_lag: usize,
4458) -> Py<PyArray2<f64>> {
4459 let data = binary_train.as_slice().unwrap();
4460 let c3 = analysis::patterns::cubic_higher_order(data, dt, max_lag);
4461 let rows: Vec<Vec<f64>> = c3.chunks(max_lag).map(|c| c.to_vec()).collect();
4462 numpy::PyArray2::from_vec2(py, &rows)
4463 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [max_lag, max_lag], false))
4464 .into()
4465}
4466
4467#[pyfunction]
4470#[pyo3(signature = (binary_train, dt=0.001))]
4471fn py_power_spectrum(
4472 py: Python<'_>,
4473 binary_train: PyReadonlyArray1<'_, i32>,
4474 dt: f64,
4475) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4476 let data = binary_train.as_slice().unwrap();
4477 let (psd, freqs) = analysis::spectral::power_spectrum(data, dt);
4478 (psd.into_pyarray(py).into(), freqs.into_pyarray(py).into())
4479}
4480
4481#[pyfunction]
4484#[pyo3(signature = (waveform, dt=3.3333333333333335e-05))]
4485fn py_waveform_width(waveform: PyReadonlyArray1<'_, f64>, dt: f64) -> f64 {
4486 analysis::waveform::waveform_width(waveform.as_slice().unwrap(), dt)
4487}
4488
4489#[pyfunction]
4490fn py_waveform_amplitude(waveform: PyReadonlyArray1<'_, f64>) -> f64 {
4491 analysis::waveform::waveform_amplitude(waveform.as_slice().unwrap())
4492}
4493
4494#[pyfunction]
4495#[pyo3(signature = (waveform, dt=3.3333333333333335e-05))]
4496fn py_waveform_repolarization_slope(waveform: PyReadonlyArray1<'_, f64>, dt: f64) -> f64 {
4497 analysis::waveform::waveform_repolarization_slope(waveform.as_slice().unwrap(), dt)
4498}
4499
4500#[pyfunction]
4501#[pyo3(signature = (waveform, dt=3.3333333333333335e-05))]
4502fn py_waveform_recovery_slope(waveform: PyReadonlyArray1<'_, f64>, dt: f64) -> f64 {
4503 analysis::waveform::waveform_recovery_slope(waveform.as_slice().unwrap(), dt)
4504}
4505
4506#[pyfunction]
4507#[pyo3(signature = (waveform, dt=3.3333333333333335e-05))]
4508fn py_waveform_halfwidth(waveform: PyReadonlyArray1<'_, f64>, dt: f64) -> f64 {
4509 analysis::waveform::waveform_halfwidth(waveform.as_slice().unwrap(), dt)
4510}
4511
4512#[pyfunction]
4513fn py_waveform_pt_ratio(waveform: PyReadonlyArray1<'_, f64>) -> f64 {
4514 analysis::waveform::waveform_pt_ratio(waveform.as_slice().unwrap())
4515}
4516
4517#[pyfunction]
4520#[pyo3(signature = (binary_train, dt=0.001, window_ms=50.0))]
4521fn py_conditional_intensity(
4522 py: Python<'_>,
4523 binary_train: PyReadonlyArray1<'_, i32>,
4524 dt: f64,
4525 window_ms: f64,
4526) -> Py<PyArray1<f64>> {
4527 let data = binary_train.as_slice().unwrap();
4528 analysis::point_process::conditional_intensity(data, dt, window_ms)
4529 .into_pyarray(py)
4530 .into()
4531}
4532
4533#[pyfunction]
4534#[pyo3(signature = (binary_train, dt=0.001, bins=30))]
4535fn py_isi_hazard_function(
4536 py: Python<'_>,
4537 binary_train: PyReadonlyArray1<'_, i32>,
4538 dt: f64,
4539 bins: usize,
4540) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4541 let data = binary_train.as_slice().unwrap();
4542 let (hazard, centres) = analysis::point_process::isi_hazard_function(data, dt, bins);
4543 (
4544 hazard.into_pyarray(py).into(),
4545 centres.into_pyarray(py).into(),
4546 )
4547}
4548
4549#[pyfunction]
4550#[pyo3(signature = (binary_train, dt=0.001, bins=30))]
4551fn py_isi_survivor_function(
4552 py: Python<'_>,
4553 binary_train: PyReadonlyArray1<'_, i32>,
4554 dt: f64,
4555 bins: usize,
4556) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4557 let data = binary_train.as_slice().unwrap();
4558 let (surv, centres) = analysis::point_process::isi_survivor_function(data, dt, bins);
4559 (
4560 surv.into_pyarray(py).into(),
4561 centres.into_pyarray(py).into(),
4562 )
4563}
4564
4565#[pyfunction]
4566#[pyo3(signature = (binary_train, dt=0.001, bins=30))]
4567fn py_renewal_density(
4568 py: Python<'_>,
4569 binary_train: PyReadonlyArray1<'_, i32>,
4570 dt: f64,
4571 bins: usize,
4572) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4573 let data = binary_train.as_slice().unwrap();
4574 let (dens, centres) = analysis::point_process::renewal_density(data, dt, bins);
4575 (
4576 dens.into_pyarray(py).into(),
4577 centres.into_pyarray(py).into(),
4578 )
4579}
4580
4581#[pyfunction]
4584#[pyo3(signature = (stimulus, binary_train, window_steps=50))]
4585fn py_spike_triggered_average(
4586 py: Python<'_>,
4587 stimulus: PyReadonlyArray1<'_, f64>,
4588 binary_train: PyReadonlyArray1<'_, i32>,
4589 window_steps: usize,
4590) -> Py<PyArray1<f64>> {
4591 analysis::stimulus::spike_triggered_average(
4592 stimulus.as_slice().unwrap(),
4593 binary_train.as_slice().unwrap(),
4594 window_steps,
4595 )
4596 .into_pyarray(py)
4597 .into()
4598}
4599
4600#[pyfunction]
4601#[pyo3(signature = (stimulus, binary_train, window_steps=50))]
4602fn py_spike_triggered_covariance(
4603 py: Python<'_>,
4604 stimulus: PyReadonlyArray1<'_, f64>,
4605 binary_train: PyReadonlyArray1<'_, i32>,
4606 window_steps: usize,
4607) -> Py<PyArray2<f64>> {
4608 let cov = analysis::stimulus::spike_triggered_covariance(
4609 stimulus.as_slice().unwrap(),
4610 binary_train.as_slice().unwrap(),
4611 window_steps,
4612 );
4613 let rows: Vec<Vec<f64>> = cov.chunks(window_steps).map(|c| c.to_vec()).collect();
4614 numpy::PyArray2::from_vec2(py, &rows)
4615 .unwrap_or_else(|_| numpy::PyArray2::zeros(py, [window_steps, window_steps], false))
4616 .into()
4617}
4618
4619#[pyfunction]
4620#[pyo3(signature = (binary_train, positions, n_bins=20, dt=0.001))]
4621fn py_spatial_information(
4622 binary_train: PyReadonlyArray1<'_, i32>,
4623 positions: PyReadonlyArray1<'_, f64>,
4624 n_bins: usize,
4625 dt: f64,
4626) -> f64 {
4627 analysis::stimulus::spatial_information(
4628 binary_train.as_slice().unwrap(),
4629 positions.as_slice().unwrap(),
4630 n_bins,
4631 dt,
4632 )
4633}
4634
4635#[pyfunction]
4636#[pyo3(signature = (binary_train, positions, n_bins=50, threshold_std=2.0, dt=0.001))]
4637fn py_place_field_detection(
4638 binary_train: PyReadonlyArray1<'_, i32>,
4639 positions: PyReadonlyArray1<'_, f64>,
4640 n_bins: usize,
4641 threshold_std: f64,
4642 dt: f64,
4643) -> Vec<(f64, f64)> {
4644 analysis::stimulus::place_field_detection(
4645 binary_train.as_slice().unwrap(),
4646 positions.as_slice().unwrap(),
4647 n_bins,
4648 threshold_std,
4649 dt,
4650 )
4651}
4652
4653#[pyfunction]
4654#[pyo3(signature = (binary_train, stimulus_values, n_bins=20, dt=0.001))]
4655fn py_tuning_curve(
4656 py: Python<'_>,
4657 binary_train: PyReadonlyArray1<'_, i32>,
4658 stimulus_values: PyReadonlyArray1<'_, f64>,
4659 n_bins: usize,
4660 dt: f64,
4661) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4662 let (rates, centres) = analysis::stimulus::tuning_curve(
4663 binary_train.as_slice().unwrap(),
4664 stimulus_values.as_slice().unwrap(),
4665 n_bins,
4666 dt,
4667 );
4668 (
4669 rates.into_pyarray(py).into(),
4670 centres.into_pyarray(py).into(),
4671 )
4672}
4673
4674#[pyfunction]
4677fn py_phase_locking_value(
4678 binary_train: PyReadonlyArray1<'_, i32>,
4679 lfp_signal: PyReadonlyArray1<'_, f64>,
4680) -> f64 {
4681 analysis::lfp::phase_locking_value(
4682 binary_train.as_slice().unwrap(),
4683 lfp_signal.as_slice().unwrap(),
4684 )
4685}
4686
4687#[pyfunction]
4688#[pyo3(signature = (binary_train, lfp_signal, dt=0.001))]
4689fn py_spike_field_coherence(
4690 py: Python<'_>,
4691 binary_train: PyReadonlyArray1<'_, i32>,
4692 lfp_signal: PyReadonlyArray1<'_, f64>,
4693 dt: f64,
4694) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4695 let (sfc, freqs) = analysis::lfp::spike_field_coherence(
4696 binary_train.as_slice().unwrap(),
4697 lfp_signal.as_slice().unwrap(),
4698 dt,
4699 );
4700 (sfc.into_pyarray(py).into(), freqs.into_pyarray(py).into())
4701}
4702
4703#[pyfunction]
4704#[pyo3(signature = (binary_train, lfp_signal, n_bins=36))]
4705fn py_spike_phase_histogram(
4706 py: Python<'_>,
4707 binary_train: PyReadonlyArray1<'_, i32>,
4708 lfp_signal: PyReadonlyArray1<'_, f64>,
4709 n_bins: usize,
4710) -> (Py<PyArray1<i64>>, Py<PyArray1<f64>>) {
4711 let (hist, centres) = analysis::lfp::spike_phase_histogram(
4712 binary_train.as_slice().unwrap(),
4713 lfp_signal.as_slice().unwrap(),
4714 n_bins,
4715 );
4716 (
4717 hist.into_pyarray(py).into(),
4718 centres.into_pyarray(py).into(),
4719 )
4720}
4721
4722#[pyfunction]
4725fn py_isolation_distance(
4726 cluster: PyReadonlyArray2<'_, f64>,
4727 noise: PyReadonlyArray2<'_, f64>,
4728) -> f64 {
4729 let c_shape = cluster.shape();
4730 let n_shape = noise.shape();
4731 let d = c_shape[1];
4732 let c_data: Vec<f64> = cluster.as_slice().unwrap().to_vec();
4733 let n_data: Vec<f64> = noise.as_slice().unwrap().to_vec();
4734 analysis::sorting_quality::isolation_distance(&c_data, c_shape[0], &n_data, n_shape[0], d)
4735}
4736
4737#[pyfunction]
4738fn py_l_ratio(cluster: PyReadonlyArray2<'_, f64>, noise: PyReadonlyArray2<'_, f64>) -> f64 {
4739 let c_shape = cluster.shape();
4740 let n_shape = noise.shape();
4741 let d = c_shape[1];
4742 let c_data: Vec<f64> = cluster.as_slice().unwrap().to_vec();
4743 let n_data: Vec<f64> = noise.as_slice().unwrap().to_vec();
4744 analysis::sorting_quality::l_ratio(&c_data, c_shape[0], &n_data, n_shape[0], d)
4745}
4746
4747#[pyfunction]
4748fn py_silhouette_score(
4749 features: PyReadonlyArray2<'_, f64>,
4750 labels: PyReadonlyArray1<'_, i64>,
4751) -> f64 {
4752 let shape = features.shape();
4753 let f_data: Vec<f64> = features.as_slice().unwrap().to_vec();
4754 let l_data: Vec<i64> = labels.as_slice().unwrap().to_vec();
4755 analysis::sorting_quality::silhouette_score(&f_data, shape[0], shape[1], &l_data)
4756}
4757
4758#[pyfunction]
4759fn py_d_prime(cluster_a: PyReadonlyArray2<'_, f64>, cluster_b: PyReadonlyArray2<'_, f64>) -> f64 {
4760 let a_shape = cluster_a.shape();
4761 let b_shape = cluster_b.shape();
4762 let d = a_shape[1];
4763 let a_data: Vec<f64> = cluster_a.as_slice().unwrap().to_vec();
4764 let b_data: Vec<f64> = cluster_b.as_slice().unwrap().to_vec();
4765 analysis::sorting_quality::d_prime(&a_data, a_shape[0], &b_data, b_shape[0], d)
4766}
4767
4768#[pyfunction]
4769#[pyo3(signature = (binary_train, dt=0.001, refractory_ms=1.5))]
4770fn py_isi_violation_rate(
4771 binary_train: PyReadonlyArray1<'_, i32>,
4772 dt: f64,
4773 refractory_ms: f64,
4774) -> f64 {
4775 analysis::sorting_quality::isi_violation_rate(
4776 binary_train.as_slice().unwrap(),
4777 dt,
4778 refractory_ms,
4779 )
4780}
4781
4782#[pyfunction]
4783#[pyo3(signature = (binary_train, n_bins=100))]
4784fn py_presence_ratio(binary_train: PyReadonlyArray1<'_, i32>, n_bins: usize) -> f64 {
4785 analysis::sorting_quality::presence_ratio(binary_train.as_slice().unwrap(), n_bins)
4786}
4787
4788#[pyfunction]
4789#[pyo3(signature = (amplitudes, bins=100))]
4790fn py_amplitude_cutoff(amplitudes: PyReadonlyArray1<'_, f64>, bins: usize) -> f64 {
4791 analysis::sorting_quality::amplitude_cutoff(amplitudes.as_slice().unwrap(), bins)
4792}
4793
4794#[pyfunction]
4795fn py_snr(waveforms: PyReadonlyArray2<'_, f64>) -> f64 {
4796 let shape = waveforms.shape();
4797 let data: Vec<f64> = waveforms.as_slice().unwrap().to_vec();
4798 analysis::sorting_quality::snr(&data, shape[0], shape[1])
4799}
4800
4801#[pyfunction]
4802#[pyo3(signature = (cluster, noise, k=4))]
4803fn py_nn_hit_rate(
4804 cluster: PyReadonlyArray2<'_, f64>,
4805 noise: PyReadonlyArray2<'_, f64>,
4806 k: usize,
4807) -> f64 {
4808 let c_shape = cluster.shape();
4809 let n_shape = noise.shape();
4810 let d = c_shape[1];
4811 let c_data: Vec<f64> = cluster.as_slice().unwrap().to_vec();
4812 let n_data: Vec<f64> = noise.as_slice().unwrap().to_vec();
4813 analysis::sorting_quality::nn_hit_rate(&c_data, c_shape[0], &n_data, n_shape[0], d, k)
4814}
4815
4816#[pyfunction]
4817#[pyo3(signature = (waveforms, timestamps, n_bins=10))]
4818fn py_drift_metric(
4819 waveforms: PyReadonlyArray2<'_, f64>,
4820 timestamps: PyReadonlyArray1<'_, f64>,
4821 n_bins: usize,
4822) -> f64 {
4823 let shape = waveforms.shape();
4824 let data: Vec<f64> = waveforms.as_slice().unwrap().to_vec();
4825 let ts: Vec<f64> = timestamps.as_slice().unwrap().to_vec();
4826 analysis::sorting_quality::drift_metric(&data, shape[0], shape[1], &ts, n_bins)
4827}
4828
4829#[pyfunction]
4832#[pyo3(signature = (trains, n_components=3, bin_size=10))]
4833fn py_spike_train_pca(
4834 py: Python<'_>,
4835 trains: Vec<PyReadonlyArray1<'_, i32>>,
4836 n_components: usize,
4837 bin_size: usize,
4838) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4839 let vecs: Vec<Vec<i32>> = trains
4840 .iter()
4841 .map(|t| t.as_slice().unwrap().to_vec())
4842 .collect();
4843 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4844 let (proj, expl) = analysis::dimensionality::spike_train_pca(&refs, n_components, bin_size);
4845 (proj.into_pyarray(py).into(), expl.into_pyarray(py).into())
4846}
4847
4848#[pyfunction]
4849#[pyo3(signature = (conditions, n_components=3, bin_size=10))]
4850fn py_demixed_pca(
4851 py: Python<'_>,
4852 conditions: Vec<Vec<PyReadonlyArray1<'_, i32>>>,
4853 n_components: usize,
4854 bin_size: usize,
4855) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4856 let vecs: Vec<Vec<Vec<i32>>> = conditions
4857 .iter()
4858 .map(|cond| {
4859 cond.iter()
4860 .map(|t| t.as_slice().unwrap().to_vec())
4861 .collect()
4862 })
4863 .collect();
4864 let refs: Vec<Vec<&[i32]>> = vecs
4865 .iter()
4866 .map(|cond| cond.iter().map(|v| v.as_slice()).collect())
4867 .collect();
4868 let (proj, expl) = analysis::dimensionality::demixed_pca(&refs, n_components, bin_size);
4869 (proj.into_pyarray(py).into(), expl.into_pyarray(py).into())
4870}
4871
4872#[pyfunction]
4873#[pyo3(signature = (trains, n_factors=3, bin_size=10, n_iter=50))]
4874fn py_factor_analysis(
4875 py: Python<'_>,
4876 trains: Vec<PyReadonlyArray1<'_, i32>>,
4877 n_factors: usize,
4878 bin_size: usize,
4879 n_iter: usize,
4880) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4881 let vecs: Vec<Vec<i32>> = trains
4882 .iter()
4883 .map(|t| t.as_slice().unwrap().to_vec())
4884 .collect();
4885 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4886 let (loadings, psi) =
4887 analysis::dimensionality::factor_analysis(&refs, n_factors, bin_size, n_iter);
4888 (
4889 loadings.into_pyarray(py).into(),
4890 psi.into_pyarray(py).into(),
4891 )
4892}
4893
4894#[pyfunction]
4899#[pyo3(signature = (mat, n_components=3))]
4900fn py_pca_components(
4901 py: Python<'_>,
4902 mat: PyReadonlyArray2<'_, f64>,
4903 n_components: usize,
4904) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4905 let shape = mat.shape();
4906 let data: Vec<f64> = mat.as_slice().unwrap().to_vec();
4907 let (proj, expl) =
4908 analysis::dimensionality::pca_from_centered(&data, shape[0], shape[1], n_components);
4909 (proj.into_pyarray(py).into(), expl.into_pyarray(py).into())
4910}
4911
4912#[pyfunction]
4913#[pyo3(signature = (mean_mat, n_components=3))]
4914fn py_demixed_components(
4915 py: Python<'_>,
4916 mean_mat: PyReadonlyArray2<'_, f64>,
4917 n_components: usize,
4918) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4919 let shape = mean_mat.shape();
4920 let data: Vec<f64> = mean_mat.as_slice().unwrap().to_vec();
4921 let (proj, expl) =
4922 analysis::dimensionality::demixed_from_centered(&data, shape[0], shape[1], n_components);
4923 (proj.into_pyarray(py).into(), expl.into_pyarray(py).into())
4924}
4925
4926#[pyfunction]
4927#[pyo3(signature = (mat, n_factors=3, n_iter=50))]
4928fn py_factor_loadings(
4929 py: Python<'_>,
4930 mat: PyReadonlyArray2<'_, f64>,
4931 n_factors: usize,
4932 n_iter: usize,
4933) -> (Py<PyArray1<f64>>, Py<PyArray1<f64>>) {
4934 let shape = mat.shape();
4935 let data: Vec<f64> = mat.as_slice().unwrap().to_vec();
4936 let (loadings, psi) =
4937 analysis::dimensionality::fa_from_centered(&data, shape[0], shape[1], n_factors, n_iter);
4938 (
4939 loadings.into_pyarray(py).into(),
4940 psi.into_pyarray(py).into(),
4941 )
4942}
4943
4944#[pyfunction]
4947#[pyo3(signature = (trains, n_latents=3, bin_ms=20.0, dt=0.001, max_iter=50, tol=1e-4, seed=42))]
4948fn py_gpfa<'py>(
4949 py: Python<'py>,
4950 trains: Vec<PyReadonlyArray1<'py, i32>>,
4951 n_latents: usize,
4952 bin_ms: f64,
4953 dt: f64,
4954 max_iter: usize,
4955 tol: f64,
4956 seed: u64,
4957) -> PyResult<Bound<'py, PyDict>> {
4958 let vecs: Vec<Vec<i32>> = trains
4959 .iter()
4960 .map(|t| t.as_slice().unwrap().to_vec())
4961 .collect();
4962 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
4963 let result = analysis::gpfa::gpfa(&refs, n_latents, bin_ms, dt, max_iter, tol, seed);
4964
4965 let dict = PyDict::new(py);
4966 dict.set_item("trajectories", result.trajectories.into_pyarray(py))?;
4967 dict.set_item("C", result.c.into_pyarray(py))?;
4968 dict.set_item("d", result.d.into_pyarray(py))?;
4969 dict.set_item("R", result.r.into_pyarray(py))?;
4970 dict.set_item("tau", result.tau.into_pyarray(py))?;
4971 dict.set_item("log_likelihoods", result.log_likelihoods.into_pyarray(py))?;
4972 dict.set_item("n_latents", result.n_latents)?;
4973 dict.set_item("n_bins", result.n_bins)?;
4974 dict.set_item("n_neurons", result.n_neurons)?;
4975 Ok(dict)
4976}
4977
4978#[pyfunction]
4984#[pyo3(signature = (y, n_neurons, n_bins, c0, d0, r0_diag, tau, n_latents, max_iter, tol))]
4985#[allow(clippy::too_many_arguments, clippy::type_complexity)]
4986fn py_gpfa_em<'py>(
4987 py: Python<'py>,
4988 y: PyReadonlyArray1<'py, f64>,
4989 n_neurons: usize,
4990 n_bins: usize,
4991 c0: PyReadonlyArray1<'py, f64>,
4992 d0: PyReadonlyArray1<'py, f64>,
4993 r0_diag: PyReadonlyArray1<'py, f64>,
4994 tau: PyReadonlyArray1<'py, f64>,
4995 n_latents: usize,
4996 max_iter: usize,
4997 tol: f64,
4998) -> PyResult<(
4999 Bound<'py, PyArray1<f64>>,
5000 Bound<'py, PyArray1<f64>>,
5001 Bound<'py, PyArray1<f64>>,
5002 Bound<'py, PyArray1<f64>>,
5003 Vec<f64>,
5004)> {
5005 let (x_post, c, d, r, log_liks) = analysis::gpfa::gpfa_em_from_init(
5006 y.as_slice()?,
5007 c0.as_slice()?,
5008 d0.as_slice()?,
5009 r0_diag.as_slice()?,
5010 tau.as_slice()?,
5011 n_neurons,
5012 n_bins,
5013 n_latents,
5014 max_iter,
5015 tol,
5016 );
5017 Ok((
5018 x_post.into_pyarray(py),
5019 c.into_pyarray(py),
5020 d.into_pyarray(py),
5021 r.into_pyarray(py),
5022 log_liks,
5023 ))
5024}
5025
5026#[pyfunction]
5027#[pyo3(signature = (new_trains, c, d, r, tau, n_latents, bin_ms=20.0, dt=0.001))]
5028fn py_gpfa_transform(
5029 py: Python<'_>,
5030 new_trains: Vec<PyReadonlyArray1<'_, i32>>,
5031 c: PyReadonlyArray1<'_, f64>,
5032 d: PyReadonlyArray1<'_, f64>,
5033 r: PyReadonlyArray1<'_, f64>,
5034 tau: PyReadonlyArray1<'_, f64>,
5035 n_latents: usize,
5036 bin_ms: f64,
5037 dt: f64,
5038) -> Py<PyArray1<f64>> {
5039 let vecs: Vec<Vec<i32>> = new_trains
5040 .iter()
5041 .map(|t| t.as_slice().unwrap().to_vec())
5042 .collect();
5043 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
5044 let proj = analysis::gpfa::gpfa_transform(
5045 &refs,
5046 c.as_slice().unwrap(),
5047 d.as_slice().unwrap(),
5048 r.as_slice().unwrap(),
5049 tau.as_slice().unwrap(),
5050 n_latents,
5051 bin_ms,
5052 dt,
5053 );
5054 proj.into_pyarray(py).into()
5055}
5056
5057#[pyfunction]
5060#[pyo3(signature = (trains, bin_ms=5.0, dt=0.001, min_support=3, max_pattern_size=5, n_surrogates=100, alpha=0.05, seed=42))]
5061fn py_spade_detect<'py>(
5062 py: Python<'py>,
5063 trains: Vec<PyReadonlyArray1<'py, i32>>,
5064 bin_ms: f64,
5065 dt: f64,
5066 min_support: usize,
5067 max_pattern_size: usize,
5068 n_surrogates: usize,
5069 alpha: f64,
5070 seed: u64,
5071) -> PyResult<Vec<Bound<'py, PyDict>>> {
5072 let vecs: Vec<Vec<i32>> = trains
5073 .iter()
5074 .map(|t| t.as_slice().unwrap().to_vec())
5075 .collect();
5076 let refs: Vec<&[i32]> = vecs.iter().map(|v| v.as_slice()).collect();
5077 let results = analysis::spade::spade_detect(
5078 &refs,
5079 bin_ms,
5080 dt,
5081 min_support,
5082 max_pattern_size,
5083 n_surrogates,
5084 alpha,
5085 seed,
5086 );
5087 let mut dicts = Vec::new();
5088 for pat in results {
5089 let dict = PyDict::new(py);
5090 dict.set_item(
5091 "neurons",
5092 pat.neurons.iter().map(|&n| n as i64).collect::<Vec<_>>(),
5093 )?;
5094 dict.set_item("lags", pat.lags.clone())?;
5095 dict.set_item("count", pat.count as i64)?;
5096 dict.set_item("p_value", pat.p_value)?;
5097 dicts.push(dict);
5098 }
5099 Ok(dicts)
5100}
5101
5102#[pyfunction]
5105fn py_dna_design_sequence(_py: Python<'_>, length: usize, seed: u64) -> String {
5106 String::from_utf8(dna::design_sequence(length, seed)).unwrap_or_default()
5107}
5108
5109#[pyfunction]
5110fn py_dna_design_orthogonal_set(
5111 _py: Python<'_>,
5112 count: usize,
5113 length: usize,
5114 seed: u64,
5115) -> Vec<String> {
5116 dna::design_orthogonal_set(count, length, seed)
5117 .into_iter()
5118 .map(|s| String::from_utf8(s).unwrap_or_default())
5119 .collect()
5120}
5121
5122#[pyfunction]
5123fn py_dna_check_cross_hybridization<'py>(
5124 py: Python<'py>,
5125 sequences: Vec<String>,
5126 threshold: usize,
5127) -> PyResult<Py<PyAny>> {
5128 let seqs: Vec<Vec<u8>> = sequences.into_iter().map(|s| s.into_bytes()).collect();
5129 let flags = dna::check_cross_hybridization(&seqs, threshold);
5130 let result: Vec<Py<PyAny>> = flags
5131 .into_iter()
5132 .map(|(i, j, score)| {
5133 let d = PyDict::new(py);
5134 d.set_item("strand_a", i).unwrap();
5135 d.set_item("strand_b", j).unwrap();
5136 d.set_item("score", score).unwrap();
5137 d.into_any().unbind()
5138 })
5139 .collect();
5140 Ok(result.into_pyobject(py)?.into())
5141}
5142
5143#[pyfunction]
5144#[pyo3(signature = (
5145 gate_types, gate_inputs, gate_outputs, gate_thresholds, gate_leaks,
5146 input_names, input_concs, duration_s=1800.0, dt=1.0,
5147 k_hyb=3e5, k_disp=1.0, temperature_c=37.0, use_rk4=true
5148))]
5149#[allow(clippy::too_many_arguments)]
5150fn py_dna_simulate_kinetics<'py>(
5151 py: Python<'py>,
5152 gate_types: Vec<String>,
5153 gate_inputs: Vec<Vec<String>>,
5154 gate_outputs: Vec<String>,
5155 gate_thresholds: Vec<f64>,
5156 gate_leaks: Vec<f64>,
5157 input_names: Vec<String>,
5158 input_concs: Vec<f64>,
5159 duration_s: f64,
5160 dt: f64,
5161 k_hyb: f64,
5162 k_disp: f64,
5163 temperature_c: f64,
5164 use_rk4: bool,
5165) -> PyResult<Py<PyAny>> {
5166 let gates: Vec<dna::DnaGateSpec> = gate_types
5167 .iter()
5168 .zip(gate_inputs.iter())
5169 .zip(gate_outputs.iter())
5170 .zip(gate_thresholds.iter())
5171 .zip(gate_leaks.iter())
5172 .map(|((((gt, gi), go), th), lk)| {
5173 let gate_type = match gt.to_uppercase().as_str() {
5174 "AND" => dna::DnaGateType::And,
5175 "OR" => dna::DnaGateType::Or,
5176 "NOT" => dna::DnaGateType::Not,
5177 "THRESHOLD" => dna::DnaGateType::Threshold,
5178 "MUX" => dna::DnaGateType::Mux,
5179 "AMPLIFIER" => dna::DnaGateType::Amplifier,
5180 "BUFFER" => dna::DnaGateType::Buffer,
5181 "NAND" => dna::DnaGateType::Nand,
5182 "XOR" => dna::DnaGateType::Xor,
5183 _ => dna::DnaGateType::And,
5184 };
5185 dna::DnaGateSpec {
5186 gate_type,
5187 input_names: gi.clone(),
5188 output_name: go.clone(),
5189 threshold: *th,
5190 leak_rate: *lk,
5191 }
5192 })
5193 .collect();
5194
5195 let mut inputs = std::collections::HashMap::new();
5196 for (name, conc) in input_names.into_iter().zip(input_concs) {
5197 inputs.insert(name, conc);
5198 }
5199
5200 let config = dna::KineticConfig {
5201 k_hyb,
5202 k_disp,
5203 temperature_c,
5204 max_conc: 200.0,
5205 use_rk4,
5206 };
5207
5208 let result = dna::simulate_kinetics(&gates, &inputs, duration_s, dt, &config);
5209
5210 let dict = PyDict::new(py);
5211 for (key, trace) in result {
5212 dict.set_item(key, trace.into_pyarray(py))?;
5213 }
5214 Ok(dict.into_any().unbind())
5215}
5216
5217#[pyfunction]
5218#[pyo3(signature = (sequence, min_stem=4, min_loop=3))]
5219fn py_dna_detect_hairpins(
5220 _py: Python<'_>,
5221 sequence: &str,
5222 min_stem: usize,
5223 min_loop: usize,
5224) -> Vec<(usize, usize, usize)> {
5225 dna::detect_hairpins(sequence.as_bytes(), min_stem, min_loop)
5226}
5227
5228#[pyfunction]
5238#[pyo3(signature = (h_indices, h_values, j_i, j_j, j_values, spins, offset=0.0))]
5239fn py_qa_ising_energy(
5240 _py: Python<'_>,
5241 h_indices: Vec<usize>,
5242 h_values: Vec<f64>,
5243 j_i: Vec<usize>,
5244 j_j: Vec<usize>,
5245 j_values: Vec<f64>,
5246 spins: Vec<i8>,
5247 offset: f64,
5248) -> f64 {
5249 let h: Vec<(usize, f64)> = h_indices.into_iter().zip(h_values).collect();
5250 let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5251 quantum::ising_energy(&h, &j, &spins, offset)
5252}
5253
5254#[pyfunction]
5256#[pyo3(signature = (h_indices, h_values, j_i, j_j, j_values, configs, offset=0.0))]
5257fn py_qa_batch_ising_energy(
5258 _py: Python<'_>,
5259 h_indices: Vec<usize>,
5260 h_values: Vec<f64>,
5261 j_i: Vec<usize>,
5262 j_j: Vec<usize>,
5263 j_values: Vec<f64>,
5264 configs: Vec<Vec<i8>>,
5265 offset: f64,
5266) -> Vec<f64> {
5267 let h: Vec<(usize, f64)> = h_indices.into_iter().zip(h_values).collect();
5268 let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5269 quantum::batch_ising_energy(&h, &j, &configs, offset)
5270}
5271
5272#[pyfunction]
5276#[pyo3(signature = (h_indices, h_values, j_i, j_j, j_values, n_qubits, offset=0.0, n_sweeps=1000, num_reads=10, beta_start=0.1, beta_end=10.0, seed=42))]
5277fn py_qa_simulated_annealing<'py>(
5278 py: Python<'py>,
5279 h_indices: Vec<usize>,
5280 h_values: Vec<f64>,
5281 j_i: Vec<usize>,
5282 j_j: Vec<usize>,
5283 j_values: Vec<f64>,
5284 n_qubits: usize,
5285 offset: f64,
5286 n_sweeps: usize,
5287 num_reads: usize,
5288 beta_start: f64,
5289 beta_end: f64,
5290 seed: u64,
5291) -> PyResult<Py<PyAny>> {
5292 let h: Vec<(usize, f64)> = h_indices.into_iter().zip(h_values).collect();
5293 let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5294
5295 let (best_spins, best_energy, all_energies, all_samples) = quantum::simulated_annealing(
5296 &h, &j, n_qubits, offset, n_sweeps, num_reads, beta_start, beta_end, seed,
5297 );
5298
5299 let dict = PyDict::new(py);
5300 dict.set_item("best_spins", best_spins)?;
5301 dict.set_item("best_energy", best_energy)?;
5302 dict.set_item("energies", all_energies)?;
5303 dict.set_item("samples", all_samples)?;
5304 dict.set_item("n_sweeps", n_sweeps)?;
5305 dict.set_item("num_reads", num_reads)?;
5306 dict.set_item("backend", "rust")?;
5307 Ok(dict.into_any().unbind())
5308}
5309
5310#[pyfunction]
5312#[allow(clippy::type_complexity)] fn py_qa_gauge_transform(
5314 _py: Python<'_>,
5315 h_indices: Vec<usize>,
5316 h_values: Vec<f64>,
5317 j_i: Vec<usize>,
5318 j_j: Vec<usize>,
5319 j_values: Vec<f64>,
5320 gauge: Vec<i8>,
5321) -> (Vec<(usize, f64)>, Vec<((usize, usize), f64)>) {
5322 let h: Vec<(usize, f64)> = h_indices.into_iter().zip(h_values).collect();
5323 let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5324 quantum::gauge_transform(&h, &j, &gauge)
5325}
5326
5327#[pyfunction]
5329#[pyo3(signature = (n_qubits, n_gauges=10, seed=42))]
5330fn py_qa_generate_gauges(
5331 _py: Python<'_>,
5332 n_qubits: usize,
5333 n_gauges: usize,
5334 seed: u64,
5335) -> Vec<Vec<i8>> {
5336 quantum::generate_gauges(n_qubits, n_gauges, seed)
5337}
5338
5339#[pyfunction]
5341#[pyo3(signature = (n_qubits, j_i, j_j, j_values, max_partition_size=64))]
5342fn py_qa_greedy_partition(
5343 _py: Python<'_>,
5344 n_qubits: usize,
5345 j_i: Vec<usize>,
5346 j_j: Vec<usize>,
5347 j_values: Vec<f64>,
5348 max_partition_size: usize,
5349) -> Vec<Vec<usize>> {
5350 let j: Vec<((usize, usize), f64)> = j_i.into_iter().zip(j_j).zip(j_values).collect();
5351 quantum::greedy_partition(n_qubits, &j, max_partition_size)
5352}
5353
5354#[pyfunction]
5358#[pyo3(signature = (adjacency_flat, n, pitch_um=250.0, loss_db_per_cm=2.0))]
5359fn py_ph_route_waveguides<'py>(
5360 py: Python<'py>,
5361 adjacency_flat: Vec<f64>,
5362 n: usize,
5363 pitch_um: f64,
5364 loss_db_per_cm: f64,
5365) -> PyResult<Py<PyAny>> {
5366 let result = photonic::route_waveguides(&adjacency_flat, n, pitch_um, loss_db_per_cm);
5367
5368 let dict = PyDict::new(py);
5369 let sources: Vec<usize> = result.iter().map(|r| r.source).collect();
5370 let targets: Vec<usize> = result.iter().map(|r| r.target).collect();
5371 let lengths: Vec<f64> = result.iter().map(|r| r.length_um).collect();
5372 let losses: Vec<f64> = result.iter().map(|r| r.loss_db).collect();
5373 let crossings: Vec<usize> = result.iter().map(|r| r.n_crossings).collect();
5374
5375 dict.set_item("sources", sources)?;
5376 dict.set_item("targets", targets)?;
5377 dict.set_item("lengths_um", lengths)?;
5378 dict.set_item("losses_db", losses)?;
5379 dict.set_item("crossings", crossings)?;
5380 dict.set_item("n_segments", result.len())?;
5381 dict.set_item("backend", "rust")?;
5382 Ok(dict.into_any().unbind())
5383}
5384
5385#[pyfunction]
5387fn py_ph_mzi_transfer_matrix(_py: Python<'_>, phase_rad: f64) -> Vec<f64> {
5388 photonic::mzi_transfer_matrix(phase_rad).to_vec()
5389}
5390
5391#[pyfunction]
5393fn py_ph_cascade_mzi(_py: Python<'_>, phases: Vec<f64>) -> Vec<f64> {
5394 photonic::cascade_mzi(&phases).to_vec()
5395}
5396
5397#[pyfunction]
5399#[pyo3(signature = (channel_ids, wavelengths, bandwidths, powers, adjacent_xt_db=-25.0))]
5400fn py_ph_analyze_crosstalk<'py>(
5401 py: Python<'py>,
5402 channel_ids: Vec<usize>,
5403 wavelengths: Vec<f64>,
5404 bandwidths: Vec<f64>,
5405 powers: Vec<f64>,
5406 adjacent_xt_db: f64,
5407) -> PyResult<Py<PyAny>> {
5408 let channels: Vec<(usize, f64, f64, f64)> = channel_ids
5409 .into_iter()
5410 .zip(wavelengths)
5411 .zip(bandwidths)
5412 .zip(powers)
5413 .map(|(((id, wl), bw), p)| (id, wl, bw, p))
5414 .collect();
5415
5416 let result = photonic::analyze_crosstalk(&channels, adjacent_xt_db);
5417
5418 let dict = PyDict::new(py);
5419 let ids: Vec<usize> = result.iter().map(|r| r.channel_id).collect();
5420 let xts: Vec<f64> = result.iter().map(|r| r.crosstalk_db).collect();
5421 let osnrs: Vec<f64> = result.iter().map(|r| r.osnr_db).collect();
5422 let adjs: Vec<usize> = result.iter().map(|r| r.n_adjacent).collect();
5423
5424 dict.set_item("channel_ids", ids)?;
5425 dict.set_item("crosstalk_db", xts)?;
5426 dict.set_item("osnr_db", osnrs)?;
5427 dict.set_item("n_adjacent", adjs)?;
5428 dict.set_item("backend", "rust")?;
5429 Ok(dict.into_any().unbind())
5430}
5431
5432#[pyfunction]
5434#[pyo3(signature = (wg_sources, wg_targets, wg_losses, laser_power_dbm=0.0, detector_sensitivity_dbm=-20.0))]
5435fn py_ph_analyze_power_budget<'py>(
5436 py: Python<'py>,
5437 wg_sources: Vec<usize>,
5438 wg_targets: Vec<usize>,
5439 wg_losses: Vec<f64>,
5440 laser_power_dbm: f64,
5441 detector_sensitivity_dbm: f64,
5442) -> PyResult<Py<PyAny>> {
5443 let wgs: Vec<(usize, usize, f64)> = wg_sources
5444 .into_iter()
5445 .zip(wg_targets)
5446 .zip(wg_losses)
5447 .map(|((s, t), l)| (s, t, l))
5448 .collect();
5449
5450 let result =
5451 photonic::analyze_power_budget(&wgs, &[], laser_power_dbm, detector_sensitivity_dbm);
5452
5453 let dict = PyDict::new(py);
5454 let margins: Vec<f64> = result.iter().map(|r| r.margin_db).collect();
5455 let passed: Vec<bool> = result.iter().map(|r| r.passed).collect();
5456 let total_losses: Vec<f64> = result.iter().map(|r| r.total_loss_db).collect();
5457
5458 dict.set_item("margins_db", margins)?;
5459 dict.set_item("passed", passed)?;
5460 dict.set_item("total_losses_db", total_losses)?;
5461 dict.set_item("n_paths", result.len())?;
5462 dict.set_item("backend", "rust")?;
5463 Ok(dict.into_any().unbind())
5464}
5465
5466#[pyfunction]
5468#[pyo3(signature = (num_waveguides, gap_nm, coupling_length_um, wavelength_nm=1550.0, core_index=3.48, cladding_index=1.45))]
5469fn py_ph_analyze_crosstalk_bank<'py>(
5470 py: Python<'py>,
5471 num_waveguides: usize,
5472 gap_nm: f64,
5473 coupling_length_um: f64,
5474 wavelength_nm: f64,
5475 core_index: f64,
5476 cladding_index: f64,
5477) -> PyResult<Py<PyAny>> {
5478 let r = photonic::analyze_crosstalk_bank(
5479 num_waveguides,
5480 gap_nm,
5481 coupling_length_um,
5482 wavelength_nm,
5483 core_index,
5484 cladding_index,
5485 );
5486 let dict = PyDict::new(py);
5487 dict.set_item("num_waveguides", r.num_waveguides)?;
5488 dict.set_item("num_pairs", r.num_near_pairs + r.num_far_pairs)?;
5489 dict.set_item("num_near_pairs", r.num_near_pairs)?;
5490 dict.set_item("num_far_pairs", r.num_far_pairs)?;
5491 dict.set_item("gap_nm", r.gap_nm)?;
5492 dict.set_item("coupling_length_um", r.coupling_length_um)?;
5493 dict.set_item("adjacent_coupling_ratio", r.adjacent_coupling_ratio)?;
5494 dict.set_item("adjacent_isolation_db", r.adjacent_isolation_db)?;
5495 dict.set_item("next_nearest_coupling_ratio", r.next_nearest_coupling_ratio)?;
5496 dict.set_item("next_nearest_isolation_db", r.next_nearest_isolation_db)?;
5497 dict.set_item("worst_isolation_db", r.worst_isolation_db)?;
5498 dict.set_item("mean_coupling_ratio", r.mean_coupling_ratio)?;
5499 dict.set_item("max_coupling_ratio", r.max_coupling_ratio)?;
5500 dict.set_item("crosstalk_safe", r.crosstalk_safe)?;
5501 dict.set_item("backend", "rust")?;
5502 Ok(dict.into_any().unbind())
5503}
5504
5505#[pyfunction]
5509#[pyo3(signature = (pairs_a, pairs_b, gaps_nm, lengths_um, wavelength_nm=1550.0, core_index=3.48, cladding_index=1.45))]
5510fn py_ph_analyze_crosstalk_pairs<'py>(
5511 py: Python<'py>,
5512 pairs_a: Vec<usize>,
5513 pairs_b: Vec<usize>,
5514 gaps_nm: Vec<f64>,
5515 lengths_um: Vec<f64>,
5516 wavelength_nm: f64,
5517 core_index: f64,
5518 cladding_index: f64,
5519) -> PyResult<Py<PyAny>> {
5520 let n = pairs_a.len();
5521 if pairs_b.len() != n || gaps_nm.len() != n || lengths_um.len() != n {
5522 return Err(pyo3::exceptions::PyValueError::new_err(
5523 "pairs_a, pairs_b, gaps_nm, lengths_um must be equal length",
5524 ));
5525 }
5526 let pairs: Vec<(usize, usize, f64, f64)> = pairs_a
5527 .into_iter()
5528 .zip(pairs_b)
5529 .zip(gaps_nm)
5530 .zip(lengths_um)
5531 .map(|(((a, b), g), l)| (a, b, g, l))
5532 .collect();
5533 let results =
5534 photonic::analyze_crosstalk_pairs(&pairs, wavelength_nm, core_index, cladding_index);
5535
5536 let dict = PyDict::new(py);
5537 let idx_a: Vec<usize> = results.iter().map(|r| r.index_a).collect();
5538 let idx_b: Vec<usize> = results.iter().map(|r| r.index_b).collect();
5539 let gaps: Vec<f64> = results.iter().map(|r| r.gap_nm).collect();
5540 let lens: Vec<f64> = results.iter().map(|r| r.coupling_length_um).collect();
5541 let kappas: Vec<f64> = results
5542 .iter()
5543 .map(|r| r.coupling_coefficient_per_um)
5544 .collect();
5545 let ratios: Vec<f64> = results.iter().map(|r| r.coupling_ratio).collect();
5546 let isos: Vec<f64> = results.iter().map(|r| r.isolation_db).collect();
5547
5548 dict.set_item("pair_a", idx_a)?;
5549 dict.set_item("pair_b", idx_b)?;
5550 dict.set_item("gap_nm", gaps)?;
5551 dict.set_item("coupling_length_um", lens)?;
5552 dict.set_item("coupling_coefficient_per_um", kappas)?;
5553 dict.set_item("coupling_ratio", ratios)?;
5554 dict.set_item("isolation_db", isos)?;
5555 dict.set_item("num_pairs", n)?;
5556 dict.set_item("backend", "rust")?;
5557 Ok(dict.into_any().unbind())
5558}
5559
5560#[pyfunction]
5568#[pyo3(signature = (mac_counts, weights, max_luts, max_power, max_latency=0, t_init=1.0, t_min=0.001, alpha=0.95, max_iter=2000, seed=42))]
5569fn py_opt_sa_search<'py>(
5570 py: Python<'py>,
5571 mac_counts: Vec<i64>,
5572 weights: Vec<f64>,
5573 max_luts: i64,
5574 max_power: f64,
5575 max_latency: i64,
5576 t_init: f64,
5577 t_min: f64,
5578 alpha: f64,
5579 max_iter: usize,
5580 seed: u64,
5581) -> PyResult<Py<PyAny>> {
5582 let candidates: Vec<Vec<optimizer::Candidate>> = mac_counts
5583 .iter()
5584 .map(|&mc| optimizer::generate_candidates(mc))
5585 .collect();
5586
5587 let result = optimizer::simulated_annealing(
5588 &candidates,
5589 &weights,
5590 max_luts,
5591 max_power,
5592 max_latency,
5593 t_init,
5594 t_min,
5595 alpha,
5596 max_iter,
5597 seed,
5598 );
5599
5600 let dict = PyDict::new(py);
5601 match result {
5602 Some(r) => {
5603 let mut luts_list = Vec::new();
5605 let mut power_list = Vec::new();
5606 let mut acc_list = Vec::new();
5607 for (i, &idx) in r.best_config.iter().enumerate() {
5608 let c = &candidates[i][idx];
5609 luts_list.push(c.luts);
5610 power_list.push(c.power);
5611 acc_list.push(c.accuracy);
5612 }
5613
5614 dict.set_item("best_config", r.best_config)?;
5615 dict.set_item("best_score", r.best_score)?;
5616 dict.set_item("pareto_luts", r.pareto_luts)?;
5617 dict.set_item("pareto_power", r.pareto_power)?;
5618 dict.set_item("pareto_score", r.pareto_score)?;
5619 dict.set_item("feasible", true)?;
5620 dict.set_item("layer_luts", luts_list)?;
5621 dict.set_item("layer_power", power_list)?;
5622 dict.set_item("layer_accuracy", acc_list)?;
5623 }
5624 None => {
5625 dict.set_item("feasible", false)?;
5626 }
5627 }
5628 dict.set_item("backend", "rust")?;
5629 Ok(dict.into_any().unbind())
5630}
5631
5632#[pyfunction]
5634fn py_opt_extract_pareto<'py>(
5635 py: Python<'py>,
5636 luts: Vec<i64>,
5637 power: Vec<f64>,
5638 score: Vec<f64>,
5639) -> PyResult<Py<PyAny>> {
5640 let indices = optimizer::extract_pareto(&luts, &power, &score);
5641 let dict = PyDict::new(py);
5642 let p_luts: Vec<i64> = indices.iter().map(|&i| luts[i]).collect();
5643 let p_power: Vec<f64> = indices.iter().map(|&i| power[i]).collect();
5644 let p_score: Vec<f64> = indices.iter().map(|&i| score[i]).collect();
5645 dict.set_item("indices", indices)?;
5646 dict.set_item("luts", p_luts)?;
5647 dict.set_item("power", p_power)?;
5648 dict.set_item("score", p_score)?;
5649 dict.set_item("backend", "rust")?;
5650 Ok(dict.into_any().unbind())
5651}
5652
5653#[pyfunction]
5657#[pyo3(signature = (genomes, mutation_rate=0.1, mutation_scale=0.1, seed=42))]
5658fn py_evo_batch_mutate(
5659 _py: Python<'_>,
5660 mut genomes: Vec<Vec<f64>>,
5661 mutation_rate: f64,
5662 mutation_scale: f64,
5663 seed: u64,
5664) -> Vec<Vec<f64>> {
5665 evo::batch_mutate_weights(&mut genomes, mutation_rate, mutation_scale, seed);
5666 genomes
5667}
5668
5669#[pyfunction]
5671fn py_evo_batch_fitness(
5672 _py: Python<'_>,
5673 genomes: Vec<Vec<f64>>,
5674 inputs: Vec<f64>,
5675 target: f64,
5676) -> Vec<f64> {
5677 evo::batch_evaluate_fitness(&genomes, &inputs, target)
5678}
5679
5680#[pyfunction]
5682#[pyo3(signature = (parents_a, parents_b, seed=42))]
5683fn py_evo_batch_crossover(
5684 _py: Python<'_>,
5685 parents_a: Vec<Vec<f64>>,
5686 parents_b: Vec<Vec<f64>>,
5687 seed: u64,
5688) -> Vec<Vec<f64>> {
5689 evo::batch_crossover(&parents_a, &parents_b, seed)
5690}
5691
5692#[pyfunction]
5694fn py_evo_diversity(_py: Python<'_>, genomes: Vec<Vec<f64>>) -> f64 {
5695 evo::population_diversity(&genomes)
5696}
5697
5698#[pyfunction]
5700#[pyo3(signature = (genomes, archive, k_nearest=5))]
5701fn py_evo_novelty(
5702 _py: Python<'_>,
5703 genomes: Vec<Vec<f64>>,
5704 archive: Vec<Vec<f64>>,
5705 k_nearest: usize,
5706) -> Vec<f64> {
5707 evo::novelty_scores(&genomes, &archive, k_nearest)
5708}
5709
5710#[pyfunction]
5712#[pyo3(signature = (fitness, n_select, tournament_size=3, seed=42))]
5713fn py_evo_tournament(
5714 _py: Python<'_>,
5715 fitness: Vec<f64>,
5716 n_select: usize,
5717 tournament_size: usize,
5718 seed: u64,
5719) -> Vec<usize> {
5720 evo::tournament_select(&fitness, n_select, tournament_size, seed)
5721}
5722
5723#[pyfunction]
5741#[pyo3(signature = (
5742 obs_flat, controls_flat, t_len, p_dim, m_dim,
5743 a_flat, b_flat, c_flat, d_flat, q_flat, r_flat,
5744 mu_0, sigma_0_flat, d_dim,
5745))]
5746#[allow(clippy::too_many_arguments)]
5747fn py_lgssm_kalman_filter<'py>(
5748 py: Python<'py>,
5749 obs_flat: Vec<f64>,
5750 controls_flat: Vec<f64>,
5751 t_len: usize,
5752 p_dim: usize,
5753 m_dim: usize,
5754 a_flat: Vec<f64>,
5755 b_flat: Vec<f64>,
5756 c_flat: Vec<f64>,
5757 d_flat: Vec<f64>,
5758 q_flat: Vec<f64>,
5759 r_flat: Vec<f64>,
5760 mu_0: Vec<f64>,
5761 sigma_0_flat: Vec<f64>,
5762 d_dim: usize,
5763) -> PyResult<Py<PyAny>> {
5764 use ndarray::Array1;
5765 use ndarray::Array2;
5766
5767 let to_2d = |flat: &[f64], rows: usize, cols: usize| -> Array2<f64> {
5768 Array2::from_shape_vec((rows, cols), flat.to_vec()).expect("shape")
5769 };
5770 let obs = to_2d(&obs_flat, t_len, p_dim);
5771 let controls = to_2d(&controls_flat, t_len, m_dim);
5772 let a = to_2d(&a_flat, d_dim, d_dim);
5773 let b = to_2d(&b_flat, d_dim, m_dim);
5774 let c = to_2d(&c_flat, p_dim, d_dim);
5775 let d = to_2d(&d_flat, p_dim, m_dim);
5776 let q = to_2d(&q_flat, d_dim, d_dim);
5777 let r = to_2d(&r_flat, p_dim, p_dim);
5778 let mu_0_arr = Array1::from(mu_0);
5779 let sigma_0 = to_2d(&sigma_0_flat, d_dim, d_dim);
5780
5781 let result = lgssm::kalman_filter(
5782 obs.view(),
5783 controls.view(),
5784 a.view(),
5785 b.view(),
5786 c.view(),
5787 d.view(),
5788 q.view(),
5789 r.view(),
5790 mu_0_arr.view(),
5791 sigma_0.view(),
5792 );
5793
5794 let means: Vec<Vec<f64>> = (0..t_len)
5796 .map(|t| (0..d_dim).map(|i| result.means[(t, i)]).collect())
5797 .collect();
5798 let covs: Vec<Vec<Vec<f64>>> = (0..t_len)
5799 .map(|t| {
5800 (0..d_dim)
5801 .map(|i| (0..d_dim).map(|j| result.covariances[(t, i, j)]).collect())
5802 .collect()
5803 })
5804 .collect();
5805 let pred_means: Vec<Vec<f64>> = (0..t_len)
5806 .map(|t| (0..d_dim).map(|i| result.pred_means[(t, i)]).collect())
5807 .collect();
5808 let pred_covs: Vec<Vec<Vec<f64>>> = (0..t_len)
5809 .map(|t| {
5810 (0..d_dim)
5811 .map(|i| {
5812 (0..d_dim)
5813 .map(|j| result.pred_covariances[(t, i, j)])
5814 .collect()
5815 })
5816 .collect()
5817 })
5818 .collect();
5819
5820 let dict = PyDict::new(py);
5821 dict.set_item("means", means)?;
5822 dict.set_item("covariances", covs)?;
5823 dict.set_item("pred_means", pred_means)?;
5824 dict.set_item("pred_covariances", pred_covs)?;
5825 dict.set_item("log_likelihood", result.log_likelihood)?;
5826 dict.set_item("backend", "rust")?;
5827 Ok(dict.into_any().unbind())
5828}
5829
5830#[pyfunction]
5840#[pyo3(signature = (knm_flat, n, i, j))]
5841fn py_ollivier_ricci_curvature(knm_flat: Vec<f64>, n: usize, i: usize, j: usize) -> PyResult<f64> {
5842 use topology::CurvatureError;
5843 match topology::ollivier_ricci_curvature(&knm_flat, n, i, j) {
5844 Ok(value) => Ok(value),
5845 Err(CurvatureError::BadShape) => Err(PyValueError::new_err(
5846 "knm must be a square coupling matrix with at least one node",
5847 )),
5848 Err(CurvatureError::BadValue) => Err(PyValueError::new_err(
5849 "knm must contain only finite, non-negative values",
5850 )),
5851 Err(CurvatureError::BadIndex) => Err(PyValueError::new_err(
5852 "node index out of range for coupling graph",
5853 )),
5854 Err(CurvatureError::Infeasible) => {
5855 Err(PyValueError::new_err("transport problem is infeasible"))
5856 }
5857 }
5858}
5859
5860#[pyfunction]
5867#[pyo3(signature = (x0, y0, a, epsilon, sigma, x_threshold, n_steps, current))]
5868#[allow(clippy::too_many_arguments)]
5869fn py_cazelles_map_simulate<'py>(
5870 py: Python<'py>,
5871 x0: f64,
5872 y0: f64,
5873 a: f64,
5874 epsilon: f64,
5875 sigma: f64,
5876 x_threshold: f64,
5877 n_steps: usize,
5878 current: f64,
5879) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
5880 let mut neuron = crate::neurons::CazellesMapNeuron {
5881 x: x0,
5882 y: y0,
5883 a,
5884 epsilon,
5885 sigma,
5886 x_threshold,
5887 };
5888 let (trace, spikes) = neuron.simulate(n_steps, current);
5889 (trace.into_pyarray(py), spikes, neuron.x, neuron.y)
5892}
5893
5894#[pyfunction]
5904#[pyo3(signature = (x0, y0, m0, m1, a, d, j, beta, eps, x_threshold, n_steps, current))]
5905#[allow(clippy::too_many_arguments)]
5906fn py_courage_nekorkin_map_simulate<'py>(
5907 py: Python<'py>,
5908 x0: f64,
5909 y0: f64,
5910 m0: f64,
5911 m1: f64,
5912 a: f64,
5913 d: f64,
5914 j: f64,
5915 beta: f64,
5916 eps: f64,
5917 x_threshold: f64,
5918 n_steps: usize,
5919 current: f64,
5920) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
5921 let mut neuron = crate::neurons::CourageNekorkinMapNeuron {
5922 x: x0,
5923 y: y0,
5924 m0,
5925 m1,
5926 a,
5927 d,
5928 j,
5929 beta,
5930 eps,
5931 x_threshold,
5932 };
5933 let (trace, spikes) = neuron.simulate(n_steps, current);
5934 (trace.into_pyarray(py), spikes, neuron.x, neuron.y)
5935}
5936
5937#[pyfunction]
5943#[pyo3(signature = (x0, y0, alpha, sigma, mu, x_threshold, n_steps, current))]
5944#[allow(clippy::too_many_arguments)]
5945fn py_rulkov_map_simulate<'py>(
5946 py: Python<'py>,
5947 x0: f64,
5948 y0: f64,
5949 alpha: f64,
5950 sigma: f64,
5951 mu: f64,
5952 x_threshold: f64,
5953 n_steps: usize,
5954 current: f64,
5955) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
5956 let mut neuron = crate::neurons::RulkovMapNeuron {
5957 x: x0,
5958 y: y0,
5959 alpha,
5960 sigma,
5961 mu,
5962 x_threshold,
5963 };
5964 let (trace, spikes) = neuron.simulate(n_steps, current);
5965 (trace.into_pyarray(py), spikes, neuron.x, neuron.y)
5966}
5967
5968#[pyfunction]
5976#[pyo3(signature = (x0, y0, alpha, beta, mu, sigma, x_threshold, x_reset, n_steps, current))]
5977#[allow(clippy::too_many_arguments)]
5978fn py_ibarz_tanaka_map_simulate<'py>(
5979 py: Python<'py>,
5980 x0: f64,
5981 y0: f64,
5982 alpha: f64,
5983 beta: f64,
5984 mu: f64,
5985 sigma: f64,
5986 x_threshold: f64,
5987 x_reset: f64,
5988 n_steps: usize,
5989 current: f64,
5990) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
5991 let mut neuron = crate::neurons::IbarzTanakaMapNeuron {
5992 x: x0,
5993 y: y0,
5994 alpha,
5995 beta,
5996 mu,
5997 sigma,
5998 x_threshold,
5999 x_reset,
6000 };
6001 let (trace, spikes) = neuron.simulate(n_steps, current);
6002 (trace.into_pyarray(py), spikes, neuron.x, neuron.y)
6003}
6004
6005#[pyfunction]
6013#[pyo3(signature = (x0, alpha, beta, x_threshold, n_steps, current))]
6014fn py_medvedev_map_simulate<'py>(
6015 py: Python<'py>,
6016 x0: f64,
6017 alpha: f64,
6018 beta: f64,
6019 x_threshold: f64,
6020 n_steps: usize,
6021 current: f64,
6022) -> (Bound<'py, PyArray1<f64>>, i64, f64) {
6023 let mut neuron = crate::neurons::MedvedevMapNeuron {
6024 x: x0,
6025 alpha,
6026 beta,
6027 x_threshold,
6028 };
6029 let (trace, spikes) = neuron.simulate(n_steps, current);
6030 (trace.into_pyarray(py), spikes, neuron.x)
6031}
6032
6033#[pyfunction]
6041#[pyo3(signature = (theta0, dt, gain, theta_threshold, n_steps, current))]
6042fn py_ermentrout_kopell_map_simulate<'py>(
6043 py: Python<'py>,
6044 theta0: f64,
6045 dt: f64,
6046 gain: f64,
6047 theta_threshold: f64,
6048 n_steps: usize,
6049 current: f64,
6050) -> (Bound<'py, PyArray1<f64>>, i64, f64) {
6051 let mut neuron = crate::neurons::ErmentroutKopellMapNeuron {
6052 theta: theta0,
6053 dt,
6054 gain,
6055 theta_threshold,
6056 };
6057 let (trace, spikes) = neuron.simulate(n_steps, current);
6058 (trace.into_pyarray(py), spikes, neuron.theta)
6059}
6060
6061#[pyfunction]
6071#[pyo3(signature = (v0, w0, a, epsilon, gamma, dt, v_peak, n_steps, current))]
6072#[allow(clippy::too_many_arguments)]
6073fn py_mckean_simulate<'py>(
6074 py: Python<'py>,
6075 v0: f64,
6076 w0: f64,
6077 a: f64,
6078 epsilon: f64,
6079 gamma: f64,
6080 dt: f64,
6081 v_peak: f64,
6082 n_steps: usize,
6083 current: f64,
6084) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6085 let mut neuron = crate::neurons::McKeanNeuron {
6086 v: v0,
6087 w: w0,
6088 a,
6089 epsilon,
6090 gamma,
6091 dt,
6092 v_peak,
6093 };
6094 let (trace, spikes) = neuron.simulate(n_steps, current);
6095 (trace.into_pyarray(py), spikes, neuron.v, neuron.w)
6096}
6097
6098#[pyfunction]
6107#[pyo3(signature = (v0, r0, tau_r, v_peak, dt, n_steps, current))]
6108#[allow(clippy::too_many_arguments)]
6109fn py_wilson_hr_simulate<'py>(
6110 py: Python<'py>,
6111 v0: f64,
6112 r0: f64,
6113 tau_r: f64,
6114 v_peak: f64,
6115 dt: f64,
6116 n_steps: usize,
6117 current: f64,
6118) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6119 let mut neuron = crate::neurons::WilsonHRNeuron {
6120 v: v0,
6121 r: r0,
6122 tau_r,
6123 v_peak,
6124 dt,
6125 };
6126 let (trace, spikes) = neuron.simulate(n_steps, current);
6127 (trace.into_pyarray(py), spikes, neuron.v, neuron.r)
6128}
6129
6130#[pyfunction]
6139#[pyo3(signature = (v0, w0, z0, alpha, beta, eps1, eps2, gamma, dt, v_threshold, n_steps, current))]
6140#[allow(clippy::too_many_arguments)]
6141fn py_pernarowski_simulate<'py>(
6142 py: Python<'py>,
6143 v0: f64,
6144 w0: f64,
6145 z0: f64,
6146 alpha: f64,
6147 beta: f64,
6148 eps1: f64,
6149 eps2: f64,
6150 gamma: f64,
6151 dt: f64,
6152 v_threshold: f64,
6153 n_steps: usize,
6154 current: f64,
6155) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64) {
6156 let mut neuron = crate::neurons::PernarowskiNeuron {
6157 v: v0,
6158 w: w0,
6159 z: z0,
6160 alpha,
6161 beta,
6162 eps1,
6163 eps2,
6164 gamma,
6165 dt,
6166 v_threshold,
6167 };
6168 let (trace, spikes) = neuron.simulate(n_steps, current);
6169 (trace.into_pyarray(py), spikes, neuron.v, neuron.w, neuron.z)
6170}
6171
6172#[pyfunction]
6183#[pyo3(signature = (v0, w0, alpha, beta, epsilon, rho, dt, v_peak, n_steps, current))]
6184#[allow(clippy::too_many_arguments)]
6185fn py_terman_wang_simulate<'py>(
6186 py: Python<'py>,
6187 v0: f64,
6188 w0: f64,
6189 alpha: f64,
6190 beta: f64,
6191 epsilon: f64,
6192 rho: f64,
6193 dt: f64,
6194 v_peak: f64,
6195 n_steps: usize,
6196 current: f64,
6197) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6198 let mut neuron = crate::neurons::TermanWangOscillator {
6199 v: v0,
6200 w: w0,
6201 alpha,
6202 beta,
6203 epsilon,
6204 rho,
6205 dt,
6206 v_peak,
6207 };
6208 let (trace, spikes) = neuron.simulate(n_steps, current);
6209 (trace.into_pyarray(py), spikes, neuron.v, neuron.w)
6210}
6211
6212#[pyfunction]
6221#[pyo3(signature = (
6222 v0, theta0, i1_0, i2_0, v_rest, v_reset, theta_reset, theta_inf,
6223 tau_v, tau_theta, tau_1, tau_2, a, b, r1, r2, dt, n_steps, current
6224))]
6225#[allow(clippy::too_many_arguments)]
6226fn py_mihalas_niebur_simulate<'py>(
6227 py: Python<'py>,
6228 v0: f64,
6229 theta0: f64,
6230 i1_0: f64,
6231 i2_0: f64,
6232 v_rest: f64,
6233 v_reset: f64,
6234 theta_reset: f64,
6235 theta_inf: f64,
6236 tau_v: f64,
6237 tau_theta: f64,
6238 tau_1: f64,
6239 tau_2: f64,
6240 a: f64,
6241 b: f64,
6242 r1: f64,
6243 r2: f64,
6244 dt: f64,
6245 n_steps: usize,
6246 current: f64,
6247) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64, f64) {
6248 let mut neuron = crate::neurons::MihalasNieburNeuron {
6249 v: v0,
6250 theta: theta0,
6251 i1: i1_0,
6252 i2: i2_0,
6253 v_rest,
6254 v_reset,
6255 theta_reset,
6256 theta_inf,
6257 tau_v,
6258 tau_theta,
6259 tau_1,
6260 tau_2,
6261 a,
6262 b,
6263 r1,
6264 r2,
6265 dt,
6266 };
6267 let (trace, spikes) = neuron.simulate(n_steps, current);
6268 (
6269 trace.into_pyarray(py),
6270 spikes,
6271 neuron.v,
6272 neuron.theta,
6273 neuron.i1,
6274 neuron.i2,
6275 )
6276}
6277
6278#[pyfunction]
6287#[pyo3(signature = (
6288 v0, theta0, theta_inf, i_asc1_0, i_asc2_0, v_rest, v_reset, tau_m, tau_theta,
6289 tau_asc1, tau_asc2, a_theta, delta_theta, r_asc1, r_asc2, resistance, dt,
6290 n_steps, current
6291))]
6292#[allow(clippy::too_many_arguments)]
6293fn py_glif_simulate<'py>(
6294 py: Python<'py>,
6295 v0: f64,
6296 theta0: f64,
6297 theta_inf: f64,
6298 i_asc1_0: f64,
6299 i_asc2_0: f64,
6300 v_rest: f64,
6301 v_reset: f64,
6302 tau_m: f64,
6303 tau_theta: f64,
6304 tau_asc1: f64,
6305 tau_asc2: f64,
6306 a_theta: f64,
6307 delta_theta: f64,
6308 r_asc1: f64,
6309 r_asc2: f64,
6310 resistance: f64,
6311 dt: f64,
6312 n_steps: usize,
6313 current: f64,
6314) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64, f64) {
6315 let mut neuron = crate::neurons::GLIFNeuron {
6316 v: v0,
6317 theta: theta0,
6318 theta_inf,
6319 i_asc1: i_asc1_0,
6320 i_asc2: i_asc2_0,
6321 v_rest,
6322 v_reset,
6323 tau_m,
6324 tau_theta,
6325 tau_asc1,
6326 tau_asc2,
6327 a_theta,
6328 delta_theta,
6329 r_asc1,
6330 r_asc2,
6331 resistance,
6332 dt,
6333 };
6334 let (trace, spikes) = neuron.simulate(n_steps, current);
6335 (
6336 trace.into_pyarray(py),
6337 spikes,
6338 neuron.v,
6339 neuron.theta,
6340 neuron.i_asc1,
6341 neuron.i_asc2,
6342 )
6343}
6344
6345#[pyfunction]
6353#[pyo3(signature = (v0, w0, a, b, epsilon, dt, v_threshold, n_steps, current))]
6354#[allow(clippy::too_many_arguments)]
6355fn py_fitzhugh_nagumo_simulate<'py>(
6356 py: Python<'py>,
6357 v0: f64,
6358 w0: f64,
6359 a: f64,
6360 b: f64,
6361 epsilon: f64,
6362 dt: f64,
6363 v_threshold: f64,
6364 n_steps: usize,
6365 current: f64,
6366) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6367 let mut neuron = crate::neurons::FitzHughNagumoNeuron {
6368 v: v0,
6369 w: w0,
6370 a,
6371 b,
6372 epsilon,
6373 dt,
6374 v_threshold,
6375 };
6376 let (trace, spikes) = neuron.simulate(n_steps, current);
6377 (trace.into_pyarray(py), spikes, neuron.v, neuron.w)
6378}
6379
6380#[pyfunction]
6388#[pyo3(signature = (x0, y0, z0, b, r, s, x_rest, dt, x_threshold, n_steps, current))]
6389#[allow(clippy::too_many_arguments)]
6390fn py_hindmarsh_rose_simulate<'py>(
6391 py: Python<'py>,
6392 x0: f64,
6393 y0: f64,
6394 z0: f64,
6395 b: f64,
6396 r: f64,
6397 s: f64,
6398 x_rest: f64,
6399 dt: f64,
6400 x_threshold: f64,
6401 n_steps: usize,
6402 current: f64,
6403) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64) {
6404 let mut neuron = crate::neurons::HindmarshRoseNeuron {
6405 x: x0,
6406 y: y0,
6407 z: z0,
6408 b,
6409 r,
6410 s,
6411 x_rest,
6412 dt,
6413 x_threshold,
6414 };
6415 let (trace, spikes) = neuron.simulate(n_steps, current);
6416 (trace.into_pyarray(py), spikes, neuron.x, neuron.y, neuron.z)
6417}
6418
6419#[pyfunction]
6426#[pyo3(signature = (v0, w0, y0, a, b, c, d, delta, mu, dt, v_threshold, n_steps, current))]
6427#[allow(clippy::too_many_arguments)]
6428fn py_fitzhugh_rinzel_simulate<'py>(
6429 py: Python<'py>,
6430 v0: f64,
6431 w0: f64,
6432 y0: f64,
6433 a: f64,
6434 b: f64,
6435 c: f64,
6436 d: f64,
6437 delta: f64,
6438 mu: f64,
6439 dt: f64,
6440 v_threshold: f64,
6441 n_steps: usize,
6442 current: f64,
6443) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64, f64) {
6444 let mut neuron = crate::neurons::FitzHughRinzelNeuron {
6445 v: v0,
6446 w: w0,
6447 y: y0,
6448 a,
6449 b,
6450 c,
6451 d,
6452 delta,
6453 mu,
6454 dt,
6455 v_threshold,
6456 };
6457 let (trace, spikes) = neuron.simulate(n_steps, current);
6458 (trace.into_pyarray(py), spikes, neuron.v, neuron.w, neuron.y)
6459}
6460
6461#[pyfunction]
6468#[pyo3(signature = (v0, u0, cap, k, vr, vt, vpeak, a, b, c, d, dt, n_steps, current))]
6469#[allow(clippy::too_many_arguments)]
6470fn py_izhikevich2007_simulate<'py>(
6471 py: Python<'py>,
6472 v0: f64,
6473 u0: f64,
6474 cap: f64,
6475 k: f64,
6476 vr: f64,
6477 vt: f64,
6478 vpeak: f64,
6479 a: f64,
6480 b: f64,
6481 c: f64,
6482 d: f64,
6483 dt: f64,
6484 n_steps: usize,
6485 current: f64,
6486) -> (Bound<'py, PyArray1<f64>>, i64, f64, f64) {
6487 let mut neuron = crate::rk4_neurons::Izhikevich2007Rk4 {
6488 v: v0,
6489 u: u0,
6490 cap,
6491 k,
6492 vr,
6493 vt,
6494 vpeak,
6495 a,
6496 b,
6497 c,
6498 d,
6499 dt,
6500 };
6501 let (trace, spikes) = neuron.simulate(n_steps, current);
6502 (trace.into_pyarray(py), spikes, neuron.v, neuron.u)
6503}
6504
6505#[pyfunction]
6513fn py_inject_bitflip_u8<'py>(
6514 py: Python<'py>,
6515 bitstream: PyReadonlyArray1<'_, u8>,
6516 ber: f64,
6517 seed: u64,
6518) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6519 let mut buf = bitstream.as_slice()?.to_vec();
6520 let n = fault::inject_bitflip_u8(&mut buf, ber, seed);
6521 Ok((buf.into_pyarray(py).into(), n))
6522}
6523
6524#[pyfunction]
6525fn py_inject_stuck_at_0_u8<'py>(
6526 py: Python<'py>,
6527 bitstream: PyReadonlyArray1<'_, u8>,
6528 ber: f64,
6529 seed: u64,
6530) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6531 let mut buf = bitstream.as_slice()?.to_vec();
6532 let n = fault::inject_stuck_at_0_u8(&mut buf, ber, seed);
6533 Ok((buf.into_pyarray(py).into(), n))
6534}
6535
6536#[pyfunction]
6537fn py_inject_stuck_at_1_u8<'py>(
6538 py: Python<'py>,
6539 bitstream: PyReadonlyArray1<'_, u8>,
6540 ber: f64,
6541 seed: u64,
6542) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6543 let mut buf = bitstream.as_slice()?.to_vec();
6544 let n = fault::inject_stuck_at_1_u8(&mut buf, ber, seed);
6545 Ok((buf.into_pyarray(py).into(), n))
6546}
6547
6548#[pyfunction]
6549fn py_inject_dropout_u8<'py>(
6550 py: Python<'py>,
6551 bitstream: PyReadonlyArray1<'_, u8>,
6552 ber: f64,
6553 seed: u64,
6554) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6555 let mut buf = bitstream.as_slice()?.to_vec();
6556 let n = fault::inject_dropout_u8(&mut buf, ber, seed);
6557 Ok((buf.into_pyarray(py).into(), n))
6558}
6559
6560#[pyfunction]
6561fn py_inject_gaussian_u8<'py>(
6562 py: Python<'py>,
6563 bitstream: PyReadonlyArray1<'_, u8>,
6564 ber: f64,
6565 seed: u64,
6566) -> PyResult<(Py<PyArray1<u8>>, u64)> {
6567 let mut buf = bitstream.as_slice()?.to_vec();
6568 let n = fault::inject_gaussian_u8(&mut buf, ber, seed);
6569 Ok((buf.into_pyarray(py).into(), n))
6570}
6571
6572#[pyfunction]
6579#[pyo3(signature = (
6580 adj_offsets, adj_neighbours, adj_scc_abs, vertex_weights,
6581 part_map, parts_concat, parts_offsets,
6582 n_parts, kl_iterations, correlation_penalty,
6583))]
6584#[allow(clippy::too_many_arguments)]
6585fn py_kl_refine<'py>(
6586 py: Python<'py>,
6587 adj_offsets: PyReadonlyArray1<'_, i64>,
6588 adj_neighbours: PyReadonlyArray1<'_, i32>,
6589 adj_scc_abs: PyReadonlyArray1<'_, f64>,
6590 vertex_weights: PyReadonlyArray1<'_, f64>,
6591 part_map: PyReadonlyArray1<'_, i32>,
6592 parts_concat: PyReadonlyArray1<'_, i32>,
6593 parts_offsets: PyReadonlyArray1<'_, i64>,
6594 n_parts: i32,
6595 kl_iterations: i32,
6596 correlation_penalty: f64,
6597) -> PyResult<(Py<PyArray1<i32>>, u64)> {
6598 let mut pm = part_map.as_slice()?.to_vec();
6599 let moves = partition::kl_refine(
6600 adj_offsets.as_slice()?,
6601 adj_neighbours.as_slice()?,
6602 adj_scc_abs.as_slice()?,
6603 vertex_weights.as_slice()?,
6604 &mut pm,
6605 parts_concat.as_slice()?,
6606 parts_offsets.as_slice()?,
6607 n_parts,
6608 kl_iterations,
6609 correlation_penalty,
6610 );
6611 Ok((pm.into_pyarray(py).into(), moves))
6612}
6613
6614#[pyfunction]
6627#[pyo3(signature = (
6628 v_e, g_ampa_e, g_gaba_e, refrac_e, i_drive_e, xi_e, spikes_e_out,
6629 v_i, g_ampa_i, g_gaba_i, refrac_i, i_drive_i, xi_i, spikes_i_out,
6630 e_l, e_ampa, e_gaba, g_l, c_m, v_threshold, v_reset, t_refrac,
6631 tau_ampa, tau_gaba, sigma_e, sigma_i, dt,
6632))]
6633#[allow(clippy::too_many_arguments)]
6634fn py_ping_step<'py>(
6635 _py: Python<'py>,
6636 v_e: PyReadwriteArray1<'_, f64>,
6637 g_ampa_e: PyReadwriteArray1<'_, f64>,
6638 g_gaba_e: PyReadwriteArray1<'_, f64>,
6639 refrac_e: PyReadwriteArray1<'_, f64>,
6640 i_drive_e: PyReadonlyArray1<'_, f64>,
6641 xi_e: PyReadonlyArray1<'_, f64>,
6642 spikes_e_out: PyReadwriteArray1<'_, u8>,
6643 v_i: PyReadwriteArray1<'_, f64>,
6644 g_ampa_i: PyReadwriteArray1<'_, f64>,
6645 g_gaba_i: PyReadwriteArray1<'_, f64>,
6646 refrac_i: PyReadwriteArray1<'_, f64>,
6647 i_drive_i: PyReadonlyArray1<'_, f64>,
6648 xi_i: PyReadonlyArray1<'_, f64>,
6649 spikes_i_out: PyReadwriteArray1<'_, u8>,
6650 e_l: f64,
6651 e_ampa: f64,
6652 e_gaba: f64,
6653 g_l: f64,
6654 c_m: f64,
6655 v_threshold: f64,
6656 v_reset: f64,
6657 t_refrac: f64,
6658 tau_ampa: f64,
6659 tau_gaba: f64,
6660 sigma_e: f64,
6661 sigma_i: f64,
6662 dt: f64,
6663) -> PyResult<(u32, u32)> {
6664 let mut v_e = v_e;
6665 let mut g_ampa_e = g_ampa_e;
6666 let mut g_gaba_e = g_gaba_e;
6667 let mut refrac_e = refrac_e;
6668 let mut spikes_e_out = spikes_e_out;
6669 let mut v_i = v_i;
6670 let mut g_ampa_i = g_ampa_i;
6671 let mut g_gaba_i = g_gaba_i;
6672 let mut refrac_i = refrac_i;
6673 let mut spikes_i_out = spikes_i_out;
6674 let (ne, ni) = ping::step_kernel(
6675 v_e.as_slice_mut()?,
6676 g_ampa_e.as_slice_mut()?,
6677 g_gaba_e.as_slice_mut()?,
6678 refrac_e.as_slice_mut()?,
6679 i_drive_e.as_slice()?,
6680 xi_e.as_slice()?,
6681 spikes_e_out.as_slice_mut()?,
6682 v_i.as_slice_mut()?,
6683 g_ampa_i.as_slice_mut()?,
6684 g_gaba_i.as_slice_mut()?,
6685 refrac_i.as_slice_mut()?,
6686 i_drive_i.as_slice()?,
6687 xi_i.as_slice()?,
6688 spikes_i_out.as_slice_mut()?,
6689 e_l,
6690 e_ampa,
6691 e_gaba,
6692 g_l,
6693 c_m,
6694 v_threshold,
6695 v_reset,
6696 t_refrac,
6697 tau_ampa,
6698 tau_gaba,
6699 sigma_e,
6700 sigma_i,
6701 dt,
6702 );
6703 Ok((ne, ni))
6704}
6705
6706#[pyfunction]
6716#[pyo3(signature = (indptr, indices, data, x, y))]
6717fn py_parallel_csr_spmv_add(
6718 indptr: PyReadonlyArray1<'_, i32>,
6719 indices: PyReadonlyArray1<'_, i32>,
6720 data: PyReadonlyArray1<'_, f64>,
6721 x: PyReadonlyArray1<'_, f64>,
6722 y: PyReadwriteArray1<'_, f64>,
6723) -> PyResult<()> {
6724 let mut y = y;
6725 cortical_inject::parallel_csr_spmv_add(
6726 indptr.as_slice()?,
6727 indices.as_slice()?,
6728 data.as_slice()?,
6729 x.as_slice()?,
6730 y.as_slice_mut()?,
6731 );
6732 Ok(())
6733}
6734
6735#[pyfunction]
6741#[pyo3(signature = (indptrs, indices_list, data_list, xs, y))]
6742fn py_parallel_csr_multi_spmv_add(
6743 indptrs: Vec<PyReadonlyArray1<'_, i32>>,
6744 indices_list: Vec<PyReadonlyArray1<'_, i32>>,
6745 data_list: Vec<PyReadonlyArray1<'_, f64>>,
6746 xs: Vec<PyReadonlyArray1<'_, f64>>,
6747 y: PyReadwriteArray1<'_, f64>,
6748) -> PyResult<()> {
6749 let mut y = y;
6750 let indptr_slices: Vec<&[i32]> = indptrs
6751 .iter()
6752 .map(|a| a.as_slice())
6753 .collect::<Result<_, _>>()?;
6754 let indices_slices: Vec<&[i32]> = indices_list
6755 .iter()
6756 .map(|a| a.as_slice())
6757 .collect::<Result<_, _>>()?;
6758 let data_slices: Vec<&[f64]> = data_list
6759 .iter()
6760 .map(|a| a.as_slice())
6761 .collect::<Result<_, _>>()?;
6762 let x_slices: Vec<&[f64]> = xs.iter().map(|a| a.as_slice()).collect::<Result<_, _>>()?;
6763 cortical_inject::parallel_csr_multi_spmv_add(
6764 &indptr_slices,
6765 &indices_slices,
6766 &data_slices,
6767 &x_slices,
6768 y.as_slice_mut()?,
6769 );
6770 Ok(())
6771}