From d851aad548c580b7ad4b508c9fb8f6c5da76a2a6 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 20:02:35 +0100 Subject: [PATCH] feat: add CSV and JSON data export for visualization Add to_csv(), export_csv(), and export_json() methods to Study for exporting trial data to external visualization tools. CSV export works without extra dependencies; JSON export requires the serde feature. --- src/study.rs | 147 +++++++++++++++++++++++++++ tests/export_tests.rs | 230 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 377 insertions(+) create mode 100644 tests/export_tests.rs diff --git a/src/study.rs b/src/study.rs index 0875b94..5f6d6a5 100644 --- a/src/study.rs +++ b/src/study.rs @@ -1732,6 +1732,119 @@ impl Study where V: PartialOrd + Clone + fmt::Display, { + /// Export completed trials to CSV format. + /// + /// Columns: `trial_id`, `value`, `state`, then one column per unique + /// parameter label, then one column per unique user-attribute key. + /// + /// Parameters without labels use a generated name (`param_`). + /// Pruned trials have an empty `value` cell. + /// + /// # Errors + /// + /// Returns an I/O error if writing fails. + pub fn to_csv(&self, mut writer: impl std::io::Write) -> std::io::Result<()> { + use std::collections::BTreeMap; + + let trials = self.completed_trials.read(); + + // Collect all unique parameter labels (sorted for deterministic column order). + let mut param_columns: BTreeMap = BTreeMap::new(); + for trial in trials.iter() { + for &id in trial.params.keys() { + param_columns.entry(id).or_insert_with(|| { + trial + .param_labels + .get(&id) + .cloned() + .unwrap_or_else(|| id.to_string()) + }); + } + } + // Fill in labels from other trials that might have better labels. + for trial in trials.iter() { + for (&id, label) in &trial.param_labels { + param_columns.entry(id).or_insert_with(|| label.clone()); + } + } + + // Collect all unique attribute keys (sorted). + let mut attr_keys: Vec = Vec::new(); + for trial in trials.iter() { + for key in trial.user_attrs.keys() { + if !attr_keys.contains(key) { + attr_keys.push(key.clone()); + } + } + } + attr_keys.sort(); + + let param_ids: Vec = param_columns.keys().copied().collect(); + + // Write header. + write!(writer, "trial_id,value,state")?; + for id in ¶m_ids { + write!(writer, ",{}", csv_escape(¶m_columns[id]))?; + } + for key in &attr_keys { + write!(writer, ",{}", csv_escape(key))?; + } + writeln!(writer)?; + + // Write one row per trial. + for trial in trials.iter() { + write!(writer, "{}", trial.id)?; + + // Value: empty for pruned trials. + if trial.state == TrialState::Complete { + write!(writer, ",{}", trial.value)?; + } else { + write!(writer, ",")?; + } + + write!( + writer, + ",{}", + match trial.state { + TrialState::Complete => "Complete", + TrialState::Pruned => "Pruned", + TrialState::Failed => "Failed", + TrialState::Running => "Running", + } + )?; + + for id in ¶m_ids { + if let Some(pv) = trial.params.get(id) { + write!(writer, ",{pv}")?; + } else { + write!(writer, ",")?; + } + } + + for key in &attr_keys { + if let Some(attr) = trial.user_attrs.get(key) { + write!(writer, ",{}", csv_escape(&format_attr(attr)))?; + } else { + write!(writer, ",")?; + } + } + + writeln!(writer)?; + } + + Ok(()) + } + + /// Export completed trials to a CSV file. + /// + /// # Errors + /// + /// Returns an I/O error if the file cannot be created or written. + pub fn export_csv(&self, path: impl AsRef) -> std::io::Result<()> { + let file = std::fs::File::create(path)?; + self.to_csv(std::io::BufWriter::new(file)) + } + /// Returns a human-readable summary of the study. /// /// The summary includes: @@ -2073,6 +2186,19 @@ pub struct StudySnapshot { #[cfg(feature = "serde")] impl Study { + /// Export trials as a pretty-printed JSON array to a file. + /// + /// Each element in the array is a serialized [`CompletedTrial`]. + /// + /// # Errors + /// + /// Returns an I/O error if the file cannot be created or written. + pub fn export_json(&self, path: impl AsRef) -> std::io::Result<()> { + let file = std::fs::File::create(path)?; + let trials = self.trials(); + serde_json::to_writer_pretty(file, &trials).map_err(std::io::Error::other) + } + /// Saves the study state to a JSON file. /// /// # Errors @@ -2171,3 +2297,24 @@ fn is_trial_pruned(e: &E) -> bool { any.downcast_ref::().is_some() } } + +/// Escape a string for CSV output. If the value contains a comma, quote, or +/// newline, wrap it in double-quotes and double any embedded quotes. +fn csv_escape(s: &str) -> String { + if s.contains(',') || s.contains('"') || s.contains('\n') { + format!("\"{}\"", s.replace('"', "\"\"")) + } else { + s.to_string() + } +} + +/// Format an `AttrValue` as a string for CSV cells. +fn format_attr(attr: &crate::trial::AttrValue) -> String { + use crate::trial::AttrValue; + match attr { + AttrValue::Float(v) => v.to_string(), + AttrValue::Int(v) => v.to_string(), + AttrValue::String(v) => v.clone(), + AttrValue::Bool(v) => v.to_string(), + } +} diff --git a/tests/export_tests.rs b/tests/export_tests.rs new file mode 100644 index 0000000..9a43149 --- /dev/null +++ b/tests/export_tests.rs @@ -0,0 +1,230 @@ +use optimizer::parameter::{FloatParam, IntParam, Parameter}; +use optimizer::sampler::random::RandomSampler; +use optimizer::{Direction, Study}; + +#[test] +fn csv_empty_study_produces_header_only() { + let study: Study = Study::new(Direction::Minimize); + let mut buf = Vec::new(); + study.to_csv(&mut buf).unwrap(); + let csv = String::from_utf8(buf).unwrap(); + assert_eq!(csv, "trial_id,value,state\n"); +} + +#[test] +fn csv_includes_all_trial_data() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + let x = FloatParam::new(0.0, 10.0).name("x"); + let y = IntParam::new(1, 5).name("y"); + + study + .optimize(3, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, optimizer::Error>(xv + yv as f64) + }) + .unwrap(); + + let mut buf = Vec::new(); + study.to_csv(&mut buf).unwrap(); + let csv = String::from_utf8(buf).unwrap(); + + let lines: Vec<&str> = csv.lines().collect(); + // Header + 3 data rows. + assert_eq!(lines.len(), 4); + + // Header should contain our parameter names. + let header = lines[0]; + assert!(header.starts_with("trial_id,value,state")); + assert!(header.contains("x")); + assert!(header.contains("y")); + + // Each data row should have the right number of columns. + let n_cols = header.split(',').count(); + for line in &lines[1..] { + assert_eq!(line.split(',').count(), n_cols); + } + + // All rows should have "Complete" state. + for line in &lines[1..] { + assert!(line.contains("Complete")); + } +} + +#[test] +fn csv_handles_pruned_trials() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + let x = FloatParam::new(0.0, 10.0).name("x"); + + // First trial: complete + let mut trial = study.create_trial(); + let _ = x.suggest(&mut trial).unwrap(); + study.complete_trial(trial, 1.0); + + // Second trial: pruned + let mut trial = study.create_trial(); + let _ = x.suggest(&mut trial).unwrap(); + study.prune_trial(trial); + + let mut buf = Vec::new(); + study.to_csv(&mut buf).unwrap(); + let csv = String::from_utf8(buf).unwrap(); + + let lines: Vec<&str> = csv.lines().collect(); + assert_eq!(lines.len(), 3); // header + 2 data rows + + // Pruned trial should have empty value. + let pruned_line = lines[2]; + assert!(pruned_line.contains("Pruned")); + // The value field (second column) should be empty. + let cols: Vec<&str> = pruned_line.split(',').collect(); + assert_eq!(cols[2], "Pruned"); + assert_eq!(cols[1], ""); // empty value for pruned +} + +#[test] +fn csv_handles_different_parameter_sets() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + let x = FloatParam::new(0.0, 10.0).name("x"); + let y = FloatParam::new(0.0, 10.0).name("y"); + + // First trial: only x + let mut trial = study.create_trial(); + let xv = x.suggest(&mut trial).unwrap(); + study.complete_trial(trial, xv); + + // Second trial: only y + let mut trial = study.create_trial(); + let yv = y.suggest(&mut trial).unwrap(); + study.complete_trial(trial, yv); + + let mut buf = Vec::new(); + study.to_csv(&mut buf).unwrap(); + let csv = String::from_utf8(buf).unwrap(); + + let lines: Vec<&str> = csv.lines().collect(); + assert_eq!(lines.len(), 3); + + // Both x and y columns should exist. + let header = lines[0]; + assert!(header.contains("x")); + assert!(header.contains("y")); + + // Each row has the right column count (missing params are empty). + let n_cols = header.split(',').count(); + for line in &lines[1..] { + assert_eq!(line.split(',').count(), n_cols); + } +} + +#[test] +fn csv_output_is_parseable() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + let lr = FloatParam::new(0.001, 0.1).name("learning_rate"); + let layers = IntParam::new(1, 5).name("n_layers"); + + study + .optimize(5, |trial| { + let l = lr.suggest(trial)?; + let n = layers.suggest(trial)?; + Ok::<_, optimizer::Error>(l * n as f64) + }) + .unwrap(); + + let mut buf = Vec::new(); + study.to_csv(&mut buf).unwrap(); + let csv = String::from_utf8(buf).unwrap(); + + // Parse each row: every value field should be a valid f64 for complete trials. + let lines: Vec<&str> = csv.lines().collect(); + for line in &lines[1..] { + let cols: Vec<&str> = line.split(',').collect(); + // trial_id should be a number + cols[0].parse::().unwrap(); + // value should be parseable as f64 + cols[1].parse::().unwrap(); + // state should be a known value + assert!(["Complete", "Pruned", "Failed", "Running"].contains(&cols[2])); + } +} + +#[test] +fn export_csv_writes_file() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + let x = FloatParam::new(0.0, 10.0).name("x"); + study + .optimize(3, |trial| { + let xv = x.suggest(trial)?; + Ok::<_, optimizer::Error>(xv * xv) + }) + .unwrap(); + + let dir = std::env::temp_dir().join("optimizer_export_test"); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("test_export.csv"); + + study.export_csv(&path).unwrap(); + + let contents = std::fs::read_to_string(&path).unwrap(); + assert!(contents.starts_with("trial_id,value,state")); + assert!(contents.lines().count() == 4); // header + 3 rows + + // Clean up. + let _ = std::fs::remove_dir_all(&dir); +} + +#[cfg(feature = "serde")] +#[test] +fn export_json_writes_file() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + let x = FloatParam::new(0.0, 10.0).name("x"); + study + .optimize(3, |trial| { + let xv = x.suggest(trial)?; + Ok::<_, optimizer::Error>(xv * xv) + }) + .unwrap(); + + let dir = std::env::temp_dir().join("optimizer_json_export_test"); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("test_export.json"); + + study.export_json(&path).unwrap(); + + let contents = std::fs::read_to_string(&path).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(&contents).unwrap(); + let arr = parsed.as_array().unwrap(); + assert_eq!(arr.len(), 3); + + // Each entry should have the expected fields. + for entry in arr { + assert!(entry.get("id").is_some()); + assert!(entry.get("value").is_some()); + assert!(entry.get("state").is_some()); + assert!(entry.get("params").is_some()); + } + + // Clean up. + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn csv_includes_user_attributes() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + let x = FloatParam::new(0.0, 10.0).name("x"); + + study + .optimize(2, |trial| { + let xv = x.suggest(trial)?; + trial.set_user_attr("training_time_secs", 45.2); + Ok::<_, optimizer::Error>(xv * xv) + }) + .unwrap(); + + let mut buf = Vec::new(); + study.to_csv(&mut buf).unwrap(); + let csv = String::from_utf8(buf).unwrap(); + + let header = csv.lines().next().unwrap(); + assert!(header.contains("training_time_secs")); +}