gavindiaz/rust-optimizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: unify optimize and optimize_with via blanket Objective impl

Browse Source 
- Add blanket `impl Objective<V> for Fn(&mut Trial) -> Result<V, E>`
  so closures work directly with `optimize`
- Rewrite optimize, optimize_async, optimize_parallel to accept
  `impl Objective<V>` with before_trial/after_trial hooks
- Remove optimize_with, optimize_with_async, optimize_with_parallel
- Remove max_retries and retry logic from Objective trait
- Add explicit closure type annotations for HRTB inference
- Convert FnMut test closures to Fn via RefCell/Cell
This commit is contained in:
Manuel Raimann
2026-02-12 13:09:14 +01:00
parent c20a53dfba
commit 47b5f9cec8
41 changed files with 316 additions and 793 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/study/builder.rs
+2 -2
View File
@@ -33,7 +33,7 @@ fn test_builder_with_sampler() {
let study: Study<f64> = Study::builder().sampler(TpeSampler::new()).build();
study
.optimize(10, |trial| {
.optimize(10, |trial: &mut optimizer::Trial| {
let val = x.suggest(trial)?;
Ok::<_, Error>(val * val)
})
@@ -77,7 +77,7 @@ fn test_builder_optimizes_correctly() {
.build();
study
.optimize(100, |trial| {
.optimize(100, |trial: &mut optimizer::Trial| {
let val = x.suggest(trial)?;
Ok::<_, Error>((val - 3.0) * (val - 3.0))
})
tests/study/enqueue.rs
+7 -5
View File
@@ -1,3 +1,4 @@
use std::cell::RefCell;
use std::collections::HashMap;
use optimizer::parameter::{FloatParam, IntParam, ParamValue, Parameter};
@@ -73,16 +74,17 @@ fn test_enqueue_with_optimize() {
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))]));
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
let mut values = Vec::new();
let values = RefCell::new(Vec::new());
study
.optimize(5, |trial| {
.optimize(5, |trial: &mut optimizer::Trial| {
let x_val = x.suggest(trial)?;
values.push(x_val);
values.borrow_mut().push(x_val);
Ok::<_, Error>(x_val * x_val)
})
.unwrap();
let values = values.into_inner();
// First two trials should use enqueued values
assert_eq!(values[0], 1.0);
assert_eq!(values[1], 2.0);
@@ -117,7 +119,7 @@ fn test_enqueue_trials_appear_in_completed_trials() {
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(7.0))]));
study
.optimize(1, |trial| {
.optimize(1, |trial: &mut optimizer::Trial| {
let x_val = x.suggest(trial)?;
Ok::<_, Error>(x_val)
})
@@ -178,7 +180,7 @@ fn test_enqueue_counted_in_n_trials() {
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
study
.optimize(5, |trial| {
.optimize(5, |trial: &mut optimizer::Trial| {
let x_val = x.suggest(trial)?;
Ok::<_, Error>(x_val)
})
tests/study/objective.rs
+11 -195
View File
@@ -30,7 +30,7 @@ fn test_callback_early_stopping() {
let study: Study<f64> = Study::new(Direction::Minimize);
study
.optimize_with(
.optimize(
100,
EarlyStopAfter5 {
x_param: FloatParam::new(0.0, 10.0),
@@ -69,7 +69,7 @@ fn test_callback_early_stopping_on_first_trial() {
let study: Study<f64> = Study::new(Direction::Minimize);
study
.optimize_with(
.optimize(
100,
StopImmediately {
x_param: FloatParam::new(0.0, 10.0),
@@ -109,7 +109,7 @@ fn test_callback_sampler_early_stopping() {
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
study
.optimize_with(
.optimize(
100,
StopAfter3 {
x_param: FloatParam::new(0.0, 10.0),
@@ -121,226 +121,42 @@ fn test_callback_sampler_early_stopping() {
}
#[test]
fn test_retries_successful_trials_not_retried() {
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
fn test_objective_struct_basic() {
use optimizer::Objective;
struct SuccessObj {
struct SquareObj {
x_param: FloatParam,
call_count: Arc<AtomicU32>,
}
impl Objective<f64> for SuccessObj {
impl Objective<f64> for SquareObj {
type Error = Error;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, Error> {
let x = self.x_param.suggest(trial)?;
self.call_count.fetch_add(1, Ordering::Relaxed);
Ok(x * x)
}
fn max_retries(&self) -> usize {
3
}
}
let study: Study<f64> = Study::new(Direction::Minimize);
let call_count = Arc::new(AtomicU32::new(0));
let obj = SuccessObj {
let obj = SquareObj {
x_param: FloatParam::new(0.0, 10.0),
call_count: Arc::clone(&call_count),
};
study.optimize_with(5, obj).unwrap();
study.optimize(5, obj).unwrap();
// All trials succeed on first try — exactly 5 calls
assert_eq!(call_count.load(Ordering::Relaxed), 5);
assert_eq!(study.n_trials(), 5);
}
#[test]
fn test_retries_failed_trials_retried_up_to_max() {
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use optimizer::Objective;
struct AlwaysFailObj {
x_param: FloatParam,
call_count: Arc<AtomicU32>,
}
impl Objective<f64> for AlwaysFailObj {
type Error = String;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, String> {
let _ = self.x_param.suggest(trial).map_err(|e| e.to_string())?;
self.call_count.fetch_add(1, Ordering::Relaxed);
Err("always fails".to_string())
}
fn max_retries(&self) -> usize {
3
}
}
let study: Study<f64> = Study::new(Direction::Minimize);
let call_count = Arc::new(AtomicU32::new(0));
let obj = AlwaysFailObj {
x_param: FloatParam::new(0.0, 10.0),
call_count: Arc::clone(&call_count),
};
let result = study.optimize_with(1, obj);
// 1 initial attempt + 3 retries = 4 total calls
assert_eq!(call_count.load(Ordering::Relaxed), 4);
// No trials completed
assert!(matches!(result, Err(Error::NoCompletedTrials)));
}
#[test]
fn test_retries_permanently_failed_after_exhaustion() {
use optimizer::Objective;
struct AlwaysFailObj {
x_param: FloatParam,
}
impl Objective<f64> for AlwaysFailObj {
type Error = String;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, String> {
let _ = self.x_param.suggest(trial).map_err(|e| e.to_string())?;
Err("transient error".to_string())
}
fn max_retries(&self) -> usize {
2
}
}
let study: Study<f64> = Study::new(Direction::Minimize);
let obj = AlwaysFailObj {
x_param: FloatParam::new(0.0, 10.0),
};
let result = study.optimize_with(3, obj);
assert!(
matches!(result, Err(Error::NoCompletedTrials)),
"all trials should permanently fail"
);
assert_eq!(
study.n_trials(),
0,
"no completed trials should be recorded"
);
}
#[test]
fn test_retries_uses_same_parameters() {
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};
use optimizer::Objective;
struct RetryObj {
x_param: FloatParam,
seen_values: Arc<Mutex<Vec<f64>>>,
call_count: Arc<AtomicU32>,
}
impl Objective<f64> for RetryObj {
type Error = String;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, String> {
let x = self.x_param.suggest(trial).map_err(|e| e.to_string())?;
self.seen_values.lock().unwrap().push(x);
let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1;
// Fail first two attempts, succeed on third
if count < 3 {
Err("transient".to_string())
} else {
Ok(x * x)
}
}
fn max_retries(&self) -> usize {
2
}
}
let study: Study<f64> = Study::new(Direction::Minimize);
let seen_values = Arc::new(Mutex::new(Vec::new()));
let call_count = Arc::new(AtomicU32::new(0));
let obj = RetryObj {
x_param: FloatParam::new(0.0, 10.0),
seen_values: Arc::clone(&seen_values),
call_count: Arc::clone(&call_count),
};
study.optimize_with(1, obj).unwrap();
let values = seen_values.lock().unwrap();
assert_eq!(values.len(), 3, "should be called 3 times (1 + 2 retries)");
// All three calls should have gotten the same parameter value
assert_eq!(values[0], values[1]);
assert_eq!(values[1], values[2]);
}
#[test]
fn test_retries_n_trials_counts_unique_configs() {
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use optimizer::Objective;
struct FailFirstObj {
x_param: FloatParam,
call_count: Arc<AtomicU32>,
}
impl Objective<f64> for FailFirstObj {
type Error = String;
fn evaluate(&self, trial: &mut Trial) -> Result<f64, String> {
let x = self.x_param.suggest(trial).map_err(|e| e.to_string())?;
let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1;
// Fail first attempt of each config, succeed on retry
if count % 2 == 1 {
Err("transient".to_string())
} else {
Ok(x * x)
}
}
fn max_retries(&self) -> usize {
2
}
}
let study: Study<f64> = Study::new(Direction::Minimize);
let call_count = Arc::new(AtomicU32::new(0));
let obj = FailFirstObj {
x_param: FloatParam::new(0.0, 10.0),
call_count: Arc::clone(&call_count),
};
study.optimize_with(3, obj).unwrap();
// 3 unique configs, each needing 2 calls = 6 total calls
assert_eq!(call_count.load(Ordering::Relaxed), 6);
// But only 3 completed trials
assert_eq!(study.n_trials(), 3);
}
#[test]
fn test_retries_with_zero_max_retries_same_as_optimize() {
let study: Study<f64> = Study::new(Direction::Minimize);
fn test_closure_and_objective_produce_same_results() {
let x_param = FloatParam::new(0.0, 10.0);
let call_count = std::cell::Cell::new(0u32);
let study: Study<f64> = Study::new(Direction::Minimize);
study
.optimize(5, |trial| {
.optimize(5, |trial: &mut Trial| {
let x = x_param.suggest(trial)?;
call_count.set(call_count.get() + 1);
Ok::<_, Error>(x * x)
})
.unwrap();
assert_eq!(call_count.get(), 5);
assert_eq!(study.n_trials(), 5);
}
tests/study/summary.rs
+2 -2
View File
@@ -8,7 +8,7 @@ fn test_summary_with_completed_trials() {
let x = FloatParam::new(0.0, 10.0).name("x");
study
.optimize(5, |trial| {
.optimize(5, |trial: &mut optimizer::Trial| {
let val = x.suggest(trial)?;
Ok::<_, Error>(val * val)
})
@@ -61,7 +61,7 @@ fn test_display_matches_summary() {
let x = FloatParam::new(0.0, 10.0).name("x");
study
.optimize(3, |trial| {
.optimize(3, |trial: &mut optimizer::Trial| {
let val = x.suggest(trial)?;
Ok::<_, Error>(val)
})
tests/study/workflow.rs
+16 -14
View File
@@ -8,7 +8,7 @@ fn test_study_basic_workflow() {
let x_param = FloatParam::new(-5.0, 5.0);
study
.optimize(10, |trial| {
.optimize(10, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x * x)
})
@@ -25,11 +25,11 @@ fn test_study_with_failures() {
let x_param = FloatParam::new(-5.0, 5.0);
// Every other trial fails
let mut counter = 0;
let counter = std::cell::Cell::new(0u32);
study
.optimize(10, |trial| {
counter += 1;
if counter % 2 == 0 {
.optimize(10, |trial: &mut optimizer::Trial| {
counter.set(counter.get() + 1);
if counter.get().is_multiple_of(2) {
return Err::<f64, &str>("intentional failure");
}
let x = x_param.suggest(trial).map_err(|_| "param error")?;
@@ -64,7 +64,7 @@ fn test_study_trials_iteration() {
let x_param = FloatParam::new(0.0, 1.0);
study
.optimize(5, |trial| {
.optimize(5, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x)
})
@@ -95,7 +95,7 @@ fn test_study_set_sampler() {
let x_param = FloatParam::new(-5.0, 5.0);
study
.optimize(10, |trial| {
.optimize(10, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x * x)
})
@@ -110,7 +110,7 @@ fn test_study_with_i32_value_type() {
let x_param = IntParam::new(-10, 10);
study
.optimize(10, |trial| {
.optimize(10, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x.abs() as i32)
})
@@ -125,7 +125,9 @@ fn test_study_with_i32_value_type() {
fn test_optimize_all_trials_fail() {
let study: Study<f64> = Study::new(Direction::Minimize);
let result = study.optimize(5, |_trial| Err::<f64, &str>("always fails"));
let result = study.optimize(5, |_trial: &mut optimizer::Trial| {
Err::<f64, &str>("always fails")
});
assert!(
matches!(result, Err(Error::NoCompletedTrials)),
@@ -139,7 +141,7 @@ fn test_best_value() {
let x_param = FloatParam::new(0.0, 10.0);
study
.optimize(10, |trial| {
.optimize(10, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x)
})
@@ -160,7 +162,7 @@ fn test_best_trial_with_nan_values() {
let x_param = FloatParam::new(0.0, 10.0);
study
.optimize(5, |trial| {
.optimize(5, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x)
})
@@ -199,7 +201,7 @@ fn test_multiple_params_in_optimization() {
let n_param = IntParam::new(1, 5);
study
.optimize(10, |trial| {
.optimize(10, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
let n = n_param.suggest(trial)?;
Ok::<_, Error>(x * x + n as f64)
@@ -216,7 +218,7 @@ fn test_suggest_bool_in_optimization() {
let x_param = FloatParam::new(0.0, 10.0);
study
.optimize(10, |trial| {
.optimize(10, |trial: &mut optimizer::Trial| {
let use_feature = use_feature_param.suggest(trial)?;
let x = x_param.suggest(trial)?;
@@ -235,7 +237,7 @@ fn test_completed_trial_get() {
let n_param = IntParam::new(1, 10).name("n");
study
.optimize(5, |trial| {
.optimize(5, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
let n = n_param.suggest(trial)?;
Ok::<_, Error>(x * x + n as f64)