surge_sparse/
complex_klu.rs

1// SPDX-License-Identifier: LicenseRef-PolyForm-Noncommercial-1.0.0
2//! Complex sparse direct solver backed by KLU via 2n x 2n real expansion.
3
4use num_complex::Complex64;
5
6use crate::KluSolver;
7use crate::csc::CscMatrix;
8use crate::error::{SparseError, SparseResult};
9
10#[derive(Debug, Clone, Copy)]
11struct SourceEntry {
12    top_left: usize,
13    top_right: usize,
14    bottom_left: usize,
15    bottom_right: usize,
16}
17
18type RealExpansion = (Vec<usize>, Vec<usize>, Vec<SourceEntry>, Vec<f64>);
19
20/// Complex sparse direct solver for a fixed square CSC sparsity pattern.
21pub struct ComplexKluSolver {
22    n: usize,
23    source_col_ptrs: Vec<usize>,
24    source_row_indices: Vec<usize>,
25    klu: KluSolver,
26    entries: Vec<SourceEntry>,
27    values_buf: Vec<f64>,
28    real_rhs_buf: Vec<f64>,
29}
30
31impl ComplexKluSolver {
32    pub fn new(matrix: &CscMatrix<Complex64>) -> SparseResult<Self> {
33        if !matrix.is_square() {
34            return Err(SparseError::MatrixNotSquare {
35                nrows: matrix.nrows(),
36                ncols: matrix.ncols(),
37            });
38        }
39        let n = matrix.nrows();
40        if n == 0 {
41            return Err(SparseError::EmptyMatrix);
42        }
43
44        let n2 = checked_mul(n, 2, "real-expanded dimension")?;
45        let (col_ptrs, row_indices, entries, values_buf) = build_real_expansion(matrix)?;
46        let mut klu = KluSolver::new(n2, &col_ptrs, &row_indices)?;
47        klu.factor(&values_buf)?;
48
49        Ok(Self {
50            n,
51            source_col_ptrs: matrix.col_ptrs().to_vec(),
52            source_row_indices: matrix.row_indices().to_vec(),
53            klu,
54            entries,
55            values_buf,
56            real_rhs_buf: vec![0.0; n2],
57        })
58    }
59
60    pub fn factor(&mut self, matrix: &CscMatrix<Complex64>) -> SparseResult<()> {
61        self.update_values(matrix)?;
62        self.klu.factor(&self.values_buf)
63    }
64
65    pub fn refactor(&mut self, matrix: &CscMatrix<Complex64>) -> SparseResult<()> {
66        self.update_values(matrix)?;
67        match self.klu.refactor(&self.values_buf) {
68            Ok(()) => Ok(()),
69            Err(SparseError::KluRefactorFailed) => self.klu.factor(&self.values_buf),
70            Err(error) => Err(error),
71        }
72    }
73
74    pub fn solve(&mut self, rhs: &[Complex64]) -> SparseResult<Vec<Complex64>> {
75        let mut solution = rhs.to_vec();
76        self.solve_in_place(&mut solution)?;
77        Ok(solution)
78    }
79
80    pub fn solve_in_place(&mut self, rhs: &mut [Complex64]) -> SparseResult<()> {
81        if rhs.len() != self.n {
82            return Err(SparseError::RhsLengthMismatch {
83                expected: self.n,
84                found: rhs.len(),
85            });
86        }
87
88        for (index, value) in rhs.iter().enumerate() {
89            self.real_rhs_buf[index] = value.re;
90            self.real_rhs_buf[index + self.n] = value.im;
91        }
92        self.klu.solve(&mut self.real_rhs_buf)?;
93
94        for (index, value) in rhs.iter_mut().enumerate() {
95            *value = Complex64::new(self.real_rhs_buf[index], self.real_rhs_buf[index + self.n]);
96        }
97        Ok(())
98    }
99
100    pub fn dim(&self) -> usize {
101        self.n
102    }
103
104    pub fn rcond(&self) -> f64 {
105        self.klu.rcond()
106    }
107
108    fn update_values(&mut self, matrix: &CscMatrix<Complex64>) -> SparseResult<()> {
109        if !matrix.has_same_pattern_slices(&self.source_col_ptrs, &self.source_row_indices) {
110            return Err(SparseError::PatternMismatch);
111        }
112
113        let values = matrix.values();
114        for (source_index, entry) in self.entries.iter().enumerate() {
115            let value = values[source_index];
116            self.values_buf[entry.top_left] = value.re;
117            self.values_buf[entry.top_right] = -value.im;
118            self.values_buf[entry.bottom_left] = value.im;
119            self.values_buf[entry.bottom_right] = value.re;
120        }
121        Ok(())
122    }
123}
124
125fn build_real_expansion(source: &CscMatrix<Complex64>) -> SparseResult<RealExpansion> {
126    let n = source.nrows();
127    let source_nnz = source.nnz();
128    let n2 = checked_mul(n, 2, "real-expanded dimension")?;
129    let total_nnz = checked_mul(source_nnz, 4, "real-expanded nnz")?;
130    let col_ptr_len = checked_add(n2, 1, "real-expanded column-pointer length")?;
131
132    let mut col_ptrs = vec![0usize; col_ptr_len];
133    let mut row_indices = Vec::with_capacity(total_nnz);
134    let mut entries = Vec::with_capacity(total_nnz);
135    let mut values = Vec::with_capacity(total_nnz);
136
137    entries.resize(
138        source_nnz,
139        SourceEntry {
140            top_left: 0,
141            top_right: 0,
142            bottom_left: 0,
143            bottom_right: 0,
144        },
145    );
146
147    let mut nnz = 0usize;
148    for (col, window) in source.col_ptrs().windows(2).enumerate() {
149        let start = window[0];
150        let end = window[1];
151        col_ptrs[col] = nnz;
152
153        for ((entry, &row), &value) in entries[start..end]
154            .iter_mut()
155            .zip(source.row_indices()[start..end].iter())
156            .zip(source.values()[start..end].iter())
157        {
158            row_indices.push(row);
159            entry.top_left = nnz;
160            values.push(value.re);
161            nnz += 1;
162        }
163        for ((entry, &row), &value) in entries[start..end]
164            .iter_mut()
165            .zip(source.row_indices()[start..end].iter())
166            .zip(source.values()[start..end].iter())
167        {
168            row_indices.push(row + n);
169            entry.bottom_left = nnz;
170            values.push(value.im);
171            nnz += 1;
172        }
173        col_ptrs[col + 1] = nnz;
174    }
175
176    for (col, window) in source.col_ptrs().windows(2).enumerate() {
177        let start = window[0];
178        let end = window[1];
179        let real_col = col + n;
180        col_ptrs[real_col] = nnz;
181
182        for ((entry, &row), &value) in entries[start..end]
183            .iter_mut()
184            .zip(source.row_indices()[start..end].iter())
185            .zip(source.values()[start..end].iter())
186        {
187            row_indices.push(row);
188            entry.top_right = nnz;
189            values.push(-value.im);
190            nnz += 1;
191        }
192        for ((entry, &row), &value) in entries[start..end]
193            .iter_mut()
194            .zip(source.row_indices()[start..end].iter())
195            .zip(source.values()[start..end].iter())
196        {
197            row_indices.push(row + n);
198            entry.bottom_right = nnz;
199            values.push(value.re);
200            nnz += 1;
201        }
202        col_ptrs[real_col + 1] = nnz;
203    }
204
205    Ok((col_ptrs, row_indices, entries, values))
206}
207
208fn checked_mul(lhs: usize, rhs: usize, what: &'static str) -> SparseResult<usize> {
209    lhs.checked_mul(rhs)
210        .ok_or(SparseError::SizeOverflow { what })
211}
212
213fn checked_add(lhs: usize, rhs: usize, what: &'static str) -> SparseResult<usize> {
214    lhs.checked_add(rhs)
215        .ok_or(SparseError::SizeOverflow { what })
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221    use crate::CscMatrix;
222
223    fn complex_csc(
224        n: usize,
225        col_ptrs: Vec<usize>,
226        row_indices: Vec<usize>,
227        values: Vec<Complex64>,
228    ) -> CscMatrix<Complex64> {
229        CscMatrix::try_new(n, n, col_ptrs, row_indices, values).unwrap()
230    }
231
232    #[test]
233    fn rejects_pattern_changes() {
234        let mat = complex_csc(1, vec![0, 1], vec![0], vec![Complex64::new(1.0, 2.0)]);
235        let mut solver = ComplexKluSolver::new(&mat).expect("valid factorization");
236
237        let changed = complex_csc(1, vec![0, 1], vec![0], vec![Complex64::new(2.0, 3.0)]);
238        solver
239            .refactor(&changed)
240            .expect("value-only change should be allowed");
241
242        let changed_pattern = complex_csc(
243            2,
244            vec![0, 2, 4],
245            vec![0, 1, 0, 1],
246            vec![
247                Complex64::new(1.0, 0.0),
248                Complex64::new(-1.0, 0.0),
249                Complex64::new(-1.0, 0.0),
250                Complex64::new(1.0, 0.0),
251            ],
252        );
253
254        let error = solver
255            .refactor(&changed_pattern)
256            .expect_err("pattern change must be rejected");
257        assert_eq!(error, SparseError::PatternMismatch);
258    }
259
260    #[test]
261    fn solve_in_place_reuses_caller_buffer() {
262        let mat = complex_csc(1, vec![0, 1], vec![0], vec![Complex64::new(2.0, 0.0)]);
263        let mut solver = ComplexKluSolver::new(&mat).expect("valid factorization");
264        let mut rhs = vec![Complex64::new(4.0, 0.0)];
265        solver
266            .solve_in_place(&mut rhs)
267            .expect("solve_in_place should succeed");
268        assert_eq!(rhs, vec![Complex64::new(2.0, 0.0)]);
269    }
270
271    #[test]
272    fn complex_klu_zero_rhs_returns_zero() {
273        let mat = complex_csc(
274            2,
275            vec![0, 2, 4],
276            vec![0, 1, 0, 1],
277            vec![
278                Complex64::new(4.0, -1.0),
279                Complex64::new(-1.0, 0.0),
280                Complex64::new(-1.0, 0.0),
281                Complex64::new(4.0, -1.0),
282            ],
283        );
284        let mut solver = ComplexKluSolver::new(&mat).unwrap();
285
286        let result = solver
287            .solve(&[Complex64::new(0.0, 0.0), Complex64::new(0.0, 0.0)])
288            .unwrap();
289        for val in &result {
290            assert!(val.re.abs() < 1e-14, "re should be ~0, got {}", val.re);
291            assert!(val.im.abs() < 1e-14, "im should be ~0, got {}", val.im);
292        }
293    }
294
295    #[test]
296    fn complex_klu_pure_real_rhs() {
297        let mat = complex_csc(
298            2,
299            vec![0, 2, 4],
300            vec![0, 1, 0, 1],
301            vec![
302                Complex64::new(2.0, 0.0),
303                Complex64::new(1.0, 0.0),
304                Complex64::new(1.0, 0.0),
305                Complex64::new(3.0, 0.0),
306            ],
307        );
308        let mut solver = ComplexKluSolver::new(&mat).unwrap();
309
310        let result = solver
311            .solve(&[Complex64::new(3.0, 0.0), Complex64::new(4.0, 0.0)])
312            .unwrap();
313        assert!((result[0].re - 1.0).abs() < 1e-12);
314        assert!((result[1].re - 1.0).abs() < 1e-12);
315        assert!(result[0].im.abs() < 1e-14);
316        assert!(result[1].im.abs() < 1e-14);
317    }
318
319    #[test]
320    fn complex_klu_dim_and_rcond() {
321        let mat = complex_csc(
322            3,
323            vec![0, 1, 2, 3],
324            vec![0, 1, 2],
325            vec![
326                Complex64::new(1.0, 0.0),
327                Complex64::new(2.0, 0.0),
328                Complex64::new(3.0, 0.0),
329            ],
330        );
331        let solver = ComplexKluSolver::new(&mat).unwrap();
332
333        assert_eq!(solver.dim(), 3);
334        assert!(solver.rcond() > 0.0);
335        assert!(solver.rcond() <= 1.0);
336    }
337
338    #[test]
339    fn complex_klu_full_complex_solve() {
340        let mat = complex_csc(
341            2,
342            vec![0, 1, 2],
343            vec![0, 1],
344            vec![Complex64::new(1.0, 2.0), Complex64::new(3.0, 4.0)],
345        );
346        let mut solver = ComplexKluSolver::new(&mat).unwrap();
347
348        let result = solver
349            .solve(&[Complex64::new(1.0, 2.0), Complex64::new(3.0, 4.0)])
350            .unwrap();
351        assert!((result[0].re - 1.0).abs() < 1e-12);
352        assert!(result[0].im.abs() < 1e-12);
353        assert!((result[1].re - 1.0).abs() < 1e-12);
354        assert!(result[1].im.abs() < 1e-12);
355    }
356
357    #[test]
358    fn complex_klu_pure_imaginary_rhs() {
359        let mat = complex_csc(
360            2,
361            vec![0, 1, 2],
362            vec![0, 1],
363            vec![Complex64::new(2.0, 0.0), Complex64::new(4.0, 0.0)],
364        );
365        let mut solver = ComplexKluSolver::new(&mat).unwrap();
366
367        let result = solver
368            .solve(&[Complex64::new(0.0, 2.0), Complex64::new(0.0, 8.0)])
369            .unwrap();
370        assert!(result[0].re.abs() < 1e-14);
371        assert!((result[0].im - 1.0).abs() < 1e-12);
372        assert!(result[1].re.abs() < 1e-14);
373        assert!((result[1].im - 2.0).abs() < 1e-12);
374    }
375
376    #[test]
377    fn complex_klu_wrong_rhs_length() {
378        let mat = complex_csc(
379            2,
380            vec![0, 1, 2],
381            vec![0, 1],
382            vec![Complex64::new(1.0, 0.0), Complex64::new(1.0, 0.0)],
383        );
384        let mut solver = ComplexKluSolver::new(&mat).unwrap();
385
386        let err = solver.solve(&[Complex64::new(1.0, 0.0)]).unwrap_err();
387        assert!(matches!(err, SparseError::RhsLengthMismatch { .. }));
388    }
389
390    #[test]
391    fn complex_klu_factor_refactor() {
392        let mat1 = complex_csc(1, vec![0, 1], vec![0], vec![Complex64::new(2.0, 0.0)]);
393        let mut solver = ComplexKluSolver::new(&mat1).unwrap();
394
395        let result = solver.solve(&[Complex64::new(6.0, 0.0)]).unwrap();
396        assert!((result[0].re - 3.0).abs() < 1e-12);
397
398        let mat2 = complex_csc(1, vec![0, 1], vec![0], vec![Complex64::new(3.0, 0.0)]);
399        solver.factor(&mat2).unwrap();
400
401        let result = solver.solve(&[Complex64::new(6.0, 0.0)]).unwrap();
402        assert!((result[0].re - 2.0).abs() < 1e-12);
403    }
404}
surge_sparse/complex_klu.rs

surge_sparse/
complex_klu.rs
