fpm_rs/reconstruction/
result.rs

1use std::{
2    collections::BTreeMap,
3    fs::File,
4    io::{BufReader, BufWriter},
5    path::Path,
6};
7
8use image::{GrayImage, Luma};
9use num_complex::Complex64;
10use serde::{Deserialize, Serialize};
11
12use crate::{
13    Array2, Result,
14    backend::FftDirection,
15    complex,
16    diagnostics::ReconstructionHistory,
17    error::Error,
18    model::{Pupil, ifftshift_copy},
19};
20
21use super::ReconstructionState;
22
23pub const RESULT_BUNDLE_FORMAT_VERSION: u32 = 1;
24
25#[derive(Clone, Debug, Default, Serialize, Deserialize)]
26pub struct RuntimeInfo {
27    pub elapsed_seconds: f64,
28    pub completed_iterations: usize,
29    pub stopped_early: bool,
30    pub algorithm: String,
31}
32
33#[derive(Clone, Debug, Serialize, Deserialize)]
34pub struct ReconstructionResult {
35    pub object: Array2<Complex64>,
36    pub amplitude: Array2<f64>,
37    pub phase: Array2<f64>,
38    pub object_spectrum: Array2<Complex64>,
39    pub recovered_pupil: Pupil,
40    /// Per-source `(row, column)` corrections in Fourier-grid pixels.
41    pub calibrated_illumination: Option<Vec<(f64, f64)>>,
42    pub recovered_frame_gains: Option<Vec<f64>>,
43    pub recovered_background: Option<Vec<f64>>,
44    pub history: ReconstructionHistory,
45    pub diagnostics: BTreeMap<String, f64>,
46    pub runtime: RuntimeInfo,
47    pub metadata: BTreeMap<String, String>,
48}
49
50#[derive(Serialize, Deserialize)]
51struct ReconstructionResultBundle {
52    format_version: u32,
53    result: ReconstructionResult,
54}
55
56impl ReconstructionResult {
57    pub(crate) fn from_state(
58        state: &mut ReconstructionState,
59        history: ReconstructionHistory,
60        runtime: RuntimeInfo,
61    ) -> Result<Self> {
62        let object = state_object(state)?;
63        let amplitude = complex::amplitude(&object);
64        let phase = complex::phase(&object);
65        let mut diagnostics = BTreeMap::new();
66        if let Some(loss) = history.final_loss() {
67            diagnostics.insert("final_loss".into(), loss);
68        }
69        if let Some(record) = history.iterations.last() {
70            if let Some(residual) = record.admm_primal_residual_rms {
71                diagnostics.insert("final_admm_primal_residual_rms".into(), residual);
72            }
73            if let Some(residual) = record.admm_dual_residual_rms {
74                diagnostics.insert("final_admm_dual_residual_rms".into(), residual);
75            }
76        }
77        Ok(Self {
78            object,
79            amplitude,
80            phase,
81            object_spectrum: state.object_spectrum.clone(),
82            recovered_pupil: state.pupil.clone(),
83            calibrated_illumination: state.illumination_corrections.clone(),
84            recovered_frame_gains: state.frame_gains.clone(),
85            recovered_background: state.background.clone(),
86            history,
87            diagnostics,
88            runtime,
89            metadata: BTreeMap::new(),
90        })
91    }
92
93    pub fn save_amplitude(&self, path: impl AsRef<Path>) -> Result<()> {
94        save_grayscale(&self.amplitude, path, false)
95    }
96
97    pub fn save_phase(&self, path: impl AsRef<Path>) -> Result<()> {
98        save_grayscale(&self.phase, path, true)
99    }
100
101    pub fn save_complex_object(&self, path: impl AsRef<Path>) -> Result<()> {
102        let writer = BufWriter::new(File::create(path)?);
103        serde_json::to_writer(writer, &self.object)?;
104        Ok(())
105    }
106
107    pub fn save_pupil(&self, path: impl AsRef<Path>) -> Result<()> {
108        let writer = BufWriter::new(File::create(path)?);
109        serde_json::to_writer(writer, &self.recovered_pupil)?;
110        Ok(())
111    }
112
113    pub fn save_loss_csv(&self, path: impl AsRef<Path>) -> Result<()> {
114        let mut writer = csv::Writer::from_path(path)?;
115        writer.write_record([
116            "iteration",
117            "loss",
118            "elapsed_seconds",
119            "admm_primal_residual_rms",
120            "admm_dual_residual_rms",
121        ])?;
122        for record in &self.history.iterations {
123            writer.serialize((
124                record.iteration,
125                record.loss,
126                record.elapsed_seconds,
127                record.admm_primal_residual_rms,
128                record.admm_dual_residual_rms,
129            ))?;
130        }
131        writer.flush()?;
132        Ok(())
133    }
134
135    /// Saves every result array, calibration value, diagnostic, history entry,
136    /// runtime field, and metadata value in a versioned JSON bundle.
137    pub fn save_bundle(&self, path: impl AsRef<Path>) -> Result<()> {
138        self.validate()?;
139        let writer = BufWriter::new(File::create(path)?);
140        serde_json::to_writer(
141            writer,
142            &ReconstructionResultBundle {
143                format_version: RESULT_BUNDLE_FORMAT_VERSION,
144                result: self.clone(),
145            },
146        )?;
147        Ok(())
148    }
149
150    pub fn load_bundle(path: impl AsRef<Path>) -> Result<Self> {
151        let reader = BufReader::new(File::open(path)?);
152        let bundle: ReconstructionResultBundle = serde_json::from_reader(reader)?;
153        if bundle.format_version != RESULT_BUNDLE_FORMAT_VERSION {
154            return Err(Error::InvalidParameter {
155                name: "result bundle format_version",
156                reason: format!(
157                    "expected {RESULT_BUNDLE_FORMAT_VERSION}, got {}",
158                    bundle.format_version
159                ),
160            });
161        }
162        bundle.result.validate()?;
163        Ok(bundle.result)
164    }
165
166    pub fn validate(&self) -> Result<()> {
167        let shape = self.object.shape();
168        if self.amplitude.shape() != shape
169            || self.phase.shape() != shape
170            || self.object_spectrum.shape() != shape
171        {
172            return Err(Error::InvalidShape(
173                "result object, amplitude, phase, and spectrum shapes must match".into(),
174            ));
175        }
176        if self.recovered_pupil.support.len() != self.recovered_pupil.values.len() {
177            return Err(Error::InvalidShape(
178                "result pupil support and values have different lengths".into(),
179            ));
180        }
181        if self
182            .object
183            .as_slice()
184            .iter()
185            .chain(self.object_spectrum.as_slice())
186            .chain(self.recovered_pupil.values.as_slice())
187            .any(|value| !value.re.is_finite() || !value.im.is_finite())
188            || self
189                .amplitude
190                .as_slice()
191                .iter()
192                .chain(self.phase.as_slice())
193                .any(|value| !value.is_finite())
194        {
195            return Err(Error::InvalidModel(
196                "result arrays contain non-finite values".into(),
197            ));
198        }
199        if self.calibrated_illumination.as_ref().is_some_and(|values| {
200            values
201                .iter()
202                .any(|&(row, column)| !row.is_finite() || !column.is_finite())
203        }) || self.recovered_frame_gains.as_ref().is_some_and(|values| {
204            values
205                .iter()
206                .any(|value| !value.is_finite() || *value <= 0.0)
207        }) || self
208            .recovered_background
209            .as_ref()
210            .is_some_and(|values| values.iter().any(|value| !value.is_finite()))
211        {
212            return Err(Error::InvalidModel(
213                "result calibration values are invalid".into(),
214            ));
215        }
216        if self.diagnostics.values().any(|value| !value.is_finite()) {
217            return Err(Error::InvalidModel(
218                "result diagnostics contain non-finite values".into(),
219            ));
220        }
221        if !self.runtime.elapsed_seconds.is_finite()
222            || self.runtime.elapsed_seconds < 0.0
223            || self.runtime.algorithm.is_empty()
224            || self.runtime.completed_iterations != self.history.iterations.len()
225            || self
226                .history
227                .iterations
228                .iter()
229                .enumerate()
230                .any(|(index, record)| {
231                    record.iteration != index + 1
232                        || !record.loss.is_finite()
233                        || !record.elapsed_seconds.is_finite()
234                        || record.elapsed_seconds < 0.0
235                        || record
236                            .admm_primal_residual_rms
237                            .is_some_and(|value| !value.is_finite() || value < 0.0)
238                        || record
239                            .admm_dual_residual_rms
240                            .is_some_and(|value| !value.is_finite() || value < 0.0)
241                })
242            || self
243                .history
244                .iterations
245                .windows(2)
246                .any(|pair| pair[1].elapsed_seconds < pair[0].elapsed_seconds)
247            || self
248                .history
249                .iterations
250                .last()
251                .is_some_and(|record| record.elapsed_seconds > self.runtime.elapsed_seconds)
252        {
253            return Err(Error::InvalidModel(
254                "result runtime and history are inconsistent".into(),
255            ));
256        }
257        Ok(())
258    }
259}
260
261pub(crate) fn state_object(state: &mut ReconstructionState) -> Result<Array2<Complex64>> {
262    if let Some(cached) = &state.object_real_space_cache {
263        return Ok(cached.clone());
264    }
265    let shape = state.object_spectrum.shape();
266    let mut unshifted = vec![Complex64::default(); state.object_spectrum.len()];
267    ifftshift_copy(state.object_spectrum.as_slice(), &mut unshifted, shape);
268    state.backend.fft2(
269        &mut unshifted,
270        shape,
271        FftDirection::Inverse,
272        &mut state.scratch.column,
273    )?;
274    let object = Array2::from_vec(shape, unshifted)?;
275    state.object_real_space_cache = Some(object.clone());
276    Ok(object)
277}
278
279pub(crate) fn save_grayscale(
280    values: &Array2<f64>,
281    path: impl AsRef<Path>,
282    phase: bool,
283) -> Result<()> {
284    let range = if phase {
285        (-std::f64::consts::PI, std::f64::consts::PI)
286    } else {
287        let minimum = values
288            .as_slice()
289            .iter()
290            .copied()
291            .fold(f64::INFINITY, f64::min);
292        let maximum = values
293            .as_slice()
294            .iter()
295            .copied()
296            .fold(f64::NEG_INFINITY, f64::max);
297        (minimum, maximum)
298    };
299    save_grayscale_with_range(values, path, range)
300}
301
302pub(crate) fn save_signed_grayscale(values: &Array2<f64>, path: impl AsRef<Path>) -> Result<()> {
303    let maximum_absolute = values
304        .as_slice()
305        .iter()
306        .map(|value| value.abs())
307        .fold(0.0, f64::max)
308        .max(f64::EPSILON);
309    save_grayscale_with_range(values, path, (-maximum_absolute, maximum_absolute))
310}
311
312fn save_grayscale_with_range(
313    values: &Array2<f64>,
314    path: impl AsRef<Path>,
315    (minimum, maximum): (f64, f64),
316) -> Result<()> {
317    let width = u32::try_from(values.width()).map_err(|_| {
318        Error::InvalidShape("image width does not fit the PNG dimension type".into())
319    })?;
320    let height = u32::try_from(values.height()).map_err(|_| {
321        Error::InvalidShape("image height does not fit the PNG dimension type".into())
322    })?;
323    let range = (maximum - minimum).max(f64::EPSILON);
324    let mut image = GrayImage::new(width, height);
325    for row in 0..values.height() {
326        for column in 0..values.width() {
327            let normalized = ((values[(row, column)] - minimum) / range).clamp(0.0, 1.0);
328            image.put_pixel(
329                column as u32,
330                row as u32,
331                Luma([(normalized * 255.0).round() as u8]),
332            );
333        }
334    }
335    image.save(path)?;
336    Ok(())
337}
fpm_rs/reconstruction/result.rs

fpm_rs/reconstruction/
result.rs
