Skip to content

Optimize GP predictions for Egor optimizer usage #332

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Nov 7, 2025
Merged

Optimize GP predictions for Egor optimizer usage #332

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
a172825
Add correlation models bench
relf Nov 6, 2025
721e281
Simplify diff matrix computation
relf Nov 6, 2025
13a0197
Implement predict_valvar for GP to compute value and var in one go
relf Nov 6, 2025
c28112b
Use predict_valvar where needed
relf Nov 6, 2025
c034cb7
Implement predict_valvar_gradients
relf Nov 7, 2025
0f3e24c
Add sanity check test for valvar variant prediction
relf Nov 7, 2025
4a4ca0c
Use predict_valvar_gradients
relf Nov 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions crates/ego/benches/ego.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use criterion::{Criterion, criterion_group, criterion_main};
use egobox_ego::{EgorBuilder, InfillStrategy};
use egobox_ego::{EGOBOX_LOG, EgorBuilder, InfillStrategy};
use egobox_moe::{CorrelationSpec, RegressionSpec};
use env_logger::{Builder, Env};
use ndarray::{Array2, ArrayView2, Zip, array};

/// Ackley test function: min f(x)=0 at x=(0, 0, 0)
Expand All @@ -15,15 +16,21 @@ fn ackley(x: &ArrayView2<f64>) -> Array2<f64> {
fn criterion_ego(c: &mut Criterion) {
let xlimits = array![[-32.768, 32.768], [-32.768, 32.768], [-32.768, 32.768]];
let mut group = c.benchmark_group("ego");
group.sample_size(20);
group.bench_function("ego ackley", |b| {
let env = Env::new().filter_or(EGOBOX_LOG, "error");
let mut builder = Builder::from_env(env);
let builder = builder.target(env_logger::Target::Stdout);
builder.try_init().ok();

b.iter(|| {
std::hint::black_box(
EgorBuilder::optimize(ackley)
.configure(|config| {
config
.configure_gp(|conf| {
conf.regression_spec(RegressionSpec::CONSTANT)
.correlation_spec(CorrelationSpec::ABSOLUTEEXPONENTIAL)
.correlation_spec(CorrelationSpec::MATERN52)
})
.infill_strategy(InfillStrategy::WB2S)
.max_iters(10)
Expand Down
158 changes: 72 additions & 86 deletions crates/ego/src/criteria/ei.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,23 +28,20 @@ impl InfillCriterion for ExpectedImprovement {
_scale: Option<f64>,
) -> f64 {
let pt = ArrayView::from_shape((1, x.len()), x).unwrap();
match obj_model.predict(&pt) {
Ok(p) => match obj_model.predict_var(&pt) {
Ok(s) => {
if s[0] < f64::EPSILON {
0.0
} else {
let pred = p[0];
let k = sigma_weight.unwrap_or(1.0);
let sigma = k * s[0].sqrt();
let args0 = (fmin - pred) / sigma;
let args1 = args0 * norm_cdf(args0);
let args2 = norm_pdf(args0);
sigma * (args1 + args2)
}
match obj_model.predict_valvar(&pt) {
Ok((p, s)) => {
if s[0] < f64::EPSILON {
0.0
} else {
let pred = p[0];
let k = sigma_weight.unwrap_or(1.0);
let sigma = k * s[0].sqrt();
let args0 = (fmin - pred) / sigma;
let args1 = args0 * norm_cdf(args0);
let args2 = norm_pdf(args0);
sigma * (args1 + args2)
}
_ => 0.0,
},
}
_ => 0.0,
}
}
Expand All @@ -60,36 +57,32 @@ impl InfillCriterion for ExpectedImprovement {
_scale: Option<f64>,
) -> Array1<f64> {
let pt = ArrayView::from_shape((1, x.len()), x).unwrap();
match obj_model.predict(&pt) {
Ok(p) => match obj_model.predict_var(&pt) {
Ok(s) => {
if s[0] < f64::EPSILON {
Array1::zeros(pt.len())
} else {
let pred = p[0];
let diff_y = fmin - pred;
let k = sigma_weight.unwrap_or(1.0);
let sigma = s[0].sqrt();
let arg = (fmin - pred) / (k * sigma);
let y_prime = obj_model.predict_gradients(&pt).unwrap();
let y_prime = y_prime.row(0);
let sig_2_prime = obj_model.predict_var_gradients(&pt).unwrap();

let sig_2_prime = sig_2_prime.row(0);
let sig_prime = sig_2_prime.mapv(|v| k * v / (2. * sigma));
let arg_prime = y_prime.mapv(|v| v / (-k * sigma))
- diff_y.to_owned() * sig_prime.mapv(|v| v / (k * sigma * k * sigma));
let factor = k * sigma * (-arg / SQRT_2PI) * (-(arg * arg) / 2.).exp();

let arg1 = y_prime.mapv(|v| v * (-norm_cdf(arg)));
let arg2 = diff_y * norm_pdf(arg) * arg_prime.to_owned();
let arg3 = sig_prime.to_owned() * norm_pdf(arg);
let arg4 = factor * arg_prime;
arg1 + arg2 + arg3 + arg4
}
match obj_model.predict_valvar(&pt) {
Ok((p, s)) => {
if s[0] < f64::EPSILON {
Array1::zeros(pt.len())
} else {
let pred = p[0];
let diff_y = fmin - pred;
let k = sigma_weight.unwrap_or(1.0);
let sigma = s[0].sqrt();
let arg = (fmin - pred) / (k * sigma);

let (y_prime, var_prime) = obj_model.predict_valvar_gradients(&pt).unwrap();
let y_prime = y_prime.row(0);
let sig_2_prime = var_prime.row(0);
let sig_prime = sig_2_prime.mapv(|v| k * v / (2. * sigma));
let arg_prime = y_prime.mapv(|v| v / (-k * sigma))
- diff_y.to_owned() * sig_prime.mapv(|v| v / (k * sigma * k * sigma));
let factor = k * sigma * (-arg / SQRT_2PI) * (-(arg * arg) / 2.).exp();

let arg1 = y_prime.mapv(|v| v * (-norm_cdf(arg)));
let arg2 = diff_y * norm_pdf(arg) * arg_prime.to_owned();
let arg3 = sig_prime.to_owned() * norm_pdf(arg);
let arg4 = factor * arg_prime;
arg1 + arg2 + arg3 + arg4
}
_ => Array1::zeros(pt.len()),
},
}
_ => Array1::zeros(pt.len()),
}
}
Expand Down Expand Up @@ -120,20 +113,17 @@ impl InfillCriterion for LogExpectedImprovement {
) -> f64 {
let pt = ArrayView::from_shape((1, x.len()), x).unwrap();

match obj_model.predict(&pt) {
Ok(p) => match obj_model.predict_var(&pt) {
Ok(s) => {
if s[0] < f64::EPSILON {
f64::MIN
} else {
let pred = p[0];
let sigma = s[0].sqrt();
let u = (fmin - pred) / sigma;
log_ei_helper(u) + sigma.ln()
}
match obj_model.predict_valvar(&pt) {
Ok((p, s)) => {
if s[0] < f64::EPSILON {
f64::MIN
} else {
let pred = p[0];
let sigma = s[0].sqrt();
let u = (fmin - pred) / sigma;
log_ei_helper(u) + sigma.ln()
}
_ => f64::MIN,
},
}
_ => f64::MIN,
}
}
Expand All @@ -150,35 +140,31 @@ impl InfillCriterion for LogExpectedImprovement {
) -> Array1<f64> {
let pt = ArrayView::from_shape((1, x.len()), x).unwrap();

match obj_model.predict(&pt) {
Ok(p) => match obj_model.predict_var(&pt) {
Ok(s) => {
if s[0] < f64::EPSILON {
Array1::from_elem(pt.len(), f64::MIN)
} else {
let pred = p[0];
let diff_y = fmin - pred;
let sigma = s[0].sqrt();
let arg = diff_y / sigma;

let y_prime = obj_model.predict_gradients(&pt).unwrap();
let y_prime = y_prime.row(0);
let sig_2_prime = obj_model.predict_var_gradients(&pt).unwrap();
let sig_2_prime = sig_2_prime.row(0);
let sig_prime = sig_2_prime.mapv(|v| v / (2. * sigma));

let arg_prime = y_prime.mapv(|v| v / (-sigma))
- diff_y.to_owned() * sig_prime.mapv(|v| v / (sigma * sigma));

let dhelper = d_log_ei_helper(arg);
let arg1 = arg_prime.mapv(|v| dhelper * v);

let arg2 = sig_prime / sigma;
arg1 + arg2
}
match obj_model.predict_valvar(&pt) {
Ok((p, s)) => {
if s[0] < f64::EPSILON {
Array1::from_elem(pt.len(), f64::MIN)
} else {
let pred = p[0];
let diff_y = fmin - pred;
let sigma = s[0].sqrt();
let arg = diff_y / sigma;

let (y_prime, var_prime) = obj_model.predict_valvar_gradients(&pt).unwrap();
let y_prime = y_prime.row(0);
let sig_2_prime = var_prime.row(0);
let sig_prime = sig_2_prime.mapv(|v| v / (2. * sigma));

let arg_prime = y_prime.mapv(|v| v / (-sigma))
- diff_y.to_owned() * sig_prime.mapv(|v| v / (sigma * sigma));

let dhelper = d_log_ei_helper(arg);
let arg1 = arg_prime.mapv(|v| dhelper * v);

let arg2 = sig_prime / sigma;
arg1 + arg2
}
_ => Array1::from_elem(pt.len(), f64::MIN),
},
}
_ => Array1::from_elem(pt.len(), f64::MIN),
}
}
Expand Down
26 changes: 26 additions & 0 deletions crates/ego/src/gpmix/mixint.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -606,6 +606,19 @@ impl GpSurrogate for MixintGpMixture {
self.moe.predict_var(&xcast)
}

fn predict_valvar(
&self,
x: &ArrayView2<f64>,
) -> egobox_moe::Result<(Array1<f64>, Array1<f64>)> {
let mut xcast = if self.work_in_folded_space {
unfold_with_enum_mask(&self.xtypes, x)
} else {
x.to_owned()
};
cast_to_discrete_values_mut(&self.xtypes, &mut xcast);
self.moe.predict_valvar(&xcast)
}

/// Save Moe model in given file.
#[cfg(feature = "persistent")]
fn save(&self, path: &str, format: GpFileFormat) -> egobox_moe::Result<()> {
Expand Down Expand Up @@ -660,6 +673,19 @@ impl GpSurrogateExt for MixintGpMixture {
self.moe.predict_var_gradients(&xcast)
}

fn predict_valvar_gradients(
&self,
x: &ArrayView2<f64>,
) -> egobox_moe::Result<(Array2<f64>, Array2<f64>)> {
let mut xcast = if self.work_in_folded_space {
unfold_with_enum_mask(&self.xtypes, x)
} else {
x.to_owned()
};
cast_to_discrete_values_mut(&self.xtypes, &mut xcast);
self.moe.predict_valvar_gradients(&xcast)
}

fn sample(&self, x: &ArrayView2<f64>, n_traj: usize) -> egobox_moe::Result<Array2<f64>> {
let mut xcast = if self.work_in_folded_space {
unfold_with_enum_mask(&self.xtypes, x)
Expand Down
11 changes: 7 additions & 4 deletions crates/ego/src/solver/coego.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,14 +198,17 @@ where
) -> Result<Array1<f64>> {
let mut res: Vec<f64> = vec![];
let x = &x.view().insert_axis(Axis(0));
let sigma = obj_model.predict_var(&x.view()).unwrap()[0].sqrt();

let (pred, var) = obj_model.predict_valvar(x)?;
let sigma = var[0].sqrt();
// Use lower trust bound for a minimization
let pred = obj_model.predict(x)?[0] - CSTR_DOUBT * sigma;
let pred = pred[0] - CSTR_DOUBT * sigma;
res.push(pred);
for cstr_model in cstr_models {
let sigma = cstr_model.predict_var(&x.view()).unwrap()[0].sqrt();
let (pred, var) = cstr_model.predict_valvar(x)?;
let sigma = var[0].sqrt();
// Use upper trust bound
res.push(cstr_model.predict(x)?[0] + CSTR_DOUBT * sigma);
res.push(pred[0] + CSTR_DOUBT * sigma);
}
Ok(Array1::from_vec(res))
}
Expand Down
13 changes: 7 additions & 6 deletions crates/ego/src/solver/solver_computations.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,18 +223,19 @@ where
active: &[usize],
) -> f64 {
let x = Array::from_shape_vec((1, x.len()), x.to_vec()).unwrap();
let sigma = cstr_model.predict_var(&x.view()).unwrap()[0].sqrt();
let cstr_val = cstr_model.predict(&x.view()).unwrap()[0];

let (pred, var) = cstr_model.predict_valvar(&x.view()).unwrap();
let sigma = var[0].sqrt();
let cstr_val = pred[0];

if let Some(grad) = gradient {
let (pred_grad, var_grad) = cstr_model.predict_valvar_gradients(&x.view()).unwrap();
let sigma_prime = if sigma < f64::EPSILON {
0.
} else {
cstr_model.predict_var_gradients(&x.view()).unwrap()[[0, 0]] / (2. * sigma)
var_grad[[0, 0]] / (2. * sigma)
};
let grd = cstr_model
.predict_gradients(&x.view())
.unwrap()
let grd = pred_grad
.row(0)
.mapv(|v| (v + CSTR_DOUBT * sigma_prime) / scale_cstr)
.to_vec();
Expand Down
59 changes: 26 additions & 33 deletions crates/ego/src/utils/cstr_pof.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,17 @@ use ndarray::{Array1, ArrayView};
/// a constraint function wrt the tolerance (ie cstr <= cstr_tol)
fn pof(x: &[f64], cstr_model: &dyn MixtureGpSurrogate, cstr_tol: f64) -> f64 {
let pt = ArrayView::from_shape((1, x.len()), x).unwrap();
match cstr_model.predict(&pt) {
Ok(p) => match cstr_model.predict_var(&pt) {
Ok(s) => {
if s[0] < f64::EPSILON {
0.0
} else {
let pred = p[0];
let sigma = s[0].sqrt();
let args0 = (cstr_tol - pred) / sigma;
norm_cdf(args0)
}
match cstr_model.predict_valvar(&pt) {
Ok((p, s)) => {
if s[0] < f64::EPSILON {
0.0
} else {
let pred = p[0];
let sigma = s[0].sqrt();
let args0 = (cstr_tol - pred) / sigma;
norm_cdf(args0)
}
_ => 0.0,
},
}
_ => 0.0,
}
}
Expand All @@ -30,27 +27,23 @@ fn pof(x: &[f64], cstr_model: &dyn MixtureGpSurrogate, cstr_tol: f64) -> f64 {
/// constraint surrogate model.
fn pof_grad(x: &[f64], cstr_model: &dyn MixtureGpSurrogate, cstr_tol: f64) -> Array1<f64> {
let pt = ArrayView::from_shape((1, x.len()), x).unwrap();
match cstr_model.predict(&pt) {
Ok(p) => match cstr_model.predict_var(&pt) {
Ok(s) => {
if s[0] < f64::EPSILON {
Array1::zeros(pt.len())
} else {
let pred = p[0];
let sigma = s[0].sqrt();
let arg = (cstr_tol - pred) / sigma;
let y_prime = cstr_model.predict_gradients(&pt).unwrap();
let y_prime = y_prime.row(0);
let sig_2_prime = cstr_model.predict_var_gradients(&pt).unwrap();
let sig_2_prime = sig_2_prime.row(0);
let sig_prime = sig_2_prime.mapv(|v| v / (2. * sigma));
let arg_prime = y_prime.mapv(|v| v / (-sigma))
+ sig_prime.mapv(|v| v * pred / (sigma * sigma));
norm_pdf(arg) * arg_prime.to_owned()
}
match cstr_model.predict_valvar(&pt) {
Ok((p, s)) => {
if s[0] < f64::EPSILON {
Array1::zeros(pt.len())
} else {
let pred = p[0];
let sigma = s[0].sqrt();
let arg = (cstr_tol - pred) / sigma;
let (y_prime, var_prime) = cstr_model.predict_valvar_gradients(&pt).unwrap();
let y_prime = y_prime.row(0);
let sig_2_prime = var_prime.row(0);
let sig_prime = sig_2_prime.mapv(|v| v / (2. * sigma));
let arg_prime =
y_prime.mapv(|v| v / (-sigma)) + sig_prime.mapv(|v| v * pred / (sigma * sigma));
norm_pdf(arg) * arg_prime.to_owned()
}
_ => Array1::zeros(pt.len()),
},
}
_ => Array1::zeros(pt.len()),
}
}
Expand Down
4 changes: 4 additions & 0 deletions crates/gp/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,3 +61,7 @@ argmin_testfunctions.workspace = true
[[bench]]
name = "gp"
harness = false

[[bench]]
name = "corr"
harness = false
Loading
Loading