Skip to main content

sc_neurocore_engine/bindings/
partition.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Commercial license available
3// © Concepts 1996–2026 Miroslav Šotek. All rights reserved.
4// © Code 2020–2026 Miroslav Šotek. All rights reserved.
5// ORCID: 0009-0009-3560-0851
6// Contact: www.anulum.li | protoscience@anulum.li
7// SC-NeuroCore — hierarchical partition PyO3 binding
8
9//! Python binding for correlation-aware Kernighan-Lin partition refinement.
10
11use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1};
12use pyo3::prelude::*;
13
14use crate::partition;
15
16/// Register hierarchical partition refinement with the extension module.
17pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
18    module.add_function(wrap_pyfunction!(py_kl_refine, module)?)?;
19    Ok(())
20}
21
22// Caller passes flat numpy arrays (CSR adjacency + flat scc weights +
23// flat vertex_weights + initial part_map). The kernel mutates a copy
24// of part_map in-place and returns (new_part_map, num_moves).
25#[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}