fpm_rs/model/
pupil.rs

1use num_complex::Complex64;
2use serde::{Deserialize, Serialize};
3
4use crate::{Array2, Result, error::Error, experiment::Optics};
5
6use super::Sampling;
7
8#[derive(Clone, Debug, Serialize, Deserialize)]
9pub struct Pupil {
10    pub values: Array2<Complex64>,
11    pub support: Vec<bool>,
12}
13
14impl Pupil {
15    pub fn new(values: Array2<Complex64>, support: Vec<bool>) -> Result<Self> {
16        if support.len() != values.len() {
17            return Err(Error::InvalidShape(format!(
18                "pupil support length {} does not match pupil length {}",
19                support.len(),
20                values.len()
21            )));
22        }
23        if values
24            .as_slice()
25            .iter()
26            .any(|value| !value.re.is_finite() || !value.im.is_finite())
27        {
28            return Err(Error::InvalidModel(
29                "pupil contains non-finite complex values".into(),
30            ));
31        }
32        Ok(Self { values, support })
33    }
34
35    /// Builds a sampled circular pupil from microscope optics.
36    ///
37    /// Samples outside `sqrt(kx^2 + ky^2) <= 2 pi NA / lambda` are set to zero.
38    /// Inside the support, optional defocus contributes the paraxial phase
39    /// `-z * (kx^2 + ky^2) / (2 k0)`, where `k0 = 2 pi n / lambda`.
40    ///
41    /// If [`crate::experiment::PupilAberration`] is present, this method adds
42    /// crate-specific phase terms
43    /// `astigmatism * rho^2 * cos(2 theta)`,
44    /// `coma * (3 rho^3 - 2 rho) * cos(theta)`, and
45    /// `spherical * (6 rho^4 - 6 rho^2 + 1)`, with amplitude
46    /// `exp(-edge_apodization * rho^2)`. These coefficients are direct radian
47    /// weights, not normalized Zernike coefficients.
48    pub fn circular(shape: (usize, usize), sampling: &Sampling, optics: &Optics) -> Result<Self> {
49        sampling.validate()?;
50        optics.validate()?;
51        let cutoff = std::f64::consts::TAU * optics.objective_na / optics.wavelength;
52        let medium_k = optics.medium_wavenumber();
53        let mut values = Vec::with_capacity(shape.0 * shape.1);
54        let mut support = Vec::with_capacity(shape.0 * shape.1);
55        let aberration = optics.pupil_aberration.as_ref();
56        for row in 0..shape.0 {
57            let ky = (row as f64 - (shape.0 / 2) as f64) * sampling.dky;
58            for column in 0..shape.1 {
59                let kx = (column as f64 - (shape.1 / 2) as f64) * sampling.dkx;
60                let radius = kx.hypot(ky);
61                let inside = radius <= cutoff;
62                support.push(inside);
63                if !inside {
64                    values.push(Complex64::new(0.0, 0.0));
65                    continue;
66                }
67                let rho = if cutoff > 0.0 { radius / cutoff } else { 0.0 };
68                let theta = ky.atan2(kx);
69                let mut phase = 0.0;
70                if let Some(defocus) = optics.defocus_distance {
71                    phase -= defocus * (kx * kx + ky * ky) / (2.0 * medium_k);
72                }
73                let mut amplitude = 1.0;
74                if let Some(aberration) = aberration {
75                    phase += aberration.astigmatism * rho * rho * (2.0 * theta).cos();
76                    phase += aberration.coma * (3.0 * rho.powi(3) - 2.0 * rho) * theta.cos();
77                    phase += aberration.spherical * (6.0 * rho.powi(4) - 6.0 * rho * rho + 1.0);
78                    amplitude = (-aberration.edge_apodization * rho * rho).exp();
79                }
80                values.push(Complex64::from_polar(amplitude, phase));
81            }
82        }
83        Self::new(Array2::from_vec(shape, values)?, support)
84    }
85
86    pub fn shape(&self) -> (usize, usize) {
87        self.values.shape()
88    }
89
90    pub fn apply_support(&mut self) {
91        for (value, &inside) in self.values.as_mut_slice().iter_mut().zip(&self.support) {
92            if !inside {
93                *value = Complex64::new(0.0, 0.0);
94            }
95        }
96    }
97}
fpm_rs/model/pupil.rs

fpm_rs/model/
pupil.rs
