Skip to main content

fpm_rs/algorithms/
fpie.rs

1use crate::{
2    Result,
3    algorithms::{NoIterationMetrics, StepOutput, objective::LossType},
4    error::Error,
5    measurements::MeasurementRead,
6    reconstruction::{Batch, ReconstructionProblem, ReconstructionState},
7};
8
9use super::{
10    ReconstructionAlgorithm,
11    common::{ObjectDenominator, UpdateConfiguration, projection_update},
12};
13
14/// Regularized ptychographic iterative-engine reconstruction adapted to FPM.
15///
16/// # Method
17///
18/// `Fpie` performs the same detector-amplitude projection as
19/// [`super::AlternatingProjection`], but preconditions each object correction
20/// with the rPIE denominator
21/// `(1 - stability) * |pupil|^2 + stability * max(|pupil|^2)`. This blends
22/// local inverse-pupil weighting with a global power bound, reducing unstable
23/// updates where the pupil transfer is weak. The corrected low-resolution
24/// spectrum is inserted into the corresponding overlapping patch of the
25/// high-resolution object spectrum.
26///
27/// This crate adapts the rPIE update, originally formulated for scanned
28/// ptychography, to image-plane Fourier ptychography.
29///
30/// # Reference
31///
32/// [A. Maiden, D. Johnson, and P. Li, “Further improvements to the
33/// ptychographical iterative engine” (2017)](https://doi.org/10.1364/OPTICA.4.000736),
34/// *Optica* **4**(7), 736–745.
35#[derive(Clone, Debug)]
36pub struct Fpie {
37    /// Number of complete passes through the acquisition schedule.
38    pub iterations: usize,
39    /// Relaxation factor applied to each object-spectrum correction.
40    pub object_step: f64,
41    /// Blend between local pupil power (`0`) and maximum pupil power (`1`) in
42    /// the rPIE denominator.
43    pub stability: f64,
44    /// Number of measured frames supplied to each reconstruction step.
45    pub batch_size: usize,
46    /// Positive numerical floor added to the rPIE denominator.
47    pub epsilon: f64,
48    /// Loss used for diagnostics; the projection itself always enforces the
49    /// measured amplitude.
50    pub loss_type: LossType,
51}
52
53impl Default for Fpie {
54    fn default() -> Self {
55        Self {
56            iterations: 50,
57            object_step: 0.8,
58            stability: 0.1,
59            batch_size: 1,
60            epsilon: 1e-10,
61            loss_type: LossType::AmplitudeMse,
62        }
63    }
64}
65
66impl Fpie {
67    /// Sets the number of complete acquisition-schedule passes; validation requires non-zero.
68    pub fn iterations(mut self, iterations: usize) -> Self {
69        self.iterations = iterations;
70        self
71    }
72
73    /// Sets the finite positive relaxation applied to object-spectrum corrections.
74    pub fn object_step(mut self, step: f64) -> Self {
75        self.object_step = step;
76        self
77    }
78
79    /// Sets the rPIE pupil-power blend after clamping the value to `[0, 1]`.
80    pub fn stability(mut self, stability: f64) -> Self {
81        self.stability = stability.clamp(0.0, 1.0);
82        self
83    }
84}
85
86impl ReconstructionAlgorithm for Fpie {
87    type IterationMetrics = NoIterationMetrics;
88
89    fn validate(&self) -> Result<()> {
90        if !self.object_step.is_finite() || self.object_step <= 0.0 {
91            return Err(Error::InvalidParameter {
92                name: "object_step",
93                reason: "must be finite and positive".into(),
94            });
95        }
96        if !self.stability.is_finite() || !(0.0..=1.0).contains(&self.stability) {
97            return Err(Error::InvalidParameter {
98                name: "stability",
99                reason: "must be finite and between zero and one".into(),
100            });
101        }
102        if !self.epsilon.is_finite() || self.epsilon <= 0.0 || self.batch_size == 0 {
103            return Err(Error::InvalidParameter {
104                name: "epsilon/batch_size",
105                reason: "epsilon must be positive and batch size non-zero".into(),
106            });
107        }
108        Ok(())
109    }
110
111    fn step<M: MeasurementRead>(
112        &mut self,
113        problem: &ReconstructionProblem<M>,
114        state: &mut ReconstructionState,
115        batch: &Batch,
116        _iteration: usize,
117    ) -> Result<StepOutput<Self::IterationMetrics>> {
118        Ok(projection_update(
119            problem,
120            state,
121            batch,
122            UpdateConfiguration {
123                object_step: self.object_step,
124                pupil_step: None,
125                epsilon: self.epsilon,
126                loss_type: self.loss_type,
127                object_denominator: ObjectDenominator::Rpie(self.stability),
128                constrain_pupil: true,
129                gain_update: None,
130                background_update: None,
131                momentum: None,
132            },
133        )?
134        .into())
135    }
136
137    fn iterations(&self) -> usize {
138        self.iterations
139    }
140
141    fn batch_size(&self) -> usize {
142        self.batch_size
143    }
144}