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 let variance = variance_sum / T::from(self.rows).unwrap();
248 let std_dev = variance.sqrt();
249
250 if std_dev <= T::zero() {
251 continue;
252 }
253
254 for row in 0..self.rows {
255 self[(row, col)] = (self[(row, col)] - mean) / std_dev;
256 }
257 }
258 }
259
260 pub fn normalize(&mut self) {
261 for col in 0..self.cols {
262 let mut min = self[(0, col)];
263 let mut max = self[(0, col)];
264
265 for row in 1..self.rows {
266 if self[(row, col)] < min {
267 min = self[(row, col)];
268 }
269 if self[(row, col)] > max {
270 max = self[(row, col)];
271 }
272 }
273
274 let range = max - min;
275
276 if range <= T::zero() {
277 continue;
278 }
279
280 for row in 0..self.rows {
281 self[(row, col)] = (self[(row, col)] - min) / range;
282 }
283 }
284 }
285}
286
287impl<T> AsRef<[T]> for Matrix<T> {
288 fn as_ref(&self) -> &[T] {
289 &self.data
290 }
291}
292
293impl<T> Index<usize> for Matrix<T> {
294 type Output = [T];
295
296 fn index(&self, row: usize) -> &Self::Output {
297 self.row(row)
298 }
299}
300
301impl<T> Index<(usize, usize)> for Matrix<T> {
302 type Output = T;
303
304 fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
305 let flat_index = row * self.cols + col;
306 &self.data[flat_index]
307 }
308}
309
310impl<T> Index<[usize; 2]> for Matrix<T> {
311 type Output = T;
312
313 fn index(&self, index: [usize; 2]) -> &Self::Output {
314 let flat_index = index[0] * self.cols + index[1];
315 &self.data[flat_index]
316 }
317}
318
319impl<T> IndexMut<usize> for Matrix<T> {
320 fn index_mut(&mut self, row: usize) -> &mut Self::Output {
321 self.row_mut(row)
322 }
323}
324
325impl<T> IndexMut<(usize, usize)> for Matrix<T> {
326 fn index_mut(&mut self, (row, col): (usize, usize)) -> &mut Self::Output {
327 let flat_index = row * self.cols + col;
328 &mut self.data[flat_index]
329 }
330}
331
332impl<T> IndexMut<[usize; 2]> for Matrix<T> {
333 fn index_mut(&mut self, index: [usize; 2]) -> &mut Self::Output {
334 let flat_index = index[0] * self.cols + index[1];
335 &mut self.data[flat_index]
336 }
337}
338
339impl<T: Default + Clone> From<(usize, usize)> for Matrix<T> {
340 fn from((rows, cols): (usize, usize)) -> Self {
341 let data = vec![T::default(); rows * cols];
342 Matrix { data, rows, cols }
343 }
344}
345
346impl<T> From<Vec<T>> for Matrix<T> {
347 fn from(vec: Vec<T>) -> Self {
348 let rows = vec.len();
349 let cols = 1;
350 Matrix {
351 data: vec,
352 rows,
353 cols,
354 }
355 }
356}
357
358impl<T> From<Vec<Vec<T>>> for Matrix<T> {
359 fn from(vec_of_vecs: Vec<Vec<T>>) -> Self {
360 let rows = vec_of_vecs.len();
361 let cols = if rows > 0 { vec_of_vecs[0].len() } else { 0 };
362 let mut data = Vec::with_capacity(rows * cols);
363
364 for row in vec_of_vecs.into_iter() {
365 assert!(
366 row.len() == cols,
367 "All rows must have the same number of columns"
368 );
369 data.extend(row);
370 }
371
372 Matrix { data, rows, cols }
373 }
374}
375
376impl<T> From<(usize, usize, Vec<T>)> for Matrix<T> {
377 fn from((rows, cols, data): (usize, usize, Vec<T>)) -> Self {
378 assert!(
379 rows * cols == data.len(),
380 "Data length must match rows * cols"
381 );
382 Matrix { data, rows, cols }
383 }
384}
385
386impl<T> FromIterator<Vec<T>> for Matrix<T> {
387 fn from_iter<I: IntoIterator<Item = Vec<T>>>(iter: I) -> Self {
388 Matrix::from_rows(iter)
389 }
390}
391
392impl<T: Clone> Clone for Matrix<T> {
393 fn clone(&self) -> Self {
394 Matrix {
395 data: self.data.clone(),
396 rows: self.rows,
397 cols: self.cols,
398 }
399 }
400}
401
402impl<T: Default> Default for Matrix<T> {
403 fn default() -> Self {
404 Matrix {
405 data: Vec::new(),
406 rows: 0,
407 cols: 0,
408 }
409 }
410}
411
412impl<T: PartialEq> PartialEq for Matrix<T> {
413 fn eq(&self, other: &Self) -> bool {
414 self.rows == other.rows && self.cols == other.cols && self.data == other.data
415 }
416}
417
418impl<T: Debug> Debug for Matrix<T> {
419 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
420 const MAX_ROWS: usize = 8;
421 const MAX_COLS: usize = 8;
422
423 writeln!(f, "Matrix({} × {}) {{", self.rows, self.cols)?;
424
425 let rows = self.rows.min(MAX_ROWS);
426 let cols = self.cols.min(MAX_COLS);
427
428 for row in 0..rows {
429 write!(f, " [")?;
430
431 for col in 0..cols {
432 if col > 0 {
433 write!(f, ", ")?;
434 }
435
436 write!(f, "{:?}", self[(row, col)])?;
437 }
438
439 if cols < self.cols {
440 write!(f, ", ...")?;
441 }
442
443 writeln!(f, "]")?;
444 }
445
446 if rows < self.rows {
447 writeln!(f, " ...")?;
448 }
449
450 write!(f, "}}")
451 }
452}
453
454#[cfg(test)]
455mod tests {
456 use super::*;
457
458 #[test]
459 fn test_matrix_indexing() {
460 let data = vec![1, 2, 3, 4, 5, 6];
461 let matrix = Matrix::new(data).reshape(2, 3);
462
463 assert_eq!(matrix[(0, 0)], 1);
464 assert_eq!(matrix[(0, 1)], 2);
465 assert_eq!(matrix[(1, 0)], 4);
466 assert_eq!(matrix[[1, 2]], 6);
467 }
468
469 #[test]
470 fn test_matrix_row_access() {
471 let data = vec![1, 2, 3, 4, 5, 6];
472 let matrix = Matrix::new(data).reshape(2, 3);
473
474 assert_eq!(matrix.row(0), &[1, 2, 3]);
475 assert_eq!(matrix.row(1), &[4, 5, 6]);
476 }
477
478 #[test]
479 fn test_matrix_append_row() {
480 let mut matrix = Matrix::new(vec![1, 2, 3, 4, 5, 6]).reshape(2, 3);
481 matrix.append_row(vec![7, 8, 9]);
482
483 assert_eq!(matrix.rows(), 3);
484 assert_eq!(matrix.row(2), &[7, 8, 9]);
485 }
486
487 #[test]
488 fn test_matrix_append_column() {
489 let mut matrix = Matrix::new([1, 2, 3, 4]).reshape(2, 2);
490 matrix.append_column(vec![5, 6]);
491
492 assert_eq!(matrix.cols(), 3);
493 assert_eq!(matrix.row(0), &[1, 2, 5]);
494 assert_eq!(matrix.row(1), &[3, 4, 6]);
495 }
496
497 #[test]
498 fn test_matrix_reshape() {
499 let matrix = Matrix::new(vec![1, 2, 3, 4, 5, 6]).reshape(3, 2);
500
501 assert_eq!(matrix.rows(), 3);
502 assert_eq!(matrix.cols(), 2);
503 assert_eq!(matrix.row(0), &[1, 2]);
504 assert_eq!(matrix.row(1), &[3, 4]);
505 assert_eq!(matrix.row(2), &[5, 6]);
506 }
507
508 #[test]
509 fn test_matrix_split_at_row() {
510 let matrix = Matrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8]).reshape(4, 2);
511 let (first_part, second_part) = matrix.split_at_row(2);
512
513 assert_eq!(first_part.rows(), 2);
514 assert_eq!(first_part.cols(), 2);
515 assert_eq!(first_part.row(0), &[1, 2]);
516 assert_eq!(first_part.row(1), &[3, 4]);
517 assert_eq!(second_part.rows(), 2);
518 assert_eq!(second_part.cols(), 2);
519 assert_eq!(second_part.row(0), &[5, 6]);
520 assert_eq!(second_part.row(1), &[7, 8]);
521 }
522}