use crate::core::math::Scalar;
use super::geometry::{mesh_and_poll_size, poll_directions};
use super::halton::halton_seed;
pub(crate) enum Transition {
Continue,
Converged,
}
pub(crate) struct StepOutcome<F> {
pub(crate) transition: Transition,
pub(crate) evaluated: Vec<(Vec<F>, F)>,
}
pub(crate) struct MadsWork<F> {
n: usize,
x: Vec<F>,
fx: F,
ell: i32,
t0: usize,
ell_max: i32,
t_max: usize,
scale: F,
poll_size_min: F,
}
impl<F: Scalar> MadsWork<F> {
pub(crate) fn try_init<E>(
x0: Vec<F>,
poll_size_init: F,
poll_size_min: F,
eval: &mut dyn FnMut(&[F]) -> Result<F, E>,
) -> Result<(Self, Vec<F>, F), E> {
let n = x0.len();
assert!(n >= 1, "Mads requires a non-empty start point");
let fx = eval(&x0)?;
let work = Self {
n,
x: x0.clone(),
fx,
ell: 0,
t0: halton_seed(n),
ell_max: 0,
t_max: 0,
scale: poll_size_init,
poll_size_min,
};
Ok((work, x0, fx))
}
pub(crate) fn poll_size(&self) -> F {
let (_, poll) = mesh_and_poll_size(self.ell);
self.scale * F::from_f64(poll).expect("poll size representable")
}
pub(crate) fn mesh_index(&self) -> i32 {
self.ell
}
pub(crate) fn step<E>(
&mut self,
eval: &mut dyn FnMut(&[F]) -> Result<F, E>,
) -> Result<StepOutcome<F>, E> {
let t = if self.ell >= self.ell_max {
self.ell_max = self.ell;
(self.ell + self.t0 as i32) as usize } else {
self.t_max + 1
};
self.t_max = self.t_max.max(t);
let (mesh_f64, _) = mesh_and_poll_size(self.ell);
let mesh = self.scale * F::from_f64(mesh_f64).expect("mesh size representable");
let dirs = poll_directions(t, self.ell, self.n);
let mut evaluated = Vec::new();
let mut improved = false;
for d in &dirs {
let mut trial = self.x.clone();
for i in 0..self.n {
let di = F::from_i64(d[i]).expect("integer direction representable");
trial[i] = trial[i] + mesh * di;
}
let f = eval(&trial)?;
if f < self.fx {
self.x = trial.clone();
self.fx = f;
evaluated.push((trial, f));
improved = true;
break;
}
}
if improved {
self.ell -= 1;
} else {
self.ell += 1;
}
let transition = if self.poll_size() <= self.poll_size_min {
Transition::Converged
} else {
Transition::Continue
};
Ok(StepOutcome {
transition,
evaluated,
})
}
}