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.
This commit is contained in:
+147
@@ -1732,6 +1732,119 @@ impl<V> Study<V>
|
|||||||
where
|
where
|
||||||
V: PartialOrd + Clone + fmt::Display,
|
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_<id>`).
|
||||||
|
/// 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<ParamId, String> = 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<String> = 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<ParamId> = 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::path::Path>) -> 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.
|
/// Returns a human-readable summary of the study.
|
||||||
///
|
///
|
||||||
/// The summary includes:
|
/// The summary includes:
|
||||||
@@ -2073,6 +2186,19 @@ pub struct StudySnapshot<V> {
|
|||||||
|
|
||||||
#[cfg(feature = "serde")]
|
#[cfg(feature = "serde")]
|
||||||
impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
|
impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
|
||||||
|
/// 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::path::Path>) -> 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.
|
/// Saves the study state to a JSON file.
|
||||||
///
|
///
|
||||||
/// # Errors
|
/// # Errors
|
||||||
@@ -2171,3 +2297,24 @@ fn is_trial_pruned<E: 'static>(e: &E) -> bool {
|
|||||||
any.downcast_ref::<crate::error::TrialPruned>().is_some()
|
any.downcast_ref::<crate::error::TrialPruned>().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(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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<f64> = 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<f64> = 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<f64> = 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<f64> = 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<f64> = 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::<u64>().unwrap();
|
||||||
|
// value should be parseable as f64
|
||||||
|
cols[1].parse::<f64>().unwrap();
|
||||||
|
// state should be a known value
|
||||||
|
assert!(["Complete", "Pruned", "Failed", "Running"].contains(&cols[2]));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn export_csv_writes_file() {
|
||||||
|
let study: Study<f64> = 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<f64> = 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<f64> = 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"));
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user