differential_equations/linalg/matrix/
macros.rs

1//! Convenience macros for constructing matrices.
2//!
3//! - `matrix![ [a, b], [c, d] ]` or `matrix![ a, b; c, d ]` constructs a full dense matrix (row-major)
4//! - `banded_matrix!( k => [vals...], k2 => [vals...] )` constructs a banded matrix by diagonals, inferring size
5//!   from the provided diagonal lengths. Here k is the diagonal offset (i - j):
6//!   0 is the main diagonal, 1 is the first subdiagonal, -1 is the first superdiagonal, etc.
7//!
8//! Examples:
9//! let m: Matrix<f64> = matrix![ [1.0, 2.0], [3.0, 4.0] ];
10//! let b: Matrix<f64> = banded_matrix!(3, 1, 1; (0,0,1.0), (1,0,2.0), (0,1,3.0));
11
12// no local imports required; macros reference items via $crate paths
13
14/// Create a full dense matrix from rows.
15/// Usage:
16/// - matrix![ [a, b, c], [d, e, f], [g, h, i] ]
17/// - matrix![ a, b, c; d, e, f; g, h, i ]
18#[macro_export]
19macro_rules! matrix {
20    // Semicolon-separated rows form: matrix![ a, b; c, d ]
21    ( $( $( $x:expr ),+ ) ;+ $(;)? ) => {{
22        let rows_vec = vec![ $( vec![ $( $x ),+ ] ),+ ];
23        let n = rows_vec.len();
24        assert!(rows_vec.iter().all(|r| r.len() == n), "matrix! requires a square n x n list of rows");
25        let mut data = Vec::with_capacity(n*n);
26        for r in rows_vec.into_iter() { data.extend(r.into_iter()); }
27        $crate::linalg::matrix::Matrix::from_vec(n, n, data)
28    }};
29    ( $( [ $( $x:expr ),* $(,)? ] ),+ $(,)? ) => {{
30        // Collect rows into a Vec<Vec<_>> first
31        let rows_vec = vec![ $( vec![ $( $x ),* ] ),+ ];
32        let n = rows_vec.len();
33        // Ensure square
34        assert!(rows_vec.iter().all(|r| r.len() == n), "matrix! requires a square n x n list of rows");
35        let mut data = Vec::with_capacity(n*n);
36        for r in rows_vec.into_iter() { data.extend(r.into_iter()); }
37        $crate::linalg::matrix::Matrix::full(n, data)
38    }};
39}
40
41/// Create a banded matrix by specifying diagonals. Size and bands are inferred.
42/// Usage: banded_matrix!( 0 => [d0...], 1 => [d1...], -1 => [u1...], k => [..], ... )
43#[macro_export]
44macro_rules! banded_matrix {
45    ( $( $k:expr => [ $( $v:expr ),* $(,)? ] ),+ $(,)? ) => {{
46        // First pass: determine n, ml, mu from provided diagonals
47        let mut n: usize = 0usize;
48        let mut ml: usize = 0usize;
49        let mut mu: usize = 0usize;
50        $( {
51            let k: isize = $k as isize;
52            let vals = [ $( $v ),* ];
53            let len = vals.len();
54            let kk: usize = if k < 0 { (-k) as usize } else { k as usize };
55            let candidate = len + kk;
56            if candidate > n { n = candidate; }
57            if k < 0 { if kk > mu { mu = kk; } } else { if kk > ml { ml = kk; } }
58        } )+;
59        let mut m = $crate::linalg::matrix::Matrix::banded(n, ml, mu);
60        // Second pass: fill values; allow shorter diagonals (len <= n - |k|)
61        $( {
62            let k: isize = $k as isize;
63            let vals = [ $( $v ),* ];
64            let len = vals.len();
65            if k >= 0 {
66                let kk = k as usize;
67                assert!(len <= n - kk, "diagonal length {} too long for offset {} with inferred n={}", len, k, n);
68                for t in 0..len {
69                    let j = t;
70                    let i = t + kk;
71                    m[(i, j)] = vals[t];
72                }
73            } else {
74                let kk = (-k) as usize;
75                assert!(len <= n - kk, "diagonal length {} too long for offset {} with inferred n={}", len, k, n);
76                for t in 0..len {
77                    let i = t;
78                    let j = t + kk;
79                    m[(i, j)] = vals[t];
80                }
81            }
82        } )+;
83        m
84    }};
85}
86
87#[cfg(test)]
88mod tests {
89    use crate::linalg::matrix::Matrix;
90
91    #[test]
92    fn macro_full_matrix() {
93        let m: Matrix<f64> = matrix![ 1.0, 2.0; 3.0, 4.0 ];
94        assert_eq!(m.n, 2);
95        assert_eq!(m[(0, 0)], 1.0);
96        assert_eq!(m[(0, 1)], 2.0);
97        assert_eq!(m[(1, 0)], 3.0);
98        assert_eq!(m[(1, 1)], 4.0);
99    }
100
101    #[test]
102    fn macro_banded_matrix() {
103        // Inferred n=3 from main diag len=3
104        let b: Matrix<f64> = banded_matrix!( 0 => [1.0,1.0,1.0], 1 => [2.0,2.0], -1 => [3.0,3.0] );
105        // verify some values
106        assert_eq!(b[(0, 0)], 1.0);
107        assert_eq!(b[(1, 1)], 1.0);
108        assert_eq!(b[(2, 2)], 1.0);
109        assert_eq!(b[(1, 0)], 2.0);
110        assert_eq!(b[(2, 1)], 2.0);
111        assert_eq!(b[(0, 1)], 3.0);
112        assert_eq!(b[(1, 2)], 3.0);
113        // out of band read yields zero
114        assert_eq!(b[(0, 2)], 0.0);
115    }
116
117    #[test]
118    fn macro_banded_by_diagonals() {
119        // 4x4: main diag 1s, first upper 2s, first lower 3s
120        let b: Matrix<f64> =
121            banded_matrix!( 0 => [1.0,1.0,1.0,1.0], 1 => [2.0,2.0,2.0], -1 => [3.0,3.0,3.0] );
122        assert_eq!(b[(0, 0)], 1.0);
123        assert_eq!(b[(1, 1)], 1.0);
124        assert_eq!(b[(2, 2)], 1.0);
125        assert_eq!(b[(3, 3)], 1.0);
126        assert_eq!(b[(1, 0)], 2.0);
127        assert_eq!(b[(2, 1)], 2.0);
128        assert_eq!(b[(3, 2)], 2.0);
129        assert_eq!(b[(0, 1)], 3.0);
130        assert_eq!(b[(1, 2)], 3.0);
131        assert_eq!(b[(2, 3)], 3.0);
132        assert_eq!(b[(0, 3)], 0.0);
133    }
134}