fpm_rs/reconstruction/
schedule.rs

1use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
2use serde::{Deserialize, Serialize};
3
4use crate::{Result, experiment::KVector, measurements::MeasurementRead, model::ImagePlaneModel};
5
6use super::ReconstructionProblem;
7
8#[derive(Clone, Debug, Default, Serialize, Deserialize)]
9pub enum FrameSchedule {
10    #[default]
11    Sequential,
12    BrightfieldFirst,
13    SpiralOut,
14    RandomShuffle {
15        seed: u64,
16    },
17    /// Highest empirical shot-noise SNR first. This requires measurements and
18    /// is applied by [`Self::order_for_problem`]; [`Self::order`] retains
19    /// sequential order because a compiled model alone contains no signal data.
20    SnrWeighted,
21}
22
23impl FrameSchedule {
24    /// Returns a model-only frame order.
25    ///
26    /// [`Self::SnrWeighted`] cannot be estimated from model geometry, so this
27    /// method returns sequential order for that variant. Reconstruction runners
28    /// use [`Self::order_for_problem`] and therefore have access to measurements.
29    pub fn order(&self, model: &ImagePlaneModel, iteration: usize) -> Vec<usize> {
30        let mut order: Vec<usize> = (0..model.frame_count()).collect();
31        match self {
32            Self::Sequential | Self::SnrWeighted => {}
33            Self::BrightfieldFirst => order.sort_by(|&left, &right| {
34                let left_vector = frame_vector(model, left);
35                let right_vector = frame_vector(model, right);
36                left_vector
37                    .kx
38                    .hypot(left_vector.ky)
39                    .total_cmp(&right_vector.kx.hypot(right_vector.ky))
40            }),
41            Self::SpiralOut => order.sort_by(|&left, &right| {
42                let left_vector = frame_vector(model, left);
43                let right_vector = frame_vector(model, right);
44                left_vector
45                    .kx
46                    .hypot(left_vector.ky)
47                    .total_cmp(&right_vector.kx.hypot(right_vector.ky))
48                    .then_with(|| {
49                        left_vector
50                            .ky
51                            .atan2(left_vector.kx)
52                            .rem_euclid(std::f64::consts::TAU)
53                            .total_cmp(
54                                &right_vector
55                                    .ky
56                                    .atan2(right_vector.kx)
57                                    .rem_euclid(std::f64::consts::TAU),
58                            )
59                    })
60            }),
61            Self::RandomShuffle { seed } => {
62                let mut rng = StdRng::seed_from_u64(
63                    seed.wrapping_add((iteration as u64).wrapping_mul(0x9e3779b97f4a7c15)),
64                );
65                order.shuffle(&mut rng);
66            }
67        }
68        order
69    }
70
71    /// Returns a frame order using both the model and measured intensities.
72    ///
73    /// The SNR-weighted schedule ranks frames by
74    /// `frame_weight * sqrt(mean(max(measured - known_background, 0)))` over
75    /// unmasked pixels. This is an empirical Poisson shot-noise proxy, not a
76    /// camera-noise calibration. Zero-weight or fully masked frames are placed
77    /// last, and equal scores retain ascending frame-index order.
78    pub fn order_for_problem<M: MeasurementRead>(
79        &self,
80        problem: &ReconstructionProblem<M>,
81        iteration: usize,
82    ) -> Result<Vec<usize>> {
83        if !matches!(self, Self::SnrWeighted) {
84            return Ok(self.order(&problem.model, iteration));
85        }
86        let mut scored = Vec::with_capacity(problem.model.frame_count());
87        for frame in 0..problem.model.frame_count() {
88            scored.push((frame, frame_snr_score(problem, frame)?));
89        }
90        scored.sort_by(|&(left_frame, left_score), &(right_frame, right_score)| {
91            right_score
92                .total_cmp(&left_score)
93                .then_with(|| left_frame.cmp(&right_frame))
94        });
95        Ok(scored.into_iter().map(|(frame, _)| frame).collect())
96    }
97}
98
99fn frame_snr_score<M: MeasurementRead>(
100    problem: &ReconstructionProblem<M>,
101    frame: usize,
102) -> Result<f64> {
103    let weight = problem.measurements.frame_weight(frame)?;
104    if weight == 0.0 {
105        return Ok(f64::NEG_INFINITY);
106    }
107    let measured = problem.measurements.frame(frame)?;
108    let mask = problem.measurements.frame_mask(frame)?;
109    let mut signal_sum = 0.0;
110    let mut valid_pixels = 0;
111    for (pixel, &value) in measured.iter().enumerate() {
112        if mask.is_none_or(|values| values[pixel] != 0) {
113            let background = problem.model.background_value(frame, pixel)?;
114            signal_sum += (value - background).max(0.0);
115            valid_pixels += 1;
116        }
117    }
118    if valid_pixels == 0 {
119        Ok(f64::NEG_INFINITY)
120    } else {
121        Ok(weight * (signal_sum / valid_pixels as f64).sqrt())
122    }
123}
124
125fn frame_vector(model: &ImagePlaneModel, frame: usize) -> KVector {
126    if let Some(matrix) = &model.multiplexing_matrix {
127        let mut vector = KVector::default();
128        let mut total_weight = 0.0;
129        for &(source, weight) in &matrix[frame] {
130            vector.kx += weight * model.k_vectors[source].kx;
131            vector.ky += weight * model.k_vectors[source].ky;
132            total_weight += weight;
133        }
134        vector.kx /= total_weight;
135        vector.ky /= total_weight;
136        vector
137    } else {
138        model.k_vectors[frame]
139    }
140}
fpm_rs/reconstruction/schedule.rs

fpm_rs/reconstruction/
schedule.rs
