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

use ndarray::{prelude::*, Shape};

pub struct NdIndexIterator<D: Dimension> {
    shape: Shape<D>,
    counter: usize,
}
impl<D> Iterator for NdIndexIterator<D>
where
    D: Dimension,
{
    type Item = Array1<usize>;

    fn next(&mut self) -> Option<Self::Item> {
        let max = self.shape.raw_dim().size();
        if self.counter < max {
            let dimensions = self.shape.raw_dim().slice();
            let n_dimensions = dimensions.len();
            let mut result = Array::<usize, Ix1>::zeros(n_dimensions);
            let mut counter = self.counter;
            for (i, d) in dimensions.iter().rev().enumerate() {
                // println!("{}: {} / {} = {} + {}", i, counter, d, counter /d , counter % d);

                result[n_dimensions - i - 1] = counter % d;
                counter /= d;
            }
            self.counter += 1;
            Some(result)
        } else {
            None
        }
    }
}

/// Creates an iterator that generates indices for an array of given `shape`
pub fn ndindex<Sh>(shape: Sh) -> NdIndexIterator<Sh::Dim>
where
    Sh: ShapeBuilder,
{
    let shape = shape.into_shape(); //.raw_dim();
    NdIndexIterator {
        shape: shape,
        counter: 0,
    }
}

// checkout indexed_iter
/// Creates an array with rows that hold the indices generated by `ndindex`
pub fn get_ndindex_array<D>(shape: &Shape<D>) -> Array2<f64>
where
    D: Dimension,
{
    // let shape = shape.into_shape();
    let dim = shape.raw_dim();
    let (m, n) = (dim.size(), dim.slice().len());
    let mut result = Array2::<f64>::zeros((m, n));

    let index_iterator = NdIndexIterator {
        shape: shape.clone(),
        counter: 0,
    };

    for (mut r, i) in result.outer_iter_mut().zip(index_iterator) {
        // it would be nicer to use f64::from + u32::try_from ... learn more about error/result handling! .. there's something in num_traits
        r.assign(&i.mapv(|e| e as f64));
    }

    result
}

#[cfg(test)]
mod tests {

    #[test]
    fn it_works() {
        let result = 2 + 2;
        assert_eq!(result, 4);
    }
}