1mod saved;
13
14use faer::MatRef;
15use nalgebra::DMatrix;
16
17use crate::{
18 Complex, Error, Result,
19 linalg::{GmresOptions, Lu, gmres, norm, product_views, view, view_mut},
20 numerics::finite,
21};
22
23fn internal_operator(lower: MatRef<'_, Complex>, upper: MatRef<'_, Complex>) -> DMatrix<Complex> {
25 let mut value = -product_views(lower, upper);
26 for i in 0..value.nrows() {
27 value[(i, i)] += 1.0;
28 }
29 value
30}
31
32#[derive(Debug)]
38enum InternalFactor {
39 Reflections,
41 Krylov(DMatrix<Complex>),
43 Lu(Lu),
45}
46
47#[derive(Debug)]
49pub(super) struct InternalSolve {
50 pub(super) value: DMatrix<Complex>,
52 factor: InternalFactor,
53}
54
55fn is_contraction(operator: &DMatrix<Complex>) -> bool {
60 let n = operator.nrows();
61 let mut rows = vec![0.0_f64; n];
62 for (j, column) in operator.as_slice().chunks_exact(n).enumerate() {
63 for (i, (&z, row)) in column.iter().zip(&mut rows).enumerate() {
64 let delta = if i == j {
65 Complex::new(1.0, 0.0) - z
66 } else {
67 z
68 };
69 *row = (*row + (delta.re.abs() + delta.im.abs()).next_up()).next_up();
72 }
73 }
74 rows.iter().all(|&sum| sum < 1.0)
75}
76
77fn absolute_row_sums(matrix: MatRef<'_, Complex>, weights: &[f64]) -> Vec<f64> {
79 let accumulate = |sum: &mut f64, z: Complex, weight: f64| {
80 if z != Complex::default() && weight != 0.0 {
83 let term = ((z.re.abs() + z.im.abs()).next_up() * weight).next_up();
84 *sum = (*sum + term).next_up();
85 }
86 };
87 let mut sums = vec![0.0_f64; matrix.nrows()];
88 if matrix.col_stride() == 1 {
89 for (row, sum) in matrix.row_iter().zip(&mut sums) {
90 for (&z, &weight) in row.iter().zip(weights) {
91 accumulate(sum, z, weight);
92 }
93 }
94 } else {
95 for (column, &weight) in matrix.col_iter().zip(weights) {
96 for (&z, sum) in column.iter().zip(&mut sums) {
97 accumulate(sum, z, weight);
98 }
99 }
100 }
101 sums
102}
103
104fn reflections_contract(lower: MatRef<'_, Complex>, upper: MatRef<'_, Complex>) -> bool {
109 let upper_rows = absolute_row_sums(upper, &vec![1.0; upper.ncols()]);
110 absolute_row_sums(lower, &upper_rows)
111 .iter()
112 .all(|&sum| sum < 1.0)
113}
114
115fn reflection_action(
117 lower: MatRef<'_, Complex>,
118 upper: MatRef<'_, Complex>,
119 x: MatRef<'_, Complex>,
120 adjoint: bool,
121) -> DMatrix<Complex> {
122 let mut result = if adjoint {
123 let inner = product_views(lower.adjoint(), x);
124 product_views(upper.adjoint(), view(&inner))
125 } else {
126 let inner = product_views(upper, x);
127 product_views(lower, view(&inner))
128 };
129 for j in 0..result.ncols() {
130 for i in 0..result.nrows() {
131 result[(i, j)] = x[(i, j)] - result[(i, j)];
132 }
133 }
134 result
135}
136
137fn internal_iteration(
140 rhs: &DMatrix<Complex>,
141 apply: impl Fn(MatRef<'_, Complex>) -> DMatrix<Complex>,
142) -> Result<Option<DMatrix<Complex>>> {
143 let column_norm = |j| norm(rhs.column(j).as_slice());
147 let minimum = (0..rhs.ncols())
148 .map(column_norm)
149 .filter(|&v| v > 0.0)
150 .fold(f64::INFINITY, f64::min);
151 if rhs.iter().all(|&z| z == Complex::default()) {
152 return Ok(Some(DMatrix::zeros(rhs.nrows(), rhs.ncols())));
153 }
154 if !minimum.is_finite() {
155 return Ok(None);
156 }
157 let options = GmresOptions {
158 rtol: 0.0,
159 atol: 8.0 * f64::EPSILON * minimum,
160 restart: KRYLOV_STEPS,
161 max_iterations: KRYLOV_STEPS,
162 };
163 match gmres(rhs.as_slice(), options, |x| {
164 let x_view = MatRef::from_column_major_slice(x, rhs.nrows(), rhs.ncols());
165 let result = apply(x_view);
166 Ok(Vec::from(result.data))
167 }) {
168 Ok((answer, _)) => Ok(Some(DMatrix::from_vec(rhs.nrows(), rhs.ncols(), answer))),
169 Err(error) if error.is_numerical() => Ok(None),
171 Err(error) => Err(error),
172 }
173}
174
175const ITERATIVE_MIN_ROWS: usize = 512;
178
179const ITERATIVE_MAX_COLUMNS: usize = 8;
181
182const KRYLOV_STEPS: usize = 16;
184
185impl InternalSolve {
186 pub(super) fn new(
187 lower: MatRef<'_, Complex>,
188 upper: MatRef<'_, Complex>,
189 rhs: DMatrix<Complex>,
190 ) -> Result<Self> {
191 Self::with_threshold(lower, upper, rhs, ITERATIVE_MIN_ROWS)
192 }
193
194 pub(super) fn direct(
197 lower: MatRef<'_, Complex>,
198 upper: MatRef<'_, Complex>,
199 rhs: DMatrix<Complex>,
200 ) -> Result<Self> {
201 if rhs.iter().any(|&z| !finite(z)) {
202 return Err(Error::Singular);
203 }
204 Self::lu(internal_operator(lower, upper), rhs)
205 }
206
207 fn with_threshold(
210 lower: MatRef<'_, Complex>,
211 upper: MatRef<'_, Complex>,
212 rhs: DMatrix<Complex>,
213 min_rows: usize,
214 ) -> Result<Self> {
215 if rhs.iter().any(|&z| !finite(z)) {
216 return Err(Error::Singular);
217 }
218 let thin = rhs.nrows() >= min_rows && rhs.ncols() <= ITERATIVE_MAX_COLUMNS;
219 let factored = thin && reflections_contract(lower, upper);
220 if factored
221 && let Some(value) =
222 internal_iteration(&rhs, |x| reflection_action(lower, upper, x, false))?
223 {
224 return Ok(Self {
225 value,
226 factor: InternalFactor::Reflections,
227 });
228 }
229 let operator = internal_operator(lower, upper);
230 if thin
231 && !factored
232 && is_contraction(&operator)
233 && let Some(value) = internal_iteration(&rhs, |x| product_views(view(&operator), x))?
234 {
235 return Ok(Self {
236 value,
237 factor: InternalFactor::Krylov(operator),
238 });
239 }
240 Self::lu(operator, rhs)
241 }
242
243 fn lu(operator: DMatrix<Complex>, mut rhs: DMatrix<Complex>) -> Result<Self> {
245 let lu = Lu::new(operator)?;
246 lu.solve_in_place(view_mut(&mut rhs))?;
247 if rhs.iter().any(|&z| !finite(z)) {
248 return Err(Error::Singular);
249 }
250 Ok(Self {
251 value: rhs,
252 factor: InternalFactor::Lu(lu),
253 })
254 }
255
256 pub(super) fn solve_adjoint(
259 &self,
260 lower: MatRef<'_, Complex>,
261 upper: MatRef<'_, Complex>,
262 rhs: DMatrix<Complex>,
263 ) -> Result<AdjointSolve<'_>> {
264 let adjoint = self.solve_rhs(lower, upper, rhs, true)?;
265 Ok(AdjointSolve {
266 adjoint,
267 value: &self.value,
268 })
269 }
270
271 pub(super) fn solve_forward(
273 &self,
274 lower: MatRef<'_, Complex>,
275 upper: MatRef<'_, Complex>,
276 rhs: DMatrix<Complex>,
277 ) -> Result<DMatrix<Complex>> {
278 self.solve_rhs(lower, upper, rhs, false)
279 }
280
281 fn solve_rhs(
283 &self,
284 lower: MatRef<'_, Complex>,
285 upper: MatRef<'_, Complex>,
286 mut rhs: DMatrix<Complex>,
287 adjoint: bool,
288 ) -> Result<DMatrix<Complex>> {
289 let fallback;
290 let lu = match &self.factor {
291 InternalFactor::Lu(lu) => lu,
292 InternalFactor::Reflections => {
293 if let Some(value) =
294 internal_iteration(&rhs, |x| reflection_action(lower, upper, x, adjoint))?
295 {
296 return Ok(value);
297 }
298 fallback = Lu::new(internal_operator(lower, upper))?;
299 &fallback
300 }
301 InternalFactor::Krylov(operator) => {
302 if let Some(value) = internal_iteration(&rhs, |x| {
303 let operator = view(operator);
304 if adjoint {
305 product_views(operator.adjoint(), x)
306 } else {
307 product_views(operator, x)
308 }
309 })? {
310 return Ok(value);
311 }
312 fallback = Lu::new(operator.clone())?;
313 &fallback
314 }
315 };
316 if adjoint {
317 lu.solve_adjoint_in_place(view_mut(&mut rhs))?;
318 } else {
319 lu.solve_in_place(view_mut(&mut rhs))?;
320 }
321 if rhs.iter().any(|&z| !finite(z)) {
322 return Err(Error::Singular);
323 }
324 Ok(rhs)
325 }
326}
327
328#[derive(Debug)]
330pub(super) struct AdjointSolve<'a> {
331 pub(super) adjoint: DMatrix<Complex>,
333 pub(super) value: &'a DMatrix<Complex>,
335}
336
337#[cfg(test)]
338mod tests {
339 use super::*;
340 use crate::linalg;
341 use crate::test_support::{patterned, prop_assert_close, re_dot};
342 use proptest::{prelude::*, test_runner::TestCaseError};
343 use std::result::Result;
344
345 fn pairing(a: &DMatrix<Complex>, b: &DMatrix<Complex>) -> Complex {
347 let rotated = a * Complex::i();
348 Complex::new(re_dot(a.iter(), b.iter()), re_dot(rotated.iter(), b.iter()))
349 }
350
351 fn check_krylov_matches_dense_solve(
355 n: usize,
356 p: usize,
357 seed: f64,
358 ) -> Result<(), TestCaseError> {
359 let weight = Complex::new(0.1 / f64::from(u32::try_from(n).unwrap()).sqrt(), 0.0);
360 let lower = patterned(n, n, seed) * weight;
361 let upper = patterned(n, n, seed + 0.4) * weight;
362 let rhs = patterned(n, p, seed + 0.2);
363 let g = patterned(n, p, seed + 0.7);
364 let operator = internal_operator(view(&lower), view(&upper));
365 let dense = linalg::solve_owned(operator.clone(), rhs.clone()).unwrap();
366 prop_assert!(is_contraction(&operator));
367 let value = internal_iteration(&rhs, |x| product_views(view(&operator), x)).unwrap();
368 let adjoint =
369 internal_iteration(&g, |x| product_views(view(&operator).adjoint(), x)).unwrap();
370 let (value, adjoint) = (value.unwrap(), adjoint.unwrap());
371 prop_assert_close!(&value, dense.value(), 2e-14 * dense.value().norm());
372 let dense_adjoint = dense.adjoint_rhs(g.clone()).unwrap();
373 prop_assert_close!(&adjoint, &dense_adjoint, 2e-14 * g.norm());
374 let factored = internal_iteration(&rhs, |x| {
375 reflection_action(view(&lower), view(&upper), x, false)
376 })
377 .unwrap()
378 .unwrap();
379 let factored_adjoint = internal_iteration(&g, |x| {
380 reflection_action(view(&lower), view(&upper), x, true)
381 })
382 .unwrap()
383 .unwrap();
384 prop_assert_close!(&factored, &value, 2e-14 * value.norm());
385 prop_assert_close!(&factored_adjoint, &adjoint, 2e-14 * g.norm());
386 prop_assert_close!(
387 pairing(&g, &value),
388 pairing(&adjoint, &rhs),
389 2e-13 * rhs.norm() * g.norm()
390 );
391 Ok(())
392 }
393
394 proptest! {
395 #![proptest_config(ProptestConfig::with_cases(16))]
396 #[test]
397 fn reflection_actions_match_dense_solve_and_adjoint(
398 n in 3_usize..12, p in 1_usize..4, seed in -2.0_f64..2.0,
399 ) {
400 check_krylov_matches_dense_solve(n, p, seed)?;
401 }
402 }
403
404 #[test]
405 fn production_threshold_solves_thin_contractions_iteratively() -> Result<(), TestCaseError> {
406 let coefficient = Complex::new(1.0, 0.0) - Complex::new(0.1, 0.03).powi(2);
409 for (n, iterates) in [(ITERATIVE_MIN_ROWS, true), (ITERATIVE_MIN_ROWS - 1, false)] {
410 let weak = DMatrix::identity(n, n) * Complex::new(0.1, 0.03);
411 let rhs = patterned(n, 2, 0.4);
412 let solve = InternalSolve::new(view(&weak), view(&weak), rhs.clone()).unwrap();
413 prop_assert_eq!(
414 matches!(solve.factor, InternalFactor::Reflections),
415 iterates
416 );
417 prop_assert_close!(&solve.value * coefficient, &rhs, 2e-13 * rhs.norm());
418 let tangent = solve
419 .solve_forward(view(&weak), view(&weak), rhs.clone())
420 .unwrap();
421 prop_assert_close!(tangent * coefficient, &rhs, 2e-13 * rhs.norm());
422 let gradient = solve
423 .solve_adjoint(view(&weak), view(&weak), rhs.clone())
424 .unwrap()
425 .adjoint;
426 prop_assert_close!(gradient * coefficient.conj(), &rhs, 2e-13 * rhs.norm());
427 }
428 Ok(())
429 }
430
431 #[test]
432 fn cancellation_contractions_keep_the_dense_product_path() -> Result<(), TestCaseError> {
433 let n = 24;
434 let lower = DMatrix::from_fn(n, n, |_, j| {
435 Complex::new(if j % 2 == 0 { 0.1 } else { -0.1 }, 0.0)
436 });
437 let upper = DMatrix::from_element(n, n, Complex::new(0.1, 0.0));
438 let rhs = patterned(n, 2, 0.4);
439 prop_assert!(!reflections_contract(view(&lower), view(&upper)));
440 let operator = internal_operator(view(&lower), view(&upper));
441 prop_assert!(is_contraction(&operator));
442 let solve =
443 InternalSolve::with_threshold(view(&lower), view(&upper), rhs.clone(), n).unwrap();
444 prop_assert!(matches!(solve.factor, InternalFactor::Krylov(_)));
445 prop_assert_close!(&operator * &solve.value, &rhs, 2e-13 * rhs.norm());
446 let tangent = solve
447 .solve_forward(view(&lower), view(&upper), rhs.clone())
448 .unwrap();
449 prop_assert_close!(&operator * tangent, &rhs, 2e-13 * rhs.norm());
450 let adjoint = solve
451 .solve_adjoint(view(&lower), view(&upper), rhs.clone())
452 .unwrap()
453 .adjoint;
454 prop_assert_close!(operator.adjoint() * adjoint, &rhs, 2e-13 * rhs.norm());
455 Ok(())
456 }
457
458 #[test]
459 fn reflection_certificate_handles_layouts_and_rejects_unit_gain() -> Result<(), TestCaseError> {
460 let n = 12;
461 let lower = patterned(n, n, 0.4) * Complex::new(0.02, 0.0);
462 let upper = patterned(n, n, 0.9) * Complex::new(0.02, 0.0);
463 let lower_rows = lower.transpose();
464 let upper_rows = upper.transpose();
465 for lower in [view(&lower), view(&lower_rows).transpose()] {
466 for upper in [view(&upper), view(&upper_rows).transpose()] {
467 prop_assert!(reflections_contract(lower, upper));
468 prop_assert!(is_contraction(&internal_operator(lower, upper)));
469 }
470 }
471 let identity = DMatrix::identity(n, n);
472 prop_assert!(!reflections_contract(view(&identity), view(&identity)));
473 let huge = &identity * Complex::new(f64::MAX, 0.0);
475 prop_assert!(!reflections_contract(view(&huge), view(&huge)));
476 Ok(())
477 }
478
479 #[test]
480 fn thin_large_solve_skips_lu_and_keeps_dense_fallback() -> Result<(), TestCaseError> {
481 let (n, min_rows) = (24, 16);
486 prop_assert!(n > KRYLOV_STEPS);
487 let solve = |lower: &DMatrix<Complex>, upper: &DMatrix<Complex>, rhs, min_rows| {
488 InternalSolve::with_threshold(view(lower), view(upper), rhs, min_rows)
489 };
490 let identity = DMatrix::identity(n, n);
491 let weak = &identity * Complex::new(0.1, 0.03);
492 let rhs = patterned(n, 2, 0.4);
493 let iterative = solve(&weak, &weak, rhs.clone(), min_rows).unwrap();
494 prop_assert!(matches!(iterative.factor, InternalFactor::Reflections));
495 let coefficient = Complex::new(1.0, 0.0) - Complex::new(0.1, 0.03).powi(2);
496 prop_assert_close!(&iterative.value * coefficient, &rhs, 2e-13 * rhs.norm());
497 let gradient = iterative
498 .solve_adjoint(view(&weak), view(&weak), rhs.clone())
499 .unwrap()
500 .adjoint;
501 prop_assert_close!(gradient * coefficient.conj(), &rhs, 2e-13 * rhs.norm());
502
503 let at = solve(&weak, &weak, rhs.clone(), n).unwrap();
505 prop_assert!(matches!(at.factor, InternalFactor::Reflections));
506 let below = solve(&weak, &weak, rhs.clone(), n + 1).unwrap();
507 prop_assert!(matches!(below.factor, InternalFactor::Lu(_)));
508 prop_assert_close!(&below.value, &at.value, 2e-13 * rhs.norm());
509 let wide = patterned(n, ITERATIVE_MAX_COLUMNS + 1, 0.4);
510 let batch = solve(&weak, &weak, wide.clone(), min_rows).unwrap();
511 prop_assert!(matches!(batch.factor, InternalFactor::Lu(_)));
512 prop_assert_close!(&batch.value * coefficient, &wide, 2e-13 * wide.norm());
513
514 let mut cyclic = DMatrix::zeros(n, n);
517 for i in 0..n {
518 cyclic[(i, (i + 1) % n)] = Complex::new(0.95, 0.0);
519 }
520 let dense = solve(&cyclic, &identity, rhs.clone(), min_rows).unwrap();
521 prop_assert!(matches!(dense.factor, InternalFactor::Lu(_)));
522 let operator = internal_operator(view(&cyclic), view(&identity));
523 prop_assert!(is_contraction(&operator));
524 prop_assert_close!(&operator * &dense.value, &rhs, 2e-13 * rhs.norm());
525 let eigenvector = DMatrix::from_element(n, 2, Complex::new(1.0, 0.0));
526 let easy_forward = solve(&cyclic, &identity, eigenvector, min_rows).unwrap();
527 prop_assert!(matches!(easy_forward.factor, InternalFactor::Reflections));
528 let tangent = easy_forward
530 .solve_forward(view(&cyclic), view(&identity), rhs.clone())
531 .unwrap();
532 prop_assert_close!(&operator * tangent, &rhs, 2e-13 * rhs.norm());
533 let reverse = easy_forward
536 .solve_adjoint(view(&cyclic), view(&identity), rhs.clone())
537 .unwrap()
538 .adjoint;
539 prop_assert_close!(operator.adjoint() * reverse, &rhs, 2e-13 * rhs.norm());
540
541 prop_assert!(solve(&identity, &identity, rhs, min_rows).is_err());
543 let mut partial = DMatrix::zeros(n, n);
544 partial[(0, 0)] = Complex::new(1.0, 0.0);
545 let mut in_range = DMatrix::zeros(n, 1);
546 in_range[(1, 0)] = Complex::new(1.0, 0.0);
547 prop_assert!(solve(&partial, &identity, in_range, min_rows).is_err());
548 prop_assert!(solve(&identity, &identity, DMatrix::zeros(n, 1), min_rows).is_err());
549 Ok(())
550 }
551}