sc_neurocore_engine/bindings/
partition.rs1use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1};
12use pyo3::prelude::*;
13
14use crate::partition;
15
16pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
18 module.add_function(wrap_pyfunction!(py_kl_refine, module)?)?;
19 Ok(())
20}
21
22#[pyfunction]
26#[pyo3(signature = (
27 adj_offsets, adj_neighbours, adj_scc_abs, vertex_weights,
28 part_map, parts_concat, parts_offsets,
29 n_parts, kl_iterations, correlation_penalty,
30))]
31#[allow(clippy::too_many_arguments)]
32fn py_kl_refine<'py>(
33 py: Python<'py>,
34 adj_offsets: PyReadonlyArray1<'_, i64>,
35 adj_neighbours: PyReadonlyArray1<'_, i32>,
36 adj_scc_abs: PyReadonlyArray1<'_, f64>,
37 vertex_weights: PyReadonlyArray1<'_, f64>,
38 part_map: PyReadonlyArray1<'_, i32>,
39 parts_concat: PyReadonlyArray1<'_, i32>,
40 parts_offsets: PyReadonlyArray1<'_, i64>,
41 n_parts: i32,
42 kl_iterations: i32,
43 correlation_penalty: f64,
44) -> PyResult<(Py<PyArray1<i32>>, u64)> {
45 let mut pm = part_map.as_slice()?.to_vec();
46 let moves = partition::kl_refine(
47 adj_offsets.as_slice()?,
48 adj_neighbours.as_slice()?,
49 adj_scc_abs.as_slice()?,
50 vertex_weights.as_slice()?,
51 &mut pm,
52 parts_concat.as_slice()?,
53 parts_offsets.as_slice()?,
54 n_parts,
55 kl_iterations,
56 correlation_penalty,
57 );
58 Ok((pm.into_pyarray(py).into(), moves))
59}