Refactoring: move SurrogateBuilder and mixed integer moe from ego to moe (#382)

* feat: add SurrogateBuilder trait and XType enum for surrogate model configuration

- Introduced the `SurrogateBuilder` trait in `surrogate_builder.rs` to define methods for configuring and training surrogate models used by the Egor optimizer.
- Implemented the `SurrogateBuilder` trait for `GpMixtureParams<f64>`, providing functionality for setting regression specifications, correlation models, clustering, and training methods.
- Added the `XType` enum in `xtypes.rs` to represent different variable types (Float, Int, Ord, Enum) and a helper function `discrete` to check for discrete types in a given list of `XType`.
- Included unit tests for the `discrete` function to ensure correct identification of discrete variable types.

* Deprecate re-export from moe

* Linting

* Documention

* fix: improve documentation for LHS design and correlation model trait

Author: relf

crates/doe/src/lhs.rs

@@ -34,9 +34,9 @@ pub enum LhsKind {
 
 type RngRef<R> = Arc<RwLock<R>>;
 
-/// The LHS design is built as follows: each dimension space is divided into ns sections
-/// where ns is the number of sampling points, and one point in selected in each section.
-/// The selection method gives different kind of LHS (see [LhsKind])
+/// The LHS design is built as follows: each dimension space is divided into `ns` sections
+/// where `ns` is the number of sampling points, and one point is selected in each section.
+/// The selection method corresponds to different kinds of LHS: see [`LhsKind`]
 #[derive(Clone, Debug)]
 #[cfg_attr(feature = "serializable", derive(Serialize, Deserialize))]
 pub struct Lhs<F: Float, R: Rng> {

crates/ego/src/criteria/ei.rs

@@ -176,11 +176,9 @@ pub const LOG_EI: LogExpectedImprovement = LogExpectedImprovement {};
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{
-        gpmix::mixint::{MixintContext, MoeBuilder},
-        types::*,
-    };
+    use crate::types::*;
     use approx::assert_abs_diff_eq;
+    use egobox_moe::{MixintContext, MoeBuilder};
     // use egobox_moe::GpSurrogate;
     use finitediff::vec;
     use linfa::Dataset;

crates/ego/src/egor.rs

@@ -105,10 +105,10 @@ use crate::EgorConfig;
 use crate::EgorState;
 use crate::HotStartMode;
 use crate::errors::Result;
-use crate::gpmix::mixint::*;
 use crate::types::*;
 use crate::{CHECKPOINT_FILE, CheckpointingFrequency, HotStartCheckpoint};
 use crate::{EgorSolver, to_xtypes};
+use egobox_moe::{MixintGpMixtureParams, to_discrete_space};
 
 use argmin::core::observers::ObserverMode;
 
@@ -117,7 +117,7 @@ use log::info;
 use ndarray::{Array2, ArrayBase, Axis, Data, Ix2, concatenate};
 
 use argmin::core::{Error, Executor, KV, State, observers::Observe};
-use serde::de::DeserializeOwned;
+use serde::{Serialize, de::DeserializeOwned};
 
 use ndarray_npy::write_npy;
 use std::path::PathBuf;
@@ -219,14 +219,14 @@ impl<O: GroupFunc, C: CstrFn> EgorFactory<O, C> {
 pub struct Egor<
     O: GroupFunc,
     C: CstrFn = Cstr,
-    SB: SurrogateBuilder + DeserializeOwned = GpMixtureParams<f64>,
+    SB: SurrogateBuilder + Serialize + DeserializeOwned = GpMixtureParams<f64>,
 > {
     fobj: ObjFunc<O, C>,
     solver: EgorSolver<SB, C>,
     run_info: Option<RunInfo>,
 }
 
-impl<O: GroupFunc, C: CstrFn, SB: SurrogateBuilder + DeserializeOwned> Egor<O, C, SB> {
+impl<O: GroupFunc, C: CstrFn, SB: SurrogateBuilder + Serialize + DeserializeOwned> Egor<O, C, SB> {
     /// Runs the (constrained) optimization of the objective function.
     pub fn run(&self) -> Result<OptimResult<f64>> {
         let xtypes = self.solver.config.xtypes.clone();
@@ -422,7 +422,7 @@ mod tests {
     use argmin::core::{TerminationReason, TerminationStatus};
     use argmin_testfunctions::rosenbrock;
     use egobox_doe::{Lhs, SamplingMethod};
-    use egobox_moe::NbClusters;
+    use egobox_moe::{NbClusters, as_continuous_limits};
     use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Ix1, Zip, array, s};
     use ndarray_rand::rand::SeedableRng;
     use rand_xoshiro::Xoshiro256Plus;
@@ -432,7 +432,8 @@ mod tests {
     use serial_test::serial;
     use std::time::Instant;
 
-    use crate::{CoegoStatus, DOE_FILE, DOE_INITIAL_FILE, gpmix::spec::*, utils::EGOBOX_LOG};
+    use crate::{CoegoStatus, DOE_FILE, DOE_INITIAL_FILE, utils::EGOBOX_LOG};
+    use egobox_moe::{CorrelationSpec, RegressionSpec};
 
     #[cfg(not(feature = "blas"))]
     use linfa_linalg::norm::*;

crates/ego/src/gpmix/spec.rs

@@ -314,7 +314,6 @@
 #![warn(rustdoc::broken_intra_doc_links)]
 
 pub mod criteria;
-pub mod gpmix;
 
 mod egor;
 mod errors;
@@ -323,14 +322,14 @@ mod types;
 
 pub use crate::egor::*;
 pub use crate::errors::*;
-pub use crate::gpmix::spec::{CorrelationSpec, RegressionSpec};
 pub use crate::solver::*;
 pub use crate::types::*;
 pub use crate::utils::{
     CHECKPOINT_FILE, Checkpoint, CheckpointingFrequency, EGOBOX_LOG, EGOR_GP_FILENAME,
     EGOR_INITIAL_GP_FILENAME, EGOR_USE_GP_RECORDER, EGOR_USE_GP_VAR_PORTFOLIO,
     EGOR_USE_MAX_PROBA_OF_FEASIBILITY, HotStartCheckpoint, HotStartMode, find_best_result_index,
 };
+pub use egobox_moe::{CorrelationSpec, RegressionSpec};
 
 mod optimizers;
 mod utils;

crates/ego/src/lib.rs

@@ -7,7 +7,7 @@ use egobox_gp::ThetaTuning;
 use egobox_moe::MixtureGpSurrogate;
 use ndarray::{Array, Array1, Array2, ArrayBase, Axis, Data, Ix1, RemoveAxis, s};
 use rand_xoshiro::Xoshiro256Plus;
-use serde::de::DeserializeOwned;
+use serde::{Serialize, de::DeserializeOwned};
 
 use ndarray_rand::rand::seq::SliceRandom;
 
@@ -45,7 +45,7 @@ where
 
 impl<SB, C> EgorSolver<SB, C>
 where
-    SB: SurrogateBuilder + DeserializeOwned,
+    SB: SurrogateBuilder + Serialize + DeserializeOwned,
     C: CstrFn,
 {
     /// Compute array of components indices each row being used as

crates/ego/src/solver/coego.rs

@@ -44,12 +44,11 @@
 //!
 use std::marker::PhantomData;
 
-use crate::{EgorConfig, EgorSolver, errors::Result, gpmix::mixint::*, to_xtypes, types::*};
-
-use egobox_moe::GpMixtureParams;
+use crate::{EgorConfig, EgorSolver, errors::Result, to_xtypes, types::*};
+use egobox_moe::{GpMixtureParams, MixintGpMixtureParams, to_continuous_space, to_discrete_space};
 use ndarray::{Array2, ArrayBase, Data, Ix2};
 
-use serde::de::DeserializeOwned;
+use serde::{Serialize, de::DeserializeOwned};
 
 /// EGO optimizer service builder allowing to use Egor optimizer
 /// as a service.
@@ -109,11 +108,11 @@ impl<C: CstrFn> EgorServiceFactory<C> {
 
 /// Egor optimizer service API.
 #[derive(Clone)]
-pub struct EgorServiceApi<SB: SurrogateBuilder + DeserializeOwned, C: CstrFn = Cstr> {
+pub struct EgorServiceApi<SB: SurrogateBuilder + Serialize + DeserializeOwned, C: CstrFn = Cstr> {
     solver: EgorSolver<SB, C>,
 }
 
-impl<SB: SurrogateBuilder + DeserializeOwned, C: CstrFn> EgorServiceApi<SB, C> {
+impl<SB: SurrogateBuilder + Serialize + DeserializeOwned, C: CstrFn> EgorServiceApi<SB, C> {
     /// Given an evaluated doe (x, y) data, return the next promising x point
     /// where optimum may be located with regard to the infill criterion.
     /// This function inverses the control of the optimization and can be used
@@ -136,8 +135,8 @@ pub type EgorServiceBuilder = EgorServiceFactory<Cstr>;
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::gpmix::spec::*;
     use approx::assert_abs_diff_eq;
+    use egobox_moe::{CorrelationSpec, RegressionSpec};
     use ndarray::{ArrayView2, Axis, array, concatenate};
 
     use ndarray_stats::QuantileExt;

crates/ego/src/solver/egor_service.rs

@@ -173,7 +173,7 @@ impl<O, SB, C> Solver<O, EgorState<f64>> for EgorSolver<SB, C>
 where
     O: CostFunction<Param = Array2<f64>, Output = Array2<f64>> + DomainConstraints<C>,
     C: CstrFn,
-    SB: SurrogateBuilder + DeserializeOwned,
+    SB: SurrogateBuilder + Serialize + DeserializeOwned,
 {
     fn name(&self) -> &str {
         "Egor"
@@ -446,7 +446,7 @@ where
 
 impl<SB, C: CstrFn> EgorSolver<SB, C>
 where
-    SB: SurrogateBuilder + DeserializeOwned,
+    SB: SurrogateBuilder + Serialize + DeserializeOwned,
 {
     /// Iteration of EGO algorithm
     fn ego_iteration<

crates/ego/src/solver/egor_solver.rs

@@ -1,6 +1,6 @@
 use crate::errors::Result;
-use crate::gpmix::mixint::to_discrete_space;
 use crate::{types::*, utils};
+use egobox_moe::to_discrete_space;
 
 use crate::utils::{
     EGOR_DO_NOT_USE_MIDDLEPICKER_MULTISTARTER, compute_cstr_scales, logpofs, logpofs_grad, pofs,

crates/ego/src/solver/solver_computations.rs

@@ -2,7 +2,6 @@ use std::marker::PhantomData;
 
 use crate::errors::{EgoError, Result};
 use crate::find_best_result_index;
-use crate::gpmix::mixint::{as_continuous_limits, to_discrete_space};
 use crate::solver::solver_computations::MiddlePickerMultiStarter;
 use crate::solver::solver_infill_optim::InfillOptProblem;
 use crate::utils::{
@@ -11,6 +10,7 @@ use crate::utils::{
 };
 use crate::{DEFAULT_CSTR_TOL, EgorSolver, MAX_POINT_ADDITION_RETRY, ValidEgorConfig};
 use crate::{EgorState, types::*};
+use egobox_moe::{as_continuous_limits, to_discrete_space};
 
 use argmin::argmin_error_closure;
 use argmin::core::{CostFunction, Problem, State};
@@ -25,11 +25,11 @@ use ndarray::{Array1, Array2, ArrayBase, Axis, Data, Ix1, Ix2, Zip, concatenate,
 use ndarray_rand::rand::{Rng, SeedableRng};
 use rand_xoshiro::Xoshiro256Plus;
 use rayon::prelude::*;
-use serde::de::DeserializeOwned;
+use serde::{Serialize, de::DeserializeOwned};
 
 use super::coego::COEGO_IMPROVEMENT_CHECK;
 
-impl<SB: SurrogateBuilder + DeserializeOwned, C: CstrFn> EgorSolver<SB, C> {
+impl<SB: SurrogateBuilder + Serialize + DeserializeOwned, C: CstrFn> EgorSolver<SB, C> {
     /// Constructor of the optimization of the function `f` with specified random generator
     /// to get reproducibility.
     ///
@@ -104,7 +104,7 @@ impl<SB: SurrogateBuilder + DeserializeOwned, C: CstrFn> EgorSolver<SB, C> {
 
 impl<SB, C> EgorSolver<SB, C>
 where
-    SB: SurrogateBuilder + DeserializeOwned,
+    SB: SurrogateBuilder + Serialize + DeserializeOwned,
     C: CstrFn,
 {
     /// Whether we have to recluster the data