rustnum 0.1.0

NumPy-style n-dimensional arrays for Rust: one array type, NumPy names, pure Rust.
Documentation 
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

//! Kernels for NumPy's statistics, signal, set and tensor-contraction functions.
//!
//! The kernels are pure slice code: `&[T]` plus a shape, no `Array`. The public
//! `Array`-level wrappers live in `api.rs` and are re-exported from the crate root.
//!
//! Conventions shared by all kernels:
//! * multi-dimensional buffers are C-ordered (row-major) and contiguous;
//! * misuse panics with NumPy's wording where NumPy has one;
//! * NaN sorts last (NumPy's sort order).

pub(crate) mod api;
mod einsum_kernels;
mod set_kernels;
mod signal_kernels;
mod stats_kernels;

use std::cmp::Ordering;

/// `true` for values that are not equal to themselves (NaN), without a float-only API.
pub(crate) fn is_nan_like<T: PartialOrd>(x: &T) -> bool {
    x.partial_cmp(x).is_none()
}

/// Total order used by NumPy's `sort`: ordinary order, NaN after everything, NaNs equal.
pub(crate) fn cmp_nan_last<T: PartialOrd>(a: &T, b: &T) -> Ordering {
    match a.partial_cmp(b) {
        Some(o) => o,
        None => is_nan_like(a).cmp(&is_nan_like(b)),
    }
}

/// `np.searchsorted(sorted, v, side="left")` for NaN-last sorted input.
pub(crate) fn searchsorted_left<T: PartialOrd>(sorted: &[T], v: &T) -> usize {
    sorted.partition_point(|e| cmp_nan_last(e, v) == Ordering::Less)
}

/// `np.searchsorted(sorted, v, side="right")` for NaN-last sorted input.
pub(crate) fn searchsorted_right<T: PartialOrd>(sorted: &[T], v: &T) -> usize {
    sorted.partition_point(|e| cmp_nan_last(e, v) != Ordering::Greater)
}

/// Resolve a possibly negative axis, panicking like NumPy.
pub(crate) fn normalize_axis(axis: isize, ndim: usize) -> usize {
    let n = ndim as isize;
    if axis < -n || axis >= n {
        panic!("axis {axis} is out of bounds for array of dimension {ndim}");
    }
    if axis < 0 {
        (axis + n) as usize
    } else {
        axis as usize
    }
}

/// Split a C-ordered shape around `axis` into `(outer, len, inner)`.
///
/// Element `(o, k, i)` lives at flat index `(o * len + k) * inner + i`.
pub(crate) fn axis_split(shape: &[usize], axis: usize) -> (usize, usize, usize) {
    let outer = shape[..axis].iter().product();
    let inner = shape[axis + 1..].iter().product();
    (outer, shape[axis], inner)
}

/// Test helper: elementwise comparison with 1e-12 relative tolerance; NaN equals NaN.
#[cfg(test)]
pub(crate) fn assert_close(got: &[f64], want: &[f64]) {
    assert_eq!(got.len(), want.len(), "length: got {got:?}, want {want:?}");
    for (i, (g, w)) in got.iter().zip(want).enumerate() {
        let ok = (g.is_nan() && w.is_nan()) || g == w || (g - w).abs() <= 1e-12 * w.abs().max(1.0);
        assert!(
            ok,
            "element {i}: got {g}, want {w}
 got  {got:?}
 want {want:?}"
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn nan_last_order() {
        let mut v = [3.0, f64::NAN, 1.0, f64::NAN, 2.0];
        v.sort_by(cmp_nan_last);
        assert_eq!(&v[..3], &[1.0, 2.0, 3.0]);
        assert!(v[3].is_nan() && v[4].is_nan());
    }

    #[test]
    fn searchsorted_sides() {
        let s = [1.0, 2.0, 2.0, 4.0, f64::NAN];
        assert_eq!(searchsorted_left(&s, &2.0), 1);
        assert_eq!(searchsorted_right(&s, &2.0), 3);
        assert_eq!(searchsorted_left(&s, &9.0), 4);
        assert_eq!(searchsorted_left(&s, &f64::NAN), 4);
        assert_eq!(searchsorted_right(&s, &f64::NAN), 5);
    }

    #[test]
    fn axes() {
        assert_eq!(normalize_axis(-1, 3), 2);
        assert_eq!(axis_split(&[2, 3, 4], 1), (2, 3, 4));
    }

    #[test]
    #[should_panic(expected = "axis 2 is out of bounds for array of dimension 2")]
    fn bad_axis() {
        normalize_axis(2, 2);
    }
}