Refactoring to fit and eval PORTALS GPs all at once, not with ModelList

src/mitim_modules/freegsu/FREEGSUmain.py

@@ -33,7 +33,7 @@ def default_namelist(optimization_options):
 
 
 class freegsu(STRATEGYtools.opt_evaluator):
-    def __init__(self, folder, namelist=None, function_parameters={}):
+    def __init__(self, folder, namelist=None, function_parameters={}, **kwargs):
         print(
             "\n-----------------------------------------------------------------------------------------"
         )
@@ -49,6 +49,7 @@ def __init__(self, folder, namelist=None, function_parameters={}):
             folder,
             namelist=namelist,
             default_namelist_function=default_namelist if (namelist is None) else None,
+            **kwargs
         )
 
     def prep(

src/mitim_modules/maestro/utils/PORTALSbeat.py

@@ -253,7 +253,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr
         if use_previous_surrogate_data and ('portals_surrogate_data_file' in self.maestro_instance.parameters_trans_beat):
             if 'surrogateOptions' not in self.optimization_options:
                 self.optimization_options['surrogateOptions'] = {}
-            self.optimization_options['surrogateOptions']["extrapointsFile"] = self.maestro_instance.parameters_trans_beat['portals_surrogate_data_file']
+            self.optimization_options['surrogateOptions']["add_data_from_file"] = self.maestro_instance.parameters_trans_beat['portals_surrogate_data_file']
 
             self.folder_starting_point = self.maestro_instance.parameters_trans_beat['portals_last_run_folder']
 
@@ -295,7 +295,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr
         # In the situation where the last radial location moves, I cannot reuse that surrogate data
         if last_radial_location_moved and reusing_surrogate_data:
             print('\t\t- Last radial location was moved, so surrogate data will not be reused for that specific location')
-            self.optimization_options['surrogateOptions']["extrapointsModelsAvoidContent"] = ['Tar',f'_{len(self.MODELparameters[strKeys])}']
+            self.optimization_options['surrogateOptions']["add_data_to_modelsAvoidContent"] = ['Tar',f'_{len(self.MODELparameters[strKeys])}']
             self.try_flux_match_only_for_first_point = False
 
     def _inform_save(self):

src/mitim_modules/portals/PORTALSmain.py

@@ -63,8 +63,8 @@ def default_namelist(optimization_options, CGYROrun=False):
             }
 
     # Surrogate
-    optimization_options["surrogateOptions"]["selectSurrogate"] = partial(
-        PORTALStools.selectSurrogate, CGYROrun=CGYROrun
+    optimization_options["surrogateOptions"]["selectSurrogates"] = partial(
+        PORTALStools.selectSurrogates, CGYROrun=CGYROrun
     )
 
     optimization_options["surrogateOptions"]["ensure_within_bounds"] = True
@@ -86,7 +86,10 @@ def __init__(
         self, 
         folder,                             # Folder where the PORTALS workflow will be run
         namelist=None,                      # If None, default namelist will be used. If not None, it will be read and used
-        TensorsType=torch.double,           # Type of tensors to be used (torch.float, torch.double)
+        tensor_opts = {
+            "dtype": torch.double,
+            "device": torch.device("cpu"),
+        },
         CGYROrun=False,                     # If True, use CGYRO defaults for best optimization practices
         portals_transformation_variables = None,          # If None, use defaults for both main and trace
         portals_transformation_variables_trace = None,
@@ -109,7 +112,7 @@ def __init__(
         super().__init__(
             folder,
             namelist=namelist,
-            TensorsType=TensorsType,
+            tensor_opts=tensor_opts,
             default_namelist_function=(
                 partial(default_namelist, CGYROrun=CGYROrun)
                 if (namelist is None)
@@ -186,8 +189,6 @@ def __init__(
 		---------------------------------------------
 		"""
 
-
-
         (
             portals_transformation_variables,
             portals_transformation_variables_trace,
@@ -315,7 +316,7 @@ def prep(
             limitsAreRelative=limitsAreRelative,
             cold_start=cold_start,
             hardGradientLimits=hardGradientLimits,
-            dfT=self.dfT,
+            tensor_opts = self.tensor_opts,
             seedInitial=seedInitial,
             checkForSpecies=askQuestions,
             ModelOptions=ModelOptions,
@@ -335,14 +336,14 @@ def prep(
         # Ignore targets in surrogate_data.csv
         # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-        if 'extrapointsModels' not in self.optimization_options['surrogateOptions'] or \
-            self.optimization_options['surrogateOptions']['extrapointsModels'] is None or \
-            len(self.optimization_options['surrogateOptions']['extrapointsModels'])==0:
+        if 'add_data_to_models' not in self.optimization_options['surrogateOptions'] or \
+            self.optimization_options['surrogateOptions']['add_data_to_models'] is None or \
+            len(self.optimization_options['surrogateOptions']['add_data_to_models'])==0:
 
             self._define_reuse_models()
 
         else:
-            print("\t- extrapointsModels already defined, not changing")
+            print("\t- add_data_to_models already defined, not changing")
 
     def _define_reuse_models(self):
         '''
@@ -352,21 +353,21 @@ def _define_reuse_models(self):
             '_5' to avoid reusing position 5
         '''
 
-        self.optimization_options['surrogateOptions']['extrapointsModels'] = []
+        self.optimization_options['surrogateOptions']['add_data_to_models'] = []
 
         # Define avoiders
-        if self.optimization_options['surrogateOptions']['extrapointsModelsAvoidContent'] is None:
-            self.optimization_options['surrogateOptions']['extrapointsModelsAvoidContent'] = ['Tar']
+        if self.optimization_options['surrogateOptions']['add_data_to_modelsAvoidContent'] is None:
+            self.optimization_options['surrogateOptions']['add_data_to_modelsAvoidContent'] = ['Tar']
 
-        # Define extrapointsModels
+        # Define add_data_to_models
         for key in self.surrogate_parameters['surrogate_transformation_variables_lasttime'].keys():
             add_key = True
-            for avoid in self.optimization_options['surrogateOptions']['extrapointsModelsAvoidContent']:
+            for avoid in self.optimization_options['surrogateOptions']['add_data_to_modelsAvoidContent']:
                 if avoid in key:
                     add_key = False
                     break
             if add_key:
-                self.optimization_options['surrogateOptions']['extrapointsModels'].append(key)
+                self.optimization_options['surrogateOptions']['add_data_to_models'].append(key)
 
     def run(self, paramsfile, resultsfile):
         # Read what PORTALS sends

src/mitim_modules/portals/PORTALStools.py

@@ -7,24 +7,48 @@
 from mitim_tools.misc_tools.LOGtools import printMsg as print
 from IPython import embed
 
-def selectSurrogate(output, surrogateOptions, CGYROrun=False):
+def selectSurrogates(outputs, surrogateOptions, CGYROrun=False):
+    '''
+    This divides potentially different outputs into different
+    surrogates to be joined with ModelList 
+    '''
+
+    # Find transition to Targets
+    for iTar, output in enumerate(outputs):
+        if output[2:5] == "Tar":
+            break
 
-    print(
-        f'\t- Selecting surrogate options for "{output}" to be run'
-    )
+    # Find transition to last location of transport
+    for iTra, output in enumerate(outputs):
+        if output.split('_')[-1] == outputs[iTar-1].split('_')[-1]:
+            break
 
-    if output is not None:
-        # If it's a target, just linear
-        if output[2:5] == "Tar":
-            surrogateOptions["TypeMean"] = 1
-            surrogateOptions["TypeKernel"] = 2  # Constant kernel
-        # If it's not, stndard
-        else:
-            surrogateOptions["TypeMean"] = 2  # Linear in gradients, constant in rest
-            surrogateOptions["TypeKernel"] = 1  # RBF
-            # surrogateOptions['ExtraNoise']  = True
+    surrogateOptions_dict = {}
+
+    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
 
-    return surrogateOptions
+    # # 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
 
 def default_portals_transformation_variables(additional_params = []):
     """
@@ -91,7 +115,7 @@ def default_portals_transformation_variables(additional_params = []):
 
     return portals_transformation_variables, portals_transformation_variables_trace
 
-def produceNewInputs(Xorig, output, surrogate_parameters, surrogate_transformation_variables):
+def input_transformation_portals(Xorig, output, surrogate_parameters, surrogate_transformation_variables):
 
     """
     - Xorig will be a tensor (batch1...N,dim) unnormalized (with or without gradients).
@@ -123,9 +147,7 @@ def produceNewInputs(Xorig, output, surrogate_parameters, surrogate_transformati
 	"""
 
     _, num = output.split("_")
-    index = powerstate.indexes_simulation[
-        int(num)
-    ]  # num=1 -> pos=1, so that it takes the second value in vectors
+    index = powerstate.indexes_simulation[int(num)]  # num=1 -> pos=1, so that it takes the second value in vectors
 
     xFit = torch.Tensor().to(X)
     for ikey in surrogate_transformation_variables[output]:
@@ -151,12 +173,14 @@ def produceNewInputs(Xorig, output, surrogate_parameters, surrogate_transformati
 # ----------------------------------------------------------------------
 
 
-def transformPORTALS(X, surrogate_parameters, output):
+def output_transformation_portals(X, surrogate_parameters, outputs):
     """
     1. Make sure all batches are squeezed into a single dimension
     ------------------------------------------------------------------
             E.g.: (batch1,batch2,batch3,dim) -> (batch1*batch2*batch3,dim)
+    Output is a factor that account for all the outputs
     """
+
     shape_orig = np.array(X.shape)
     X = X.view(np.prod(shape_orig[:-1]), shape_orig[-1])
 
@@ -169,14 +193,16 @@ def transformPORTALS(X, surrogate_parameters, output):
     # Produce relevant quantities here (in particular, GB will be used)
     powerstate = constructEvaluationProfiles(X, surrogate_parameters)
 
-    # --- Original model output is in real units, transform to GB here b/c that's how GK codes work
-    factorGB = GBfromXnorm(X, output, powerstate)
-    # --- Ratio of fluxes (quasilinear)
-    factorRat = ratioFactor(X, surrogate_parameters, output, powerstate)
-    # --- Specific to output
-    factorImp = ImpurityGammaTrick(X, surrogate_parameters, output, powerstate)
+    compounded = torch.Tensor().to(X)
+    for output in outputs:
+        # --- Original model output is in real units, transform to GB here b/c that's how GK codes work
+        factorGB = GBfromXnorm(X, output, powerstate)
+        # --- Ratio of fluxes (quasilinear)
+        factorRat = ratioFactor(X, surrogate_parameters, output, powerstate)
+        # --- Specific to output
+        factorImp = ImpurityGammaTrick(X, surrogate_parameters, output, powerstate)
 
-    compounded = factorGB * factorRat * factorImp
+        compounded = torch.cat((compounded, factorGB * factorRat * factorImp), dim=-1)
 
     """
 	3. Go back to the original batching system
@@ -185,7 +211,7 @@ def transformPORTALS(X, surrogate_parameters, output):
 	"""
     shape_orig[-1] = compounded.shape[-1]
     compounded = compounded.view(tuple(shape_orig))
-
+    
     return compounded
 
 
@@ -226,7 +252,7 @@ def computeTurbExchangeIndividual(PexchTurb, powerstate):
     return PexchTurb_integrated
 
 
-# def transformPORTALS(X,Y,Yvar,surrogate_parameters,output):
+# def output_transformation_portals(X,Y,Yvar,surrogate_parameters,output):
 # 	'''
 # 	Transform direct evaluation output to something that the model understands better.
 
@@ -245,14 +271,14 @@ def computeTurbExchangeIndividual(PexchTurb, powerstate):
 # 	return Ytr,Ytr_var
 
 
-# def untransformPORTALS(X, mean, upper, lower, surrogate_parameters, output):
+# def unoutput_transformation_portals(X, mean, upper, lower, surrogate_parameters, output):
 # 	'''
-# 	Transform direct model output to the actual evaluation output (must be the opposite to transformPORTALS)
+# 	Transform direct model output to the actual evaluation output (must be the opposite to output_transformation_portals)
 
 # 		- Receives unnormalized X (batch1,...,dim) to construct QGB (batch1,...,1) corresponding to what output I'm looking at
 # 		- Transforms and produces Y and confidence bounds (batch1,...,)
 
-# 	This untransforms whatever has happened in the transformPORTALS function
+# 	This untransforms whatever has happened in the output_transformation_portals function
 # 	'''
 
 # 	factor = factorProducer(X,surrogate_parameters,output).squeeze(-1)
@@ -375,9 +401,11 @@ def constructEvaluationProfiles(X, surrogate_parameters, recalculateTargets=Fals
     if ("parameters_combined" in surrogate_parameters) and (
         "powerstate" in surrogate_parameters["parameters_combined"]
     ):
+
         powerstate = surrogate_parameters["parameters_combined"]["powerstate"]
 
     else:
+
         powerstate = surrogate_parameters["powerstate"]
 
         if X.shape[0] > 0:

src/mitim_modules/portals/utils/PORTALSanalysis.py

@@ -730,15 +730,15 @@ def __init__(self, gpdict):
                     self._output_variables.append(key)
         for key in self._models:
             if hasattr(self._models[key], 'gpmodel'):
-                if hasattr(self._models[key].gpmodel, 'train_X_usedToTrain'):
-                    xtrain = self._models[key].gpmodel.train_X_usedToTrain.detach().cpu().numpy()
+                if hasattr(self._models[key].gpmodel, 'train_X_use'):
+                    xtrain = self._models[key].gpmodel.train_X_use.detach().cpu().numpy()
                     if len(xtrain.shape) < 2:
                         xtrain = np.atleast_2d(xtrain)
                     if xtrain.shape[1] != len(self._input_variables):
                         xtrain = xtrain.T
                     self._training_inputs[key] = pd.DataFrame(xtrain, columns=self._input_variables)
-                if hasattr(self._models[key].gpmodel, 'train_Y_usedToTrain'):
-                    ytrain = self._models[key].gpmodel.train_Y_usedToTrain.detach().cpu().numpy()
+                if hasattr(self._models[key].gpmodel, 'train_Y_use'):
+                    ytrain = self._models[key].gpmodel.train_Y_use.detach().cpu().numpy()
                     if len(ytrain.shape) < 2:
                         ytrain = np.atleast_2d(ytrain)
                     if ytrain.shape[1] != 1: