Skip to main content

fpm_rs/model/
crop.rs

1use serde::{Deserialize, Serialize};
2
3use ndarray::{ArrayView2, ArrayViewMut2};
4use num_complex::Complex64;
5
6use crate::{
7    Result,
8    array_layout::{StandardView2, StandardViewMut2, checked_len_2d},
9    error::Error,
10};
11
12/// Fractional Fourier-grid displacement relative to an integer crop origin.
13#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
14pub struct FourierOffset {
15    /// Fractional displacement along the Fourier-array row axis, in grid pixels.
16    pub row: f64,
17    /// Fractional displacement along the Fourier-array column axis, in grid pixels.
18    pub column: f64,
19}
20
21impl FourierOffset {
22    /// Creates a `(row, column)` fractional displacement in Fourier-grid pixels.
23    pub const fn new(row: f64, column: f64) -> Self {
24        Self { row, column }
25    }
26
27    /// Returns whether both components are within `1e-12` pixel of zero.
28    pub fn is_zero(self) -> bool {
29        self.row.abs() <= 1e-12 && self.column.abs() <= 1e-12
30    }
31}
32
33/// Integer rectangular region within a centered high-resolution Fourier array.
34#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
35pub struct FourierCrop {
36    /// Zero-based first row in the high-resolution spectrum.
37    pub start_row: usize,
38    /// Zero-based first column in the high-resolution spectrum.
39    pub start_col: usize,
40    /// Number of rows, normally equal to the low-resolution image height.
41    pub height: usize,
42    /// Number of columns, normally equal to the low-resolution image width.
43    pub width: usize,
44}
45
46impl FourierCrop {
47    /// Creates an unchecked integer crop; call [`Self::validate_inside`] before use.
48    pub fn new(start_row: usize, start_col: usize, height: usize, width: usize) -> Self {
49        Self {
50            start_row,
51            start_col,
52            height,
53            width,
54        }
55    }
56
57    /// Checks non-zero crop dimensions and containment in `(height, width)` `shape`.
58    pub fn validate_inside(&self, shape: (usize, usize)) -> Result<()> {
59        let end_row = self.start_row.checked_add(self.height);
60        let end_col = self.start_col.checked_add(self.width);
61        if self.height == 0
62            || self.width == 0
63            || end_row.is_none_or(|end| end > shape.0)
64            || end_col.is_none_or(|end| end > shape.1)
65        {
66            return Err(Error::InvalidModel(format!(
67                "crop {self:?} lies outside reconstruction shape {shape:?}"
68            )));
69        }
70        Ok(())
71    }
72
73    /// Extracts from a C-contiguous standard row-major view without copying the
74    /// input. Nonstandard views are rejected.
75    pub fn extract<T: Copy>(&self, source: ArrayView2<'_, T>, destination: &mut [T]) -> Result<()> {
76        self.extract_standard(StandardView2::try_from(source)?, destination)
77    }
78
79    pub(crate) fn extract_standard<T: Copy>(
80        &self,
81        source: StandardView2<'_, T>,
82        destination: &mut [T],
83    ) -> Result<()> {
84        let source_shape = source.dim();
85        self.validate_inside(source_shape)?;
86        let crop_len = checked_len_2d((self.height, self.width))?;
87        if destination.len() != crop_len {
88            return Err(Error::LengthMismatch {
89                actual: destination.len(),
90                expected: crop_len,
91                shape: (self.height, self.width),
92            });
93        }
94        let source_values = source.as_slice();
95        for row in 0..self.height {
96            let source_start = (self.start_row + row) * source_shape.1 + self.start_col;
97            let destination_start = row * self.width;
98            destination[destination_start..destination_start + self.width]
99                .copy_from_slice(&source_values[source_start..source_start + self.width]);
100        }
101        Ok(())
102    }
103
104    /// Checks that a bilinear stencil for `offset` remains inside `(height, width)` `shape`.
105    pub fn validate_subpixel_inside(
106        &self,
107        shape: (usize, usize),
108        offset: FourierOffset,
109    ) -> Result<()> {
110        interpolation_axis(self.start_row, self.height, shape.0, offset.row)?;
111        interpolation_axis(self.start_col, self.width, shape.1, offset.column)?;
112        Ok(())
113    }
114
115    /// Bilinearly samples a potentially fractional crop from `source`.
116    ///
117    /// This is a local Fourier-grid interpolation, not a bandlimited shift
118    /// operator. Its approximation error grows with grid-frequency content and
119    /// fractional displacement; low-bandwidth objects are the intended regime.
120    pub fn extract_subpixel(
121        &self,
122        source: ArrayView2<'_, Complex64>,
123        destination: &mut [Complex64],
124        offset: FourierOffset,
125    ) -> Result<()> {
126        self.extract_subpixel_standard(StandardView2::try_from(source)?, destination, offset)
127    }
128
129    pub(crate) fn extract_subpixel_standard(
130        &self,
131        source: StandardView2<'_, Complex64>,
132        destination: &mut [Complex64],
133        offset: FourierOffset,
134    ) -> Result<()> {
135        let crop_len = checked_len_2d((self.height, self.width))?;
136        if destination.len() != crop_len {
137            return Err(Error::LengthMismatch {
138                actual: destination.len(),
139                expected: crop_len,
140                shape: (self.height, self.width),
141            });
142        }
143        if offset.is_zero() {
144            return self.extract_standard(source, destination);
145        }
146        let shape = source.dim();
147        let rows = interpolation_axis(self.start_row, self.height, shape.0, offset.row)?;
148        let columns = interpolation_axis(self.start_col, self.width, shape.1, offset.column)?;
149        let source = source.as_slice();
150        for row in 0..self.height {
151            let lower_row = rows.lower_start + row;
152            let upper_row = rows.upper_start + row;
153            for column in 0..self.width {
154                let lower_column = columns.lower_start + column;
155                let upper_column = columns.upper_start + column;
156                destination[row * self.width + column] = source[lower_row * shape.1 + lower_column]
157                    * (rows.lower_weight * columns.lower_weight)
158                    + source[lower_row * shape.1 + upper_column]
159                        * (rows.lower_weight * columns.upper_weight)
160                    + source[upper_row * shape.1 + lower_column]
161                        * (rows.upper_weight * columns.lower_weight)
162                    + source[upper_row * shape.1 + upper_column]
163                        * (rows.upper_weight * columns.upper_weight);
164            }
165        }
166        Ok(())
167    }
168
169    /// Adds the exact adjoint of [`Self::extract_subpixel`] to `destination`.
170    pub fn insert_subpixel_adjoint(
171        &self,
172        destination: ArrayViewMut2<'_, Complex64>,
173        update: &[Complex64],
174        scale: f64,
175        offset: FourierOffset,
176    ) -> Result<()> {
177        let mut destination = StandardViewMut2::try_from(destination)?;
178        let destination_shape = destination.dim();
179        self.insert_subpixel_adjoint_slice(
180            destination.as_slice_mut(),
181            destination_shape,
182            update,
183            scale,
184            offset,
185        )
186    }
187
188    pub(crate) fn insert_subpixel_adjoint_slice(
189        &self,
190        destination: &mut [Complex64],
191        destination_shape: (usize, usize),
192        update: &[Complex64],
193        scale: f64,
194        offset: FourierOffset,
195    ) -> Result<()> {
196        let destination_len = checked_len_2d(destination_shape)?;
197        if destination.len() != destination_len {
198            return Err(Error::LengthMismatch {
199                actual: destination.len(),
200                expected: destination_len,
201                shape: destination_shape,
202            });
203        }
204        let update_len = checked_len_2d((self.height, self.width))?;
205        if update.len() != update_len {
206            return Err(Error::LengthMismatch {
207                actual: update.len(),
208                expected: update_len,
209                shape: (self.height, self.width),
210            });
211        }
212        if !scale.is_finite() {
213            return Err(Error::InvalidParameter {
214                name: "scale",
215                reason: "must be finite".into(),
216            });
217        }
218        if offset.is_zero() {
219            self.validate_inside(destination_shape)?;
220            for row in 0..self.height {
221                for column in 0..self.width {
222                    destination
223                        [(self.start_row + row) * destination_shape.1 + self.start_col + column] +=
224                        scale * update[row * self.width + column];
225                }
226            }
227            return Ok(());
228        }
229        let rows =
230            interpolation_axis(self.start_row, self.height, destination_shape.0, offset.row)?;
231        let columns = interpolation_axis(
232            self.start_col,
233            self.width,
234            destination_shape.1,
235            offset.column,
236        )?;
237        for row in 0..self.height {
238            let lower_row = rows.lower_start + row;
239            let upper_row = rows.upper_start + row;
240            for column in 0..self.width {
241                let lower_column = columns.lower_start + column;
242                let upper_column = columns.upper_start + column;
243                let value = scale * update[row * self.width + column];
244                destination[lower_row * destination_shape.1 + lower_column] +=
245                    value * (rows.lower_weight * columns.lower_weight);
246                destination[lower_row * destination_shape.1 + upper_column] +=
247                    value * (rows.lower_weight * columns.upper_weight);
248                destination[upper_row * destination_shape.1 + lower_column] +=
249                    value * (rows.upper_weight * columns.lower_weight);
250                destination[upper_row * destination_shape.1 + upper_column] +=
251                    value * (rows.upper_weight * columns.upper_weight);
252            }
253        }
254        Ok(())
255    }
256}
257
258#[derive(Clone, Copy, Debug)]
259struct AxisInterpolation {
260    lower_start: usize,
261    upper_start: usize,
262    lower_weight: f64,
263    upper_weight: f64,
264}
265
266fn interpolation_axis(
267    start: usize,
268    length: usize,
269    bound: usize,
270    offset: f64,
271) -> Result<AxisInterpolation> {
272    if length == 0 || bound == 0 || !offset.is_finite() {
273        return Err(Error::InvalidModel(
274            "subpixel crop dimensions and offset must be finite and non-zero".into(),
275        ));
276    }
277    let nearest = offset.round();
278    let (integer_offset, fraction) = if (offset - nearest).abs() <= 1e-12 {
279        (nearest, 0.0)
280    } else {
281        let floor = offset.floor();
282        (floor, offset - floor)
283    };
284    if integer_offset < isize::MIN as f64 || integer_offset > isize::MAX as f64 {
285        return Err(Error::InvalidModel(
286            "subpixel crop offset is outside the supported index range".into(),
287        ));
288    }
289    let start = isize::try_from(start).map_err(|_| {
290        Error::InvalidModel("crop origin is outside the supported index range".into())
291    })?;
292    let lower_start = start
293        .checked_add(integer_offset as isize)
294        .filter(|value| *value >= 0)
295        .ok_or_else(|| Error::InvalidModel("subpixel crop starts outside the grid".into()))?;
296    let lower_start = usize::try_from(lower_start)
297        .map_err(|_| Error::InvalidModel("subpixel crop starts outside the grid".into()))?;
298    let lower_end = lower_start
299        .checked_add(length)
300        .ok_or_else(|| Error::InvalidModel("subpixel crop dimensions overflow the grid".into()))?;
301    let needs_upper = fraction > 0.0;
302    if lower_end > bound || (needs_upper && lower_end >= bound) {
303        return Err(Error::InvalidModel(
304            "subpixel crop interpolation stencil lies outside the grid".into(),
305        ));
306    }
307    Ok(AxisInterpolation {
308        lower_start,
309        upper_start: lower_start + usize::from(needs_upper),
310        lower_weight: 1.0 - fraction,
311        upper_weight: fraction,
312    })
313}
314
315/// Source-order collection of Fourier crops.
316#[derive(Clone, Debug, Serialize, Deserialize)]
317pub struct CropIndices {
318    pub(crate) crops: Vec<FourierCrop>,
319}
320
321impl CropIndices {
322    /// Stores crop rectangles in individual illumination-source order.
323    pub fn new(crops: Vec<FourierCrop>) -> Self {
324        Self { crops }
325    }
326
327    /// Returns the number of individual source crops.
328    pub fn len(&self) -> usize {
329        self.crops.len()
330    }
331
332    /// Returns whether no source crops are stored.
333    pub fn is_empty(&self) -> bool {
334        self.crops.is_empty()
335    }
336
337    /// Returns the crop for `frame`, interpreted here as an individual source index.
338    pub fn get(&self, frame: usize) -> Result<FourierCrop> {
339        self.crops
340            .get(frame)
341            .copied()
342            .ok_or(Error::FrameOutOfRange {
343                index: frame,
344                frames: self.crops.len(),
345            })
346    }
347
348    /// Borrows all crops in individual source order.
349    pub fn as_slice(&self) -> &[FourierCrop] {
350        &self.crops
351    }
352}