1use std::{
8 collections::BTreeMap,
9 fs::{self, File},
10 io::BufWriter,
11 path::{Path, PathBuf},
12 time::Instant,
13};
14
15use serde::{Deserialize, Serialize};
16
17use crate::{
18 Array2, Complex64, Result,
19 algorithms::ReconstructionAlgorithm,
20 evaluation::{evaluate_frame_intensity, evaluate_reconstruction_with_problem},
21 measurements::MeasurementRead,
22 model::ImagePlaneModel,
23 reconstruction::{ReconstructionProblem, ReconstructionResult},
24};
25
26pub const BENCHMARK_RECORD_FORMAT_VERSION: u32 = 1;
27pub const SMOKE_BENCHMARK_PROFILE: &str = "smoke";
28pub const CPU_BENCHMARK_PROFILE: &str = "cpu";
29
30#[derive(Clone, Copy, Debug, Eq, PartialEq)]
36pub struct BenchmarkProfile {
37 pub name: &'static str,
38 pub description: &'static str,
39 pub expected_runtime: &'static str,
40 pub output_directory: &'static str,
41 pub algorithms: &'static [&'static str],
42}
43
44impl BenchmarkProfile {
45 pub fn output_path(&self) -> PathBuf {
46 PathBuf::from(self.output_directory)
47 }
48}
49
50pub const BENCHMARK_PROFILES: &[BenchmarkProfile] = &[
51 BenchmarkProfile {
52 name: SMOKE_BENCHMARK_PROFILE,
53 description: "offline synthetic sanity profile for all implemented CPU algorithms",
54 expected_runtime: "under 1 minute on a typical laptop CPU",
55 output_directory: "target/benchmark-results/smoke",
56 algorithms: &[
57 "AlternatingProjection",
58 "Fpie",
59 "Epry",
60 "Admm",
61 "GradientDescent",
62 ],
63 },
64 BenchmarkProfile {
65 name: CPU_BENCHMARK_PROFILE,
66 description: "offline synthetic CPU comparison with longer iteration counts",
67 expected_runtime: "1-5 minutes on a typical laptop CPU",
68 output_directory: "target/benchmark-results/cpu",
69 algorithms: &[
70 "AlternatingProjection",
71 "Fpie",
72 "Epry",
73 "Admm",
74 "GradientDescent",
75 ],
76 },
77];
78
79pub fn benchmark_profile(name: &str) -> Option<&'static BenchmarkProfile> {
80 BENCHMARK_PROFILES
81 .iter()
82 .find(|profile| profile.name == name)
83}
84
85pub fn annotate_benchmark_profile(record: &mut BenchmarkRecord, profile: &BenchmarkProfile) {
87 record
88 .metadata
89 .insert("benchmark_profile".into(), profile.name.into());
90 record.metadata.insert(
91 "benchmark_profile_expected_runtime".into(),
92 profile.expected_runtime.into(),
93 );
94 record.metadata.insert(
95 "benchmark_profile_output_directory".into(),
96 profile.output_directory.into(),
97 );
98}
99
100#[derive(Clone, Debug, Serialize, Deserialize)]
102#[serde(deny_unknown_fields)]
103pub struct BenchmarkRecord {
104 pub format_version: u32,
105 pub dataset_name: String,
106 pub dataset_version: Option<String>,
107 pub preset_name: Option<String>,
108 pub crate_version: String,
109 pub random_seed: Option<u64>,
110 pub spatial_crop: Option<[usize; 4]>,
112 pub algorithm_name: String,
113 pub algorithm_configuration: String,
114 pub success: bool,
115 pub error: Option<String>,
116 pub frame_count: usize,
117 pub image_shape: [usize; 2],
118 pub reconstruction_shape: [usize; 2],
119 pub completed_iterations: usize,
120 pub runtime_seconds: f64,
121 pub initial_loss: Option<f64>,
122 pub final_loss: Option<f64>,
123 pub final_to_initial_loss_ratio: Option<f64>,
125 pub amplitude_rmse: Option<f64>,
126 pub phase_rmse: Option<f64>,
127 pub complex_field_relative_error: Option<f64>,
128 pub fourier_domain_relative_error: Option<f64>,
129 pub pupil_amplitude_rmse: Option<f64>,
130 pub pupil_phase_rmse: Option<f64>,
131 pub illumination_position_rmse: Option<f64>,
132 pub per_frame_residuals: Option<Vec<f64>>,
133 pub per_frame_residual_mean: Option<f64>,
134 pub per_frame_residual_max: Option<f64>,
135 pub selected_original_frame_indices: Vec<usize>,
136 pub selected_original_illumination_indices: Vec<Option<usize>>,
137 pub output_paths: Vec<PathBuf>,
138 pub metadata: BTreeMap<String, String>,
139}
140
141pub fn run_benchmark_case<A, M>(
145 dataset_name: impl Into<String>,
146 algorithm_configuration: impl Into<String>,
147 algorithm: A,
148 problem: &ReconstructionProblem<M>,
149 ground_truth: Option<&Array2<Complex64>>,
150 true_model: Option<&ImagePlaneModel>,
151 valid_object_mask: Option<&Array2<u8>>,
152) -> (BenchmarkRecord, Option<ReconstructionResult>)
153where
154 A: ReconstructionAlgorithm,
155 M: MeasurementRead,
156{
157 let algorithm_name = short_type_name::<A>().to_owned();
158 let image_shape = problem.measurements.image_shape();
159 let reconstruction_shape = problem.model.reconstruction_shape;
160 let mut record = BenchmarkRecord {
161 format_version: BENCHMARK_RECORD_FORMAT_VERSION,
162 dataset_name: dataset_name.into(),
163 dataset_version: None,
164 preset_name: None,
165 crate_version: env!("CARGO_PKG_VERSION").into(),
166 random_seed: None,
167 spatial_crop: None,
168 algorithm_name,
169 algorithm_configuration: algorithm_configuration.into(),
170 success: false,
171 error: None,
172 frame_count: problem.measurements.frame_count(),
173 image_shape: [image_shape.0, image_shape.1],
174 reconstruction_shape: [reconstruction_shape.0, reconstruction_shape.1],
175 completed_iterations: 0,
176 runtime_seconds: 0.0,
177 initial_loss: None,
178 final_loss: None,
179 final_to_initial_loss_ratio: None,
180 amplitude_rmse: None,
181 phase_rmse: None,
182 complex_field_relative_error: None,
183 fourier_domain_relative_error: None,
184 pupil_amplitude_rmse: None,
185 pupil_phase_rmse: None,
186 illumination_position_rmse: None,
187 per_frame_residuals: None,
188 per_frame_residual_mean: None,
189 per_frame_residual_max: None,
190 selected_original_frame_indices: problem
191 .measurements
192 .frame_metadata()
193 .iter()
194 .enumerate()
195 .map(|(index, metadata)| metadata.original_frame_index.unwrap_or(index))
196 .collect(),
197 selected_original_illumination_indices: problem
198 .measurements
199 .frame_metadata()
200 .iter()
201 .map(|metadata| {
202 metadata
203 .original_illumination_index
204 .or(metadata.illumination_index)
205 })
206 .collect(),
207 output_paths: Vec::new(),
208 metadata: BTreeMap::new(),
209 };
210
211 let started = Instant::now();
212 let result = match algorithm.run(problem) {
213 Ok(result) => result,
214 Err(error) => {
215 record.runtime_seconds = started.elapsed().as_secs_f64();
216 record.error = Some(error.to_string());
217 return (record, None);
218 }
219 };
220 record.runtime_seconds = started.elapsed().as_secs_f64();
221 record.algorithm_name = result.runtime.algorithm.clone();
222 record.completed_iterations = result.runtime.completed_iterations;
223 record.initial_loss = result.history.iterations.first().map(|entry| entry.loss);
224 record.final_loss = result.history.final_loss();
225 record.final_to_initial_loss_ratio = record.initial_loss.and_then(|initial| {
226 record
227 .final_loss
228 .filter(|_| initial.abs() > f64::EPSILON)
229 .map(|final_loss| final_loss / initial)
230 });
231
232 let residuals: Vec<f64> = if let Some(truth) = ground_truth {
233 let metrics = evaluate_reconstruction_with_problem(
234 &result,
235 problem,
236 truth,
237 true_model,
238 valid_object_mask,
239 );
240 match metrics {
241 Ok(metrics) => {
242 record.amplitude_rmse = Some(metrics.object.amplitude_rmse);
243 record.phase_rmse = Some(metrics.object.phase_rmse);
244 record.complex_field_relative_error = Some(metrics.object.complex_nrmse);
245 record.fourier_domain_relative_error = Some(metrics.object.fourier_nrmse);
246 record.pupil_amplitude_rmse =
247 metrics.pupil.as_ref().map(|value| value.amplitude_rmse);
248 record.pupil_phase_rmse = metrics.pupil.as_ref().map(|value| value.phase_rmse);
249 record.illumination_position_rmse = metrics
250 .illumination
251 .as_ref()
252 .map(|value| value.position_rmse);
253 metrics
254 .intensity
255 .map(|value| {
256 value
257 .per_frame
258 .into_iter()
259 .map(|frame| frame.normalized_l2)
260 .collect()
261 })
262 .unwrap_or_default()
263 }
264 Err(error) => {
265 record.error = Some(format!("benchmark metric calculation failed: {error}"));
266 return (record, Some(result));
267 }
268 }
269 } else {
270 match evaluate_frame_intensity(&result, problem) {
271 Ok(metrics) => metrics
272 .per_frame
273 .into_iter()
274 .map(|frame| frame.normalized_l2)
275 .collect(),
276 Err(error) => {
277 record.error = Some(format!("benchmark residual calculation failed: {error}"));
278 return (record, Some(result));
279 }
280 }
281 };
282 if !residuals.is_empty() {
283 record.per_frame_residual_mean =
284 Some(residuals.iter().sum::<f64>() / residuals.len() as f64);
285 record.per_frame_residual_max = residuals.iter().copied().reduce(f64::max);
286 }
287 record.per_frame_residuals = Some(residuals);
288 record.success = true;
289 (record, Some(result))
290}
291
292pub fn save_benchmark_outputs(
295 record: &mut BenchmarkRecord,
296 result: &ReconstructionResult,
297 directory: impl AsRef<Path>,
298) -> Result<()> {
299 let directory = directory.as_ref();
300 fs::create_dir_all(directory)?;
301 let stem = format!(
302 "{}-{}-{:016x}",
303 safe_stem(&record.dataset_name),
304 safe_stem(&record.algorithm_name),
305 case_hash(record),
306 );
307 let outputs = [
308 directory.join(format!("{stem}-amplitude.png")),
309 directory.join(format!("{stem}-phase.png")),
310 directory.join(format!("{stem}-result.json")),
311 directory.join(format!("{stem}-loss.csv")),
312 ];
313 result.save_amplitude(&outputs[0])?;
314 result.save_phase(&outputs[1])?;
315 result.save_bundle(&outputs[2])?;
316 result.save_loss_csv(&outputs[3])?;
317 record.output_paths.extend(outputs);
318 Ok(())
319}
320
321pub fn write_benchmark_json(records: &[BenchmarkRecord], path: impl AsRef<Path>) -> Result<()> {
322 #[derive(Serialize)]
323 struct Report<'a> {
324 format_version: u32,
325 records: &'a [BenchmarkRecord],
326 }
327
328 let writer = BufWriter::new(File::create(path)?);
329 serde_json::to_writer_pretty(
330 writer,
331 &Report {
332 format_version: BENCHMARK_RECORD_FORMAT_VERSION,
333 records,
334 },
335 )?;
336 Ok(())
337}
338
339pub fn write_benchmark_csv(records: &[BenchmarkRecord], path: impl AsRef<Path>) -> Result<()> {
340 let mut writer = csv::Writer::from_path(path)?;
341 writer.write_record([
342 "format_version",
343 "dataset_name",
344 "dataset_version",
345 "preset_name",
346 "crate_version",
347 "random_seed",
348 "spatial_crop",
349 "algorithm_name",
350 "algorithm_configuration",
351 "success",
352 "error",
353 "frame_count",
354 "image_height",
355 "image_width",
356 "reconstruction_height",
357 "reconstruction_width",
358 "completed_iterations",
359 "runtime_seconds",
360 "initial_loss",
361 "final_loss",
362 "final_to_initial_loss_ratio",
363 "amplitude_rmse",
364 "phase_rmse",
365 "complex_field_relative_error",
366 "fourier_domain_relative_error",
367 "pupil_amplitude_rmse",
368 "pupil_phase_rmse",
369 "illumination_position_rmse",
370 "per_frame_residuals",
371 "per_frame_residual_mean",
372 "per_frame_residual_max",
373 "selected_original_frame_indices",
374 "selected_original_illumination_indices",
375 "output_paths",
376 "metadata_json",
377 ])?;
378 for record in records {
379 writer.write_record([
380 record.format_version.to_string(),
381 record.dataset_name.clone(),
382 record.dataset_version.clone().unwrap_or_default(),
383 record.preset_name.clone().unwrap_or_default(),
384 record.crate_version.clone(),
385 record
386 .random_seed
387 .map_or_else(String::new, |seed| seed.to_string()),
388 record.spatial_crop.map_or_else(String::new, |crop| {
389 crop.iter()
390 .map(usize::to_string)
391 .collect::<Vec<_>>()
392 .join(";")
393 }),
394 record.algorithm_name.clone(),
395 record.algorithm_configuration.clone(),
396 record.success.to_string(),
397 record.error.clone().unwrap_or_default(),
398 record.frame_count.to_string(),
399 record.image_shape[0].to_string(),
400 record.image_shape[1].to_string(),
401 record.reconstruction_shape[0].to_string(),
402 record.reconstruction_shape[1].to_string(),
403 record.completed_iterations.to_string(),
404 record.runtime_seconds.to_string(),
405 optional_number(record.initial_loss),
406 optional_number(record.final_loss),
407 optional_number(record.final_to_initial_loss_ratio),
408 optional_number(record.amplitude_rmse),
409 optional_number(record.phase_rmse),
410 optional_number(record.complex_field_relative_error),
411 optional_number(record.fourier_domain_relative_error),
412 optional_number(record.pupil_amplitude_rmse),
413 optional_number(record.pupil_phase_rmse),
414 optional_number(record.illumination_position_rmse),
415 record
416 .per_frame_residuals
417 .as_deref()
418 .map_or_else(String::new, join_f64),
419 optional_number(record.per_frame_residual_mean),
420 optional_number(record.per_frame_residual_max),
421 join_usize(&record.selected_original_frame_indices),
422 record
423 .selected_original_illumination_indices
424 .iter()
425 .map(|value| value.map_or_else(String::new, |value| value.to_string()))
426 .collect::<Vec<_>>()
427 .join(";"),
428 record
429 .output_paths
430 .iter()
431 .map(|path| path.display().to_string())
432 .collect::<Vec<_>>()
433 .join(";"),
434 serde_json::to_string(&record.metadata)?,
435 ])?;
436 }
437 writer.flush()?;
438 Ok(())
439}
440
441fn short_type_name<T>() -> &'static str {
442 std::any::type_name::<T>()
443 .rsplit("::")
444 .next()
445 .unwrap_or("reconstruction algorithm")
446}
447
448fn optional_number(value: Option<f64>) -> String {
449 value.map_or_else(String::new, |value| value.to_string())
450}
451
452fn join_usize(values: &[usize]) -> String {
453 values
454 .iter()
455 .map(usize::to_string)
456 .collect::<Vec<_>>()
457 .join(";")
458}
459
460fn join_f64(values: &[f64]) -> String {
461 values
462 .iter()
463 .map(f64::to_string)
464 .collect::<Vec<_>>()
465 .join(";")
466}
467
468fn safe_stem(value: &str) -> String {
469 let stem: String = value
470 .chars()
471 .map(|character| {
472 if character.is_ascii_alphanumeric() || matches!(character, '-' | '_') {
473 character
474 } else {
475 '_'
476 }
477 })
478 .collect();
479 if stem.is_empty() {
480 "benchmark".into()
481 } else {
482 stem
483 }
484}
485
486fn case_hash(record: &BenchmarkRecord) -> u64 {
487 let mut hash = 0xcbf29ce484222325_u64;
490 let mut update = |bytes: &[u8]| {
491 for &byte in bytes {
492 hash ^= u64::from(byte);
493 hash = hash.wrapping_mul(0x100000001b3);
494 }
495 hash ^= 0xff;
496 hash = hash.wrapping_mul(0x100000001b3);
497 };
498 update(record.dataset_name.as_bytes());
499 update(
500 record
501 .dataset_version
502 .as_deref()
503 .unwrap_or_default()
504 .as_bytes(),
505 );
506 update(record.preset_name.as_deref().unwrap_or_default().as_bytes());
507 update(record.algorithm_name.as_bytes());
508 update(record.algorithm_configuration.as_bytes());
509 for index in &record.selected_original_frame_indices {
510 update(&index.to_le_bytes());
511 }
512 if let Some(crop) = record.spatial_crop {
513 for value in crop {
514 update(&value.to_le_bytes());
515 }
516 }
517 hash
518}