gam_problem/

psi_terms.rs

//! Exact-Newton joint-ψ term carriers and the joint-ψ workspace trait.
//!
//! These ψ-hyperparameter term types and the [`ExactNewtonJointPsiWorkspace`]
//! trait are neutral carriers in the criterion contract: they reference only
//! [`HyperOperator`] / [`DriftDerivResult`] (defined in this crate) and ndarray
//! arrays, so `gam-problem` is their public namespace.

use crate::{DriftDerivResult, HyperOperator};
use ndarray::{Array1, Array2};
use std::sync::Arc;

#[derive(Clone)]
pub struct ExactNewtonJointPsiTerms {
    pub objective_psi: f64,
    pub score_psi: Array1<f64>,
    pub hessian_psi: Array2<f64>,
    pub hessian_psi_operator: Option<Arc<dyn HyperOperator>>,
}

impl std::fmt::Debug for ExactNewtonJointPsiTerms {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ExactNewtonJointPsiTerms")
            .field("objective_psi", &self.objective_psi)
            .field("score_psi", &self.score_psi)
            .field("hessian_psi", &self.hessian_psi)
            .field(
                "hessian_psi_operator",
                &self.hessian_psi_operator.as_ref().map(|_| "<operator>"),
            )
            .finish()
    }
}

impl ExactNewtonJointPsiTerms {
    pub fn zeros(total: usize) -> Self {
        Self {
            objective_psi: 0.0,
            score_psi: Array1::zeros(total),
            hessian_psi: Array2::zeros((total, total)),
            hessian_psi_operator: None,
        }
    }
}

#[derive(Clone)]
pub struct ExactNewtonJointPsiSecondOrderTerms {
    pub objective_psi_psi: f64,
    pub score_psi_psi: Array1<f64>,
    pub hessian_psi_psi: Array2<f64>,
    pub hessian_psi_psi_operator: Option<Arc<dyn HyperOperator>>,
}

/// Direction-contracted second-order ψ terms for the profiled θ-HVP (#740).
///
/// The per-pair [`ExactNewtonJointPsiSecondOrderTerms`] are the `(ψ_i, ψ_j)`
/// entries of the joint hyper-Hessian; assembling the full outer Hessian from
/// them costs one O(n) family row pass per pair, i.e. `K²·n`. A matrix-free
/// profiled θ-HVP never needs the individual pairs — it needs, for one applied
/// outer direction with ψ-weights `α_ψ`, the `α`-contraction of those pairs
/// against the combined ψ-direction `ψ(α) = Σ_j α_j ψ_j`:
///
/// ```text
///   objective[i] = Σ_j α_j V_{ψ_i ψ_j}
///   score[i]     = Σ_j α_j g_{ψ_i ψ_j}          (a p-vector per output row i)
///   hessian[i]   = Σ_j α_j D²_β H_L[ψ_i, ψ_j]
///                = D²_β H_L[ψ_i, ψ(α)]            (bilinearity)
/// ```
///
/// All `psi_dim` output rows share the SAME contracted second leg `ψ(α)`, so a
/// family that streams its rows once over `ψ(α)` (carrying every fixed first
/// leg `ψ_i` as a batched factor column) produces every row in a SINGLE n-pass.
/// That is the cost the profiled θ-HVP turns into `K·n`-to-densify /
/// `m·n`-in-CG instead of the dense path's `K²·n`.
///
/// Indexing is over the flattened ψ coordinates in the same order as
/// [`ExactNewtonJointPsiWorkspace::second_order_terms`]; `hessian[i]` carries
/// the `D²_β H_L[ψ_i, ψ(α)]` drift as a [`DriftDerivResult`] (dense or
/// operator-backed) plus any block-local `S_{ψ_i ψ_j}` penalty motion folded by
/// the family, exactly mirroring the per-pair `hessian_psi_psi(_operator)`.
pub struct ExactNewtonJointPsiSecondOrderContracted {
    /// `objective[i] = Σ_j α_j V_{ψ_i ψ_j}`, one scalar per ψ output row.
    pub objective: Array1<f64>,
    /// `score[i] = Σ_j α_j g_{ψ_i ψ_j}`, the `psi_dim × total` matrix whose
    /// row `i` is the contracted fixed-β score derivative for output row `i`.
    pub score: Array2<f64>,
    /// `hessian[i] = D²_β H_L[ψ_i, ψ(α)]` for each ψ output row `i`.
    pub hessian: Vec<DriftDerivResult>,
}

pub trait ExactNewtonJointPsiWorkspace: Send + Sync {
    fn first_order_terms(
        &self,
        psi_index: usize,
    ) -> Result<Option<ExactNewtonJointPsiTerms>, String> {
        // A workspace that materializes every axis in one pass still answers
        // the per-index query: select row `psi_index` out of that batch. A
        // workspace that implements neither method falls through to `Ok(None)`
        // and the caller uses the family's own per-index path.
        let Some(all) = self.first_order_terms_all()? else {
            return Ok(None);
        };
        let materialized = all.len();
        match all.into_iter().nth(psi_index) {
            Some(terms) => Ok(Some(terms)),
            None => Err(format!(
                "ExactNewtonJointPsiWorkspace: psi index {psi_index} is out of range for the \
                 {materialized} axes this workspace materialized"
            )),
        }
    }

    fn first_order_terms_all(&self) -> Result<Option<Vec<ExactNewtonJointPsiTerms>>, String> {
        Ok(None)
    }

    fn second_order_terms(
        &self,
        psi_i: usize,
        psi_j: usize,
    ) -> Result<Option<ExactNewtonJointPsiSecondOrderTerms>, String>;

    /// Direction-contracted second-order ψ terms for the profiled θ-HVP (#740).
    ///
    /// Given the ψ-block weights `alpha_psi` (length `psi_dim`, the ψ slice of
    /// one applied outer direction α), return the `α`-contraction of every
    /// `(ψ_i, ψ_j)` second-order term against the combined ψ-direction
    /// `ψ(α) = Σ_j alpha_psi[j] · ψ_j`, as
    /// [`ExactNewtonJointPsiSecondOrderContracted`]. A family that can stream
    /// its rows once over `ψ(α)` overrides this so the profiled outer-Hessian
    /// operator applies one combined-direction n-pass per matvec instead of the
    /// dense path's `K²` per-pair [`Self::second_order_terms`] passes.
    ///
    /// Default returns `None`: the profiled θ-HVP operator is then not built and
    /// the evaluator keeps the exact per-pair assembly (dense
    /// `compute_outer_hessian` / `build_outer_hessian_operator`). Overriding
    /// this method is purely a representation/cost choice — it must produce the
    /// exact same contraction the per-pair terms would, which the
    /// `profiled_theta_hvp_outer_hessian_fd` finite-difference cross-check
    /// guards.
    fn second_order_terms_contracted(
        &self,
        _: &[f64],
    ) -> Result<Option<ExactNewtonJointPsiSecondOrderContracted>, String> {
        // Default implementation ignores this parameter.
        Ok(None)
    }

    fn hessian_directional_derivative(
        &self,
        psi_index: usize,
        d_beta_flat: &Array1<f64>,
    ) -> Result<Option<DriftDerivResult>, String>;

    /// Materialize the psi-Hessian derivative along each joint coefficient
    /// basis vector, in coefficient order. Row-streaming workspaces can compute
    /// this same tensor in one pass by implementing this method directly.
    fn hessian_directional_derivatives_all_beta_axes(
        &self,
        psi_index: usize,
        total: usize,
    ) -> Result<Option<Vec<DriftDerivResult>>, String> {
        per_axis_psi_hessian_directional_derivatives(self, psi_index, total)
    }

    /// {D_beta_axis D_beta_direction D_psi H}, under this workspace's row measure.
    fn hessian_second_directional_derivative_all_beta_axes(
        &self,
        psi_index: usize,
        d_beta_flat: &Array1<f64>,
    ) -> Result<Option<Vec<Array2<f64>>>, String> {
        Err(format!("exact third information derivatives are unavailable for psi axis {psi_index} and coefficient direction of length {}", d_beta_flat.len()))
    }

    /// {D_beta_axis D_psi_i D_psi_j H}, under this workspace's row measure.
    fn second_order_hessian_directional_derivative_all_beta_axes(
        &self,
        psi_i: usize,
        psi_j: usize,
    ) -> Result<Option<Vec<Array2<f64>>>, String> {
        Err(format!("exact third information derivatives are unavailable for psi pair ({psi_i}, {psi_j})"))
    }
}

/// Assemble the coefficient-axis tensor from exact directional derivatives.
/// A missing axis makes the whole tensor unavailable; errors propagate intact.
pub fn per_axis_psi_hessian_directional_derivatives(
    workspace: &(impl ExactNewtonJointPsiWorkspace + ?Sized),
    psi_index: usize,
    total: usize,
) -> Result<Option<Vec<DriftDerivResult>>, String> {
    let mut axes = Vec::with_capacity(total);
    let mut direction = Array1::<f64>::zeros(total);
    for axis in 0..total {
        direction[axis] = 1.0;
        let Some(derivative) = workspace.hessian_directional_derivative(psi_index, &direction)?
        else {
            return Ok(None);
        };
        axes.push(derivative);
        direction[axis] = 0.0;
    }
    Ok(Some(axes))
}

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

    #[test]
    fn zeros_has_zero_objective() {
        let t = ExactNewtonJointPsiTerms::zeros(3);
        assert_eq!(t.objective_psi, 0.0);
    }

    #[test]
    fn zeros_has_correct_score_dimension() {
        let t = ExactNewtonJointPsiTerms::zeros(5);
        assert_eq!(t.score_psi.len(), 5);
        assert!(t.score_psi.iter().all(|&v| v == 0.0));
    }

    #[test]
    fn zeros_has_square_hessian_of_correct_size() {
        let t = ExactNewtonJointPsiTerms::zeros(4);
        assert_eq!(t.hessian_psi.nrows(), 4);
        assert_eq!(t.hessian_psi.ncols(), 4);
        assert!(t.hessian_psi.iter().all(|&v| v == 0.0));
    }

    #[test]
    fn zeros_has_no_operator() {
        let t = ExactNewtonJointPsiTerms::zeros(2);
        assert!(t.hessian_psi_operator.is_none());
    }

    #[test]
    fn zeros_with_dimension_zero_does_not_panic() {
        let t = ExactNewtonJointPsiTerms::zeros(0);
        assert_eq!(t.score_psi.len(), 0);
        assert_eq!(t.hessian_psi.nrows(), 0);
    }
}