numrs2 0.3.3

A Rust implementation inspired by NumPy for numerical computing (NumRS2)
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

use numrs2::array::Array;
use numrs2::indexing::*;

#[test]
fn test_ellipsis_indexing() {
    // Create a 3D array
    let mut array = Array::<f64>::zeros(&[2, 3, 4]);

    // Fill with test data
    for i in 0..2 {
        for j in 0..3 {
            for k in 0..4 {
                let value = (i * 100 + j * 10 + k) as f64;
                array.set(&[i, j, k], value).unwrap();
            }
        }
    }

    // Test ellipsis at the beginning
    // This means: [:, :, 0] - select the first slice of the last dimension
    let result1 = array
        .index(&[IndexSpec::Ellipsis, IndexSpec::Index(0)])
        .unwrap();
    assert_eq!(result1.shape(), vec![2, 3]);

    // Check values
    for i in 0..2 {
        for j in 0..3 {
            assert_eq!(result1.get(&[i, j]).unwrap(), (i * 100 + j * 10) as f64);
        }
    }

    // Test ellipsis in the middle
    // This means: [0, :, :] - select the first slice of the first dimension
    let result2 = array
        .index(&[IndexSpec::Index(0), IndexSpec::Ellipsis])
        .unwrap();
    assert_eq!(result2.shape(), vec![3, 4]);

    // Check values
    for j in 0..3 {
        for k in 0..4 {
            assert_eq!(result2.get(&[j, k]).unwrap(), (j * 10 + k) as f64);
        }
    }

    // Test ellipsis with multiple indices
    // This means: [:, 0, 0] - select the first slice of second and third dimensions
    let result3 = array
        .index(&[
            IndexSpec::Ellipsis,
            IndexSpec::Index(0),
            IndexSpec::Index(0),
        ])
        .unwrap();
    assert_eq!(result3.shape(), vec![2]);

    // Check values
    for i in 0..2 {
        assert_eq!(result3.get(&[i]).unwrap(), (i * 100) as f64);
    }
}

#[test]
fn test_fancy_indexing() {
    // Create a 2D array
    let mut array = Array::<i32>::zeros(&[3, 4]);

    // Fill with test data
    for i in 0..3 {
        for j in 0..4 {
            let value = (i * 10 + j) as i32;
            array.set(&[i, j], value).unwrap();
        }
    }

    // Test fancy indexing with array of indices
    // Select rows 0 and 2
    let result1 = array
        .index(&[IndexSpec::Indices(vec![0, 2]), IndexSpec::All])
        .unwrap();
    assert_eq!(result1.shape(), vec![2, 4]);

    // Check values
    for j in 0..4 {
        assert_eq!(result1.get(&[0, j]).unwrap(), j as i32);
        assert_eq!(result1.get(&[1, j]).unwrap(), (20 + j) as i32);
    }

    // Test fancy indexing with array of indices for columns
    // Select columns 1 and 3
    let result2 = array
        .index(&[IndexSpec::All, IndexSpec::Indices(vec![1, 3])])
        .unwrap();
    assert_eq!(result2.shape(), vec![3, 2]);

    // Check values
    for i in 0..3 {
        assert_eq!(result2.get(&[i, 0]).unwrap(), (i * 10 + 1) as i32);
        assert_eq!(result2.get(&[i, 1]).unwrap(), (i * 10 + 3) as i32);
    }
}

#[test]
fn test_mask_indices() {
    // Test creating mask indices for upper triangle of a 3x3 matrix
    let indices = mask_indices(&[3, 3], |idx| idx[0] <= idx[1]).unwrap();

    // Check that we got indices for two dimensions
    assert_eq!(indices.len(), 2);

    // The first array contains row indices: [0,0,0,1,1,2]
    assert_eq!(indices[0].to_vec(), vec![0, 0, 0, 1, 1, 2]);

    // The second array contains column indices: [0,1,2,1,2,2]
    assert_eq!(indices[1].to_vec(), vec![0, 1, 2, 1, 2, 2]);
}

#[test]
fn test_ix_function() {
    // Create indices for a 2x3 grid
    let indices = indices_grid::<usize>(&[2, 3]).unwrap();

    // Check that we got indices for two dimensions
    assert_eq!(indices.len(), 2);

    // Check the shapes
    assert_eq!(indices[0].shape(), vec![2, 1]);
    assert_eq!(indices[1].shape(), vec![1, 3]);

    // Check values for first dimension: [[0], [1]]
    assert_eq!(indices[0].get(&[0, 0]).unwrap(), 0);
    assert_eq!(indices[0].get(&[1, 0]).unwrap(), 1);

    // Check values for second dimension: [[0, 1, 2]]
    assert_eq!(indices[1].get(&[0, 0]).unwrap(), 0);
    assert_eq!(indices[1].get(&[0, 1]).unwrap(), 1);
    assert_eq!(indices[1].get(&[0, 2]).unwrap(), 2);
}