1use crate::{
2 Array2, Complex64, Result,
3 configuration::{ExperimentDescription, SimulationConfiguration},
4 error::Error,
5 experiment::Illumination,
6 measurements::MeasurementStack,
7 model::{ImagePlaneModel, ReconstructionShape},
8 reconstruction::ReconstructionProblem,
9};
10
11use super::Dataset;
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub struct Rect {
15 pub row: usize,
16 pub column: usize,
17 pub height: usize,
18 pub width: usize,
19}
20
21impl Rect {
22 pub fn new(row: usize, column: usize, height: usize, width: usize) -> Result<Self> {
23 if height == 0 || width == 0 {
24 return Err(Error::Dataset(
25 "dataset crop height and width must be non-zero".into(),
26 ));
27 }
28 Ok(Self {
29 row,
30 column,
31 height,
32 width,
33 })
34 }
35}
36
37#[derive(Clone, Debug, PartialEq, Eq)]
38pub enum FrameSelector {
39 All,
40 EveryNth(usize),
41 Indices(Vec<usize>),
42}
43
44#[derive(Clone, Debug)]
45pub struct DatasetSubset {
46 measurements: MeasurementStack,
47 configuration: SimulationConfiguration,
48 ground_truth_object: Option<Array2<Complex64>>,
49 valid_object_mask: Option<Array2<u8>>,
50 provenance: std::collections::BTreeMap<String, String>,
51 measurement_units: Option<String>,
52}
53
54impl DatasetSubset {
55 pub fn measurements(&self) -> &MeasurementStack {
56 &self.measurements
57 }
58
59 pub fn configuration(&self) -> &SimulationConfiguration {
60 &self.configuration
61 }
62
63 pub fn ground_truth_object(&self) -> Option<&Array2<Complex64>> {
64 self.ground_truth_object.as_ref()
65 }
66
67 pub fn valid_object_mask(&self) -> Option<&Array2<u8>> {
68 self.valid_object_mask.as_ref()
69 }
70
71 pub fn provenance(&self) -> &std::collections::BTreeMap<String, String> {
72 &self.provenance
73 }
74
75 pub fn measurement_units(&self) -> Option<&str> {
76 self.measurement_units.as_deref()
77 }
78
79 pub fn reconstruction_problem(&self) -> Result<ReconstructionProblem<MeasurementStack>> {
80 ReconstructionProblem::new(
81 self.measurements.clone(),
82 self.configuration
83 .compiled_models
84 .reconstruction_model
85 .clone(),
86 )
87 }
88}
89
90#[derive(Clone, Debug)]
91pub struct DatasetSubsetBuilder<'a> {
92 dataset: &'a Dataset,
93 frames: FrameSelector,
94 crop: Option<Rect>,
95}
96
97impl<'a> DatasetSubsetBuilder<'a> {
98 pub(crate) fn new(dataset: &'a Dataset) -> Self {
99 Self {
100 dataset,
101 frames: FrameSelector::All,
102 crop: None,
103 }
104 }
105
106 pub fn frames(mut self, selector: FrameSelector) -> Self {
107 self.frames = selector;
108 self
109 }
110
111 pub fn every_nth_frame(self, step: usize) -> Self {
112 self.frames(FrameSelector::EveryNth(step))
113 }
114
115 pub fn crop(mut self, crop: Rect) -> Self {
116 self.crop = Some(crop);
117 self
118 }
119
120 pub fn crop_pixels(
121 self,
122 row: usize,
123 column: usize,
124 height: usize,
125 width: usize,
126 ) -> Result<Self> {
127 Ok(self.crop(Rect::new(row, column, height, width)?))
128 }
129
130 pub fn build(self) -> Result<DatasetSubset> {
131 let source = self.dataset.measurements();
132 let model = &self
133 .dataset
134 .configuration()
135 .compiled_models
136 .reconstruction_model;
137 let indices = selected_indices(&self.frames, source.frame_count())?;
138 let source_shape = source.image_shape();
139 let crop = self.crop.unwrap_or(Rect {
140 row: 0,
141 column: 0,
142 height: source_shape.0,
143 width: source_shape.1,
144 });
145 validate_crop(crop, source_shape)?;
146
147 let mut data = Vec::with_capacity(indices.len() * crop.height * crop.width);
148 let mut metadata = Vec::with_capacity(indices.len());
149 for (new_index, &source_index) in indices.iter().enumerate() {
150 crop_frame(source.frame(source_index)?, source_shape, crop, &mut data);
151 let mut frame_metadata = source.frame_metadata[source_index].clone();
152 frame_metadata.original_frame_index =
153 Some(frame_metadata.original_frame_index.unwrap_or(source_index));
154 frame_metadata.original_illumination_index = frame_metadata
155 .original_illumination_index
156 .or(frame_metadata.illumination_index);
157 frame_metadata.frame_index = new_index;
158 frame_metadata.illumination_index = (!model.is_multiplexed()).then_some(new_index);
159 metadata.push(frame_metadata);
160 }
161 let mut measurements =
162 MeasurementStack::from_vec(data, (crop.height, crop.width), metadata)?;
163 measurements.dark_frame =
164 crop_optional_shared(source.dark_frame.as_deref(), source_shape, crop);
165 measurements.flat_field =
166 crop_optional_shared(source.flat_field.as_deref(), source_shape, crop);
167 measurements.background = crop_optional_frames(
168 source.background.as_deref(),
169 source_shape,
170 source.frame_count(),
171 &indices,
172 crop,
173 )?;
174 measurements.masks = crop_optional_masks(
175 source.masks.as_deref(),
176 source_shape,
177 source.frame_count(),
178 &indices,
179 crop,
180 )?;
181 measurements.preprocessing = source.preprocessing.clone();
182 measurements.validate()?;
183
184 let scale_y = model.reconstruction_shape.0 as f64 / model.image_shape.0 as f64;
185 let scale_x = model.reconstruction_shape.1 as f64 / model.image_shape.1 as f64;
186 let reconstruction_shape = (
187 (crop.height as f64 * scale_y).round() as usize,
188 (crop.width as f64 * scale_x).round() as usize,
189 );
190 let object_crop = scaled_object_crop(crop, scale_y, scale_x)?;
191 let ground_truth_object = self
192 .dataset
193 .ground_truth_object()
194 .map(|truth| crop_array(truth, object_crop))
195 .transpose()?;
196 let valid_object_mask = self
197 .dataset
198 .valid_object_mask()
199 .map(|mask| crop_array(mask, object_crop))
200 .transpose()?;
201 let source_configuration = self.dataset.configuration();
202 let true_experiment = subset_experiment(
203 &source_configuration.true_experiment,
204 &source_configuration.compiled_models.true_model,
205 &indices,
206 crop,
207 )?;
208 let reconstruction_experiment = subset_experiment(
209 &source_configuration.reconstruction_experiment,
210 &source_configuration.compiled_models.reconstruction_model,
211 &indices,
212 crop,
213 )?;
214 let configuration = SimulationConfiguration::new(
215 true_experiment,
216 reconstruction_experiment,
217 (crop.height, crop.width),
218 ReconstructionShape::Exact(reconstruction_shape),
219 )?
220 .with_random_seed(source_configuration.random_seed);
221 Ok(DatasetSubset {
222 measurements,
223 configuration,
224 ground_truth_object,
225 valid_object_mask,
226 provenance: self.dataset.provenance().clone(),
227 measurement_units: self.dataset.measurement_units().map(str::to_owned),
228 })
229 }
230}
231
232fn subset_experiment(
233 description: &ExperimentDescription,
234 model: &ImagePlaneModel,
235 indices: &[usize],
236 crop: Rect,
237) -> Result<ExperimentDescription> {
238 let (k_vectors, frame_weights) = match &model.multiplexing_matrix {
239 Some(matrix) => (
240 model.k_vectors.clone(),
241 Some(indices.iter().map(|&index| matrix[index].clone()).collect()),
242 ),
243 None => (
244 indices
245 .iter()
246 .map(|&index| model.k_vectors[index])
247 .collect(),
248 None,
249 ),
250 };
251 let frame_gains = model
252 .frame_gains
253 .as_ref()
254 .map(|gains| indices.iter().map(|&index| gains[index]).collect());
255 let illumination = Illumination::Calibrated {
256 k_vectors,
257 frame_gains,
258 frame_weights,
259 };
260 let mut subset = ExperimentDescription::new(description.optics.clone(), illumination);
261 subset.optical_background = crop_optional_frames(
262 model.background.as_deref(),
263 model.image_shape,
264 model.frame_count(),
265 indices,
266 crop,
267 )?;
268 subset.validate()?;
269 Ok(subset)
270}
271
272fn selected_indices(selector: &FrameSelector, frame_count: usize) -> Result<Vec<usize>> {
273 let indices = match selector {
274 FrameSelector::All => (0..frame_count).collect(),
275 FrameSelector::EveryNth(0) => {
276 return Err(Error::Dataset(
277 "frame subset step must be greater than zero".into(),
278 ));
279 }
280 FrameSelector::EveryNth(step) => (0..frame_count).step_by(*step).collect(),
281 FrameSelector::Indices(indices) => indices.clone(),
282 };
283 if indices.is_empty() {
284 return Err(Error::Dataset(
285 "dataset frame subset must contain at least one frame".into(),
286 ));
287 }
288 let mut seen = vec![false; frame_count];
289 for &index in &indices {
290 if index >= frame_count {
291 return Err(Error::FrameOutOfRange {
292 index,
293 frames: frame_count,
294 });
295 }
296 if seen[index] {
297 return Err(Error::Dataset(format!(
298 "dataset frame subset contains duplicate index {index}"
299 )));
300 }
301 seen[index] = true;
302 }
303 Ok(indices)
304}
305
306fn validate_crop(crop: Rect, shape: (usize, usize)) -> Result<()> {
307 if crop
308 .row
309 .checked_add(crop.height)
310 .is_none_or(|end| end > shape.0)
311 || crop
312 .column
313 .checked_add(crop.width)
314 .is_none_or(|end| end > shape.1)
315 {
316 return Err(Error::Dataset(format!(
317 "crop {crop:?} is outside measurement shape {shape:?}"
318 )));
319 }
320 Ok(())
321}
322
323fn crop_frame<T: Copy>(source: &[T], shape: (usize, usize), crop: Rect, output: &mut Vec<T>) {
324 for row in crop.row..crop.row + crop.height {
325 let start = row * shape.1 + crop.column;
326 output.extend_from_slice(&source[start..start + crop.width]);
327 }
328}
329
330fn scaled_object_crop(crop: Rect, scale_y: f64, scale_x: f64) -> Result<Rect> {
331 let values = [
332 crop.row as f64 * scale_y,
333 crop.column as f64 * scale_x,
334 crop.height as f64 * scale_y,
335 crop.width as f64 * scale_x,
336 ];
337 if values
338 .iter()
339 .any(|value| (value - value.round()).abs() > 1e-9)
340 {
341 return Err(Error::Dataset(
342 "measurement crop does not map to integer reconstruction pixels".into(),
343 ));
344 }
345 Rect::new(
346 values[0].round() as usize,
347 values[1].round() as usize,
348 values[2].round() as usize,
349 values[3].round() as usize,
350 )
351}
352
353fn crop_array<T: Copy>(source: &Array2<T>, crop: Rect) -> Result<Array2<T>> {
354 validate_crop(crop, source.shape())?;
355 let mut output = Vec::with_capacity(crop.height * crop.width);
356 crop_frame(source.as_slice(), source.shape(), crop, &mut output);
357 Array2::from_vec((crop.height, crop.width), output)
358}
359
360fn crop_optional_shared(
361 source: Option<&[f64]>,
362 shape: (usize, usize),
363 crop: Rect,
364) -> Option<Vec<f64>> {
365 source.map(|source| {
366 let mut output = Vec::with_capacity(crop.height * crop.width);
367 crop_frame(source, shape, crop, &mut output);
368 output
369 })
370}
371
372fn crop_optional_frames(
373 source: Option<&[f64]>,
374 shape: (usize, usize),
375 frame_count: usize,
376 indices: &[usize],
377 crop: Rect,
378) -> Result<Option<Vec<f64>>> {
379 let Some(source) = source else {
380 return Ok(None);
381 };
382 let frame_len = shape.0 * shape.1;
383 if source.len() == frame_len {
384 return Ok(crop_optional_shared(Some(source), shape, crop));
385 }
386 if source.len() != frame_count * frame_len {
387 return Err(Error::InvalidMeasurements(
388 "source background length is inconsistent".into(),
389 ));
390 }
391 let mut output = Vec::with_capacity(indices.len() * crop.height * crop.width);
392 for &index in indices {
393 crop_frame(
394 &source[index * frame_len..(index + 1) * frame_len],
395 shape,
396 crop,
397 &mut output,
398 );
399 }
400 Ok(Some(output))
401}
402
403fn crop_optional_masks(
404 source: Option<&[u8]>,
405 shape: (usize, usize),
406 frame_count: usize,
407 indices: &[usize],
408 crop: Rect,
409) -> Result<Option<Vec<u8>>> {
410 let Some(source) = source else {
411 return Ok(None);
412 };
413 let frame_len = shape.0 * shape.1;
414 if source.len() == frame_len {
415 let mut output = Vec::with_capacity(crop.height * crop.width);
416 crop_frame(source, shape, crop, &mut output);
417 return Ok(Some(output));
418 }
419 if source.len() != frame_count * frame_len {
420 return Err(Error::InvalidMeasurements(
421 "source mask length is inconsistent".into(),
422 ));
423 }
424 let mut output = Vec::with_capacity(indices.len() * crop.height * crop.width);
425 for &index in indices {
426 crop_frame(
427 &source[index * frame_len..(index + 1) * frame_len],
428 shape,
429 crop,
430 &mut output,
431 );
432 }
433 Ok(Some(output))
434}