Astra run3

src/mitim_modules/freegsu/utils/FREEGSUparams.py

@@ -108,7 +108,7 @@ def createProblemParameters(
         dvs_min.extend(dvs_min2)
         dvs_max.extend(dvs_max2)
 
-        transformation = produceNewInputs
+        transformation = input_transform_freegs
     else:
         transformation = None
 
@@ -376,7 +376,7 @@ def extractCont(x, cont):
     return v
 
 
-def produceNewInputs(X, output, bounds, ParamProfile):
+def input_transform_freegs(X, output, bounds, ParamProfile):
     """
     X will be a tensor (with or without gradients) batch*dim,   unnormalized
     """

src/mitim_modules/maestro/tmp_tests/maestro_test1.py

@@ -46,7 +46,7 @@
                         "MODELparameters": { "RoaLocations": [0.35,0.55,0.75,0.875,0.9],
                                             "ProfilesPredicted": ["te", "ti", "ne"],
                                             "Physics_options": {"TypeTarget": 3},
-                                             "transport_model": {"turbulence":'TGLF',"TGLFsettings": 6, "extraOptionsTGLF": {'USE_BPER':True}}},
+                                             "transport_model": {"TGLFsettings": 6, "extraOptionsTGLF": {'USE_BPER':True}}},
                         "INITparameters": {"FastIsThermal": True, "removeIons": [5,6], "quasineutrality": True},
                         "optimization_options": {
                             "convergence_options": {

src/mitim_modules/maestro/utils/EPEDbeat.py

@@ -10,7 +10,7 @@
 from mitim_tools.popcon_tools import FunctionalForms
 from mitim_tools.misc_tools.LOGtools import printMsg as print
 from mitim_modules.maestro.utils.MAESTRObeat import beat
-from mitim_modules.powertorch.physics import CALCtools
+from mitim_modules.powertorch.utils import CALCtools
 from IPython import embed
 
 # <> Function to interpolate a curve <> 

src/mitim_modules/maestro/utils/PORTALSbeat.py

@@ -44,7 +44,7 @@ def prepare(self,
 
             profiles = self.profiles_current
             for i in range(len(profiles.Species)):
-                data_df = pd.read_csv(__mitimroot__ / "src" / "mitim_modules" / "powertorch" / "physics" / "radiation_chebyshev.csv")
+                data_df = pd.read_csv(__mitimroot__ / "src" / "mitim_modules" / "powertorch" / "physics_models" / "radiation_chebyshev.csv")
                 if not (data_df['Ion'].str.lower()==profiles.Species[i]["N"].lower()).any():
                     print(f"\t\t- {profiles.Species[i]['N']} not found in radiation table, looking for closest Z (+- 5) USING THE Z SPECIFIED IN THE INPUT.GACODE (fully stripped assumption)",typeMsg='w')
                     # Find closest Z
@@ -56,7 +56,7 @@ def prepare(self,
 
                     new_name = data_df['Ion'][iZ]
 
-                    print(f"\t\t\t- Changing name of ion from {profiles.Species[i]["N"]} ({profiles.Species[i]["Z"]}) to {new_name} ({Z[iZ]})")
+                    print(f'\t\t\t- Changing name of ion from {profiles.Species[i]["N"]} ({profiles.Species[i]["Z"]}) to {new_name} ({Z[iZ]})')
 
                     profiles.profiles['name'][i] = profiles.Species[i]["N"] = new_name
 
@@ -344,7 +344,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['surrogate_options']["extrapointsModelsAvoidContent"] = ['Tar',f'_{len(self.MODELparameters[strKeys])}']
+            self.optimization_options['surrogate_options']["extrapointsModelsAvoidContent"] = ['_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

@@ -7,14 +7,13 @@
 from collections import OrderedDict
 from mitim_tools.misc_tools import IOtools
 from mitim_tools.gacode_tools import PROFILEStools
-from mitim_tools.gacode_tools.utils import PORTALSinteraction
 from mitim_modules.portals import PORTALStools
 from mitim_modules.portals.utils import (
     PORTALSinit,
     PORTALSoptimization,
     PORTALSanalysis,
 )
-from mitim_modules.powertorch.physics import TRANSPORTtools, TARGETStools
+from mitim_modules.powertorch.physics_models import targets_analytic, transport_tgyro, transport_cgyro
 from mitim_tools.opt_tools import STRATEGYtools
 from mitim_tools.opt_tools.utils import BOgraphics
 from mitim_tools.misc_tools.LOGtools import printMsg as print
@@ -64,7 +63,7 @@ def default_namelist(optimization_options, CGYROrun=False):
     optimization_options['acquisition_options']['relative_improvement_for_stopping'] = 1e-2
 
     # Surrogate
-    optimization_options["surrogate_options"]["selectSurrogate"] = partial(PORTALStools.selectSurrogate, CGYROrun=CGYROrun)
+    optimization_options["surrogate_options"]["surrogate_selection"] = partial(PORTALStools.surrogate_selection_portals, CGYROrun=CGYROrun)
 
     if CGYROrun:
         # CGYRO runs should prioritize accuracy
@@ -171,7 +170,7 @@ def __init__(
                 "Tfast_ratio": False,  # Keep the ratio of Tfast/Te constant throughout the Te evolution
                 "ensureMachNumber": None,  # Change w0 to match this Mach number when Ti varies
             },
-            "transport_model": {"turbulence":'TGLF',"TGLFsettings": 6, "extraOptionsTGLF": {}}
+            "transport_model": {"TGLFsettings": 6, "extraOptionsTGLF": {}}
         }
 
         for key in self.MODELparameters.keys():
@@ -193,8 +192,12 @@ def __init__(
 		"""
 
         # Selection of model
-        transport_evaluator = TRANSPORTtools.tgyro_model
-        targets_evaluator = TARGETStools.analytical_model
+        if CGYROrun:
+            transport_evaluator = transport_cgyro.cgyro_model
+        else:
+            transport_evaluator = transport_tgyro.tgyro_model
+
+        targets_evaluator = targets_analytic.analytical_model
 
         self.PORTALSparameters = {
             "percentError": [5,10,1],  # (%) Error (std, in percent) of model evaluation [TGLF (treated as minimum if scan trick), NEO, TARGET]
@@ -341,15 +344,15 @@ def _define_reuse_models(self):
         '''
         The user can define a list of strings to avoid reusing surrogates.
         e.g. 
-            'Tar' to avoid reusing targets
+            '_tar' to avoid reusing targets
             '_5' to avoid reusing position 5
         '''
 
         self.optimization_options['surrogate_options']['extrapointsModels'] = []
 
         # Define avoiders
         if self.optimization_options['surrogate_options']['extrapointsModelsAvoidContent'] is None:
-            self.optimization_options['surrogate_options']['extrapointsModelsAvoidContent'] = ['Tar']
+            self.optimization_options['surrogate_options']['extrapointsModelsAvoidContent'] = ['_tar']
 
         # Define extrapointsModels
         for key in self.surrogate_parameters['surrogate_transformation_variables_lasttime'].keys():
@@ -416,17 +419,16 @@ def scalarized_objective(self, Y):
 		-------------------------------------------------------------------------
 		Prepare transport dictionary
 		-------------------------------------------------------------------------
-			Note: var_dict['QeTurb'] must have shape (dim1...N, num_radii)
+			Note: var_dict['Qe_tr_turb'] must have shape (dim1...N, num_radii)
 		"""
 
         var_dict = {}
         for of in ofs_ordered_names:
-            var, _ = of.split("_")
+
+            var = '_'.join(of.split("_")[:-1])
             if var not in var_dict:
                 var_dict[var] = torch.Tensor().to(Y)
-            var_dict[var] = torch.cat(
-                (var_dict[var], Y[..., ofs_ordered_names == of]), dim=-1
-            )
+            var_dict[var] = torch.cat((var_dict[var], Y[..., ofs_ordered_names == of]), dim=-1)
 
         """
 		-------------------------------------------------------------------------
@@ -436,7 +438,7 @@ def scalarized_objective(self, Y):
 				  res must have shape (dim1...N)
 		"""
 
-        of, cal, _, res = PORTALSinteraction.calculate_residuals(self.powerstate, self.PORTALSparameters,specific_vars=var_dict)
+        of, cal, _, res = PORTALStools.calculate_residuals(self.powerstate, self.PORTALSparameters,specific_vars=var_dict)
 
         return of, cal, res
 
@@ -468,7 +470,7 @@ def check_flags(self):
                 print("\t- Requested fineTargetsResolution, so running powerstate target calculations",typeMsg="w")
                 self.PORTALSparameters["TargetCalc"] = "powerstate"
 
-        if not issubclass(self.PORTALSparameters["transport_evaluator"], TRANSPORTtools.tgyro_model) and (self.PORTALSparameters["TargetCalc"] == "tgyro"):
+        if not issubclass(self.PORTALSparameters["transport_evaluator"], transport_tgyro.tgyro_model) and (self.PORTALSparameters["TargetCalc"] == "tgyro"):
             print("\t- Requested TGYRO targets, but transport evaluator is not tgyro, so changing to powerstate",typeMsg="w")
             self.PORTALSparameters["TargetCalc"] = "powerstate"
 
@@ -480,9 +482,6 @@ def check_flags(self):
             print("\t- In PORTALS TGYRO evaluations, we need to not recompute gradients (GradientsType=0)",typeMsg="i")
             self.MODELparameters["Physics_options"]["GradientsType"] = 0
 
-        if 'TargetType' in self.MODELparameters["Physics_options"]:
-            raise Exception("\t- TargetType is not used in PORTALS anymore")
-
         if self.PORTALSparameters["TargetCalc"] == "tgyro" and self.PORTALSparameters['profiles_postprocessing_fun'] is not None:
             print("\t- Requested custom modification of postprocessing function but targets from tgyro... are you sure?",typeMsg="q")
 
@@ -555,7 +554,7 @@ def reuseTrainingTabular(
                     self_copy.powerstate.TransportOptions["transport_evaluator"] = None
                     self_copy.powerstate.TargetOptions["ModelOptions"]["TypeTarget"] = "powerstate"
                 else:
-                    self_copy.powerstate.TransportOptions["transport_evaluator"] = TRANSPORTtools.tgyro_model
+                    self_copy.powerstate.TransportOptions["transport_evaluator"] = transport_tgyro.tgyro_model
 
                 _, dictOFs = runModelEvaluator(
                     self_copy,
@@ -571,7 +570,7 @@ def reuseTrainingTabular(
                 # ------------------------------------------------------------------------------------
 
                 for i in dictOFs:
-                    if "Tar" in i:
+                    if "_tar" in i:
                         print(f"Changing {i} in file")
 
                         optimization_data.data[i].iloc[numPORTALS] = dictOFs[i]["value"].cpu().numpy().item()
@@ -646,67 +645,49 @@ def runModelEvaluator(
     return powerstate, dictOFs
 
 def map_powerstate_to_portals(powerstate, dictOFs):
-    """
-    """
 
     for var in powerstate.ProfilesPredicted:
         # Write in OFs
         for i in range(powerstate.plasma["rho"].shape[1] - 1): # Ignore position 0, which is rho=0
             if var == "te":
-                var0, var1 = "Qe", "Pe"
+                var0, var1 = "Qe", "QeMWm2"
             elif var == "ti":
-                var0, var1 = "Qi", "Pi"
+                var0, var1 = "Qi", "QiMWm2"
             elif var == "ne":
                 var0, var1 = "Ge", "Ce"
             elif var == "nZ":
                 var0, var1 = "GZ", "CZ"
             elif var == "w0":
-                var0, var1 = "Mt", "Mt"
+                var0, var1 = "Mt", "MtJm2"
 
             """
             TRANSPORT calculation
             ---------------------
             """
 
-            dictOFs[f"{var0}Turb_{i+1}"]["value"] = powerstate.plasma[
-                f"{var1}_tr_turb"
-            ][0, i+1]
-            dictOFs[f"{var0}Turb_{i+1}"]["error"] = powerstate.plasma[
-                f"{var1}_tr_turb_stds"
-            ][0, i+1]
+            dictOFs[f"{var0}_tr_turb_{i+1}"]["value"] = powerstate.plasma[f"{var1}_tr_turb"][0, i+1]
+            dictOFs[f"{var0}_tr_turb_{i+1}"]["error"] = powerstate.plasma[f"{var1}_tr_turb_stds"][0, i+1]
 
-            dictOFs[f"{var0}Neo_{i+1}"]["value"] = powerstate.plasma[
-                f"{var1}_tr_neo"
-            ][0, i+1]
-            dictOFs[f"{var0}Neo_{i+1}"]["error"] = powerstate.plasma[
-                f"{var1}_tr_neo_stds"
-            ][0, i+1]
+            dictOFs[f"{var0}_tr_neo_{i+1}"]["value"] = powerstate.plasma[f"{var1}_tr_neo"][0, i+1]
+            dictOFs[f"{var0}_tr_neo_{i+1}"]["error"] = powerstate.plasma[f"{var1}_tr_neo_stds"][0, i+1]
 
             """
             TARGET calculation
             ---------------------
                 If that radius & profile position has target, evaluate
             """
 
-            dictOFs[f"{var0}Tar_{i+1}"]["value"] = powerstate.plasma[f"{var1}"][
-                0, i+1
-            ]
-            dictOFs[f"{var0}Tar_{i+1}"]["error"] = powerstate.plasma[
-                f"{var1}_stds"
-            ][0, i+1]
+            dictOFs[f"{var0}_tar_{i+1}"]["value"] = powerstate.plasma[f"{var1}"][0, i+1]
+            dictOFs[f"{var0}_tar_{i+1}"]["error"] = powerstate.plasma[f"{var1}_stds"][0, i+1]
 
     """
     Turbulent Exchange
     ------------------
     """
     if 'PexchTurb_1' in dictOFs:
         for i in range(powerstate.plasma["rho"].shape[1] - 1):
-            dictOFs[f"PexchTurb_{i+1}"]["value"] = powerstate.plasma["PexchTurb"][
-                0, i+1
-            ]
-            dictOFs[f"PexchTurb_{i+1}"]["error"] = powerstate.plasma[
-                "PexchTurb_stds"
-            ][0, i+1]
+            dictOFs[f"PexchTurb_{i+1}"]["value"] = powerstate.plasma["PexchTurb"][0, i+1]
+            dictOFs[f"PexchTurb_{i+1}"]["error"] = powerstate.plasma["PexchTurb_stds"][0, i+1]
 
     return dictOFs