treams_core/smatrix/solve/
saved.rs

1//! The original internal solve strategy and its numerical state, with a fixed
2//! factor slot that never builds a reflection product solely for serialization.
3//!
4//! treams-rs extension.
5
6use super::{InternalFactor, InternalSolve};
7use crate::{
8    Result,
9    linalg::Lu,
10    saved::{Reader, SavedState, Writer},
11    smatrix::saved::{invalid, product, read_matrix, total, write_matrix},
12};
13
14impl InternalSolve {
15    pub(in crate::smatrix) fn state_size(n: usize, columns: usize) -> Result<usize> {
16        if n == 0 || columns == 0 {
17            return Err(invalid());
18        }
19        total(&[17, product(&[16, n, columns])?, Lu::state_size(n)?])
20    }
21}
22
23fn write_zeros(writer: &mut Writer, count: usize) {
24    const ZEROES: [u8; 256] = [0; 256];
25    for _ in 0..count / ZEROES.len() {
26        writer.raw(&ZEROES);
27    }
28    writer.raw(&ZEROES[..count % ZEROES.len()]);
29}
30
31impl SavedState for InternalSolve {
32    fn save_state(&self) -> Result<Vec<u8>> {
33        let (n, columns) = self.value.shape();
34        let mut writer = Writer::new(Self::state_size(n, columns)?);
35        writer.usize(n);
36        writer.usize(columns);
37        writer.byte(match &self.factor {
38            InternalFactor::Reflections => 0,
39            InternalFactor::Krylov(_) => 1,
40            InternalFactor::Lu(_) => 2,
41        });
42        write_matrix(&mut writer, &self.value);
43        match &self.factor {
44            InternalFactor::Reflections => write_zeros(&mut writer, Lu::state_size(n)?),
45            InternalFactor::Krylov(operator) => {
46                write_matrix(&mut writer, operator);
47                write_zeros(&mut writer, Lu::state_size(n)? - product(&[16, n, n])?);
48            }
49            InternalFactor::Lu(lu) => lu.write_state(&mut writer),
50        }
51        Ok(writer.finish())
52    }
53
54    fn from_state(bytes: &[u8]) -> Result<Self> {
55        let mut reader = Reader::new(bytes);
56        let (n, columns, tag) = (reader.usize()?, reader.usize()?, reader.byte()?);
57        if bytes.len() != Self::state_size(n, columns)? || tag > 2 {
58            return Err(invalid());
59        }
60        let value = read_matrix(&mut reader, n, columns)?;
61        let factor = match tag {
62            0 => InternalFactor::Reflections,
63            1 => InternalFactor::Krylov(read_matrix(&mut reader, n, n)?),
64            _ => InternalFactor::Lu(Lu::read_state(&mut reader, n)?),
65        };
66        // Unused slots are canonical zeros; no matrix is reconstructed for them.
67        if reader
68            .raw(reader.remaining_len())?
69            .iter()
70            .any(|&byte| byte != 0)
71        {
72            return Err(invalid());
73        }
74        reader.finish()?;
75        Ok(Self { value, factor })
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use crate::{Complex, linalg::view, test_support::patterned};
83    use nalgebra::DMatrix;
84
85    #[test]
86    fn saved_state_preserves_lu_and_both_krylov_strategies() {
87        let n = 24;
88        let weak = DMatrix::identity(n, n) * Complex::new(0.1, 0.03);
89        let cancellation = DMatrix::from_fn(n, n, |_, j| {
90            Complex::new(if j % 2 == 0 { 0.1 } else { -0.1 }, 0.0)
91        });
92        let constant = DMatrix::from_element(n, n, Complex::new(0.1, 0.0));
93        let rhs = patterned(n, 2, 0.4);
94        let direction = patterned(n, 2, 0.9);
95        for (lower, upper, threshold, tag) in [
96            (&weak, &weak, n + 1, 2),
97            (&weak, &weak, n, 0),
98            (&cancellation, &constant, n, 1),
99        ] {
100            let original =
101                InternalSolve::with_threshold(view(lower), view(upper), rhs.clone(), threshold)
102                    .unwrap();
103            let bytes = original.save_state().unwrap();
104            assert_eq!(bytes.len(), InternalSolve::state_size(n, 2).unwrap());
105            assert_eq!(bytes[16], tag);
106            let loaded = InternalSolve::from_state(&bytes).unwrap();
107            assert_eq!(loaded.save_state().unwrap(), bytes);
108            assert_eq!(loaded.value, original.value);
109            assert_eq!(
110                loaded
111                    .solve_forward(view(lower), view(upper), direction.clone())
112                    .unwrap(),
113                original
114                    .solve_forward(view(lower), view(upper), direction.clone())
115                    .unwrap(),
116            );
117            assert_eq!(
118                loaded
119                    .solve_adjoint(view(lower), view(upper), direction.clone())
120                    .unwrap()
121                    .adjoint,
122                original
123                    .solve_adjoint(view(lower), view(upper), direction.clone())
124                    .unwrap()
125                    .adjoint,
126            );
127            assert!(InternalSolve::from_state(&bytes[..bytes.len() - 1]).is_err());
128            let mut extended = bytes.clone();
129            extended.push(0);
130            assert!(InternalSolve::from_state(&extended).is_err());
131            let mut oversized = bytes;
132            oversized[..8].fill(u8::MAX);
133            assert!(InternalSolve::from_state(&oversized).is_err());
134        }
135    }
136}
treams_core/smatrix/solve/saved.rs

treams_core/smatrix/solve/
saved.rs
