parry3d/utils/
array2.rs

1use alloc::{vec, vec::Vec};
2use core::ops::{Index, IndexMut};
3use num::{Bounded, Zero};
4
5#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
6#[derive(Debug, Clone)]
7/// A column-major 2D array.
8pub struct Array2<T> {
9    nrows: usize,
10    ncols: usize,
11    data: Vec<T>,
12}
13
14impl<T> Array2<T> {
15    /// Creates a new 2D array from the given dimensions and data.
16    #[inline]
17    pub fn new(nrows: usize, ncols: usize, data: Vec<T>) -> Self {
18        assert_eq!(data.len(), ncols * nrows);
19        Array2 { ncols, nrows, data }
20    }
21
22    /// Creates a new 2D array filled with copies of `elt`.
23    #[inline]
24    pub fn repeat(nrows: usize, ncols: usize, elt: T) -> Self
25    where
26        T: Copy,
27    {
28        Array2 {
29            nrows,
30            ncols,
31            data: vec![elt; ncols * nrows],
32        }
33    }
34
35    /// Creates a new 2D array filled with zeros.
36    #[inline]
37    pub fn zeros(nrows: usize, ncols: usize) -> Self
38    where
39        T: Copy + Zero,
40    {
41        Self::repeat(nrows, ncols, T::zero())
42    }
43
44    /// Computes the flat index for the given row and column.
45    #[inline]
46    pub fn flat_index(&self, i: usize, j: usize) -> usize {
47        i + j * self.nrows
48    }
49
50    /// Returns a reference to the underlying data.
51    #[inline]
52    pub fn data(&self) -> &[T] {
53        &self.data
54    }
55
56    /// Returns a mutable reference to the underlying data.
57    #[inline]
58    pub fn data_mut(&mut self) -> &mut [T] {
59        &mut self.data
60    }
61
62    /// Returns the number of rows.
63    #[inline]
64    pub fn nrows(&self) -> usize {
65        self.nrows
66    }
67
68    /// Returns the number of columns.
69    #[inline]
70    pub fn ncols(&self) -> usize {
71        self.ncols
72    }
73
74    /// Returns the minimum element of the array.
75    #[inline]
76    pub fn min(&self) -> T
77    where
78        T: Copy + Bounded + PartialOrd,
79    {
80        self.data
81            .iter()
82            .fold(Bounded::max_value(), |a, b| if a < *b { a } else { *b })
83    }
84
85    /// Returns the maximum element of the array.
86    #[inline]
87    pub fn max(&self) -> T
88    where
89        T: Copy + Bounded + PartialOrd,
90    {
91        self.data
92            .iter()
93            .fold(Bounded::min_value(), |a, b| if a > *b { a } else { *b })
94    }
95
96    /// Creates a new 2D array from a function.
97    #[inline]
98    pub fn from_fn(nrows: usize, ncols: usize, f: impl Fn(usize, usize) -> T) -> Self {
99        let mut data = Vec::with_capacity(ncols * nrows);
100
101        for j in 0..ncols {
102            for i in 0..nrows {
103                data.push(f(i, j));
104            }
105        }
106        data.shrink_to_fit();
107        Self { nrows, ncols, data }
108    }
109}
110
111impl<T> Index<(usize, usize)> for Array2<T> {
112    type Output = T;
113
114    #[inline]
115    fn index(&self, index: (usize, usize)) -> &Self::Output {
116        let fid = self.flat_index(index.0, index.1);
117        &self.data[fid]
118    }
119}
120
121impl<T> IndexMut<(usize, usize)> for Array2<T> {
122    #[inline]
123    fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
124        let fid = self.flat_index(index.0, index.1);
125        &mut self.data[fid]
126    }
127}