solver/numeric/

ergodic.rs

//! Ergonomic stationary eigenproblems solved by Perron iteration.
//!
//! Infinite-horizon control problems without discounting have a one-parameter
//! translation invariance and therefore no unique solution of
//! `sup_u { f + L^u V } = 0`. Their exact `T -> infinity` limit is instead the
//! principal (Perron) eigenpair of a linear operator.
//!
//! This module provides a symmetric-operator contract and a power-iteration
//! solver for that eigenpair. It is deliberately separate from
//! [`crate::numeric::finite_difference::elliptic`], which solves the
//! discounted elliptic problem `-T V + r V = source` with `r > 0`. Those two
//! problems coincide only when an explicit normalization is imposed.
//!
//! See `docs/src/reference/soc_exact.md` for the derivation and the
//! sign convention.

use crate::core::grid::Grid;
use crate::linalg::csr::CsrMatrix;

/// A symmetric linear operator whose principal eigenvector is a stationary
/// value.
///
/// The operator is stored in the sign convention used by the exact reference:
/// for the Avellaneda-Stoikov GLT matrix this is the negated matrix
/// `diag(-alpha q^2) + offdiag(+eta)` whose Perron eigenvector is the
/// long-horizon limit.
pub trait StationaryEigenProblem<const N: usize> {
    /// Builds the symmetric operator on `grid`.
    ///
    /// The returned matrix must have `grid.total_size()` rows and columns.
    fn operator(&self, grid: &Grid<N>) -> CsrMatrix;
}

/// The principal eigenpair returned by [`PerronSolver::solve`].
#[derive(Clone, Debug)]
pub struct PerronSolution {
    /// Largest algebraic eigenvalue of the operator.
    pub eigenvalue: f64,
    /// Unit-norm eigenvector associated with `eigenvalue`.
    pub eigenvector: Vec<f64>,
}

/// Power-iteration solver for the principal eigenpair of a symmetric operator.
#[derive(Clone, Copy, Debug)]
pub struct PerronSolver {
    /// Convergence tolerance on the L2 distance between consecutive iterates.
    pub tol: f64,
    /// Maximum number of power iterations.
    pub max_iter: usize,
    /// Extra shift added to the Gershgorin lower bound. A positive margin
    /// guarantees a strictly positive definite shifted operator.
    pub shift_margin: f64,
}

impl Default for PerronSolver {
    fn default() -> Self {
        Self {
            tol: 1e-10,
            max_iter: 10_000,
            shift_margin: 1.0,
        }
    }
}

impl PerronSolver {
    /// Creates a solver with default settings.
    ///
    /// # Examples
    ///
    /// ```
    /// use solver::numeric::ergodic::PerronSolver;
    /// let solver = PerronSolver::new();
    /// assert_eq!(solver.tol, 1e-10);
    /// ```
    pub fn new() -> Self {
        Self::default()
    }

    /// Selects the power-iteration convergence tolerance.
    ///
    /// # Panics
    ///
    /// Panics if `tol` is not positive.
    ///
    /// # Examples
    ///
    /// ```
    /// use solver::numeric::ergodic::PerronSolver;
    /// let solver = PerronSolver::new().with_tol(1e-12);
    /// assert_eq!(solver.tol, 1e-12);
    /// ```
    pub fn with_tol(mut self, tol: f64) -> Self {
        assert!(tol > 0.0, "tolerance must be positive");
        self.tol = tol;
        self
    }

    /// Selects the maximum number of power iterations.
    ///
    /// # Panics
    ///
    /// Panics if `max_iter` is zero.
    ///
    /// # Examples
    ///
    /// ```
    /// use solver::numeric::ergodic::PerronSolver;
    /// let solver = PerronSolver::new().with_max_iter(500);
    /// assert_eq!(solver.max_iter, 500);
    /// ```
    pub fn with_max_iter(mut self, max_iter: usize) -> Self {
        assert!(max_iter > 0, "max_iter must be positive");
        self.max_iter = max_iter;
        self
    }

    /// Selects the shift margin added to the Gershgorin lower bound.
    ///
    /// # Panics
    ///
    /// Panics if `shift_margin` is not positive.
    ///
    /// # Examples
    ///
    /// ```
    /// use solver::numeric::ergodic::PerronSolver;
    /// let solver = PerronSolver::new().with_shift_margin(0.5);
    /// assert_eq!(solver.shift_margin, 0.5);
    /// ```
    pub fn with_shift_margin(mut self, shift_margin: f64) -> Self {
        assert!(shift_margin > 0.0, "shift margin must be positive");
        self.shift_margin = shift_margin;
        self
    }

    /// Solves for the principal eigenpair of `problem` on `grid`.
    ///
    /// Power iteration is applied to `A + shift * I`, where `A` is the
    /// operator and `shift` is chosen from the Gershgorin lower bound so that
    /// the shifted operator is positive definite. Shifting preserves
    /// eigenvectors and moves every eigenvalue by the same constant, so the
    /// returned eigenvalue subtracts the shift.
    ///
    /// # Panics
    ///
    /// Panics if `problem.operator(grid)` has a different size from the grid.
    ///
    /// # Examples
    ///
    /// ```
    /// use solver::core::grid::Grid;
    /// use solver::linalg::csr::CsrMatrix;
    /// use solver::numeric::ergodic::{PerronSolver, StationaryEigenProblem};
    ///
    /// struct Single;
    /// impl StationaryEigenProblem<1> for Single {
    ///     fn operator(&self, _grid: &Grid<1>) -> CsrMatrix {
    ///         let mut mat = CsrMatrix::new(1, 1);
    ///         mat.add_entry(0, 5.0);
    ///         mat.finish_row();
    ///         mat
    ///     }
    /// }
    ///
    /// let grid = Grid::<1>::new([1], [0.0], [0.0]);
    /// let solution = PerronSolver::new().solve(&grid, &Single);
    /// assert!((solution.eigenvalue - 5.0).abs() < 1e-12);
    /// ```
    pub fn solve<const N: usize, P: StationaryEigenProblem<N>>(
        &self,
        grid: &Grid<N>,
        problem: &P,
    ) -> PerronSolution {
        let mat = problem.operator(grid);
        let n = mat.size;
        assert_eq!(n, grid.total_size(), "operator size must match the grid");

        let shift = self.shift(&mat);
        let mut v = vec![(n as f64).sqrt().recip(); n];

        for _ in 0..self.max_iter {
            let w = shifted_mat_vec(&mat, shift, &v);
            let norm = l2_norm(&w);
            if norm == 0.0 {
                break;
            }

            let mut next = w;
            for entry in next.iter_mut() {
                *entry /= norm;
            }

            let delta = l2_distance(&v, &next);
            v = next;
            if delta < self.tol {
                break;
            }
        }

        let av = mat_vec(&mat, &v);
        let vv = dot(&v, &v);
        let eigenvalue = if vv > 0.0 { dot(&v, &av) / vv } else { 0.0 };

        PerronSolution {
            eigenvalue,
            eigenvector: v,
        }
    }

    /// Computes the shift that makes `A + shift * I` positive definite.
    ///
    /// The Gershgorin theorem gives a lower bound on the smallest eigenvalue.
    /// Subtracting that bound and adding a margin guarantees non-negative
    /// diagonal entries and a non-negative shifted operator when the
    /// off-diagonal entries are already non-negative.
    fn shift(&self, mat: &CsrMatrix) -> f64 {
        let mut lower = f64::INFINITY;
        for i in 0..mat.size {
            let row_start = mat.row_ptr[i];
            let row_end = mat.row_ptr[i + 1];

            let mut diagonal = 0.0;
            let mut off_sum = 0.0;
            for idx in row_start..row_end {
                let val = mat.values[idx];
                if mat.col_indices[idx] == i {
                    diagonal = val;
                } else {
                    off_sum += val.abs();
                }
            }

            lower = lower.min(diagonal - off_sum);
        }

        (-lower).max(0.0) + self.shift_margin
    }
}

/// Dense matrix-vector product for a CSR matrix.
fn mat_vec(mat: &CsrMatrix, x: &[f64]) -> Vec<f64> {
    let mut y = vec![0.0; mat.size];
    for (i, y_i) in y.iter_mut().enumerate() {
        let mut acc = 0.0;
        for idx in mat.row_ptr[i]..mat.row_ptr[i + 1] {
            acc += mat.values[idx] * x[mat.col_indices[idx]];
        }
        *y_i = acc;
    }
    y
}

/// Applies `(A + shift * I) x`.
fn shifted_mat_vec(mat: &CsrMatrix, shift: f64, x: &[f64]) -> Vec<f64> {
    let mut y = mat_vec(mat, x);
    for (entry, &x_i) in y.iter_mut().zip(x) {
        *entry += shift * x_i;
    }
    y
}

fn dot(a: &[f64], b: &[f64]) -> f64 {
    a.iter().zip(b).map(|(x, y)| x * y).sum()
}

fn l2_norm(x: &[f64]) -> f64 {
    x.iter().map(|v| v * v).sum::<f64>().sqrt()
}

fn l2_distance(a: &[f64], b: &[f64]) -> f64 {
    a.iter()
        .zip(b)
        .map(|(x, y)| (x - y).powi(2))
        .sum::<f64>()
        .sqrt()
}

#[cfg(test)]
mod tests {
    use super::*;

    struct KnownTridiagonal;

    impl StationaryEigenProblem<1> for KnownTridiagonal {
        fn operator(&self, _grid: &Grid<1>) -> CsrMatrix {
            // Diagonal -2, off-diagonal +1. Eigenvalues are -1 and -3; the
            // Perron eigenvector for the largest eigenvalue -1 is [1, 1].
            let mut mat = CsrMatrix::new(2, 6);
            mat.add_entry(0, -2.0);
            mat.add_entry(1, 1.0);
            mat.finish_row();
            mat.add_entry(0, 1.0);
            mat.add_entry(1, -2.0);
            mat.finish_row();
            mat
        }
    }

    struct SingleNode;

    impl StationaryEigenProblem<1> for SingleNode {
        fn operator(&self, _grid: &Grid<1>) -> CsrMatrix {
            let mut mat = CsrMatrix::new(1, 1);
            mat.add_entry(0, 5.0);
            mat.finish_row();
            mat
        }
    }

    #[test]
    fn perron_iteration_recovers_largest_eigenpair() {
        let grid = Grid::<1>::new([2], [0.0], [1.0]);
        let solution = PerronSolver::new()
            .with_tol(1e-12)
            .solve(&grid, &KnownTridiagonal);

        assert!((solution.eigenvalue - (-1.0)).abs() < 1e-8);
        assert!((solution.eigenvector[0] - solution.eigenvector[1]).abs() < 1e-8);
        assert!(solution.eigenvector[0] > 0.0);
    }

    #[test]
    fn perron_single_node_recovers_its_scalar() {
        let grid = Grid::<1>::new([1], [0.0], [0.0]);
        let solution = PerronSolver::new().solve(&grid, &SingleNode);

        assert_eq!(solution.eigenvector.len(), 1);
        assert!((solution.eigenvalue - 5.0).abs() < 1e-12);
        assert!(solution.eigenvector[0] > 0.0);
    }
}