rusty_psf/core/
utils.rs

1use ndarray::{Array1, Array2, Array3, Axis};
2use crate::types::OpticalParameters;
3use num_traits::Float;
4
5/// Generate a centered coordinate grid.
6/// 
7/// Creates a 1D array of coordinates centered around zero with specified size and pixel spacing.
8/// 
9/// # Arguments
10/// * `size` - Number of points in the grid
11/// * `pixel_size` - Physical size of each pixel
12/// 
13/// # Returns
14/// A 1D array of evenly spaced coordinates
15pub fn make_grid<T: Float>(size: usize, pixel_size: T) -> Array1<T> {
16    let half = T::from(size).unwrap() / (T::from(2).unwrap());
17    let range = Array1::linspace(-half * pixel_size, (half - T::one()) * pixel_size, size);
18    range
19}
20
21/// Generate a radial coordinate grid.
22/// 
23/// Creates a 2D array of radial coordinates (distance from center) for each point.
24/// 
25/// # Arguments
26/// * `size` - Size of the grid (will be size × size)
27/// * `pixel_size` - Physical size of each pixel
28/// 
29/// # Returns
30/// A 2D array containing the radial distance from the center for each point
31pub fn make_radial_grid(size: usize, pixel_size: f64) -> Array2<f64> {
32    let coords = make_grid(size, pixel_size);
33    let mut radial = Array2::zeros((size, size));
34    
35    for i in 0..size {
36        for j in 0..size {
37            radial[[i, j]] = (coords[i].powi(2) + coords[j].powi(2)).sqrt();
38        }
39    }
40    radial
41}
42
43/// Normalize PSF to unit sum.
44/// 
45/// Normalizes a PSF array so that the sum of all elements equals 1.0.
46/// This ensures energy conservation in the PSF.
47/// 
48/// # Arguments
49/// * `psf` - 3D array containing PSF values
50/// 
51/// # Returns
52/// The normalized PSF array
53pub fn normalize_psf(mut psf: Array3<f64>) -> Array3<f64> {
54    let sum = psf.sum();
55    if sum != 0.0 {
56        psf /= sum;
57    }
58    psf
59}
60
61/// Calculate the optimal grid size based on optical parameters.
62/// 
63/// Determines the optimal grid size for PSF calculations based on the optical
64/// parameters and desired oversampling factor. The size is rounded up to the
65/// next power of 2 for efficient FFT calculations.
66/// 
67/// # Arguments
68/// * `params` - Optical parameters of the system
69/// * `oversampling` - Oversampling factor for the grid
70/// 
71/// # Returns
72/// The optimal grid size (power of 2)
73pub fn optimal_grid_size(params: &OpticalParameters, oversampling: f64) -> usize {
74    let min_size = (2.0 * params.na / params.wavelength * oversampling).ceil() as usize;
75    // Round up to next power of 2 for FFT efficiency
76    let mut size = 1;
77    while size < min_size {
78        size *= 2;
79    }
80    size
81}
82
83/// Calculate the center of mass for a 3D array.
84/// 
85/// Computes the center of mass coordinates for each dimension of a 3D array.
86/// This is useful for finding the central position of a PSF.
87/// 
88/// # Arguments
89/// * `arr` - 3D array to analyze
90/// 
91/// # Returns
92/// Array of [x, y, z] coordinates of the center of mass
93pub fn center_of_mass(arr: &Array3<f64>) -> [f64; 3] {
94    let total_mass = arr.sum();
95    if total_mass == 0.0 {
96        return [0.0, 0.0, 0.0];
97    }
98
99    let dims = arr.shape();
100    let mut com = [0.0; 3];
101
102    for i in 0..3 {
103        let axis_coords: Array1<f64> = Array1::linspace(0.0, (dims[i] - 1) as f64, dims[i]);
104        let sum = arr.sum_axis(Axis(i));
105        com[i] = (axis_coords * sum).sum() / total_mass;
106    }
107    
108    com
109}
rusty_psf/core/utils.rs

rusty_psf/core/
utils.rs
