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