1use num_traits::Float;
2#[cfg(feature = "serde")]
3use serde::{Deserialize, Serialize};
4use std::{
5 fmt::Debug,
6 ops::{Index, IndexMut},
7};
8
9#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
10pub struct Matrix<T> {
11 data: Vec<T>,
12 rows: usize,
13 cols: usize,
14}
15
16impl<T> Matrix<T> {
17 pub fn new(data: impl Into<Vec<T>>) -> Self {
18 let data = data.into();
19 let rows = data.len();
20
21 Matrix {
22 data,
23 rows,
24 cols: 1,
25 }
26 }
27
28 pub fn empty() -> Self {
29 Matrix {
30 data: Vec::new(),
31 rows: 0,
32 cols: 0,
33 }
34 }
35
36 pub fn from_rows<I>(rows: I) -> Self
37 where
38 I: IntoIterator<Item = Vec<T>>,
39 {
40 let mut iter = rows.into_iter();
41
42 let Some(first) = iter.next() else {
43 return Self {
44 data: Vec::new(),
45 rows: 0,
46 cols: 0,
47 };
48 };
49
50 let cols = first.len();
51 let (lower, upper) = iter.size_hint();
52
53 let mut data = Vec::with_capacity(cols.saturating_mul(1 + upper.unwrap_or(lower)));
54
55 data.extend(first);
56
57 let mut row_count = 1;
58
59 for row in iter {
60 debug_assert!(row.len() == cols);
61 data.extend(row);
62 row_count += 1;
63 }
64
65 Self {
66 data,
67 rows: row_count,
68 cols,
69 }
70 }
71
72 pub fn rows(&self) -> usize {
73 self.rows
74 }
75
76 pub fn cols(&self) -> usize {
77 self.cols
78 }
79
80 pub fn data(&self) -> &[T] {
81 &self.data
82 }
83
84 pub fn is_empty(&self) -> bool {
85 self.data.is_empty()
86 }
87
88 pub fn size(&self) -> usize {
89 self.data.len()
90 }
91
92 pub fn clear(&mut self) {
93 self.data.clear();
94 self.rows = 0;
95 self.cols = 0;
96 }
97
98 pub fn iter(&self) -> impl Iterator<Item = &[T]> {
99 self.data.chunks(self.cols)
100 }
101
102 pub fn row(&self, row: usize) -> &[T] {
103 let start = row * self.cols;
104 let end = start + self.cols;
105
106 &self.data[start..end]
107 }
108
109 pub fn row_mut(&mut self, row: usize) -> &mut [T] {
110 let start = row * self.cols;
111 let end = start + self.cols;
112
113 &mut self.data[start..end]
114 }
115
116 pub fn append_row(&mut self, row_data: Vec<T>) {
117 if self.is_empty() {
118 self.cols = row_data.len();
119 } else {
120 debug_assert!(
121 row_data.len() == self.cols,
122 "Row length must match the number of columns"
123 );
124 }
125
126 self.data.extend(row_data);
127 self.rows += 1;
128 }
129
130 pub fn reshape(mut self, new_rows: usize, new_cols: usize) -> Self {
131 assert!(
132 new_rows * new_cols == self.data.len(),
133 "New dimensions must match the total number of elements"
134 );
135 self.rows = new_rows;
136 self.cols = new_cols;
137 self
138 }
139
140 pub fn reshape_in_place(&mut self, new_rows: usize, new_cols: usize) {
141 assert!(
142 new_rows * new_cols == self.data.len(),
143 "New dimensions must match the total number of elements"
144 );
145 self.rows = new_rows;
146 self.cols = new_cols;
147 }
148}
149
150impl<T: Clone> Matrix<T> {
151 pub fn append_column(&mut self, col_data: Vec<T>) {
152 debug_assert!(col_data.len() == self.rows);
153 let mut new_data = Vec::with_capacity((self.rows + 1) * self.cols);
154 for (row, col) in self.data.chunks(self.cols).zip(col_data) {
155 new_data.extend_from_slice(row);
156 new_data.push(col);
157 }
158 self.data = new_data;
159 self.cols += 1;
160 }
161
162 pub fn reshape_and_fill(&mut self, rows: usize, cols: usize, default_value: T) {
163 let new_size = rows * cols;
164 self.data.clear();
165 self.data.resize(new_size, default_value);
166 self.rows = rows;
167 self.cols = cols;
168 }
169
170 pub fn sort_by_indices(&self, indices: &[usize]) -> Self {
171 debug_assert!(
172 indices.iter().all(|&i| i < self.rows),
173 "row index out of bounds"
174 );
175
176 let mut data = Vec::with_capacity(indices.len() * self.cols);
177
178 for &row in indices {
179 let start = row * self.cols;
180 let end = start + self.cols;
181 data.extend_from_slice(&self.data[start..end]);
182 }
183
184 Matrix {
185 data,
186 rows: indices.len(),
187 cols: self.cols,
188 }
189 }
190
191 pub fn split_at_row(&self, row: usize) -> (Self, Self) {
192 debug_assert!(row <= self.rows);
193
194 let first_part_data = self.data[0..row * self.cols].to_vec();
195 let second_part_data = self.data[row * self.cols..].to_vec();
196
197 let first_part = Matrix {
198 data: first_part_data,
199 rows: row,
200 cols: self.cols,
201 };
202
203 let second_part = Matrix {
204 data: second_part_data,
205 rows: self.rows - row,
206 cols: self.cols,
207 };
208
209 (first_part, second_part)
210 }
211
212 pub fn fill(&mut self, value: T) {
213 self.data.fill(value);
214 }
215
216 pub fn transpose(&self) -> Self {
217 let mut transposed_data = Vec::with_capacity(self.data.len());
218
219 for col in 0..self.cols {
220 for row in 0..self.rows {
221 transposed_data.push(self[(row, col)].clone());
222 }
223 }
224
225 Matrix {
226 data: transposed_data,
227 rows: self.cols,
228 cols: self.rows,
229 }
230 }
231}
232
233impl<T: Float> Matrix<T> {
234 pub fn standardize(&mut self) {
235 for col in 0..self.cols {
236 let mut sum = T::zero();
237 for row in 0..self.rows {
238 sum = sum + self[(row, col)];
239 }
240 let mean = sum / T::from(self.rows).unwrap();
241
242 let mut variance_sum = T::zero();
243 for row in 0..self.rows {
244 let diff = self[(row, col)] - mean;
245 variance_sum = variance_sum + diff * diff;
246 }
247
248 let variance = variance_sum / T::from(self.rows).unwrap();
249 let std_dev = variance.sqrt();
250
251 if std_dev <= T::zero() {
252 continue;
253 }
254
255 for row in 0..self.rows {
256 self[(row, col)] = (self[(row, col)] - mean) / std_dev;
257 }
258 }
259 }
260
261 pub fn normalize(&mut self) {
262 for col in 0..self.cols {
263 let mut min = self[(0, col)];
264 let mut max = self[(0, col)];
265
266 for row in 1..self.rows {
267 if self[(row, col)] < min {
268 min = self[(row, col)];
269 }
270 if self[(row, col)] > max {
271 max = self[(row, col)];
272 }
273 }
274
275 let range = max - min;
276
277 if range <= T::zero() {
278 continue;
279 }
280
281 for row in 0..self.rows {
282 self[(row, col)] = (self[(row, col)] - min) / range;
283 }
284 }
285 }
286}
287
288impl<T> AsRef<[T]> for Matrix<T> {
289 fn as_ref(&self) -> &[T] {
290 &self.data
291 }
292}
293
294impl<T> Index<usize> for Matrix<T> {
295 type Output = [T];
296
297 fn index(&self, row: usize) -> &Self::Output {
298 self.row(row)
299 }
300}
301
302impl<T> Index<(usize, usize)> for Matrix<T> {
303 type Output = T;
304
305 fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
306 let flat_index = row * self.cols + col;
307 &self.data[flat_index]
308 }
309}
310
311impl<T> Index<[usize; 2]> for Matrix<T> {
312 type Output = T;
313
314 fn index(&self, index: [usize; 2]) -> &Self::Output {
315 let flat_index = index[0] * self.cols + index[1];
316 &self.data[flat_index]
317 }
318}
319
320impl<T> IndexMut<usize> for Matrix<T> {
321 fn index_mut(&mut self, row: usize) -> &mut Self::Output {
322 self.row_mut(row)
323 }
324}
325
326impl<T> IndexMut<(usize, usize)> for Matrix<T> {
327 fn index_mut(&mut self, (row, col): (usize, usize)) -> &mut Self::Output {
328 let flat_index = row * self.cols + col;
329 &mut self.data[flat_index]
330 }
331}
332
333impl<T> IndexMut<[usize; 2]> for Matrix<T> {
334 fn index_mut(&mut self, index: [usize; 2]) -> &mut Self::Output {
335 let flat_index = index[0] * self.cols + index[1];
336 &mut self.data[flat_index]
337 }
338}
339
340impl<T: Default + Clone> From<(usize, usize)> for Matrix<T> {
341 fn from((rows, cols): (usize, usize)) -> Self {
342 let data = vec![T::default(); rows * cols];
343 Matrix { data, rows, cols }
344 }
345}
346
347impl<T> From<Vec<T>> for Matrix<T> {
348 fn from(vec: Vec<T>) -> Self {
349 let rows = vec.len();
350 let cols = 1;
351 Matrix {
352 data: vec,
353 rows,
354 cols,
355 }
356 }
357}
358
359impl<T> From<Vec<Vec<T>>> for Matrix<T> {
360 fn from(vec_of_vecs: Vec<Vec<T>>) -> Self {
361 let rows = vec_of_vecs.len();
362 let cols = if rows > 0 { vec_of_vecs[0].len() } else { 0 };
363 let mut data = Vec::with_capacity(rows * cols);
364
365 for row in vec_of_vecs.into_iter() {
366 assert!(
367 row.len() == cols,
368 "All rows must have the same number of columns"
369 );
370 data.extend(row);
371 }
372
373 Matrix { data, rows, cols }
374 }
375}
376
377impl<T> From<(usize, usize, Vec<T>)> for Matrix<T> {
378 fn from((rows, cols, data): (usize, usize, Vec<T>)) -> Self {
379 assert!(
380 rows * cols == data.len(),
381 "Data length must match rows * cols"
382 );
383 Matrix { data, rows, cols }
384 }
385}
386
387impl<T> FromIterator<Vec<T>> for Matrix<T> {
388 fn from_iter<I: IntoIterator<Item = Vec<T>>>(iter: I) -> Self {
389 Matrix::from_rows(iter)
390 }
391}
392
393impl<T: Clone> Clone for Matrix<T> {
394 fn clone(&self) -> Self {
395 Matrix {
396 data: self.data.clone(),
397 rows: self.rows,
398 cols: self.cols,
399 }
400 }
401}
402
403impl<T: Default> Default for Matrix<T> {
404 fn default() -> Self {
405 Matrix {
406 data: Vec::new(),
407 rows: 0,
408 cols: 0,
409 }
410 }
411}
412
413impl<T: PartialEq> PartialEq for Matrix<T> {
414 fn eq(&self, other: &Self) -> bool {
415 self.rows == other.rows && self.cols == other.cols && self.data == other.data
416 }
417}
418
419impl<T: Debug> Debug for Matrix<T> {
420 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
421 const MAX_ROWS: usize = 8;
422 const MAX_COLS: usize = 8;
423
424 writeln!(f, "Matrix({} × {}) {{", self.rows, self.cols)?;
425
426 let rows = self.rows.min(MAX_ROWS);
427 let cols = self.cols.min(MAX_COLS);
428
429 for row in 0..rows {
430 write!(f, " [")?;
431
432 for col in 0..cols {
433 if col > 0 {
434 write!(f, ", ")?;
435 }
436
437 write!(f, "{:?}", self[(row, col)])?;
438 }
439
440 if cols < self.cols {
441 write!(f, ", ...")?;
442 }
443
444 writeln!(f, "]")?;
445 }
446
447 if rows < self.rows {
448 writeln!(f, " ...")?;
449 }
450
451 write!(f, "}}")
452 }
453}
454
455#[cfg(test)]
456mod tests {
457 use super::*;
458
459 #[test]
460 fn test_matrix_indexing() {
461 let data = vec![1, 2, 3, 4, 5, 6];
462 let matrix = Matrix::new(data).reshape(2, 3);
463
464 assert_eq!(matrix[(0, 0)], 1);
465 assert_eq!(matrix[(0, 1)], 2);
466 assert_eq!(matrix[(1, 0)], 4);
467 assert_eq!(matrix[[1, 2]], 6);
468 }
469
470 #[test]
471 fn test_matrix_row_access() {
472 let data = vec![1, 2, 3, 4, 5, 6];
473 let matrix = Matrix::new(data).reshape(2, 3);
474
475 assert_eq!(matrix.row(0), &[1, 2, 3]);
476 assert_eq!(matrix.row(1), &[4, 5, 6]);
477 }
478
479 #[test]
480 fn test_matrix_append_row() {
481 let mut matrix = Matrix::new(vec![1, 2, 3, 4, 5, 6]).reshape(2, 3);
482 matrix.append_row(vec![7, 8, 9]);
483
484 assert_eq!(matrix.rows(), 3);
485 assert_eq!(matrix.row(2), &[7, 8, 9]);
486 }
487
488 #[test]
489 fn test_matrix_append_column() {
490 let mut matrix = Matrix::new([1, 2, 3, 4]).reshape(2, 2);
491 matrix.append_column(vec![5, 6]);
492
493 assert_eq!(matrix.cols(), 3);
494 assert_eq!(matrix.row(0), &[1, 2, 5]);
495 assert_eq!(matrix.row(1), &[3, 4, 6]);
496 }
497
498 #[test]
499 fn test_matrix_reshape() {
500 let matrix = Matrix::new(vec![1, 2, 3, 4, 5, 6]).reshape(3, 2);
501
502 assert_eq!(matrix.rows(), 3);
503 assert_eq!(matrix.cols(), 2);
504 assert_eq!(matrix.row(0), &[1, 2]);
505 assert_eq!(matrix.row(1), &[3, 4]);
506 assert_eq!(matrix.row(2), &[5, 6]);
507 }
508
509 #[test]
510 fn test_matrix_split_at_row() {
511 let matrix = Matrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8]).reshape(4, 2);
512 let (first_part, second_part) = matrix.split_at_row(2);
513
514 assert_eq!(first_part.rows(), 2);
515 assert_eq!(first_part.cols(), 2);
516 assert_eq!(first_part.row(0), &[1, 2]);
517 assert_eq!(first_part.row(1), &[3, 4]);
518 assert_eq!(second_part.rows(), 2);
519 assert_eq!(second_part.cols(), 2);
520 assert_eq!(second_part.row(0), &[5, 6]);
521 assert_eq!(second_part.row(1), &[7, 8]);
522 }
523}