diff --git a/crates/ml-hyperopt/src/optimizer.rs b/crates/ml-hyperopt/src/optimizer.rs index 43f3d0d17..d8ec488a5 100644 --- a/crates/ml-hyperopt/src/optimizer.rs +++ b/crates/ml-hyperopt/src/optimizer.rs @@ -354,15 +354,35 @@ impl ArgminOptimizer { let trials_used = trials.len(); let remaining_trials = self.max_trials.saturating_sub(trials_used); - // FIX: Each PSO iteration evaluates n_particles candidates (sequentially via mutex) - // PSO evaluates ALL particles in swarm per iteration, so divide remaining budget - // by swarm size to prevent trial count overflow (fixes 962 trial bug) - // CRITICAL FIX (2025-11-07): Use CEILING division to ensure all trials complete - // Example: 8 remaining ÷ 20 particles = 0.4 → ceil to 1 iteration (not 0) + // For small budgets (< 2× swarm size), PSO can't complete a full iteration. + // Fall back to additional LHS samples which give better space coverage. + if remaining_trials < self.n_particles * 2 { + info!( + "Budget too small for PSO ({} remaining < {} particles × 2). Using additional LHS.", + remaining_trials, self.n_particles + ); + let extra_samples = Self::latin_hypercube_sampling(remaining_trials, &bounds, &mut rng); + let mut model_guard = cost_fn.model.lock().map_err(|e| anyhow::anyhow!("lock: {e}"))?; + for i in 0..remaining_trials { + let continuous_vec: Vec = extra_samples.row(i).to_vec(); + Self::evaluate_point( + &continuous_vec, + &mut *model_guard, + &trial_results, + &trial_counter, + ¶m_names, + )?; + } + drop(model_guard); + } + let max_iters_by_budget = ((remaining_trials as f64) / (self.n_particles as f64)).ceil() as usize; - - let max_iters = max_iters_by_budget.min(self.max_iters_per_restart); + let max_iters = if remaining_trials < self.n_particles * 2 { + 0 // Already handled by LHS fallback above + } else { + max_iters_by_budget.min(self.max_iters_per_restart) + }; info!( "PSO Budget: {} iterations ({} remaining trials ÷ {} particles = {} max iters)",