ennbo-py 0.4.0

Python bindings for ENN core algorithms
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

//! Hash-based RNG Python bindings.

use ndarray::IxDyn;
use numpy::{IntoPyArray, PyArrayDyn, PyReadonlyArray1, PyReadonlyArrayDyn};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;

/// Python wrapper for the NumPy Philox reference hash.
#[pyfunction(name = "normal_hash_batch_multi_seed")]
#[doc = "kiss-coverage-off"]
pub fn normal_hash_batch_multi_seed_py<'py>(
    py: Python<'py>,
    function_seeds: PyReadonlyArray1<i64>,
    data_indices: PyReadonlyArrayDyn<i64>,
    num_metrics: i64,
) -> PyResult<Bound<'py, PyArrayDyn<f64>>> {
    let seeds: Vec<i64> = function_seeds.as_array().iter().copied().collect();
    let indices_arr = data_indices.as_array();
    let indices: Vec<i64> = indices_arr.iter().copied().collect();
    let input_shape = indices_arr.shape().to_vec();
    let mut output_shape = Vec::with_capacity(2 + input_shape.len());
    output_shape.push(seeds.len());
    output_shape.extend(input_shape.iter().copied());
    output_shape.push(num_metrics.max(0) as usize);
    let result = py.allow_threads(|| ennbo::normal_hash_batch_multi_seed(&seeds, &indices, num_metrics));
    match result {
        Ok(arr) => {
            let reshaped = arr
                .into_shape_with_order(IxDyn(&output_shape))
                .map_err(|e| PyValueError::new_err(format!("Shape error: {e}")))?;
            Ok(reshaped.into_pyarray_bound(py))
        }
        Err(e) => Err(PyValueError::new_err(e.to_string())),
    }
}

/// Python wrapper for normal_hash_batch_multi_seed_fast
#[pyfunction(name = "normal_hash_batch_multi_seed_fast")]
#[doc = "kiss-coverage-off"]
pub fn normal_hash_batch_multi_seed_fast_py<'py>(
    py: Python<'py>,
    function_seeds: PyReadonlyArray1<i64>,
    data_indices: PyReadonlyArrayDyn<i64>,
    num_metrics: i64,
) -> PyResult<Bound<'py, PyArrayDyn<f64>>> {
    let seeds_arr = function_seeds.as_array();
    let indices_arr = data_indices.as_array();

    let seeds: Vec<i64> = seeds_arr.iter().copied().collect();
    let indices: Vec<i64> = indices_arr.iter().copied().collect();
    let input_shape = indices_arr.shape().to_vec();
    let mut output_shape = Vec::with_capacity(2 + input_shape.len());
    output_shape.push(seeds.len());
    output_shape.extend(input_shape.iter().copied());
    output_shape.push(num_metrics.max(0) as usize);


    let result = py
        .allow_threads(|| ennbo::normal_hash_batch_multi_seed_fast(&seeds, &indices, num_metrics));

    match result {
        Ok(arr) => {


            let reshaped = arr
                .into_shape_with_order(IxDyn(&output_shape))
                .map_err(|e| PyValueError::new_err(format!("Shape error: {}", e)))?;
            Ok(reshaped.into_pyarray_bound(py))
        }
        Err(e) => Err(PyValueError::new_err(e.to_string())),
    }
}