Skip to main content

tensor4all_core/
krylov.rs

1//! Krylov subspace methods for solving linear equations with abstract tensors.
2//!
3//! This module provides iterative solvers that work with any type implementing [`TensorVectorSpace`],
4//! enabling their use in tensor network algorithms without requiring dense vector representations.
5//!
6//! # Solvers
7//!
8//! - [`gmres`]: Generalized Minimal Residual Method (GMRES) for non-symmetric systems
9//! - [`hermitian_lanczos_lowest_eigenpair`]: Matrix-free lowest eigenpair solver for Hermitian operators
10//!
11//! # Future Extensions
12//!
13//! - CG (Conjugate Gradient) for symmetric positive definite systems
14//! - BiCGSTAB for non-symmetric systems with better convergence properties
15//!
16//! # Example
17//!
18//! ```
19//! use tensor4all_core::{
20//!     krylov::{gmres, GmresOptions},
21//!     DynIndex, TensorDynLen, TensorVectorSpace,
22//! };
23//!
24//! # fn main() -> anyhow::Result<()> {
25//! let i = DynIndex::new_dyn(2);
26//! let rhs = TensorDynLen::from_dense(vec![i.clone()], vec![1.0, -1.0])?;
27//! let initial_guess = TensorDynLen::from_dense(vec![i.clone()], vec![0.0, 0.0])?;
28//!
29//! let apply_operator = |x: &TensorDynLen| Ok(x.clone());
30//! let result = gmres(apply_operator, &rhs, &initial_guess, &GmresOptions::default())?;
31//!
32//! assert!(result.converged);
33//! assert!(result.solution.sub(&rhs)?.maxabs() < 1e-12);
34//! # Ok(())
35//! # }
36//! ```
37
38use crate::any_scalar::AnyScalar;
39use crate::TensorVectorSpace;
40use anyhow::Result;
41use num_complex::Complex64;
42use std::sync::atomic::{AtomicUsize, Ordering};
43use std::time::{Duration, Instant};
44use tensor4all_tensorbackend::{
45    hermitian_exponential_first_column, lowest_hermitian_eigenpair, Matrix,
46};
47
48static GMRES_OP_PROFILE_COUNTER: AtomicUsize = AtomicUsize::new(0);
49
50#[derive(Debug, Clone)]
51struct GmresOpProfile {
52    started: Instant,
53    b_norm: Duration,
54    apply: Duration,
55    inner_product: Duration,
56    axpby: Duration,
57    norm: Duration,
58    scale: Duration,
59    triangular_solve: Duration,
60    solution_update: Duration,
61    apply_calls: usize,
62    inner_product_calls: usize,
63    axpby_calls: usize,
64    norm_calls: usize,
65    scale_calls: usize,
66    triangular_solve_calls: usize,
67    solution_update_calls: usize,
68}
69
70impl Default for GmresOpProfile {
71    fn default() -> Self {
72        Self {
73            started: Instant::now(),
74            b_norm: Duration::ZERO,
75            apply: Duration::ZERO,
76            inner_product: Duration::ZERO,
77            axpby: Duration::ZERO,
78            norm: Duration::ZERO,
79            scale: Duration::ZERO,
80            triangular_solve: Duration::ZERO,
81            solution_update: Duration::ZERO,
82            apply_calls: 0,
83            inner_product_calls: 0,
84            axpby_calls: 0,
85            norm_calls: 0,
86            scale_calls: 0,
87            triangular_solve_calls: 0,
88            solution_update_calls: 0,
89        }
90    }
91}
92
93impl GmresOpProfile {
94    fn measured(&self) -> Duration {
95        self.b_norm
96            + self.apply
97            + self.inner_product
98            + self.axpby
99            + self.norm
100            + self.scale
101            + self.triangular_solve
102            + self.solution_update
103    }
104
105    fn print(&self, id: usize, iterations: usize, residual_norm: f64, converged: bool) {
106        let total = self.started.elapsed();
107        let other = total.saturating_sub(self.measured());
108        eprintln!(
109            "T4A gmres_op_profile #{id}: iterations={iterations} residual={residual_norm:.6e} converged={converged} total_ms={:.3} apply_ms={:.3} apply_calls={} inner_ms={:.3} inner_calls={} axpby_ms={:.3} axpby_calls={} norm_ms={:.3} norm_calls={} scale_ms={:.3} scale_calls={} update_ms={:.3} update_calls={} triangular_ms={:.3} triangular_calls={} b_norm_ms={:.3} other_ms={:.3}",
110            total.as_secs_f64() * 1000.0,
111            self.apply.as_secs_f64() * 1000.0,
112            self.apply_calls,
113            self.inner_product.as_secs_f64() * 1000.0,
114            self.inner_product_calls,
115            self.axpby.as_secs_f64() * 1000.0,
116            self.axpby_calls,
117            self.norm.as_secs_f64() * 1000.0,
118            self.norm_calls,
119            self.scale.as_secs_f64() * 1000.0,
120            self.scale_calls,
121            self.solution_update.as_secs_f64() * 1000.0,
122            self.solution_update_calls,
123            self.triangular_solve.as_secs_f64() * 1000.0,
124            self.triangular_solve_calls,
125            self.b_norm.as_secs_f64() * 1000.0,
126            other.as_secs_f64() * 1000.0,
127        );
128    }
129}
130
131/// Options for GMRES solver.
132///
133/// # Examples
134///
135/// ```
136/// use tensor4all_core::krylov::GmresOptions;
137///
138/// let opts = GmresOptions {
139///     max_iter: 50,
140///     rtol: 1e-8,
141///     max_restarts: 5,
142///     verbose: false,
143///     check_true_residual: true,
144/// };
145/// assert_eq!(opts.max_iter, 50);
146/// assert_eq!(opts.rtol, 1e-8);
147/// ```
148#[derive(Debug, Clone)]
149pub struct GmresOptions {
150    /// Maximum number of iterations (restart cycle length).
151    /// Default: 100
152    pub max_iter: usize,
153
154    /// Convergence tolerance for relative residual norm.
155    /// The solver stops when `||r|| / ||b|| < rtol`.
156    /// Default: 1e-10
157    pub rtol: f64,
158
159    /// Maximum number of restarts.
160    /// Total iterations = max_iter * max_restarts.
161    /// Default: 10
162    pub max_restarts: usize,
163
164    /// Whether to print convergence information.
165    /// Default: false
166    pub verbose: bool,
167
168    /// When true, verify convergence by computing the true residual `||b - A*x|| / ||b||`
169    /// before declaring convergence. This prevents false convergence caused by
170    /// truncation corrupting the Krylov basis orthogonality (see Issue #207).
171    /// Costs one additional `apply_a` call when convergence is detected.
172    /// Default: false
173    pub check_true_residual: bool,
174}
175
176impl Default for GmresOptions {
177    fn default() -> Self {
178        Self {
179            max_iter: 100,
180            rtol: 1e-10,
181            max_restarts: 10,
182            verbose: false,
183            check_true_residual: false,
184        }
185    }
186}
187
188#[derive(Debug, Clone, Copy, PartialEq)]
189enum GmresTolerance {
190    Relative(f64),
191    Absolute(f64),
192}
193
194impl GmresTolerance {
195    fn residual_value(self, residual_norm: f64, b_norm: f64) -> f64 {
196        match self {
197            Self::Relative(_) => residual_norm / b_norm,
198            Self::Absolute(_) => residual_norm,
199        }
200    }
201
202    fn is_converged(self, residual_norm: f64, b_norm: f64) -> bool {
203        match self {
204            Self::Relative(rtol) => residual_norm / b_norm < rtol,
205            Self::Absolute(atol) => residual_norm < atol,
206        }
207    }
208}
209
210/// Result of GMRES solver.
211///
212/// Contains the solution, iteration count, final residual norm, and
213/// convergence status.
214///
215/// # Examples
216///
217/// ```
218/// use tensor4all_core::{DynIndex, TensorDynLen, TensorVectorSpace};
219/// use tensor4all_core::krylov::{gmres, GmresOptions};
220///
221/// let i = DynIndex::new_dyn(2);
222/// let b = TensorDynLen::from_dense(vec![i.clone()], vec![3.0, 7.0]).unwrap();
223/// let x0 = TensorDynLen::from_dense(vec![i.clone()], vec![0.0, 0.0]).unwrap();
224///
225/// let result = gmres(|x: &TensorDynLen| Ok(x.clone()), &b, &x0, &GmresOptions::default()).unwrap();
226/// assert!(result.converged);
227/// assert!(result.residual_norm < 1e-10);
228/// ```
229#[derive(Debug, Clone)]
230pub struct GmresResult<T> {
231    /// The solution vector.
232    pub solution: T,
233
234    /// Number of iterations performed.
235    pub iterations: usize,
236
237    /// Final relative residual norm.
238    pub residual_norm: f64,
239
240    /// Whether the solver converged.
241    pub converged: bool,
242}
243
244/// Options for [`hermitian_lanczos_lowest_eigenpair`].
245///
246/// These options control the maximum Krylov subspace dimension, convergence
247/// tolerance, and validation tolerance for the small projected Hermitian matrix.
248/// When in doubt, use [`Default::default`].
249///
250/// # Examples
251///
252/// ```
253/// use tensor4all_core::krylov::HermitianLanczosOptions;
254///
255/// let opts = HermitianLanczosOptions {
256///     max_iter: 32,
257///     rtol: 1.0e-9,
258///     ..HermitianLanczosOptions::default()
259/// };
260/// assert_eq!(opts.max_iter, 32);
261/// assert_eq!(opts.rtol, 1.0e-9);
262/// ```
263#[derive(Debug, Clone)]
264pub struct HermitianLanczosOptions {
265    /// Maximum Krylov subspace dimension.
266    ///
267    /// Larger values usually improve convergence but increase memory and the
268    /// cost of the small Rayleigh-Ritz eigensolve. Default: `64`.
269    pub max_iter: usize,
270
271    /// Relative residual tolerance.
272    ///
273    /// Convergence requires `||A v - lambda v|| <= max(atol, rtol *
274    /// max(1, |lambda|))`. Default: `1e-10`.
275    pub rtol: f64,
276
277    /// Absolute residual tolerance.
278    ///
279    /// This is useful when targeting eigenvalues near zero. Default: `0.0`.
280    pub atol: f64,
281
282    /// Breakdown threshold for the next Krylov vector norm.
283    ///
284    /// If the orthogonalized vector norm is at or below this value, the current
285    /// Krylov subspace is treated as invariant and iteration stops. Default:
286    /// `1e-14`.
287    pub breakdown_tol: f64,
288
289    /// Hermitian validation tolerance for the projected Rayleigh-Ritz matrix.
290    ///
291    /// The projected operator is rejected when `V^dagger A V` is not Hermitian
292    /// within this absolute tolerance. Default: `1e-10`.
293    pub hermitian_tol: f64,
294
295    /// Whether to print per-iteration residual information.
296    ///
297    /// Default: `false`.
298    pub verbose: bool,
299}
300
301impl Default for HermitianLanczosOptions {
302    fn default() -> Self {
303        Self {
304            max_iter: 64,
305            rtol: 1.0e-10,
306            atol: 0.0,
307            breakdown_tol: 1.0e-14,
308            hermitian_tol: 1.0e-10,
309            verbose: false,
310        }
311    }
312}
313
314/// Result of [`hermitian_lanczos_lowest_eigenpair`].
315///
316/// Contains the lowest Ritz eigenpair, the final true residual norm, iteration
317/// count, and convergence status.
318///
319/// # Examples
320///
321/// ```
322/// use tensor4all_core::krylov::HermitianLanczosResult;
323///
324/// let result = HermitianLanczosResult {
325///     eigenvalue: 1.0,
326///     eigenvector: vec![1.0_f64],
327///     iterations: 1,
328///     residual_norm: 0.0,
329///     converged: true,
330/// };
331/// assert!(result.converged);
332/// assert_eq!(result.eigenvalue, 1.0);
333/// ```
334#[derive(Debug, Clone)]
335pub struct HermitianLanczosResult<T> {
336    /// Lowest Ritz eigenvalue.
337    pub eigenvalue: f64,
338
339    /// Corresponding Ritz eigenvector in the original vector space.
340    pub eigenvector: T,
341
342    /// Number of Krylov operator applications used to build the final subspace.
343    pub iterations: usize,
344
345    /// True residual norm `||A v - lambda v||`.
346    pub residual_norm: f64,
347
348    /// Whether the true residual satisfied the requested tolerance.
349    pub converged: bool,
350}
351
352/// Options for [`hermitian_krylov_expm_multiply`].
353///
354/// The routine builds a Hermitian Lanczos basis for a matrix-free operator,
355/// exponentiates the small projected Hermitian matrix, and combines the basis
356/// vectors without materializing the full operator. Scalar convergence checks
357/// and projected eigendecompositions are explicit non-differentiable
358/// boundaries; vector-space tensor operations preserve backend metadata when
359/// the underlying tensor implementation supports it.
360///
361/// # Examples
362///
363/// ```
364/// use tensor4all_core::krylov::HermitianKrylovExpmOptions;
365///
366/// let options = HermitianKrylovExpmOptions {
367///     max_iter: 32,
368///     tol: 1.0e-10,
369///     ..HermitianKrylovExpmOptions::default()
370/// };
371/// assert_eq!(options.max_iter, 32);
372/// assert_eq!(options.tol, 1.0e-10);
373/// ```
374#[derive(Debug, Clone)]
375pub struct HermitianKrylovExpmOptions {
376    /// Maximum Krylov subspace dimension for one exponential application.
377    pub max_iter: usize,
378    /// Maximum number of equal time splits attempted after non-convergence.
379    pub max_time_splits: usize,
380    /// Relative convergence tolerance for the Krylov residual estimate.
381    pub tol: f64,
382    /// Breakdown threshold for zero vectors and invariant Krylov subspaces.
383    pub breakdown_tol: f64,
384    /// Hermitian validation tolerance for projected matrices.
385    pub hermitian_tol: f64,
386    /// Whether to print per-attempt diagnostics.
387    pub verbose: bool,
388}
389
390impl Default for HermitianKrylovExpmOptions {
391    fn default() -> Self {
392        Self {
393            max_iter: 30,
394            max_time_splits: 100,
395            tol: 1.0e-12,
396            breakdown_tol: 1.0e-14,
397            hermitian_tol: 1.0e-10,
398            verbose: false,
399        }
400    }
401}
402
403/// Result of [`hermitian_krylov_expm_multiply`].
404///
405/// # Examples
406///
407/// ```
408/// use tensor4all_core::krylov::HermitianKrylovExpmResult;
409///
410/// let result = HermitianKrylovExpmResult {
411///     output: vec![1.0_f64, 0.0],
412///     iterations: 1,
413///     matvecs: 1,
414///     error_estimate: 0.0,
415///     converged: true,
416///     time_splits: 1,
417/// };
418/// assert!(result.converged);
419/// assert_eq!(result.output[0], 1.0);
420/// ```
421#[derive(Debug, Clone)]
422pub struct HermitianKrylovExpmResult<T> {
423    /// Approximation to `exp(exponent * A) * initial`.
424    pub output: T,
425    /// Total Krylov iterations over all accepted substeps.
426    pub iterations: usize,
427    /// Total matrix-vector products over all accepted substeps.
428    pub matvecs: usize,
429    /// Maximum Krylov residual estimate over accepted substeps.
430    pub error_estimate: f64,
431    /// Whether the requested tolerance was met.
432    pub converged: bool,
433    /// Number of equal time splits used.
434    pub time_splits: usize,
435}
436
437#[derive(Debug, Clone)]
438struct HermitianRitzState {
439    eigenvalue: f64,
440    coefficients: Vec<Complex64>,
441    iterations: usize,
442    residual_estimate: f64,
443}
444
445/// Compute the lowest eigenpair of a Hermitian matrix-free operator.
446///
447/// The operator is supplied as `apply_a`, which maps a vector to `A * vector`.
448/// The algorithm builds an orthonormal Krylov basis using modified
449/// Gram-Schmidt with a second reorthogonalization pass, forms the small
450/// projected matrix `V^dagger A V`, and solves that small Hermitian problem via
451/// tensorbackend. The full operator matrix is never materialized.
452///
453/// # Arguments
454/// * `apply_a` - Matrix-free Hermitian operator application.
455/// * `initial` - Nonzero initial vector that defines the starting Krylov vector.
456/// * `options` - Krylov dimension, convergence tolerances, and Hermitian
457///   validation settings.
458///
459/// # Returns
460/// The lowest Ritz eigenpair and the true residual norm.
461///
462/// # Errors
463/// Returns an error if the initial vector is zero, a vector-space operation
464/// fails, `apply_a` fails, the projected operator is not Hermitian, or the small
465/// projected eigensolver fails.
466///
467/// # Examples
468///
469/// ```
470/// use tensor4all_core::{DynIndex, TensorDynLen, TensorVectorSpace};
471/// use tensor4all_core::krylov::{hermitian_lanczos_lowest_eigenpair, HermitianLanczosOptions};
472///
473/// let i = DynIndex::new_dyn(2);
474/// let initial = TensorDynLen::from_dense(vec![i.clone()], vec![1.0_f64, 1.0]).unwrap();
475/// let result = hermitian_lanczos_lowest_eigenpair(
476///     |x: &TensorDynLen| Ok(x.clone()),
477///     &initial,
478///     &HermitianLanczosOptions::default(),
479/// ).unwrap();
480///
481/// assert!(result.converged);
482/// assert!((result.eigenvalue - 1.0).abs() < 1.0e-12);
483/// ```
484pub fn hermitian_lanczos_lowest_eigenpair<T, F>(
485    apply_a: F,
486    initial: &T,
487    options: &HermitianLanczosOptions,
488) -> Result<HermitianLanczosResult<T>>
489where
490    T: TensorVectorSpace,
491    F: Fn(&T) -> Result<T>,
492{
493    validate_hermitian_lanczos_options(options)?;
494    initial.validate()?;
495
496    let initial_norm = initial.norm();
497    anyhow::ensure!(
498        initial_norm > options.breakdown_tol,
499        "hermitian_lanczos_lowest_eigenpair: zero initial vector"
500    );
501
502    let mut basis = Vec::with_capacity(options.max_iter + 1);
503    basis.push(initial.scale(AnyScalar::new_real(1.0 / initial_norm))?);
504    let mut h_cols: Vec<Vec<Complex64>> = Vec::with_capacity(options.max_iter);
505    let mut last_ritz: Option<HermitianRitzState> = None;
506
507    for j in 0..options.max_iter {
508        let w = apply_a(&basis[j])?;
509        if j == 0 {
510            w.validate()?;
511        }
512
513        let mut h_col = Vec::with_capacity(j + 2);
514        let mut w_orth = w;
515
516        for v_i in basis.iter().take(j + 1) {
517            let h_ij = v_i.inner_product(&w_orth)?;
518            h_col.push(any_scalar_to_complex(&h_ij));
519            let neg_h_ij = AnyScalar::new_real(0.0) - h_ij;
520            w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_h_ij)?;
521        }
522
523        for (i, v_i) in basis.iter().take(j + 1).enumerate() {
524            let correction = v_i.inner_product(&w_orth)?;
525            h_col[i] += any_scalar_to_complex(&correction);
526            let neg_correction = AnyScalar::new_real(0.0) - correction;
527            w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_correction)?;
528        }
529
530        let beta = w_orth.norm();
531        h_col.push(Complex64::new(beta, 0.0));
532        h_cols.push(h_col);
533
534        let subspace_dim = j + 1;
535        let projected = projected_matrix_from_columns(&h_cols, subspace_dim);
536        let ritz = lowest_hermitian_eigenpair(&projected, options.hermitian_tol)
537            .map_err(|err| anyhow::anyhow!("projected operator is not Hermitian: {err}"))?;
538        let residual_estimate = beta * ritz.eigenvector[subspace_dim - 1].norm();
539        let threshold = hermitian_lanczos_threshold(ritz.eigenvalue, options);
540        let estimate_converged = residual_estimate <= threshold;
541
542        if options.verbose {
543            eprintln!(
544                "Hermitian Lanczos iter {}: eigenvalue={:.16e} residual_estimate={:.6e} threshold={:.6e}",
545                subspace_dim, ritz.eigenvalue, residual_estimate, threshold
546            );
547        }
548
549        let ritz_state = HermitianRitzState {
550            eigenvalue: ritz.eigenvalue,
551            iterations: subspace_dim,
552            residual_estimate,
553            coefficients: ritz.eigenvector,
554        };
555        if estimate_converged || beta <= options.breakdown_tol {
556            let result = finalize_hermitian_lanczos_result(
557                &apply_a,
558                &basis[..subspace_dim],
559                &ritz_state,
560                options,
561            )?;
562            if result.converged || beta <= options.breakdown_tol {
563                return Ok(result);
564            }
565        }
566        last_ritz = Some(ritz_state);
567        basis.push(w_orth.scale(AnyScalar::new_real(1.0 / beta))?);
568    }
569
570    let ritz_state = last_ritz.ok_or_else(|| {
571        anyhow::anyhow!("hermitian_lanczos_lowest_eigenpair: max_iter must be greater than zero")
572    })?;
573    finalize_hermitian_lanczos_result(
574        &apply_a,
575        &basis[..ritz_state.iterations],
576        &ritz_state,
577        options,
578    )
579}
580
581/// Apply `exp(exponent * A)` to a vector using a matrix-free Hermitian Krylov method.
582///
583/// This routine only materializes small projected Krylov matrices. It never
584/// forms the full matrix for `A`, so it can be used by tensor-network local
585/// solvers whose operator application is available as a closure.
586///
587/// # Arguments
588/// * `apply_a` - Matrix-free Hermitian operator application.
589/// * `exponent` - Scalar exponent, for example `-im * dt` for real-time TDVP.
590/// * `initial` - Initial vector.
591/// * `options` - Krylov dimension, tolerance, and Hermitian validation options.
592///
593/// # Returns
594/// A [`HermitianKrylovExpmResult`] containing the evolved vector and diagnostics.
595///
596/// # Errors
597/// Returns an error if an option is invalid, a vector-space operation fails,
598/// the projected matrix is not Hermitian, or the requested tolerance is not met
599/// after the configured time splits.
600///
601/// # Examples
602///
603/// ```
604/// use num_complex::Complex64;
605/// use tensor4all_core::krylov::{
606///     hermitian_krylov_expm_multiply, HermitianKrylovExpmOptions,
607/// };
608/// use tensor4all_core::{DynIndex, TensorDynLen};
609///
610/// # fn main() -> anyhow::Result<()> {
611/// let index = DynIndex::new_dyn(2);
612/// let initial = TensorDynLen::from_dense(
613///     vec![index.clone()],
614///     vec![Complex64::new(1.0, 0.0), Complex64::new(0.0, 0.0)],
615/// )?;
616/// let options = HermitianKrylovExpmOptions {
617///     max_iter: 4,
618///     tol: 1.0e-12,
619///     ..Default::default()
620/// };
621/// let result = hermitian_krylov_expm_multiply(
622///     |x: &TensorDynLen| {
623///         let data = x.to_vec::<Complex64>()?;
624///         TensorDynLen::from_dense(
625///             vec![index.clone()],
626///             vec![data[0], Complex64::new(2.0, 0.0) * data[1]],
627///         )
628///     },
629///     Complex64::new(0.0, -0.25),
630///     &initial,
631///     &options,
632/// )?;
633/// let evolved = result.output.to_vec::<Complex64>()?;
634/// let expected = Complex64::new(0.25_f64.cos(), -0.25_f64.sin());
635/// assert!((evolved[0] - expected).norm() < 1.0e-10);
636/// assert!(evolved[1].norm() < 1.0e-12);
637/// # Ok(())
638/// # }
639/// ```
640pub fn hermitian_krylov_expm_multiply<T, F>(
641    mut apply_a: F,
642    exponent: Complex64,
643    initial: &T,
644    options: &HermitianKrylovExpmOptions,
645) -> Result<HermitianKrylovExpmResult<T>>
646where
647    T: TensorVectorSpace,
648    F: FnMut(&T) -> Result<T>,
649{
650    validate_hermitian_krylov_expm_options(options)?;
651    initial.validate()?;
652
653    if exponent == Complex64::new(0.0, 0.0) {
654        return Ok(HermitianKrylovExpmResult {
655            output: initial.clone(),
656            iterations: 0,
657            matvecs: 0,
658            error_estimate: 0.0,
659            converged: true,
660            time_splits: 1,
661        });
662    }
663    if initial.norm() <= options.breakdown_tol {
664        return Ok(HermitianKrylovExpmResult {
665            output: initial.clone(),
666            iterations: 0,
667            matvecs: 0,
668            error_estimate: 0.0,
669            converged: true,
670            time_splits: 1,
671        });
672    }
673
674    let mut splits = 1usize;
675    loop {
676        let step_exponent = exponent / splits as f64;
677        let mut output = initial.clone();
678        let mut iterations = 0usize;
679        let mut matvecs = 0usize;
680        let mut max_error = 0.0_f64;
681        let mut converged = true;
682
683        for _ in 0..splits {
684            let result =
685                hermitian_krylov_expm_multiply_once(&mut apply_a, step_exponent, &output, options)?;
686            iterations += result.iterations;
687            matvecs += result.matvecs;
688            max_error = max_error.max(result.error_estimate);
689            output = result.output;
690            if !result.converged {
691                converged = false;
692                break;
693            }
694        }
695
696        if converged {
697            return Ok(HermitianKrylovExpmResult {
698                output,
699                iterations,
700                matvecs,
701                error_estimate: max_error,
702                converged: true,
703                time_splits: splits,
704            });
705        }
706
707        if options.verbose {
708            eprintln!(
709                "Hermitian Krylov expm did not converge with {splits} time split(s); retrying"
710            );
711        }
712        if splits >= options.max_time_splits {
713            return Err(anyhow::anyhow!(
714                "hermitian_krylov_expm_multiply did not converge within max_time_splits={} and max_iter={}",
715                options.max_time_splits,
716                options.max_iter
717            ));
718        }
719        splits = splits.saturating_mul(2).min(options.max_time_splits);
720    }
721}
722
723fn hermitian_krylov_expm_multiply_once<T, F>(
724    apply_a: &mut F,
725    exponent: Complex64,
726    initial: &T,
727    options: &HermitianKrylovExpmOptions,
728) -> Result<HermitianKrylovExpmResult<T>>
729where
730    T: TensorVectorSpace,
731    F: FnMut(&T) -> Result<T>,
732{
733    let initial_norm = initial.norm();
734    if initial_norm <= options.breakdown_tol {
735        return Ok(HermitianKrylovExpmResult {
736            output: initial.clone(),
737            iterations: 0,
738            matvecs: 0,
739            error_estimate: 0.0,
740            converged: true,
741            time_splits: 1,
742        });
743    }
744
745    let mut basis = Vec::with_capacity(options.max_iter + 1);
746    basis.push(initial.scale(AnyScalar::new_real(1.0 / initial_norm))?);
747    let mut h_cols: Vec<Vec<Complex64>> = Vec::with_capacity(options.max_iter);
748    let mut last_coefficients: Option<Vec<Complex64>> = None;
749    let mut last_error_estimate = f64::INFINITY;
750    let mut last_subspace_dim = 0usize;
751    let threshold = options.tol * initial_norm.max(1.0);
752
753    for j in 0..options.max_iter {
754        let w = apply_a(&basis[j])?;
755        if j == 0 {
756            w.validate()?;
757        }
758        let mut w_orth = w;
759        let mut h_col = Vec::with_capacity(j + 2);
760
761        for v_i in basis.iter().take(j + 1) {
762            let h_ij = v_i.inner_product(&w_orth)?;
763            h_col.push(any_scalar_to_complex(&h_ij));
764            let neg_h_ij = AnyScalar::new_real(0.0) - h_ij;
765            w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_h_ij)?;
766        }
767
768        for (i, v_i) in basis.iter().take(j + 1).enumerate() {
769            let correction = v_i.inner_product(&w_orth)?;
770            h_col[i] += any_scalar_to_complex(&correction);
771            let neg_correction = AnyScalar::new_real(0.0) - correction;
772            w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_correction)?;
773        }
774
775        let beta_next = w_orth.norm();
776        h_col.push(Complex64::new(beta_next, 0.0));
777        h_cols.push(h_col);
778
779        let subspace_dim = j + 1;
780        let projected = projected_matrix_from_columns(&h_cols, subspace_dim);
781        let coefficients =
782            hermitian_exponential_first_column(&projected, exponent, options.hermitian_tol)
783                .map_err(|err| anyhow::anyhow!("projected operator is not Hermitian: {err}"))?;
784        let error_estimate = if beta_next <= options.breakdown_tol {
785            0.0
786        } else {
787            initial_norm * beta_next * coefficients[subspace_dim - 1].norm()
788        };
789        let converged = beta_next <= options.breakdown_tol || error_estimate <= threshold;
790
791        if options.verbose {
792            eprintln!(
793                "Hermitian Krylov expm iter {}: residual_estimate={:.6e} threshold={:.6e}",
794                subspace_dim, error_estimate, threshold
795            );
796        }
797
798        if converged {
799            let output = finalize_hermitian_krylov_expm_output(
800                &basis[..subspace_dim],
801                &coefficients,
802                initial_norm,
803                options.hermitian_tol,
804            )?;
805            return Ok(HermitianKrylovExpmResult {
806                output,
807                iterations: subspace_dim,
808                matvecs: subspace_dim,
809                error_estimate,
810                converged: true,
811                time_splits: 1,
812            });
813        }
814        last_coefficients = Some(coefficients);
815        last_error_estimate = error_estimate;
816        last_subspace_dim = subspace_dim;
817        basis.push(w_orth.scale(AnyScalar::new_real(1.0 / beta_next))?);
818    }
819
820    let coefficients = last_coefficients.ok_or_else(|| {
821        anyhow::anyhow!("hermitian_krylov_expm_multiply: max_iter must be greater than zero")
822    })?;
823    let output = finalize_hermitian_krylov_expm_output(
824        &basis[..last_subspace_dim],
825        &coefficients,
826        initial_norm,
827        options.hermitian_tol,
828    )?;
829    Ok(HermitianKrylovExpmResult {
830        output,
831        iterations: last_subspace_dim,
832        matvecs: last_subspace_dim,
833        error_estimate: last_error_estimate,
834        converged: false,
835        time_splits: 1,
836    })
837}
838
839fn finalize_hermitian_krylov_expm_output<T>(
840    basis: &[T],
841    coefficients: &[Complex64],
842    initial_norm: f64,
843    hermitian_tol: f64,
844) -> Result<T>
845where
846    T: TensorVectorSpace,
847{
848    let unit_output = combine_basis_with_complex_coefficients(
849        basis,
850        coefficients,
851        hermitian_tol,
852        "hermitian_krylov_expm_multiply",
853    )?;
854    unit_output.scale(AnyScalar::new_real(initial_norm))
855}
856
857/// Solve `A x = b` using GMRES (Generalized Minimal Residual Method).
858///
859/// This implements the restarted GMRES algorithm that works with abstract tensor types
860/// through the [`TensorVectorSpace`] trait's vector space operations.
861///
862/// # Algorithm
863///
864/// GMRES builds an orthonormal basis for the Krylov subspace
865/// `K_m = span{r_0, A r_0, A^2 r_0, ..., A^{m-1} r_0}` and finds the
866/// solution that minimizes `||b - A x||` over this subspace.
867///
868/// # Type Parameters
869///
870/// * `T` - A tensor type implementing `TensorVectorSpace`
871/// * `F` - A function that applies the linear operator: `F(x) = A x`
872///
873/// # Arguments
874///
875/// * `apply_a` - Function that applies the linear operator A to a tensor
876/// * `b` - Right-hand side tensor
877/// * `x0` - Initial guess
878/// * `options` - Solver options
879///
880/// # Returns
881///
882/// A `GmresResult` containing the solution and convergence information.
883///
884/// # Errors
885///
886/// Returns an error if:
887/// - Vector space operations (add, sub, scale, inner_product) fail
888/// - The linear operator application fails
889pub fn gmres<T, F>(apply_a: F, b: &T, x0: &T, options: &GmresOptions) -> Result<GmresResult<T>>
890where
891    T: TensorVectorSpace,
892    F: Fn(&T) -> Result<T>,
893{
894    gmres_impl(
895        apply_a,
896        b,
897        x0,
898        options,
899        GmresTolerance::Relative(options.rtol),
900        None,
901    )
902}
903
904/// Solve `A x = b` using GMRES with an absolute residual tolerance.
905///
906/// This variant stops when `||b - A*x|| < atol`. The default [`gmres`] API uses
907/// relative residual tolerance and is preferred for scale-independent solves.
908pub fn gmres_with_absolute_tolerance<T, F>(
909    apply_a: F,
910    b: &T,
911    x0: &T,
912    options: &GmresOptions,
913    atol: f64,
914) -> Result<GmresResult<T>>
915where
916    T: TensorVectorSpace,
917    F: Fn(&T) -> Result<T>,
918{
919    gmres_impl(
920        apply_a,
921        b,
922        x0,
923        options,
924        GmresTolerance::Absolute(atol),
925        None,
926    )
927}
928
929/// Solve `(a0 I + a1 A) x = b` using GMRES with relative residual tolerance.
930///
931/// The Arnoldi basis is built from the unshifted `A` callback, while affine
932/// coefficients are applied in the projected Hessenberg problem, matching
933/// KrylovKit's affine linear-solve convention.
934pub fn gmres_affine<T, F>(
935    apply_a: F,
936    b: &T,
937    x0: &T,
938    a0: AnyScalar,
939    a1: AnyScalar,
940    options: &GmresOptions,
941) -> Result<GmresResult<T>>
942where
943    T: TensorVectorSpace,
944    F: Fn(&T) -> Result<T>,
945{
946    gmres_affine_impl(
947        apply_a,
948        b,
949        x0,
950        a0,
951        a1,
952        options,
953        GmresTolerance::Relative(options.rtol),
954    )
955}
956
957/// Solve `(a0 I + a1 A) x = b` using GMRES with an absolute residual tolerance.
958///
959/// The Arnoldi basis is built from the unshifted `A` callback, while the affine
960/// coefficients are applied to the small Hessenberg problem. This mirrors
961/// KrylovKit's `linsolve(operator, b, a0, a1)` algorithm and avoids changing the
962/// Krylov basis when affine coefficients are present.
963pub fn gmres_affine_with_absolute_tolerance<T, F>(
964    apply_a: F,
965    b: &T,
966    x0: &T,
967    a0: AnyScalar,
968    a1: AnyScalar,
969    options: &GmresOptions,
970    atol: f64,
971) -> Result<GmresResult<T>>
972where
973    T: TensorVectorSpace,
974    F: Fn(&T) -> Result<T>,
975{
976    gmres_affine_impl(
977        apply_a,
978        b,
979        x0,
980        a0,
981        a1,
982        options,
983        GmresTolerance::Absolute(atol),
984    )
985}
986
987fn gmres_affine_impl<T, F>(
988    apply_a: F,
989    b: &T,
990    x0: &T,
991    a0: AnyScalar,
992    a1: AnyScalar,
993    options: &GmresOptions,
994    tolerance: GmresTolerance,
995) -> Result<GmresResult<T>>
996where
997    T: TensorVectorSpace,
998    F: Fn(&T) -> Result<T>,
999{
1000    b.validate()?;
1001    x0.validate()?;
1002
1003    let profile_enabled = std::env::var_os("T4A_GMRES_OP_PROFILE").is_some();
1004    let profile_id = if profile_enabled {
1005        GMRES_OP_PROFILE_COUNTER.fetch_add(1, Ordering::Relaxed)
1006    } else {
1007        0
1008    };
1009    let mut profile = GmresOpProfile::default();
1010    let mut total_iters = 0usize;
1011
1012    macro_rules! finish {
1013        ($result:expr) => {{
1014            let result = $result;
1015            if profile_enabled {
1016                profile.print(
1017                    profile_id,
1018                    result.iterations,
1019                    result.residual_norm,
1020                    result.converged,
1021                );
1022            }
1023            return Ok(result);
1024        }};
1025    }
1026
1027    let started = Instant::now();
1028    let b_norm = b.norm();
1029    if profile_enabled {
1030        profile.b_norm += started.elapsed();
1031    }
1032    if b_norm < 1e-15 {
1033        finish!(GmresResult {
1034            solution: x0.clone(),
1035            iterations: 0,
1036            residual_norm: 0.0,
1037            converged: true,
1038        });
1039    }
1040    if a0.is_zero() && a1.is_zero() {
1041        anyhow::bail!("gmres_affine: at least one affine coefficient must be nonzero");
1042    }
1043    if a1.is_zero() {
1044        let started = Instant::now();
1045        let solution = b.scale(AnyScalar::new_real(1.0) / a0)?;
1046        if profile_enabled {
1047            profile.scale += started.elapsed();
1048            profile.scale_calls += 1;
1049        }
1050        finish!(GmresResult {
1051            solution,
1052            iterations: 0,
1053            residual_norm: 0.0,
1054            converged: true,
1055        });
1056    }
1057
1058    let mut x = x0.clone();
1059
1060    for restart in 0..options.max_restarts {
1061        let started = Instant::now();
1062        let ax = apply_a(&x)?;
1063        if profile_enabled {
1064            profile.apply += started.elapsed();
1065            profile.apply_calls += 1;
1066        }
1067        if restart == 0 {
1068            ax.validate()?;
1069        }
1070        let started = Instant::now();
1071        let affine_x = x.axpby(a0.clone(), &ax, a1.clone())?;
1072        if profile_enabled {
1073            profile.axpby += started.elapsed();
1074            profile.axpby_calls += 1;
1075        }
1076        let started = Instant::now();
1077        let r = b.axpby(
1078            AnyScalar::new_real(1.0),
1079            &affine_x,
1080            AnyScalar::new_real(-1.0),
1081        )?;
1082        if profile_enabled {
1083            profile.axpby += started.elapsed();
1084            profile.axpby_calls += 1;
1085        }
1086        let started = Instant::now();
1087        let r_norm = r.norm();
1088        if profile_enabled {
1089            profile.norm += started.elapsed();
1090            profile.norm_calls += 1;
1091        }
1092        let residual_value = tolerance.residual_value(r_norm, b_norm);
1093        if options.verbose {
1094            eprintln!(
1095                "GMRES restart {}: initial residual = {:.6e}",
1096                restart, residual_value
1097            );
1098        }
1099        if tolerance.is_converged(r_norm, b_norm) {
1100            finish!(GmresResult {
1101                solution: x,
1102                iterations: total_iters,
1103                residual_norm: residual_value,
1104                converged: true,
1105            });
1106        }
1107
1108        let cycle_max_iter = options.max_iter;
1109        let mut v_basis: Vec<T> = Vec::with_capacity(cycle_max_iter + 1);
1110        let started = Instant::now();
1111        v_basis.push(r.scale(AnyScalar::new_real(1.0 / r_norm))?);
1112        if profile_enabled {
1113            profile.scale += started.elapsed();
1114            profile.scale_calls += 1;
1115        }
1116
1117        let mut h_matrix: Vec<Vec<AnyScalar>> = Vec::with_capacity(cycle_max_iter);
1118        let mut cs: Vec<AnyScalar> = Vec::with_capacity(cycle_max_iter);
1119        let mut sn: Vec<AnyScalar> = Vec::with_capacity(cycle_max_iter);
1120        let mut g: Vec<AnyScalar> = vec![AnyScalar::new_real(r_norm)];
1121        let mut solution_already_updated = false;
1122
1123        for j in 0..cycle_max_iter {
1124            total_iters += 1;
1125
1126            let started = Instant::now();
1127            let w = apply_a(&v_basis[j])?;
1128            if profile_enabled {
1129                profile.apply += started.elapsed();
1130                profile.apply_calls += 1;
1131            }
1132            let mut h_a_col: Vec<AnyScalar> = Vec::with_capacity(j + 2);
1133            let mut w_orth = w;
1134
1135            for v_i in v_basis.iter().take(j + 1) {
1136                let started = Instant::now();
1137                let h_ij = v_i.inner_product(&w_orth)?;
1138                if profile_enabled {
1139                    profile.inner_product += started.elapsed();
1140                    profile.inner_product_calls += 1;
1141                }
1142                h_a_col.push(h_ij.clone());
1143                let neg_h_ij = AnyScalar::new_real(0.0) - h_ij;
1144                let started = Instant::now();
1145                w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_h_ij)?;
1146                if profile_enabled {
1147                    profile.axpby += started.elapsed();
1148                    profile.axpby_calls += 1;
1149                }
1150            }
1151            for (i, v_i) in v_basis.iter().take(j + 1).enumerate() {
1152                let started = Instant::now();
1153                let correction = v_i.inner_product(&w_orth)?;
1154                if profile_enabled {
1155                    profile.inner_product += started.elapsed();
1156                    profile.inner_product_calls += 1;
1157                }
1158                h_a_col[i] = h_a_col[i].clone() + correction.clone();
1159                let neg_correction = AnyScalar::new_real(0.0) - correction;
1160                let started = Instant::now();
1161                w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_correction)?;
1162                if profile_enabled {
1163                    profile.axpby += started.elapsed();
1164                    profile.axpby_calls += 1;
1165                }
1166            }
1167
1168            let started = Instant::now();
1169            let h_jp1_j_real = w_orth.norm();
1170            if profile_enabled {
1171                profile.norm += started.elapsed();
1172                profile.norm_calls += 1;
1173            }
1174            h_a_col.push(AnyScalar::new_real(h_jp1_j_real));
1175
1176            let mut h_col: Vec<AnyScalar> = Vec::with_capacity(j + 2);
1177            for h in h_a_col.iter().take(j) {
1178                h_col.push(a1.clone() * h.clone());
1179            }
1180            h_col.push(a0.clone() + a1.clone() * h_a_col[j].clone());
1181            h_col.push(a1.clone() * h_a_col[j + 1].clone());
1182
1183            #[allow(clippy::needless_range_loop)]
1184            for i in 0..j {
1185                let h_i = h_col[i].clone();
1186                let h_ip1 = h_col[i + 1].clone();
1187                let (new_hi, new_hip1) = apply_givens_rotation(&cs[i], &sn[i], &h_i, &h_ip1);
1188                h_col[i] = new_hi;
1189                h_col[i + 1] = new_hip1;
1190            }
1191
1192            let (c_j, s_j) = compute_givens_rotation(&h_col[j], &h_col[j + 1]);
1193            cs.push(c_j.clone());
1194            sn.push(s_j.clone());
1195
1196            let (new_hj, _) = apply_givens_rotation(&c_j, &s_j, &h_col[j], &h_col[j + 1]);
1197            h_col[j] = new_hj;
1198            h_col[j + 1] = AnyScalar::new_real(0.0);
1199
1200            let g_j = g[j].clone();
1201            let g_jp1 = AnyScalar::new_real(0.0);
1202            let (new_gj, new_gjp1) = apply_givens_rotation(&c_j, &s_j, &g_j, &g_jp1);
1203            g[j] = new_gj;
1204            let res_norm = new_gjp1.abs();
1205            g.push(new_gjp1);
1206
1207            h_matrix.push(h_col);
1208            let residual_value = tolerance.residual_value(res_norm, b_norm);
1209            if options.verbose {
1210                eprintln!("GMRES iter {}: residual = {:.6e}", j + 1, residual_value);
1211            }
1212
1213            if tolerance.is_converged(res_norm, b_norm) {
1214                let started = Instant::now();
1215                let y = solve_upper_triangular(&h_matrix, &g[..=j])?;
1216                if profile_enabled {
1217                    profile.triangular_solve += started.elapsed();
1218                    profile.triangular_solve_calls += 1;
1219                }
1220                let started = Instant::now();
1221                x = update_solution(&x, &v_basis[..=j], &y)?;
1222                if profile_enabled {
1223                    profile.solution_update += started.elapsed();
1224                    profile.solution_update_calls += 1;
1225                }
1226                if options.check_true_residual {
1227                    let started = Instant::now();
1228                    let ax_check = apply_a(&x)?;
1229                    if profile_enabled {
1230                        profile.apply += started.elapsed();
1231                        profile.apply_calls += 1;
1232                    }
1233                    let started = Instant::now();
1234                    let affine_check = x.axpby(a0.clone(), &ax_check, a1.clone())?;
1235                    if profile_enabled {
1236                        profile.axpby += started.elapsed();
1237                        profile.axpby_calls += 1;
1238                    }
1239                    let started = Instant::now();
1240                    let r_check = b.axpby(
1241                        AnyScalar::new_real(1.0),
1242                        &affine_check,
1243                        AnyScalar::new_real(-1.0),
1244                    )?;
1245                    if profile_enabled {
1246                        profile.axpby += started.elapsed();
1247                        profile.axpby_calls += 1;
1248                    }
1249                    let started = Instant::now();
1250                    let true_abs_res = r_check.norm();
1251                    if profile_enabled {
1252                        profile.norm += started.elapsed();
1253                        profile.norm_calls += 1;
1254                    }
1255                    let true_residual_value = tolerance.residual_value(true_abs_res, b_norm);
1256                    if options.verbose {
1257                        eprintln!(
1258                            "GMRES true residual check: hessenberg={:.6e}, checked={:.6e}",
1259                            residual_value, true_residual_value
1260                        );
1261                    }
1262                    if tolerance.is_converged(true_abs_res, b_norm) {
1263                        finish!(GmresResult {
1264                            solution: x,
1265                            iterations: total_iters,
1266                            residual_norm: true_residual_value,
1267                            converged: true,
1268                        });
1269                    }
1270                    solution_already_updated = true;
1271                    break;
1272                } else {
1273                    finish!(GmresResult {
1274                        solution: x,
1275                        iterations: total_iters,
1276                        residual_norm: residual_value,
1277                        converged: true,
1278                    });
1279                }
1280            }
1281
1282            if h_jp1_j_real > 1e-14 {
1283                let started = Instant::now();
1284                v_basis.push(w_orth.scale(AnyScalar::new_real(1.0 / h_jp1_j_real))?);
1285                if profile_enabled {
1286                    profile.scale += started.elapsed();
1287                    profile.scale_calls += 1;
1288                }
1289            } else {
1290                let started = Instant::now();
1291                let y = solve_upper_triangular(&h_matrix, &g[..=j])?;
1292                if profile_enabled {
1293                    profile.triangular_solve += started.elapsed();
1294                    profile.triangular_solve_calls += 1;
1295                }
1296                let started = Instant::now();
1297                x = update_solution(&x, &v_basis[..=j], &y)?;
1298                if profile_enabled {
1299                    profile.solution_update += started.elapsed();
1300                    profile.solution_update_calls += 1;
1301                }
1302                let started = Instant::now();
1303                let ax_final = apply_a(&x)?;
1304                if profile_enabled {
1305                    profile.apply += started.elapsed();
1306                    profile.apply_calls += 1;
1307                }
1308                let started = Instant::now();
1309                let affine_final = x.axpby(a0.clone(), &ax_final, a1.clone())?;
1310                if profile_enabled {
1311                    profile.axpby += started.elapsed();
1312                    profile.axpby_calls += 1;
1313                }
1314                let started = Instant::now();
1315                let r_final = b.axpby(
1316                    AnyScalar::new_real(1.0),
1317                    &affine_final,
1318                    AnyScalar::new_real(-1.0),
1319                )?;
1320                if profile_enabled {
1321                    profile.axpby += started.elapsed();
1322                    profile.axpby_calls += 1;
1323                }
1324                let started = Instant::now();
1325                let final_abs_res = r_final.norm();
1326                if profile_enabled {
1327                    profile.norm += started.elapsed();
1328                    profile.norm_calls += 1;
1329                }
1330                let final_res = tolerance.residual_value(final_abs_res, b_norm);
1331                finish!(GmresResult {
1332                    solution: x,
1333                    iterations: total_iters,
1334                    residual_norm: final_res,
1335                    converged: tolerance.is_converged(final_abs_res, b_norm),
1336                });
1337            }
1338        }
1339
1340        if !solution_already_updated {
1341            let actual_iters = h_matrix.len();
1342            let started = Instant::now();
1343            let y = solve_upper_triangular(&h_matrix, &g[..actual_iters])?;
1344            if profile_enabled {
1345                profile.triangular_solve += started.elapsed();
1346                profile.triangular_solve_calls += 1;
1347            }
1348            let started = Instant::now();
1349            x = update_solution(&x, &v_basis[..actual_iters], &y)?;
1350            if profile_enabled {
1351                profile.solution_update += started.elapsed();
1352                profile.solution_update_calls += 1;
1353            }
1354        }
1355    }
1356
1357    let started = Instant::now();
1358    let ax_final = apply_a(&x)?;
1359    if profile_enabled {
1360        profile.apply += started.elapsed();
1361        profile.apply_calls += 1;
1362    }
1363    let started = Instant::now();
1364    let affine_final = x.axpby(a0, &ax_final, a1)?;
1365    if profile_enabled {
1366        profile.axpby += started.elapsed();
1367        profile.axpby_calls += 1;
1368    }
1369    let started = Instant::now();
1370    let r_final = b.axpby(
1371        AnyScalar::new_real(1.0),
1372        &affine_final,
1373        AnyScalar::new_real(-1.0),
1374    )?;
1375    if profile_enabled {
1376        profile.axpby += started.elapsed();
1377        profile.axpby_calls += 1;
1378    }
1379    let started = Instant::now();
1380    let final_abs_res = r_final.norm();
1381    if profile_enabled {
1382        profile.norm += started.elapsed();
1383        profile.norm_calls += 1;
1384    }
1385    let final_res = tolerance.residual_value(final_abs_res, b_norm);
1386
1387    finish!(GmresResult {
1388        solution: x,
1389        iterations: total_iters,
1390        residual_norm: final_res,
1391        converged: tolerance.is_converged(final_abs_res, b_norm),
1392    })
1393}
1394
1395/// Solve `A x = b` using GMRES while enforcing a total iteration limit.
1396///
1397/// [`GmresOptions::max_iter`] remains the restart cycle length and
1398/// [`GmresOptions::max_restarts`] remains the maximum number of restart cycles.
1399/// `max_total_iter` caps the total number of Arnoldi steps across all restart
1400/// cycles; the final cycle is shortened when necessary.
1401pub fn gmres_with_total_iteration_limit<T, F>(
1402    apply_a: F,
1403    b: &T,
1404    x0: &T,
1405    options: &GmresOptions,
1406    max_total_iter: usize,
1407) -> Result<GmresResult<T>>
1408where
1409    T: TensorVectorSpace,
1410    F: Fn(&T) -> Result<T>,
1411{
1412    gmres_impl(
1413        apply_a,
1414        b,
1415        x0,
1416        options,
1417        GmresTolerance::Relative(options.rtol),
1418        Some(max_total_iter),
1419    )
1420}
1421
1422fn gmres_impl<T, F>(
1423    apply_a: F,
1424    b: &T,
1425    x0: &T,
1426    options: &GmresOptions,
1427    tolerance: GmresTolerance,
1428    max_total_iter: Option<usize>,
1429) -> Result<GmresResult<T>>
1430where
1431    T: TensorVectorSpace,
1432    F: Fn(&T) -> Result<T>,
1433{
1434    // Validate structural consistency of inputs
1435    b.validate()?;
1436    x0.validate()?;
1437
1438    let b_norm = b.norm();
1439    if b_norm < 1e-15 {
1440        // b is effectively zero, return x0
1441        return Ok(GmresResult {
1442            solution: x0.clone(),
1443            iterations: 0,
1444            residual_norm: 0.0,
1445            converged: true,
1446        });
1447    }
1448
1449    let mut x = x0.clone();
1450    let mut total_iters = 0;
1451
1452    for _restart in 0..options.max_restarts {
1453        let cycle_max_iter = match max_total_iter {
1454            Some(limit) => {
1455                let remaining = limit.saturating_sub(total_iters);
1456                if remaining == 0 {
1457                    break;
1458                }
1459                options.max_iter.min(remaining)
1460            }
1461            None => options.max_iter,
1462        };
1463        if cycle_max_iter == 0 {
1464            break;
1465        }
1466
1467        // Compute initial residual: r = b - A*x
1468        let ax = apply_a(&x)?;
1469        // Validate operator output on first restart
1470        if _restart == 0 {
1471            ax.validate()?;
1472        }
1473        // r = 1.0 * b + (-1.0) * ax
1474        let r = b.axpby(AnyScalar::new_real(1.0), &ax, AnyScalar::new_real(-1.0))?;
1475        let r_norm = r.norm();
1476        let residual_value = tolerance.residual_value(r_norm, b_norm);
1477
1478        if options.verbose {
1479            eprintln!(
1480                "GMRES restart {}: initial residual = {:.6e}",
1481                _restart, residual_value
1482            );
1483        }
1484
1485        if tolerance.is_converged(r_norm, b_norm) {
1486            return Ok(GmresResult {
1487                solution: x,
1488                iterations: total_iters,
1489                residual_norm: residual_value,
1490                converged: true,
1491            });
1492        }
1493
1494        // Arnoldi process with modified Gram-Schmidt
1495        let mut v_basis: Vec<T> = Vec::with_capacity(cycle_max_iter + 1);
1496        let mut h_matrix: Vec<Vec<AnyScalar>> = Vec::with_capacity(cycle_max_iter);
1497
1498        // v_0 = r / ||r||
1499        let v0 = r.scale(AnyScalar::new_real(1.0 / r_norm))?;
1500        v_basis.push(v0);
1501
1502        // Initialize Givens rotation storage
1503        let mut cs: Vec<AnyScalar> = Vec::with_capacity(cycle_max_iter);
1504        let mut sn: Vec<AnyScalar> = Vec::with_capacity(cycle_max_iter);
1505        let mut g: Vec<AnyScalar> = vec![AnyScalar::new_real(r_norm)]; // residual in upper Hessenberg space
1506        let mut solution_already_updated = false;
1507
1508        for j in 0..cycle_max_iter {
1509            total_iters += 1;
1510
1511            // w = A * v_j
1512            let w = apply_a(&v_basis[j])?;
1513
1514            // Modified Gram-Schmidt orthogonalization
1515            let mut h_col: Vec<AnyScalar> = Vec::with_capacity(j + 2);
1516            let mut w_orth = w;
1517
1518            for v_i in v_basis.iter().take(j + 1) {
1519                let h_ij = v_i.inner_product(&w_orth)?;
1520                h_col.push(h_ij.clone());
1521                // w_orth = w_orth - h_ij * v_i = 1.0 * w_orth + (-h_ij) * v_i
1522                let neg_h_ij = AnyScalar::new_real(0.0) - h_ij;
1523                w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_h_ij)?;
1524            }
1525
1526            // KrylovKit's default orthogonalizer is ModifiedGramSchmidt2.
1527            // The second pass is important for long Krylov bases and complex
1528            // non-Hermitian local problems.
1529            for (i, v_i) in v_basis.iter().take(j + 1).enumerate() {
1530                let correction = v_i.inner_product(&w_orth)?;
1531                h_col[i] = h_col[i].clone() + correction.clone();
1532                let neg_correction = AnyScalar::new_real(0.0) - correction;
1533                w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_correction)?;
1534            }
1535
1536            let h_jp1_j_real = w_orth.norm();
1537            let h_jp1_j = AnyScalar::new_real(h_jp1_j_real);
1538            h_col.push(h_jp1_j);
1539
1540            // Apply previous Givens rotations to new column
1541            #[allow(clippy::needless_range_loop)]
1542            for i in 0..j {
1543                let h_i = h_col[i].clone();
1544                let h_ip1 = h_col[i + 1].clone();
1545                let (new_hi, new_hip1) = apply_givens_rotation(&cs[i], &sn[i], &h_i, &h_ip1);
1546                h_col[i] = new_hi;
1547                h_col[i + 1] = new_hip1;
1548            }
1549
1550            // Compute new Givens rotation for h_col[j] and h_col[j+1]
1551            let (c_j, s_j) = compute_givens_rotation(&h_col[j], &h_col[j + 1]);
1552            cs.push(c_j.clone());
1553            sn.push(s_j.clone());
1554
1555            // Apply new rotation to eliminate h_col[j+1]
1556            let (new_hj, _) = apply_givens_rotation(&c_j, &s_j, &h_col[j], &h_col[j + 1]);
1557            h_col[j] = new_hj;
1558            h_col[j + 1] = AnyScalar::new_real(0.0);
1559
1560            // Apply rotation to g
1561            let g_j = g[j].clone();
1562            let g_jp1 = AnyScalar::new_real(0.0);
1563            let (new_gj, new_gjp1) = apply_givens_rotation(&c_j, &s_j, &g_j, &g_jp1);
1564            g[j] = new_gj;
1565            let res_norm = new_gjp1.abs();
1566            g.push(new_gjp1);
1567
1568            h_matrix.push(h_col);
1569
1570            // Check convergence
1571            let residual_value = tolerance.residual_value(res_norm, b_norm);
1572
1573            if options.verbose {
1574                eprintln!("GMRES iter {}: residual = {:.6e}", j + 1, residual_value);
1575            }
1576
1577            if tolerance.is_converged(res_norm, b_norm) {
1578                // Solve upper triangular system and update x
1579                let y = solve_upper_triangular(&h_matrix, &g[..=j])?;
1580                x = update_solution(&x, &v_basis[..=j], &y)?;
1581                if options.check_true_residual {
1582                    let ax_check = apply_a(&x)?;
1583                    let r_check = b.axpby(
1584                        AnyScalar::new_real(1.0),
1585                        &ax_check,
1586                        AnyScalar::new_real(-1.0),
1587                    )?;
1588                    let true_abs_res = r_check.norm();
1589                    let true_residual_value = tolerance.residual_value(true_abs_res, b_norm);
1590
1591                    if options.verbose {
1592                        eprintln!(
1593                            "GMRES true residual check: hessenberg={:.6e}, checked={:.6e}",
1594                            residual_value, true_residual_value
1595                        );
1596                    }
1597
1598                    if tolerance.is_converged(true_abs_res, b_norm) {
1599                        return Ok(GmresResult {
1600                            solution: x,
1601                            iterations: total_iters,
1602                            residual_norm: true_residual_value,
1603                            converged: true,
1604                        });
1605                    }
1606                    solution_already_updated = true;
1607                    break;
1608                } else {
1609                    return Ok(GmresResult {
1610                        solution: x,
1611                        iterations: total_iters,
1612                        residual_norm: residual_value,
1613                        converged: true,
1614                    });
1615                }
1616            }
1617
1618            // Add new basis vector (if not converged and h_jp1_j is not too small)
1619            if h_jp1_j_real > 1e-14 {
1620                let v_jp1 = w_orth.scale(AnyScalar::new_real(1.0 / h_jp1_j_real))?;
1621                v_basis.push(v_jp1);
1622            } else {
1623                // Lucky breakdown - we've found the exact solution in the Krylov subspace
1624                let y = solve_upper_triangular(&h_matrix, &g[..=j])?;
1625                x = update_solution(&x, &v_basis[..=j], &y)?;
1626                let ax_final = apply_a(&x)?;
1627                let r_final = b.axpby(
1628                    AnyScalar::new_real(1.0),
1629                    &ax_final,
1630                    AnyScalar::new_real(-1.0),
1631                )?;
1632                let final_abs_res = r_final.norm();
1633                let final_res = tolerance.residual_value(final_abs_res, b_norm);
1634                return Ok(GmresResult {
1635                    solution: x,
1636                    iterations: total_iters,
1637                    residual_norm: final_res,
1638                    converged: tolerance.is_converged(final_abs_res, b_norm),
1639                });
1640            }
1641        }
1642
1643        // End of restart cycle - update x with current solution
1644        if !solution_already_updated {
1645            let actual_iters = h_matrix.len();
1646            let y = solve_upper_triangular(&h_matrix, &g[..actual_iters])?;
1647            x = update_solution(&x, &v_basis[..actual_iters], &y)?;
1648        }
1649    }
1650
1651    // Compute final residual
1652    let ax_final = apply_a(&x)?;
1653    let r_final = b.axpby(
1654        AnyScalar::new_real(1.0),
1655        &ax_final,
1656        AnyScalar::new_real(-1.0),
1657    )?;
1658    let final_abs_res = r_final.norm();
1659    let final_res = tolerance.residual_value(final_abs_res, b_norm);
1660
1661    Ok(GmresResult {
1662        solution: x,
1663        iterations: total_iters,
1664        residual_norm: final_res,
1665        converged: tolerance.is_converged(final_abs_res, b_norm),
1666    })
1667}
1668
1669/// Solve `A x = b` using GMRES with optional truncation after each iteration.
1670///
1671/// This is an extension of [`gmres`] that allows truncating Krylov basis vectors
1672/// to control bond dimension growth in tensor network representations.
1673///
1674/// # Type Parameters
1675///
1676/// * `T` - A tensor type implementing `TensorVectorSpace`
1677/// * `F` - A function that applies the linear operator: `F(x) = A x`
1678/// * `Tr` - A function that truncates a tensor in-place: `Tr(&mut x)`
1679///
1680/// # Arguments
1681///
1682/// * `apply_a` - Function that applies the linear operator A to a tensor
1683/// * `b` - Right-hand side tensor
1684/// * `x0` - Initial guess
1685/// * `options` - Solver options
1686/// * `truncate` - Function that truncates a tensor to control bond dimension
1687///
1688/// # Note
1689///
1690/// Truncation is applied after each Gram-Schmidt orthogonalization step
1691/// and after the final solution update. This helps control the bond dimension
1692/// growth that would otherwise occur in MPS/MPO representations.
1693///
1694/// # Examples
1695///
1696/// Solve `2x = b` with a no-op truncation function:
1697///
1698/// ```
1699/// use tensor4all_core::{DynIndex, TensorDynLen, TensorVectorSpace, AnyScalar};
1700/// use tensor4all_core::krylov::{gmres_with_truncation, GmresOptions};
1701///
1702/// let i = DynIndex::new_dyn(2);
1703/// let b = TensorDynLen::from_dense(vec![i.clone()], vec![4.0, 6.0]).unwrap();
1704/// let x0 = TensorDynLen::from_dense(vec![i.clone()], vec![0.0, 0.0]).unwrap();
1705///
1706/// // Operator A = 2*I (scales input by 2)
1707/// let apply_a = |x: &TensorDynLen| x.scale(AnyScalar::new_real(2.0));
1708/// // No-op truncation
1709/// let truncate = |_x: &mut TensorDynLen| Ok(());
1710///
1711/// let result = gmres_with_truncation(apply_a, &b, &x0, &GmresOptions::default(), truncate).unwrap();
1712/// assert!(result.converged);
1713/// // Solution should be [2.0, 3.0]
1714/// let expected = TensorDynLen::from_dense(vec![i], vec![2.0, 3.0]).unwrap();
1715/// assert!(result.solution.sub(&expected).unwrap().maxabs() < 1e-8);
1716/// ```
1717pub fn gmres_with_truncation<T, F, Tr>(
1718    apply_a: F,
1719    b: &T,
1720    x0: &T,
1721    options: &GmresOptions,
1722    truncate: Tr,
1723) -> Result<GmresResult<T>>
1724where
1725    T: TensorVectorSpace,
1726    F: Fn(&T) -> Result<T>,
1727    Tr: Fn(&mut T) -> Result<()>,
1728{
1729    // Validate structural consistency of inputs
1730    b.validate()?;
1731    x0.validate()?;
1732
1733    let b_norm = b.norm();
1734    if b_norm < 1e-15 {
1735        return Ok(GmresResult {
1736            solution: x0.clone(),
1737            iterations: 0,
1738            residual_norm: 0.0,
1739            converged: true,
1740        });
1741    }
1742
1743    let mut x = x0.clone();
1744    let mut total_iters = 0;
1745
1746    for _restart in 0..options.max_restarts {
1747        let ax = apply_a(&x)?;
1748        // Validate operator output on first restart
1749        if _restart == 0 {
1750            ax.validate()?;
1751        }
1752        let mut r = b.axpby(AnyScalar::new_real(1.0), &ax, AnyScalar::new_real(-1.0))?;
1753        truncate(&mut r)?;
1754        let r_norm = r.norm();
1755        let rel_res = r_norm / b_norm;
1756
1757        if options.verbose {
1758            eprintln!(
1759                "GMRES restart {}: initial residual = {:.6e}",
1760                _restart, rel_res
1761            );
1762        }
1763
1764        if rel_res < options.rtol {
1765            return Ok(GmresResult {
1766                solution: x,
1767                iterations: total_iters,
1768                residual_norm: rel_res,
1769                converged: true,
1770            });
1771        }
1772
1773        let mut v_basis: Vec<T> = Vec::with_capacity(options.max_iter + 1);
1774        let mut h_matrix: Vec<Vec<AnyScalar>> = Vec::with_capacity(options.max_iter);
1775
1776        let mut v0 = r.scale(AnyScalar::new_real(1.0 / r_norm))?;
1777        truncate(&mut v0)?;
1778        // After truncation, v0 might not be unit norm and might point in a different direction.
1779        // We need to:
1780        // 1. Renormalize v0 to unit norm for numerical stability
1781        // 2. Recompute g[0] = <r, v0> to maintain the correct relationship
1782        let v0_norm = v0.norm();
1783        let effective_g0 = if v0_norm > 1e-15 {
1784            v0 = v0.scale(AnyScalar::new_real(1.0 / v0_norm))?;
1785            // g[0] should be the component of r in the direction of v0
1786            // Since r was truncated and v0 = truncate(r/||r||)/||truncate(r/||r||)||,
1787            // g[0] = <r, v0> ≈ ||r|| * ||truncate(r/||r||)|| = r_norm * v0_norm
1788            r_norm * v0_norm
1789        } else {
1790            r_norm
1791        };
1792        v_basis.push(v0);
1793
1794        let mut cs: Vec<AnyScalar> = Vec::with_capacity(options.max_iter);
1795        let mut sn: Vec<AnyScalar> = Vec::with_capacity(options.max_iter);
1796        let mut g: Vec<AnyScalar> = vec![AnyScalar::new_real(effective_g0)];
1797        let mut solution_already_updated = false;
1798
1799        for j in 0..options.max_iter {
1800            total_iters += 1;
1801
1802            let w = apply_a(&v_basis[j])?;
1803
1804            let mut h_col: Vec<AnyScalar> = Vec::with_capacity(j + 2);
1805            let mut w_orth = w;
1806
1807            for v_i in v_basis.iter().take(j + 1) {
1808                let h_ij = v_i.inner_product(&w_orth)?;
1809                h_col.push(h_ij.clone());
1810                let neg_h_ij = AnyScalar::new_real(0.0) - h_ij;
1811                w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_h_ij)?;
1812            }
1813
1814            // Iterative reorthogonalization with truncation
1815            // Truncation can change the direction of w_orth, breaking orthogonality.
1816            // We iterate until all corrections are below a threshold to ensure
1817            // the Krylov basis remains orthogonal despite truncation.
1818            const REORTH_THRESHOLD: f64 = 1e-12;
1819            const MAX_REORTH_ITERS: usize = 10;
1820
1821            let mut reorth_iter_count = 0;
1822            for reorth_iter in 0..MAX_REORTH_ITERS {
1823                reorth_iter_count = reorth_iter + 1;
1824                let norm_before_truncate = w_orth.norm();
1825                truncate(&mut w_orth)?;
1826                let norm_after_truncate = w_orth.norm();
1827
1828                let mut max_correction = 0.0;
1829                for (i, v_i) in v_basis.iter().enumerate() {
1830                    let correction = v_i.inner_product(&w_orth)?;
1831                    let correction_abs = correction.abs();
1832                    if correction_abs > max_correction {
1833                        max_correction = correction_abs;
1834                    }
1835                    if correction_abs > REORTH_THRESHOLD {
1836                        let neg_correction = AnyScalar::new_real(0.0) - correction.clone();
1837                        w_orth = w_orth.axpby(AnyScalar::new_real(1.0), v_i, neg_correction)?;
1838                        // Update Hessenberg matrix entry to include correction
1839                        h_col[i] = h_col[i].clone() + correction;
1840                    }
1841                }
1842
1843                if options.verbose {
1844                    eprintln!(
1845                        "  reorth iter {}: norm {:.6e} -> {:.6e}, max_correction = {:.6e}",
1846                        reorth_iter, norm_before_truncate, norm_after_truncate, max_correction
1847                    );
1848                }
1849
1850                // If all corrections are small enough, we're done
1851                if max_correction < REORTH_THRESHOLD {
1852                    break;
1853                }
1854            }
1855
1856            if options.verbose && reorth_iter_count > 1 {
1857                eprintln!("  (needed {} reorth iterations)", reorth_iter_count);
1858            }
1859
1860            let h_jp1_j_real = w_orth.norm();
1861            let h_jp1_j = AnyScalar::new_real(h_jp1_j_real);
1862            h_col.push(h_jp1_j);
1863
1864            #[allow(clippy::needless_range_loop)]
1865            for i in 0..j {
1866                let h_i = h_col[i].clone();
1867                let h_ip1 = h_col[i + 1].clone();
1868                let (new_hi, new_hip1) = apply_givens_rotation(&cs[i], &sn[i], &h_i, &h_ip1);
1869                h_col[i] = new_hi;
1870                h_col[i + 1] = new_hip1;
1871            }
1872
1873            let (c_j, s_j) = compute_givens_rotation(&h_col[j], &h_col[j + 1]);
1874            cs.push(c_j.clone());
1875            sn.push(s_j.clone());
1876
1877            let (new_hj, _) = apply_givens_rotation(&c_j, &s_j, &h_col[j], &h_col[j + 1]);
1878            h_col[j] = new_hj;
1879            h_col[j + 1] = AnyScalar::new_real(0.0);
1880
1881            let g_j = g[j].clone();
1882            let g_jp1 = AnyScalar::new_real(0.0);
1883            let (new_gj, new_gjp1) = apply_givens_rotation(&c_j, &s_j, &g_j, &g_jp1);
1884            g[j] = new_gj;
1885            let res_norm = new_gjp1.abs();
1886            g.push(new_gjp1);
1887
1888            h_matrix.push(h_col);
1889
1890            let rel_res = res_norm / b_norm;
1891
1892            if options.verbose {
1893                eprintln!("GMRES iter {}: residual = {:.6e}", j + 1, rel_res);
1894            }
1895
1896            if rel_res < options.rtol {
1897                let y = solve_upper_triangular(&h_matrix, &g[..=j])?;
1898                x = update_solution_truncated(&x, &v_basis[..=j], &y, &truncate)?;
1899
1900                if options.check_true_residual {
1901                    // Verify with true residual to prevent false convergence
1902                    let ax_check = apply_a(&x)?;
1903                    let mut r_check = b.axpby(
1904                        AnyScalar::new_real(1.0),
1905                        &ax_check,
1906                        AnyScalar::new_real(-1.0),
1907                    )?;
1908                    truncate(&mut r_check)?;
1909                    let true_rel_res = r_check.norm() / b_norm;
1910
1911                    if options.verbose {
1912                        eprintln!(
1913                            "GMRES true residual check: hessenberg={:.6e}, checked={:.6e}",
1914                            rel_res, true_rel_res
1915                        );
1916                    }
1917
1918                    if true_rel_res < options.rtol {
1919                        return Ok(GmresResult {
1920                            solution: x,
1921                            iterations: total_iters,
1922                            residual_norm: true_rel_res,
1923                            converged: true,
1924                        });
1925                    }
1926                    // False convergence detected: x is already updated above,
1927                    // so skip the end-of-cycle update and go to next restart
1928                    solution_already_updated = true;
1929                    break;
1930                } else {
1931                    return Ok(GmresResult {
1932                        solution: x,
1933                        iterations: total_iters,
1934                        residual_norm: rel_res,
1935                        converged: true,
1936                    });
1937                }
1938            }
1939
1940            if h_jp1_j_real > 1e-14 {
1941                // Create v_{j+1} = w_orth / ||w_orth||
1942                // w_orth has already been truncated twice (after orthogonalization and after reorthogonalization)
1943                // so we don't need to truncate again. Scale doesn't increase bond dimensions.
1944                let v_jp1 = w_orth.scale(AnyScalar::new_real(1.0 / h_jp1_j_real))?;
1945                // v_jp1 should have norm ~1.0 by construction
1946                // The Arnoldi relation h_{j+1,j} * v_{j+1} = w_orth is maintained exactly
1947                v_basis.push(v_jp1);
1948            } else {
1949                let y = solve_upper_triangular(&h_matrix, &g[..=j])?;
1950                x = update_solution_truncated(&x, &v_basis[..=j], &y, &truncate)?;
1951                let ax_final = apply_a(&x)?;
1952                let r_final = b.axpby(
1953                    AnyScalar::new_real(1.0),
1954                    &ax_final,
1955                    AnyScalar::new_real(-1.0),
1956                )?;
1957                let final_res = r_final.norm() / b_norm;
1958                return Ok(GmresResult {
1959                    solution: x,
1960                    iterations: total_iters,
1961                    residual_norm: final_res,
1962                    converged: final_res < options.rtol,
1963                });
1964            }
1965        }
1966
1967        if !solution_already_updated {
1968            let actual_iters = v_basis.len().min(options.max_iter);
1969            let y = solve_upper_triangular(&h_matrix, &g[..actual_iters])?;
1970            x = update_solution_truncated(&x, &v_basis[..actual_iters], &y, &truncate)?;
1971        }
1972    }
1973
1974    let ax_final = apply_a(&x)?;
1975    let r_final = b.axpby(
1976        AnyScalar::new_real(1.0),
1977        &ax_final,
1978        AnyScalar::new_real(-1.0),
1979    )?;
1980    let final_res = r_final.norm() / b_norm;
1981
1982    Ok(GmresResult {
1983        solution: x,
1984        iterations: total_iters,
1985        residual_norm: final_res,
1986        converged: final_res < options.rtol,
1987    })
1988}
1989
1990/// Options for restarted GMRES with truncation.
1991///
1992/// This is used by [`restart_gmres_with_truncation`] which wraps the standard GMRES
1993/// with an outer loop that recomputes the true residual at each restart.
1994///
1995/// # Examples
1996///
1997/// ```
1998/// use tensor4all_core::krylov::RestartGmresOptions;
1999///
2000/// let opts = RestartGmresOptions::new()
2001///     .with_max_outer_iters(10)
2002///     .with_rtol(1e-6)
2003///     .with_inner_max_iter(20)
2004///     .with_inner_max_restarts(2)
2005///     .with_min_reduction(0.99)
2006///     .with_inner_rtol(0.01)
2007///     .with_verbose(false);
2008///
2009/// assert_eq!(opts.max_outer_iters, 10);
2010/// assert_eq!(opts.rtol, 1e-6);
2011/// assert_eq!(opts.inner_max_iter, 20);
2012/// assert_eq!(opts.inner_max_restarts, 2);
2013/// assert_eq!(opts.min_reduction, Some(0.99));
2014/// assert_eq!(opts.inner_rtol, Some(0.01));
2015/// ```
2016#[derive(Debug, Clone)]
2017pub struct RestartGmresOptions {
2018    /// Maximum number of outer restart iterations.
2019    /// Default: 20
2020    pub max_outer_iters: usize,
2021
2022    /// Convergence tolerance for relative residual norm (based on true residual).
2023    /// The solver stops when `||b - A*x|| / ||b|| < rtol`.
2024    /// Default: 1e-10
2025    pub rtol: f64,
2026
2027    /// Maximum iterations per inner GMRES cycle.
2028    /// Default: 10
2029    pub inner_max_iter: usize,
2030
2031    /// Number of restarts within each inner GMRES (usually 0).
2032    /// Default: 0
2033    pub inner_max_restarts: usize,
2034
2035    /// Stagnation detection threshold.
2036    /// If the residual reduction ratio exceeds this value (i.e., residual doesn't decrease enough),
2037    /// the solver considers it stagnated.
2038    /// For example, 0.99 means stagnation is detected when residual decreases by less than 1%.
2039    /// Default: None (no stagnation detection)
2040    pub min_reduction: Option<f64>,
2041
2042    /// Inner GMRES relative tolerance.
2043    /// If None, uses 0.1 (solve inner problem loosely).
2044    /// Default: None
2045    pub inner_rtol: Option<f64>,
2046
2047    /// Whether to print convergence information.
2048    /// Default: false
2049    pub verbose: bool,
2050}
2051
2052impl Default for RestartGmresOptions {
2053    fn default() -> Self {
2054        Self {
2055            max_outer_iters: 20,
2056            rtol: 1e-10,
2057            inner_max_iter: 10,
2058            inner_max_restarts: 0,
2059            min_reduction: None,
2060            inner_rtol: None,
2061            verbose: false,
2062        }
2063    }
2064}
2065
2066impl RestartGmresOptions {
2067    /// Create new options with default values.
2068    pub fn new() -> Self {
2069        Self::default()
2070    }
2071
2072    /// Set maximum number of outer iterations.
2073    pub fn with_max_outer_iters(mut self, max_outer_iters: usize) -> Self {
2074        self.max_outer_iters = max_outer_iters;
2075        self
2076    }
2077
2078    /// Set convergence tolerance.
2079    pub fn with_rtol(mut self, rtol: f64) -> Self {
2080        self.rtol = rtol;
2081        self
2082    }
2083
2084    /// Set maximum iterations per inner GMRES cycle.
2085    pub fn with_inner_max_iter(mut self, inner_max_iter: usize) -> Self {
2086        self.inner_max_iter = inner_max_iter;
2087        self
2088    }
2089
2090    /// Set number of restarts within each inner GMRES.
2091    pub fn with_inner_max_restarts(mut self, inner_max_restarts: usize) -> Self {
2092        self.inner_max_restarts = inner_max_restarts;
2093        self
2094    }
2095
2096    /// Set stagnation detection threshold.
2097    pub fn with_min_reduction(mut self, min_reduction: f64) -> Self {
2098        self.min_reduction = Some(min_reduction);
2099        self
2100    }
2101
2102    /// Set inner GMRES relative tolerance.
2103    pub fn with_inner_rtol(mut self, inner_rtol: f64) -> Self {
2104        self.inner_rtol = Some(inner_rtol);
2105        self
2106    }
2107
2108    /// Enable verbose output.
2109    pub fn with_verbose(mut self, verbose: bool) -> Self {
2110        self.verbose = verbose;
2111        self
2112    }
2113}
2114
2115/// Result of restarted GMRES solver.
2116///
2117/// # Examples
2118///
2119/// ```
2120/// use tensor4all_core::{DynIndex, TensorDynLen, AnyScalar};
2121/// use tensor4all_core::krylov::{restart_gmres_with_truncation, RestartGmresOptions};
2122///
2123/// let i = DynIndex::new_dyn(2);
2124/// let b = TensorDynLen::from_dense(vec![i.clone()], vec![3.0, 5.0]).unwrap();
2125///
2126/// let apply_a = |x: &TensorDynLen| x.scale(AnyScalar::new_real(3.0));
2127/// let truncate = |_x: &mut TensorDynLen| Ok(());
2128///
2129/// let result = restart_gmres_with_truncation(
2130///     apply_a, &b, None, &RestartGmresOptions::default(), truncate,
2131/// ).unwrap();
2132///
2133/// assert!(result.converged);
2134/// assert!(result.residual_norm < 1e-10);
2135/// assert!(result.outer_iterations <= 20);
2136/// ```
2137#[derive(Debug, Clone)]
2138pub struct RestartGmresResult<T> {
2139    /// The solution vector.
2140    pub solution: T,
2141
2142    /// Total number of inner GMRES iterations performed.
2143    pub iterations: usize,
2144
2145    /// Number of outer restart iterations performed.
2146    pub outer_iterations: usize,
2147
2148    /// Final relative residual norm (true residual).
2149    pub residual_norm: f64,
2150
2151    /// Whether the solver converged.
2152    pub converged: bool,
2153}
2154
2155/// Solve `A x = b` using restarted GMRES with truncation.
2156///
2157/// This wraps [`gmres_with_truncation`] with an outer loop that recomputes the true residual
2158/// at each restart. This is particularly useful for MPS/MPO computations where truncation
2159/// can cause the inner GMRES residual to be inaccurate.
2160///
2161/// # Algorithm
2162///
2163/// ```text
2164/// for outer_iter in 0..max_outer_iters:
2165///     r = b - A*x0          // Compute true residual
2166///     r = truncate(r)
2167///     if ||r|| / ||b|| < rtol:
2168///         return x0         // Converged
2169///     x' = gmres_with_truncation(A, r, 0, inner_options, truncate)
2170///     x0 = truncate(x0 + x')
2171/// ```
2172///
2173/// # Type Parameters
2174///
2175/// * `T` - A tensor type implementing `TensorVectorSpace`
2176/// * `F` - A function that applies the linear operator: `F(x) = A x`
2177/// * `Tr` - A function that truncates a tensor in-place: `Tr(&mut x)`
2178///
2179/// # Arguments
2180///
2181/// * `apply_a` - Function that applies the linear operator A to a tensor
2182/// * `b` - Right-hand side tensor
2183/// * `x0` - Initial guess (if None, starts from zero)
2184/// * `options` - Solver options
2185/// * `truncate` - Function that truncates a tensor to control bond dimension
2186///
2187/// # Returns
2188///
2189/// A `RestartGmresResult` containing the solution and convergence information.
2190///
2191/// # Examples
2192///
2193/// Solve `5x = b` with no truncation:
2194///
2195/// ```
2196/// use tensor4all_core::{DynIndex, TensorDynLen, TensorVectorSpace, AnyScalar};
2197/// use tensor4all_core::krylov::{restart_gmres_with_truncation, RestartGmresOptions};
2198///
2199/// let i = DynIndex::new_dyn(3);
2200/// let b = TensorDynLen::from_dense(vec![i.clone()], vec![5.0, 10.0, 15.0]).unwrap();
2201///
2202/// let apply_a = |x: &TensorDynLen| x.scale(AnyScalar::new_real(5.0));
2203/// let truncate = |_x: &mut TensorDynLen| Ok(());
2204///
2205/// let result = restart_gmres_with_truncation(
2206///     apply_a, &b, None, &RestartGmresOptions::default(), truncate,
2207/// ).unwrap();
2208///
2209/// assert!(result.converged);
2210/// let expected = TensorDynLen::from_dense(vec![i], vec![1.0, 2.0, 3.0]).unwrap();
2211/// assert!(result.solution.sub(&expected).unwrap().maxabs() < 1e-8);
2212/// ```
2213pub fn restart_gmres_with_truncation<T, F, Tr>(
2214    apply_a: F,
2215    b: &T,
2216    x0: Option<&T>,
2217    options: &RestartGmresOptions,
2218    truncate: Tr,
2219) -> Result<RestartGmresResult<T>>
2220where
2221    T: TensorVectorSpace,
2222    F: Fn(&T) -> Result<T>,
2223    Tr: Fn(&mut T) -> Result<()>,
2224{
2225    // Validate structural consistency of inputs
2226    b.validate()?;
2227    if let Some(x) = x0 {
2228        x.validate()?;
2229    }
2230
2231    let b_norm = b.norm();
2232    if b_norm < 1e-15 {
2233        // b is effectively zero, return x0 or zero
2234        let solution = match x0 {
2235            Some(x) => x.clone(),
2236            None => b.scale(AnyScalar::new_real(0.0))?,
2237        };
2238        return Ok(RestartGmresResult {
2239            solution,
2240            iterations: 0,
2241            outer_iterations: 0,
2242            residual_norm: 0.0,
2243            converged: true,
2244        });
2245    }
2246
2247    // Initialize x: use x0 if provided, otherwise start from zero.
2248    // Track whether x is zero to avoid unnecessary bond dimension doubling
2249    // when adding the first correction via axpby.
2250    let mut x_is_zero = x0.is_none();
2251    let mut x = match x0 {
2252        Some(x) => x.clone(),
2253        None => b.scale(AnyScalar::new_real(0.0))?,
2254    };
2255
2256    let mut total_inner_iters = 0;
2257    let mut prev_residual_norm = f64::INFINITY;
2258
2259    // Inner GMRES options
2260    let inner_options = GmresOptions {
2261        max_iter: options.inner_max_iter,
2262        rtol: options.inner_rtol.unwrap_or(0.1), // Solve loosely by default
2263        max_restarts: options.inner_max_restarts + 1, // +1 because max_restarts=0 means 1 cycle
2264        verbose: options.verbose,
2265        check_true_residual: true, // Always check in restart context to avoid false convergence
2266    };
2267
2268    for outer_iter in 0..options.max_outer_iters {
2269        // Compute true residual: r = b - A*x
2270        let ax = apply_a(&x)?;
2271        // Validate operator output on first outer iteration
2272        if outer_iter == 0 {
2273            ax.validate()?;
2274        }
2275        let mut r = b.axpby(AnyScalar::new_real(1.0), &ax, AnyScalar::new_real(-1.0))?;
2276        truncate(&mut r)?;
2277
2278        let r_norm = r.norm();
2279        let rel_res = r_norm / b_norm;
2280
2281        if options.verbose {
2282            eprintln!(
2283                "Restart GMRES outer iter {}: true residual = {:.6e}",
2284                outer_iter, rel_res
2285            );
2286        }
2287
2288        // Check convergence
2289        if rel_res < options.rtol {
2290            return Ok(RestartGmresResult {
2291                solution: x,
2292                iterations: total_inner_iters,
2293                outer_iterations: outer_iter,
2294                residual_norm: rel_res,
2295                converged: true,
2296            });
2297        }
2298
2299        // Check stagnation
2300        if let Some(min_reduction) = options.min_reduction {
2301            if outer_iter > 0 && rel_res > prev_residual_norm * min_reduction {
2302                if options.verbose {
2303                    eprintln!(
2304                        "Restart GMRES stagnated: residual ratio = {:.6e} > {:.6e}",
2305                        rel_res / prev_residual_norm,
2306                        min_reduction
2307                    );
2308                }
2309                return Ok(RestartGmresResult {
2310                    solution: x,
2311                    iterations: total_inner_iters,
2312                    outer_iterations: outer_iter,
2313                    residual_norm: rel_res,
2314                    converged: false,
2315                });
2316            }
2317        }
2318        prev_residual_norm = rel_res;
2319
2320        // Solve A*x' = r using inner GMRES with zero initial guess
2321        // The zero initial guess is created by scaling r by 0
2322        let zero = r.scale(AnyScalar::new_real(0.0))?;
2323        let inner_result = gmres_with_truncation(&apply_a, &r, &zero, &inner_options, &truncate)?;
2324
2325        total_inner_iters += inner_result.iterations;
2326
2327        if options.verbose {
2328            eprintln!(
2329                "  Inner GMRES: {} iterations, residual = {:.6e}, converged = {}",
2330                inner_result.iterations, inner_result.residual_norm, inner_result.converged
2331            );
2332        }
2333
2334        // Update solution: x = x + x'
2335        // When x is zero (first iteration with no initial guess), use x' directly
2336        // to avoid bond dimension doubling from axpby with a zero tensor.
2337        if x_is_zero {
2338            x = inner_result.solution;
2339            x_is_zero = false;
2340        } else {
2341            x = x.axpby(
2342                AnyScalar::new_real(1.0),
2343                &inner_result.solution,
2344                AnyScalar::new_real(1.0),
2345            )?;
2346        }
2347        truncate(&mut x)?;
2348    }
2349
2350    // Did not converge within max_outer_iters
2351    // Compute final residual
2352    let ax = apply_a(&x)?;
2353    let mut r = b.axpby(AnyScalar::new_real(1.0), &ax, AnyScalar::new_real(-1.0))?;
2354    truncate(&mut r)?;
2355    let final_rel_res = r.norm() / b_norm;
2356
2357    Ok(RestartGmresResult {
2358        solution: x,
2359        iterations: total_inner_iters,
2360        outer_iterations: options.max_outer_iters,
2361        residual_norm: final_rel_res,
2362        converged: false,
2363    })
2364}
2365
2366fn validate_hermitian_lanczos_options(options: &HermitianLanczosOptions) -> Result<()> {
2367    anyhow::ensure!(
2368        options.max_iter > 0,
2369        "hermitian_lanczos_lowest_eigenpair: max_iter must be greater than zero"
2370    );
2371    anyhow::ensure!(
2372        options.rtol.is_finite() && options.rtol >= 0.0,
2373        "hermitian_lanczos_lowest_eigenpair: rtol must be finite and non-negative"
2374    );
2375    anyhow::ensure!(
2376        options.atol.is_finite() && options.atol >= 0.0,
2377        "hermitian_lanczos_lowest_eigenpair: atol must be finite and non-negative"
2378    );
2379    anyhow::ensure!(
2380        options.breakdown_tol.is_finite() && options.breakdown_tol >= 0.0,
2381        "hermitian_lanczos_lowest_eigenpair: breakdown_tol must be finite and non-negative"
2382    );
2383    anyhow::ensure!(
2384        options.hermitian_tol.is_finite() && options.hermitian_tol >= 0.0,
2385        "hermitian_lanczos_lowest_eigenpair: hermitian_tol must be finite and non-negative"
2386    );
2387    Ok(())
2388}
2389
2390fn validate_hermitian_krylov_expm_options(options: &HermitianKrylovExpmOptions) -> Result<()> {
2391    anyhow::ensure!(
2392        options.max_iter > 0,
2393        "hermitian_krylov_expm_multiply: max_iter must be greater than zero"
2394    );
2395    anyhow::ensure!(
2396        options.max_time_splits > 0,
2397        "hermitian_krylov_expm_multiply: max_time_splits must be greater than zero"
2398    );
2399    anyhow::ensure!(
2400        options.tol.is_finite() && options.tol >= 0.0,
2401        "hermitian_krylov_expm_multiply: tol must be finite and non-negative"
2402    );
2403    anyhow::ensure!(
2404        options.breakdown_tol.is_finite() && options.breakdown_tol >= 0.0,
2405        "hermitian_krylov_expm_multiply: breakdown_tol must be finite and non-negative"
2406    );
2407    anyhow::ensure!(
2408        options.hermitian_tol.is_finite() && options.hermitian_tol >= 0.0,
2409        "hermitian_krylov_expm_multiply: hermitian_tol must be finite and non-negative"
2410    );
2411    Ok(())
2412}
2413
2414fn projected_matrix_from_columns(h_cols: &[Vec<Complex64>], dim: usize) -> Matrix<Complex64> {
2415    let mut data = vec![Complex64::new(0.0, 0.0); dim * dim];
2416    for col in 0..dim {
2417        for row in 0..dim {
2418            if let Some(value) = h_cols.get(col).and_then(|h_col| h_col.get(row)) {
2419                data[row + dim * col] = *value;
2420            }
2421        }
2422    }
2423    Matrix::from_col_major_vec(dim, dim, data)
2424}
2425
2426fn any_scalar_to_complex(value: &AnyScalar) -> Complex64 {
2427    value
2428        .as_c64()
2429        .unwrap_or_else(|| Complex64::new(value.real(), 0.0))
2430}
2431
2432fn any_scalar_from_complex(value: Complex64, tolerance: f64) -> Result<AnyScalar> {
2433    if value.im.abs() <= tolerance {
2434        Ok(AnyScalar::new_real(value.re))
2435    } else {
2436        Ok(AnyScalar::new_complex(value.re, value.im))
2437    }
2438}
2439
2440fn hermitian_lanczos_threshold(eigenvalue: f64, options: &HermitianLanczosOptions) -> f64 {
2441    options.atol.max(options.rtol * eigenvalue.abs().max(1.0))
2442}
2443
2444fn hermitian_true_residual_norm<T, F>(apply_a: &F, eigenvector: &T, eigenvalue: f64) -> Result<f64>
2445where
2446    T: TensorVectorSpace,
2447    F: Fn(&T) -> Result<T>,
2448{
2449    let av = apply_a(eigenvector)?;
2450    let lambda_v = eigenvector.scale(AnyScalar::new_real(eigenvalue))?;
2451    let residual = av.axpby(
2452        AnyScalar::new_real(1.0),
2453        &lambda_v,
2454        AnyScalar::new_real(-1.0),
2455    )?;
2456    Ok(residual.norm())
2457}
2458
2459fn finalize_hermitian_lanczos_result<T, F>(
2460    apply_a: &F,
2461    basis: &[T],
2462    ritz: &HermitianRitzState,
2463    options: &HermitianLanczosOptions,
2464) -> Result<HermitianLanczosResult<T>>
2465where
2466    T: TensorVectorSpace,
2467    F: Fn(&T) -> Result<T>,
2468{
2469    let eigenvector = combine_basis_with_coefficients(basis, &ritz.coefficients, options)?;
2470    let residual_norm = hermitian_true_residual_norm(apply_a, &eigenvector, ritz.eigenvalue)?;
2471    let threshold = hermitian_lanczos_threshold(ritz.eigenvalue, options);
2472    let converged = residual_norm <= threshold;
2473
2474    if options.verbose {
2475        eprintln!(
2476            "Hermitian Lanczos final check iter {}: residual_estimate={:.6e} true_residual={:.6e} threshold={:.6e}",
2477            ritz.iterations, ritz.residual_estimate, residual_norm, threshold
2478        );
2479    }
2480
2481    Ok(HermitianLanczosResult {
2482        eigenvalue: ritz.eigenvalue,
2483        eigenvector,
2484        iterations: ritz.iterations,
2485        residual_norm,
2486        converged,
2487    })
2488}
2489
2490fn combine_basis_with_coefficients<T>(
2491    basis: &[T],
2492    coefficients: &[Complex64],
2493    options: &HermitianLanczosOptions,
2494) -> Result<T>
2495where
2496    T: TensorVectorSpace,
2497{
2498    anyhow::ensure!(
2499        !basis.is_empty(),
2500        "hermitian_lanczos_lowest_eigenpair: empty Krylov basis"
2501    );
2502    anyhow::ensure!(
2503        basis.len() == coefficients.len(),
2504        "hermitian_lanczos_lowest_eigenpair: coefficient length {} does not match basis length {}",
2505        coefficients.len(),
2506        basis.len()
2507    );
2508
2509    let mut result = basis[0].scale(any_scalar_from_complex(
2510        coefficients[0],
2511        options.hermitian_tol,
2512    )?)?;
2513    for (basis_vector, coefficient) in basis.iter().zip(coefficients.iter()).skip(1) {
2514        result = result.axpby(
2515            AnyScalar::new_real(1.0),
2516            basis_vector,
2517            any_scalar_from_complex(*coefficient, options.hermitian_tol)?,
2518        )?;
2519    }
2520    Ok(result)
2521}
2522
2523fn combine_basis_with_complex_coefficients<T>(
2524    basis: &[T],
2525    coefficients: &[Complex64],
2526    hermitian_tol: f64,
2527    context: &'static str,
2528) -> Result<T>
2529where
2530    T: TensorVectorSpace,
2531{
2532    anyhow::ensure!(!basis.is_empty(), "{context}: empty Krylov basis");
2533    anyhow::ensure!(
2534        basis.len() == coefficients.len(),
2535        "{context}: coefficient length {} does not match basis length {}",
2536        coefficients.len(),
2537        basis.len()
2538    );
2539
2540    let mut result = basis[0].scale(any_scalar_from_complex(coefficients[0], hermitian_tol)?)?;
2541    for (basis_vector, coefficient) in basis.iter().zip(coefficients.iter()).skip(1) {
2542        result = result.axpby(
2543            AnyScalar::new_real(1.0),
2544            basis_vector,
2545            any_scalar_from_complex(*coefficient, hermitian_tol)?,
2546        )?;
2547    }
2548    Ok(result)
2549}
2550
2551/// Compute Givens rotation coefficients to eliminate b in (a, b).
2552///
2553/// This function keeps computation in `AnyScalar` space to preserve AD metadata
2554/// as much as possible.
2555fn compute_givens_rotation(a: &AnyScalar, b: &AnyScalar) -> (AnyScalar, AnyScalar) {
2556    let a_abs = a.abs();
2557    let b_abs = b.abs();
2558    let r = (a_abs * a_abs + b_abs * b_abs).sqrt();
2559    if r < 1e-15 {
2560        (AnyScalar::new_real(1.0), AnyScalar::new_real(0.0))
2561    } else if a_abs < 1e-15 {
2562        (
2563            AnyScalar::new_real(0.0),
2564            b.clone().conj() / AnyScalar::new_real(r),
2565        )
2566    } else {
2567        let phase = a.clone() / AnyScalar::new_real(a_abs);
2568        (
2569            AnyScalar::new_real(a_abs / r),
2570            phase * b.clone().conj() / AnyScalar::new_real(r),
2571        )
2572    }
2573}
2574
2575/// Apply Givens rotation: (c, s) @ (x, y) -> (c*x + s*y, -conj(s)*x + c*y) for complex
2576/// or (c*x + s*y, -s*x + c*y) for real.
2577///
2578/// This function keeps computation in `AnyScalar` space to preserve AD metadata
2579/// as much as possible.
2580fn apply_givens_rotation(
2581    c: &AnyScalar,
2582    s: &AnyScalar,
2583    x: &AnyScalar,
2584    y: &AnyScalar,
2585) -> (AnyScalar, AnyScalar) {
2586    let new_x = c.clone() * x.clone() + s.clone() * y.clone();
2587    let new_y = -(s.clone().conj() * x.clone()) + c.clone() * y.clone();
2588    (new_x, new_y)
2589}
2590
2591/// Solve upper triangular system R y = g using back substitution.
2592fn solve_upper_triangular(h: &[Vec<AnyScalar>], g: &[AnyScalar]) -> Result<Vec<AnyScalar>> {
2593    let n = g.len();
2594    if n == 0 {
2595        return Ok(vec![]);
2596    }
2597
2598    let mut y = vec![AnyScalar::new_real(0.0); n];
2599
2600    for i in (0..n).rev() {
2601        let mut sum = g[i].clone();
2602
2603        for j in (i + 1)..n {
2604            // sum = sum - h[j][i] * y[j]
2605            let prod = h[j][i].clone() * y[j].clone();
2606            sum = sum - prod;
2607        }
2608
2609        let h_ii = &h[i][i];
2610        if h_ii.abs() < 1e-15 {
2611            return Err(anyhow::anyhow!(
2612                "Near-singular upper triangular matrix in GMRES"
2613            ));
2614        }
2615
2616        y[i] = sum / h_ii.clone();
2617    }
2618
2619    Ok(y)
2620}
2621
2622/// Update solution: x_new = x + sum_i y_i * v_i
2623fn update_solution<T: TensorVectorSpace>(x: &T, v_basis: &[T], y: &[AnyScalar]) -> Result<T> {
2624    let mut result = x.clone();
2625
2626    for (vi, yi) in v_basis.iter().zip(y.iter()) {
2627        let scaled_vi = vi.scale(yi.clone())?;
2628        // result = result + scaled_vi = 1.0 * result + 1.0 * scaled_vi
2629        result = result.axpby(
2630            AnyScalar::new_real(1.0),
2631            &scaled_vi,
2632            AnyScalar::new_real(1.0),
2633        )?;
2634    }
2635
2636    Ok(result)
2637}
2638
2639/// Update solution with truncation: x_new = truncate(x + sum_i y_i * v_i)
2640fn update_solution_truncated<T, Tr>(
2641    x: &T,
2642    v_basis: &[T],
2643    y: &[AnyScalar],
2644    truncate: &Tr,
2645) -> Result<T>
2646where
2647    T: TensorVectorSpace,
2648    Tr: Fn(&mut T) -> Result<()>,
2649{
2650    let mut result = x.clone();
2651    // Detect if x is effectively zero.
2652    // When x is created via scale(0.0), it preserves the original bond structure
2653    // (e.g., bond dim 4), causing axpby to double bond dimensions unnecessarily.
2654    // By detecting zero, we can use scaled_vi directly, avoiding the doubling.
2655    let mut result_is_zero = x.norm() == 0.0;
2656
2657    for (vi, yi) in v_basis.iter().zip(y.iter()) {
2658        let scaled_vi = vi.scale(yi.clone())?;
2659        if result_is_zero {
2660            result = scaled_vi;
2661            result_is_zero = false;
2662        } else {
2663            result = result.axpby(
2664                AnyScalar::new_real(1.0),
2665                &scaled_vi,
2666                AnyScalar::new_real(1.0),
2667            )?;
2668        }
2669        // Truncate after each addition to control bond dimension growth
2670        truncate(&mut result)?;
2671    }
2672
2673    Ok(result)
2674}
2675
2676#[cfg(test)]
2677mod tests;