fpm_rs/diagnostics/
recorder.rs

1use std::{
2    sync::{Arc, Mutex, MutexGuard},
3    time::Instant,
4};
5
6use crate::{
7    Array2, Result,
8    callbacks::{Callback, CallbackAction, StepContext},
9    reconstruction::ReconstructionResult,
10};
11
12use super::{
13    DiagnosticRequest, IterationDiagnostics, ReconstructionDiagnostics, compute_fourier_coverage,
14};
15
16#[derive(Clone, Debug)]
17pub struct DiagnosticRecorderConfig {
18    pub every: usize,
19
20    pub record_iteration_history: bool,
21    pub record_frame_summaries: bool,
22    pub record_raw_stack_stats: bool,
23    pub record_coverage: bool,
24
25    pub record_object_snapshots: bool,
26    pub record_pupil_snapshots: bool,
27    pub snapshot_every: usize,
28}
29
30impl Default for DiagnosticRecorderConfig {
31    fn default() -> Self {
32        Self {
33            every: 1,
34            record_iteration_history: true,
35            record_frame_summaries: false,
36            record_raw_stack_stats: false,
37            record_coverage: false,
38            record_object_snapshots: false,
39            record_pupil_snapshots: false,
40            snapshot_every: 10,
41        }
42    }
43}
44
45#[derive(Default)]
46struct DiagnosticRecorderState {
47    diagnostics: ReconstructionDiagnostics,
48    started_at: Option<Instant>,
49    previous_object: Option<Array2<Complex64Proxy>>,
50    previous_pupil: Option<Array2<Complex64Proxy>>,
51    object_snapshots: Vec<(usize, Array2<f64>, Array2<f64>)>,
52    pupil_snapshots: Vec<(usize, Array2<f64>, Array2<f64>)>,
53}
54
55/// A cloneable callback whose recorded output remains accessible after a run.
56///
57/// Pass a clone to a runner and retain the original handle:
58///
59/// ```no_run
60/// # use fpm_rs::diagnostics::{DiagnosticRecorder, DiagnosticRecorderConfig};
61/// let recorder = DiagnosticRecorder::new(DiagnosticRecorderConfig::default());
62/// let callback = Box::new(recorder.clone());
63/// # let _ = callback;
64/// // Run with `callback`, then read `recorder.diagnostics()`.
65/// ```
66///
67/// Clones share one recorder state and may be read from another thread while a
68/// run is active, but they must not be installed in multiple concurrent runs:
69/// each run resets the shared state at its start. Recorder state is
70/// best-effort diagnostics, so a poisoned mutex is recovered and reset on the
71/// next run rather than preventing later diagnostic reads or reuse.
72#[derive(Clone)]
73pub struct DiagnosticRecorder {
74    config: DiagnosticRecorderConfig,
75    state: Arc<Mutex<DiagnosticRecorderState>>,
76}
77
78type Complex64Proxy = num_complex::Complex64;
79
80impl DiagnosticRecorder {
81    pub fn new(config: DiagnosticRecorderConfig) -> Self {
82        Self {
83            config,
84            state: Arc::new(Mutex::new(DiagnosticRecorderState::default())),
85        }
86    }
87
88    pub fn diagnostics(&self) -> ReconstructionDiagnostics {
89        self.lock_state().diagnostics.clone()
90    }
91
92    pub fn into_diagnostics(self) -> ReconstructionDiagnostics {
93        self.diagnostics()
94    }
95
96    pub fn object_snapshots(&self) -> Vec<(usize, Array2<f64>, Array2<f64>)> {
97        self.lock_state().object_snapshots.clone()
98    }
99
100    pub fn pupil_snapshots(&self) -> Vec<(usize, Array2<f64>, Array2<f64>)> {
101        self.lock_state().pupil_snapshots.clone()
102    }
103
104    /// Clears all values collected by this recorder and its clones.
105    pub fn reset(&self) {
106        *self.lock_state() = DiagnosticRecorderState::default();
107    }
108
109    fn lock_state(&self) -> MutexGuard<'_, DiagnosticRecorderState> {
110        self.state
111            .lock()
112            .unwrap_or_else(|poisoned| poisoned.into_inner())
113    }
114}
115
116impl Callback for DiagnosticRecorder {
117    fn requires_for(
118        &self,
119        hook: crate::callbacks::CallbackHook,
120        iteration: usize,
121    ) -> Vec<DiagnosticRequest> {
122        let should_record = cadence_matches(iteration, self.config.every);
123        let should_snapshot = cadence_matches(iteration, self.config.snapshot_every);
124        match hook {
125            crate::callbacks::CallbackHook::Start => {
126                let mut requests = Vec::new();
127                if self.config.record_raw_stack_stats {
128                    requests.push(DiagnosticRequest::RawFrameStats);
129                }
130                requests
131            }
132            crate::callbacks::CallbackHook::IterationEnd if should_record || should_snapshot => {
133                let mut requests = Vec::new();
134                if should_record && self.config.record_iteration_history {
135                    requests.push(DiagnosticRequest::Loss);
136                    requests.push(DiagnosticRequest::PerFrameError);
137                }
138                if should_record && self.config.record_frame_summaries {
139                    requests.push(DiagnosticRequest::FrameSummaries);
140                }
141                if should_snapshot && self.config.record_object_snapshots {
142                    requests.push(DiagnosticRequest::ObjectAmplitude);
143                    requests.push(DiagnosticRequest::ObjectPhase);
144                }
145                if should_snapshot && self.config.record_pupil_snapshots {
146                    requests.push(DiagnosticRequest::Pupil);
147                }
148                requests
149            }
150            _ => Vec::new(),
151        }
152    }
153
154    fn on_start(&mut self, context: &StepContext<'_>) -> Result<CallbackAction> {
155        let mut state = self.lock_state();
156        *state = DiagnosticRecorderState::default();
157        state.started_at = Some(Instant::now());
158        if self.config.record_coverage {
159            state.diagnostics.coverage = Some(compute_fourier_coverage(context.model)?);
160        }
161        if let Some(raw_frame_stats) = &context.diagnostics.raw_frame_stats {
162            state
163                .diagnostics
164                .raw_frame_stats
165                .extend(raw_frame_stats.iter().cloned());
166        }
167        Ok(CallbackAction::Continue)
168    }
169
170    fn on_iteration_end(&mut self, context: &StepContext<'_>) -> Result<CallbackAction> {
171        let should_record = cadence_matches(context.iteration, self.config.every);
172        let should_snapshot = cadence_matches(context.iteration, self.config.snapshot_every);
173        let mut state = self.lock_state();
174
175        if should_record && self.config.record_iteration_history {
176            let elapsed_ms = state
177                .started_at
178                .map(|started| started.elapsed().as_secs_f64() * 1e3);
179            let object_relative_change = relative_change(
180                state.previous_object.as_ref(),
181                &context.state.object_spectrum,
182            );
183            let pupil_relative_change =
184                relative_change(state.previous_pupil.as_ref(), &context.state.pupil.values);
185            state
186                .diagnostics
187                .iteration_history
188                .push(IterationDiagnostics {
189                    iteration: context.iteration,
190                    total_loss: context.diagnostics.loss,
191                    data_loss: context.diagnostics.loss,
192                    regularization_loss: None,
193                    object_relative_change,
194                    pupil_relative_change,
195                    median_frame_loss: context
196                        .diagnostics
197                        .per_frame_error
198                        .as_ref()
199                        .and_then(|values| median(values)),
200                    worst_frame_loss: context
201                        .diagnostics
202                        .per_frame_error
203                        .as_ref()
204                        .and_then(|values| values.iter().copied().reduce(f64::max)),
205                    elapsed_ms,
206                });
207        }
208        if should_record {
209            state.previous_object = Some(context.state.object_spectrum.clone());
210            state.previous_pupil = Some(context.state.pupil.values.clone());
211        }
212
213        if should_record && let Some(frame_diagnostics) = &context.diagnostics.frame_diagnostics {
214            state
215                .diagnostics
216                .frame_diagnostics
217                .extend(frame_diagnostics.iter().cloned());
218        }
219        if self.config.record_object_snapshots
220            && should_snapshot
221            && let (Some(amplitude), Some(phase)) = (
222                context.diagnostics.object_amplitude.clone(),
223                context.diagnostics.object_phase.clone(),
224            )
225        {
226            state
227                .object_snapshots
228                .push((context.iteration, amplitude, phase));
229        }
230        if self.config.record_pupil_snapshots
231            && should_snapshot
232            && let (Some(amplitude), Some(phase)) = (
233                context.diagnostics.pupil_amplitude.clone(),
234                context.diagnostics.pupil_phase.clone(),
235            )
236        {
237            state
238                .pupil_snapshots
239                .push((context.iteration, amplitude, phase));
240        }
241        Ok(CallbackAction::Continue)
242    }
243
244    fn on_finish(&mut self, _result: &ReconstructionResult) -> Result<()> {
245        Ok(())
246    }
247}
248
249fn cadence_matches(iteration: usize, every: usize) -> bool {
250    every > 0 && iteration.is_multiple_of(every)
251}
252
253fn relative_change(
254    previous: Option<&Array2<Complex64Proxy>>,
255    current: &Array2<Complex64Proxy>,
256) -> Option<f64> {
257    let previous = previous?;
258    if previous.shape() != current.shape() {
259        return None;
260    }
261    let mut difference = 0.0;
262    let mut reference = 0.0;
263    for (&previous, &current) in previous.as_slice().iter().zip(current.as_slice()) {
264        difference += (current - previous).norm_sqr();
265        reference += previous.norm_sqr();
266    }
267    Some(difference.sqrt() / reference.sqrt().max(f64::EPSILON))
268}
269
270fn median(values: &[f64]) -> Option<f64> {
271    if values.is_empty() {
272        return None;
273    }
274    let mut sorted = values.to_vec();
275    sorted.sort_by(|a, b| a.total_cmp(b));
276    let mid = sorted.len() / 2;
277    Some(if sorted.len().is_multiple_of(2) {
278        0.5 * (sorted[mid - 1] + sorted[mid])
279    } else {
280        sorted[mid]
281    })
282}
283
284#[cfg(test)]
285mod tests {
286    use std::panic::{AssertUnwindSafe, catch_unwind};
287
288    use super::*;
289
290    #[test]
291    fn poisoned_recorder_state_remains_resettable_and_readable() {
292        let recorder = DiagnosticRecorder::new(DiagnosticRecorderConfig::default());
293        let state = recorder.state.clone();
294        let result = catch_unwind(AssertUnwindSafe(|| {
295            let _guard = state.lock().unwrap();
296            panic!("intentional recorder-lock poison for test");
297        }));
298        assert!(result.is_err());
299
300        recorder.reset();
301        assert!(recorder.diagnostics().iteration_history.is_empty());
302        assert!(recorder.object_snapshots().is_empty());
303        assert!(recorder.pupil_snapshots().is_empty());
304    }
305}
fpm_rs/diagnostics/recorder.rs

fpm_rs/diagnostics/
recorder.rs
