fpm_rs/reconstruction/
checkpoint.rs

1use std::{
2    fs::File,
3    io::{BufReader, BufWriter},
4    path::Path,
5};
6
7use num_complex::Complex64;
8use serde::{Deserialize, Serialize};
9
10use crate::{
11    Array2, Result, diagnostics::ReconstructionHistory, error::Error,
12    measurements::MeasurementRead, model::Pupil,
13};
14
15use super::{AlgorithmAuxiliaryState, ReconstructionProblem, ReconstructionState};
16
17pub const CHECKPOINT_FORMAT_VERSION: u32 = 1;
18
19/// Serializable algorithm state used to resume a reconstruction exactly.
20#[derive(Clone, Debug, Serialize, Deserialize)]
21pub struct ReconstructionCheckpoint {
22    pub format_version: u32,
23    pub completed_iterations: usize,
24    pub object_spectrum: Array2<Complex64>,
25    pub pupil: Pupil,
26    /// Per-source `(row, column)` corrections in Fourier-grid pixels.
27    pub illumination_corrections: Option<Vec<(f64, f64)>>,
28    pub frame_gains: Option<Vec<f64>>,
29    pub background: Option<Vec<f64>>,
30    #[serde(default)]
31    pub algorithm_auxiliary: Option<AlgorithmAuxiliaryState>,
32    pub history: ReconstructionHistory,
33}
34
35impl ReconstructionCheckpoint {
36    pub fn capture(
37        completed_iterations: usize,
38        state: &ReconstructionState,
39        history: &ReconstructionHistory,
40    ) -> Self {
41        Self {
42            format_version: CHECKPOINT_FORMAT_VERSION,
43            completed_iterations,
44            object_spectrum: state.object_spectrum.clone(),
45            pupil: state.pupil.clone(),
46            illumination_corrections: state.illumination_corrections.clone(),
47            frame_gains: state.frame_gains.clone(),
48            background: state.background.clone(),
49            algorithm_auxiliary: state.algorithm_auxiliary.clone(),
50            history: history.clone(),
51        }
52    }
53
54    pub fn save(&self, path: impl AsRef<Path>) -> Result<()> {
55        self.validate()?;
56        let writer = BufWriter::new(File::create(path)?);
57        serde_json::to_writer(writer, self)?;
58        Ok(())
59    }
60
61    pub fn load(path: impl AsRef<Path>) -> Result<Self> {
62        let reader = BufReader::new(File::open(path)?);
63        let checkpoint: Self = serde_json::from_reader(reader)?;
64        checkpoint.validate()?;
65        Ok(checkpoint)
66    }
67
68    /// Loads a checkpoint and verifies all dimensions and calibration counts
69    /// against the problem that will resume it.
70    pub fn load_for_problem<M: MeasurementRead>(
71        path: impl AsRef<Path>,
72        problem: &ReconstructionProblem<M>,
73    ) -> Result<Self> {
74        let checkpoint = Self::load(path)?;
75        checkpoint.validate_for_problem(problem)?;
76        Ok(checkpoint)
77    }
78
79    /// Validates integrity that does not depend on a reconstruction problem.
80    pub fn validate(&self) -> Result<()> {
81        if self.format_version != CHECKPOINT_FORMAT_VERSION {
82            return Err(Error::InvalidParameter {
83                name: "checkpoint format_version",
84                reason: format!(
85                    "expected {CHECKPOINT_FORMAT_VERSION}, got {}",
86                    self.format_version
87                ),
88            });
89        }
90        if self.pupil.support.len() != self.pupil.values.len() {
91            return Err(Error::InvalidShape(
92                "checkpoint pupil support and values have different lengths".into(),
93            ));
94        }
95        if self
96            .object_spectrum
97            .as_slice()
98            .iter()
99            .chain(self.pupil.values.as_slice())
100            .any(|value| !value.re.is_finite() || !value.im.is_finite())
101        {
102            return Err(Error::InvalidModel(
103                "checkpoint contains non-finite complex values".into(),
104            ));
105        }
106        if self
107            .illumination_corrections
108            .as_ref()
109            .is_some_and(|values| {
110                values
111                    .iter()
112                    .any(|&(row, column)| !row.is_finite() || !column.is_finite())
113            })
114        {
115            return Err(Error::InvalidModel(
116                "checkpoint illumination corrections contain non-finite values".into(),
117            ));
118        }
119        if self.frame_gains.as_ref().is_some_and(|values| {
120            values
121                .iter()
122                .any(|value| !value.is_finite() || *value <= 0.0)
123        }) {
124            return Err(Error::InvalidModel(
125                "checkpoint frame gains must be finite and positive".into(),
126            ));
127        }
128        if self
129            .background
130            .as_ref()
131            .is_some_and(|values| values.iter().any(|value| !value.is_finite()))
132        {
133            return Err(Error::InvalidModel(
134                "checkpoint background contains non-finite values".into(),
135            ));
136        }
137        if self
138            .algorithm_auxiliary
139            .as_ref()
140            .is_some_and(|auxiliary| match auxiliary {
141                AlgorithmAuxiliaryState::Admm(admm) => {
142                    admm.auxiliary_fields.len() != admm.dual_fields.len()
143                        || admm
144                            .auxiliary_fields
145                            .iter()
146                            .chain(&admm.dual_fields)
147                            .any(|value| !value.re.is_finite() || !value.im.is_finite())
148                }
149            })
150        {
151            return Err(Error::InvalidModel(
152                "checkpoint algorithm auxiliary state is inconsistent or non-finite".into(),
153            ));
154        }
155        let records = &self.history.iterations;
156        if records.len() != self.completed_iterations
157            || records.iter().enumerate().any(|(index, record)| {
158                record.iteration != index + 1
159                    || !record.loss.is_finite()
160                    || !record.elapsed_seconds.is_finite()
161                    || record.elapsed_seconds < 0.0
162                    || record
163                        .admm_primal_residual_rms
164                        .is_some_and(|value| !value.is_finite() || value < 0.0)
165                    || record
166                        .admm_dual_residual_rms
167                        .is_some_and(|value| !value.is_finite() || value < 0.0)
168            })
169            || records
170                .windows(2)
171                .any(|pair| pair[1].elapsed_seconds < pair[0].elapsed_seconds)
172        {
173            return Err(Error::InvalidModel(
174                "checkpoint history is incomplete, non-finite, or non-monotonic".into(),
175            ));
176        }
177        Ok(())
178    }
179
180    /// Validates checkpoint dimensions and calibration variables for `problem`.
181    pub fn validate_for_problem<M: MeasurementRead>(
182        &self,
183        problem: &ReconstructionProblem<M>,
184    ) -> Result<()> {
185        problem.validate()?;
186        self.validate()?;
187        if self.object_spectrum.shape() != problem.model.reconstruction_shape {
188            return Err(Error::InvalidShape(format!(
189                "checkpoint spectrum shape {:?} differs from reconstruction shape {:?}",
190                self.object_spectrum.shape(),
191                problem.model.reconstruction_shape
192            )));
193        }
194        if self.pupil.shape() != problem.model.image_shape {
195            return Err(Error::InvalidShape(
196                "checkpoint pupil does not match the model image shape".into(),
197            ));
198        }
199        if self
200            .illumination_corrections
201            .as_ref()
202            .is_some_and(|values| values.len() != problem.model.source_count())
203        {
204            return Err(Error::InvalidModel(
205                "checkpoint illumination correction count does not match the model".into(),
206            ));
207        }
208        if self
209            .frame_gains
210            .as_ref()
211            .is_some_and(|values| values.len() != problem.model.frame_count())
212        {
213            return Err(Error::InvalidModel(
214                "checkpoint frame gain count does not match the model".into(),
215            ));
216        }
217        let image_len = problem.measurements.frame_len();
218        let stack_len = image_len
219            .checked_mul(problem.model.frame_count())
220            .ok_or_else(|| Error::InvalidShape("checkpoint stack length overflows".into()))?;
221        if self
222            .background
223            .as_ref()
224            .is_some_and(|values| values.len() != image_len && values.len() != stack_len)
225        {
226            return Err(Error::InvalidModel(
227                "checkpoint background dimensions do not match the model".into(),
228            ));
229        }
230        let mode_count = problem.model.multiplexing_matrix.as_ref().map_or_else(
231            || Ok(problem.model.frame_count()),
232            |matrix| {
233                matrix.iter().try_fold(0_usize, |count, row| {
234                    count.checked_add(row.len()).ok_or_else(|| {
235                        Error::InvalidShape("checkpoint source mode count overflows".into())
236                    })
237                })
238            },
239        )?;
240        let auxiliary_len = image_len
241            .checked_mul(mode_count)
242            .ok_or_else(|| Error::InvalidShape("checkpoint auxiliary length overflows".into()))?;
243        if self
244            .algorithm_auxiliary
245            .as_ref()
246            .is_some_and(|auxiliary| match auxiliary {
247                AlgorithmAuxiliaryState::Admm(admm) => admm.auxiliary_fields.len() != auxiliary_len,
248            })
249        {
250            return Err(Error::InvalidModel(
251                "checkpoint algorithm auxiliary dimensions do not match the model".into(),
252            ));
253        }
254        Ok(())
255    }
256}
fpm_rs/reconstruction/checkpoint.rs

fpm_rs/reconstruction/
checkpoint.rs
