Extended Lengyel model as standalone module and MAESTRO beat

pyproject.toml

@@ -46,9 +46,11 @@ dependencies = [
   "pyyaml",
   "megpy>=2.0.3",
   "fibe>=0.3",
+  "periodictable",
   "vmecpp",
-  # "quends @ git+https://github.com/sandialabs/quends.git"
   "scikit-image",   # Stricly not for MITIM, but good to have for pygacode
+  # "extended-lengyel @ git+https://github.com/cfs-energy/extended-lengyel.git",
+  # "quends @ git+https://github.com/sandialabs/quends.git"
   #"onnx2pytorch",   # Stricly not for MITIM, but good to use ONNX models
 ]
 

src/mitim_modules/maestro/MAESTROmain.py

@@ -9,9 +9,10 @@
 from mitim_tools import __version__, __mitimroot__
 from IPython import embed
 
+from mitim_modules.maestro.utils.EPEDbeat import eped_beat
 from mitim_modules.maestro.utils.TRANSPbeat import transp_beat
 from mitim_modules.maestro.utils.PORTALSbeat import portals_beat
-from mitim_modules.maestro.utils.EPEDbeat import eped_beat
+from mitim_modules.maestro.utils.LENGYELbeat import lengyel_beat
 from mitim_modules.maestro.utils.MAESTRObeat import creator_from_eped, creator_from_parameterization, creator
 from mitim_modules.maestro.utils.MAESTRObeat import beat as beat_generic
 
@@ -110,6 +111,9 @@ def define_beat(self, beat, initializer = None, cold_start = False):
         elif beat == 'eped':
             print(f'\n- Beat {self.counter_current}: EPED ******************************* {timeBeginning.strftime("%Y-%m-%d %H:%M:%S")}')
             self.beats[self.counter_current] = eped_beat(self)
+        elif beat == 'lengyel':
+            print(f'\n- Beat {self.counter_current}: LENGYEL ******************************* {timeBeginning.strftime("%Y-%m-%d %H:%M:%S")}')
+            self.beats[self.counter_current] = lengyel_beat(self)
 
         # Access current beat easily
         self.beat = self.beats[self.counter_current]
@@ -315,10 +319,14 @@ def _plot_beats(self, fn, num_beats = 2, only_beats = None, full_plot = True):
             if only_beats is None or only_beats == beat.name:
 
                 print(f'\t- Plotting beat #{counter}...')
-                log_file = self.folder_logs / f'plot_{counter}.log' if (not self.terminal_outputs) else None
-                with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file):
-                    msg = beat.plot(fn = self.fn, counter = i, full_plot = full_plot)
-                print(msg)
+                try:
+                    log_file = self.folder_logs / f'plot_{counter}.log' if (not self.terminal_outputs) else None
+                    with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file):
+                        msg = beat.plot(fn = self.fn, counter = i, full_plot = full_plot)
+                    print(msg)
+                except FileNotFoundError:
+                    print(f'\t\t- Could not plot beat #{counter} because some files are missing', typeMsg = 'w')
+
 
     def _plot_results(self, fn):
 

src/mitim_modules/maestro/scripts/plot_maestro.py

@@ -3,6 +3,8 @@
 from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools
 from mitim_tools.opt_tools import STRATEGYtools
 from mitim_tools.misc_tools.utils import remote_tools
+from mitim_tools.plasmastate_tools.utils import state_plotting
+from mitim_tools.misc_tools.LOGtools import printMsg as print
 from pathlib import Path
 from IPython import embed
 
@@ -115,9 +117,11 @@ def main():
 
         colors = GRAPHICStools.listColors()
 
+
     ms = []
     x, scripts = [], []
     x0, scripts0 = [], []
+    ps_final = []
     for i,folder in enumerate(folders):
         m, ps, ps_lab = MAESTROplot.plotMAESTRO(folder, fn = fn, num_beats=beats, only_beats = only, full_plot = full)
         ms.append(m)
@@ -132,6 +136,16 @@ def main():
         if len(x0) > len(x):
             x = x0
             scripts = scripts0
+
+        ps_final.append(m.final_state)
+
+    # Plot all profiles
+    if any(ps_final):
+        figsProfs = state_plotting.add_figures(fn,fnlab_pre = "MAESTRO Profiles ALL - ", tab_color=4)
+        state_plotting.plotAll(ps_final, figs=figsProfs)
+    else:
+        print("No final profiles to plot, no simulation finished", typeMsg="w")
+
 
     if len(folders) > 1:
         for let in ['A','B']:

src/mitim_modules/maestro/scripts/run_maestro.py

@@ -7,6 +7,7 @@
 from mitim_modules.maestro.utils import TRANSPbeat, PORTALSbeat
 from mitim_tools.misc_tools.IOtools import mitim_timer
 from mitim_tools.misc_tools import PLASMAtools
+from mitim_tools.misc_tools.LOGtools import printMsg as print
 from IPython import embed
 
 def profiles_postprocessing_fun(file_profs, lumpImpurities = True, enforce_same_density_gradients = True):
@@ -110,7 +111,7 @@ def parse_maestro_nml(file_path):
 
     beat_namelists = {}
 
-    for beat_type in ["eped","eped_initializer", "transp", "transp_soft", "portals", "portals_soft"]:
+    for beat_type in ["eped","eped_initializer", "transp", "transp_soft", "portals", "portals_soft", "lengyel"]:
 
         if f"{beat_type}_beat" in maestro_namelist["maestro"]:
 
@@ -139,8 +140,6 @@ def parse_maestro_nml(file_path):
 
                 beat_namelist = maestro_namelist["maestro"][f"{beat_type}_beat"][f"{beat_type}_namelist"]
 
-
-
             # ***************************************************************************
             # Nothin yet
             # ***************************************************************************
@@ -165,7 +164,8 @@ def parse_maestro_nml(file_path):
                 beat_namelist = maestro_namelist["maestro"]["eped_beat"]["eped_namelist"]
 
         else:
-            raise ValueError(f"[MITIM] {beat_type} beat not found in the MAESTRO namelist")
+            beat_namelist = None
+            print(f"[MITIM] {beat_type} beat not found in the MAESTRO namelist", typeMsg='w')
 
         beat_namelists[beat_type] = beat_namelist
 
@@ -220,6 +220,8 @@ def run_maestro_local(
             label_beat = "eped"
         elif maestro_beats["beats"][0] in ["portals", "portals_soft"]:
             label_beat = "portals"
+        elif maestro_beats["beats"][0] in ["lengyel"]:
+            label_beat = "lengyel"
 
         m.define_beat(label_beat, initializer=None if creator_added else parameters_initialize["initializer"])
 

src/mitim_modules/maestro/utils/EPEDbeat.py

@@ -30,12 +30,15 @@ def prepare(
             BetaN = None,           # Force this BetaN (e.g. at creator stage), otherwise from the profiles_current
             Tesep_keV = None,       # Force this Te at the separatrix, otherwise from the profiles_current
             nesep_20 = None,        # Force this ne at the separatrix, otherwise from the profiles_current
-            corrections_set = {},   # Force these inputs to the NN (e.g. exact delta, Rmajor, etc)
+            corrections_set = None,   # Force these inputs to the NN (e.g. exact delta, Rmajor, etc)
             ptop_multiplier = 1.0,  # Multiplier for the ptop, useful for sensitivity studies
             TioverTe = 1.0,        # Ratio of Ti/Te at the top of the pedestal
             **kwargs
             ):
 
+        if corrections_set is None:
+            corrections_set = {}
+
         if nn_location is not None:
 
             print(f'\t- Choice of EPED: NN from {IOtools.clipstr(nn_location)}', typeMsg='i')
@@ -348,6 +351,7 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F
                 scan_results[key]['wtop_psipol'] = np.array(scan_results[key]['wtop_psipol'])
                 scan_results[key]['value'] = np.array(scan_results[key]['value'])
 
+                self.nn.force_within_range = None # Do not throw warnings during the scan
                 scan_results[key]['ptop_kPa_nominal'], scan_results[key]['wtop_psipol_nominal'] = self.nn(*inputs_to_eped)
 
         # ---------------------------------
@@ -472,8 +476,8 @@ def plot(self,  fn = None, counter = 0, full_plot = True):
         fig = fn.add_figure(label='EPED', tab_color=counter)
         axs = fig.subplot_mosaic(
             """
-            ABCDH
-            AEFGI
+            ABCDHJ
+            AEFGIK
             """
         )
         axs = [ ax for ax in axs.values() ]