9b4395b0df
increment_trial() was never called — budget enforcement was broken because TrialBudgetObserver had its own internal Arc<AtomicUsize> disconnected from the trial_counter used by evaluate_point() and cost(). With num_trials=2, PSO ran 20 particle evaluations instead of stopping at 2. Now TrialBudgetObserver::new() takes the caller-owned Arc<AtomicUsize> so all evaluation paths (LHS via evaluate_point + PSO via cost()) share one counter. ObjectiveFunction::cost() and ParallelObjectiveFunction::cost() increment trial_counter directly (fetch_add SeqCst) and check budget inline; observe_iter() reads the same counter to terminate PSO between iterations. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
59 lines
1.8 KiB
Rust
59 lines
1.8 KiB
Rust
use argmin::core::observers::Observe;
|
|
use argmin::core::{Error, State};
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::sync::Arc;
|
|
|
|
/// Observer that enforces strict trial budget limits
|
|
///
|
|
/// Argmin's PSO does not reliably respect max_iters when set to small values (e.g., 1).
|
|
/// This observer shares the same `Arc<AtomicUsize>` counter that the cost function and
|
|
/// LHS evaluation phase use, so every evaluation (LHS + PSO) counts toward the budget.
|
|
#[derive(Clone, Debug)]
|
|
pub struct TrialBudgetObserver {
|
|
max_trials: usize,
|
|
trial_counter: Arc<AtomicUsize>,
|
|
}
|
|
|
|
impl TrialBudgetObserver {
|
|
/// Create a new observer that shares the caller-owned trial counter.
|
|
///
|
|
/// The same `trial_counter` must be passed to the cost function and incremented
|
|
/// on every evaluation (LHS and PSO) so the budget is accurately tracked.
|
|
pub fn new(max_trials: usize, trial_counter: Arc<AtomicUsize>) -> Self {
|
|
Self {
|
|
max_trials,
|
|
trial_counter,
|
|
}
|
|
}
|
|
|
|
pub fn should_terminate(&self) -> bool {
|
|
self.trial_counter.load(Ordering::SeqCst) >= self.max_trials
|
|
}
|
|
|
|
pub fn get_trials_used(&self) -> usize {
|
|
self.trial_counter.load(Ordering::SeqCst)
|
|
}
|
|
}
|
|
|
|
impl<I> Observe<I> for TrialBudgetObserver
|
|
where
|
|
I: State,
|
|
{
|
|
fn observe_iter(&mut self, _state: &I, _kv: &argmin::core::KV) -> Result<(), Error> {
|
|
let used = self.trial_counter.load(Ordering::SeqCst);
|
|
if used >= self.max_trials {
|
|
tracing::warn!(
|
|
"Trial budget exhausted: {}/{} trials used. Terminating optimization.",
|
|
used,
|
|
self.max_trials
|
|
);
|
|
return Err(Error::msg(format!(
|
|
"Trial budget exhausted: {used}/{}",
|
|
self.max_trials
|
|
)));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|