slain_num/
mat.rs

1use crate::simd;
2
3#[derive(Clone, Debug, PartialEq)]
4pub struct Mat {
5    rows: usize,
6    cols: usize,
7    data: Vec<f64>,
8}
9
10impl Mat {
11    #[must_use]
12    pub fn zeros(rows: usize, cols: usize) -> Self {
13        Self { rows, cols, data: vec![0.0; rows * cols] }
14    }
15
16    #[must_use]
17    pub fn identity(n: usize) -> Self {
18        let mut m = Self::zeros(n, n);
19        for i in 0..n {
20            m[(i, i)] = 1.0;
21        }
22        m
23    }
24
25    #[must_use]
26    pub fn from_vec(rows: usize, cols: usize, data: Vec<f64>) -> Self {
27        assert_eq!(data.len(), rows * cols, "matrix data length");
28        Self { rows, cols, data }
29    }
30
31    #[must_use]
32    pub fn rows(&self) -> usize {
33        self.rows
34    }
35
36    #[must_use]
37    pub fn cols(&self) -> usize {
38        self.cols
39    }
40
41    #[must_use]
42    pub fn row(&self, r: usize) -> &[f64] {
43        &self.data[r * self.cols..(r + 1) * self.cols]
44    }
45
46    pub fn row_mut(&mut self, r: usize) -> &mut [f64] {
47        &mut self.data[r * self.cols..(r + 1) * self.cols]
48    }
49
50    #[must_use]
51    pub fn as_slice(&self) -> &[f64] {
52        &self.data
53    }
54
55    #[must_use]
56    pub fn transpose(&self) -> Self {
57        let mut t = Self::zeros(self.cols, self.rows);
58        for r in 0..self.rows {
59            for c in 0..self.cols {
60                t.data[c * self.rows + r] = self.data[r * self.cols + c];
61            }
62        }
63        t
64    }
65
66    #[must_use]
67    pub fn matmul(&self, b: &Self) -> Self {
68        assert_eq!(self.cols, b.rows, "inner dimensions");
69        let mut out = Self::zeros(self.rows, b.cols);
70        for r in 0..self.rows {
71            let orow = &mut out.data[r * b.cols..(r + 1) * b.cols];
72            for (k, &a) in self.row(r).iter().enumerate() {
73                if a != 0.0 {
74                    simd::axpy(a, b.row(k), orow);
75                }
76            }
77        }
78        out
79    }
80
81    #[must_use]
82    pub fn mul_vec(&self, v: &[f64]) -> Vec<f64> {
83        (0..self.rows).map(|r| simd::dot(self.row(r), v)).collect()
84    }
85
86    #[must_use]
87    pub fn gram(&self) -> Self {
88        let t = self.transpose();
89        let n = self.cols;
90        let mut g = Self::zeros(n, n);
91        for i in 0..n {
92            for j in i..n {
93                let v = simd::dot(t.row(i), t.row(j));
94                g.data[i * n + j] = v;
95                g.data[j * n + i] = v;
96            }
97        }
98        g
99    }
100
101    #[must_use]
102    pub fn t_mul_vec(&self, v: &[f64]) -> Vec<f64> {
103        let mut out = vec![0.0; self.cols];
104        for (r, &k) in v.iter().enumerate().take(self.rows) {
105            simd::axpy(k, self.row(r), &mut out);
106        }
107        out
108    }
109}
110
111impl std::ops::Index<(usize, usize)> for Mat {
112    type Output = f64;
113    fn index(&self, (r, c): (usize, usize)) -> &f64 {
114        &self.data[r * self.cols + c]
115    }
116}
117
118impl std::ops::IndexMut<(usize, usize)> for Mat {
119    fn index_mut(&mut self, (r, c): (usize, usize)) -> &mut f64 {
120        &mut self.data[r * self.cols + c]
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    fn filled(rows: usize, cols: usize, seed: f64) -> Mat {
129        Mat::from_vec(rows, cols, (0..rows * cols).map(|i| ((i as f64 + seed) * 0.713).sin()).collect())
130    }
131
132    #[test]
133    fn products_match_naive() {
134        let (a, b) = (filled(7, 13, 1.0), filled(13, 9, 2.0));
135        let c = a.matmul(&b);
136        for r in 0..7 {
137            for k in 0..9 {
138                let want: f64 = (0..13).map(|i| a[(r, i)] * b[(i, k)]).sum();
139                assert!((c[(r, k)] - want).abs() < 1e-12);
140            }
141        }
142        let g = a.gram();
143        let want = a.transpose().matmul(&a);
144        assert!(g.as_slice().iter().zip(want.as_slice()).all(|(x, y)| (x - y).abs() < 1e-12));
145        let v: Vec<f64> = (0..7).map(|i| i as f64).collect();
146        let tv = a.t_mul_vec(&v);
147        let want = a.transpose().mul_vec(&v);
148        assert!(tv.iter().zip(&want).all(|(x, y)| (x - y).abs() < 1e-12));
149    }
150}
slain_num/mat.rs

slain_num/
mat.rs
