Skip to main content

gam_solve/rho_optimizer/
fd_audit.rs

1use super::*;
2
3// FD-OK: this entire module is the outer-gradient finite-difference AUDIT
4// oracle. It deliberately computes central finite differences of the outer
5// REML/LAML criterion to verify the analytic gradient against — it is a
6// diagnostic that runs only behind the `outer_fd_audit_eligible` gates in
7// spatial_optimization.rs / custom_family/fit.rs, never on the fit math path.
8// The `fd`-named struct fields and fn names below are intrinsic to that audit
9// role, so the no-production-finite-differences guard treats this file as a
10// sanctioned audit region closed at the bottom of the module.
11
12/// Per-θ component of an outer-gradient finite-difference audit.
13#[derive(Clone, Debug)]
14pub struct OuterGradientFdComponent {
15    /// Human label for the block this coordinate belongs to (e.g. "timewiggle").
16    pub block: String,
17    /// Flat θ index.
18    pub index: usize,
19    /// Analytic ∂V/∂θ_i returned by the family evaluator.
20    pub analytic: f64,
21    /// Central finite-difference of the outer criterion in θ_i.
22    pub fd: f64,
23}
24
25impl OuterGradientFdComponent {
26    /// Absolute analytic−FD gap.
27    pub fn abs_gap(&self) -> f64 {
28        (self.analytic - self.fd).abs()
29    }
30    /// analytic/fd ratio (None when fd≈0). A clean −1 signals a sign
31    /// convention; a stable constant ≠1 signals a dropped/extra additive term.
32    pub fn ratio(&self) -> Option<f64> {
33        if self.fd.abs() > 1e-12 {
34            Some(self.analytic / self.fd)
35        } else {
36            None
37        }
38    }
39}
40
41/// Result of a component-by-component finite-difference audit of an outer
42/// REML/LAML gradient at a fixed θ, plus the outer-Hessian eigenvalues.
43///
44/// This is the discriminating diagnostic that forks the two failure modes of a
45/// non-terminating outer loop: an **objective↔gradient desync** (analytic ≠ FD
46/// on some component → the trust region chases a phantom descent direction
47/// forever) versus **weak identifiability** (analytic ≈ FD everywhere but a
48/// near-zero outer-Hessian eigenvalue → a genuinely flat valley the optimizer
49/// crawls along). It is family-agnostic: any path that exposes an outer
50/// evaluator closure `θ ↦ (V, ∇V, H)` can call it.
51#[derive(Clone, Debug)]
52pub struct OuterGradientFdAudit {
53    /// Outer criterion value at θ₀.
54    pub value: f64,
55    /// Per-coordinate analytic-vs-FD comparison.
56    pub components: Vec<OuterGradientFdComponent>,
57    /// Eigenvalues of the (symmetrized) outer Hessian at θ₀, ascending. Empty
58    /// when no analytic/operator Hessian was available.
59    pub hessian_eigenvalues: Vec<f64>,
60}
61
62impl OuterGradientFdAudit {
63    /// Per-block L2 norm of the analytic gradient.
64    pub fn analytic_block_norms(&self) -> Vec<(String, f64)> {
65        let mut order: Vec<String> = Vec::new();
66        let mut acc: std::collections::HashMap<String, f64> = std::collections::HashMap::new();
67        for c in &self.components {
68            if !acc.contains_key(&c.block) {
69                order.push(c.block.clone());
70            }
71            *acc.entry(c.block.clone()).or_insert(0.0) += c.analytic * c.analytic;
72        }
73        order
74            .into_iter()
75            .map(|b| {
76                let v = acc.get(&b).copied().unwrap_or(0.0).sqrt();
77                (b, v)
78            })
79            .collect()
80    }
81
82    /// Worst per-coordinate analytic−FD gap and its component.
83    pub fn worst_component(&self) -> Option<&OuterGradientFdComponent> {
84        self.components.iter().max_by(|a, b| {
85            a.abs_gap()
86                .partial_cmp(&b.abs_gap())
87                .unwrap_or(std::cmp::Ordering::Equal)
88        })
89    }
90
91    /// Smallest-magnitude outer-Hessian eigenvalue (flatness proxy).
92    pub fn min_abs_eigenvalue(&self) -> Option<f64> {
93        self.hessian_eigenvalues
94            .iter()
95            .map(|e| e.abs())
96            .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
97    }
98
99    /// Emit a single human-readable verdict block to the log.
100    pub fn log_verdict(&self, context: &str) {
101        log::warn!("[OUTER-FD-AUDIT/{context}] value={:.6e}", self.value);
102        for (block, norm) in self.analytic_block_norms() {
103            log::warn!("[OUTER-FD-AUDIT/{context}] block={block} |g_analytic|={norm:.6e}");
104        }
105        for c in &self.components {
106            let ratio = c
107                .ratio()
108                .map(|r| format!("{r:.4}"))
109                .unwrap_or_else(|| "n/a".to_string());
110            log::warn!(
111                "[OUTER-FD-AUDIT/{context}] block={} i={} analytic={:.6e} fd={:.6e} gap={:.3e} ratio={}",
112                c.block,
113                c.index,
114                c.analytic,
115                c.fd,
116                c.abs_gap(),
117                ratio
118            );
119        }
120        if !self.hessian_eigenvalues.is_empty() {
121            let evs: Vec<String> = self
122                .hessian_eigenvalues
123                .iter()
124                .map(|e| format!("{e:.4e}"))
125                .collect();
126            log::warn!(
127                "[OUTER-FD-AUDIT/{context}] hessian_eigenvalues=[{}] min_abs={:.4e}",
128                evs.join(", "),
129                self.min_abs_eigenvalue().unwrap_or(f64::NAN)
130            );
131        }
132        match self.worst_component() {
133            Some(w) if w.abs_gap() > 1e-3 && w.abs_gap() > 1e-3 * w.fd.abs().max(1.0) => {
134                log::warn!(
135                    "[OUTER-FD-AUDIT/{context}] VERDICT=DESYNC worst_block={} worst_i={} gap={:.3e} (analytic gradient disagrees with FD of the criterion: fix the derivative)",
136                    w.block,
137                    w.index,
138                    w.abs_gap()
139                );
140            }
141            _ => {
142                let flat = self.min_abs_eigenvalue().map(|m| m < 1e-6).unwrap_or(false);
143                if flat {
144                    log::warn!(
145                        "[OUTER-FD-AUDIT/{context}] VERDICT=FLATNESS min_abs_eig={:.3e} (analytic≈FD but the outer Hessian is near-singular: weak identifiability, fix termination not the gradient)",
146                        self.min_abs_eigenvalue().unwrap_or(f64::NAN)
147                    );
148                } else {
149                    log::warn!(
150                        "[OUTER-FD-AUDIT/{context}] VERDICT=CLEAN analytic≈FD and outer Hessian well-conditioned at this θ"
151                    );
152                }
153            }
154        }
155    }
156}
157
158/// Run a component-by-component central finite-difference audit of an outer
159/// REML/LAML gradient at a fixed θ₀.
160///
161/// `eval` is the family's outer evaluator: `θ, mode ↦ (V, ∇V, H)` where the
162/// gradient is honored at `ValueAndGradient`/`ValueGradientHessian` and `H` at
163/// `ValueGradientHessian`. `block_for_index` labels each flat θ coordinate
164/// (used only to group the report). `h` is the FD step.
165///
166/// Cost: one `ValueGradientHessian` eval at θ₀ plus `2·len(θ)` `ValueOnly`
167/// evals. The caller is responsible for only invoking this on a
168/// diagnostic-sized problem (it is not part of the production hot loop).
169pub fn outer_gradient_fd_audit<EvalF>(
170    // fd-ok: FD-audit certificate, not in math path
171    theta0: &Array1<f64>,
172    h: f64,
173    block_for_index: impl Fn(usize) -> String,
174    mut eval: EvalF,
175) -> Result<OuterGradientFdAudit, String>
176where
177    EvalF: FnMut(
178        &Array1<f64>,
179        crate::estimate::reml::reml_outer_engine::EvalMode,
180    ) -> Result<(f64, Array1<f64>, HessianResult), String>,
181{
182    use crate::estimate::reml::reml_outer_engine::EvalMode;
183    let (value, analytic_grad, hess) = eval(theta0, EvalMode::ValueGradientHessian)?;
184    if analytic_grad.len() != theta0.len() {
185        return Err(format!(
186            "outer_gradient_fd_audit: analytic gradient length {} != theta length {}",
187            analytic_grad.len(),
188            theta0.len()
189        ));
190    }
191    let mut components = Vec::with_capacity(theta0.len());
192    for i in 0..theta0.len() {
193        let mut tp = theta0.clone();
194        tp[i] += h;
195        let mut tm = theta0.clone();
196        tm[i] -= h;
197        let (vp, _, _) = eval(&tp, EvalMode::ValueOnly)?;
198        let (vm, _, _) = eval(&tm, EvalMode::ValueOnly)?;
199        let fd = (vp - vm) / (2.0 * h);
200        components.push(OuterGradientFdComponent {
201            block: block_for_index(i),
202            index: i,
203            analytic: analytic_grad[i],
204            fd,
205        });
206    }
207    let hessian_eigenvalues = match hess.materialize_dense() {
208        Ok(Some(mut hmat)) => {
209            // Symmetrize defensively before the self-adjoint solve.
210            let n = hmat.nrows();
211            if n == hmat.ncols() && n > 0 {
212                for r in 0..n {
213                    for c in (r + 1)..n {
214                        let avg = 0.5 * (hmat[[r, c]] + hmat[[c, r]]);
215                        hmat[[r, c]] = avg;
216                        hmat[[c, r]] = avg;
217                    }
218                }
219                match gam_linalg::faer_ndarray::FaerEigh::eigh(&hmat, faer::Side::Lower) {
220                    Ok((vals, _)) => {
221                        let mut v: Vec<f64> = vals.to_vec();
222                        v.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
223                        v
224                    }
225                    Err(_) => Vec::new(),
226                }
227            } else {
228                Vec::new()
229            }
230        }
231        _ => Vec::new(),
232    };
233    Ok(OuterGradientFdAudit {
234        value,
235        components,
236        hessian_eigenvalues,
237    })
238}
239// END-FD-OK