fpm_rs/reconstruction/
runner.rs

1use std::{collections::BTreeSet, sync::Arc, time::Instant};
2
3use crate::{
4    Result,
5    algorithms::ReconstructionAlgorithm,
6    backend::Backend,
7    callbacks::{Callback, CallbackAction, CallbackHook, StepContext},
8    complex,
9    diagnostics::{
10        DiagnosticRequest, Diagnostics, FrameDiagnosticRecord, IterationRecord,
11        RawFrameStatisticsRecord, ReconstructionHistory, StepDiagnostics,
12    },
13    error::Error,
14    measurements::MeasurementRead,
15    metrics::intensity::{compare_intensity_masked, intensity_statistics},
16    model::ForwardModel,
17};
18
19use super::{
20    Batch, ReconstructionCheckpoint, ReconstructionProblem, ReconstructionResult,
21    ReconstructionState, RunOptions, RuntimeInfo, state_object,
22};
23
24pub struct Runner<A> {
25    algorithm: A,
26    options: RunOptions,
27    callbacks: Vec<Box<dyn Callback>>,
28    initial_checkpoint: Option<ReconstructionCheckpoint>,
29    backend: Option<Arc<dyn Backend>>,
30}
31
32impl<A: ReconstructionAlgorithm> Runner<A> {
33    pub fn new(algorithm: A, options: RunOptions) -> Self {
34        Self {
35            algorithm,
36            options,
37            callbacks: Vec::new(),
38            initial_checkpoint: None,
39            backend: None,
40        }
41    }
42
43    pub fn with_callback(mut self, callback: Box<dyn Callback>) -> Self {
44        self.callbacks.push(callback);
45        self
46    }
47
48    pub fn with_callbacks(mut self, callbacks: Vec<Box<dyn Callback>>) -> Self {
49        self.callbacks.extend(callbacks);
50        self
51    }
52
53    pub fn resume_from(mut self, checkpoint: ReconstructionCheckpoint) -> Self {
54        self.initial_checkpoint = Some(checkpoint);
55        self
56    }
57
58    pub fn with_backend(mut self, backend: Arc<dyn Backend>) -> Self {
59        self.backend = Some(backend);
60        self
61    }
62
63    pub fn run<M: MeasurementRead>(
64        mut self,
65        problem: &ReconstructionProblem<M>,
66    ) -> Result<ReconstructionResult> {
67        self.algorithm.validate()?;
68        problem.validate()?;
69        self.algorithm.validate_problem(problem)?;
70        if self.options.batch_size == 0 {
71            return Err(Error::InvalidParameter {
72                name: "batch_size",
73                reason: "must be greater than zero".into(),
74            });
75        }
76        let started = Instant::now();
77        let algorithm_name = std::any::type_name::<A>()
78            .rsplit("::")
79            .next()
80            .unwrap_or("reconstruction algorithm")
81            .to_owned();
82        let (mut state, mut history, starting_iteration) =
83            if let Some(checkpoint) = &self.initial_checkpoint {
84                (
85                    if let Some(backend) = &self.backend {
86                        ReconstructionState::from_checkpoint_with_backend(
87                            problem,
88                            checkpoint,
89                            backend.clone(),
90                        )?
91                    } else {
92                        ReconstructionState::from_checkpoint(problem, checkpoint)?
93                    },
94                    checkpoint.history.clone(),
95                    checkpoint.completed_iterations,
96                )
97            } else {
98                (
99                    if let Some(backend) = &self.backend {
100                        self.algorithm
101                            .initialize_with_backend(problem, backend.clone())?
102                    } else {
103                        self.algorithm.initialize(problem)?
104                    },
105                    ReconstructionHistory::default(),
106                    0,
107                )
108            };
109        let previous_elapsed = history
110            .iterations
111            .last()
112            .map_or(0.0, |record| record.elapsed_seconds);
113        let start_requests =
114            callback_requests(&self.callbacks, CallbackHook::Start, starting_iteration);
115        let start_diagnostics = build_diagnostics(
116            problem,
117            &mut state,
118            &start_requests,
119            None,
120            None,
121            starting_iteration,
122        )?;
123        let start_context = StepContext {
124            iteration: starting_iteration,
125            frame_index: None,
126            batch_index: None,
127            state: &state,
128            diagnostics: &start_diagnostics,
129            history: &history,
130            model: &problem.model,
131            problem_name: problem.name.as_deref(),
132        };
133        let mut stopped_early = false;
134        for callback in &mut self.callbacks {
135            if callback.on_start(&start_context)? == CallbackAction::Stop {
136                stopped_early = true;
137            }
138        }
139        if !stopped_early {
140            for zero_based_iteration in starting_iteration..self.options.max_iterations {
141                let current_iteration = zero_based_iteration + 1;
142                let frame_requests =
143                    callback_requests(&self.callbacks, CallbackHook::FrameEnd, current_iteration);
144                let order = self
145                    .options
146                    .schedule
147                    .order_for_problem(problem, zero_based_iteration)?;
148                let mut iteration_step = StepDiagnostics::default();
149                for (batch_index, indices) in order.chunks(self.options.batch_size).enumerate() {
150                    let batch = Batch::new(indices.to_vec(), batch_index);
151                    let batch_step =
152                        self.algorithm
153                            .step(problem, &mut state, &batch, zero_based_iteration)?;
154                    state.object_real_space_cache = None;
155                    if self.options.enable_frame_callbacks {
156                        let batch_loss = batch_step.mean_loss();
157                        let mut diagnostics = build_diagnostics(
158                            problem,
159                            &mut state,
160                            &frame_requests,
161                            batch_loss,
162                            Some(&batch_step),
163                            current_iteration,
164                        )?;
165                        // A batch update completes every frame in the batch at
166                        // once. Emit one frame hook per completed frame, with
167                        // the frame's own natural loss and the shared post-batch
168                        // state. Expensive diagnostics are computed only once.
169                        for &frame in &batch.indices {
170                            if frame_requests.contains(&DiagnosticRequest::Loss) {
171                                diagnostics.loss = batch_step.per_frame_loss.get(&frame).copied();
172                            }
173                            let context = StepContext {
174                                iteration: current_iteration,
175                                frame_index: Some(frame),
176                                batch_index: Some(batch_index),
177                                state: &state,
178                                diagnostics: &diagnostics,
179                                history: &history,
180                                model: &problem.model,
181                                problem_name: problem.name.as_deref(),
182                            };
183                            for callback in &mut self.callbacks {
184                                if callback.on_frame_end(&context)? == CallbackAction::Stop {
185                                    stopped_early = true;
186                                }
187                            }
188                            if stopped_early {
189                                break;
190                            }
191                        }
192                    }
193                    iteration_step.merge(batch_step);
194                    if stopped_early {
195                        break;
196                    }
197                }
198                let loss = iteration_step.mean_loss().unwrap_or(f64::NAN);
199                history.iterations.push(IterationRecord {
200                    iteration: current_iteration,
201                    loss,
202                    elapsed_seconds: previous_elapsed + started.elapsed().as_secs_f64(),
203                    admm_primal_residual_rms: iteration_step.admm_primal_residual_rms(),
204                    admm_dual_residual_rms: iteration_step.admm_dual_residual_rms(),
205                });
206                let iteration_requests = callback_requests(
207                    &self.callbacks,
208                    CallbackHook::IterationEnd,
209                    current_iteration,
210                );
211                let diagnostics = build_diagnostics(
212                    problem,
213                    &mut state,
214                    &iteration_requests,
215                    Some(loss),
216                    Some(&iteration_step),
217                    current_iteration,
218                )?;
219                let context = StepContext {
220                    iteration: current_iteration,
221                    frame_index: None,
222                    batch_index: None,
223                    state: &state,
224                    diagnostics: &diagnostics,
225                    history: &history,
226                    model: &problem.model,
227                    problem_name: problem.name.as_deref(),
228                };
229                for callback in &mut self.callbacks {
230                    if callback.on_iteration_end(&context)? == CallbackAction::Stop {
231                        stopped_early = true;
232                    }
233                }
234                if stopped_early {
235                    break;
236                }
237            }
238        }
239
240        let runtime = RuntimeInfo {
241            elapsed_seconds: previous_elapsed + started.elapsed().as_secs_f64(),
242            completed_iterations: history
243                .iterations
244                .last()
245                .map_or(starting_iteration, |record| record.iteration),
246            stopped_early,
247            algorithm: algorithm_name,
248        };
249        let mut result = ReconstructionResult::from_state(&mut state, history, runtime)?;
250        if let Some(name) = &problem.name {
251            result.metadata.insert("problem_name".into(), name.clone());
252        }
253        for callback in &mut self.callbacks {
254            callback.on_finish(&result)?;
255        }
256        Ok(result)
257    }
258}
259
260fn callback_requests(
261    callbacks: &[Box<dyn Callback>],
262    hook: CallbackHook,
263    iteration: usize,
264) -> BTreeSet<DiagnosticRequest> {
265    callbacks
266        .iter()
267        .flat_map(|callback| callback.requires_for(hook, iteration))
268        .collect()
269}
270
271fn build_diagnostics<M: MeasurementRead>(
272    problem: &ReconstructionProblem<M>,
273    state: &mut ReconstructionState,
274    requests: &BTreeSet<DiagnosticRequest>,
275    natural_loss: Option<f64>,
276    step: Option<&StepDiagnostics>,
277    iteration: usize,
278) -> Result<Diagnostics> {
279    let mut diagnostics = Diagnostics::default();
280    if requests.contains(&DiagnosticRequest::Loss) {
281        diagnostics.loss = natural_loss;
282    }
283    if requests.contains(&DiagnosticRequest::RawFrameStats) {
284        let mut values = Vec::with_capacity(problem.model.frame_count());
285        for frame in 0..problem.model.frame_count() {
286            let measured = problem.measurements.frame(frame)?;
287            values.push(RawFrameStatisticsRecord {
288                frame_index: frame,
289                metrics: intensity_statistics(&measured, None)?,
290            });
291        }
292        diagnostics.raw_frame_stats = Some(values);
293    }
294    if requests.contains(&DiagnosticRequest::PerFrameError)
295        && let Some(step) = step
296    {
297        let mut values = vec![f64::NAN; problem.model.frame_count()];
298        for (&frame, &loss) in &step.per_frame_loss {
299            values[frame] = loss;
300        }
301        diagnostics.per_frame_error = Some(values);
302    }
303    if requests.contains(&DiagnosticRequest::FrameSummaries)
304        || requests.contains(&DiagnosticRequest::ResidualImages)
305        || (requests.contains(&DiagnosticRequest::PerFrameError) && step.is_none())
306    {
307        let diagnostic_model = model_with_state_calibration(problem, state)?;
308        let forward = ForwardModel::with_backend(&diagnostic_model, state.backend.clone())?;
309        let mut workspace = forward.workspace();
310        let mut predicted = vec![0.0; problem.model.image_shape.0 * problem.model.image_shape.1];
311        let mut frame_diagnostics = Vec::with_capacity(problem.model.frame_count());
312        let mut residual_images = if requests.contains(&DiagnosticRequest::ResidualImages) {
313            Some(Vec::with_capacity(problem.model.frame_count()))
314        } else {
315            None
316        };
317        let mut per_frame_error =
318            if requests.contains(&DiagnosticRequest::PerFrameError) && step.is_none() {
319                Some(Vec::with_capacity(problem.model.frame_count()))
320            } else {
321                None
322            };
323        for frame in 0..problem.model.frame_count() {
324            forward.forward_intensity_into(
325                &state.object_spectrum,
326                &state.pupil,
327                frame,
328                &mut workspace,
329                &mut predicted,
330            )?;
331            let measured = problem.measurements.frame(frame)?;
332            let mask = problem.measurements.frame_mask(frame)?;
333            let metadata = problem
334                .measurements
335                .frame_metadata()
336                .get(frame)
337                .cloned()
338                .unwrap_or_else(|| crate::measurements::FrameMetadata::new(frame));
339            if let Some(values) = per_frame_error.as_mut() {
340                let mut loss_sum = 0.0;
341                let mut valid_pixels = 0;
342                for pixel in 0..predicted.len() {
343                    if mask.is_none_or(|values| values[pixel] != 0) {
344                        let residual =
345                            predicted[pixel].max(0.0).sqrt() - measured[pixel].max(0.0).sqrt();
346                        loss_sum += residual * residual;
347                        valid_pixels += 1;
348                    }
349                }
350                values.push(if valid_pixels == 0 {
351                    0.0
352                } else {
353                    loss_sum / valid_pixels as f64
354                });
355            }
356            if let Some(images) = residual_images.as_mut() {
357                let mut residual: Vec<_> = predicted
358                    .iter()
359                    .zip(measured.iter())
360                    .map(|(&predicted, &measured)| predicted - measured)
361                    .collect();
362                if let Some(mask) = mask {
363                    for (value, &valid) in residual.iter_mut().zip(mask) {
364                        if valid == 0 {
365                            *value = 0.0;
366                        }
367                    }
368                }
369                images.push(crate::Array2::from_vec(
370                    problem.model.image_shape,
371                    residual,
372                )?);
373            }
374            if requests.contains(&DiagnosticRequest::FrameSummaries) {
375                let metrics = compare_intensity_masked(&measured, &predicted, mask, None)?;
376                frame_diagnostics.push(FrameDiagnosticRecord {
377                    iteration: Some(iteration),
378                    frame_index: frame,
379                    illumination_index: metadata.illumination_index.unwrap_or(frame),
380                    metrics,
381                });
382            }
383        }
384        if diagnostics.per_frame_error.is_none() {
385            diagnostics.per_frame_error = per_frame_error;
386        }
387        diagnostics.frame_diagnostics = Some(frame_diagnostics);
388        if let Some(images) = residual_images {
389            diagnostics.residual_images = Some(images);
390        }
391    }
392    if requests.contains(&DiagnosticRequest::ObjectAmplitude)
393        || requests.contains(&DiagnosticRequest::ObjectPhase)
394    {
395        let object = state_object(state)?;
396        if requests.contains(&DiagnosticRequest::ObjectAmplitude) {
397            diagnostics.object_amplitude = Some(complex::amplitude(&object));
398        }
399        if requests.contains(&DiagnosticRequest::ObjectPhase) {
400            diagnostics.object_phase = Some(complex::phase(&object));
401        }
402    }
403    if requests.contains(&DiagnosticRequest::Pupil) {
404        diagnostics.pupil_amplitude = Some(complex::amplitude(&state.pupil.values));
405        diagnostics.pupil_phase = Some(complex::phase(&state.pupil.values));
406    }
407    Ok(diagnostics)
408}
409
410fn model_with_state_calibration<M: MeasurementRead>(
411    problem: &ReconstructionProblem<M>,
412    state: &ReconstructionState,
413) -> Result<crate::model::ImagePlaneModel> {
414    let mut model = problem.model.clone();
415    model.frame_gains = state.frame_gains.clone();
416    model.background = state.background.clone();
417    if state.illumination_corrections.is_some() {
418        model.subpixel_offsets = Some(
419            (0..model.source_count())
420                .map(|source| state.effective_source_offset(&problem.model, source))
421                .collect::<Result<Vec<_>>>()?,
422        );
423    }
424    model.validate()?;
425    Ok(model)
426}
fpm_rs/reconstruction/runner.rs

fpm_rs/reconstruction/
runner.rs
