fpm_rs/simulation/
simulator.rs

1use num_complex::Complex64;
2use rand::{SeedableRng, rngs::StdRng};
3use rand_distr::{Distribution, Normal};
4
5use crate::{
6    Result,
7    backend::{Backend, CpuBackend, FftDirection},
8    error::Error,
9    measurements::{FrameMetadata, MeasurementStack},
10    model::{ForwardModel, ImagePlaneModel, fftshift_copy},
11};
12
13use super::{
14    CameraModel, IlluminationAcquisitionErrors, SimulationResult, SyntheticObject,
15    result::SimulationParameters,
16};
17
18pub struct Simulator {
19    true_model: ImagePlaneModel,
20    reconstruction_model: Option<ImagePlaneModel>,
21    object: Option<SyntheticObject>,
22    camera: Option<CameraModel>,
23    illumination_acquisition_errors: Option<IlluminationAcquisitionErrors>,
24    seed: u64,
25    ideal: bool,
26}
27
28impl Simulator {
29    /// Creates a simulator whose input model describes the true experiment.
30    ///
31    /// Use [`Self::reconstruction_model`] with a separately compiled assumed
32    /// optical model to simulate geometry or pupil mismatch.
33    pub fn new(true_model: ImagePlaneModel) -> Self {
34        Self {
35            true_model,
36            reconstruction_model: None,
37            object: None,
38            camera: None,
39            illumination_acquisition_errors: None,
40            seed: 0,
41            ideal: false,
42        }
43    }
44
45    pub fn ideal(model: ImagePlaneModel) -> Self {
46        Self {
47            ideal: true,
48            ..Self::new(model)
49        }
50    }
51
52    pub fn object(mut self, object: impl Into<SyntheticObject>) -> Self {
53        self.object = Some(object.into());
54        self
55    }
56
57    pub fn camera(mut self, camera: CameraModel) -> Self {
58        self.camera = Some(camera);
59        self.ideal = false;
60        self
61    }
62
63    /// Adds non-geometric illumination acquisition errors.
64    pub fn illumination_acquisition_errors(
65        mut self,
66        errors: IlluminationAcquisitionErrors,
67    ) -> Self {
68        self.illumination_acquisition_errors = Some(errors);
69        self.ideal = false;
70        self
71    }
72
73    /// Sets the optical model supplied to reconstruction after simulation.
74    ///
75    /// For an optical mismatch, compile this model and the true model from
76    /// separate experiment descriptions. When a camera model is present,
77    /// [`Self::simulate`] compiles the known linear camera response into this
78    /// returned reconstruction model.
79    pub fn reconstruction_model(mut self, model: ImagePlaneModel) -> Self {
80        self.reconstruction_model = Some(model);
81        self
82    }
83
84    pub fn seed(mut self, seed: u64) -> Self {
85        self.seed = seed;
86        self
87    }
88
89    /// Simulates measurements and returns detector counts plus a count-space
90    /// reconstruction model.
91    ///
92    /// This differs from serialized [`crate::configuration::SimulationConfiguration`],
93    /// which stores a strict optical reconstruction model and exposes
94    /// `reconstruction_model_for_counts()` for callers that load detector
95    /// counts later.
96    pub fn simulate(self) -> Result<SimulationResult> {
97        self.true_model.validate()?;
98        let object = self.object.ok_or(Error::InvalidParameter {
99            name: "object",
100            reason: "a ground-truth object is required".into(),
101        })?;
102        if object.shape() != self.true_model.reconstruction_shape {
103            return Err(Error::InvalidShape(format!(
104                "object shape {:?} differs from reconstruction shape {:?}",
105                object.shape(),
106                self.true_model.reconstruction_shape
107            )));
108        }
109        let mut reconstruction_model = self
110            .reconstruction_model
111            .unwrap_or_else(|| self.true_model.clone());
112        reconstruction_model.validate()?;
113        if reconstruction_model.image_shape != self.true_model.image_shape
114            || reconstruction_model.reconstruction_shape != self.true_model.reconstruction_shape
115            || reconstruction_model.frame_count() != self.true_model.frame_count()
116        {
117            return Err(Error::InvalidModel(
118                "true and reconstruction models must have matching image, object, and frame dimensions"
119                    .into(),
120            ));
121        }
122        if let Some(camera) = &self.camera {
123            reconstruction_model = camera.compile_reconstruction_model(reconstruction_model)?;
124        }
125        let mut true_model = self.true_model;
126        let mut rng = StdRng::seed_from_u64(self.seed);
127        let missing_frames = self
128            .illumination_acquisition_errors
129            .as_ref()
130            .map_or_else(Vec::new, |errors| errors.missing_frames.clone());
131        if let Some(errors) = &self.illumination_acquisition_errors {
132            apply_illumination_acquisition_errors(&mut true_model, errors, &mut rng)?;
133        }
134        true_model.validate()?;
135
136        let object_spectrum = object_spectrum(&object, &true_model)?;
137        let forward = ForwardModel::new(&true_model)?;
138        let image_len = true_model.image_shape.0 * true_model.image_shape.1;
139        let worker_count = std::thread::available_parallelism().map_or(1, |count| count.get());
140        let mut data = vec![0.0; image_len * true_model.frame_count()];
141        forward.forward_intensity_stack_into(
142            &object_spectrum,
143            &true_model.pupil,
144            &mut data,
145            worker_count,
146        )?;
147        // Keep camera processing serial so a seed produces the same
148        // random stream regardless of the machine's available parallelism.
149        for (frame, frame_data) in data.chunks_exact_mut(image_len).enumerate() {
150            if missing_frames.contains(&frame) {
151                // A failed source contributes no coherent field, but optical
152                // background remains present and is still measured by the camera.
153                for (pixel, value) in frame_data.iter_mut().enumerate() {
154                    *value = true_model.background_value(frame, pixel)?;
155                }
156            }
157            if let Some(camera) = &self.camera {
158                camera.measure_frame(frame_data, &mut rng)?;
159            }
160        }
161        let metadata = (0..true_model.frame_count())
162            .map(|frame| {
163                let mut metadata = FrameMetadata::new(frame);
164                if missing_frames.contains(&frame) {
165                    metadata.weight = 0.0;
166                }
167                metadata
168            })
169            .collect();
170        let measurements = MeasurementStack::from_vec(data, true_model.image_shape, metadata)?;
171        Ok(SimulationResult {
172            measurements,
173            ground_truth_object: object.field,
174            true_model: true_model.clone(),
175            reconstruction_model,
176            camera: self.camera,
177            illumination_acquisition_errors: self.illumination_acquisition_errors,
178            parameters: SimulationParameters {
179                ideal: self.ideal,
180                frame_count: true_model.frame_count(),
181                image_shape: true_model.image_shape,
182                missing_frames,
183            },
184            random_seed: self.seed,
185        })
186    }
187}
188
189fn object_spectrum(
190    object: &SyntheticObject,
191    model: &ImagePlaneModel,
192) -> Result<crate::Array2<Complex64>> {
193    let backend = CpuBackend::new(model.image_shape, model.reconstruction_shape)?;
194    let mut values = object.field.as_slice().to_vec();
195    let mut column =
196        vec![Complex64::default(); model.image_shape.0.max(model.reconstruction_shape.0)];
197    backend.fft2(
198        &mut values,
199        model.reconstruction_shape,
200        FftDirection::Forward,
201        &mut column,
202    )?;
203    let mut centered = vec![Complex64::default(); values.len()];
204    fftshift_copy(&values, &mut centered, model.reconstruction_shape);
205    crate::Array2::from_vec(model.reconstruction_shape, centered)
206}
207
208fn apply_illumination_acquisition_errors(
209    model: &mut ImagePlaneModel,
210    errors: &IlluminationAcquisitionErrors,
211    rng: &mut StdRng,
212) -> Result<()> {
213    if !errors.frame_gain_relative_std.is_finite() || errors.frame_gain_relative_std < 0.0 {
214        return Err(Error::InvalidParameter {
215            name: "frame_gain_relative_std",
216            reason: "must be finite and non-negative".into(),
217        });
218    }
219    let frame_count = model.frame_count();
220    let source_count = model.source_count();
221    let multiplexed = model.is_multiplexed();
222    let mut sorted_missing = errors.missing_frames.clone();
223    sorted_missing.sort_unstable();
224    if sorted_missing.iter().any(|&frame| frame >= frame_count)
225        || sorted_missing.windows(2).any(|pair| pair[0] == pair[1])
226    {
227        return Err(Error::InvalidParameter {
228            name: "missing_frames",
229            reason: format!("must contain unique frame indices below {frame_count}"),
230        });
231    }
232    if let Some(permutation) = &errors.source_permutation {
233        let mut sorted = permutation.clone();
234        sorted.sort_unstable();
235        if permutation.len() != source_count
236            || sorted
237                .iter()
238                .enumerate()
239                .any(|(expected, &actual)| expected != actual)
240        {
241            return Err(Error::InvalidParameter {
242                name: "source_permutation",
243                reason: format!("must be a permutation of 0..{source_count}"),
244            });
245        }
246    }
247    if let Some(permutation) = &errors.source_permutation {
248        model.k_vectors = permutation
249            .iter()
250            .map(|&index| model.k_vectors[index])
251            .collect();
252        model.crop_indices.crops = permutation
253            .iter()
254            .map(|&index| model.crop_indices.crops[index])
255            .collect();
256        model.subpixel_offsets = model
257            .subpixel_offsets
258            .as_ref()
259            .map(|offsets| permutation.iter().map(|&index| offsets[index]).collect());
260        if !multiplexed && let Some(gains) = &model.frame_gains {
261            model.frame_gains = Some(permutation.iter().map(|&index| gains[index]).collect());
262        }
263        let image_len = model.image_shape.0 * model.image_shape.1;
264        if !multiplexed
265            && let Some(background) = &model.background
266            && background.len() == image_len * frame_count
267        {
268            let mut reordered = Vec::with_capacity(background.len());
269            for &index in permutation {
270                let start = index * image_len;
271                reordered.extend_from_slice(&background[start..start + image_len]);
272            }
273            model.background = Some(reordered);
274        }
275    }
276    if errors.frame_gain_relative_std > 0.0 {
277        let distribution = Normal::new(1.0, errors.frame_gain_relative_std)
278            .map_err(|error| Error::Numerical(error.to_string()))?;
279        let gains = model
280            .frame_gains
281            .get_or_insert_with(|| vec![1.0; frame_count]);
282        for gain in gains {
283            *gain *= distribution.sample(rng).max(0.01);
284        }
285    }
286    Ok(())
287}
fpm_rs/simulation/simulator.rs

fpm_rs/simulation/
simulator.rs
