From 3a5b108d2c1164c72a6c9c6f2e630d6d084a47c6 Mon Sep 17 00:00:00 2001
From: Saves Paul <paul.saves@alumni.enac.fr>
Date: Tue, 14 Nov 2023 16:36:42 +0100
Subject: [PATCH] Update NestedLHS to use design space (#473)

* Update nestedMFK for design_sapce

* Implement a test of mi-mf sampling

* Change options order
---
 smt/applications/mfk.py                       | 43 ++++++++++++++++---
 .../tests/test_mfk_mfkpls_mixed.py            | 12 +++---
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/smt/applications/mfk.py b/smt/applications/mfk.py
index b6bf8538e..bc01a4eac 100644
--- a/smt/applications/mfk.py
+++ b/smt/applications/mfk.py
@@ -34,10 +34,11 @@
 from smt.utils.misc import standardization
 
 from smt.surrogate_models.krg_based import compute_n_param
+from smt.utils.design_space import ensure_design_space
 
 
 class NestedLHS(object):
-    def __init__(self, nlevel, xlimits, random_state=None):
+    def __init__(self, nlevel, xlimits=None, design_space=None, random_state=None):
         """
         Constructor where values of options can be passed in.
 
@@ -49,12 +50,26 @@ def __init__(self, nlevel, xlimits, random_state=None):
         xlimits : ndarray
             The interval of the domain in each dimension with shape (nx, 2)
 
+        design_space : DesignSpace
+            The design space with type and bounds for every design variable
+
         random_state : Numpy RandomState object or seed number which controls random draws
 
         """
         self.nlevel = nlevel
-        self.xlimits = xlimits
         self.random_state = random_state
+        if xlimits is None and design_space is None:
+            raise ValueError(
+                "Either xlimits or design_space should be specified to have bounds for the sampling."
+            )
+        elif xlimits is not None and design_space is not None:
+            raise ValueError(
+                "Use either design_space for mixed inputs or xlimits for continuous one. Please avoid overspecification."
+            )
+        elif xlimits is not None:
+            self.design_space = ensure_design_space(xlimits=xlimits)
+        else:
+            self.design_space = design_space
 
     def __call__(self, nb_samples_hifi):
         """
@@ -85,14 +100,30 @@ def __call__(self, nb_samples_hifi):
             raise ValueError("nt must be a list of decreasing integers")
 
         doe = []
-        p0 = LHS(xlimits=self.xlimits, criterion="ese", random_state=self.random_state)
-        doe.append(p0(nt[0]))
+        p0 = LHS(
+            xlimits=np.array(self.design_space.get_unfolded_num_bounds()),
+            criterion="ese",
+            random_state=self.random_state,
+        )
+        p0nt0 = p0(nt[0])
+        if self.design_space:
+            p0nt0, _ = self.design_space.correct_get_acting(p0nt0)
+            p0nt0, _ = self.design_space.fold_x(p0nt0)
+
+        doe.append(p0nt0)
 
         for i in range(1, self.nlevel):
             p = LHS(
-                xlimits=self.xlimits, criterion="ese", random_state=self.random_state
+                xlimits=np.array(self.design_space.get_unfolded_num_bounds()),
+                criterion="ese",
+                random_state=self.random_state,
             )
-            doe.append(p(nt[i]))
+            pnti = p(nt[i])
+            if self.design_space:
+                pnti, _ = self.design_space.correct_get_acting(pnti)
+                pnti, _ = self.design_space.fold_x(pnti)
+
+            doe.append(pnti)
 
         for i in range(1, self.nlevel)[::-1]:
             ind = []
diff --git a/smt/applications/tests/test_mfk_mfkpls_mixed.py b/smt/applications/tests/test_mfk_mfkpls_mixed.py
index ccbb5587f..c209b42b1 100644
--- a/smt/applications/tests/test_mfk_mfkpls_mixed.py
+++ b/smt/applications/tests/test_mfk_mfkpls_mixed.py
@@ -27,6 +27,8 @@
     MixedIntegerSurrogateModel,
 )
 
+from smt.applications import NestedLHS
+
 from smt.sampling_methods import LHS
 
 from smt.surrogate_models import (
@@ -387,7 +389,9 @@ def run_mfk_mixed_example(self):
 
         # HF training data:
         n_train_HF = 20  # training set size
-        xt_HF, is_acting_t_HF = ds.sample_valid_x(n_train_HF)
+        n_train_LF = 40
+        nlhs = NestedLHS(nlevel=2, design_space=ds)
+        xt_LF, xt_HF = nlhs(n_train_HF)
 
         y1t_HF = np.zeros(n_train_HF)  # obj 1
         y2t_HF = np.zeros(n_train_HF)  # obj 2
@@ -400,12 +404,6 @@ def run_mfk_mixed_example(self):
 
         for krg_method in KRG_METHODS:
             if "mfk" in krg_method:  # if multifi compute LF validation data
-                n_train_LF = 40  # training set size   IF MULTIFI
-                sampling = MixedIntegerSamplingMethod(LHS, ds, criterion="ese")
-                xt_LF = np.concatenate(
-                    (xt_HF, sampling.expand_lhs(xt_HF, n_train_LF - n_train_HF)), axis=0
-                )
-
                 y1t_LF = np.zeros(n_train_LF)  # obj 1
                 y2t_LF = np.zeros(n_train_LF)  # obj 2
                 for i in range(n_train_LF):