Skip to main content

lace/utils/
matrix.rs

1use std::ops::Index;
2
3use rayon::prelude::ParallelSlice;
4use rayon::prelude::ParallelSliceMut;
5
6/// A lightweight Matrix abstraction that does almost nothing.
7#[derive(Clone, Debug)]
8pub struct Matrix<T: Send + Sync> {
9    n_rows: usize,
10    n_cols: usize,
11    values: Vec<T>,
12}
13
14impl<T: Send + Sync> Matrix<T> {
15    pub fn from_raw_parts(values: Vec<T>, n_rows: usize) -> Self {
16        let n_cols = values.len() / n_rows;
17        assert_eq!(values.len(), n_rows * n_cols);
18        Matrix {
19            n_rows,
20            n_cols,
21            values,
22        }
23    }
24
25    /// Create a new Matrix from a vector of vectors
26    pub fn from_vecs(mut vecs: Vec<Vec<T>>) -> Self {
27        let n_rows = vecs.len();
28        let n_cols = vecs[0].len();
29        let mut values = Vec::with_capacity(n_rows * n_cols);
30
31        vecs.drain(..).for_each(|mut row| {
32            row.drain(..).for_each(|x| values.push(x));
33        });
34
35        Matrix {
36            n_rows,
37            n_cols,
38            values,
39        }
40    }
41
42    #[inline]
43    pub fn nelem(&self) -> usize {
44        self.n_cols * self.n_rows
45    }
46
47    #[inline]
48    pub fn raw_values(&self) -> &Vec<T> {
49        &self.values
50    }
51
52    #[inline]
53    pub fn raw_values_mut(&mut self) -> &mut Vec<T> {
54        &mut self.values
55    }
56
57    /// Create a mutable iterator through rows
58    ///
59    /// # Example
60    ///
61    /// ```rust
62    /// # use lace::utils::Matrix;
63    /// let vecs: Vec<Vec<u8>> = vec![
64    ///     vec![0, 1, 2],
65    ///     vec![3, 4, 5],
66    /// ];
67    ///
68    /// let mut mat = Matrix::from_vecs(vecs);
69    ///
70    /// mat.rows_mut().for_each(|mut row| {
71    ///     row.iter_mut().for_each(|mut x| *x += 1 );
72    /// });
73    ///
74    /// assert_eq!(mat.raw_values(), &vec![1, 2, 3, 4, 5, 6])
75    /// ```
76    #[inline]
77    pub fn rows_mut<'a>(&'a mut self) -> std::slice::ChunksExactMut<'a, T> {
78        self.values.chunks_exact_mut(self.n_cols)
79    }
80
81    /// Parallel version of `rows_mut`
82    #[inline]
83    pub fn par_rows_mut<'a>(
84        &'a mut self,
85    ) -> rayon::slice::ChunksExactMut<'a, T> {
86        self.values.par_chunks_exact_mut(self.n_cols)
87    }
88
89    /// Create an iterator through rows
90    ///
91    /// # Example
92    ///
93    /// ```rust
94    /// # use lace::utils::Matrix;
95    /// let vecs: Vec<Vec<u8>> = vec![
96    ///     vec![0, 1, 2],
97    ///     vec![3, 4, 5],
98    /// ];
99    ///
100    /// let mut mat = Matrix::from_vecs(vecs);
101    ///
102    /// let rowsum: Vec<u8> = mat.rows().map(|row| {
103    ///     row.iter().sum::<u8>()
104    /// })
105    /// .collect();
106    ///
107    /// assert_eq!(rowsum, vec![3_u8, 12_u8])
108    /// ```
109    #[inline]
110    pub fn rows<'a>(&'a self) -> std::slice::ChunksExact<'a, T> {
111        self.values.chunks_exact(self.n_cols)
112    }
113
114    /// Parallel version of `rows`
115    #[inline]
116    pub fn par_rows<'a>(&'a self) -> rayon::slice::ChunksExact<'a, T> {
117        self.values.par_chunks_exact(self.n_cols)
118    }
119
120    /// Does an implicit transpose by inverting coordinates.
121    ///
122    /// # Notes
123    /// The matrix is not rebuild, so if you are attempting to access the
124    /// transposed matrix row-wise many times, there will be a lot of cache
125    /// misses. This method is best used if you plan to traverse the transposed
126    /// matrix a small number of times.
127    ///
128    /// # Example
129    ///
130    /// ```rust
131    /// # use lace::utils::Matrix;
132    /// use lace::utils::Shape;
133    ///
134    /// let vecs: Vec<Vec<u32>> = vec![
135    ///     vec![0, 1, 2],
136    ///     vec![3, 4, 5],
137    ///     vec![6, 7, 8],
138    ///     vec![9, 10, 11],
139    /// ];
140    ///
141    /// let mat = Matrix::from_vecs(vecs);
142    /// let mat_t = mat.clone().implicit_transpose();
143    ///
144    /// assert_eq!(mat.n_rows(), mat_t.n_cols());
145    /// assert_eq!(mat.n_cols(), mat_t.n_rows());
146    ///
147    /// for i in 0..3 {
148    ///     for j in 0..3 {
149    ///         assert_eq!(mat[(i, j)], mat_t[(j, i)]);
150    ///     }
151    /// }
152    ///
153    /// ```
154    #[inline]
155    pub fn implicit_transpose(self) -> ImplicitlyTransposedMatrix<T> {
156        ImplicitlyTransposedMatrix {
157            n_rows: self.n_cols,
158            n_cols: self.n_rows,
159            values: self.values,
160        }
161    }
162}
163
164impl<T: Send + Sync + Clone> Matrix<T> {
165    /// Treat the input vector, `col` like a column vector and replicate it
166    /// `n_cols` times.
167    ///
168    /// # Example
169    ///
170    /// ```rust
171    /// # use lace::utils::Matrix;
172    /// let col: Vec<u32> = vec![0, 1, 2];
173    ///
174    /// let mat = Matrix::vtile(col, 12);
175    ///
176    /// assert_eq!(mat[(0, 0)], 0);
177    /// assert_eq!(mat[(0, 11)], 0);
178    ///
179    /// assert_eq!(mat[(1, 0)], 1);
180    /// assert_eq!(mat[(1, 11)], 1);
181    ///
182    /// assert_eq!(mat[(2, 0)], 2);
183    /// assert_eq!(mat[(2, 11)], 2);
184    /// ```
185    pub fn vtile(col: Vec<T>, n_cols: usize) -> Self {
186        let n_rows = col.len();
187        let mut values: Vec<T> = Vec::with_capacity(n_rows * n_cols);
188        col.iter().for_each(|x| {
189            (0..n_cols).for_each(|_| values.push(x.clone()));
190        });
191
192        Matrix {
193            n_rows,
194            n_cols,
195            values,
196        }
197    }
198
199    /// Get a reference to the item at (i, j) without bounds checks
200    ///
201    /// # Safety
202    /// Behavior is undefined if (i, j) is out of bounds
203    pub unsafe fn get_unchecked(&self, (i, j): (usize, usize)) -> &T {
204        self.values.get_unchecked(self.n_cols * i + j)
205    }
206
207    /// Get a mutable reference to the item at (i, j) without bounds checks
208    ///
209    /// # Safety
210    /// Behavior is undefined if (i, j) is out of bounds
211    pub unsafe fn get_unchecked_mut(
212        &mut self,
213        (i, j): (usize, usize),
214    ) -> &mut T {
215        self.values.get_unchecked_mut(self.n_cols * i + j)
216    }
217}
218
219impl<T> Index<(usize, usize)> for Matrix<T>
220where
221    T: Send + Sync,
222{
223    type Output = T;
224
225    #[inline]
226    fn index(&self, (i, j): (usize, usize)) -> &Self::Output {
227        &self.values[self.n_cols * i + j]
228    }
229}
230
231impl<T> Index<(usize, usize)> for &Matrix<T>
232where
233    T: Send + Sync,
234{
235    type Output = T;
236
237    #[inline]
238    fn index(&self, (i, j): (usize, usize)) -> &Self::Output {
239        &self.values[self.n_cols * i + j]
240    }
241}
242
243#[derive(Clone, Debug)]
244pub struct ImplicitlyTransposedMatrix<T: Send + Sync> {
245    n_rows: usize,
246    n_cols: usize,
247    values: Vec<T>,
248}
249
250impl<T> Index<(usize, usize)> for ImplicitlyTransposedMatrix<T>
251where
252    T: Send + Sync,
253{
254    type Output = T;
255
256    #[inline]
257    fn index(&self, ix: (usize, usize)) -> &Self::Output {
258        let (i, j) = ix;
259        &self.values[self.n_rows * j + i]
260    }
261}
262
263impl<T> Index<(usize, usize)> for &ImplicitlyTransposedMatrix<T>
264where
265    T: Send + Sync,
266{
267    type Output = T;
268
269    #[inline]
270    fn index(&self, ix: (usize, usize)) -> &Self::Output {
271        let (i, j) = ix;
272        &self.values[self.n_rows * j + i]
273    }
274}
275
276impl<T: Send + Sync> ImplicitlyTransposedMatrix<T> {
277    #[inline]
278    pub fn nelem(&self) -> usize {
279        self.n_cols * self.n_rows
280    }
281
282    #[inline]
283    pub fn raw_values(&self) -> &Vec<T> {
284        &self.values
285    }
286
287    #[inline]
288    pub fn raw_values_mut(&mut self) -> &mut Vec<T> {
289        &mut self.values
290    }
291
292    pub fn transpose(self) -> Matrix<T> {
293        Matrix {
294            n_rows: self.n_cols,
295            n_cols: self.n_rows,
296            values: self.values,
297        }
298    }
299
300    /// Get a reference to the entry at index without bounds checks
301    ///
302    /// # Safety
303    /// Undefined behavior if the index is out of bounds
304    pub unsafe fn get_unchecked(&self, (i, j): (usize, usize)) -> &T {
305        self.values.get_unchecked(self.n_rows * j + i)
306    }
307
308    /// Get a mutable reference to the entry at index without bounds checks
309    ///
310    /// # Safety
311    /// Undefined behavior if the index is out of bounds
312    pub unsafe fn get_unchecked_mut(
313        &mut self,
314        (i, j): (usize, usize),
315    ) -> &mut T {
316        self.values.get_unchecked_mut(self.n_rows * j + i)
317    }
318}
319
320pub trait Shape {
321    fn n_rows(&self) -> usize;
322    fn n_cols(&self) -> usize;
323    fn shape(&self) -> (usize, usize) {
324        (self.n_rows(), self.n_cols())
325    }
326}
327
328impl<T: Send + Sync> Shape for Matrix<T> {
329    fn n_rows(&self) -> usize {
330        self.n_rows
331    }
332
333    fn n_cols(&self) -> usize {
334        self.n_cols
335    }
336}
337
338impl<T: Send + Sync> Shape for &Matrix<T> {
339    fn n_rows(&self) -> usize {
340        self.n_rows
341    }
342
343    fn n_cols(&self) -> usize {
344        self.n_cols
345    }
346}
347
348impl<T: Send + Sync> Shape for ImplicitlyTransposedMatrix<T> {
349    fn n_rows(&self) -> usize {
350        self.n_rows
351    }
352
353    fn n_cols(&self) -> usize {
354        self.n_cols
355    }
356}
357
358impl<T: Send + Sync> Shape for &ImplicitlyTransposedMatrix<T> {
359    fn n_rows(&self) -> usize {
360        self.n_rows
361    }
362
363    fn n_cols(&self) -> usize {
364        self.n_cols
365    }
366}