diff --git a/src/sampler/bohb.rs b/src/sampler/bohb.rs index a7121c1..2444e8f 100644 --- a/src/sampler/bohb.rs +++ b/src/sampler/bohb.rs @@ -178,6 +178,7 @@ impl BohbSampler { intermediate_values: trial.intermediate_values.clone(), state: trial.state, user_attrs: trial.user_attrs.clone(), + constraints: trial.constraints.clone(), }) }) .collect() @@ -396,6 +397,7 @@ mod tests { intermediate_values, state: TrialState::Complete, user_attrs: HashMap::new(), + constraints: Vec::new(), } } diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index a7ca85c..883d834 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -41,6 +41,9 @@ pub struct CompletedTrial { pub state: TrialState, /// User-defined attributes stored during the trial. pub user_attrs: HashMap, + /// Constraint values for this trial (<=0.0 means feasible). + #[cfg_attr(feature = "serde", serde(default))] + pub constraints: Vec, } impl CompletedTrial { @@ -61,6 +64,7 @@ impl CompletedTrial { intermediate_values: Vec::new(), state: TrialState::Complete, user_attrs: HashMap::new(), + constraints: Vec::new(), } } @@ -83,6 +87,7 @@ impl CompletedTrial { intermediate_values, state: TrialState::Complete, user_attrs, + constraints: Vec::new(), } } @@ -127,6 +132,14 @@ impl CompletedTrial { }) } + /// Returns `true` if all constraints are satisfied (values <= 0.0). + /// + /// A trial with no constraints is considered feasible. + #[must_use] + pub fn is_feasible(&self) -> bool { + self.constraints.iter().all(|&c| c <= 0.0) + } + /// Gets a user attribute by key. #[must_use] pub fn user_attr(&self, key: &str) -> Option<&AttrValue> { diff --git a/src/study.rs b/src/study.rs index 82fdacf..a5d2dba 100644 --- a/src/study.rs +++ b/src/study.rs @@ -343,18 +343,11 @@ where /// Returns the trial ID of the current best trial from the given slice. #[cfg(feature = "tracing")] fn best_id(&self, trials: &[CompletedTrial]) -> Option { + let direction = self.direction; trials .iter() .filter(|t| t.state == TrialState::Complete) - .max_by(|a, b| { - let ordering = a.value.partial_cmp(&b.value); - match self.direction { - Direction::Minimize => { - ordering.map_or(core::cmp::Ordering::Equal, core::cmp::Ordering::reverse) - } - Direction::Maximize => ordering.unwrap_or(core::cmp::Ordering::Equal), - } - }) + .max_by(|a, b| Self::compare_trials(a, b, direction)) .map(|t| t.id) } @@ -461,6 +454,7 @@ where trial.user_attrs().clone(), ); completed.state = TrialState::Complete; + completed.constraints = trial.constraint_values().to_vec(); self.completed_trials.write().push(completed); } @@ -570,6 +564,7 @@ where trial.user_attrs().clone(), ); completed.state = TrialState::Pruned; + completed.constraints = trial.constraint_values().to_vec(); self.completed_trials.write().push(completed); } @@ -636,12 +631,46 @@ where .count() } + /// Compares two completed trials using constraint-aware ranking. + /// + /// 1. Feasible trials always rank above infeasible trials. + /// 2. Among feasible trials, rank by objective value (respecting direction). + /// 3. Among infeasible trials, rank by total constraint violation (lower is better). + fn compare_trials( + a: &CompletedTrial, + b: &CompletedTrial, + direction: Direction, + ) -> core::cmp::Ordering { + match (a.is_feasible(), b.is_feasible()) { + (true, false) => core::cmp::Ordering::Greater, + (false, true) => core::cmp::Ordering::Less, + (false, false) => { + let va: f64 = a.constraints.iter().map(|c| c.max(0.0)).sum(); + let vb: f64 = b.constraints.iter().map(|c| c.max(0.0)).sum(); + vb.partial_cmp(&va).unwrap_or(core::cmp::Ordering::Equal) + } + (true, true) => { + let ordering = a.value.partial_cmp(&b.value); + match direction { + Direction::Minimize => { + ordering.map_or(core::cmp::Ordering::Equal, core::cmp::Ordering::reverse) + } + Direction::Maximize => ordering.unwrap_or(core::cmp::Ordering::Equal), + } + } + } + } + /// Returns the trial with the best objective value. /// /// The "best" trial depends on the optimization direction: /// - `Direction::Minimize`: Returns the trial with the lowest objective value. /// - `Direction::Maximize`: Returns the trial with the highest objective value. /// + /// When constraints are present, feasible trials always rank above infeasible + /// trials. Among infeasible trials, those with lower total constraint violation + /// are preferred. + /// /// # Errors /// /// Returns `Error::NoCompletedTrials` if no trials have been completed. @@ -675,25 +704,12 @@ where V: Clone, { let trials = self.completed_trials.read(); + let direction = self.direction; let best = trials .iter() .filter(|t| t.state == TrialState::Complete) - .max_by(|a, b| { - // For Minimize, we want the smallest value to be "max" in ordering - // For Maximize, we want the largest value to be "max" in ordering - let ordering = a.value.partial_cmp(&b.value); - match self.direction { - Direction::Minimize => { - // Reverse ordering: smaller values are "greater" for max_by - ordering.map_or(core::cmp::Ordering::Equal, core::cmp::Ordering::reverse) - } - Direction::Maximize => { - // Normal ordering: larger values are "greater" for max_by - ordering.unwrap_or(core::cmp::Ordering::Equal) - } - } - }) + .max_by(|a, b| Self::compare_trials(a, b, direction)) .ok_or(crate::Error::NoCompletedTrials)?; Ok(best.clone()) @@ -752,21 +768,14 @@ where V: Clone, { let trials = self.completed_trials.read(); + let direction = self.direction; let mut completed: Vec<_> = trials .iter() .filter(|t| t.state == TrialState::Complete) .cloned() .collect(); - completed.sort_by(|a, b| match self.direction { - Direction::Minimize => a - .value - .partial_cmp(&b.value) - .unwrap_or(core::cmp::Ordering::Equal), - Direction::Maximize => b - .value - .partial_cmp(&a.value) - .unwrap_or(core::cmp::Ordering::Equal), - }); + // Sort best-first: reverse the compare_trials ordering (which is designed for max_by) + completed.sort_by(|a, b| Self::compare_trials(b, a, direction)); completed.truncate(n); completed } diff --git a/src/trial.rs b/src/trial.rs index cbb7f65..4f90fd3 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -89,6 +89,8 @@ pub struct Trial { user_attrs: HashMap, /// Pre-filled parameter values from enqueue (used instead of sampling). fixed_params: HashMap, + /// Constraint values for this trial (<=0.0 means feasible). + constraint_values: Vec, } impl core::fmt::Debug for Trial { @@ -105,6 +107,7 @@ impl core::fmt::Debug for Trial { .field("has_pruner", &self.pruner.is_some()) .field("user_attrs", &self.user_attrs) .field("fixed_params", &self.fixed_params) + .field("constraint_values", &self.constraint_values) .finish() } } @@ -144,6 +147,7 @@ impl Trial { pruner: None, user_attrs: HashMap::new(), fixed_params: HashMap::new(), + constraint_values: Vec::new(), } } @@ -175,6 +179,7 @@ impl Trial { pruner: Some(pruner), user_attrs: HashMap::new(), fixed_params: HashMap::new(), + constraint_values: Vec::new(), } } @@ -292,6 +297,20 @@ impl Trial { &self.user_attrs } + /// Sets constraint values for this trial. + /// + /// Each value represents a constraint; a value <= 0.0 means the constraint + /// is satisfied (feasible). A value > 0.0 means the constraint is violated. + pub fn set_constraints(&mut self, values: Vec) { + self.constraint_values = values; + } + + /// Returns the constraint values for this trial. + #[must_use] + pub fn constraint_values(&self) -> &[f64] { + &self.constraint_values + } + /// Sets the trial state to Complete. pub(crate) fn set_complete(&mut self) { self.state = TrialState::Complete; diff --git a/tests/integration.rs b/tests/integration.rs index 618af9e..0167650 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -2141,3 +2141,107 @@ fn test_into_iterator_preserves_insertion_order() { let ids: Vec = (&study).into_iter().map(|t| t.id).collect(); assert_eq!(ids, vec![0, 1, 2]); } + +// ============================================================================= +// Tests: Constraint handling +// ============================================================================= + +#[test] +fn test_is_feasible_all_satisfied() { + let study: Study = Study::new(Direction::Minimize); + let mut trial = study.create_trial(); + trial.set_constraints(vec![-1.0, 0.0, -0.5]); + study.complete_trial(trial, 1.0); + + let completed = study.best_trial().unwrap(); + assert!(completed.is_feasible()); +} + +#[test] +fn test_is_feasible_one_violated() { + let study: Study = Study::new(Direction::Minimize); + let mut trial = study.create_trial(); + trial.set_constraints(vec![-1.0, 0.5, -0.5]); + study.complete_trial(trial, 1.0); + + let completed = study.best_trial().unwrap(); + assert!(!completed.is_feasible()); +} + +#[test] +fn test_is_feasible_empty_constraints() { + let study: Study = Study::new(Direction::Minimize); + let trial = study.create_trial(); + study.complete_trial(trial, 1.0); + + let completed = study.best_trial().unwrap(); + assert!(completed.is_feasible()); +} + +#[test] +fn test_best_trial_prefers_feasible() { + let study: Study = Study::new(Direction::Minimize); + + // Infeasible trial with better objective + let mut trial1 = study.create_trial(); + trial1.set_constraints(vec![1.0]); + study.complete_trial(trial1, 0.1); + + // Feasible trial with worse objective + let mut trial2 = study.create_trial(); + trial2.set_constraints(vec![-1.0]); + study.complete_trial(trial2, 100.0); + + let best = study.best_trial().unwrap(); + assert_eq!(best.id, 1); // feasible trial wins + assert_eq!(best.value, 100.0); +} + +#[test] +fn test_best_trial_feasible_by_objective() { + let study: Study = Study::new(Direction::Minimize); + + // Feasible, worse objective + let mut trial1 = study.create_trial(); + trial1.set_constraints(vec![-1.0]); + study.complete_trial(trial1, 10.0); + + // Feasible, better objective + let mut trial2 = study.create_trial(); + trial2.set_constraints(vec![-0.5]); + study.complete_trial(trial2, 2.0); + + let best = study.best_trial().unwrap(); + assert_eq!(best.id, 1); // lower objective wins among feasible + assert_eq!(best.value, 2.0); +} + +#[test] +fn test_top_trials_ranks_feasible_above_infeasible() { + let study: Study = Study::new(Direction::Minimize); + + // Infeasible, low violation + let mut t0 = study.create_trial(); + t0.set_constraints(vec![0.5]); + study.complete_trial(t0, 1.0); + + // Feasible, worst objective among feasible + let mut t1 = study.create_trial(); + t1.set_constraints(vec![-1.0]); + study.complete_trial(t1, 50.0); + + // Feasible, best objective among feasible + let mut t2 = study.create_trial(); + t2.set_constraints(vec![-0.1]); + study.complete_trial(t2, 5.0); + + // Infeasible, high violation + let mut t3 = study.create_trial(); + t3.set_constraints(vec![3.0]); + study.complete_trial(t3, 0.5); + + let top = study.top_trials(4); + let ids: Vec = top.iter().map(|t| t.id).collect(); + // Feasible sorted by objective first (5.0, 50.0), then infeasible by violation (0.5, 3.0) + assert_eq!(ids, vec![2, 1, 0, 3]); +}