gam_solve/estimate/
outer_eval_capture.rs

1//! Structured capture of outer-objective evidence for integration tests.
2//!
3//! The raw-evaluation window serves flexible-link measurements (#1876). The
4//! finite-difference record serves end-to-end gradient gates (#2460): when
5//! explicitly enabled, the generic outer runner compares the analytic gradient
6//! at its first bounded seed with a finite difference of that same objective.
7//! Tests consume typed arrays rather than scraping formatted production logs.
8//!
9//! Both channels are disabled by default. The raw window is process-global
10//! because its flexible-link measurements intentionally span helper calls. The
11//! finite-difference request is thread-local: a parallel integration test can
12//! neither consume nor overwrite another test's one-shot audit.
13
14use ndarray::{Array1, Array2};
15use std::cell::RefCell;
16use std::sync::atomic::{AtomicBool, Ordering};
17use std::sync::{Mutex, OnceLock};
18
19/// One captured outer evaluation: the outer coordinate `theta = (ρ ‖ link)`, the
20/// scalar cost, and the analytic outer gradient in the same layout.
21#[derive(Clone, Debug)]
22pub struct OuterEvalRecord {
23    pub theta: Array1<f64>,
24    pub cost: f64,
25    pub gradient: Array1<f64>,
26}
27
28/// Analytic-vs-finite-difference evidence for the ψ block at one real outer
29/// seed.
30///
31/// `theta` retains the complete outer seed and `rho_dim` locates the ψ block in
32/// that seed. Every gradient and scalar-stencil array contains exactly
33/// `psi_dim` entries in ψ-local order. Smoothing-parameter ρ coordinates are
34/// deliberately excluded: the κ/geometry gates that request this record do not
35/// grade them, and each unnecessary finite-difference coordinate costs two
36/// complete inner profiles.
37#[derive(Clone, Debug)]
38pub struct OuterGradientFdRecord {
39    pub theta: Array1<f64>,
40    pub rho_dim: usize,
41    pub psi_dim: usize,
42    pub cost: f64,
43    pub analytic_psi_gradient: Array1<f64>,
44    pub finite_difference_psi_gradient: Array1<f64>,
45    pub psi_steps: Array1<f64>,
46    pub fixed_beta_psi_gradient: Array1<f64>,
47    pub logdet_h_psi_gradient: Array1<f64>,
48    pub frozen_logdet_h_psi_gradient: Array1<f64>,
49    pub mode_response_logdet_h_psi_gradient: Array1<f64>,
50    pub analytic_mode_response_norm: Array1<f64>,
51    pub finite_difference_mode_response_norm: Array1<f64>,
52    pub mode_response_relative_error: Array1<f64>,
53    pub mode_response_max_abs_error: Array1<f64>,
54    pub logdet_s_psi_gradient: Array1<f64>,
55    pub kkt_psi_gradient: Array1<f64>,
56    pub finite_difference_fixed_beta_psi_gradient: Array1<f64>,
57    pub finite_difference_logdet_h_psi_gradient: Array1<f64>,
58    pub finite_difference_logdet_s_psi_gradient: Array1<f64>,
59    pub finite_difference_kkt_psi_gradient: Array1<f64>,
60}
61
62/// Maximum evaluations retained per capture window (opening iterates only).
63const MAX_CAPTURED: usize = 8;
64
65static ENABLED: AtomicBool = AtomicBool::new(false);
66
67struct OuterGradientFdCapture {
68    min_psi_dim: usize,
69    record: Option<OuterGradientFdRecord>,
70    components: Vec<(f64, f64, f64, f64, f64, f64)>,
71    criterion_components: Option<(f64, [f64; 4])>,
72    selected_mode: Option<(Array1<f64>, Option<Array2<f64>>)>,
73}
74
75thread_local! {
76    static FD_CAPTURE: RefCell<Option<OuterGradientFdCapture>> = const { RefCell::new(None) };
77}
78
79fn buffer() -> &'static Mutex<Vec<OuterEvalRecord>> {
80    static BUFFER: OnceLock<Mutex<Vec<OuterEvalRecord>>> = OnceLock::new();
81    BUFFER.get_or_init(|| Mutex::new(Vec::new()))
82}
83
84/// Start capturing outer evaluations, clearing any prior window.
85pub fn enable_outer_eval_capture() {
86    buffer().lock().expect("outer-eval capture buffer").clear();
87    ENABLED.store(true, Ordering::Relaxed);
88}
89
90/// Stop capturing and drain the recorded opening evaluations (in eval order).
91pub fn take_outer_eval_capture() -> Vec<OuterEvalRecord> {
92    ENABLED.store(false, Ordering::Relaxed);
93    std::mem::take(&mut *buffer().lock().expect("outer-eval capture buffer"))
94}
95
96/// Request one structured audit at the next outer seed with enough ψ axes.
97pub fn enable_outer_gradient_fd_capture(min_psi_dim: usize) {
98    FD_CAPTURE.with(|capture| {
99        *capture.borrow_mut() = Some(OuterGradientFdCapture {
100            min_psi_dim,
101            record: None,
102            components: Vec::new(),
103            criterion_components: None,
104            selected_mode: None,
105        });
106    });
107}
108
109pub(crate) fn begin_outer_gradient_component_capture() {
110    FD_CAPTURE.with(|capture| {
111        if let Some(state) = capture.borrow_mut().as_mut() {
112            state.components.clear();
113        }
114    });
115}
116
117pub(crate) fn outer_gradient_component_capture_enabled() -> bool {
118    FD_CAPTURE.with(|capture| {
119        capture
120            .borrow()
121            .as_ref()
122            .is_some_and(|state| state.record.is_none())
123    })
124}
125
126pub(crate) fn record_outer_gradient_component(
127    fixed_beta: f64,
128    logdet_h: f64,
129    frozen_logdet_h: f64,
130    mode_response_logdet_h: f64,
131    logdet_s: f64,
132    kkt: f64,
133) {
134    FD_CAPTURE.with(|capture| {
135        if let Some(state) = capture.borrow_mut().as_mut()
136            && state.record.is_none()
137        {
138            state.components.push((
139                fixed_beta,
140                logdet_h,
141                frozen_logdet_h,
142                mode_response_logdet_h,
143                logdet_s,
144                kkt,
145            ));
146        }
147    });
148}
149
150pub(crate) fn take_outer_gradient_components() -> Vec<(f64, f64, f64, f64, f64, f64)> {
151    FD_CAPTURE.with(|capture| {
152        capture
153            .borrow_mut()
154            .as_mut()
155            .map_or_else(Vec::new, |state| std::mem::take(&mut state.components))
156    })
157}
158
159pub(crate) fn begin_outer_criterion_component_capture() {
160    FD_CAPTURE.with(|capture| {
161        if let Some(state) = capture.borrow_mut().as_mut() {
162            state.criterion_components = None;
163            state.selected_mode = None;
164        }
165    });
166}
167
168/// Retain the final selected scalar-criterion decomposition for an armed
169/// outer-gradient audit.
170///
171/// This is public only so sibling workspace evaluators can report through the
172/// same typed sink after their own nonconvex mode selection. It is a no-op
173/// unless [`enable_outer_gradient_fd_capture`] armed the calling thread.
174pub fn record_outer_criterion_components(cost: f64, components: [f64; 4]) {
175    FD_CAPTURE.with(|capture| {
176        if let Some(state) = capture.borrow_mut().as_mut()
177            && state.record.is_none()
178        {
179            state.criterion_components = Some((cost, components));
180        }
181    });
182}
183
184pub(crate) fn take_outer_criterion_components() -> Option<(f64, [f64; 4])> {
185    FD_CAPTURE.with(|capture| {
186        capture
187            .borrow_mut()
188            .as_mut()
189            .and_then(|state| state.criterion_components.take())
190    })
191}
192
193/// Retain the selected coefficient mode and its analytic extended-coordinate
194/// response columns for an armed finite-difference audit.
195///
196/// Sibling workspace evaluators call this only after nonconvex candidate
197/// selection, beside [`record_outer_criterion_components`]. Value-only
198/// evaluations pass no response columns but still retain their selected
199/// coefficients for the scalar stencil.
200pub fn record_outer_selected_mode(
201    beta: Array1<f64>,
202    ext_mode_response_cols: Option<Array2<f64>>,
203) {
204    FD_CAPTURE.with(|capture| {
205        if let Some(state) = capture.borrow_mut().as_mut()
206            && state.record.is_none()
207        {
208            state.selected_mode = Some((beta, ext_mode_response_cols));
209        }
210    });
211}
212
213pub(crate) fn take_outer_selected_mode() -> Option<(Array1<f64>, Option<Array2<f64>>)> {
214    FD_CAPTURE.with(|capture| {
215        capture
216            .borrow_mut()
217            .as_mut()
218            .and_then(|state| state.selected_mode.take())
219    })
220}
221
222/// Stop the audit window and take its single record.
223pub fn take_outer_gradient_fd_capture() -> Option<OuterGradientFdRecord> {
224    FD_CAPTURE.with(|capture| capture.borrow_mut().take().and_then(|state| state.record))
225}
226
227pub(crate) fn outer_gradient_fd_capture_enabled(psi_dim: usize) -> bool {
228    FD_CAPTURE.with(|capture| {
229        capture
230            .borrow()
231            .as_ref()
232            .is_some_and(|state| state.record.is_none() && psi_dim >= state.min_psi_dim)
233    })
234}
235
236pub(crate) fn record_outer_gradient_fd(record: OuterGradientFdRecord) {
237    FD_CAPTURE.with(|capture| {
238        if let Some(state) = capture.borrow_mut().as_mut()
239            && state.record.is_none()
240            && record.psi_dim >= state.min_psi_dim
241        {
242            state.record = Some(record);
243        }
244    });
245}
246
247/// Record one outer evaluation when capture is enabled (no-op otherwise). Only
248/// the first [`MAX_CAPTURED`] evaluations of a window are retained.
249pub(crate) fn record_outer_eval(theta: &Array1<f64>, cost: f64, gradient: &Array1<f64>) {
250    if !ENABLED.load(Ordering::Relaxed) {
251        return;
252    }
253    let mut b = buffer().lock().expect("outer-eval capture buffer");
254    if b.len() < MAX_CAPTURED {
255        b.push(OuterEvalRecord {
256            theta: theta.clone(),
257            cost,
258            gradient: gradient.clone(),
259        });
260    }
261}
gam_solve/estimate/outer_eval_capture.rs

gam_solve/estimate/
outer_eval_capture.rs
