1use std::ops::Index;
2
3use rayon::prelude::ParallelSlice;
4use rayon::prelude::ParallelSliceMut;
5
6#[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 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 #[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 #[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 #[inline]
110 pub fn rows<'a>(&'a self) -> std::slice::ChunksExact<'a, T> {
111 self.values.chunks_exact(self.n_cols)
112 }
113
114 #[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 #[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 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 pub unsafe fn get_unchecked(&self, (i, j): (usize, usize)) -> &T {
204 self.values.get_unchecked(self.n_cols * i + j)
205 }
206
207 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 pub unsafe fn get_unchecked(&self, (i, j): (usize, usize)) -> &T {
305 self.values.get_unchecked(self.n_rows * j + i)
306 }
307
308 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}