feat: add trial user attributes for logging and analysis

Add AttrValue enum (Float, Int, String, Bool) with From impls and
set_user_attr/user_attr/user_attrs methods on both Trial and
CompletedTrial. Attributes set during optimization are propagated
through complete_trial and prune_trial. AttrValue is re-exported
at crate root and in the prelude.

src/lib.rs

@@ -210,7 +210,7 @@ pub use sampler::grid::GridSearchSampler;
 pub use sampler::random::RandomSampler;
 pub use sampler::tpe::TpeSampler;
 pub use study::Study;
-pub use trial::Trial;
+pub use trial::{AttrValue, Trial};
 pub use types::{Direction, TrialState};
 
 /// Convenient wildcard import for the most common types.
@@ -236,6 +236,6 @@ pub mod prelude {
     pub use crate::sampler::random::RandomSampler;
     pub use crate::sampler::tpe::TpeSampler;
     pub use crate::study::Study;
-    pub use crate::trial::Trial;
+    pub use crate::trial::{AttrValue, Trial};
     pub use crate::types::Direction;
 }

src/pruner/hyperband.rs

@@ -278,6 +278,7 @@ mod tests {
             HashMap::new(),
             0.0,
             values.to_vec(),
+            HashMap::new(),
         )
     }
 

src/pruner/percentile.rs

@@ -204,6 +204,7 @@ mod tests {
             HashMap::new(),
             0.0,
             values.to_vec(),
+            HashMap::new(),
         )
     }
 

src/pruner/successive_halving.rs

@@ -240,6 +240,7 @@ mod tests {
             HashMap::new(),
             0.0,
             values.to_vec(),
+            HashMap::new(),
         )
     }
 

src/pruner/wilcoxon.rs

@@ -311,6 +311,7 @@ mod tests {
             HashMap::new(),
             value,
             intermediate_values,
+            HashMap::new(),
         )
     }
 

src/sampler/mod.rs

@@ -9,6 +9,7 @@ use std::collections::HashMap;
 use crate::distribution::Distribution;
 use crate::param::ParamValue;
 use crate::parameter::{ParamId, Parameter};
+use crate::trial::AttrValue;
 use crate::types::TrialState;
 
 /// A completed trial with its parameters, distributions, and objective value.
@@ -32,6 +33,8 @@ pub struct CompletedTrial<V = f64> {
     pub intermediate_values: Vec<(u64, f64)>,
     /// The state of the trial (Complete, Pruned, or Failed).
     pub state: TrialState,
+    /// User-defined attributes stored during the trial.
+    pub user_attrs: HashMap<String, AttrValue>,
 }
 
 impl<V> CompletedTrial<V> {
@@ -51,10 +54,11 @@ impl<V> CompletedTrial<V> {
             value,
             intermediate_values: Vec::new(),
             state: TrialState::Complete,
+            user_attrs: HashMap::new(),
         }
     }
 
-    /// Creates a new completed trial with intermediate values.
+    /// Creates a new completed trial with intermediate values and user attributes.
     pub fn with_intermediate_values(
         id: u64,
         params: HashMap<ParamId, ParamValue>,
@@ -62,6 +66,7 @@ impl<V> CompletedTrial<V> {
         param_labels: HashMap<ParamId, String>,
         value: V,
         intermediate_values: Vec<(u64, f64)>,
+        user_attrs: HashMap<String, AttrValue>,
     ) -> Self {
         Self {
             id,
@@ -71,6 +76,7 @@ impl<V> CompletedTrial<V> {
             value,
             intermediate_values,
             state: TrialState::Complete,
+            user_attrs,
         }
     }
 
@@ -114,6 +120,18 @@ impl<V> CompletedTrial<V> {
                 .expect("parameter type mismatch: stored value incompatible with parameter")
         })
     }
+
+    /// Gets a user attribute by key.
+    #[must_use]
+    pub fn user_attr(&self, key: &str) -> Option<&AttrValue> {
+        self.user_attrs.get(key)
+    }
+
+    /// Returns all user attributes.
+    #[must_use]
+    pub fn user_attrs(&self) -> &HashMap<String, AttrValue> {
+        &self.user_attrs
+    }
 }
 
 /// A pending (running) trial with its parameters and distributions, but no objective value yet.

src/study.rs

@@ -363,6 +363,7 @@ where
             trial.param_labels().clone(),
             value,
             trial.intermediate_values().to_vec(),
+            trial.user_attrs().clone(),
         );
         completed.state = TrialState::Complete;
         self.completed_trials.write().push(completed);
@@ -471,6 +472,7 @@ where
             trial.param_labels().clone(),
             V::default(),
             trial.intermediate_values().to_vec(),
+            trial.user_attrs().clone(),
         );
         completed.state = TrialState::Pruned;
         self.completed_trials.write().push(completed);

src/trial.rs

@@ -13,6 +13,49 @@ use crate::pruner::Pruner;
 use crate::sampler::{CompletedTrial, Sampler};
 use crate::types::TrialState;
 
+/// A user attribute value that can be stored on a trial.
+#[derive(Clone, Debug, PartialEq)]
+pub enum AttrValue {
+    /// A floating-point attribute.
+    Float(f64),
+    /// An integer attribute.
+    Int(i64),
+    /// A string attribute.
+    String(String),
+    /// A boolean attribute.
+    Bool(bool),
+}
+
+impl From<f64> for AttrValue {
+    fn from(v: f64) -> Self {
+        Self::Float(v)
+    }
+}
+
+impl From<i64> for AttrValue {
+    fn from(v: i64) -> Self {
+        Self::Int(v)
+    }
+}
+
+impl From<String> for AttrValue {
+    fn from(v: String) -> Self {
+        Self::String(v)
+    }
+}
+
+impl From<&str> for AttrValue {
+    fn from(v: &str) -> Self {
+        Self::String(v.to_owned())
+    }
+}
+
+impl From<bool> for AttrValue {
+    fn from(v: bool) -> Self {
+        Self::Bool(v)
+    }
+}
+
 /// A trial represents a single evaluation of the objective function.
 ///
 /// Each trial has a unique ID and stores the sampled parameters along with
@@ -41,6 +84,8 @@ pub struct Trial {
     intermediate_values: Vec<(u64, f64)>,
     /// The pruner used to decide whether to stop this trial early.
     pruner: Option<Arc<dyn Pruner>>,
+    /// User-defined attributes for logging, debugging, and analysis.
+    user_attrs: HashMap<String, AttrValue>,
 }
 
 impl core::fmt::Debug for Trial {
@@ -55,6 +100,7 @@ impl core::fmt::Debug for Trial {
             .field("has_history", &self.history.is_some())
             .field("intermediate_values", &self.intermediate_values)
             .field("has_pruner", &self.pruner.is_some())
+            .field("user_attrs", &self.user_attrs)
             .finish()
     }
 }
@@ -92,6 +138,7 @@ impl Trial {
             history: None,
             intermediate_values: Vec::new(),
             pruner: None,
+            user_attrs: HashMap::new(),
         }
     }
 
@@ -121,6 +168,7 @@ impl Trial {
             history: Some(history),
             intermediate_values: Vec::new(),
             pruner: Some(pruner),
+            user_attrs: HashMap::new(),
         }
     }
 
@@ -209,6 +257,23 @@ impl Trial {
         &self.intermediate_values
     }
 
+    /// Sets a user attribute on this trial.
+    pub fn set_user_attr(&mut self, key: impl Into<String>, value: impl Into<AttrValue>) {
+        self.user_attrs.insert(key.into(), value.into());
+    }
+
+    /// Gets a user attribute by key.
+    #[must_use]
+    pub fn user_attr(&self, key: &str) -> Option<&AttrValue> {
+        self.user_attrs.get(key)
+    }
+
+    /// Returns all user attributes.
+    #[must_use]
+    pub fn user_attrs(&self) -> &HashMap<String, AttrValue> {
+        &self.user_attrs
+    }
+
     /// Sets the trial state to Complete.
     pub(crate) fn set_complete(&mut self) {
         self.state = TrialState::Complete;

tests/median_pruner_tests.rs

@@ -13,6 +13,7 @@ fn trial_with_values(id: u64, intermediate_values: Vec<(u64, f64)>) -> Completed
         HashMap::new(),
         0.0,
         intermediate_values,
+        HashMap::new(),
     )
 }
 

tests/user_attr_tests.rs

@@ -0,0 +1,148 @@
+use optimizer::parameter::{FloatParam, Parameter};
+use optimizer::{AttrValue, Direction, Study};
+
+#[test]
+fn set_and_get_float_attr() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+    let x = FloatParam::new(0.0, 1.0);
+
+    study
+        .optimize(1, |trial| {
+            let _ = x.suggest(trial)?;
+            trial.set_user_attr("score", 42.5);
+            assert_eq!(trial.user_attr("score"), Some(&AttrValue::Float(42.5)));
+            Ok::<_, optimizer::Error>(1.0)
+        })
+        .unwrap();
+}
+
+#[test]
+fn set_and_get_int_attr() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+    let x = FloatParam::new(0.0, 1.0);
+
+    study
+        .optimize(1, |trial| {
+            let _ = x.suggest(trial)?;
+            trial.set_user_attr("epoch", 42_i64);
+            assert_eq!(trial.user_attr("epoch"), Some(&AttrValue::Int(42)));
+            Ok::<_, optimizer::Error>(1.0)
+        })
+        .unwrap();
+}
+
+#[test]
+fn set_and_get_string_attr() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+    let x = FloatParam::new(0.0, 1.0);
+
+    study
+        .optimize(1, |trial| {
+            let _ = x.suggest(trial)?;
+            trial.set_user_attr("model", "resnet50");
+            assert_eq!(
+                trial.user_attr("model"),
+                Some(&AttrValue::String("resnet50".to_owned()))
+            );
+            Ok::<_, optimizer::Error>(1.0)
+        })
+        .unwrap();
+}
+
+#[test]
+fn set_and_get_bool_attr() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+    let x = FloatParam::new(0.0, 1.0);
+
+    study
+        .optimize(1, |trial| {
+            let _ = x.suggest(trial)?;
+            trial.set_user_attr("converged", true);
+            assert_eq!(trial.user_attr("converged"), Some(&AttrValue::Bool(true)));
+            Ok::<_, optimizer::Error>(1.0)
+        })
+        .unwrap();
+}
+
+#[test]
+fn attrs_propagate_to_completed_trial() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+    let x = FloatParam::new(0.0, 1.0);
+
+    study
+        .optimize(1, |trial| {
+            let _ = x.suggest(trial)?;
+            trial.set_user_attr("time_secs", 1.5);
+            trial.set_user_attr("tag", "baseline");
+            Ok::<_, optimizer::Error>(1.0)
+        })
+        .unwrap();
+
+    let best = study.best_trial().unwrap();
+    assert_eq!(best.user_attr("time_secs"), Some(&AttrValue::Float(1.5)));
+    assert_eq!(
+        best.user_attr("tag"),
+        Some(&AttrValue::String("baseline".to_owned()))
+    );
+}
+
+#[test]
+fn overwrite_attr_replaces_value() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+    let x = FloatParam::new(0.0, 1.0);
+
+    study
+        .optimize(1, |trial| {
+            let _ = x.suggest(trial)?;
+            trial.set_user_attr("key", "old");
+            trial.set_user_attr("key", "new");
+            assert_eq!(
+                trial.user_attr("key"),
+                Some(&AttrValue::String("new".to_owned()))
+            );
+            Ok::<_, optimizer::Error>(1.0)
+        })
+        .unwrap();
+
+    let best = study.best_trial().unwrap();
+    assert_eq!(
+        best.user_attr("key"),
+        Some(&AttrValue::String("new".to_owned()))
+    );
+}
+
+#[test]
+fn missing_attr_returns_none() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+    let x = FloatParam::new(0.0, 1.0);
+
+    study
+        .optimize(1, |trial| {
+            let _ = x.suggest(trial)?;
+            assert_eq!(trial.user_attr("nonexistent"), None);
+            Ok::<_, optimizer::Error>(1.0)
+        })
+        .unwrap();
+
+    let best = study.best_trial().unwrap();
+    assert_eq!(best.user_attr("nonexistent"), None);
+}
+
+#[test]
+fn user_attrs_map_returns_all() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+    let x = FloatParam::new(0.0, 1.0);
+
+    study
+        .optimize(1, |trial| {
+            let _ = x.suggest(trial)?;
+            trial.set_user_attr("a", 1.0);
+            trial.set_user_attr("b", true);
+            assert_eq!(trial.user_attrs().len(), 2);
+            Ok::<_, optimizer::Error>(1.0)
+        })
+        .unwrap();
+
+    let best = study.best_trial().unwrap();
+    assert_eq!(best.user_attrs().len(), 2);
+}