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#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
14pub struct FourierOffset {
15 pub row: f64,
17 pub column: f64,
19}
20
21impl FourierOffset {
22 pub const fn new(row: f64, column: f64) -> Self {
24 Self { row, column }
25 }
26
27 pub fn is_zero(self) -> bool {
29 self.row.abs() <= 1e-12 && self.column.abs() <= 1e-12
30 }
31}
32
33#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
35pub struct FourierCrop {
36 pub start_row: usize,
38 pub start_col: usize,
40 pub height: usize,
42 pub width: usize,
44}
45
46impl FourierCrop {
47 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 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 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 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 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 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#[derive(Clone, Debug, Serialize, Deserialize)]
317pub struct CropIndices {
318 pub(crate) crops: Vec<FourierCrop>,
319}
320
321impl CropIndices {
322 pub fn new(crops: Vec<FourierCrop>) -> Self {
324 Self { crops }
325 }
326
327 pub fn len(&self) -> usize {
329 self.crops.len()
330 }
331
332 pub fn is_empty(&self) -> bool {
334 self.crops.is_empty()
335 }
336
337 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 pub fn as_slice(&self) -> &[FourierCrop] {
350 &self.crops
351 }
352}