trying to recover the modellist

pabloprf · Dec 11, 2024 · 8ad1247 · 8ad1247
1 parent 7156b7f
commit 8ad1247
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 37 deletions.
diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py
@@ -25,22 +25,28 @@ def selectSurrogates(outputs, surrogateOptions, CGYROrun=False):
 
     surrogateOptions_dict = {}
 
-    # Turbulent and Neoclassical
-
-    surrogateOptions_dict[iTra] = copy.deepcopy(surrogateOptions)
-    surrogateOptions_dict[iTra]["TypeMean"] = 2  # Linear in gradients, constant in rest
-    surrogateOptions_dict[iTra]["TypeKernel"] = 1  # RBF
-    # surrogateOptions_dict[len(output)]['ExtraNoise']  = True
-
-    surrogateOptions_dict[iTar] = copy.deepcopy(surrogateOptions)
-    surrogateOptions_dict[iTar]["TypeMean"] = 2  # Linear in gradients, constant in rest
-    surrogateOptions_dict[iTar]["TypeKernel"] = 1  # RBF
-    # surrogateOptions_dict[len(output)]['ExtraNoise']  = True
-
-    # Targets (If it's a target, just linear)
-    surrogateOptions_dict[len(outputs)] = copy.deepcopy(surrogateOptions)
-    surrogateOptions_dict[len(outputs)]["TypeMean"] = 1
-    surrogateOptions_dict[len(outputs)]["TypeKernel"] = 2  # Constant kernel
+    for i in range(len(outputs)):
+        # Turbulent and Neoclassical at inner locations
+        surrogateOptions_dict[i+1] = copy.deepcopy(surrogateOptions)
+        surrogateOptions_dict[i+1]["TypeMean"] = 1  # Linear
+        surrogateOptions_dict[i+1]["TypeKernel"] = 1  # RBF
+
+
+
+    # # Turbulent and Neoclassical at inner locations
+    # surrogateOptions_dict[iTra] = copy.deepcopy(surrogateOptions)
+    # surrogateOptions_dict[iTra]["TypeMean"] = 1  # Linear
+    # surrogateOptions_dict[iTra]["TypeKernel"] = 1  # RBF
+
+    # # Turbulent and Neoclassical at outer location (generally less variables)
+    # surrogateOptions_dict[iTar] = copy.deepcopy(surrogateOptions)
+    # surrogateOptions_dict[iTar]["TypeMean"] = 1  # Linear
+    # surrogateOptions_dict[iTar]["TypeKernel"] = 1  # RBF
+
+    # # Targets (If it's a target, just linear)
+    # surrogateOptions_dict[len(outputs)] = copy.deepcopy(surrogateOptions)
+    # surrogateOptions_dict[len(outputs)]["TypeMean"] = 1 # Linear
+    # surrogateOptions_dict[len(outputs)]["TypeKernel"] = 2  # Constant kernel
 
     return surrogateOptions_dict
 

diff --git a/src/mitim_tools/opt_tools/BOTORCHtools.py b/src/mitim_tools/opt_tools/BOTORCHtools.py
@@ -40,7 +40,7 @@ def __init__(
         input_transform=None,
         outcome_transform=None,
         surrogateOptions={},
-        variables=None,
+        #variables=None,
         train_X_added=torch.Tensor([]),
         train_Y_added=torch.Tensor([]),
         train_Yvar_added=torch.Tensor([]),
@@ -214,12 +214,12 @@ def __init__(
         elif TypeMean == 1:
             self.mean_module = gpytorch.means.linear_mean.LinearMean(
                 self.ard_num_dims, batch_shape=self._aug_batch_shape, bias=True )
-        elif TypeMean == 2:
-            self.mean_module = MITIM_LinearMeanGradients(
-                batch_shape=self._aug_batch_shape, variables=variables )
-        elif TypeMean == 3:
-            self.mean_module = MITIM_CriticalGradient(
-                batch_shape=self._aug_batch_shape, variables=variables )
+        # elif TypeMean == 2:
+        #     self.mean_module = MITIM_LinearMeanGradients(
+        #         batch_shape=self._aug_batch_shape, variables=variables )
+        # elif TypeMean == 3:
+        #     self.mean_module = MITIM_CriticalGradient(
+        #         batch_shape=self._aug_batch_shape, variables=variables )
 
         """
 		-----------------------------------------------------------------------
@@ -890,7 +890,7 @@ class MITIM_LinearMeanGradients(gpytorch.means.mean.Mean):
     def __init__(self, batch_shape=torch.Size(), variables=None, **kwargs):
         super().__init__()
 
-        # Indeces of variables that are gradient, so subject to CG behavior
+        # Indeces of variables that are gradient, so subject to critical gradient behavior
         grad_vector = []
         if variables is not None:
             for i, variable in enumerate(variables):

diff --git a/src/mitim_tools/opt_tools/STEPtools.py b/src/mitim_tools/opt_tools/STEPtools.py
@@ -103,7 +103,7 @@ def fit_step(self, avoidPoints=None, fit_output_contains=None):
     def _fit_multioutput_model(self):
 
         # Base model
-        self.GP["mo_model"] = SURROGATEtools.surrogate_model.simple(
+        self.GP["mo_model"] = SURROGATEtools.surrogate_model.only_define(
             self.x,
             self.y,
             self.yvar,

diff --git a/src/mitim_tools/opt_tools/SURROGATEtools.py b/src/mitim_tools/opt_tools/SURROGATEtools.py
@@ -27,7 +27,7 @@ class surrogate_model:
     """
 
     @classmethod
-    def simple(cls, *args, **kwargs):
+    def only_define(cls, *args, **kwargs):
         # Create an instance of the class
         instance = cls.__new__(cls)
         # Initialize the parameters manually
@@ -202,15 +202,15 @@ def __init__(
 
         self.normalization_pass(input_transform, outcome_transform)
 
-        self.variables = (
-            self.surrogate_transformation_variables[self.outputs[0]]
-            if (
-                (self.outputs is not None)
-                and ("surrogate_transformation_variables" in self.__dict__)
-                and (self.surrogate_transformation_variables is not None)
-            )
-            else None
-        )
+        # self.variables = (
+        #     self.surrogate_transformation_variables[self.outputs[0]]
+        #     if (
+        #         (self.outputs is not None)
+        #         and ("surrogate_transformation_variables" in self.__dict__)
+        #         and (self.surrogate_transformation_variables is not None)
+        #     )
+        #     else None
+        # )
 
         # *************************************************************************************
         # Model
@@ -225,7 +225,7 @@ def __init__(
             input_transform=input_transform,
             outcome_transform=outcome_transform,
             surrogateOptions=self.surrogateOptions,
-            variables=self.variables,
+            #variables=self.variables,
             train_X_added=self.train_X_added,
             train_Y_added=self.train_Y_added,
             train_Yvar_added=self.train_Yvar_added,

diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py
@@ -24,7 +24,7 @@
 
 # Initialize class
 portals_fun = PORTALSmain.portals(folderWork)
-portals_fun.optimization_options["BO_iterations"] = 1
+portals_fun.optimization_options["BO_iterations"] = 5
 portals_fun.optimization_options["initial_training"] = 3
 portals_fun.MODELparameters["RhoLocations"] = [0.25, 0.45, 0.65, 0.85]
 portals_fun.INITparameters["removeFast"] = True