Merge pull request #11 from MenloSystems/feature/out-of-bounds-resample

pengowen123 · web-flow · commit 2bb2b3f53495 · 2025-05-19T08:51:38.000Z
Resample when a point is out-of-bounds
diff --git a/src/lib.rs b/src/lib.rs
@@ -128,6 +128,8 @@ pub use crate::options::CMAESOptions;
 pub use crate::parameters::Weights;
 #[cfg(feature = "plotters")]
 pub use crate::plotting::PlotOptions;
+pub use crate::sampling::Bounds;
+pub use crate::sampling::Constraints;
 pub use crate::termination::TerminationReason;
 
 use std::f64;
@@ -243,18 +245,21 @@ impl<F> CMAES<F> {
             return Err(InvalidOptionsError::Cm);
         }
 
-        // Initialize point sampler
         let seed = options.seed.unwrap_or_else(rand::random);
+
+        // Initialize constant parameters according to the options
+        let parameters = Parameters::from_options(&options, seed);
+
+        // Initialize point sampler
         let sampler = Sampler::new(
             dimensions,
+            options.constraints,
+            options.max_resamples,
             options.population_size,
             objective_function,
             seed,
         );
 
-        // Initialize constant parameters according to the options
-        let parameters = Parameters::from_options(&options, seed);
-
         // Initialize variable parameters
         let state = State::new(options.initial_mean, options.initial_step_size);
 
diff --git a/src/options.rs b/src/options.rs
@@ -6,6 +6,8 @@ use std::time::Duration;
 
 use crate::mode::Mode;
 use crate::parameters::Weights;
+use crate::sampling::Bounds;
+use crate::sampling::Constraints;
 #[cfg(feature = "plotters")]
 use crate::PlotOptions;
 use crate::CMAES;
@@ -52,6 +54,12 @@ pub struct CMAESOptions {
     /// [`BIPOP`][crate::restart::BIPOP] restart strategies (see the [`restart`][crate::restart]
     /// module).
     pub population_size: usize,
+    /// The constraints for the search. Resamples until all samples satisfy the constraints or
+    /// `max_resamples` are hit. Default value is `None`.
+    pub constraints: Option<Box<dyn Constraints>>,
+    /// How many times to resample points in order to stay inside `bounds`.
+    /// `None` disables the limit. Defaults to 10.
+    pub max_resamples: Option<usize>,
     /// The distribution to use when assigning weights to individuals. Default value is
     /// [`Weights::Negative`].
     pub weights: Weights,
@@ -142,6 +150,8 @@ impl CMAESOptions {
             initial_mean,
             initial_step_size,
             population_size: 4 + (3.0 * (dimensions as f64).ln()).floor() as usize,
+            constraints: None,
+            max_resamples: Some(10),
             weights: Weights::Negative,
             parallel_update: false,
             cm: 1.0,
@@ -187,6 +197,25 @@ impl CMAESOptions {
         self
     }
 
+    /// Changes the bounds from the defalt value. Vector length must match the number of dimensions.
+    /// Convenience method for constraints().
+    pub fn bounds(mut self, lower: Vec<f64>, upper: Vec<f64>) -> Self {
+        self.constraints = Some(Box::new(Bounds { lower, upper }));
+        self
+    }
+
+    /// Changes the constraints from the defalt value.
+    pub fn constraints(mut self, constraints: Box<dyn Constraints>) -> Self {
+        self.constraints = Some(constraints);
+        self
+    }
+
+    /// Changes the maximum number of resamples from the default value.
+    pub fn max_resamples(mut self, max_resamples: Option<usize>) -> Self {
+        self.max_resamples = max_resamples;
+        self
+    }
+
     /// Changes the weight distribution from the default value. See [`Weights`] for
     /// possible distributions.
     pub fn weights(mut self, weights: Weights) -> Self {
diff --git a/src/restart/mod.rs b/src/restart/mod.rs
@@ -317,7 +317,9 @@ impl Restarter {
             let seed = self.rng.gen();
 
             // Apply default configuration (may be overridden by individual restart strategies)
-            let mut options = self.default_options.clone()
+            let mut options = self
+                .default_options
+                .clone()
                 .initial_mean(initial_mean)
                 .mode(self.mode)
                 .parallel_update(self.parallel_update)
diff --git a/src/restart/options.rs b/src/restart/options.rs
@@ -4,8 +4,7 @@
 use std::ops::RangeInclusive;
 use std::time::Duration;
 
-
-use super::{DEFAULT_INITIAL_STEP_SIZE, RestartStrategy, Restarter};
+use super::{RestartStrategy, Restarter, DEFAULT_INITIAL_STEP_SIZE};
 use crate::{CMAESOptions, Mode};
 
 /// Represents invalid options for a `Restarter`.
diff --git a/src/sampling.rs b/src/sampling.rs
@@ -11,10 +11,42 @@ use crate::mode::Mode;
 use crate::state::State;
 use crate::{ObjectiveFunction, ParallelObjectiveFunction};
 
+pub trait Constraints: Sync + std::fmt::Debug {
+    fn meets_constraints(&self, x: &DVector<f64>) -> bool;
+    fn clone_box(&self) -> Box<dyn Constraints>;
+}
+
+impl Clone for Box<dyn Constraints> {
+    fn clone(&self) -> Box<dyn Constraints> {
+        self.clone_box()
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Bounds {
+    pub lower: Vec<f64>,
+    pub upper: Vec<f64>,
+}
+
+impl Constraints for Bounds {
+    fn meets_constraints(&self, x: &DVector<f64>) -> bool {
+        (0..x.len()).all(|i| x[i] >= self.lower[i] && x[i] <= self.upper[i])
+    }
+
+    fn clone_box(&self) -> Box<dyn Constraints> {
+        Box::new(self.clone())
+    }
+}
+
 /// A type for sampling and evaluating points from the distribution for each generation
 pub struct Sampler<F> {
     /// Number of dimensions to sample from
     dim: usize,
+    /// If set, resamples until all points satisfy the constraints
+    constraints: Option<Box<dyn Constraints>>,
+    /// The maximum number of resamples.
+    /// If this limit is hit, uses points even if they violate the constraints
+    max_resamples: Option<usize>,
     /// Number of points to sample each generation
     population_size: usize,
     /// RNG from which all random numbers are sourced
@@ -26,9 +58,18 @@ pub struct Sampler<F> {
 }
 
 impl<F> Sampler<F> {
-    pub fn new(dim: usize, population_size: usize, objective_function: F, rng_seed: u64) -> Self {
+    pub fn new(
+        dim: usize,
+        constraints: Option<Box<dyn Constraints>>,
+        max_resamples: Option<usize>,
+        population_size: usize,
+        objective_function: F,
+        rng_seed: u64,
+    ) -> Self {
         Self {
             dim,
+            constraints,
+            max_resamples,
             population_size,
             rng: ChaCha12Rng::seed_from_u64(rng_seed),
             objective_function,
@@ -49,22 +90,59 @@ impl<F> Sampler<F> {
         let normal = Normal::new(0.0, 1.0).unwrap();
 
         // Random steps in the distribution N(0, I)
-        let z = (0..self.population_size)
-            .map(|_| {
-                DVector::from_iterator(
-                    self.dim,
-                    (0..self.dim).map(|_| normal.sample(&mut self.rng)),
-                )
-            })
-            .collect::<Vec<_>>();
-        let transform = |zk| state.cov_transform() * zk;
-        let y = if parallel_update {
-            z.into_par_iter().map(transform).collect()
-        } else {
-            z.into_iter().map(transform).collect()
+        let mut sample = |n: usize, constraints: Option<&dyn Constraints>| {
+            let z = (0..n)
+                .map(|_| {
+                    DVector::from_iterator(
+                        self.dim,
+                        (0..self.dim).map(|_| normal.sample(&mut self.rng)),
+                    )
+                })
+                .collect::<Vec<_>>();
+            let transform = |zk| state.cov_transform() * zk;
+
+            let ok_constraints = |yk: &DVector<f64>| match constraints {
+                Some(constraints) => {
+                    constraints.meets_constraints(&to_point(&yk, state.mean(), state.sigma()))
+                }
+                None => true,
+            };
+
+            if parallel_update {
+                z.into_par_iter()
+                    .map(transform)
+                    .filter(ok_constraints)
+                    .collect()
+            } else {
+                z.into_iter()
+                    .map(transform)
+                    .filter(ok_constraints)
+                    .collect()
+            }
         };
 
-        // Evaluate and rank points
+        let mut y: Vec<DVector<f64>> = vec![];
+
+        let mut i: usize = 0;
+        loop {
+            let remain = self.population_size - y.len();
+            if remain == 0 {
+                break;
+            }
+
+            let mut constraints = self.constraints.as_ref().map(|x| x.as_ref());
+            if let Some(max) = self.max_resamples {
+                if i >= max {
+                    constraints = None;
+                }
+            }
+
+            let mut new_samps: Vec<_> = sample(remain, constraints);
+            y.append(&mut new_samps);
+
+            i += 1;
+        }
+
         let mut points = evaluate_points(y, &mut self.objective_function)?;
 
         self.function_evals += points.len();
@@ -138,6 +216,10 @@ pub struct EvaluatedPoint {
     value: f64,
 }
 
+fn to_point(unscaled_step: &DVector<f64>, mean: &DVector<f64>, sigma: f64) -> DVector<f64> {
+    mean + sigma * unscaled_step
+}
+
 impl EvaluatedPoint {
     /// Returns a new `EvaluatedPoint` from the unscaled step from the mean, the mean, and the step
     /// size
@@ -149,7 +231,7 @@ impl EvaluatedPoint {
         sigma: f64,
         mut objective_function: F,
     ) -> Result<Self, InvalidFunctionValueError> {
-        let point = mean + sigma * &unscaled_step;
+        let point = to_point(&unscaled_step, &mean, sigma);
         let value = objective_function(&point);
 
         if value.is_nan() {
@@ -206,7 +288,14 @@ mod tests {
     fn test_sample() {
         let dim = 10;
         let population_size = 12;
-        let mut sampler = Sampler::new(dim, population_size, Box::new(|_: &DVector<f64>| 0.0), 1);
+        let mut sampler = Sampler::new(
+            dim,
+            None,
+            None,
+            population_size,
+            Box::new(|_: &DVector<f64>| 0.0),
+            1,
+        );
         let state = State::new(vec![0.0; dim].into(), 2.0);
 
         let n = 5;
@@ -220,6 +309,8 @@ mod tests {
 
         let mut sampler_nan = Sampler::new(
             dim,
+            None,
+            None,
             population_size,
             Box::new(|_: &DVector<f64>| f64::NAN),
             1,
@@ -228,6 +319,72 @@ mod tests {
         assert!(sampler_nan.sample(&state, Mode::Minimize, false).is_err());
     }
 
+    #[test]
+    fn test_resample() {
+        let dim = 1;
+        let population_size = 1;
+
+        let bounds = Bounds {
+            lower: vec![0.0],
+            upper: vec![1.0],
+        };
+
+        let objective_function = |_: &DVector<f64>| 0.0;
+
+        // No resampling: Value should be out-of-bounds
+        {
+            let mut sampler = Sampler::new(
+                dim,
+                Some(Box::new(bounds.clone())),
+                Some(0),
+                population_size,
+                objective_function,
+                1,
+            );
+            let state = State::new(vec![0.0; dim].into(), 2.0);
+            let individuals = sampler.sample(&state, Mode::Minimize, false).unwrap();
+
+            assert!(
+                individuals[0].point[0] < bounds.lower[0]
+                    || individuals[0].point[0] > bounds.upper[0]
+            );
+        }
+
+        // With limited resampling: Value should be in bounds
+        {
+            let mut sampler = Sampler::new(
+                dim,
+                Some(Box::new(bounds.clone())),
+                Some(10),
+                population_size,
+                objective_function,
+                1,
+            );
+            let state = State::new(vec![0.0; dim].into(), 2.0);
+            let individuals = sampler.sample(&state, Mode::Minimize, false).unwrap();
+
+            assert!(individuals[0].point[0] >= bounds.lower[0]);
+            assert!(individuals[0].point[0] <= bounds.upper[0]);
+        }
+
+        // With unlimited resampling: Value should be in bounds
+        {
+            let mut sampler = Sampler::new(
+                dim,
+                Some(Box::new(bounds.clone())),
+                None,
+                population_size,
+                objective_function,
+                1,
+            );
+            let state = State::new(vec![0.0; dim].into(), 2.0);
+            let individuals = sampler.sample(&state, Mode::Minimize, false).unwrap();
+
+            assert!(individuals[0].point[0] >= bounds.lower[0]);
+            assert!(individuals[0].point[0] <= bounds.upper[0]);
+        }
+    }
+
     fn sample_sort(mode: Mode, expected: [f64; 5]) {
         let mut counter = 0.0;
         let function = |_: &DVector<f64>| {
@@ -243,7 +400,7 @@ mod tests {
         let dim = 10;
         let population_size = expected.len();
 
-        let mut sampler = Sampler::new(dim, population_size, function, 1);
+        let mut sampler = Sampler::new(dim, None, None, population_size, function, 1);
         let state = State::new(vec![0.0; dim].into(), 2.0);
 
         let individuals = sampler.sample(&state, mode, false).unwrap();
