arpack 0.0.3

//! Arnoldi-iteration eigenvalue drivers (`*naupd_c` / `*neupd_c`).
//!
//! ARPACK-NG ships two Arnoldi families: real-non-symmetric
//! (`{s,d}{na,ne}upd_c`) and complex (`{c,z}{na,ne}upd_c`). This
//! module wraps the complex family for both `Complex<f64>` and
//! `Complex<f32>` scalars; the real-non-symmetric family is not
//! wrapped yet.
//!
//! Hermitian operators have real eigenvalues but are still driven
//! through the complex Arnoldi routine; the returned eigenvalue
//! comes back complex and callers verify / discard the imaginary
//! part themselves.
//!
//! The crate exposes two layers per scalar type:
//!
//! - [`eigenpairs_c64`] / [`eigenpairs_c32`] — general entry point
//!   accepting `nev >= 1` and a [`Which`] selector.
//! - [`smallest_eigenpair_c64`] / [`smallest_eigenpair_c32`] —
//!   convenience wrappers fixed to `nev = 1` and
//!   [`Which::SmallestRealPart`].
//!
//! Thread-safety: every entry point acquires a process-wide mutex
//! so the entire `*aupd_c` + `*eupd_c` sequence runs atomically
//! against ARPACK's Fortran SAVE state.

use std::os::raw::c_int;

use arpack_sys::{__BindgenComplex, cnaupd_c, cneupd_c, znaupd_c, zneupd_c};
use num_complex::{Complex32, Complex64};

use crate::error::{Error, aupd_error, eupd_error};
use crate::lock::lock;
use crate::solution::{
    EigSolution, MultiEigSolution, c_int_from_usize, singular_from_multi, tol_as_f32, tol_as_f64,
    usize_from_iparam,
};
use crate::which::Which;

/// Tunable parameters for the complex Arnoldi driver. The fields
/// have the same meaning as in [`crate::symmetric::Options`]; the
/// type is duplicated rather than shared so the two driver families
/// can evolve independently without breaking each other's API.
#[derive(Clone, Debug)]
pub struct Options {
    /// Convergence tolerance. `0.0` accepts ARPACK's default
    /// (machine epsilon for the working precision).
    pub tol: f64,
    /// Maximum number of restart iterations.
    pub max_iter: usize,
    /// Krylov-subspace dimension `ncv`. Must satisfy
    /// `nev + 2 <= ncv < n` — `zneupd` requires at least two extra
    /// Krylov vectors for restart-deflation (`ncv - nev >= 2`,
    /// stricter than the symmetric driver's `ncv > nev`), and
    /// `ncv == n` would leave IRLM no room to restart. `None`
    /// selects `min(2*nev + 4, n - 1)` floored at `nev + 2`.
    pub ncv: Option<usize>,
}

impl Default for Options {
    fn default() -> Self {
        Self {
            tol: 0.0,
            max_iter: 300,
            ncv: None,
        }
    }
}

/// Compute up to `nev` eigenpairs of a complex linear operator.
///
/// `matvec(x, y)` must compute `y <- A x` where both slices have
/// length `n`. The [`Which`] selector controls which Ritz values
/// to retain — must be one of [`Which::SmallestRealPart`],
/// [`Which::LargestRealPart`], [`Which::SmallestImagPart`],
/// [`Which::LargestImagPart`], [`Which::SmallestMagnitude`],
/// [`Which::LargestMagnitude`] (per-family restriction enforced
/// by the wrapper).
///
/// Returns a [`MultiEigSolution<Complex64>`] holding up to `nev`
/// converged eigenpairs plus iparam diagnostics. The
/// `eigenvalues` / `eigenvectors` arrays both have length
/// `min(nconv, nev)`; the raw ARPACK count is preserved in
/// `nconv` for diagnostics (it occasionally exceeds `nev` when
/// extra Ritz values converge to tolerance, but the extra
/// values are not surfaced because the extraction buffer is
/// `nev`-sized).
///
/// # Ordering
///
/// Unlike the real-symmetric drivers, the complex Arnoldi
/// extraction routine (`zneupd`) does **not** sort its output —
/// the order in `eigenvalues` depends on ARPACK's internal
/// selection state and is not contracted by this wrapper.
/// Callers that need a stable order must sort the returned
/// vectors themselves (e.g. by `eigenvalues[k].re`).
///
/// # Partial convergence
///
/// On `Options::max_iter` exhaustion with `nconv == 0`, returns
/// [`Error::MaxIterReached`]; on exhaustion with
/// `0 < nconv < nev`, returns `Ok(MultiEigSolution { nconv, .. })`
/// carrying the partial set (ARPACK's `zneupd` guards `-14` on
/// `nconv .le. 0`, so extraction succeeds for any `nconv >= 1`).
///
/// # Allocation
///
/// Workspace size scales as `O(n * ncv + ncv^2)`. Inputs whose
/// byte size exceeds `isize::MAX` cause `Vec` allocations to
/// panic rather than return [`Error::InvalidParam`] — same
/// convention as the standard library and the `symmetric` module.
pub fn eigenpairs_c64<F>(
    n: usize,
    nev: usize,
    which: Which,
    matvec: F,
    options: &Options,
) -> Result<MultiEigSolution<Complex64>, Error>
where
    F: FnMut(&[Complex64], &mut [Complex64]),
{
    eigenpairs_c64_impl(n, nev, which, matvec, options)
}

/// Compute up to `nev` eigenpairs of a complex linear operator,
/// f32 precision. See [`eigenpairs_c64`] for the long-form
/// contract.
pub fn eigenpairs_c32<F>(
    n: usize,
    nev: usize,
    which: Which,
    matvec: F,
    options: &Options,
) -> Result<MultiEigSolution<Complex32>, Error>
where
    F: FnMut(&[Complex32], &mut [Complex32]),
{
    eigenpairs_c32_impl(n, nev, which, matvec, options)
}

/// Smallest-real-part eigenpair of a complex linear operator.
///
/// Thin wrapper around [`eigenpairs_c64`] with `nev = 1` and
/// [`Which::SmallestRealPart`]. For Hermitian operators the
/// returned eigenvalue's imaginary part is numerically zero and
/// the real part is the smallest eigenvalue.
///
/// On `Options::max_iter` exhaustion, returns
/// [`Error::MaxIterReached`].
pub fn smallest_eigenpair_c64<F>(
    n: usize,
    matvec: F,
    options: &Options,
) -> Result<EigSolution<Complex64>, Error>
where
    F: FnMut(&[Complex64], &mut [Complex64]),
{
    let multi = eigenpairs_c64(n, 1, Which::SmallestRealPart, matvec, options)?;
    Ok(singular_from_multi(multi))
}

/// Smallest-real-part eigenpair of a complex linear operator,
/// f32 precision. See [`smallest_eigenpair_c64`] for the
/// long-form contract.
pub fn smallest_eigenpair_c32<F>(
    n: usize,
    matvec: F,
    options: &Options,
) -> Result<EigSolution<Complex32>, Error>
where
    F: FnMut(&[Complex32], &mut [Complex32]),
{
    let multi = eigenpairs_c32(n, 1, Which::SmallestRealPart, matvec, options)?;
    Ok(singular_from_multi(multi))
}

/// Generate a complex Arnoldi driver (`{c,z}{na,ne}upd_c`) for one
/// scalar precision. The two ARPACK precisions differ only in the
/// `Complex` scalar type, the underlying float used for the
/// `__BindgenComplex<_>` pointer casts and `rwork`, the
/// `*naupd_c` / `*neupd_c` symbol pair, and how `Options::tol` is
/// narrowed; everything else — the `nev + 2 <= ncv < n` bounds,
/// `lworkl = 3 * ncv^2 + 5 * ncv`, the `rwork[ncv]` / `workev[2*ncv]`
/// scratch, `ipntr(14)`, and the `info = 1` decision tree — is
/// identical, so the body is written once here and instantiated per
/// precision. Named after the type set it spans (`Complex<f32>` /
/// `Complex<f64>`), not "Scalar".
macro_rules! impl_complex_arnoldi_driver {
    ($fn:ident, $cty:ty, $float:ty, $aupd:path, $eupd:path, $tol:path) => {
        fn $fn<F>(
            n: usize,
            nev: usize,
            which: Which,
            mut matvec: F,
            options: &Options,
        ) -> Result<MultiEigSolution<$cty>, Error>
        where
            F: FnMut(&[$cty], &mut [$cty]),
        {
            if nev == 0 {
                return Err(Error::InvalidParam("nev must be positive"));
            }
            if !which.accepted_by_complex_arnoldi() {
                return Err(Error::InvalidParam(
                    "Which selector not accepted by the complex Arnoldi driver",
                ));
            }
            // Bound `nev` (caller-controlled) to the `c_int` range before
            // using it in `usize` arithmetic (`nev + 3`, `2 * nev + 4`,
            // `nev + 2`, `nev_i32 + 2`). On 64-bit targets — the only ones
            // supported per the workspace's `compile_error!` — the bounded
            // value cannot overflow in those expressions; without this
            // upfront check, `nev = usize::MAX` panics in debug builds.
            let nev_i32 = c_int_from_usize(nev)?;
            // Complex Arnoldi has a tighter constraint than real symmetric:
            // `zneupd` rejects `ncv - nev < 2` with `info = -3`, so the
            // smallest legal `ncv` is `nev + 2` and the precondition on
            // `n` is `n >= nev + 3`.
            if n < nev + 3 {
                return Err(Error::InvalidParam(
                    "n too small for complex Arnoldi (require n >= nev + 3)",
                ));
            }
            let ncv = options
                .ncv
                .unwrap_or_else(|| (2 * nev + 4).min(n - 1).max(nev + 2));

            let n_i32 = c_int_from_usize(n)?;
            let ncv_i32 = c_int_from_usize(ncv)?;
            let max_iter_i32 = c_int_from_usize(options.max_iter)?;

            if !(ncv_i32 >= nev_i32 + 2 && ncv_i32 < n_i32) {
                return Err(Error::InvalidParam("require nev + 2 <= ncv and ncv < n"));
            }
            if max_iter_i32 <= 0 {
                return Err(Error::InvalidParam("max_iter must be positive"));
            }

            let v_len = n
                .checked_mul(ncv)
                .ok_or(Error::InvalidParam("n * ncv overflows usize"))?;
            let workd_len = n
                .checked_mul(3)
                .ok_or(Error::InvalidParam("3 * n overflows usize"))?;
            let ncv_sq = ncv
                .checked_mul(ncv)
                .ok_or(Error::InvalidParam("ncv * ncv overflows usize"))?;
            let three_ncv_sq = ncv_sq
                .checked_mul(3)
                .ok_or(Error::InvalidParam("3 * ncv^2 overflows usize"))?;
            let five_ncv = ncv
                .checked_mul(5)
                .ok_or(Error::InvalidParam("5 * ncv overflows usize"))?;
            let lworkl = three_ncv_sq
                .checked_add(five_ncv)
                .ok_or(Error::InvalidParam("3*ncv^2 + 5*ncv overflows usize"))?;
            let workev_len = ncv
                .checked_mul(2)
                .ok_or(Error::InvalidParam("2 * ncv overflows usize"))?;

            let lworkl_i32 = c_int_from_usize(lworkl)?;

            let tol = $tol(options.tol);
            let zero = <$cty>::new(0.0, 0.0);
            let mut resid = vec![zero; n];
            let mut v = vec![zero; v_len];
            let ldv = n_i32;
            let mut iparam = [0i32; 11];
            iparam[0] = 1; // exact shifts via ARPACK
            iparam[2] = max_iter_i32;
            iparam[3] = 1; // NB block size; ARPACK only supports NB = 1
            iparam[6] = 1; // mode 1: standard problem A x = lambda x
            // znaupd's ICB Fortran wrapper declares `ipntr(14)`, unlike
            // dsaupd's `ipntr(11)`. Allocating shorter is an OOB write.
            let mut ipntr = [0i32; 14];
            let mut workd = vec![zero; workd_len];
            let mut workl = vec![zero; lworkl];
            let mut rwork: Vec<$float> = vec![0.0; ncv];

            let bmat = c"I".as_ptr();
            let which_ptr = which.as_c_str().as_ptr();

            let _guard = lock();

            let mut ido: c_int = 0;
            let mut info: c_int = 0;
            let mut x_buf = vec![zero; n];

            loop {
                // SAFETY: Every pointer aliases a Vec whose length matches
                // (or exceeds) what ARPACK reads/writes; the `Complex<_>` ↔
                // `__BindgenComplex<_>` cast is sound because both are
                // `#[repr(C)] { re, im }`. The process-wide lock serializes
                // ARPACK's Fortran SAVE state.
                unsafe {
                    $aupd(
                        &mut ido,
                        bmat,
                        n_i32,
                        which_ptr,
                        nev_i32,
                        tol,
                        resid.as_mut_ptr() as *mut __BindgenComplex<$float>,
                        ncv_i32,
                        v.as_mut_ptr() as *mut __BindgenComplex<$float>,
                        ldv,
                        iparam.as_mut_ptr(),
                        ipntr.as_mut_ptr(),
                        workd.as_mut_ptr() as *mut __BindgenComplex<$float>,
                        workl.as_mut_ptr() as *mut __BindgenComplex<$float>,
                        lworkl_i32,
                        rwork.as_mut_ptr(),
                        &mut info,
                    );
                }

                match ido {
                    -1 | 1 => {
                        let x_off = (ipntr[0] - 1) as usize;
                        let y_off = (ipntr[1] - 1) as usize;
                        debug_assert!(x_off + n <= workd.len() && y_off + n <= workd.len());
                        x_buf.copy_from_slice(&workd[x_off..x_off + n]);
                        matvec(&x_buf, &mut workd[y_off..y_off + n]);
                    }
                    99 => break,
                    other => return Err(Error::UnexpectedIdo(other)),
                }
            }

            // info handling (uniform with the symmetric driver per the
            // research finding that zneupd guards `-14` on `nconv .le. 0`,
            // not on `nconv < nev`), mapped by `aupd_error`:
            // - info = 0           : full convergence, extract nev pairs.
            // - info = 1, nconv=0  : MaxIterReached (skip *eupd; zneupd
            //                        would return info = -14).
            // - info = 1, nconv>=1 : call *eupd, extract `nconv` pairs.
            // - info = 3 / -9999 / other non-zero → typed error.
            let nconv = usize_from_iparam(iparam[4]);
            let iters = usize_from_iparam(iparam[2]);
            let n_matvec = usize_from_iparam(iparam[8]);

            if let Some(err) = aupd_error(info, iters, nconv, n_matvec) {
                return Err(err);
            }

            let rvec: c_int = 1;
            let howmny = c"A".as_ptr();
            let mut select = vec![0i32; ncv];
            let mut d = vec![zero; nev];
            let sigma = __BindgenComplex::<$float> { re: 0.0, im: 0.0 };
            let mut workev = vec![zero; workev_len];
            let mut info_eup: c_int = 0;

            // SAFETY: as for the *naupd_c call above; v doubles as z (output
            // eigenvector storage), which is the standard ARPACK pattern.
            unsafe {
                $eupd(
                    rvec,
                    howmny,
                    select.as_mut_ptr(),
                    d.as_mut_ptr() as *mut __BindgenComplex<$float>,
                    v.as_mut_ptr() as *mut __BindgenComplex<$float>,
                    ldv,
                    sigma,
                    workev.as_mut_ptr() as *mut __BindgenComplex<$float>,
                    bmat,
                    n_i32,
                    which_ptr,
                    nev_i32,
                    tol,
                    resid.as_mut_ptr() as *mut __BindgenComplex<$float>,
                    ncv_i32,
                    v.as_mut_ptr() as *mut __BindgenComplex<$float>,
                    ldv,
                    iparam.as_mut_ptr(),
                    ipntr.as_mut_ptr(),
                    workd.as_mut_ptr() as *mut __BindgenComplex<$float>,
                    workl.as_mut_ptr() as *mut __BindgenComplex<$float>,
                    lworkl_i32,
                    rwork.as_mut_ptr(),
                    &mut info_eup,
                );
            }

            if let Some(err) = eupd_error(info_eup, iters, nconv, n_matvec) {
                return Err(err);
            }

            // Cap the surfaced count at `nev`: ARPACK may report
            // `nconv > nev` ("bonus" Ritz values satisfying the convergence
            // bound), but `d` is sized to `nev` per the documented `*eupd`
            // interface, so slots beyond `nev` are not safely indexable.
            // Preserve the raw `iparam[4]` count in `nconv` as a diagnostic;
            // truncate the eigenvalue / eigenvector arrays to `nev`.
            let extracted = nconv.min(nev);
            let eigenvalues = d[..extracted].to_vec();
            let mut eigenvectors = Vec::with_capacity(extracted);
            for k in 0..extracted {
                eigenvectors.push(v[k * n..(k + 1) * n].to_vec());
            }

            Ok(MultiEigSolution {
                eigenvalues,
                eigenvectors,
                nev_requested: nev,
                nconv,
                iters,
                n_matvec,
            })
        }
    };
}

impl_complex_arnoldi_driver!(
    eigenpairs_c64_impl,
    Complex64,
    f64,
    znaupd_c,
    zneupd_c,
    tol_as_f64
);
impl_complex_arnoldi_driver!(
    eigenpairs_c32_impl,
    Complex32,
    f32,
    cnaupd_c,
    cneupd_c,
    tol_as_f32
);