fpm_rs/reconstruction/
state.rs

1use num_complex::Complex64;
2use serde::{Deserialize, Serialize};
3use std::sync::Arc;
4
5use crate::{
6    Array2, Result,
7    backend::{Backend, CpuBackend, FftDirection},
8    error::Error,
9    measurements::MeasurementRead,
10    model::{FourierOffset, ImagePlaneModel, Pupil, fftshift_copy},
11};
12
13use super::{ReconstructionCheckpoint, ReconstructionProblem};
14
15#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
16pub struct AdmmAuxiliaryState {
17    pub auxiliary_fields: Vec<Complex64>,
18    pub dual_fields: Vec<Complex64>,
19}
20
21#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
22pub enum AlgorithmAuxiliaryState {
23    Admm(AdmmAuxiliaryState),
24}
25
26#[derive(Clone, Debug)]
27pub struct ReconstructionScratch {
28    pub patch: Vec<Complex64>,
29    pub exit_spectrum: Vec<Complex64>,
30    pub field: Vec<Complex64>,
31    pub projected_field: Vec<Complex64>,
32    pub projected_spectrum: Vec<Complex64>,
33    pub difference: Vec<Complex64>,
34    /// High-resolution accumulator used by mini-batch algorithms.
35    pub object_gradient: Vec<Complex64>,
36    pub regularization_field: Vec<Complex64>,
37    pub pupil_gradient: Vec<Complex64>,
38    pub calibration_reference: Vec<f64>,
39    pub illumination_gradient: Vec<(f64, f64)>,
40    pub illumination_curvature: Vec<(f64, f64)>,
41    pub illumination_weight: Vec<f64>,
42    /// Per-source low-resolution fields for an incoherently multiplexed frame.
43    pub multiplex_fields: Vec<Complex64>,
44    /// Matching pre-update object patches for multiplexed projection updates.
45    pub multiplex_patches: Vec<Complex64>,
46    pub multiplex_offsets: Vec<FourierOffset>,
47    pub column: Vec<Complex64>,
48}
49
50impl ReconstructionScratch {
51    fn new(low_shape: (usize, usize), high_shape: (usize, usize)) -> Self {
52        let low_len = low_shape.0 * low_shape.1;
53        Self {
54            patch: vec![Complex64::default(); low_len],
55            exit_spectrum: vec![Complex64::default(); low_len],
56            field: vec![Complex64::default(); low_len],
57            projected_field: vec![Complex64::default(); low_len],
58            projected_spectrum: vec![Complex64::default(); low_len],
59            difference: vec![Complex64::default(); low_len],
60            object_gradient: Vec::new(),
61            regularization_field: Vec::new(),
62            pupil_gradient: vec![Complex64::default(); low_len],
63            calibration_reference: vec![0.0; low_len],
64            illumination_gradient: Vec::new(),
65            illumination_curvature: Vec::new(),
66            illumination_weight: Vec::new(),
67            multiplex_fields: Vec::new(),
68            multiplex_patches: Vec::new(),
69            multiplex_offsets: Vec::new(),
70            column: vec![Complex64::default(); low_shape.0.max(high_shape.0)],
71        }
72    }
73}
74
75#[derive(Clone)]
76pub struct ReconstructionState {
77    pub object_spectrum: Array2<Complex64>,
78    pub object_real_space_cache: Option<Array2<Complex64>>,
79    pub pupil: Pupil,
80    /// Per-source `(row, column)` corrections in Fourier-grid pixels.
81    pub illumination_corrections: Option<Vec<(f64, f64)>>,
82    pub frame_gains: Option<Vec<f64>>,
83    pub background: Option<Vec<f64>>,
84    pub algorithm_auxiliary: Option<AlgorithmAuxiliaryState>,
85    pub scratch: ReconstructionScratch,
86    pub(crate) backend: Arc<dyn Backend>,
87}
88
89impl std::fmt::Debug for ReconstructionState {
90    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        formatter
92            .debug_struct("ReconstructionState")
93            .field("object_spectrum", &self.object_spectrum)
94            .field("pupil", &self.pupil)
95            .field("illumination_corrections", &self.illumination_corrections)
96            .field("frame_gains", &self.frame_gains)
97            .field("background", &self.background)
98            .field("algorithm_auxiliary", &self.algorithm_auxiliary)
99            .finish_non_exhaustive()
100    }
101}
102
103impl ReconstructionState {
104    pub fn effective_source_offset(
105        &self,
106        model: &ImagePlaneModel,
107        source: usize,
108    ) -> Result<FourierOffset> {
109        let base = model.source_offset(source)?;
110        let correction = match &self.illumination_corrections {
111            None => (0.0, 0.0),
112            Some(values) => *values.get(source).ok_or_else(|| {
113                Error::InvalidModel(
114                    "illumination correction count does not match source count".into(),
115                )
116            })?,
117        };
118        if !correction.0.is_finite() || !correction.1.is_finite() {
119            return Err(Error::InvalidModel(
120                "illumination corrections must be finite".into(),
121            ));
122        }
123        let effective = FourierOffset::new(base.row + correction.0, base.column + correction.1);
124        model.validate_source_offset(source, effective)?;
125        Ok(effective)
126    }
127
128    pub fn initialize<M: MeasurementRead>(problem: &ReconstructionProblem<M>) -> Result<Self> {
129        let backend: Arc<dyn Backend> = Arc::new(CpuBackend::new(
130            problem.model.image_shape,
131            problem.model.reconstruction_shape,
132        )?);
133        Self::initialize_with_backend(problem, backend)
134    }
135
136    pub fn initialize_with_backend<M: MeasurementRead>(
137        problem: &ReconstructionProblem<M>,
138        backend: Arc<dyn Backend>,
139    ) -> Result<Self> {
140        problem.validate()?;
141        let low_shape = problem.model.image_shape;
142        let high_shape = problem.model.reconstruction_shape;
143        let low_len = low_shape.0 * low_shape.1;
144        let mut average_amplitude = vec![0.0; low_len];
145        let mut amplitude_weight = vec![0.0; low_len];
146        let mut total_amplitude = 0.0;
147        let mut total_weight = 0.0;
148        for frame_index in 0..problem.measurements.frame_count() {
149            let frame = problem.measurements.frame(frame_index)?;
150            let frame_weight = problem.measurements.frame_weight(frame_index)?;
151            if frame_weight == 0.0 {
152                continue;
153            }
154            let mask = problem.measurements.frame_mask(frame_index)?;
155            let gain = problem.model.frame_gain(frame_index)?;
156            for (pixel, (average, &intensity)) in
157                average_amplitude.iter_mut().zip(frame.iter()).enumerate()
158            {
159                if mask.is_some_and(|values| values[pixel] == 0) {
160                    continue;
161                }
162                let background = problem.model.background_value(frame_index, pixel)?;
163                let amplitude = ((intensity - background) / gain).max(0.0).sqrt();
164                *average += frame_weight * amplitude;
165                amplitude_weight[pixel] += frame_weight;
166                total_amplitude += frame_weight * amplitude;
167                total_weight += frame_weight;
168            }
169        }
170        if total_weight == 0.0 {
171            return Err(Error::InvalidMeasurements(
172                "no positive-weight, unmasked measurements are available".into(),
173            ));
174        }
175        let fallback_amplitude = total_amplitude / total_weight;
176        for (value, &weight) in average_amplitude.iter_mut().zip(&amplitude_weight) {
177            *value = if weight > 0.0 {
178                *value / weight
179            } else {
180                fallback_amplitude
181            };
182        }
183        let mut object = vec![Complex64::default(); high_shape.0 * high_shape.1];
184        for row in 0..high_shape.0 {
185            let low_row = row * low_shape.0 / high_shape.0;
186            for column in 0..high_shape.1 {
187                let low_column = column * low_shape.1 / high_shape.1;
188                object[row * high_shape.1 + column] =
189                    Complex64::new(average_amplitude[low_row * low_shape.1 + low_column], 0.0);
190            }
191        }
192        let mut column = vec![Complex64::default(); low_shape.0.max(high_shape.0)];
193        backend.fft2(&mut object, high_shape, FftDirection::Forward, &mut column)?;
194        let mut centered = vec![Complex64::default(); object.len()];
195        fftshift_copy(&object, &mut centered, high_shape);
196        let object_spectrum = Array2::from_vec(high_shape, centered)?;
197        Ok(Self {
198            object_spectrum,
199            object_real_space_cache: None,
200            pupil: problem.model.pupil.clone(),
201            illumination_corrections: None,
202            frame_gains: problem.model.frame_gains.clone(),
203            background: problem.model.background.clone(),
204            algorithm_auxiliary: None,
205            scratch: ReconstructionScratch::new(low_shape, high_shape),
206            backend,
207        })
208    }
209
210    pub fn from_object<M: MeasurementRead>(
211        problem: &ReconstructionProblem<M>,
212        object: Array2<Complex64>,
213    ) -> Result<Self> {
214        if object.shape() != problem.model.reconstruction_shape {
215            return Err(Error::InvalidShape(format!(
216                "initial object shape {:?} differs from reconstruction shape {:?}",
217                object.shape(),
218                problem.model.reconstruction_shape
219            )));
220        }
221        let mut state = Self::initialize(problem)?;
222        let mut raw_spectrum = object.into_vec();
223        state.backend.fft2(
224            &mut raw_spectrum,
225            problem.model.reconstruction_shape,
226            FftDirection::Forward,
227            &mut state.scratch.column,
228        )?;
229        fftshift_copy(
230            &raw_spectrum,
231            state.object_spectrum.as_mut_slice(),
232            problem.model.reconstruction_shape,
233        );
234        Ok(state)
235    }
236
237    pub fn from_checkpoint<M: MeasurementRead>(
238        problem: &ReconstructionProblem<M>,
239        checkpoint: &ReconstructionCheckpoint,
240    ) -> Result<Self> {
241        let backend: Arc<dyn Backend> = Arc::new(CpuBackend::new(
242            problem.model.image_shape,
243            problem.model.reconstruction_shape,
244        )?);
245        Self::from_checkpoint_with_backend(problem, checkpoint, backend)
246    }
247
248    pub fn from_checkpoint_with_backend<M: MeasurementRead>(
249        problem: &ReconstructionProblem<M>,
250        checkpoint: &ReconstructionCheckpoint,
251        backend: Arc<dyn Backend>,
252    ) -> Result<Self> {
253        checkpoint.validate_for_problem(problem)?;
254        let low_shape = problem.model.image_shape;
255        let high_shape = problem.model.reconstruction_shape;
256        Ok(Self {
257            object_spectrum: checkpoint.object_spectrum.clone(),
258            object_real_space_cache: None,
259            pupil: checkpoint.pupil.clone(),
260            illumination_corrections: checkpoint.illumination_corrections.clone(),
261            frame_gains: checkpoint.frame_gains.clone(),
262            background: checkpoint.background.clone(),
263            algorithm_auxiliary: checkpoint.algorithm_auxiliary.clone(),
264            scratch: ReconstructionScratch::new(low_shape, high_shape),
265            backend,
266        })
267    }
268}
fpm_rs/reconstruction/state.rs

fpm_rs/reconstruction/
state.rs
