neural_network_study/
matrix.rs

1use rand::{Rng, rngs::StdRng};
2use rayon::prelude::*;
3use serde::{Deserialize, Deserializer, Serialize, de};
4use std::ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign};
5
6/// A simple 2-dimensional matrix with basic operations
7#[derive(Clone, Debug, Default, PartialEq, Serialize)]
8pub struct Matrix {
9    rows: usize,
10    cols: usize,
11    data: Vec<f64>,
12}
13
14#[derive(Deserialize)]
15struct MatrixRepr {
16    rows: usize,
17    cols: usize,
18    data: Vec<f64>,
19}
20
21impl<'de> Deserialize<'de> for Matrix {
22    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
23    where
24        D: Deserializer<'de>,
25    {
26        let repr = MatrixRepr::deserialize(deserializer)?;
27        let expected_len = repr
28            .rows
29            .checked_mul(repr.cols)
30            .ok_or_else(|| de::Error::custom("matrix dimensions overflow usize"))?;
31
32        if repr.data.len() != expected_len {
33            return Err(de::Error::custom(format!(
34                "matrix data length mismatch: expected {}, got {}",
35                expected_len,
36                repr.data.len()
37            )));
38        }
39
40        Ok(Self {
41            rows: repr.rows,
42            cols: repr.cols,
43            data: repr.data,
44        })
45    }
46}
47
48impl Matrix {
49    fn element_count(rows: usize, cols: usize) -> usize {
50        rows.checked_mul(cols)
51            .expect("matrix dimensions overflow usize")
52    }
53
54    /// Creates a new matrix with the given number of rows and columns,
55    /// initialized to zero.
56    pub fn new(rows: usize, cols: usize) -> Self {
57        let data = vec![0.0; Self::element_count(rows, cols)];
58        Self { rows, cols, data }
59    }
60
61    /// Creates a new matrix with the given number of rows and columns,
62    /// initialized with random values between -1.0 and 1.0.
63    pub fn random(rng: &mut StdRng, rows: usize, cols: usize) -> Self {
64        Self::random_range(rng, rows, cols, -1.0, 1.0)
65    }
66
67    /// Creates a new matrix with the given number of rows and columns,
68    /// initialized with random values in the provided half-open range.
69    pub fn random_range(rng: &mut StdRng, rows: usize, cols: usize, min: f64, max: f64) -> Self {
70        let data = (0..Self::element_count(rows, cols))
71            .map(|_| rng.random_range(min..max))
72            .collect();
73        Self { rows, cols, data }
74    }
75
76    /// Creates a new matrix from a 2D vector.
77    /// The outer vector represents the rows, and the inner vectors represent the columns.
78    /// Panics if the inner vectors have different lengths.
79    pub fn from_vec(rows: usize, cols: usize, data: Vec<f64>) -> Self {
80        if data.len() != Self::element_count(rows, cols) {
81            panic!("data length does not match row and col count")
82        }
83        Self { rows, cols, data }
84    }
85
86    /// Creates a new matrix from a column vector.
87    pub fn from_col_vec(data: Vec<f64>) -> Self {
88        let rows = data.len();
89        let cols = 1;
90        Self::from_vec(rows, cols, data)
91    }
92
93    /// Transposes the matrix.
94    pub fn transpose(&self) -> Self {
95        let mut transposed_data = vec![0.0; self.rows * self.cols];
96        for i in 0..self.rows {
97            for j in 0..self.cols {
98                transposed_data[j * self.rows + i] = self.data[i * self.cols + j];
99            }
100        }
101        Self::from_vec(self.cols, self.rows, transposed_data)
102    }
103
104    /// Returns the number of rows in the matrix.
105    pub fn rows(&self) -> usize {
106        self.rows
107    }
108
109    /// Returns the number of columns in the matrix.
110    pub fn cols(&self) -> usize {
111        self.cols
112    }
113
114    /// Returns the column at the given index as a vector.
115    /// Panics if the index is out of bounds.
116    pub fn col(&self, col: usize) -> Vec<f64> {
117        if col >= self.cols {
118            panic!("Index out of bounds");
119        }
120        (0..self.rows)
121            .map(|i| self.data[i * self.cols + col])
122            .collect()
123    }
124
125    /// Returns the data in the matrix.
126    pub fn data(&self) -> &[f64] {
127        &self.data
128    }
129
130    /// Returns the mutable data in the matrix.
131    pub fn data_mut(&mut self) -> &mut [f64] {
132        &mut self.data
133    }
134
135    /// Returns the value at the given row and column.
136    /// Panics if the indices are out of bounds.
137    pub fn get(&self, row: usize, col: usize) -> f64 {
138        if row >= self.rows || col >= self.cols {
139            panic!("Index out of bounds");
140        }
141        self.data[row * self.cols + col]
142    }
143
144    /// Returns a mutable reference to the value at the given row and column.
145    /// Panics if the indices are out of bounds.
146    pub fn get_mut(&mut self, row: usize, col: usize) -> &mut f64 {
147        if row >= self.rows || col >= self.cols {
148            panic!("Index out of bounds");
149        }
150        &mut self.data[row * self.cols + col]
151    }
152
153    /// Sets the value at the given row and column.
154    /// Panics if the indices are out of bounds.
155    pub fn set(&mut self, row: usize, col: usize, value: f64) {
156        if row >= self.rows || col >= self.cols {
157            panic!("Index out of bounds");
158        }
159        self.data[row * self.cols + col] = value;
160    }
161
162    pub fn apply<F>(&mut self, f: F)
163    where
164        F: Fn(f64) -> f64,
165    {
166        for i in 0..self.rows {
167            for j in 0..self.cols {
168                let index = i * self.cols + j;
169                self.data[index] = f(self.data[index]);
170            }
171        }
172    }
173
174    pub fn hadamard_product(&mut self, other: &Matrix) {
175        if self.rows != other.rows || self.cols != other.cols {
176            panic!("Matrices must have the same dimensions for Hadamard product");
177        }
178        for i in 0..self.rows {
179            for j in 0..self.cols {
180                self.set(i, j, self.get(i, j) * other.get(i, j));
181            }
182        }
183    }
184
185    fn multiply_matrix_parallelized(&self, other: &Matrix) -> Matrix {
186        if self.cols != other.rows {
187            panic!("Matrices have incompatible dimensions for multiplication");
188        }
189
190        let other_t = Arc::new(other.transpose()); // wrap in Arc
191        let self_data = &self.data;
192        let other_data = &other_t.data;
193        let self_cols = self.cols;
194        let other_cols = other.cols;
195
196        let result_data: Vec<f64> = (0..self.rows)
197            .into_par_iter()
198            .flat_map_iter(|i| {
199                (0..other_t.rows).map(move |j| {
200                    let mut sum = 0.0;
201                    let row_start = i * self_cols;
202                    let col_start = j * self_cols;
203                    for k in 0..self_cols {
204                        sum += self_data[row_start + k] * other_data[col_start + k];
205                    }
206                    sum
207                })
208            })
209            .collect();
210
211        Matrix::from_vec(self.rows, other_cols, result_data)
212    }
213
214    fn multiply_matrix_naive(&self, other: &Matrix) -> Matrix {
215        if self.cols != other.rows {
216            panic!("Matrices have incompatible dimensions for multiplication");
217        }
218
219        let other_t = other.transpose(); // for better cache locality
220        let mut result = Matrix::new(self.rows, other.cols);
221
222        let self_data = &self.data;
223        let other_data = &other_t.data;
224        let result_data = &mut result.data;
225
226        let m = self.rows;
227        let n = self.cols;
228        let p = other.cols;
229
230        for i in 0..m {
231            for j in 0..p {
232                let mut sum = 0.0;
233                let a_row = i * n;
234                let b_row = j * n; // because other_t has shape [p x n]
235                for k in 0..n {
236                    sum += self_data[a_row + k] * other_data[b_row + k];
237                }
238                result_data[i * p + j] = sum;
239            }
240        }
241
242        result
243    }
244
245    pub fn multiply_matrix(&self, other: &Matrix) -> Matrix {
246        if self.rows * other.cols >= 128 * 128 {
247            self.multiply_matrix_parallelized(other)
248        } else {
249            self.multiply_matrix_naive(other)
250        }
251    }
252}
253
254impl Add<&Matrix> for Matrix {
255    type Output = Matrix;
256
257    /// Adds two matrices together, component-wise.
258    /// Panics if the matrices have different dimensions.
259    fn add(self, other: &Matrix) -> Matrix {
260        if self.rows != other.rows || self.cols != other.cols {
261            panic!("Matrices must have the same dimensions to be added");
262        }
263        let mut result = Matrix::new(self.rows, self.cols);
264        for i in 0..self.rows {
265            for j in 0..self.cols {
266                result.set(i, j, self.get(i, j) + other.get(i, j));
267            }
268        }
269        result
270    }
271}
272
273impl AddAssign<&Matrix> for Matrix {
274    /// Adds another matrix to this matrix, component-wise.
275    /// Panics if the matrices have different dimensions.
276    /// This is an in-place operation.
277    fn add_assign(&mut self, other: &Matrix) {
278        if self.rows != other.rows || self.cols != other.cols {
279            panic!("Matrices must have the same dimensions to be added");
280        }
281        for i in 0..self.rows {
282            for j in 0..self.cols {
283                self.set(i, j, self.get(i, j) + other.get(i, j));
284            }
285        }
286    }
287}
288
289impl Sub<&Matrix> for Matrix {
290    type Output = Matrix;
291
292    /// Subtracts another matrix from this matrix, component-wise.
293    /// Panics if the matrices have different dimensions.
294    fn sub(self, rhs: &Matrix) -> Self::Output {
295        if self.rows != rhs.rows || self.cols != rhs.cols {
296            panic!("Matrices must have the same dimensions to be subtracted");
297        }
298        let mut result = Matrix::new(self.rows, self.cols);
299        for i in 0..self.rows {
300            for j in 0..self.cols {
301                result.set(i, j, self.get(i, j) - rhs.get(i, j));
302            }
303        }
304        result
305    }
306}
307
308impl SubAssign<&Matrix> for Matrix {
309    fn sub_assign(&mut self, other: &Matrix) {
310        if self.rows != other.rows || self.cols != other.cols {
311            panic!("Matrices must have the same dimensions to be added");
312        }
313        for i in 0..self.rows {
314            for j in 0..self.cols {
315                self.set(i, j, self.get(i, j) - other.get(i, j));
316            }
317        }
318    }
319}
320
321impl Mul<f64> for Matrix {
322    type Output = Matrix;
323
324    /// Multiplies the matrix by a scalar.
325    fn mul(self, scalar: f64) -> Matrix {
326        let mut result = Matrix::new(self.rows, self.cols);
327        for i in 0..self.rows {
328            for j in 0..self.cols {
329                result.data[i * self.cols + j] = self.data[i * self.cols + j] * scalar;
330            }
331        }
332        result
333    }
334}
335
336impl MulAssign<f64> for Matrix {
337    /// Multiplies the matrix by a scalar in-place.
338    fn mul_assign(&mut self, scalar: f64) {
339        for i in 0..self.rows {
340            for j in 0..self.cols {
341                self.data[i * self.cols + j] *= scalar;
342            }
343        }
344    }
345}
346
347use std::sync::Arc;
348
349impl Mul<&Matrix> for &Matrix {
350    type Output = Matrix;
351
352    fn mul(self, other: &Matrix) -> Matrix {
353        self.multiply_matrix(other)
354    }
355}
356
357#[cfg(test)]
358mod matrix_tests {
359    use rand::SeedableRng;
360
361    use super::*;
362
363    #[test]
364    fn it_works() {
365        let m = Matrix::new(2, 3);
366        assert_eq!(m.rows(), 2);
367        assert_eq!(m.cols(), 3);
368        assert_eq!(m.data().len(), 2 * 3);
369    }
370
371    #[test]
372    fn it_creates_random_matrix() {
373        let mut rng = StdRng::from_os_rng();
374        let m = Matrix::random(&mut rng, 2, 3);
375        assert_eq!(m.rows, 2);
376        assert_eq!(m.cols, 3);
377        assert_eq!(m.data.len(), 2 * 3);
378        for i in 0..2 {
379            for j in 0..3 {
380                assert!(m.get(i, j) >= -1.0 && m.get(i, j) <= 1.0);
381            }
382        }
383    }
384
385    #[test]
386    fn it_creates_a_matrix_from_a_vector() {
387        let v = vec![1.0, 2.0, 5.0, 3.0, 4.0, 6.0];
388        let m = Matrix::from_vec(2, 3, v.clone());
389        assert_eq!(m.rows, 2);
390        assert_eq!(m.cols, 3);
391        assert_eq!(m.data, v);
392    }
393
394    #[test]
395    fn serde_rejects_matrix_data_with_invalid_length() {
396        let json = r#"{ "rows": 2, "cols": 3, "data": [1.0, 2.0, 3.0] }"#;
397        let result = serde_json::from_str::<Matrix>(json);
398
399        assert!(
400            result.is_err(),
401            "deserialization should reject data whose length does not match rows * cols"
402        );
403    }
404
405    #[test]
406    fn it_transposes_matrix() {
407        let m = Matrix::from_vec(
408            3,
409            2,
410            vec![
411                /* row 0 */ 1.0, 2.0, /* row 1 */ 5.0, 3.0, /* row 2 */ 4.0, 6.0,
412            ],
413        );
414        let transposed = m.transpose();
415        assert_eq!(transposed.rows, 2);
416        assert_eq!(transposed.cols, 3);
417        assert_eq!(transposed.get(0, 0), 1.0);
418        assert_eq!(transposed.get(0, 1), 5.0);
419        assert_eq!(transposed.get(0, 2), 4.0);
420        assert_eq!(transposed.get(1, 0), 2.0);
421        assert_eq!(transposed.get(1, 1), 3.0);
422        assert_eq!(transposed.get(1, 2), 6.0);
423    }
424
425    #[test]
426    fn it_gets_and_sets_values() {
427        let mut m = Matrix::new(2, 3);
428        m.set(0, 0, 1.0);
429        m.set(1, 2, 2.0);
430        assert_eq!(m.get(0, 0), 1.0);
431        assert_eq!(m.get(1, 2), 2.0);
432        assert_eq!(m.get(0, 1), 0.0);
433        assert_eq!(m.get(1, 0), 0.0);
434    }
435
436    #[test]
437    #[should_panic(expected = "Index out of bounds")]
438    fn it_panics_on_out_of_bounds_get() {
439        let m = Matrix::new(2, 3);
440        m.get(2, 0);
441    }
442
443    #[test]
444    #[should_panic(expected = "Index out of bounds")]
445    fn it_panics_on_out_of_bounds_set() {
446        let mut m = Matrix::new(2, 3);
447        m.set(2, 0, 1.0);
448    }
449
450    #[test]
451    #[should_panic(expected = "Index out of bounds")]
452    fn it_panics_on_out_of_bounds_get_mut() {
453        let mut m = Matrix::new(2, 3);
454        m.get_mut(2, 0);
455    }
456
457    #[test]
458    #[should_panic(expected = "Index out of bounds")]
459    fn it_panics_on_out_of_bounds_set_mut() {
460        let mut m = Matrix::new(2, 3);
461        m.get_mut(2, 0);
462    }
463
464    #[test]
465    fn it_gets_and_sets_mutable_values() {
466        let mut m = Matrix::new(2, 3);
467        *m.get_mut(0, 0) = 1.0;
468        *m.get_mut(1, 2) = 2.0;
469        assert_eq!(m.get(0, 0), 1.0);
470        assert_eq!(m.get(1, 2), 2.0);
471        assert_eq!(m.get(0, 1), 0.0);
472        assert_eq!(m.get(1, 0), 0.0);
473    }
474
475    #[test]
476    fn it_returns_mutable_data() {
477        let mut m = Matrix::new(2, 3);
478        m.data_mut()[0] = 1.0;
479        m.data_mut()[5] = 2.0;
480        assert_eq!(m.get(0, 0), 1.0);
481        assert_eq!(m.get(1, 2), 2.0);
482        assert_eq!(m.get(0, 1), 0.0);
483        assert_eq!(m.get(1, 0), 0.0);
484    }
485
486    #[test]
487    fn it_adds_matrices() {
488        let m1 = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]);
489        let m2 = Matrix::from_vec(2, 2, vec![5.0, 6.0, 7.0, 8.0]);
490        let result = m1 + &m2;
491        assert_eq!(result.get(0, 0), 6.0);
492        assert_eq!(result.get(0, 1), 8.0);
493        assert_eq!(result.get(1, 0), 10.0);
494        assert_eq!(result.get(1, 1), 12.0);
495    }
496
497    #[test]
498    fn it_adds_and_assigns() {
499        let mut m1 = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]);
500        let m2 = Matrix::from_vec(2, 2, vec![5.0, 6.0, 7.0, 8.0]);
501        m1 += &m2;
502        assert_eq!(m1.get(0, 0), 6.0);
503        assert_eq!(m1.get(0, 1), 8.0);
504        assert_eq!(m1.get(1, 0), 10.0);
505        assert_eq!(m1.get(1, 1), 12.0);
506    }
507
508    #[test]
509    fn it_multiplies_by_scalar() {
510        let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]);
511        let result = m * 2.0;
512        assert_eq!(result.get(0, 0), 2.0);
513        assert_eq!(result.get(0, 1), 4.0);
514        assert_eq!(result.get(1, 0), 6.0);
515        assert_eq!(result.get(1, 1), 8.0);
516    }
517
518    #[test]
519    fn it_multiplies_by_scalar_in_place() {
520        let mut m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]);
521        m *= 2.0;
522        assert_eq!(m.get(0, 0), 2.0);
523        assert_eq!(m.get(0, 1), 4.0);
524        assert_eq!(m.get(1, 0), 6.0);
525        assert_eq!(m.get(1, 1), 8.0);
526    }
527
528    #[test]
529    fn it_multiplies_matrices() {
530        let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
531        let n = Matrix::from_vec(3, 2, vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0]);
532        let e = Matrix::from_vec(2, 2, vec![58.0, 64.0, 139.0, 154.0]);
533        let r = &m * &n;
534        assert_eq!(r, e);
535    }
536
537    #[test]
538    fn it_maps() {
539        let mut m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]);
540        m.apply(|x| x * 2.0);
541        assert_eq!(m.get(0, 0), 2.0);
542        assert_eq!(m.get(0, 1), 4.0);
543        assert_eq!(m.get(1, 0), 6.0);
544        assert_eq!(m.get(1, 1), 8.0);
545    }
546}
neural_network_study/matrix.rs

neural_network_study/
matrix.rs
