Skip to main content

fpm_rs/
benchmark.rs

1//! Reproducible single-case reconstruction benchmarks.
2//!
3//! The runner is generic over the existing reconstruction and measurement
4//! traits. Calling it once per concrete algorithm avoids a second algorithm
5//! registry or trait-object hierarchy.
6
7use 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/// Stable metadata for a named benchmark profile.
31///
32/// Profiles describe runtime and output expectations only. Examples still list
33/// concrete algorithms explicitly so the benchmark layer does not become an
34/// algorithm registry.
35#[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
85/// Adds the selected profile metadata to one benchmark record.
86pub 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/// Serializable summary of one algorithm run on one immutable problem.
101#[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    /// Original-image crop as `[row, column, height, width]`, when applicable.
111    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    /// Final loss divided by initial loss; lower values indicate a larger decrease.
124    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
141/// Runs one concrete algorithm and always returns a record. Reconstruction or
142/// metric failures are stored in `record.error`; successful reconstruction data
143/// is returned separately so callers may inspect or save it.
144pub 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
292/// Saves the standard reconstruction artifacts for a benchmark case and adds
293/// their paths to the record.
294pub 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    // Stable FNV-1a rather than a process-seeded map hasher, so output names are
488    // reproducible across runs and platforms.
489    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}