Change GP training Rust API to use one dimensional array

.gitignore

@@ -19,7 +19,7 @@
 
 # Project
 **/*.json
-**/.bin
+**/*.bin
 *.npy
 input.txt
 output.txt

ego/src/criteria/ei.rs

@@ -28,7 +28,7 @@ impl InfillCriterion for ExpectedImprovement {
         let pt = ArrayView::from_shape((1, x.len()), x).unwrap();
         if let Ok(p) = obj_model.predict(&pt) {
             if let Ok(s) = obj_model.predict_var(&pt) {
-                let pred = p[[0, 0]];
+                let pred = p[0];
                 let sigma = s[[0, 0]].sqrt();
                 let args0 = (fmin - pred) / sigma;
                 let args1 = (fmin - pred) * norm_cdf(args0);
@@ -58,7 +58,7 @@ impl InfillCriterion for ExpectedImprovement {
                 if sigma.abs() < 1e-12 {
                     Array1::zeros(pt.len())
                 } else {
-                    let pred = p[[0, 0]];
+                    let pred = p[0];
                     let diff_y = fmin - pred;
                     let arg = (fmin - pred) / sigma;
                     let y_prime = obj_model.predict_gradients(&pt).unwrap();

ego/src/criteria/wb2.rs

@@ -31,7 +31,7 @@ impl InfillCriterion for WB2Criterion {
         let scale = scale.unwrap_or(self.0.unwrap_or(1.0));
         let pt = ArrayView::from_shape((1, x.len()), x).unwrap();
         let ei = EI.value(x, obj_model, fmin, None);
-        scale * ei - obj_model.predict(&pt).unwrap()[[0, 0]]
+        scale * ei - obj_model.predict(&pt).unwrap()[0]
     }
 
     /// Computes derivatives of WB2S infill criterion wrt to x components at given `x` point
@@ -78,7 +78,7 @@ pub(crate) fn compute_wb2s_scale(
     let i_max = ei_x.argmax().unwrap();
     let pred_max = obj_model
         .predict(&x.row(i_max).insert_axis(Axis(0)))
-        .unwrap()[[0, 0]];
+        .unwrap()[0];
     let ei_max = ei_x[i_max];
     if ei_max.abs() > 100. * f64::EPSILON {
         ratio * pred_max / ei_max
@@ -113,7 +113,7 @@ mod tests {
             .regression_spec(RegressionSpec::CONSTANT)
             .correlation_spec(CorrelationSpec::SQUAREDEXPONENTIAL)
             .recombination(Recombination::Hard)
-            .fit(&Dataset::new(xt, yt))
+            .fit(&Dataset::new(xt, yt.remove_axis(Axis(1))))
             .expect("GP fitting");
         let bgp = Box::new(gp) as Box<dyn MixtureGpSurrogate>;
 
@@ -153,12 +153,7 @@ mod tests {
         let fdiff2 = (bgp.predict(&xtest21.view()).unwrap()
             - bgp.predict(&xtest22.view()).unwrap())
             / (2. * h);
-        println!(
-            "gp fdiff({}) = [[{}, {}]]",
-            xtest,
-            fdiff1[[0, 0]],
-            fdiff2[[0, 0]]
-        );
+        println!("gp fdiff({}) = [[{}, {}]]", xtest, fdiff1[0], fdiff2[0]);
         println!(
             "GP predict derivatives({}) = {}",
             xtest,

ego/src/gpmix/mixint.rs

@@ -14,7 +14,9 @@ use egobox_moe::{
 };
 use linfa::traits::{Fit, PredictInplace};
 use linfa::{DatasetBase, Float, ParamGuard};
-use ndarray::{s, Array, Array2, ArrayBase, ArrayView2, Axis, Data, DataMut, Ix2, Zip};
+use ndarray::{
+    s, Array, Array1, Array2, ArrayBase, ArrayView1, ArrayView2, Axis, Data, DataMut, Ix1, Ix2, Zip,
+};
 use ndarray_rand::rand::SeedableRng;
 use ndarray_stats::QuantileExt;
 use rand_xoshiro::Xoshiro256Plus;
@@ -343,7 +345,7 @@ impl MixintGpMixtureValidParams {
     fn _train(
         &self,
         xt: &ArrayBase<impl Data<Elem = f64>, Ix2>,
-        yt: &ArrayBase<impl Data<Elem = f64>, Ix2>,
+        yt: &ArrayBase<impl Data<Elem = f64>, Ix1>,
     ) -> Result<MixintGpMixture> {
         let mut xcast = if self.work_in_folded_space {
             unfold_with_enum_mask(&self.xtypes, &xt.view())
@@ -369,7 +371,7 @@ impl MixintGpMixtureValidParams {
     fn _train_on_clusters(
         &self,
         xt: &ArrayBase<impl Data<Elem = f64>, Ix2>,
-        yt: &ArrayBase<impl Data<Elem = f64>, Ix2>,
+        yt: &ArrayBase<impl Data<Elem = f64>, Ix1>,
         clustering: &egobox_moe::Clustering,
     ) -> Result<MixintGpMixture> {
         let mut xcast = if self.work_in_folded_space {
@@ -458,32 +460,32 @@ impl SurrogateBuilder for MixintGpMixtureParams {
 
     fn train(
         &self,
-        xt: &ArrayView2<f64>,
-        yt: &ArrayView2<f64>,
+        xt: ArrayView2<f64>,
+        yt: ArrayView1<f64>,
     ) -> Result<Box<dyn MixtureGpSurrogate>> {
-        let mixmoe = self.check_ref()?._train(xt, yt)?;
+        let mixmoe = self.check_ref()?._train(&xt, &yt)?;
         Ok(mixmoe).map(|mixmoe| Box::new(mixmoe) as Box<dyn MixtureGpSurrogate>)
     }
 
     fn train_on_clusters(
         &self,
-        xt: &ArrayView2<f64>,
-        yt: &ArrayView2<f64>,
+        xt: ArrayView2<f64>,
+        yt: ArrayView1<f64>,
         clustering: &Clustering,
     ) -> Result<Box<dyn MixtureGpSurrogate>> {
-        let mixmoe = self.check_ref()?._train_on_clusters(xt, yt, clustering)?;
+        let mixmoe = self.check_ref()?._train_on_clusters(&xt, &yt, clustering)?;
         Ok(mixmoe).map(|mixmoe| Box::new(mixmoe) as Box<dyn MixtureGpSurrogate>)
     }
 }
 
-impl<D: Data<Elem = f64>> Fit<ArrayBase<D, Ix2>, ArrayBase<D, Ix2>, EgoError>
+impl<D: Data<Elem = f64>> Fit<ArrayBase<D, Ix2>, ArrayBase<D, Ix1>, EgoError>
     for MixintGpMixtureValidParams
 {
     type Object = MixintGpMixture;
 
     fn fit(
         &self,
-        dataset: &DatasetBase<ArrayBase<D, Ix2>, ArrayBase<D, Ix2>>,
+        dataset: &DatasetBase<ArrayBase<D, Ix2>, ArrayBase<D, Ix1>>,
     ) -> Result<Self::Object> {
         let x = dataset.records();
         let y = dataset.targets();
@@ -522,7 +524,7 @@ pub struct MixintGpMixture {
     /// i.e for "blue" in ["red", "green", "blue"] either \[2\] or [0, 0, 1]
     work_in_folded_space: bool,
     /// Training inputs
-    training_data: (Array2<f64>, Array2<f64>),
+    training_data: (Array2<f64>, Array1<f64>),
     /// Parameters used to trin this model
     params: MixintGpMixtureValidParams,
 }
@@ -559,7 +561,7 @@ impl GpSurrogate for MixintGpMixture {
         self.moe.dims()
     }
 
-    fn predict(&self, x: &ArrayView2<f64>) -> egobox_moe::Result<Array2<f64>> {
+    fn predict(&self, x: &ArrayView2<f64>) -> egobox_moe::Result<Array1<f64>> {
         let mut xcast = if self.work_in_folded_space {
             unfold_with_enum_mask(&self.xtypes, x)
         } else {
@@ -628,7 +630,7 @@ impl GpSurrogateExt for MixintGpMixture {
 }
 
 impl CrossValScore<f64, EgoError, MixintGpMixtureParams, Self> for MixintGpMixture {
-    fn training_data(&self) -> &(Array2<f64>, Array2<f64>) {
+    fn training_data(&self) -> &(Array2<f64>, Array1<f64>) {
         &self.training_data
     }
 
@@ -643,20 +645,20 @@ impl MixtureGpSurrogate for MixintGpMixture {
     }
 }
 
-impl<D: Data<Elem = f64>> PredictInplace<ArrayBase<D, Ix2>, Array2<f64>> for MixintGpMixture {
-    fn predict_inplace(&self, x: &ArrayBase<D, Ix2>, y: &mut Array2<f64>) {
+impl<D: Data<Elem = f64>> PredictInplace<ArrayBase<D, Ix2>, Array1<f64>> for MixintGpMixture {
+    fn predict_inplace(&self, x: &ArrayBase<D, Ix2>, y: &mut Array1<f64>) {
         assert_eq!(
             x.nrows(),
-            y.nrows(),
+            y.len(),
             "The number of data points must match the number of output targets."
         );
 
         let values = self.moe.predict(x).expect("MixintGpMixture prediction");
         *y = values;
     }
 
-    fn default_target(&self, x: &ArrayBase<D, Ix2>) -> Array2<f64> {
-        Array2::zeros((x.nrows(), self.moe.dims().1))
+    fn default_target(&self, x: &ArrayBase<D, Ix2>) -> Array1<f64> {
+        Array1::zeros((x.nrows(),))
     }
 }
 
@@ -760,7 +762,7 @@ impl MixintContext {
     pub fn create_surrogate(
         &self,
         surrogate_builder: &MoeBuilder,
-        dataset: &DatasetBase<Array2<f64>, Array2<f64>>,
+        dataset: &DatasetBase<Array2<f64>, Array1<f64>>,
     ) -> Result<MixintGpMixture> {
         let mut params = MixintGpMixtureParams::new(&self.xtypes, surrogate_builder);
         let params = params.work_in_folded_space(self.work_in_folded_space);
@@ -870,7 +872,7 @@ mod tests {
 
         let surrogate_builder = MoeBuilder::new();
         let xt = array![[0.], [2.], [3.0], [4.]];
-        let yt = array![[0.], [1.5], [0.9], [1.]];
+        let yt = array![0., 1.5, 0.9, 1.];
         let ds = Dataset::new(xt, yt);
         let mixi_moe = mixi
             .create_surrogate(&surrogate_builder, &ds)
@@ -884,7 +886,7 @@ mod tests {
             .expect("Predict var fail");
         println!("{ytest:?}");
         assert_abs_diff_eq!(
-            array![[0.], [0.7872696212255119], [1.5], [0.9], [1.]],
+            array![0., 0.7872696212255119, 1.5, 0.9, 1.],
             ytest,
             epsilon = 1e-3
         );
@@ -894,13 +896,13 @@ mod tests {
             yvar,
             epsilon = 1e-3
         );
-        println!("LOOCV = {}", mixi_moe.loocv_score());
+        //println!("LOOCV = {}", mixi_moe.loocv_score());
     }
 
-    fn ftest(x: &Array2<f64>) -> Array2<f64> {
-        let mut y = (x.column(0).to_owned() * x.column(0)).insert_axis(Axis(1));
-        y = &y + (x.column(1).to_owned() * x.column(1)).insert_axis(Axis(1));
-        y = &y * (x.column(2).insert_axis(Axis(1)).mapv(|v| v + 1.));
+    fn ftest(x: &Array2<f64>) -> Array1<f64> {
+        let mut y = x.column(0).to_owned() * x.column(0);
+        y = &y + (x.column(1).to_owned() * x.column(1));
+        y = &y * (x.column(2).mapv(|v| v + 1.));
         y
     }
 

ego/src/gpmix/mod.rs

@@ -7,7 +7,7 @@ use egobox_gp::ThetaTuning;
 use egobox_moe::{
     Clustering, CorrelationSpec, GpMixtureParams, MixtureGpSurrogate, RegressionSpec,
 };
-use ndarray::ArrayView2;
+use ndarray::{ArrayView1, ArrayView2};
 
 use linfa::ParamGuard;
 
@@ -51,22 +51,22 @@ impl SurrogateBuilder for GpMixtureParams<f64> {
 
     fn train(
         &self,
-        xt: &ArrayView2<f64>,
-        yt: &ArrayView2<f64>,
+        xt: ArrayView2<f64>,
+        yt: ArrayView1<f64>,
     ) -> Result<Box<dyn MixtureGpSurrogate>> {
         let checked = self.check_ref()?;
-        let moe = checked.train(xt, yt)?;
+        let moe = checked.train(&xt, &yt)?;
         Ok(moe).map(|moe| Box::new(moe) as Box<dyn MixtureGpSurrogate>)
     }
 
     fn train_on_clusters(
         &self,
-        xt: &ArrayView2<f64>,
-        yt: &ArrayView2<f64>,
+        xt: ArrayView2<f64>,
+        yt: ArrayView1<f64>,
         clustering: &Clustering,
     ) -> Result<Box<dyn MixtureGpSurrogate>> {
         let checked = self.check_ref()?;
-        let moe = checked.train_on_clusters(xt, yt, clustering)?;
+        let moe = checked.train_on_clusters(&xt, &yt, clustering)?;
         Ok(moe).map(|moe| Box::new(moe) as Box<dyn MixtureGpSurrogate>)
     }
 }

ego/src/solver/egor_impl.rs

@@ -97,7 +97,7 @@ where
         &self,
         model_name: &str,
         xt: &ArrayBase<impl Data<Elem = f64>, Ix2>,
-        yt: &ArrayBase<impl Data<Elem = f64>, Ix2>,
+        yt: &ArrayBase<impl Data<Elem = f64>, Ix1>,
         make_clustering: bool,
         optimize_theta: bool,
         clustering: Option<&Clustering>,
@@ -114,7 +114,7 @@ where
         {
             info!("{} Clustering and training...", model_name);
             let model = builder
-                .train(&xt.view(), &yt.view())
+                .train(xt.view(), yt.view())
                 .expect("GP training failure");
             info!(
                 "... {} trained ({} / {})",
@@ -155,7 +155,7 @@ where
             builder.set_theta_tunings(&theta_tunings);
 
             let model = builder
-                .train_on_clusters(&xt.view(), &yt.view(), clustering)
+                .train_on_clusters(xt.view(), yt.view(), clustering)
                 .expect("GP training failure");
             model
         }
@@ -204,13 +204,7 @@ where
                 self.make_clustered_surrogate(
                     &name,
                     &state.data.as_ref().unwrap().0,
-                    &state
-                        .data
-                        .as_ref()
-                        .unwrap()
-                        .1
-                        .slice(s![.., k..k + 1])
-                        .to_owned(),
+                    &state.data.as_ref().unwrap().1.slice(s![.., k]).to_owned(),
                     false,
                     true,
                     state.clusterings.as_ref().unwrap()[k].as_ref(),
@@ -412,7 +406,7 @@ where
                     self.make_clustered_surrogate(
                         &name,
                         &xt,
-                        &yt.slice(s![.., k..k + 1]).to_owned(),
+                        &yt.slice(s![.., k]).to_owned(),
                         make_clustering,
                         optimize_theta,
                         clusterings[k].as_ref(),
@@ -595,7 +589,7 @@ where
                                 .unwrap()
                                 .view(),
                         )
-                        .unwrap()[[0, 0]]
+                        .unwrap()[0]
                         / scale_cstr
                 };
                 #[cfg(feature = "nlopt")]
@@ -684,7 +678,7 @@ where
             Ok(res)
         } else {
             let x = &xk.view().insert_axis(Axis(0));
-            let pred = obj_model.predict(x)?[[0, 0]];
+            let pred = obj_model.predict(x)?[0];
             let var = obj_model.predict_var(x)?[[0, 0]];
             let conf = match self.config.q_ei {
                 QEiStrategy::KrigingBeliever => 0.,
@@ -694,7 +688,7 @@ where
             };
             res.push(pred + conf * f64::sqrt(var));
             for cstr_model in cstr_models {
-                res.push(cstr_model.predict(x)?[[0, 0]]);
+                res.push(cstr_model.predict(x)?[0]);
             }
             Ok(res)
         }

ego/src/solver/trego.rs

@@ -156,7 +156,7 @@ impl<SB: SurrogateBuilder + DeserializeOwned> EgorSolver<SB> {
                                 .unwrap()
                                 .view(),
                         )
-                        .unwrap()[[0, 0]]
+                        .unwrap()[0]
                         / scale_cstr
                 };
                 #[cfg(feature = "nlopt")]

ego/src/types.rs

@@ -3,7 +3,7 @@ use crate::{errors::Result, EgorState};
 use argmin::core::CostFunction;
 use egobox_moe::{Clustering, MixtureGpSurrogate, ThetaTuning};
 use linfa::Float;
-use ndarray::{Array1, Array2, ArrayView2};
+use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
 use serde::{Deserialize, Serialize};
 
 /// Optimization result
@@ -129,15 +129,15 @@ pub trait SurrogateBuilder: Clone + Serialize + Sync {
     /// Train the surrogate with given training dataset (x, y)
     fn train(
         &self,
-        xt: &ArrayView2<f64>,
-        yt: &ArrayView2<f64>,
+        xt: ArrayView2<f64>,
+        yt: ArrayView1<f64>,
     ) -> Result<Box<dyn MixtureGpSurrogate>>;
 
     /// Train the surrogate with given training dataset (x, y) and given clustering
     fn train_on_clusters(
         &self,
-        xt: &ArrayView2<f64>,
-        yt: &ArrayView2<f64>,
+        xt: ArrayView2<f64>,
+        yt: ArrayView1<f64>,
         clustering: &Clustering,
     ) -> Result<Box<dyn MixtureGpSurrogate>>;
 }