diff --git a/config/config_user_example.json b/config/config_user_example.json index 7e5fbb61..9fad0060 100644 --- a/config/config_user_example.json +++ b/config/config_user_example.json @@ -20,10 +20,26 @@ "scratch": "/Users/exampleusername/scratch/", "modules": "" }, + "perlmutter": { + "machine": "perlmutter.nersc.gov", + "username": "exampleusername", + "scratch": "/pscratch/sd/p/exampleusername/", + "identity": "$HOME/.ssh/nersc", + "slurm": { + "account": "m3195_g", + "partition": "debug", + "constraint": "gpu", + "mem": "4GB", + "email": "optional@email" + } + }, "engaging": { "machine": "eofe7.mit.edu", "username": "exampleusername", - "partition": "sched_mit_psfc", + "slurm": { + "partition": "sched_mit_psfc", + "exclude": "node584" + }, "identity": "~/.ssh/id_rsa", "scratch": "/nobackup1/exampleusername/", "modules": "" @@ -39,7 +55,6 @@ "machine": "mfews02.psfc.mit.edu", "username": "exampleusername", "tunnel": "mferws01.psfc.mit.edu", - "scratch_tunnel": "/home/exampleusername/scratch/", "port": 9224, "scratch": "/home/exampleusername/scratch/", "modules": "" @@ -48,7 +63,6 @@ "machine": "irisa.gat.com", "username": "exampleusername", "tunnel": "cybele.gat.com", - "scratch_tunnel": "/home/exampleusername/scratch/", "port": 2039, "scratch": "/cluster-scratch/exampleusername/", "modules": "" diff --git a/config/machines_sources/aliases.bashrc b/config/machines_sources/aliases.bashrc index aedca619..9fafc54d 100644 --- a/config/machines_sources/aliases.bashrc +++ b/config/machines_sources/aliases.bashrc @@ -3,23 +3,34 @@ # --------------------------------------------------------------------------------------------------------------------- # mitim_tools interfaces: read, run, plot -alias mitim_plot_gacode="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gacode_tools/exe/read_gacodes.py --files $1" -alias mitim_plot_tgyro="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gacode_tools/exe/read_tgyros.py --folders $1" -alias mitim_plot_tglf="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gacode_tools/exe/read_tglf.py --folder $1" # [--suffix _0.55] [--gacode input.gacode] -alias mitim_plot_eq="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gs_tools/exe/read_eqs.py --files $1" -alias mitim_read_transp="ipython3 -i -- $MITIM_PATH/src/mitim_tools/transp_tools/exe/read_transp.py --files $1" -alias mitim_run_tglf="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gacode_tools/exe/run_tglf.py --folder $1 --tglf $2" #[--gacode input.gacode] [--scan RLTS_2] [--drives True] + +alias mitim_plot_gacode="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gacode_tools/exe/read_gacode.py $1" +alias mitim_plot_tgyro="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gacode_tools/exe/read_tgyro.py $1" +alias mitim_plot_tglf="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gacode_tools/exe/read_tglf.py $1" # [--suffix _0.55] [--gacode input.gacode] +alias mitim_plot_cgyro="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gacode_tools/exe/read_cgyro.py $1" +alias mitim_plot_eq="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gs_tools/exe/read_eq.py $1" +alias mitim_plot_transp="ipython3 -i -- $MITIM_PATH/src/mitim_tools/transp_tools/exe/read_transp.py $1" + +alias mitim_run_tglf="ipython3 -i -- $MITIM_PATH/src/mitim_tools/gacode_tools/exe/run_tglf.py $1 $2" # (folder input.tglf) [--gacode input.gacode] [--scan RLTS_2] [--drives True] # Optimizations -alias mitim_plot_opt="ipython3 -i -- $MITIM_PATH/src/mitim_tools/opt_tools/exe/read.py --type 4 --resolution 20 --folders $1" -alias mitim_plot_portalsSR="ipython3 -i -- $MITIM_PATH/src/mitim_modules/portals/exe/check_initialization.py $1" -alias mitim_plot_portals="mitim_plot_opt $1" +alias mitim_plot_opt="ipython3 -i -- $MITIM_PATH/src/mitim_tools/opt_tools/exe/read.py --type 4 --resolution 20 $1" +alias mitim_plot_portals="ipython3 -i -- $MITIM_PATH/src/mitim_modules/portals/exe/readMetrics.py $1" +alias mitim_slurm="python3 $MITIM_PATH/src/mitim_tools/opt_tools/exe/slurm.py $1" # TRANSP -alias mitim_trcheck="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_check.py $1" # mitim_trcheck pablorf -alias mitim_trclean="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_clean.py $1" # mitim_trclean 88664P CMOD --numbers 1,2,3 -alias mitim_trlook="ipython3 -i -- $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_look.py $1" # mitim_trlook 152895P01 CMOD --nofull --plot --remove +alias mitim_trcheck="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_check.py $1" # mitim_trcheck pablorf +alias mitim_trcheck_p="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_check_periodic.py $1" # mitim_trcheck_p pablorf +alias mitim_trclean="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_clean.py $1" # mitim_trclean 88664P CMOD --numbers 1,2,3 +alias mitim_trlook="ipython3 -i -- $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_look.py $1" # mitim_trlook 152895P01 CMOD --nofull --plot --remove + +# To run TRANSP (in folder with required files): transp 88664 P01 CMOD --version tshare --trmpi 32 --toricmpi 32 --ptrmpi 32 +alias transp="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_transp.py" + +# IM Aliases +alias runim="python3 $MITIM_PATH/src/mitim_tools/im_tools/exe/run_im.py ./" # To run complete IM evaluation: runim 7 [DebugOption: --debug 0] +alias runmitim="python3 $MITIM_PATH/mitim_opt/scenarios_tools/routines/runMITIM_BO.py" # To peform scenario optimization # Others alias compare_nml="python3 $MITIM_PATH/src/mitim_tools/transp_tools/routines/compareNML.py $1" -alias eff_job="python3 $MITIM_PATH/src/mitim_tools/misc_tools/PARALLELtools.py $1" # Give sbatch.out or slurm_output.dat +alias eff_job="python3 $MITIM_PATH/src/mitim_tools/misc_tools/PARALLELtools.py $1" # Give mitim.out or slurm_output.dat diff --git a/config/machines_sources/eofe.bashrc b/config/machines_sources/eofe.bashrc index 0143389c..3687c6c2 100644 --- a/config/machines_sources/eofe.bashrc +++ b/config/machines_sources/eofe.bashrc @@ -1,3 +1,7 @@ +# This is a setup file for the EOFE cluster. Remember (see Installation guide) that this is not +# stricly necessary. However, it is convenient to have it here, so that you can load the modules +# with a single command. If you do not want to use this file, you can load the modules manually +# if you have these lines in your .bashrc (or equivalent) module use /home/software/psfc/modulefiles/ @@ -64,7 +68,7 @@ then # -------------------------------------------------------------------------------- export GACODE_PLATFORM=PSFC_EOFE_RPP - export GACODE_ROOT=/home/$USER/gacode_sparc + export GACODE_ROOT=/home/$USER/gacode_rpp # -------------------------------------------------------------------------------- # ASTRA diff --git a/config/machines_sources/iris.bashrc b/config/machines_sources/iris.bashrc index 3109789b..12a7e7c9 100644 --- a/config/machines_sources/iris.bashrc +++ b/config/machines_sources/iris.bashrc @@ -1,17 +1,19 @@ -# Stuff only for IRIS machine +# This is a setup file for the IRIS cluster. Remember (see Installation guide) that this is not +# stricly necessary. However, it is convenient to have it here, so that you can load the modules +# with a single command. If you do not want to use this file, you can load the modules manually +# if you have these lines in your .bashrc (or equivalent) # ------------------------------------------------------------------------------------------------ # GACODE # ------------------------------------------------------------------------------------------------ -# module load atom/pygacode - ! [ -z "$PS1" ] && echo " * Proceeding to load GACODE modules" export GACODE_PLATFORM=IRIS export GACODE_ROOT=/home/$USER/gacode - . ${GACODE_ROOT}/shared/bin/gacode_setup +module load ntcc + #. ${GACODE_ROOT}/platform/env/env.${GACODE_PLATFORM} # IMPORTANT to not do this, otherwise it will then point to atom if [ $? -eq 124 ] @@ -20,10 +22,3 @@ then else ! [ -z "$PS1" ] && echo " * GACODE modules loaded" fi - -# ------------------------------------------------------------------------------------------------ -# TRANSP/NTCC -# ------------------------------------------------------------------------------------------------ - -module load omfit/unstable -module load ntcc diff --git a/config/machines_sources/nersc.bashrc b/config/machines_sources/nersc.bashrc new file mode 100644 index 00000000..b94e8ad1 --- /dev/null +++ b/config/machines_sources/nersc.bashrc @@ -0,0 +1,20 @@ +# This is a setup file for NERSC Perlmutter. Remember (see Installation guide) that this is not +# stricly necessary. However, it is convenient to have it here, so that you can load the modules +# with a single command. If you do not want to use this file, you can load the modules manually +# if you have these lines in your .bashrc (or equivalent) + +if [ "$NERSC_HOST" = perlmutter ] +then + ! [ -z "$PS1" ] && echo " * PERLMUTTER" + + # ------------------------------------------------------------------------------------------------ + # GACODE + # ------------------------------------------------------------------------------------------------ + + export GACODE_PLATFORM=PERLMUTTER_GPU + export GACODE_ROOT=$HOME/gacode + . $GACODE_ROOT/shared/bin/gacode_setup + . ${GACODE_ROOT}/platform/env/env.$GACODE_PLATFORM +fi + +source $MITIM_PATH/config/machines_sources/slurm_aliases.bashrc \ No newline at end of file diff --git a/config/mitim.bashrc b/config/mitim.bashrc index cb9bc4bf..9a889bb6 100644 --- a/config/mitim.bashrc +++ b/config/mitim.bashrc @@ -40,13 +40,15 @@ then elif [ "${HOSTNAME:0:3}" = toki ] then source $MITIM_PATH/config/machines_sources/toki.bashrc +# NERSC +elif [ "$NERSC_HOST" = perlmutter ] +then + source $MITIM_PATH/config/machines_sources/nersc.bashrc # None of the above else ! [ -z "$PS1" ] && echo " - No specific environment file loaded" fi - - # ------------------------------------------------------------------------------------------------------- # Aliases for quick tools (plotting, reading) # ------------------------------------------------------------------------------------------------------- @@ -55,44 +57,4 @@ fi source $MITIM_PATH/config/machines_sources/aliases.bashrc -# ------------------------------------------------------------------------------------------------------- -# TRANSP Aliases -# ------------------------------------------------------------------------------------------------------- - -# To run TRANSP (in folder with required files): transp 88664 P01 CMOD --version tshare --trmpi 32 --toricmpi 32 --ptrmpi 32 -alias transp="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_transp.py" - -# To send and get LOOK or do a FETCH: look 88664P01 CMOD --nofull --plot --remove -alias look="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_look.py" - -# To check the grid: trcheck pablorf -alias trcheck="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_check.py" - -# To check the grid periodically: trcheckp pablorf -alias trcheckp="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_check_periodic.py" - -# To delete runs from TRANSP grid: cleanruns 88664P CMOD --numbers 1 --to 5 OR 88664P CMOD --numbers 1,2,3 -alias cleanruns="python3 $MITIM_PATH/src/mitim_tools/transp_tools/exe/run_clean.py" - -# Plot transp: trplot transp.cdf -alias trplot="ipython3 -i $MITIM_PATH/src/mitim_tools/transp_tools/outputs/exploreCDFs.py" - -# ------------------------------------------------------------------------------------------------------- -# IM Aliases -# ------------------------------------------------------------------------------------------------------- - -# To run complete IM evaluation: runim 7 [DebugOption: --debug 0] -alias runim="python3 $MITIM_PATH/src/mitim_tools/im_tools/exe/run_im.py ./" - -# To peform scenario optimization -alias runmitim="python3 $MITIM_PATH/mitim_opt/scenarios_tools/routines/runMITIM_BO.py" - -# ------------------------------------------------------------------------------------------------------- -# Other Useful Aliases -# ------------------------------------------------------------------------------------------------------- - -# e.g.: grepy hola -function grepy { clear; grep -rno "$1" . --color --include \*.py; } -function grepf90 { clear; grep -rno "$1" . --color --include \*.{f90,F90}; } - -! [ -z "$PS1" ] && echo -e "\033[32m>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\033[0m" +! [ -z "$PS1" ] && echo -e "\033[32m>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\033[0m" \ No newline at end of file diff --git a/docs/capabilities/portals_capabilities.rst b/docs/capabilities/portals_capabilities.rst index e02cfcb8..8d5acaf7 100644 --- a/docs/capabilities/portals_capabilities.rst +++ b/docs/capabilities/portals_capabilities.rst @@ -1,7 +1,7 @@ PORTALS ======= -The PORTALS method, described in `P. Rodriguez-Fernandez et al., Nulc. Fusion (2022) `_ consists of using Bayesian Optimization techniques to find steady-state solutions of transport codes of arbitrary fidelity. +The PORTALS method, described in `P. Rodriguez-Fernandez et al.,arXiv (2023) `_ consists of using Bayesian Optimization techniques to find steady-state solutions of transport codes of arbitrary fidelity. Once setup has been successful, the following regression test should run smoothly: diff --git a/docs/capabilities/tglf_capabilities.rst b/docs/capabilities/tglf_capabilities.rst index 1998361d..985fc90d 100644 --- a/docs/capabilities/tglf_capabilities.rst +++ b/docs/capabilities/tglf_capabilities.rst @@ -76,7 +76,8 @@ To generate the input files (input.tglf) to TGLF at each radial location, MITIM Now, we are ready to run TGLF. Once the ``prep()`` command has finished, one can run TGLF with different settings and assumptions. That is why, at this point, a sub-folder name for this specific run can be provided. Similarly to the ``prep()`` command, a ``restart`` flag can be provided. The set of control inputs to TGLF (like saturation rule, electromagnetic effects, etc.) are provided in two ways. -First, the argument ``TGLFsettings`` (which goes from 1 to 5 as of now) indicates the base case to start with. The user is referred to ``GACODEdefaults.py`` to understand the meaning of each setting. +First, the argument ``TGLFsettings`` indicates the base case to start with. +The user is referred to ``templates/input.tglf.models.json`` to understand the meaning of each setting, and ``templates/input.tglf.controls`` for the default setup. Second, the argument ``extraOptions`` can be passed as a dictionary of variables to change. For example, the following two commands will run TGLF with saturation rule number 2 with and without electromagnetic effets. After each ``run()`` command, a ``read()`` is needed, to populate the *tglf.results* dictionary with the TGLF outputs (``label`` refers to the dictionary key for each run): @@ -96,19 +97,6 @@ For example, the following two commands will run TGLF with saturation rule numbe tglf.read( label = 'no_em' ) -.. note:: - - One can change every TGLF input with the ``extraOptions = {}`` dictionary, as shown earlier. However, ``GACODEdefaults.py`` contains a list of presets for TGLF that can be selected by simply passing the argument ``TGLFsettings`` to the ``.run()`` method. Available preset are: - - - TGLFsettings = 0: Minimal working example - - TGLFsettings = 1: "Old" ES SAT1 - - TGLFsettings = 2: ES SAT0 - - TGLFsettings = 3: ES SAT1 (a.k.a. SAT1geo) - - TGLFsettings = 4: ES SAT2 - - TGLFsettings = 5: EM SAT2 - - The user is not limited to use those combinations. One can start with a given ``TGLFsettings`` option, and then modify as many parameters as needed with the ``extraOptions`` dictionary. - .. tip:: In this example, ``tglf.results['yes_em']`` and ``tglf.results['no_em']`` are themselves dictionaries, so please do ``.keys()`` to get all the possible results that have been obtained. @@ -117,7 +105,7 @@ TGLF results can be plotted together by indicating what labels to plot: .. code-block:: python - tglf.plotRun( labels = ['yes_em', 'no_em'] ) + tglf.plot( labels = ['yes_em', 'no_em'] ) As a result, a TGLF notebook with different tabs will be opened with all relevant output quantities: @@ -162,7 +150,7 @@ Similarly as in the previous section, you need to run the ``prep()`` command, bu - **PROFILES_GEN** to generate an *input.gacode* file from the *plasmastate.cdf* and *.geq* files. This file is standard within the GACODE suite and contains all plasma information that is required to run core transport codes. -The rest of the workflow is identical to the previous section, including ``.run()``, ``.read()`` and ``.plotRun()``. +The rest of the workflow is identical to the previous section, including ``.run()``, ``.read()`` and ``.plot()``. Run TGLF from input.tglf file @@ -182,7 +170,7 @@ If you have a input.tglf file already, you can still use this script to run it. tglf = TGLFtools.TGLF() tglf.prep_from_tglf( folder, inputtglf_file, input_gacode = inputgacode_file ) -The rest of the workflow is identical, including ``.run()``, ``.read()`` and ``.plotRun()``. +The rest of the workflow is identical, including ``.run()``, ``.read()`` and ``.plot()``. .. tip:: @@ -204,7 +192,7 @@ The rest of the workflow is identical, including ``.run()``, ``.read()`` and ``. tglf = TGLFtools.TGLF() tglf.prep_from_tglf( folder, inputtglf_file ) tglf.read (folder = f'{folder}/', label = 'yes_em' ) - tglf.plotRun( labels = ['yes_em'] ) + tglf.plot( labels = ['yes_em'] ) Please note that the previous code will only work is TGLF was run using MITIM. This is because MITIM stores the results with a suffix that indicates the radial location (``rho``) where the run was performed. @@ -219,7 +207,10 @@ The rest of the workflow is identical, including ``.run()``, ``.read()`` and ``. Run 1D scans of TGLF input parameter ------------------------------------ -*Nothing here yet* +*Under Development* + +*(In the meantime, please checkout* `tutorials/TGLF_tutorial.py `_ *)* + TGLF aliases ------------ diff --git a/docs/capabilities/tgyro_capabilities.rst b/docs/capabilities/tgyro_capabilities.rst index 747c3ae3..168f038b 100644 --- a/docs/capabilities/tgyro_capabilities.rst +++ b/docs/capabilities/tgyro_capabilities.rst @@ -95,7 +95,7 @@ And plot: .. code-block:: python - tgyro.plotRun(labels=['run1']) + tgyro.plot(labels=['run1']) As a result, a TGYRO notebook with different tabs will be opened with all relevant output quantities: diff --git a/docs/capabilities/transp_capabilities.rst b/docs/capabilities/transp_capabilities.rst index 94c2d8ce..471bda1f 100644 --- a/docs/capabilities/transp_capabilities.rst +++ b/docs/capabilities/transp_capabilities.rst @@ -91,7 +91,7 @@ If TRANSP has already been run and the .CDF results file already exists (``cdf_f transp_results = CDFtools.CDFreactor( cdf_file ) - transp_results.plotRun() + transp_results.plot() .. tip:: @@ -114,7 +114,7 @@ If TRANSP has already been run and the .CDF results file already exists (``cdf_f .. code-block:: python - transp_results.plotRun( time = t1 ) + transp_results.plot( time = t1 ) .. note:: diff --git a/docs/faq.rst b/docs/faq.rst index 0005ae50..39f70061 100644 --- a/docs/faq.rst +++ b/docs/faq.rst @@ -5,8 +5,8 @@ Frequently Asked Questions :local: :depth: 1 -Issues during MITIM setup -------------------------- +Issues during MITIM installation +-------------------------------- .. dropdown:: ``pyqt`` fails to install @@ -24,6 +24,14 @@ Issues during MITIM setup pip3 install -e $MITIM_PATH\[pyqt\] +.. dropdown:: ``ERROR: Wheel`` error in pip: + + Make sure you are getting the fresh packages, by using the ``--no-cache`` option: + + .. code-block:: console + + pip3 install -e $MITIM_PATH\[pyqt\] --no-cache + Issues during MITIM tests ------------------------- @@ -39,6 +47,8 @@ Issues during MITIM tests . $GACODE_ROOT/shared/bin/gacode_setup . ${GACODE_ROOT}/platform/env/env.${GACODE_PLATFORM} + \# Add also modules that are required to run MPI instances in your machine + If you still have problems with MITIM execution of TGLF and you have checked that by manually logging-in to the machine you can run TGLF, then it is possible that you have print or echo statements in your ``.bashrc`` or ``.zshrc`` files. Please remove them or add the following: diff --git a/docs/installation.rst b/docs/installation.rst index 9dc4ce06..6f410f6c 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -13,7 +13,7 @@ Installation Instructions ------------ -Clone the `GitHub repository `_ (do not forget to select **Watch All activity** to receive notifications of new releases and announcements, and **Star** the repository to increase its visibility): +Clone the `GitHub repository `_ (do not forget to select the appropriate settings to receive notifications of new releases and announcements, and **Star** the repository to increase its visibility): .. code-block:: console @@ -67,38 +67,93 @@ In ``$MITIM_PATH/config/``, there is a ``config_user_example.json`` with specifi .. code-block:: console cp $MITIM_PATH/config/config_user_example.json $MITIM_PATH/config/config_user.json - vim $MITIM_PATH/config/config_user.json Apart from machine configurations, ``preferences`` in ``config_user.json`` also includes a ``verbose_level`` flag, which indicates the amount of messages that are printed to the terminal when running MITIM. For debugging purposes, it is recommended a maximum verbose level of ``5``. For production runs, a minimum verbose level of ``1`` is recommended so that you only get important messages. -``preferences`` also allows a ``dpi_notebook`` value (in percent from standard), which should be adjusted for each user's screen configuration. +``preferences`` also allows a ``dpi_notebook`` value (in percent from standard), which should be adjusted for each user's screen configuration if the MITIM notebook figures are too small or too large. -For example, if TGLF is set up to run in the *eofe7.mit.edu* machine, this means that, every time in the MITIM workflow when TGLF needs to run, it will access *eofe7.mit.edu* machine to do so, and therefore you must specify how to access the engaging machine: +This is an example of a ``config_user.json`` file that specifies that TGLF should be run in the *eofe7.mit.edu* machine and TGYRO in the *perlmutter.nersc.gov* machine. +The ``slurm`` options are only required if you are running in a computing cluster that uses the SLURM scheduler, and you can specify the partition, account, nodes to exclude and default memory requirements. +In this example, the ``identity`` option is only required if you are running in a computing cluster that requires a specific SSH key to access it. .. code-block:: console - + + { "preferences": { "tglf": "engaging", + "tgyro": "perlmutter", "verbose_level": "5", "dpi_notebook": "80" }, "engaging": { "machine": "eofe7.mit.edu", - "username": "pablorf", - "partition": "sched_mit_psfc", - "identity": "~/.ssh/id_rsa", - "scratch": "/nobackup1/pablorf/scratch/" - } + "username": "YOUR_USERNAME", + "scratch": "/pool001/YOUR_USERNAME/scratch/", + "slurm": { + "partition": "sched_mit_psfc", + "exclude": "node584" + } + }, + "perlmutter": { + "machine": "perlmutter.nersc.gov", + "username": "YOUR_USERNAME", + "scratch": "/pscratch/sd/p/YOUR_USERNAME/scratch/", + "identity": "/Users/YOUR_USERNAME/.ssh/id_rsa_nersc", + "slurm": { + "account": "YOUR_ACCOUNT", + "partition": "YOUR_PARTITION", + "constraint": "gpu", + "mem": "4GB" + } + } + } If you select to run a code in a given machine, please make sure you have ssh rights to that machine with the login instructions specified, unless you are running it locally. -MITIM will attempt to secure-copy and access that machine through a standard SSH connection and it must therefore be set-up prior to launching MITIM jobs. -Make sure that you can ssh with ``ssh username@machine`` with no password for the SSH keys or via a proxy connection (otherwise MITIM will ask for the password very often). +MITIM will attempt to create SSH and SFTP connections to that machine, and will ask for the password if it is not available in the SSH keys or via a proxy connection. .. attention:: Note that MITIM does not maintain or develop the simulation codes that are used within it, such as those from `GACODE `_ or `TRANSP `_. It assumes that proper permissions have been obtained and that working versions of those codes exist in the machine configured to run them. +Please note that MITIM will try to run the codes with standard commands that the shell must understand. +For example, to run the TGLF code, MITIM will want to execute the command ``tglf`` in the *eofe7.mit.edu* machine as specified in the example above. +There are several ways to make sure that the shell understands the command: + +.. dropdown:: 1. Use MITIM automatic machine configuration (limited) + + First, MITIM will automatically try to source ``$MITIM_PATH/config/mitim.bashrc``. If the machine you are running the code + in is listed in ``$MITIM_PATH/config/machine_sources/``, this will also source a machine-specific file. + Specifications to run certain codes (e.g. the GACODE suite) could be available there, but the number of machines that are + automatically available in the public repository is, obviously, limited. + + Please note that this option is only valid if the **MITIM repository and the environment variable $MITIM_PATH are available also in that machine** (not only from the one you are launching the code). + However, if you prefer to use options 2 or 3, the sourcing of a non-existing file will not cause any issues. + +.. dropdown:: 2. Source at shell initialization (recommended) + + Is the commands are available upon login in that machine (e.g. in your personal ``.bashrc`` file), MITIM will be able to run them. + Please note that aliases are usually not available in non-interactive shells, and it is recommended to use full paths and to avoid print (echo) statements. + +.. dropdown:: 3. Send specific commands per code + + Finally, you can populate the ``modules`` option per machine in your ``config_user.json`` file. For example: + + .. code-block:: console + + "engaging": { + ... + "modules": "export GACODE_ROOT=/home/$USER/gacode && . ${GACODE_ROOT}/shared/bin/gacode_setup" + ... + } + + MITIM will execute those commands immediately after ``source $MITIM_PATH/config/mitim.bashrc`` and before running any code. + + Note that you can the same machine listed several times in your ``config_user.json`` file, with different ``modules`` options per code. + You just need to give it a different name per code. + + + License and contributions ------------------------- diff --git a/setup.py b/setup.py index c86c78cd..4da5cc1e 100644 --- a/setup.py +++ b/setup.py @@ -31,6 +31,8 @@ "pyDOE", "multiprocessing_on_dill", "deap", + "paramiko", + "tqdm", "botorch==0.9.4", # Comes w/ gpytorch==1.11, torch>=1.13.1. PRF also tested w/ torch-2.1.1 "scikit-image", # Stricly not for MITIM, but good to have for pygacode ], diff --git a/src/mitim_modules/__init__.py b/src/mitim_modules/__init__.py index 01a6773c..ee494c89 100644 --- a/src/mitim_modules/__init__.py +++ b/src/mitim_modules/__init__.py @@ -1 +1 @@ -from mitim_tools import __version__ \ No newline at end of file +from mitim_tools import __version__ diff --git a/src/mitim_modules/freegsu/FREEGSUmain.py b/src/mitim_modules/freegsu/FREEGSUmain.py index c66181f4..f0a92e3d 100644 --- a/src/mitim_modules/freegsu/FREEGSUmain.py +++ b/src/mitim_modules/freegsu/FREEGSUmain.py @@ -389,8 +389,7 @@ def combined_analysis( # Plot if fn is None: - plt.ioff() - fn = GUItools.FigureNotebook(0, "Combined analysis") + fn = GUItools.FigureNotebook("Combined analysis") fig1 = fn.add_figure(label="FreeGSU - Eq. & Coils") fig2 = fn.add_figure(label="FreeGSU - Metrics") @@ -422,6 +421,8 @@ def combined_analysis( fn=fn, ) + return fn + def apply_rangeVar( setCoils, coilLimits, function_parameters, rangeVar, turns_real=None diff --git a/src/mitim_modules/freegsu/exe/readGS_combine.py b/src/mitim_modules/freegsu/exe/readGS_combine.py index df3ca2c3..37469643 100644 --- a/src/mitim_modules/freegsu/exe/readGS_combine.py +++ b/src/mitim_modules/freegsu/exe/readGS_combine.py @@ -64,7 +64,7 @@ # Combine # ------------------------------------------------------------------------------------------------------------------------------------- -FREEGSUmain.combined_analysis( +fn = FREEGSUmain.combined_analysis( opt_funs, n=n, times=times, diff --git a/src/mitim_modules/freegsu/exe/readGS_pickled.py b/src/mitim_modules/freegsu/exe/readGS_pickled.py index 5d205992..d23a04a4 100644 --- a/src/mitim_modules/freegsu/exe/readGS_pickled.py +++ b/src/mitim_modules/freegsu/exe/readGS_pickled.py @@ -44,6 +44,6 @@ # Plot # --------------------------------------------------- -GSplotting.plotResult( +fn = GSplotting.plotResult( prfs, metrics, m["function_parameters"]["Constraints"], ProblemExtras=params ) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 20a622c5..988dc703 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -1,17 +1,18 @@ -import os, torch, copy +import os +import torch +import copy import numpy as np import dill as pickle_dill -import matplotlib.pyplot as plt from IPython import embed from collections import OrderedDict -from mitim_tools.misc_tools import IOtools, GRAPHICStools +from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.gacode_tools.aux import PORTALSinteraction from mitim_modules.portals import PORTALStools from mitim_modules.portals.aux import ( PORTALSinit, - PORTALSplot, PORTALSoptimization, + PORTALSanalysis, ) from mitim_tools.opt_tools import STRATEGYtools from mitim_tools.opt_tools.aux import BOgraphics @@ -158,60 +159,12 @@ def __init__(self, folder, namelist=None, TensorsType=torch.double): """ Physics-informed parameters to fit surrogates --------------------------------------------- - Note: Dict value indicates what variables need to change at this location to add this one (only one of them is needed) - Note 2: index key indicates when to transition to next (in terms of number of individuals available for fitting) - Things to add: - 'aLte': ['aLte'], 'aLti': ['aLti'], 'aLne': ['aLne'], - 'nuei': ['te','ne'],'tite': ['te','ti'], 'c_s': ['te'], 'w0_n': ['w0'], - 'beta_e': ['te','ne'] - - transition_evaluations is the number of points to be fitted that require a parameter transition. - Note that this ignores ExtraData or ExtraPoints. - - transition_evaluations[0]: max to only consider gradients - - transition_evaluations[1]: no beta_e - - transition_evaluations[2]: full """ - transition_evaluations = [10, 30, 100] - physicsBasedParams = { - transition_evaluations[0]: OrderedDict( - { - "aLte": ["aLte"], - "aLti": ["aLti"], - "aLne": ["aLne"], - "aLw0_n": ["aLw0"], - } - ), - transition_evaluations[1]: OrderedDict( - { - "aLte": ["aLte"], - "aLti": ["aLti"], - "aLne": ["aLne"], - "aLw0_n": ["aLw0"], - "nuei": ["te", "ne"], - "tite": ["te", "ti"], - "w0_n": ["w0"], - } - ), - transition_evaluations[2]: OrderedDict( - { - "aLte": ["aLte"], - "aLti": ["aLti"], - "aLne": ["aLne"], - "aLw0_n": ["aLw0"], - "nuei": ["te", "ne"], - "tite": ["te", "ti"], - "w0_n": ["w0"], - "beta_e": ["te", "ne"], - } - ), - } - - # If doing trace impurities, alnZ only affects that channel, but the rest of turbulent state depends on the rest of parameters - physicsBasedParams_trace = copy.deepcopy(physicsBasedParams) - physicsBasedParams_trace[transition_evaluations[0]]["aLnZ"] = ["aLnZ"] - physicsBasedParams_trace[transition_evaluations[1]]["aLnZ"] = ["aLnZ"] - physicsBasedParams_trace[transition_evaluations[2]]["aLnZ"] = ["aLnZ"] + ( + physicsBasedParams, + physicsBasedParams_trace, + ) = PORTALStools.default_physicsBasedParams() """ Parameters to run PORTALS @@ -370,16 +323,14 @@ def run(self, paramsfile, resultsfile): ) a, b = IOtools.reducePathLevel(self.folder, level=1) - name = f"portals_{b}_ev{numPORTALS}".format( - b, numPORTALS - ) # e.g. portals_jet37_ev0 + name = f"portals_{b}_ev{numPORTALS}" # e.g. portals_jet37_ev0 # Specify the number of PORTALS evaluation. Copy in case of parallel run extra_params_model = copy.deepcopy(self.extra_params) extra_params_model["numPORTALS"] = numPORTALS # Run - _, tgyro, powerstate, dictOFs = runModelEvaluator( + _, tgyro, _, dictOFs = runModelEvaluator( self, FolderEvaluation, numPORTALS, @@ -432,7 +383,7 @@ def scalarized_objective(self, Y): var_dict = {} for of in ofs_ordered_names: - var, pos = of.split("_") + var, _ = of.split("_") if var not in var_dict: var_dict[var] = torch.Tensor().to(Y) var_dict[var] = torch.cat( @@ -454,152 +405,14 @@ def scalarized_objective(self, Y): return of, cal, res def analyze_results(self, plotYN=True, fn=None, restart=False, analysis_level=2): - self_complete = analyze_results( + """ + analysis_level = 2: Standard as other classes. Run best case, plot TGYRO + analysis_level = 3: Read from Execution and also calculate metrics (4: full metrics) + """ + return analyze_results( self, plotYN=plotYN, fn=fn, restart=restart, analysis_level=analysis_level ) - return self_complete - - def plotPORTALS_metrics( - self, - self_parent, - MITIMextra_dict=None, - plotAllFluxes=False, - stds=2, - plotExpected=True, - ): - indecesPlot = [ - self_parent.res.best_absolute_index, - 0, - -1, - ] # [self_parent.res.best_absolute_index,0,None] - - print("- Proceeding to read PORTALS results") - self.portals_plot = PORTALSplot.PORTALSresults( - self.folder, - self_parent.prfs_model, - self_parent.res, - MITIMextra_dict=MITIMextra_dict, - indecesPlot=indecesPlot, - ) - - # It may have changed - indecesPlot = [ - self.portals_plot.numBest, - self.portals_plot.numOrig, - self.portals_plot.numExtra, - ] - - indexToReferMaxValues = ( - self.portals_plot.numBest - ) # self.portals_plot.numOrig #self.portals_plot.numBest - - print(f"- Proceeding to plot PORTALS results ({stds}sigma models)") - self.portals_plot.tgyros[indecesPlot[1]].plot( - fn=self_parent.fn, prelabel=f"({indecesPlot[1]}) TGYRO - " - ) - if indecesPlot[0] < len(self.portals_plot.tgyros): - self.portals_plot.tgyros[indecesPlot[0]].plot( - fn=self_parent.fn, prelabel=f"({indecesPlot[0]}) TGYRO - " - ) - - figs = [ - self_parent.fn.add_figure(label="PROFILES - Profiles"), - self_parent.fn.add_figure(label="PROFILES - Powers"), - self_parent.fn.add_figure(label="PROFILES - Geometry"), - self_parent.fn.add_figure(label="PROFILES - Gradients"), - self_parent.fn.add_figure(label="PROFILES - Flows"), - self_parent.fn.add_figure(label="PROFILES - Other"), - self_parent.fn.add_figure(label="PROFILES - Impurities"), - ] - - if indecesPlot[0] < len(self.portals_plot.profiles): - PROFILEStools.plotAll( - [ - self.portals_plot.profiles[indecesPlot[1]], - self.portals_plot.profiles[indecesPlot[0]], - ], - figs=figs, - extralabs=[f"{indecesPlot[1]}", f"{indecesPlot[0]}"], - ) - - fig = self_parent.fn.add_figure(label="PORTALS Metrics") - PORTALSplot.plotConvergencePORTALS( - self.portals_plot, - fig=fig, - plotAllFluxes=plotAllFluxes, - indexToMaximize=indexToReferMaxValues, - stds=stds, - ) - - # ---------------------------------------------------------------------------------------------------------------- - # Next analysis - # ---------------------------------------------------------------------------------------------------------------- - - if plotExpected: - self.plotPORTALS_expected( - self_parent, - labelsFluxes=self.portals_plot.labelsFluxes, - MITIMextra_dict=MITIMextra_dict, - stds=stds, - ) - - def plotPORTALS_expected( - self, - self_parent, - labelsFluxes={}, - max_plot_points=4, - plotNext=True, - MITIMextra_dict=None, - stds=2, - ): - print( - f"- Proceeding to plot PORTALS expected next variations ({stds}sigma models)" - ) - - trained_points = self_parent.prfs_model.steps[-1].train_X.shape[0] - indexBest = self_parent.res.best_absolute_index - - # Best point - plotPoints = [indexBest] - labelAssigned = [f"#{indexBest} (best)"] - - # Last point - if (trained_points - 1) != indexBest: - plotPoints.append(trained_points - 1) - labelAssigned.append(f"#{trained_points-1} (last)") - - # Last ones - i = 0 - while len(plotPoints) < max_plot_points: - if (trained_points - 2 - i) < 1: - break - if (trained_points - 2 - i) != indexBest: - plotPoints.append(trained_points - 2 - i) - labelAssigned.append(f"#{trained_points-2-i}") - i += 1 - - # First point - if 0 not in plotPoints: - if len(plotPoints) == max_plot_points: - plotPoints[-1] = 0 - labelAssigned[-1] = "#0 (base)" - else: - plotPoints.append(0) - labelAssigned.append("#0 (base)") - - PORTALSplot.plotExpected( - self_parent.prfs_model, - self.folder, - self_parent.fn, - plotPoints=plotPoints, - plotNext=plotNext, - labelAssigned=labelAssigned, - labelsFluxes=labelsFluxes, - MITIMextra_dict=MITIMextra_dict, - stds=stds, - ) - def reuseTrainingTabular( self, folderRead, folderNew, reevaluateTargets=0, restartIfExists=False ): @@ -716,332 +529,6 @@ def reuseTrainingTabular( TabularErrors_Read.updateFile() -def analyze_results( - self, - plotYN=True, - fn=None, - restart=False, - onlyBest=False, - analysis_level=2, - useMITIMextra=True, -): - """ - analysis_level = 2: Standard as other classes. Run best case, plot TGYRO - analysis_level = 3: Read from Execution and also calculate metrics (4: full metrics) - """ - - if plotYN: - print( - "\n ***************************************************************************" - ) - print("* PORTALS plotting module - PORTALS") - print( - "***************************************************************************\n" - ) - - # ---------------------------------------------------------------------------------------------------------------- - # Interpret stuff - # ---------------------------------------------------------------------------------------------------------------- - - ( - variations_original, - variations_best, - self_complete, - ) = self.analyze_optimization_results() - - # ---------------------------------------------------------------------------------------------------------------- - # Running cases: Original and Best - # ---------------------------------------------------------------------------------------------------------------- - - if analysis_level in [2, 5]: - if not onlyBest: - print("\t- Running original case") - FolderEvaluation = f"{self.folder}/Outputs/final_analysis_original/" - if not os.path.exists(FolderEvaluation): - IOtools.askNewFolder(FolderEvaluation, force=True) - - dictDVs = {} - for i in variations_best: - dictDVs[i] = {"value": variations_original[i]} - - # Run - a, b = IOtools.reducePathLevel(self.folder, level=1) - name0 = f"portals_{b}_ev{0}" # e.g. portals_jet37_ev0 - - resultsO, tgyroO, powerstateO, _ = runModelEvaluator( - self_complete, FolderEvaluation, 0, dictDVs, name0, restart=restart - ) - - # Run TGLF - powerstateO.tgyro_current.nameRuns_default = name0 - powerstateO.tgyro_current.runTGLF( - fromlabel=name0, - rhos=self_complete.TGYROparameters["RhoLocations"], - restart=restart, - ) - - print(f"\t- Running best case #{self.res.best_absolute_index}") - FolderEvaluation = f"{self.folder}/Outputs/final_analysis_best/" - if not os.path.exists(FolderEvaluation): - IOtools.askNewFolder(FolderEvaluation, force=True) - - dictDVs = {} - for i in variations_best: - dictDVs[i] = {"value": variations_best[i]} - - # Run - a, b = IOtools.reducePathLevel(self.folder, level=1) - name = f"portals_{b}_ev{self.res.best_absolute_index}" # e.g. portals_jet37_ev0 - resultsB, tgyroB, powerstateB, _ = runModelEvaluator( - self_complete, - FolderEvaluation, - self.res.best_absolute_index, - dictDVs, - name, - restart=restart, - ) - - # Run TGLF - powerstateB.tgyro_current.nameRuns_default = name - powerstateB.tgyro_current.runTGLF( - fromlabel=name, - rhos=self_complete.TGYROparameters["RhoLocations"], - restart=restart, - ) - - # ---------------------------------------------------------------------------------------------------------------- - # Plotting - # ---------------------------------------------------------------------------------------------------------------- - - # Plot - if plotYN: - if not onlyBest: - tgyroO.plotRun(fn=fn, labels=[name0]) - tgyroB.plotRun(fn=fn, labels=[name]) - - # ---------------------------------------------------------------------------------------------------------------- - # Ranges of variation - # ---------------------------------------------------------------------------------------------------------------- - - if plotYN: - figR = fn.add_figure(label="PROFILES Ranges") - pps = np.max( - [3, len(self_complete.TGYROparameters["ProfilesPredicted"])] - ) # Because plotGradients require at least Te, Ti, ne - grid = plt.GridSpec(2, pps, hspace=0.3, wspace=0.3) - axsR = [] - for i in range(pps): - axsR.append(figR.add_subplot(grid[0, i])) - axsR.append(figR.add_subplot(grid[1, i])) - - # Ranges --------------------------------------------------------------------------------------------------------- - produceInfoRanges( - self_complete, - self.prfs_model.bounds_orig, - axsR=axsR, - color="k", - lw=0.2, - alpha=0.05, - label="original", - ) - produceInfoRanges( - self_complete, - self.prfs_model.bounds, - axsR=axsR, - color="c", - lw=0.2, - alpha=0.05, - label="final", - ) - # ---------------------------------------------------------------------------------------------------------------- - - if analysis_level > 2: - if useMITIMextra: - with open(self_complete.MITIMextra, "rb") as handle: - MITIMextra_dict = pickle_dill.load(handle) - else: - MITIMextra_dict = None - - # ---------------------------------------------------------------------------------------------------------------- - # Metrics analysis - # ---------------------------------------------------------------------------------------------------------------- - self_complete.plotPORTALS_metrics( - self, - MITIMextra_dict=MITIMextra_dict, - plotAllFluxes=True, - plotExpected=analysis_level > 3, - ) - - # ---------------------------------------------------------------------------------------------------------------- - # Ranges of variation - # ---------------------------------------------------------------------------------------------------------------- - if useMITIMextra: - plotImpurity = ( - self_complete.PORTALSparameters["ImpurityOfInterest"] - if "nZ" in self_complete.TGYROparameters["ProfilesPredicted"] - else None - ) - plotRotation = ( - "w0" in self_complete.TGYROparameters["ProfilesPredicted"] - ) - - try: - avoidPoints = self.prfs_model.avoidPoints - except: - avoidPoints = [] - - profiles = ( - MITIMextra_dict[0]["tgyro"] - .results[ - [ik for ik in MITIMextra_dict[0]["tgyro"].results.keys()][0] - ] - .profiles - ) - profiles.plotGradients( - axsR, - color="b", - lastRho=self_complete.TGYROparameters["RhoLocations"][-1], - ms=0, - lw=1.0, - label="#0", - ls="-o" if avoidPoints else "--o", - plotImpurity=plotImpurity, - plotRotation=plotRotation, - ) - - for i in range(100): - try: - profiles = ( - MITIMextra_dict[i]["tgyro"] - .results[ - [ - ik - for ik in MITIMextra_dict[i]["tgyro"].results.keys() - ][0] - ] - .profiles - ) - except: - break - profiles.plotGradients( - axsR, - color="r", - lastRho=self_complete.TGYROparameters["RhoLocations"][-1], - ms=0, - lw=0.3, - ls="-o" if avoidPoints else "-.o", - plotImpurity=plotImpurity, - plotRotation=plotRotation, - ) - - profiles.plotGradients( - axsR, - color="g", - lastRho=self_complete.TGYROparameters["RhoLocations"][-1], - ms=0, - lw=1.0, - label=f"#{self.res.best_absolute_index} (best)", - plotImpurity=plotImpurity, - plotRotation=plotRotation, - ) - - axsR[0].legend(loc="best") - - # ---------------------------------------------------------------------------------------------------------------- - # Compare profiles - # ---------------------------------------------------------------------------------------------------------------- - if useMITIMextra: - profile_original = ( - MITIMextra_dict[0]["tgyro"] - .results[ - [ik for ik in MITIMextra_dict[0]["tgyro"].results.keys()][0] - ] - .profiles - ) - profile_best = ( - MITIMextra_dict[self.res.best_absolute_index]["tgyro"] - .results[ - [ - ik - for ik in MITIMextra_dict[self.res.best_absolute_index][ - "tgyro" - ].results.keys() - ][0] - ] - .profiles - ) - profile_original_unCorrected = MITIMextra_dict["profiles_original_un"] - profile_original_0 = MITIMextra_dict["profiles_original"] - - fig4 = fn.add_figure(label="PROFILES Comparison") - grid = plt.GridSpec( - 2, - np.max( - [3, len(self_complete.TGYROparameters["ProfilesPredicted"])] - ), - hspace=0.3, - wspace=0.3, - ) - axs4 = [ - fig4.add_subplot(grid[0, 0]), - fig4.add_subplot(grid[1, 0]), - fig4.add_subplot(grid[0, 1]), - fig4.add_subplot(grid[1, 1]), - fig4.add_subplot(grid[0, 2]), - fig4.add_subplot(grid[1, 2]), - ] - - cont = 1 - if "nZ" in self_complete.TGYROparameters["ProfilesPredicted"]: - axs4.append(fig4.add_subplot(grid[0, 2 + cont])) - axs4.append(fig4.add_subplot(grid[1, 2 + cont])) - cont += 1 - if "w0" in self_complete.TGYROparameters["ProfilesPredicted"]: - axs4.append(fig4.add_subplot(grid[0, 2 + cont])) - axs4.append(fig4.add_subplot(grid[1, 2 + cont])) - - colors = GRAPHICStools.listColors() - - plotImpurity = ( - self_complete.PORTALSparameters["ImpurityOfInterest"] - if "nZ" in self_complete.TGYROparameters["ProfilesPredicted"] - else None - ) - plotRotation = ( - "w0" in self_complete.TGYROparameters["ProfilesPredicted"] - ) - - for i, (profiles, label, alpha) in enumerate( - zip( - [ - profile_original_unCorrected, - profile_original_0, - profile_original, - profile_best, - ], - ["Original", "Corrected", "Initial", "Final"], - [0.2, 1.0, 1.0, 1.0], - ) - ): - profiles.plotGradients( - axs4, - color=colors[i], - label=label, - lastRho=self_complete.TGYROparameters["RhoLocations"][-1], - alpha=alpha, - useRoa=True, - RhoLocationsPlot=self_complete.TGYROparameters["RhoLocations"], - plotImpurity=plotImpurity, - plotRotation=plotRotation, - ) - - axs4[0].legend(loc="best") - - # self.fn.show() - - return self_complete - - def runModelEvaluator( self, FolderEvaluation, @@ -1171,245 +658,33 @@ def runModelEvaluator( return tgyro_current_results, tgyro_current, powerstate, dictOFs -def produceInfoRanges( - self_complete, bounds, axsR, label="", color="k", lw=0.2, alpha=0.05 +def analyze_results( + self, plotYN=True, fn=None, restart=False, analysis_level=2, onlyBest=False ): - rhos = np.append([0], self_complete.TGYROparameters["RhoLocations"]) - aLTe, aLTi, aLne, aLnZ, aLw0 = ( - np.zeros((len(rhos), 2)), - np.zeros((len(rhos), 2)), - np.zeros((len(rhos), 2)), - np.zeros((len(rhos), 2)), - np.zeros((len(rhos), 2)), - ) - for i in range(len(rhos) - 1): - if f"aLte_{i+1}" in bounds: - aLTe[i + 1, :] = bounds[f"aLte_{i+1}"] - if f"aLti_{i+1}" in bounds: - aLTi[i + 1, :] = bounds[f"aLti_{i+1}"] - if f"aLne_{i+1}" in bounds: - aLne[i + 1, :] = bounds[f"aLne_{i+1}"] - if f"aLnZ_{i+1}" in bounds: - aLnZ[i + 1, :] = bounds[f"aLnZ_{i+1}"] - if f"aLw0_{i+1}" in bounds: - aLw0[i + 1, :] = bounds[f"aLw0_{i+1}"] - - X = torch.zeros( - ((len(rhos) - 1) * len(self_complete.TGYROparameters["ProfilesPredicted"]), 2) - ) - l = len(rhos) - 1 - X[0:l, :] = torch.from_numpy(aLTe[1:, :]) - X[l : 2 * l, :] = torch.from_numpy(aLTi[1:, :]) - - cont = 0 - if "ne" in self_complete.TGYROparameters["ProfilesPredicted"]: - X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLne[1:, :]) - cont += 1 - if "nZ" in self_complete.TGYROparameters["ProfilesPredicted"]: - X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLnZ[1:, :]) - cont += 1 - if "w0" in self_complete.TGYROparameters["ProfilesPredicted"]: - X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLw0[1:, :]) - cont += 1 - - X = X.transpose(0, 1) - - powerstate = PORTALStools.constructEvaluationProfiles( - X, self_complete.surrogate_parameters, recalculateTargets=False - ) - - GRAPHICStools.fillGraph( - axsR[0], - powerstate.plasma["rho"][0], - powerstate.plasma["te"][0], - y_up=powerstate.plasma["te"][1], - alpha=alpha, - color=color, - lw=lw, - label=label, - ) - GRAPHICStools.fillGraph( - axsR[1], - rhos, - aLTe[:, 0], - y_up=aLTe[:, 1], - alpha=alpha, - color=color, - label=label, - lw=lw, - ) - - GRAPHICStools.fillGraph( - axsR[2], - powerstate.plasma["rho"][0], - powerstate.plasma["ti"][0], - y_up=powerstate.plasma["ti"][1], - alpha=alpha, - color=color, - label=label, - lw=lw, - ) - GRAPHICStools.fillGraph( - axsR[3], - rhos, - aLTi[:, 0], - y_up=aLTi[:, 1], - alpha=alpha, - color=color, - label=label, - lw=lw, - ) - - cont = 0 - if "ne" in self_complete.TGYROparameters["ProfilesPredicted"]: - GRAPHICStools.fillGraph( - axsR[3 + cont + 1], - powerstate.plasma["rho"][0], - powerstate.plasma["ne"][0] * 0.1, - y_up=powerstate.plasma["ne"][1] * 0.1, - alpha=alpha, - color=color, - label=label, - lw=lw, - ) - GRAPHICStools.fillGraph( - axsR[3 + cont + 2], - rhos, - aLne[:, 0], - y_up=aLne[:, 1], - alpha=alpha, - color=color, - label=label, - lw=lw, - ) - cont += 2 - - if "nZ" in self_complete.TGYROparameters["ProfilesPredicted"]: - GRAPHICStools.fillGraph( - axsR[3 + cont + 1], - powerstate.plasma["rho"][0], - powerstate.plasma["nZ"][0] * 0.1, - y_up=powerstate.plasma["nZ"][1] * 0.1, - alpha=alpha, - color=color, - label=label, - lw=lw, - ) - GRAPHICStools.fillGraph( - axsR[3 + cont + 2], - rhos, - aLnZ[:, 0], - y_up=aLnZ[:, 1], - alpha=alpha, - color=color, - label=label, - lw=lw, - ) - cont += 2 - - if "w0" in self_complete.TGYROparameters["ProfilesPredicted"]: - GRAPHICStools.fillGraph( - axsR[3 + cont + 1], - powerstate.plasma["rho"][0], - powerstate.plasma["w0"][0] * 1e-3, - y_up=powerstate.plasma["w0"][1] * 1e-3, - alpha=alpha, - color=color, - label=label, - lw=lw, - ) - GRAPHICStools.fillGraph( - axsR[3 + cont + 2], - rhos, - aLw0[:, 0], - y_up=aLw0[:, 1], - alpha=alpha, - color=color, - label=label, - lw=lw, - ) - cont += 2 - - -def grabProfilesFromObjects(folder, based_on_last=False, true_original=True): - with open(f"{folder}/Outputs/MITIMextra.pkl", "rb") as f: - auxE = pickle_dill.load(f) - - if not based_on_last: - with open(f"{folder}/Outputs/MITIMstate.pkl", "rb") as f: - aux = pickle_dill.load(f) - i = aux.BOmetrics["overall"]["indBest"] - - else: - for ikey in auxE: - if type(auxE[ikey]) != dict: - break - i = ikey - 1 - - p_new = ( - auxE[i]["tgyro"] - .results[[j for j in auxE[i]["tgyro"].results.keys()][0]] - .profiles_final - ) - - if true_original: - p_orig = auxE["profiles_original"] - else: - p_orig = ( - auxE[0]["tgyro"] - .results[[j for j in auxE[0]["tgyro"].results.keys()][0]] - .profiles_final - ) - - rhos_simulated = auxE[i]["tgyro"].rhosToSimulate - - return p_orig, p_new, rhos_simulated - - -def fix_pickledstate(state_to_mod, powerstate_to_add): - """ - If I have modified the source code of powerstate, it won't be able to load an old PORTALS - What you can do here is to insert an updated class. You can do this by running a bit the code restarting, to get portals_fun - and then apply this: - - E.g.: - fix_pickledstate('sparc_cgyro1/Outputs/MITIMstate.pkl',portals_fun.surrogate_parameters['powerstate']) - """ - - with open(state_to_mod, "rb") as f: - aux = pickle_dill.load(f) - - # Powerstate is stored at the highest level - aux.surrogate_parameters["powerstate"] = powerstate_to_add - aux.mainFunction.powerstate = powerstate_to_add - - for i in range(len(aux.steps)): - # Surrogate parameters are used in stepSettings at each step - aux.steps[i].stepSettings["surrogate_parameters"][ - "powerstate" - ] = powerstate_to_add + if plotYN: + print("\n *****************************************************") + print("* MITIM plotting module - PORTALS") + print("*****************************************************\n") - # Surrogate parameters are passed to each individual surrogate model, at each step, in surrogate_parameters_extended - for j in range(len(aux.steps[i].GP["individual_models"])): - aux.steps[i].GP["individual_models"][j].surrogate_parameters_step[ - "powerstate" - ] = powerstate_to_add + # ---------------------------------------------------------------------------------------------------------------- + # Interpret stuff + # ---------------------------------------------------------------------------------------------------------------- - with open(state_to_mod, "wb") as f: - pickle_dill.dump(aux, f) + portals_full = PORTALSanalysis.PORTALSanalyzer(self) + # ---------------------------------------------------------------------------------------------------------------- + # Plot full information from analysis class + # ---------------------------------------------------------------------------------------------------------------- -def readFullPORTALS(folder): - print( - f">> Proceeding to read as much information from PORTALS run as possible from {folder}" - ) + if plotYN: + portals_full.fn = fn + portals_full.plotPORTALS() - opt_fun = STRATEGYtools.FUNmain(folder) - # opt_fun.read_optimization_results(plotYN=False,analysis_level=4) - self_complete = opt_fun.plot_optimization_results(analysis_level=4, plotYN=False) - portals = opt_fun.prfs_model.mainFunction + # ---------------------------------------------------------------------------------------------------------------- + # Running cases: Original and Best + # ---------------------------------------------------------------------------------------------------------------- - with open(self_complete.MITIMextra, "rb") as handle: - MITIMextra_dict = pickle_dill.load(handle) + if analysis_level in [2, 5]: + portals_full.runCases(onlyBest=onlyBest, restart=restart, fn=fn) - return opt_fun, portals, self_complete, MITIMextra_dict + return portals_full.opt_fun.prfs_model.mainFunction diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 472cb8d0..4ccb8d22 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -1,11 +1,9 @@ -import torch, pickle, copy -from collections import OrderedDict +import torch +import copy import numpy as np -from IPython import embed -from mitim_tools.opt_tools import SURROGATEtools -from mitim_tools.opt_tools.aux import BOgraphics, SAMPLINGtools -from mitim_tools.misc_tools import IOtools, PLASMAtools +from collections import OrderedDict from mitim_tools.misc_tools.IOtools import printMsg as print +from IPython import embed def selectSurrogate(output, surrogateOptions): @@ -103,7 +101,7 @@ def transformmitim(X, surrogate_parameters, output): ) # --- Original model output is in real units, transform to GB here b/c that's how GK codes work - factorGB = GBfromXnorm(X, surrogate_parameters, output, powerstate) + factorGB = GBfromXnorm(X, output, powerstate) # --- Ratio of fluxes (quasilinear) factorRat = ratioFactor(X, surrogate_parameters, output, powerstate) # --- Specific to output @@ -197,7 +195,7 @@ def computeTurbExchangeIndividual(PexchTurb, powerstate): # return mean, upper, lower -def GBfromXnorm(x, surrogate_parameters, output, powerstate): +def GBfromXnorm(x, output, powerstate): # Decide, depending on the output here, which to use as normalization and at what location varFull = output.split("_")[0] pos = int(output.split("_")[1]) @@ -210,9 +208,9 @@ def GBfromXnorm(x, surrogate_parameters, output, powerstate): elif varFull[:2] == "Mt": quantity = "Pgb" elif varFull[:2] == "Ge": - quantity = "Ggb" if (not powerstate.useConvectiveFluxes) else "Qgb" + quantity = "Ggb" if (not powerstate.useConvectiveFluxes) else "Qgb_convection" elif varFull[:2] == "GZ": - quantity = "Ggb" if (not powerstate.useConvectiveFluxes) else "Qgb" + quantity = "Ggb" if (not powerstate.useConvectiveFluxes) else "Qgb_convection" elif varFull[:5] == "Pexch": quantity = "Sgb" @@ -231,7 +229,7 @@ def ImpurityGammaTrick(x, surrogate_parameters, output, powerstate): pos = int(output.split("_")[1]) if ("GZ" in output) and surrogate_parameters["applyImpurityGammaTrick"]: - factor = nZ = powerstate.plasma["ni"][ + factor = powerstate.plasma["ni"][ : x.shape[0], powerstate.indexes_simulation[pos], powerstate.impurityPosition - 1, @@ -252,44 +250,44 @@ def ratioFactor(X, surrogate_parameters, output, powerstate): v = torch.ones(tuple(X.shape[:-1]) + (1,)).to(X) - """ - Apply diffusivities (not real value, just capturing dependencies, - work on normalization, like e_J). Or maybe calculate gradients within powerstate - Remember that for Ti I'm using ne... - """ - if surrogate_parameters["useDiffusivities"]: - pos = int(output.split("_")[-1]) - var = output.split("_")[0] - - if var == "te": - grad = x[:, i] * ( - powerstate.plasma["te"][:, powerstate.indexes_simulation[pos]] - / powerstate.plasma["a"] - ) # keV/m - v[:] = grad * powerstate.plasma["ne"][:, powerstate.indexes_simulation[pos]] - - if var == "ti": - grad = x[:, i] * ( - powerstate.plasma["ti"][:, powerstate.indexes_simulation[pos]] - / powerstate.plasma["a"] - ) # keV/m - v[:] = grad * powerstate.plasma["ne"][:, powerstate.indexes_simulation[pos]] - - # if var == 'ne': - # grad = x[:,i] * ( powerstate.plasma['ne'][:,pos]/powerstate.plasma['a']) # keV/m - # v[:] = grad - - """ - Apply flux ratios - For example [1,Qi,Qi] means I will fit to [Qi, Qe/Qi, Ge/Qi] - """ - - if surrogate_parameters["useFluxRatios"]: - """ - Not ready yet... since my code is not dealing with other outputs at a time so - I don't know Qi if I'm evaluating other fluxes... - """ - pass + # """ + # Apply diffusivities (not real value, just capturing dependencies, + # work on normalization, like e_J). Or maybe calculate gradients within powerstate + # Remember that for Ti I'm using ne... + # """ + # if surrogate_parameters["useDiffusivities"]: + # pos = int(output.split("_")[-1]) + # var = output.split("_")[0] + + # if var == "te": + # grad = x[:, i] * ( + # powerstate.plasma["te"][:, powerstate.indexes_simulation[pos]] + # / powerstate.plasma["a"] + # ) # keV/m + # v[:] = grad * powerstate.plasma["ne"][:, powerstate.indexes_simulation[pos]] + + # if var == "ti": + # grad = x[:, i] * ( + # powerstate.plasma["ti"][:, powerstate.indexes_simulation[pos]] + # / powerstate.plasma["a"] + # ) # keV/m + # v[:] = grad * powerstate.plasma["ne"][:, powerstate.indexes_simulation[pos]] + + # # if var == 'ne': + # # grad = x[:,i] * ( powerstate.plasma['ne'][:,pos]/powerstate.plasma['a']) # keV/m + # # v[:] = grad + + # """ + # Apply flux ratios + # For example [1,Qi,Qi] means I will fit to [Qi, Qe/Qi, Ge/Qi] + # """ + + # if surrogate_parameters["useFluxRatios"]: + # """ + # Not ready yet... since my code is not dealing with other outputs at a time so + # I don't know Qi if I'm evaluating other fluxes... + # """ + # pass return v @@ -342,3 +340,65 @@ def constructEvaluationProfiles(X, surrogate_parameters, recalculateTargets=True powerstate.calculateTargets() return powerstate + + +def default_physicsBasedParams(): + """ + Physics-informed parameters to fit surrogates + --------------------------------------------- + Note: Dict value indicates what variables need to change at this location to add this one (only one of them is needed) + Note 2: index key indicates when to transition to next (in terms of number of individuals available for fitting) + Things to add: + 'aLte': ['aLte'], 'aLti': ['aLti'], 'aLne': ['aLne'], + 'nuei': ['te','ne'],'tite': ['te','ti'], 'c_s': ['te'], 'w0_n': ['w0'], + 'beta_e': ['te','ne'] + + transition_evaluations is the number of points to be fitted that require a parameter transition. + Note that this ignores ExtraData or ExtraPoints. + - transition_evaluations[0]: max to only consider gradients + - transition_evaluations[1]: no beta_e + - transition_evaluations[2]: full + """ + + transition_evaluations = [10, 30, 100] + physicsBasedParams = { + transition_evaluations[0]: OrderedDict( + { + "aLte": ["aLte"], + "aLti": ["aLti"], + "aLne": ["aLne"], + "aLw0_n": ["aLw0"], + } + ), + transition_evaluations[1]: OrderedDict( + { + "aLte": ["aLte"], + "aLti": ["aLti"], + "aLne": ["aLne"], + "aLw0_n": ["aLw0"], + "nuei": ["te", "ne"], + "tite": ["te", "ti"], + "w0_n": ["w0"], + } + ), + transition_evaluations[2]: OrderedDict( + { + "aLte": ["aLte"], + "aLti": ["aLti"], + "aLne": ["aLne"], + "aLw0_n": ["aLw0"], + "nuei": ["te", "ne"], + "tite": ["te", "ti"], + "w0_n": ["w0"], + "beta_e": ["te", "ne"], + } + ), + } + + # If doing trace impurities, alnZ only affects that channel, but the rest of turbulent state depends on the rest of parameters + physicsBasedParams_trace = copy.deepcopy(physicsBasedParams) + physicsBasedParams_trace[transition_evaluations[0]]["aLnZ"] = ["aLnZ"] + physicsBasedParams_trace[transition_evaluations[1]]["aLnZ"] = ["aLnZ"] + physicsBasedParams_trace[transition_evaluations[2]]["aLnZ"] = ["aLnZ"] + + return physicsBasedParams, physicsBasedParams_trace diff --git a/src/mitim_modules/portals/aux/PORTALSanalysis.py b/src/mitim_modules/portals/aux/PORTALSanalysis.py new file mode 100644 index 00000000..73eef970 --- /dev/null +++ b/src/mitim_modules/portals/aux/PORTALSanalysis.py @@ -0,0 +1,987 @@ +import os +import copy +import torch +import numpy as np +import dill as pickle_dill +import matplotlib.pyplot as plt +from mitim_tools.opt_tools import STRATEGYtools +from mitim_tools.misc_tools import IOtools, PLASMAtools, GRAPHICStools +from mitim_tools.gacode_tools import TGLFtools, TGYROtools, PROFILEStools +from mitim_tools.gacode_tools.aux import PORTALSinteraction +from mitim_modules.portals.aux import PORTALSplot +from mitim_modules.powertorch import STATEtools +from mitim_tools.misc_tools.IOtools import printMsg as print + +from IPython import embed + + +class PORTALSanalyzer: + # **************************************************************************** + # INITIALIZATION + # **************************************************************************** + + def __init__(self, opt_fun, folderAnalysis=None): + print("\n************************************") + print("* Initializing PORTALS analyzer...") + print("************************************") + + self.opt_fun = opt_fun + + self.folder = ( + folderAnalysis + if folderAnalysis is not None + else f"{self.opt_fun.folder}/Analysis/" + ) + if not os.path.exists(self.folder): + os.system(f"mkdir {self.folder}") + + self.fn = None + + # Preparation + print("- Grabbing model") + self.step = self.opt_fun.prfs_model.steps[-1] + self.gp = self.step.GP["combined_model"] + + self.powerstate = self.opt_fun.prfs_model.mainFunction.surrogate_parameters[ + "powerstate" + ] + + # Read dictionaries + with open(self.opt_fun.prfs_model.mainFunction.MITIMextra, "rb") as f: + self.mitim_runs = pickle_dill.load(f) + + self.prep_metrics() + + @classmethod + def from_folder(cls, folder, folderRemote=None, folderAnalysis=None): + print(f"\n...Opening PORTALS class from folder {IOtools.clipstr(folder)}") + + opt_fun = STRATEGYtools.FUNmain(folder) + + try: + opt_fun.read_optimization_results( + analysis_level=4, plotYN=False, folderRemote=folderRemote + ) + + return cls(opt_fun, folderAnalysis=folderAnalysis) + + except (FileNotFoundError, AttributeError) as e: + print("- Could not read optimization results due to error:", typeMsg="w") + print(e) + print("- Trying to read PORTALS initialization...", typeMsg="w") + return PORTALSinitializer(folder) + + @classmethod + def merge_instances(cls, instances, folderAnalysis=None, base_index=0): + print( + "\t- Merging PORTALSanalyzer instances by tricking evaluations counter", + typeMsg="w", + ) + + merged_mitim_runs = {} + merged_profiles = [] + merged_tgyros = [] + cont = 0 + for instance in instances: + for key in range(0, instance.ilast + 1): + merged_mitim_runs[cont + key] = instance.mitim_runs[key] + merged_profiles.append(instance.profiles[key]) + merged_tgyros.append(instance.tgyros[key]) + cont += instance.ilast + 1 + + base_instance = instances[base_index] + + merged_instance = cls(base_instance.opt_fun, folderAnalysis) + merged_instance.mitim_runs = merged_mitim_runs + merged_instance.mitim_runs["profiles_original"] = base_instance.mitim_runs[ + "profiles_original" + ] + merged_instance.mitim_runs["profiles_original_un"] = base_instance.mitim_runs[ + "profiles_original_un" + ] + + merged_instance.prep_metrics(ilast=cont - 1) + + return merged_instance + + @classmethod + def merge_from_folders(cls, folders, folderAnalysis=None, base_index=0): + instances = [cls.from_folder(folder) for folder in folders] + return cls.merge_instances( + instances, folderAnalysis=folderAnalysis, base_index=base_index + ) + + # **************************************************************************** + # PREPARATION + # **************************************************************************** + + def prep_metrics(self, calculateRicci={"d0": 2.0, "l": 1.0}, ilast=None): + print("- Interpreting PORTALS results") + + # What's the last iteration? + if ilast is None: + # self.opt_fun.prfs_model.train_Y.shape[0] + for ikey in self.mitim_runs: + if not isinstance(self.mitim_runs[ikey], dict): + break + self.ilast = ikey - 1 + else: + self.ilast = ilast + + # Store indeces + self.ibest = self.opt_fun.res.best_absolute_index + self.i0 = 0 + + if self.ilast == self.ibest: + self.iextra = None + else: + self.iextra = self.ilast + + # Store setup of TGYRO run + self.rhos = self.mitim_runs[0]["tgyro"].results["tglf_neo"].rho[0, 1:] + self.roa = self.mitim_runs[0]["tgyro"].results["tglf_neo"].roa[0, 1:] + + self.PORTALSparameters = self.opt_fun.prfs_model.mainFunction.PORTALSparameters + self.TGYROparameters = self.opt_fun.prfs_model.mainFunction.TGYROparameters + self.TGLFparameters = self.opt_fun.prfs_model.mainFunction.TGLFparameters + + # Useful flags + self.ProfilesPredicted = self.TGYROparameters["ProfilesPredicted"] + + self.runWithImpurity = ( + self.PORTALSparameters["ImpurityOfInterest"] - 1 + if "nZ" in self.ProfilesPredicted + else None + ) + + self.runWithRotation = "w0" in self.ProfilesPredicted + self.includeFast = self.PORTALSparameters["includeFastInQi"] + self.useConvectiveFluxes = self.PORTALSparameters["useConvectiveFluxes"] + self.forceZeroParticleFlux = self.PORTALSparameters["forceZeroParticleFlux"] + + # Profiles and tgyro results + print("\t- Reading profiles and tgyros for each evaluation") + + self.profiles, self.tgyros = [], [] + for i in range(self.ilast + 1): + t = self.mitim_runs[i]["tgyro"].results["use"] + p = t.profiles_final + + self.tgyros.append(t) + self.profiles.append(p) + + if len(self.profiles) <= self.ibest: + print( + "\t- PORTALS was read after new residual was computed but before pickle was written!", + typeMsg="w", + ) + self.ibest -= 1 + self.iextra = None + + self.profiles_next = None + x_train_num = self.step.train_X.shape[0] + file = f"{self.opt_fun.folder}/Execution/Evaluation.{x_train_num}/model_complete/input.gacode" + if os.path.exists(file): + print("\t\t- Reading next profile to evaluate (from folder)") + self.profiles_next = PROFILEStools.PROFILES_GACODE( + file, calculateDerived=False + ) + + file = f"{self.opt_fun.folder}/Execution/Evaluation.{x_train_num}/model_complete/input.gacode.new" + if os.path.exists(file): + self.profiles_next_new = PROFILEStools.PROFILES_GACODE( + file, calculateDerived=False + ) + self.profiles_next_new.printInfo(label="NEXT") + else: + self.profiles_next_new = self.profiles_next + self.profiles_next_new.deriveQuantities() + else: + print("\t\t- Could not read next profile to evaluate (from folder)") + + print("\t- Processing metrics") + + self.evaluations, self.resM = [], [] + self.FusionGain, self.tauE, self.FusionPower = [], [], [] + self.resTe, self.resTi, self.resne, self.resnZ, self.resw0 = [], [], [], [], [] + if calculateRicci is not None: + self.qR_Ricci, self.chiR_Ricci, self.points_Ricci = [], [], [] + else: + self.qR_Ricci, self.chiR_Ricci, self.points_Ricci = None, None, None + + for i, (p, t) in enumerate(zip(self.profiles, self.tgyros)): + print(f"\t\t- Processing evaluation {i}/{len(self.profiles)-1}") + + self.evaluations.append(i) + self.FusionGain.append(p.derived["Q"]) + self.FusionPower.append(p.derived["Pfus"]) + self.tauE.append(p.derived["tauE"]) + + # ------------------------------------------------ + # Residual definitions + # ------------------------------------------------ + + powerstate = self.opt_fun.prfs_model.mainFunction.powerstate + + try: + OriginalFimp = powerstate.TransportOptions["ModelOptions"][ + "OriginalFimp" + ] + except: + OriginalFimp = 1.0 + + impurityPosition = ( + self.runWithImpurity + 1 if self.runWithImpurity is not None else 1 + ) + + portals_variables = t.TGYROmodeledVariables( + useConvectiveFluxes=self.useConvectiveFluxes, + includeFast=self.includeFast, + impurityPosition=impurityPosition, + UseFineGridTargets=self.PORTALSparameters["fineTargetsResolution"], + OriginalFimp=OriginalFimp, + forceZeroParticleFlux=self.PORTALSparameters["forceZeroParticleFlux"], + ) + + if ( + len(powerstate.plasma["volp"].shape) > 1 + and powerstate.plasma["volp"].shape[1] > 1 + ): + powerstate.unrepeat(do_fine=False) + powerstate.repeat(do_fine=False) + + _, _, source, res = PORTALSinteraction.calculatePseudos( + portals_variables["var_dict"], + self.PORTALSparameters, + self.TGYROparameters, + powerstate, + ) + + # Make sense of tensor "source" which are defining the entire predictive set in + Qe_resR = np.zeros(self.rhos.shape[0]) + Qi_resR = np.zeros(self.rhos.shape[0]) + Ge_resR = np.zeros(self.rhos.shape[0]) + GZ_resR = np.zeros(self.rhos.shape[0]) + Mt_resR = np.zeros(self.rhos.shape[0]) + cont = 0 + for prof in self.TGYROparameters["ProfilesPredicted"]: + for ix in range(self.rhos.shape[0]): + if prof == "te": + Qe_resR[ix] = source[0, cont].abs() + if prof == "ti": + Qi_resR[ix] = source[0, cont].abs() + if prof == "ne": + Ge_resR[ix] = source[0, cont].abs() + if prof == "nZ": + GZ_resR[ix] = source[0, cont].abs() + if prof == "w0": + Mt_resR[ix] = source[0, cont].abs() + + cont += 1 + + res = -res.item() + + self.resTe.append(Qe_resR) + self.resTi.append(Qi_resR) + self.resne.append(Ge_resR) + self.resnZ.append(GZ_resR) + self.resw0.append(Mt_resR) + self.resM.append(res) + + # Ricci Metrics + if calculateRicci is not None: + try: + ( + y1, + y2, + y1_std, + y2_std, + ) = PORTALSinteraction.calculatePseudos_distributions( + portals_variables["var_dict"], + self.PORTALSparameters, + self.TGYROparameters, + powerstate, + ) + + QR, chiR = PLASMAtools.RicciMetric( + y1, + y2, + y1_std, + y2_std, + d0=calculateRicci["d0"], + l=calculateRicci["l"], + ) + + self.qR_Ricci.append(QR[0]) + self.chiR_Ricci.append(chiR[0]) + self.points_Ricci.append( + [ + y1.cpu().numpy()[0, :], + y2.cpu().numpy()[0, :], + y1_std.cpu().numpy()[0, :], + y2_std.cpu().numpy()[0, :], + ] + ) + except: + print("\t- Could not calculate Ricci metric", typeMsg="w") + calculateRicci = None + self.qR_Ricci, self.chiR_Ricci, self.points_Ricci = None, None, None + + self.labelsFluxes = portals_variables["labels"] + + self.FusionGain = np.array(self.FusionGain) + self.FusionPower = np.array(self.FusionPower) + self.tauE = np.array(self.tauE) + self.resM = np.array(self.resM) + self.evaluations = np.array(self.evaluations) + self.resTe, self.resTi, self.resne, self.resnZ, self.resw0 = ( + np.array(self.resTe), + np.array(self.resTi), + np.array(self.resne), + np.array(self.resnZ), + np.array(self.resw0), + ) + + if calculateRicci is not None: + self.chiR_Ricci = np.array(self.chiR_Ricci) + self.qR_Ricci = np.array(self.qR_Ricci) + self.points_Ricci = np.array(self.points_Ricci) + + # Normalized L1 norms + self.resTeM = np.abs(self.resTe).mean(axis=1) + self.resTiM = np.abs(self.resTi).mean(axis=1) + self.resneM = np.abs(self.resne).mean(axis=1) + self.resnZM = np.abs(self.resnZ).mean(axis=1) + self.resw0M = np.abs(self.resw0).mean(axis=1) + + self.resCheck = ( + self.resTeM + self.resTiM + self.resneM + self.resnZM + self.resw0M + ) / len(self.TGYROparameters["ProfilesPredicted"]) + + # --------------------------------------------------------------------------------------------------------------------- + # Jacobian + # --------------------------------------------------------------------------------------------------------------------- + + DeltaQ1 = [] + for i in self.TGYROparameters["ProfilesPredicted"]: + if i == "te": + DeltaQ1.append(-self.resTe) + if i == "ti": + DeltaQ1.append(-self.resTi) + if i == "ne": + DeltaQ1.append(-self.resne) + DeltaQ1 = np.array(DeltaQ1) + self.DeltaQ = DeltaQ1[0, :, :] + for i in range(DeltaQ1.shape[0] - 1): + self.DeltaQ = np.append(self.DeltaQ, DeltaQ1[i + 1, :, :], axis=1) + + self.aLTn_perc = None + # try: self.aLTn_perc = calcLinearizedModel(self.opt_fun.prfs_model,self.DeltaQ,numChannels=self.numChannels,numRadius=self.numRadius,sepers=[self.i0, self.ibest]) + # except: print('\t- Jacobian calculation failed',typeMsg='w') + + self.DVdistMetric_x = self.opt_fun.res.DVdistMetric_x + self.DVdistMetric_y = self.opt_fun.res.DVdistMetric_y + + # **************************************************************************** + # PLOTTING + # **************************************************************************** + + def plotPORTALS(self): + if self.fn is None: + from mitim_tools.misc_tools.GUItools import FigureNotebook + + self.fn = FigureNotebook("PORTALS Summary", geometry="1700x1000") + + fig = self.fn.add_figure(label="PROFILES Ranges", tab_color=0) + self.plotRanges(fig=fig) + + self.plotSummary(fn=self.fn, fn_color=1) + + fig = self.fn.add_figure(label="PORTALS Metrics", tab_color=2) + self.plotMetrics(fig=fig) + + fig = self.fn.add_figure(label="PORTALS Expected", tab_color=3) + self.plotExpected(fig=fig) + + fig = self.fn.add_figure(label="PORTALS Simulation", tab_color=4) + self.plotModelComparison(fig=fig) + + def plotMetrics(self, **kwargs): + PORTALSplot.PORTALSanalyzer_plotMetrics(self, **kwargs) + + def plotExpected(self, **kwargs): + PORTALSplot.PORTALSanalyzer_plotExpected(self, **kwargs) + + def plotSummary(self, **kwargs): + PORTALSplot.PORTALSanalyzer_plotSummary(self, **kwargs) + + def plotRanges(self, **kwargs): + PORTALSplot.PORTALSanalyzer_plotRanges(self, **kwargs) + + def plotModelComparison(self, UseThisTGLFfull=None, **kwargs): + UseTGLFfull_x = None + + if UseThisTGLFfull is not None: + """ + UseThisTGLFfull should be a tuple (folder,label) where to read the + results of running the method runTGLFfull() below. + Note that it could be [None, label] + """ + folder, label = UseThisTGLFfull + if folder is None: + folder = f"{self.folder}/tglf_full/" + self.tglf_full = TGLFtools.TGLF(rhos=self.rhos) + for ev in range(self.ilast + 1): + self.tglf_full.read( + folder=f"{folder}/Evaluation.{ev}/tglf_{label}/", label=f"ev{ev}" + ) + UseTGLFfull_x = label + + return PORTALSplot.PORTALSanalyzer_plotModelComparison( + self, UseTGLFfull_x=UseTGLFfull_x, **kwargs + ) + + # **************************************************************************** + # UTILITIES to extract aspects of PORTALS + # **************************************************************************** + + def extractProfiles(self, evaluation=None, correct_targets=True): + if evaluation is None: + evaluation = self.ibest + elif evaluation < 0: + evaluation = self.ilast + + p0 = self.mitim_runs[evaluation]["tgyro"].results["use"].profiles + + p = copy.deepcopy(p0) + + if correct_targets: + print( + "\t- Replacing powers from input.gacode to have the same as input.gacode.new", + typeMsg="i", + ) + + p1 = self.mitim_runs[evaluation]["tgyro"].results["use"].profiles_final + + for ikey in [ + "qei(MW/m^3)", + "qbrem(MW/m^3)", + "qsync(MW/m^3)", + "qline(MW/m^3)", + "qfuse(MW/m^3)", + "qfusi(MW/m^3)", + ]: + p.profiles[ikey] = p1.profiles[ikey] + + return p + + def extractModels(self, step=-1): + if step < 0: + step = len(self.opt_fun.prfs_model.steps) - 1 + + gps = self.opt_fun.prfs_model.steps[step].GP["individual_models"] + + # Make dictionary + models = {} + for gp in gps: + models[gp.output] = simple_model_portals(gp) + + # PRINTING + print( + f""" +**************************************************************************************************** +> MITIM has extracted {len(models)} GP models as a dictionary (only returned variable), to proceed: + 1. Look at the dictionary keys to see which models are available: + models.keys() + 2. Select one model and print its information (e.g. variable labels and order): + m = models['QeTurb_1'] + m.printInfo() + 3. Trained points are stored as m.x, m.y, m.yvar, and you can make predictions with: + x_test = m.x + mean, upper, lower = m(x_test) + 4. Extract samples from the GP with: + x_test = m.x + samples = m(x_test,samples=100) +**************************************************************************************************** +""", + typeMsg="i", + ) + + return models + + def extractPORTALS(self, evaluation=None, folder=None): + if evaluation is None: + evaluation = self.ibest + elif evaluation < 0: + evaluation = self.ilast + + if folder is None: + folder = f"{self.folder}/portals_step{evaluation}/" + + folder = IOtools.expandPath(folder) + if not os.path.exists(folder): + os.system(f"mkdir {folder}") + + # Original class + portals_fun_original = self.opt_fun.prfs_model.mainFunction + + # Start from the profiles of that step + fileGACODE = f"{folder}/input.gacode_transferred" + p = self.extractProfiles(evaluation=evaluation) + p.writeCurrentStatus(file=fileGACODE) + + # New class + from mitim_modules.portals.PORTALSmain import evaluatePORTALS + + portals_fun = evaluatePORTALS(folder) + + # Transfer settings + portals_fun.PORTALSparameters = portals_fun_original.PORTALSparameters + portals_fun.TGYROparameters = portals_fun_original.TGYROparameters + portals_fun.TGLFparameters = portals_fun_original.TGLFparameters + + # PRINTING + print( + f""" +**************************************************************************************************** +> MITIM has extracted PORTALS class to run in {IOtools.clipstr(folder)}, to proceed: + 1. Modify any parameter as required + portals_fun.PORTALSparameters, portals_fun.TGYROparameters, portals_fun.TGLFparameters, portals_fun.Optim + 2. Take the class portals_fun (arg #0) and prepare it with fileGACODE (arg #1) and folder (arg #2) with: + portals_fun.prep(fileGACODE,folder) + 3. Run PORTALS with: + prf_bo = STRATEGYtools.PRF_BO(portals_fun); prf_bo.run() +**************************************************************************************************** +""", + typeMsg="i", + ) + + return portals_fun, fileGACODE, folder + + def extractTGYRO(self, folder=None, restart=False, evaluation=0): + if evaluation is None: + evaluation = self.ibest + elif evaluation < 0: + evaluation = self.ilast + + if folder is None: + folder = f"{self.folder}/tgyro_step{evaluation}/" + + folder = IOtools.expandPath(folder) + if not os.path.exists(folder): + os.system(f"mkdir {folder}") + + print(f"> Extracting and preparing TGYRO in {IOtools.clipstr(folder)}") + + profiles = self.extractProfiles(evaluation=evaluation) + + tgyro = TGYROtools.TGYRO() + tgyro.prep( + folder, profilesclass_custom=profiles, restart=restart, forceIfRestart=True + ) + + TGLFsettings = self.TGLFparameters["TGLFsettings"] + extraOptionsTGLF = self.TGLFparameters["extraOptionsTGLF"] + PredictionSet = [ + int("te" in self.TGYROparameters["ProfilesPredicted"]), + int("ti" in self.TGYROparameters["ProfilesPredicted"]), + int("ne" in self.TGYROparameters["ProfilesPredicted"]), + ] + + return tgyro, self.rhos, PredictionSet, TGLFsettings, extraOptionsTGLF + + def extractTGLF(self, folder=None, positions=None, evaluation=None, restart=False): + if evaluation is None: + evaluation = self.ibest + elif evaluation < 0: + evaluation = self.ilast + + """ + NOTE on radial location extraction: + Two possible options for the rho locations to use: + 1. self.TGYROparameters["RhoLocations"] -> the ones PORTALS sent to TGYRO + 2. self.rhos (came from TGYRO's t.rho[0, 1:]) -> the ones written by the TGYRO run (clipped to 7 decimal places) + Because we want here to run TGLF *exactly* as TGYRO did, we use the first option. + However, this should be fixed in the future, we should never send to TGYRO more than 7 decimal places of any variable # TO FIX + """ + rhos_considered = self.TGYROparameters["RhoLocations"] + + if positions is None: + rhos = rhos_considered + else: + rhos = [] + for i in positions: + rhos.append(rhos_considered[i]) + + if folder is None: + folder = f"{self.folder}/tglf_ev{evaluation}/" + + folder = IOtools.expandPath(folder) + + if not os.path.exists(folder): + os.system(f"mkdir {folder}") + + print( + f"> Extracting and preparing TGLF in {IOtools.clipstr(folder)} from evaluation #{evaluation}" + ) + + inputgacode = f"{folder}/input.gacode.start" + p = self.extractProfiles(evaluation=evaluation) + p.writeCurrentStatus(file=inputgacode) + + tglf = TGLFtools.TGLF(rhos=rhos) + _ = tglf.prep(folder, restart=restart, inputgacode=inputgacode) + + TGLFsettings = self.TGLFparameters["TGLFsettings"] + extraOptions = self.TGLFparameters["extraOptionsTGLF"] + + return tglf, TGLFsettings, extraOptions + + # **************************************************************************** + # UTILITIES for post-analysis + # **************************************************************************** + + def runTGLFfull( + self, + folder=None, + restart=False, + label="default", + tglf_object=None, + **kwargsTGLF, + ): + """ + This runs TGLF for all evaluations, all radii. + This is convenient if I want to re=run TGLF with different settings, e.g. different TGLFsettings, + that you can provide as keyword arguments. + """ + + if folder is None: + folder = f"{self.folder}/tglf_full/" + + if not os.path.exists(folder): + os.system(f"mkdir {folder}") + + for ev in range(self.ilast + 1): + tglf, TGLFsettings, extraOptions = self.extractTGLF( + folder=f"{folder}/Evaluation.{ev}/", evaluation=ev, restart=restart + ) + + kwargsTGLF_this = copy.deepcopy(kwargsTGLF) + + if "TGLFsettings" not in kwargsTGLF_this: + kwargsTGLF_this["TGLFsettings"] = TGLFsettings + if "extraOptions" not in kwargsTGLF_this: + kwargsTGLF_this["extraOptions"] = extraOptions + + tglf.run(subFolderTGLF=f"tglf_{label}/", restart=restart, **kwargsTGLF_this) + + # Read all previously run cases into a single class + if tglf_object is None: + tglf_object = copy.deepcopy(tglf) + + for ev in range(self.ilast + 1): + tglf_object.read( + folder=f"{folder}/Evaluation.{ev}/tglf_{label}/", + label=f"{label}_ev{ev}", + ) + + return tglf_object + + def runCases(self, onlyBest=False, restart=False, fn=None): + from mitim_modules.portals.PORTALSmain import runModelEvaluator + + variations_best = self.opt_fun.res.best_absolute_full["x"] + variations_original = self.opt_fun.res.evaluations[0]["x"] + + if not onlyBest: + print("\t- Running original case") + FolderEvaluation = f"{self.folder}/final_analysis_original/" + if not os.path.exists(FolderEvaluation): + IOtools.askNewFolder(FolderEvaluation, force=True) + + dictDVs = {} + for i in variations_best: + dictDVs[i] = {"value": variations_original[i]} + + # Run + a, b = IOtools.reducePathLevel(self.folder, level=1) + name0 = f"portals_{b}_ev{0}" # e.g. portals_jet37_ev0 + + resultsO, tgyroO, powerstateO, _ = runModelEvaluator( + self.opt_fun.prfs_model.mainFunction, + FolderEvaluation, + 0, + dictDVs, + name0, + restart=restart, + ) + + print(f"\t- Running best case #{self.opt_fun.res.best_absolute_index}") + FolderEvaluation = f"{self.folder}/Outputs/final_analysis_best/" + if not os.path.exists(FolderEvaluation): + IOtools.askNewFolder(FolderEvaluation, force=True) + + dictDVs = {} + for i in variations_best: + dictDVs[i] = {"value": variations_best[i]} + + # Run + a, b = IOtools.reducePathLevel(self.folder, level=1) + name = f"portals_{b}_ev{self.res.best_absolute_index}" # e.g. portals_jet37_ev0 + resultsB, tgyroB, powerstateB, _ = runModelEvaluator( + self.opt_fun.prfs_model.mainFunction, + FolderEvaluation, + self.res.best_absolute_index, + dictDVs, + name, + restart=restart, + ) + + # Plot + if fn is not None: + if not onlyBest: + tgyroO.plot(fn=fn, labels=[name0]) + tgyroB.plot(fn=fn, labels=[name]) + + +# **************************************************************************** +# Helpers +# **************************************************************************** + + +class simple_model_portals: + def __init__(self, gp): + self.gp = gp + + self.x = self.gp.gpmodel.train_X_usedToTrain + self.y = self.gp.gpmodel.train_Y_usedToTrain + self.yvar = self.gp.gpmodel.train_Yvar_usedToTrain + + # self.printInfo() + + def printInfo(self): + print(f"> Model for {self.gp.output} created") + print( + f"\t- Fitted to {len(self.gp.variables)} variables in this order: {self.gp.variables}" + ) + print(f"\t- Trained with {self.x.shape[0]} points") + + def __call__(self, x, samples=None): + numpy_provided = False + if isinstance(x, np.ndarray): + x = torch.Tensor(x) + numpy_provided = True + + mean, upper, lower, samples = self.gp.predict( + x, produceFundamental=True, nSamples=samples + ) + + if samples is None: + if numpy_provided: + return ( + mean[..., 0].detach().cpu().numpy(), + upper[..., 0].detach().cpu().numpy(), + lower[..., 0].detach().cpu().numpy(), + ) + else: + return ( + mean[..., 0].detach(), + upper[..., 0].detach(), + lower[..., 0].detach(), + ) + else: + if numpy_provided: + return samples[..., 0].detach().cpu().numpy() + else: + return samples[..., 0].detach() + + +def calcLinearizedModel( + prfs_model, DeltaQ, posBase=-1, numChannels=3, numRadius=4, sepers=[] +): + """ + posBase = 1 is aLTi, 0 is aLTe, if the order is [a/LTe,aLTi] + -1 is diagonal + -2 is + + NOTE for PRF: THIS ONLY WORKS FOR TURBULENCE, nOT NEO! + """ + + trainx = prfs_model.steps[-1].GP["combined_model"].train_X.cpu().numpy() + + istep, aLTn_est, aLTn_base = 0, [], [] + for i in range(trainx.shape[0]): + if i >= prfs_model.Optim["initialPoints"]: + istep += 1 + + # Jacobian + J = ( + prfs_model.steps[istep] + .GP["combined_model"] + .localBehavior(trainx[i, :], plotYN=False) + ) + + J = 1e-3 * J[: trainx.shape[1], : trainx.shape[1]] # Only turbulence + + print(f"\t- Reading Jacobian for step {istep}") + + Q = DeltaQ[i, :] + + if posBase < 0: + # All channels together ------------------------------------------------ + mult = torch.Tensor() + for i in range(12): + if posBase == -1: + a = torch.Tensor([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # Diagonal + elif posBase == -2: + a = torch.Tensor( + [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0] + ) # Block diagonal + a = torch.roll(a, i) + mult = torch.cat((mult, a.unsqueeze(0)), dim=0) + + J_reduced = J * mult + aLTn = (J_reduced.inverse().cpu().numpy()).dot(Q) + aLTn_base0 = trainx[i, :] + # ------------------------------------------------------------------------ + + else: + # Channel per channel, only ion temperature gradient ------------------------ + J_mod = [] + aLTn_base0 = [] + cont = 0 + for c in range(numChannels): + for r in range(numRadius): + J_mod.append(J[cont, posBase * numRadius + r].cpu().numpy()) + aLTn_base0.append(trainx[i, posBase * numRadius + r]) + cont += 1 + J_mod = np.array(J_mod) + aLTn_base0 = np.array(aLTn_base0) + aLTn = Q / J_mod + # ------------------------------------------------------------------------ + + aLTn_base.append(aLTn_base0) + aLTn_est.append(aLTn) + + aLTn_est = np.array(aLTn_est) + aLTn_base = np.array(aLTn_base) + + aLTn_perc = [ + np.abs(i / j) * 100.0 if i is not None else None + for i, j in zip(aLTn_est, aLTn_base) + ] + + return aLTn_perc + + +class PORTALSinitializer: + def __init__(self, folder): + self.folder = IOtools.expandPath(folder) + + # Read powerstates + self.powerstates = [] + self.profiles = [] + for i in range(100): + try: + prof = PROFILEStools.PROFILES_GACODE( + f"{self.folder}/Outputs/ProfilesEvaluated/input.gacode.{i}" + ) + p = STATEtools.read_saved_state( + f"{self.folder}/Initialization/initialization_simple_relax/portals_{IOtools.reducePathLevel(self.folder)[1]}_ev{i}/powerstate.pkl" + ) + except FileNotFoundError: + break + + self.profiles.append(prof) + p.profiles.deriveQuantities() + self.powerstates.append(p) + + self.fn = None + + def plotMetrics(self, extra_lab="", **kwargs): + if self.fn is None: + from mitim_tools.misc_tools.GUItools import FigureNotebook + + self.fn = FigureNotebook("PowerState", geometry="1800x900") + + figMain = self.fn.add_figure(label=f"{extra_lab} - PowerState") + figG = self.fn.add_figure(label=f"{extra_lab} - Sequence") + + grid = plt.GridSpec(4, 6, hspace=0.3, wspace=0.4) + axs = [ + figMain.add_subplot(grid[0, 1]), + figMain.add_subplot(grid[0, 2]), + figMain.add_subplot(grid[0, 3]), + figMain.add_subplot(grid[0, 4]), + figMain.add_subplot(grid[0, 5]), + figMain.add_subplot(grid[1, 1]), + figMain.add_subplot(grid[1, 2]), + figMain.add_subplot(grid[1, 3]), + figMain.add_subplot(grid[1, 4]), + figMain.add_subplot(grid[1, 5]), + figMain.add_subplot(grid[2, 1]), + figMain.add_subplot(grid[2, 2]), + figMain.add_subplot(grid[2, 3]), + figMain.add_subplot(grid[2, 4]), + figMain.add_subplot(grid[2, 5]), + figMain.add_subplot(grid[3, 1]), + figMain.add_subplot(grid[3, 2]), + figMain.add_subplot(grid[3, 3]), + figMain.add_subplot(grid[3, 4]), + figMain.add_subplot(grid[3, 5]), + ] + + axsRes = figMain.add_subplot(grid[:, 0]) + + colors = GRAPHICStools.listColors() + + # POWERPLOT + + if len(self.powerstates) > 0: + for i in range(len(self.powerstates)): + self.powerstates[i].plot( + axs=axs, axsRes=axsRes, c=colors[i], label=f"#{i}" + ) + + axs[0].legend(prop={"size": 8}) + + axsRes.set_xlim([0, i]) + + # GRADIENTS + if len(self.profiles) > 0: + grid = plt.GridSpec(2, 5, hspace=0.3, wspace=0.3) + axsGrads = [] + for j in range(5): + for i in range(2): + axsGrads.append(figG.add_subplot(grid[i, j])) + for i, p in enumerate(self.profiles): + p.plotGradients( + axsGrads, + color=colors[i], + plotImpurity=3, + plotRotation=True, + lastRho=self.powerstates[0].plasma["rho"][-1, -1].item(), + ) + + axsGrads_extra = [ + axs[0], + axs[5], + axs[1], + axs[6], + axs[2], + axs[7], + axs[3], + axs[8], + axs[4], + axs[9], + ] + for i, p in enumerate(self.profiles): + p.plotGradients( + axsGrads_extra, + color=colors[i], + plotImpurity=3, + plotRotation=True, + lastRho=self.powerstates[0].plasma["rho"][-1, -1].item(), + lw=0.5, + ms=0, + ) diff --git a/src/mitim_modules/portals/aux/PORTALScgyro.py b/src/mitim_modules/portals/aux/PORTALScgyro.py index e656e6ab..cb9dc487 100644 --- a/src/mitim_modules/portals/aux/PORTALScgyro.py +++ b/src/mitim_modules/portals/aux/PORTALScgyro.py @@ -1,4 +1,5 @@ -import os, sys, copy +import os +import copy import numpy as np from IPython import embed from mitim_tools.misc_tools import IOtools, PLASMAtools @@ -16,13 +17,14 @@ The CGYRO file must use particle flux. Convective transformation occurs later """ +trick_hardcoded_f = None + # train_sep is the number of initial runs in it#0 results file. Now, it's usually 1 # start_num is the number of the first iteration, usually 0 # trick_harcoded_f is the name of the file until the iteration number. E.g. '/Users/pablorf/PROJECTS/project_2022_mitimalcc/iter/cgyro3/run3_recover/Outputs/cgyro_results/iter_rmp_75_' -train_sep, start_num, trick_hardcoded_f, last_one = 1, 0, None, 100 - -# includeMtAndGz_hardcoded, train_sep,start_num,last_one,trick_hardcoded_f = False, 1, 0,100, '/Users/pablorf/PROJECTS/project_2023_SPARClmodes/recover2/run13/Outputs/cgyro_results/sparc_l_high_nu_it_' +# e.g.: +# includeMtAndGz_hardcoded, train_sep,start_num,last_one,trick_hardcoded_f = True, 1, 0,100, '/Users/pablorf/PROJECTS/project_2023_PRISAd3d/cgyro_recover3/run1//Outputs/cgyro_results/d3d_5chan_it_' def evaluateCGYRO( @@ -41,7 +43,7 @@ def evaluateCGYRO( OriginalFimp = PORTALSparameters["fImp_orig"] print( - f"\t- Suggested command to send files (not .new) for CGYRO evaluation:", + "\t- Suggested command to send files (not .new) for CGYRO evaluation:", typeMsg="i", ) lapfile = f"{IOtools.expandPath(folder,ensurePathValid=True)}/Outputs/ProfilesEvaluated/input.gacode.{numPORTALS}" diff --git a/src/mitim_modules/portals/aux/PORTALSinit.py b/src/mitim_modules/portals/aux/PORTALSinit.py index 0ecebb8e..1901daad 100644 --- a/src/mitim_modules/portals/aux/PORTALSinit.py +++ b/src/mitim_modules/portals/aux/PORTALSinit.py @@ -1,12 +1,14 @@ -import os, torch, copy +import os +import torch +import copy import numpy as np from collections import OrderedDict -from IPython import embed from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.powertorch import STATEtools from mitim_modules.portals import PORTALStools from mitim_tools.misc_tools.IOtools import printMsg as print +from IPython import embed def initializeProblem( @@ -69,7 +71,7 @@ def initializeProblem( if "RoaLocations" in portals_fun.TGYROparameters: roa = portals_fun.TGYROparameters["RoaLocations"] rho = np.interp(roa, profiles.derived["roa"], profiles.profiles["rho(-)"]) - print(f"\t * r/a provided, transforming to rho:") + print("\t * r/a provided, transforming to rho:") print(f"\t\t r/a = {roa}") print(f"\t\t rho = {rho}") portals_fun.TGYROparameters["RhoLocations"] = rho diff --git a/src/mitim_modules/portals/aux/PORTALSplot.py b/src/mitim_modules/portals/aux/PORTALSplot.py index 7a60bdc1..360e8ca0 100644 --- a/src/mitim_modules/portals/aux/PORTALSplot.py +++ b/src/mitim_modules/portals/aux/PORTALSplot.py @@ -1,404 +1,45 @@ -import copy, torch +import torch import numpy as np import matplotlib.pyplot as plt from IPython import embed -from mitim_tools.misc_tools import GRAPHICStools, PLASMAtools -from mitim_tools.gacode_tools import PROFILEStools, TGYROtools -from mitim_tools.gacode_tools.aux import PORTALSinteraction +from mitim_tools.misc_tools import GRAPHICStools +from mitim_tools.gacode_tools import PROFILEStools +from mitim_modules.portals import PORTALStools + from mitim_tools.misc_tools.IOtools import printMsg as print factor_dw0dr = 1e-5 label_dw0dr = "$-d\\omega_0/dr$ (krad/s/cm)" +# --------------------------------------------------------------------------------------------------------------------- +# Plotting methods for PORTALS class +# --------------------------------------------------------------------------------------------------------------------- -class PORTALSresults: - def __init__( - self, - folderWork, - prfs_model, - ResultsOptimization, - MITIMextra_dict=None, - indecesPlot=[-1, 0, None], - calculateRicci={"d0": 2.0, "l": 1.0}, - ): - self.prfs_model = prfs_model - self.ResultsOptimization = ResultsOptimization - - includeFast = self.prfs_model.mainFunction.PORTALSparameters["includeFastInQi"] - impurityPosition = self.prfs_model.mainFunction.PORTALSparameters[ - "ImpurityOfInterest" - ] - self.useConvectiveFluxes = self.prfs_model.mainFunction.PORTALSparameters[ - "useConvectiveFluxes" - ] - - self.numChannels = len( - self.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"] - ) - self.numRadius = len( - self.prfs_model.mainFunction.TGYROparameters["RhoLocations"] - ) - self.numBest = ( - indecesPlot[0] - if (indecesPlot[0] > 0) - else (self.prfs_model.train_Y.shape[0] + indecesPlot[0]) - ) - self.numOrig = indecesPlot[1] - self.numExtra = ( - indecesPlot[2] - if (indecesPlot[2] is None or indecesPlot[2] > 0) - else (self.prfs_model.train_Y.shape[0] + indecesPlot[2]) - ) - self.sepers = [self.numOrig, self.numBest] - - if (self.numExtra is not None) and (self.numExtra == self.numBest): - self.numExtra = None - - self.posZ = prfs_model.mainFunction.PORTALSparameters["ImpurityOfInterest"] - 1 - - # Profiles and tgyro results - print("\t- Reading profiles and tgyros for each evaluation") - self.profiles, self.tgyros = [], [] - for i in range(self.prfs_model.train_Y.shape[0]): - # print(f'\t- Reading TGYRO results and PROFILES for evaluation {i}/{self.prfs_model.train_Y.shape[0]-1}') - if MITIMextra_dict is not None: - # print('\t\t* Reading from MITIMextra_dict',typeMsg='i') - self.tgyros.append(MITIMextra_dict[i]["tgyro"].results["use"]) - self.profiles.append( - MITIMextra_dict[i]["tgyro"].results["use"].profiles_final - ) - else: - # print('\t\t* Reading from scratch from folders (will surely take longer)',typeMsg='i') - folderEvaluation = folderWork + f"/Execution/Evaluation.{i}/" - self.profiles.append( - PROFILEStools.PROFILES_GACODE( - folderEvaluation + "/model_complete/input.gacode.new" - ) - ) - self.tgyros.append( - TGYROtools.TGYROoutput( - folderEvaluation + "model_complete/", profiles=self.profiles[i] - ) - ) - - if len(self.profiles) <= self.numBest: - print( - "\t- PORTALS was read after new residual was computed but before pickle was written!", - typeMsg="w", - ) - self.numBest -= 1 - self.numExtra = None - - # Create some metrics - - print("\t- Process results") - - self.evaluations, self.resM = [], [] - self.FusionGain, self.tauE, self.FusionPower = [], [], [] - self.resTe, self.resTi, self.resne, self.resnZ, self.resw0 = [], [], [], [], [] - if calculateRicci is not None: - self.QR_Ricci, self.chiR_Ricci, self.points_Ricci = [], [], [] - else: - self.QR_Ricci, self.chiR_Ricci, self.points_Ricci = None, None, None - for i, (p, t) in enumerate(zip(self.profiles, self.tgyros)): - self.evaluations.append(i) - self.FusionGain.append(p.derived["Q"]) - self.FusionPower.append(p.derived["Pfus"]) - self.tauE.append(p.derived["tauE"]) - - # ------------------------------------------------ - # Residual definitions - # ------------------------------------------------ - - powerstate = self.prfs_model.mainFunction.powerstate - - try: - OriginalFimp = powerstate.TransportOptions["ModelOptions"][ - "OriginalFimp" - ] - except: - OriginalFimp = 1.0 - - portals_variables = t.TGYROmodeledVariables( - useConvectiveFluxes=self.useConvectiveFluxes, - includeFast=includeFast, - impurityPosition=impurityPosition, - ProfilesPredicted=self.prfs_model.mainFunction.TGYROparameters[ - "ProfilesPredicted" - ], - UseFineGridTargets=self.prfs_model.mainFunction.PORTALSparameters[ - "fineTargetsResolution" - ], - OriginalFimp=OriginalFimp, - forceZeroParticleFlux=self.prfs_model.mainFunction.PORTALSparameters[ - "forceZeroParticleFlux" - ], - ) - - if ( - len(powerstate.plasma["volp"].shape) > 1 - and powerstate.plasma["volp"].shape[1] > 1 - ): - powerstate.unrepeat(do_fine=False) - powerstate.repeat(do_fine=False) - - _, _, source, res = PORTALSinteraction.calculatePseudos( - portals_variables["var_dict"], - self.prfs_model.mainFunction.PORTALSparameters, - self.prfs_model.mainFunction.TGYROparameters, - powerstate, - ) - - # Make sense of tensor "source" which are defining the entire predictive set in - Qe_resR = np.zeros( - len(self.prfs_model.mainFunction.TGYROparameters["RhoLocations"]) - ) - Qi_resR = np.zeros( - len(self.prfs_model.mainFunction.TGYROparameters["RhoLocations"]) - ) - Ge_resR = np.zeros( - len(self.prfs_model.mainFunction.TGYROparameters["RhoLocations"]) - ) - GZ_resR = np.zeros( - len(self.prfs_model.mainFunction.TGYROparameters["RhoLocations"]) - ) - Mt_resR = np.zeros( - len(self.prfs_model.mainFunction.TGYROparameters["RhoLocations"]) - ) - cont = 0 - for prof in self.prfs_model.mainFunction.TGYROparameters[ - "ProfilesPredicted" - ]: - for ix in range( - len(self.prfs_model.mainFunction.TGYROparameters["RhoLocations"]) - ): - if prof == "te": - Qe_resR[ix] = source[0, cont].abs() - if prof == "ti": - Qi_resR[ix] = source[0, cont].abs() - if prof == "ne": - Ge_resR[ix] = source[0, cont].abs() - if prof == "nZ": - GZ_resR[ix] = source[0, cont].abs() - if prof == "w0": - Mt_resR[ix] = source[0, cont].abs() - - cont += 1 - - res = -res.item() - - self.resTe.append(Qe_resR) - self.resTi.append(Qi_resR) - self.resne.append(Ge_resR) - self.resnZ.append(GZ_resR) - self.resw0.append(Mt_resR) - self.resM.append(res) - - # Ricci Metrics - if calculateRicci is not None: - try: - ( - y1, - y2, - y1_std, - y2_std, - ) = PORTALSinteraction.calculatePseudos_distributions( - portals_variables["var_dict"], - self.prfs_model.mainFunction.PORTALSparameters, - self.prfs_model.mainFunction.TGYROparameters, - powerstate, - ) - - QR, chiR = PLASMAtools.RicciMetric( - y1, - y2, - y1_std, - y2_std, - d0=calculateRicci["d0"], - l=calculateRicci["l"], - ) - self.QR_Ricci.append(QR[0]) - self.chiR_Ricci.append(chiR[0]) - self.points_Ricci.append( - [ - y1.cpu().numpy()[0, :], - y2.cpu().numpy()[0, :], - y1_std.cpu().numpy()[0, :], - y2_std.cpu().numpy()[0, :], - ] - ) - except: - print("\t- Could not calculate Ricci metric", typeMsg="w") - calculateRicci = None - self.QR_Ricci, self.chiR_Ricci, self.points_Ricci = None, None, None - - self.labelsFluxes = portals_variables["labels"] - - self.FusionGain = np.array(self.FusionGain) - self.FusionPower = np.array(self.FusionPower) - self.tauE = np.array(self.tauE) - self.resM = np.array(self.resM) - self.evaluations = np.array(self.evaluations) - self.resTe, self.resTi, self.resne, self.resnZ, self.resw0 = ( - np.array(self.resTe), - np.array(self.resTi), - np.array(self.resne), - np.array(self.resnZ), - np.array(self.resw0), - ) - - if calculateRicci is not None: - self.chiR_Ricci = np.array(self.chiR_Ricci) - self.QR_Ricci = np.array(self.QR_Ricci) - self.points_Ricci = np.array(self.points_Ricci) - - # Normalized L1 norms - self.resTeM = np.abs(self.resTe).mean(axis=1) - self.resTiM = np.abs(self.resTi).mean(axis=1) - self.resneM = np.abs(self.resne).mean(axis=1) - self.resnZM = np.abs(self.resnZ).mean(axis=1) - self.resw0M = np.abs(self.resw0).mean(axis=1) - - self.resCheck = ( - self.resTeM + self.resTiM + self.resneM + self.resnZM + self.resw0M - ) / len(self.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]) - - # --------------------------------------------------------------------------------------------------------------------- - # Jacobian - # --------------------------------------------------------------------------------------------------------------------- - - DeltaQ1 = [] - for i in self.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - if i == "te": - DeltaQ1.append(-self.resTe) - if i == "ti": - DeltaQ1.append(-self.resTi) - if i == "ne": - DeltaQ1.append(-self.resne) - DeltaQ1 = np.array(DeltaQ1) - self.DeltaQ = DeltaQ1[0, :, :] - for i in range(DeltaQ1.shape[0] - 1): - self.DeltaQ = np.append(self.DeltaQ, DeltaQ1[i + 1, :, :], axis=1) - - self.aLTn_perc = aLTi_perc = None - # try: self.aLTn_perc = aLTi_perc = calcLinearizedModel(self.prfs_model,self.DeltaQ,numChannels=self.numChannels,numRadius=self.numRadius,sepers=self.sepers) - # except: print('\t- Jacobian calculation failed',typeMsg='w') - - self.DVdistMetric_x = ResultsOptimization.DVdistMetric_x - self.DVdistMetric_y = ResultsOptimization.DVdistMetric_y - - -def calcLinearizedModel( - prfs_model, DeltaQ, posBase=-1, numChannels=3, numRadius=4, sepers=[] -): - """ - posBase = 1 is aLTi, 0 is aLTe, if the order is [a/LTe,aLTi] - -1 is diagonal - -2 is - - NOTE for PRF: THIS ONLY WORKS FOR TURBULENCE, nOT NEO! - """ - - trainx = prfs_model.steps[-1].GP["combined_model"].train_X.cpu().numpy() - - istep, aLTn_est, aLTn_base = 0, [], [] - for i in range(trainx.shape[0]): - if i >= prfs_model.Optim["initialPoints"]: - istep += 1 - - # Jacobian - J = ( - prfs_model.steps[istep] - .GP["combined_model"] - .localBehavior(trainx[i, :], plotYN=False) - ) - - J = 1e-3 * J[: trainx.shape[1], : trainx.shape[1]] # Only turbulence - - print(f"\t- Reading Jacobian for step {istep}") - - Q = DeltaQ[i, :] - - if posBase < 0: - # All channels together ------------------------------------------------ - mult = torch.Tensor() - for i in range(12): - if posBase == -1: - a = torch.Tensor([1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # Diagonal - elif posBase == -2: - a = torch.Tensor( - [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0] - ) # Block diagonal - a = torch.roll(a, i) - mult = torch.cat((mult, a.unsqueeze(0)), dim=0) - - J_reduced = J * mult - aLTn = (J_reduced.inverse().cpu().numpy()).dot(Q) - aLTn_base0 = trainx[i, :] - # ------------------------------------------------------------------------ - else: - # Channel per channel, only ion temperature gradient ------------------------ - J_mod = [] - aLTn_base0 = [] - cont = 0 - for c in range(numChannels): - for r in range(numRadius): - J_mod.append(J[cont, posBase * numRadius + r].cpu().numpy()) - aLTn_base0.append(trainx[i, posBase * numRadius + r]) - cont += 1 - J_mod = np.array(J_mod) - aLTn_base0 = np.array(aLTn_base0) - aLTn = Q / J_mod - # ------------------------------------------------------------------------ - - aLTn_base.append(aLTn_base0) - aLTn_est.append(aLTn) - - aLTn_est = np.array(aLTn_est) - aLTn_base = np.array(aLTn_base) - - aLTn_perc = [ - np.abs(i / j) * 100.0 if i is not None else None - for i, j in zip(aLTn_est, aLTn_base) - ] - - return aLTn_perc - - -def plotConvergencePORTALS( - portals, - folderWork=None, +def PORTALSanalyzer_plotMetrics( + self, fig=None, - stds=2, - plotAllFluxes=False, indexToMaximize=None, + plotAllFluxes=False, + index_extra=None, + stds=2, plotFlows=True, fontsize_leg=5, includeRicci=True, + useConvectiveFluxes=False, # By default, plot in real particle units + file_save=None, + **kwargs, # To allow pass fn that may be used in another plotMetrics method ): - """ - Either folderwork or portals need to be provided - """ - - labelsFluxes = portals.labelsFluxes # For channel residuals + print("- Plotting PORTALS Metrics") - useConvectiveFluxes = False # portals.useConvectiveFluxes - labelsFluxesF = { - "te": "$Q_e$ ($MW/m^2$)", - "ti": "$Q_i$ ($MW/m^2$)", - "ne": "$\\Gamma_e$ ($10^{20}/s/m^2$)", - "nZ": "$\\Gamma_Z$ ($10^{20}/s/m^2$)", - "w0": "$M_T$ ($J/m^2$)", - } + if index_extra is not None: + self.iextra = index_extra if fig is None: plt.ion() fig = plt.figure(figsize=(15, 8)) - numprofs = len(portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]) - - wspace = 0.3 if numprofs <= 4 else 0.5 + numprofs = len(self.ProfilesPredicted) grid = plt.GridSpec(nrows=8, ncols=numprofs + 1, hspace=0.3, wspace=0.35) @@ -414,7 +55,7 @@ def plotConvergencePORTALS( axTi_f = fig.add_subplot(grid[6:, 1]) cont = 0 - if "ne" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: + if "ne" in self.ProfilesPredicted: axne = fig.add_subplot(grid[:4, 2 + cont]) axne.set_title("Electron Density") axne_g = fig.add_subplot(grid[4:6, 2 + cont]) @@ -423,9 +64,8 @@ def plotConvergencePORTALS( else: axne = axne_g = axne_f = None - if "nZ" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - impos = portals.prfs_model.mainFunction.PORTALSparameters["ImpurityOfInterest"] - labIon = f"Ion {impos} ({portals.profiles[0].Species[impos-1]['N']}{int(portals.profiles[0].Species[impos-1]['Z'])},{int(portals.profiles[0].Species[impos-1]['A'])})" + if self.runWithImpurity: + labIon = f"Ion {self.runWithImpurity+1} ({self.profiles[0].Species[self.runWithImpurity]['N']}{int(self.profiles[0].Species[self.runWithImpurity]['Z'])},{int(self.profiles[0].Species[self.runWithImpurity]['A'])})" axnZ = fig.add_subplot(grid[:4, 2 + cont]) axnZ.set_title(f"{labIon} Density") axnZ_g = fig.add_subplot(grid[4:6, 2 + cont]) @@ -434,7 +74,7 @@ def plotConvergencePORTALS( else: axnZ = axnZ_g = axnZ_f = None - if "w0" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: + if self.runWithRotation: axw0 = fig.add_subplot(grid[:4, 2 + cont]) axw0.set_title("Rotation") axw0_g = fig.add_subplot(grid[4:6, 2 + cont]) @@ -448,9 +88,9 @@ def plotConvergencePORTALS( axR = fig.add_subplot(grid[6:8, numprofs]) if indexToMaximize is None: - indexToMaximize = portals.numBest + indexToMaximize = self.ibest if indexToMaximize < 0: - indexToMaximize = portals.prfs_model.train_Y.shape[0] + indexToMaximize + indexToMaximize = self.ilast + 1 + indexToMaximize # --------------------------------------------------------------------------------------------------------- # Plot all profiles @@ -459,7 +99,7 @@ def plotConvergencePORTALS( lwt = 0.1 lw = 0.2 alph = 0.6 - for i, p in enumerate(portals.profiles): + for i, p in enumerate(self.profiles): if p is not None: if i < 5: col = "k" @@ -474,9 +114,7 @@ def plotConvergencePORTALS( lab = "" ix = np.argmin( - np.abs( - p.profiles["rho(-)"] - portals.tgyros[portals.numOrig].rho[0][-1] - ) + np.abs(p.profiles["rho(-)"] - self.tgyros[self.i0].rho[0][-1]) ) axTe.plot( p.profiles["rho(-)"], @@ -528,7 +166,7 @@ def plotConvergencePORTALS( if axnZ is not None: axnZ.plot( p.profiles["rho(-)"], - p.profiles["ni(10^19/m^3)"][:, portals.posZ] * 1e-1, + p.profiles["ni(10^19/m^3)"][:, self.runWithImpurity] * 1e-1, lw=lw, color=col, label=lab, @@ -536,7 +174,7 @@ def plotConvergencePORTALS( ) axnZ_g.plot( p.profiles["rho(-)"][:ix], - p.derived["aLni"][:ix, portals.posZ], + p.derived["aLni"][:ix, self.runWithImpurity], lw=lw, color=col, alpha=alph, @@ -559,7 +197,7 @@ def plotConvergencePORTALS( alpha=alph, ) - t = portals.tgyros[i] + t = self.tgyros[i] if (t is not None) and plotAllFluxes: axTe_f.plot( t.rho[0], @@ -587,20 +225,27 @@ def plotConvergencePORTALS( if axne_f is not None: axne_f.plot(t.rho[0], Ge[0], "-", c=col, lw=lwt, alpha=alph) - axne_f.plot(t.rho[0], Ge_tar[0], "--", c=col, lw=lwt, alpha=alph) + axne_f.plot( + t.rho[0], + Ge_tar[0] * (1 - int(self.forceZeroParticleFlux)), + "--", + c=col, + lw=lwt, + alpha=alph, + ) if axnZ_f is not None: if useConvectiveFluxes: GZ, GZ_tar = ( - t.Ci_sim_turb[portals.posZ, :, :] - + t.Ci_sim_turb[portals.posZ, :, :], - t.Ci_tar[portals.posZ, :, :], + t.Ci_sim_turb[self.runWithImpurity, :, :] + + t.Ci_sim_turb[self.runWithImpurity, :, :], + t.Ci_tar[self.runWithImpurity, :, :], ) else: GZ, GZ_tar = ( - t.Gi_sim_turb[portals.posZ, :, :] - + t.Gi_sim_neo[portals.posZ, :, :] - ), t.Gi_tar[portals.posZ, :, :] + t.Gi_sim_turb[self.runWithImpurity, :, :] + + t.Gi_sim_neo[self.runWithImpurity, :, :] + ), t.Gi_tar[self.runWithImpurity, :, :] axnZ_f.plot(t.rho[0], GZ[0], "-", c=col, lw=lwt, alpha=alph) axnZ_f.plot(t.rho[0], GZ_tar[0], "--", c=col, lw=lwt, alpha=alph) @@ -622,20 +267,20 @@ def plotConvergencePORTALS( for cont, (indexUse, col, lab) in enumerate( zip( - [portals.numOrig, portals.numBest, portals.numExtra], + [self.i0, self.ibest, self.iextra], ["r", "g", "m"], [ - f"Initial (#{portals.numOrig})", - f"Best (#{portals.numBest})", - f"Last (#{portals.numExtra})", + f"Initial (#{self.i0})", + f"Best (#{self.ibest})", + f"Last (#{self.iextra})", ], ) ): - if (indexUse is None) or (indexUse >= len(portals.profiles)): + if (indexUse is None) or (indexUse >= len(self.profiles)): continue - p = portals.profiles[indexUse] - t = portals.tgyros[indexUse] + p = self.profiles[indexUse] + t = self.tgyros[indexUse] ix = np.argmin(np.abs(p.profiles["rho(-)"] - t.rho[0][-1])) axTe.plot( @@ -684,14 +329,14 @@ def plotConvergencePORTALS( if axnZ is not None: axnZ.plot( p.profiles["rho(-)"], - p.profiles["ni(10^19/m^3)"][:, portals.posZ] * 1e-1, + p.profiles["ni(10^19/m^3)"][:, self.runWithImpurity] * 1e-1, lw=2, color=col, label=lab, ) axnZ_g.plot( p.profiles["rho(-)"][:ix], - p.derived["aLni"][:ix, portals.posZ], + p.derived["aLni"][:ix, self.runWithImpurity], markersize=msFlux, lw=2, color=col, @@ -714,18 +359,7 @@ def plotConvergencePORTALS( color=col, ) - ( - QeBest_min, - QeBest_max, - QiBest_min, - QiBest_max, - GeBest_min, - GeBest_max, - GZBest_min, - GZBest_max, - MtBest_min, - MtBest_max, - ) = plotFluxComparison( + plotFluxComparison( p, t, axTe_f, @@ -733,16 +367,17 @@ def plotConvergencePORTALS( axne_f, axnZ_f, axw0_f, - posZ=portals.posZ, + runWithImpurity=self.runWithImpurity, fontsize_leg=fontsize_leg, stds=stds, col=col, lab=lab, msFlux=msFlux, + forceZeroParticleFlux=self.forceZeroParticleFlux, useConvectiveFluxes=useConvectiveFluxes, maxStore=indexToMaximize == indexUse, - decor=portals.numBest == indexUse, - plotFlows=plotFlows and (portals.numBest == indexUse), + decor=self.ibest == indexUse, + plotFlows=plotFlows and (self.ibest == indexUse), ) ax = axTe @@ -832,112 +467,112 @@ def plotConvergencePORTALS( ax.set_xticklabels([]) ax = axC - if "te" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resTeM + if "te" in self.ProfilesPredicted: + v = self.resTeM ax.plot( - portals.evaluations, + self.evaluations, v, "-o", lw=0.5, c="b", markersize=2, - label=labelsFluxes["te"], + label=self.labelsFluxes["te"], ) - if "ti" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resTiM + if "ti" in self.ProfilesPredicted: + v = self.resTiM ax.plot( - portals.evaluations, + self.evaluations, v, "-s", lw=0.5, c="m", markersize=2, - label=labelsFluxes["ti"], + label=self.labelsFluxes["ti"], ) - if "ne" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resneM + if "ne" in self.ProfilesPredicted: + v = self.resneM ax.plot( - portals.evaluations, + self.evaluations, v, "-*", lw=0.5, c="k", markersize=2, - label=labelsFluxes["ne"], + label=self.labelsFluxes["ne"], ) - if "nZ" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resnZM + if "nZ" in self.ProfilesPredicted: + v = self.resnZM ax.plot( - portals.evaluations, + self.evaluations, v, "-v", lw=0.5, c="c", markersize=2, - label=labelsFluxes["nZ"], + label=self.labelsFluxes["nZ"], ) - if "w0" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resw0M + if "w0" in self.ProfilesPredicted: + v = self.resw0M ax.plot( - portals.evaluations, + self.evaluations, v, "-v", lw=0.5, c="darkred", markersize=2, - label=labelsFluxes["w0"], + label=self.labelsFluxes["w0"], ) for cont, (indexUse, col, lab, mars) in enumerate( zip( - [portals.numOrig, portals.numBest, portals.numExtra], + [self.i0, self.ibest, self.iextra], ["r", "g", "m"], ["Initial", "Best", "Last"], ["o", "s", "*"], ) ): - if (indexUse is None) or (indexUse >= len(portals.profiles)): + if (indexUse is None) or (indexUse >= len(self.profiles)): continue - if "te" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resTeM + if "te" in self.ProfilesPredicted: + v = self.resTeM ax.plot( - [portals.evaluations[indexUse]], + [self.evaluations[indexUse]], [v[indexUse]], mars, color=col, markersize=4, ) - if "ti" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resTiM + if "ti" in self.ProfilesPredicted: + v = self.resTiM ax.plot( - [portals.evaluations[indexUse]], + [self.evaluations[indexUse]], [v[indexUse]], mars, color=col, markersize=4, ) - if "ne" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resneM + if "ne" in self.ProfilesPredicted: + v = self.resneM ax.plot( - [portals.evaluations[indexUse]], + [self.evaluations[indexUse]], [v[indexUse]], mars, color=col, markersize=4, ) - if "nZ" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resnZM + if "nZ" in self.ProfilesPredicted: + v = self.resnZM ax.plot( - [portals.evaluations[indexUse]], + [self.evaluations[indexUse]], [v[indexUse]], mars, color=col, markersize=4, ) - if "w0" in portals.prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - v = portals.resw0M + if "w0" in self.ProfilesPredicted: + v = self.resw0M ax.plot( - [portals.evaluations[indexUse]], + [self.evaluations[indexUse]], [v[indexUse]], mars, color=col, @@ -945,13 +580,11 @@ def plotConvergencePORTALS( ) # Plot las point as check - ax.plot( - [portals.evaluations[-1]], [portals.resCheck[-1]], "-o", markersize=2, color="k" - ) + ax.plot([self.evaluations[-1]], [self.resCheck[-1]], "-o", markersize=2, color="k") - separator = portals.prfs_model.Optim["initialPoints"] + 0.5 - 1 + separator = self.opt_fun.prfs_model.Optim["initialPoints"] + 0.5 - 1 - if portals.evaluations[-1] < separator: + if self.evaluations[-1] < separator: separator = None GRAPHICStools.addDenseAxis(ax, n=5) @@ -968,8 +601,8 @@ def plotConvergencePORTALS( ratio=0.9, withleg=True, size=fontsize_leg * 1.5, - title="Channels $\\widehat{L_1}$", - ) # ax.legend(prop={'size':fontsize_leg},loc='lower left') + title="Channels $\\frac{1}{N_c}L_1$", + ) ax.set_xticklabels([]) if separator is not None: @@ -993,25 +626,25 @@ def plotConvergencePORTALS( ax = axR for resChosen, label, c in zip( - [portals.resM, portals.resCheck], - ["$\\widehat{L_2}$", "$\\widehat{L_1}$"], + [self.resM, self.resCheck], + ["OF: $\\frac{1}{N}L_2$", "$\\frac{1}{N}L_1$"], ["olive", "rebeccapurple"], ): ax.plot( - portals.evaluations, resChosen, "-o", lw=1.0, c=c, markersize=2, label=label + self.evaluations, resChosen, "-o", lw=1.0, c=c, markersize=2, label=label ) for cont, (indexUse, col, lab, mars) in enumerate( zip( - [portals.numOrig, portals.numBest, portals.numExtra], + [self.i0, self.ibest, self.iextra], ["r", "g", "m"], ["Initial", "Best", "Last"], ["o", "s", "*"], ) ): - if (indexUse is None) or (indexUse >= len(portals.profiles)): + if (indexUse is None) or (indexUse >= len(self.profiles)): continue ax.plot( - [portals.evaluations[indexUse]], + [self.evaluations[indexUse]], [resChosen[indexUse]], "o", color=col, @@ -1038,21 +671,25 @@ def plotConvergencePORTALS( GRAPHICStools.addDenseAxis(ax, n=5) ax.set_xlabel("Iterations (calls/radius)") - ax.set_ylabel("Residual Definitions") + ax.set_ylabel("Residual") ax.set_xlim(left=0) try: ax.set_yscale("log") except: pass GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, size=fontsize_leg * 2.0 - ) # ax.legend(prop={'size':fontsize_leg},loc='lower left') + ax, + ratio=0.9, + withleg=True, + size=fontsize_leg * 2.0, + title="Residuals", + ) ax = axA ax.plot( - portals.DVdistMetric_x, - portals.DVdistMetric_y, + self.DVdistMetric_x, + self.DVdistMetric_y, "-o", c="olive", lw=1.0, @@ -1062,25 +699,25 @@ def plotConvergencePORTALS( for cont, (indexUse, col, lab, mars) in enumerate( zip( - [portals.numOrig, portals.numBest, portals.numExtra], + [self.i0, self.ibest, self.iextra], ["r", "g", "m"], ["Initial", "Best", "Last"], ["o", "s", "*"], ) ): - if (indexUse is None) or (indexUse >= len(portals.profiles)): + if (indexUse is None) or (indexUse >= len(self.profiles)): continue - v = portals.chiR_Ricci - try: - axt.plot( - [portals.evaluations[indexUse]], - [portals.DVdistMetric_y[indexUse]], - "o", - color=col, - markersize=4, - ) - except: - pass + v = self.chiR_Ricci + # try: + # axt.plot( + # [self.evaluations[indexUse]], + # [self.DVdistMetric_y[indexUse]], + # "o", + # color=col, + # markersize=4, + # ) + # except: + # pass if separator is not None: GRAPHICStools.drawLineWithTxt( @@ -1108,11 +745,11 @@ def plotConvergencePORTALS( pass ax.set_xticklabels([]) - if includeRicci and portals.chiR_Ricci is not None: + if includeRicci and self.chiR_Ricci is not None: axt = axA.twinx() (l2,) = axt.plot( - portals.DVdistMetric_x, - portals.DVdistMetric_y, + self.DVdistMetric_x, + self.DVdistMetric_y, "-o", c="olive", lw=1.0, @@ -1120,8 +757,8 @@ def plotConvergencePORTALS( label="$\\Delta$ $a/L_{X}$", ) axt.plot( - portals.evaluations, - portals.chiR_Ricci, + self.evaluations, + self.chiR_Ricci, "-o", lw=1.0, c="rebeccapurple", @@ -1130,17 +767,17 @@ def plotConvergencePORTALS( ) for cont, (indexUse, col, lab, mars) in enumerate( zip( - [portals.numOrig, portals.numBest, portals.numExtra], + [self.i0, self.ibest, self.iextra], ["r", "g", "m"], ["Initial", "Best", "Last"], ["o", "s", "*"], ) ): - if (indexUse is None) or (indexUse >= len(portals.profiles)): + if (indexUse is None) or (indexUse >= len(self.profiles)): continue - v = portals.chiR_Ricci + v = self.chiR_Ricci axt.plot( - [portals.evaluations[indexUse]], + [self.evaluations[indexUse]], [v[indexUse]], "o", color=col, @@ -1150,19 +787,19 @@ def plotConvergencePORTALS( axt.set_ylim([0, 1]) axt.legend(loc="best", prop={"size": fontsize_leg * 1.5}) l2.set_visible(False) - elif portals.aLTn_perc is not None: + elif self.aLTn_perc is not None: ax = axA # .twinx() - x = portals.evaluations + x = self.evaluations - if len(x) > len(portals.aLTn_perc): + if len(x) > len(self.aLTn_perc): x = x[:-1] x0, aLTn_perc0 = [], [] - for i in range(len(portals.aLTn_perc)): - if portals.aLTn_perc[i] is not None: + for i in range(len(self.aLTn_perc)): + if self.aLTn_perc[i] is not None: x0.append(x[i]) - aLTn_perc0.append(portals.aLTn_perc[i].mean()) + aLTn_perc0.append(self.aLTn_perc[i].mean()) ax.plot( x0, aLTn_perc0, @@ -1173,12 +810,12 @@ def plotConvergencePORTALS( label="$\\Delta$ $a/L_{X}^*$ (%)", ) - v = portals.aLTn_perc[portals.numOrig].mean() - ax.plot([portals.evaluations[portals.numOrig]], v, "o", color="r", markersize=4) + v = self.aLTn_perc[self.i0].mean() + ax.plot([self.evaluations[self.i0]], v, "o", color="r", markersize=4) try: - v = portals.aLTn_perc[portals.numBest].mean() + v = self.aLTn_perc[self.ibest].mean() ax.plot( - [portals.evaluations[portals.numBest]], + [self.evaluations[self.ibest]], [v], "o", color="g", @@ -1214,30 +851,30 @@ def plotConvergencePORTALS( ax = axQ - isThereFusion = np.nanmax(portals.FusionGain) > 0 + isThereFusion = np.nanmax(self.FusionGain) > 0 if isThereFusion: - v = portals.FusionGain + v = self.FusionGain axt6 = ax.twinx() # None else: - v = portals.tauE + v = self.tauE axt6 = None # ax.yaxis.tick_right() # ax.yaxis.set_label_position("right") - ax.plot(portals.evaluations, v, "-o", lw=1.0, c="olive", markersize=2, label="$Q$") + ax.plot(self.evaluations, v, "-o", lw=1.0, c="olive", markersize=2, label="$Q$") for cont, (indexUse, col, lab, mars) in enumerate( zip( - [portals.numOrig, portals.numBest, portals.numExtra], + [self.i0, self.ibest, self.iextra], ["r", "g", "m"], ["Initial", "Best", "Last"], ["o", "s", "*"], ) ): - if (indexUse is None) or (indexUse >= len(portals.profiles)): + if (indexUse is None) or (indexUse >= len(self.profiles)): continue ax.plot( - [portals.evaluations[indexUse]], [v[indexUse]], "o", color=col, markersize=4 + [self.evaluations[indexUse]], [v[indexUse]], "o", color=col, markersize=4 ) vmin, vmax = np.max([0, np.nanmin(v)]), np.nanmax(v) @@ -1314,9 +951,9 @@ def plotConvergencePORTALS( ) if (axt6 is not None) and (isThereFusion): - v = portals.FusionPower + v = self.FusionPower axt6.plot( - portals.evaluations, + self.evaluations, v, "-o", lw=1.0, @@ -1326,16 +963,16 @@ def plotConvergencePORTALS( ) for cont, (indexUse, col, lab, mars) in enumerate( zip( - [portals.numOrig, portals.numBest, portals.numExtra], + [self.i0, self.ibest, self.iextra], ["r", "g", "m"], ["Initial", "Best", "Last"], ["o", "s", "*"], ) ): - if (indexUse is None) or (indexUse >= len(portals.profiles)): + if (indexUse is None) or (indexUse >= len(self.profiles)): continue axt6.plot( - [portals.evaluations[indexUse]], + [self.evaluations[indexUse]], [v[indexUse]], "s", color=col, @@ -1346,7 +983,7 @@ def plotConvergencePORTALS( axt6.set_ylim(bottom=0) (l2,) = ax.plot( - portals.evaluations, + self.evaluations, v, "-o", lw=1.0, @@ -1358,7 +995,7 @@ def plotConvergencePORTALS( l2.set_visible(False) for ax in [axQ, axA, axR, axC]: - ax.set_xlim([0, len(portals.FusionGain) + 2]) + ax.set_xlim([0, len(self.FusionGain) + 2]) # for ax in [axA,axR,axC]: # ax.yaxis.tick_right() @@ -1369,707 +1006,299 @@ def plotConvergencePORTALS( # typeMsg="i", # ) + # Save plot + if file_save is not None: + plt.savefig(file_save, transparent=True, dpi=300) -def varToReal(y, prfs_model): - """ - NEO - """ - of, cal, res = prfs_model.mainFunction.scalarized_objective( - torch.Tensor(y).to(prfs_model.mainFunction.dfT).unsqueeze(0) - ) +def PORTALSanalyzer_plotExpected( + self, fig=None, stds=2, max_plot_points=4, plotNext=True +): + print("- Plotting PORTALS Expected") - cont = 0 - Qe, Qi, Ge, GZ, Mt = [], [], [], [], [] - Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = [], [], [], [], [] - for prof in prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - for rad in prfs_model.mainFunction.TGYROparameters["RhoLocations"]: - if prof == "te": - Qe.append(of[0, cont]) - Qe_tar.append(cal[0, cont]) - if prof == "ti": - Qi.append(of[0, cont]) - Qi_tar.append(cal[0, cont]) - if prof == "ne": - Ge.append(of[0, cont]) - Ge_tar.append(cal[0, cont]) - if prof == "nZ": - GZ.append(of[0, cont]) - GZ_tar.append(cal[0, cont]) - if prof == "w0": - Mt.append(of[0, cont]) - Mt_tar.append(cal[0, cont]) + if fig is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + # ---------------------------------------------------------------------- + # Plot + # ---------------------------------------------------------------------- + + trained_points = self.ilast + 1 + self.ibest = self.opt_fun.res.best_absolute_index + + # Best point + plotPoints = [self.ibest] + labelAssigned = [f"#{self.ibest} (best)"] + + # Last point + if (trained_points - 1) != self.ibest: + plotPoints.append(trained_points - 1) + labelAssigned.append(f"#{trained_points-1} (last)") + + # Last ones + i = 0 + while len(plotPoints) < max_plot_points: + if (trained_points - 2 - i) < 1: + break + if (trained_points - 2 - i) != self.ibest: + plotPoints.append(trained_points - 2 - i) + labelAssigned.append(f"#{trained_points-2-i}") + i += 1 + + # First point + if 0 not in plotPoints: + if len(plotPoints) == max_plot_points: + plotPoints[-1] = 0 + labelAssigned[-1] = "#0 (base)" + else: + plotPoints.append(0) + labelAssigned.append("#0 (base)") - cont += 1 + if fig is None: + fig = plt.figure(figsize=(12, 8)) - Qe, Qi, Ge, GZ, Mt = ( - np.array(Qe), - np.array(Qi), - np.array(Ge), - np.array(GZ), - np.array(Mt), - ) - Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = ( - np.array(Qe_tar), - np.array(Qi_tar), - np.array(Ge_tar), - np.array(GZ_tar), - np.array(Mt_tar), - ) + model = self.step.GP["combined_model"] - return Qe, Qi, Ge, GZ, Mt, Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar + x_train_num = self.step.train_X.shape[0] + # ---- Training + x_train = torch.from_numpy(self.step.train_X).to(model.train_X) + y_trainreal = torch.from_numpy(self.step.train_Y).to(model.train_X) + yL_trainreal = torch.from_numpy(self.step.train_Ystd).to(model.train_X) + yU_trainreal = torch.from_numpy(self.step.train_Ystd).to(model.train_X) -def plotVars( - prfs_model, - y, - axs, - axsR, - contP=0, - lines=["-s", "--o"], - yerr=None, - plotPoints=None, - plotResidual=True, - lab="", - color=None, - plotErr=[False] * 10, - colors=GRAPHICStools.listColors(), -): - [axTe_f, axTi_f, axne_f, axnZ_f, axw0_f] = axs - [axTe_r, axTi_r, axne_r, axnZ_r, axw0_r] = axsR + y_train = model.predict(x_train)[0] - ms, cp, lwc = 4, 2, 0.5 + # ---- Next + y_next = yU_next = yL_next = None + if plotNext: + try: + y_next, yU_next, yL_next, _ = model.predict(self.step.x_next) + except: + pass - if plotPoints is None: - plotPoints = range(y.shape[0]) + # ---- Plot - cont = -1 - for i in plotPoints: - cont += 1 + numprofs = len(self.ProfilesPredicted) - lw = 1.5 if i == 0 else 1.0 + if numprofs <= 4: + wspace = 0.3 + else: + wspace = 0.5 - contP += 1 + grid = plt.GridSpec(nrows=4, ncols=numprofs, hspace=0.2, wspace=wspace) - x_var = ( - prfs_model.mainFunction.surrogate_parameters["powerstate"] - .plasma["roa"][0, 1:] - .cpu() - .numpy() - ) # prfs_model.mainFunction.TGYROparameters['RhoLocations'] + axTe = fig.add_subplot(grid[0, 0]) + axTe.set_title("Electron Temperature") + axTe_g = fig.add_subplot(grid[1, 0], sharex=axTe) + axTe_f = fig.add_subplot(grid[2, 0], sharex=axTe) + axTe_r = fig.add_subplot(grid[3, 0], sharex=axTe) - try: - Qe, Qi, Ge, GZ, Mt, Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = varToReal( - y[i, :].detach().cpu().numpy(), prfs_model - ) - except: - continue + axTi = fig.add_subplot(grid[0, 1], sharex=axTe) + axTi.set_title("Ion Temperature") + axTi_g = fig.add_subplot(grid[1, 1], sharex=axTe) + axTi_f = fig.add_subplot(grid[2, 1], sharex=axTe) + axTi_r = fig.add_subplot(grid[3, 1], sharex=axTe) - if yerr is not None: - ( - QeEl, - QiEl, - GeEl, - GZEl, - MtEl, - Qe_tarEl, - Qi_tarEl, - Ge_tarEl, - GZ_tarEl, - Mt_tarEl, - ) = varToReal(yerr[0][i, :].detach().cpu().numpy(), prfs_model) - ( - QeEu, - QiEu, - GeEu, - GZEu, - MtEu, - Qe_tarEu, - Qi_tarEu, - Ge_tarEu, - GZ_tarEu, - Mt_tarEu, - ) = varToReal(yerr[1][i, :].detach().cpu().numpy(), prfs_model) + cont = 0 + if "ne" in self.ProfilesPredicted: + axne = fig.add_subplot(grid[0, 2 + cont], sharex=axTe) + axne.set_title("Electron Density") + axne_g = fig.add_subplot(grid[1, 2 + cont], sharex=axTe) + axne_f = fig.add_subplot(grid[2, 2 + cont], sharex=axTe) + axne_r = fig.add_subplot(grid[3, 2 + cont], sharex=axTe) + cont += 1 + else: + axne = axne_g = axne_f = axne_r = None + if self.runWithImpurity: + labIon = f"Ion {self.runWithImpurity+1} ({self.profiles[0].Species[self.runWithImpurity]['N']}{int(self.profiles[0].Species[self.runWithImpurity]['Z'])},{int(self.profiles[0].Species[self.runWithImpurity]['A'])})" + axnZ = fig.add_subplot(grid[0, 2 + cont], sharex=axTe) + axnZ.set_title(f"{labIon} Density") + axnZ_g = fig.add_subplot(grid[1, 2 + cont], sharex=axTe) + axnZ_f = fig.add_subplot(grid[2, 2 + cont], sharex=axTe) + axnZ_r = fig.add_subplot(grid[3, 2 + cont], sharex=axTe) + cont += 1 + else: + axnZ = axnZ_g = axnZ_f = axnZ_r = None - ax = axTe_f + if self.runWithRotation: + axw0 = fig.add_subplot(grid[0, 2 + cont], sharex=axTe) + axw0.set_title("Rotation") + axw0_g = fig.add_subplot(grid[1, 2 + cont], sharex=axTe) + axw0_f = fig.add_subplot(grid[2, 2 + cont], sharex=axTe) + axw0_r = fig.add_subplot(grid[3, 2 + cont], sharex=axTe) + cont += 1 + else: + axw0 = axw0_g = axw0_f = axw0_r = None - if lines[0] is not None: - ax.plot( - x_var, - Qe, - lines[0], - c=colors[contP] if color is None else color, - label="$Q$" + lab if i == 0 else "", - lw=lw, - markersize=ms, - ) - if lines[1] is not None: - ax.plot( - x_var, - Qe_tar, - lines[1], - c=colors[contP] if color is None else color, + colorsA = GRAPHICStools.listColors() + colors = [] + coli = -1 + for label in labelAssigned: + if "best" in label: + colors.append("g") + elif "last" in label: + colors.append("m") + elif "base" in label: + colors.append("r") + else: + coli += 1 + while colorsA[coli] in ["g", "m", "r"]: + coli += 1 + colors.append(colorsA[coli]) + + rho = self.profiles[0].profiles["rho(-)"] + roa = self.profiles[0].derived["roa"] + rhoVals = self.TGYROparameters["RhoLocations"] + roaVals = np.interp(rhoVals, rho, roa) + lastX = roaVals[-1] + + # ---- Plot profiles + cont = -1 + for i in plotPoints: + cont += 1 + + p = self.profiles[i] + + ix = np.argmin(np.abs(p.derived["roa"] - lastX)) + 1 + + lw = 1.0 if cont > 0 else 1.5 + + ax = axTe + ax.plot( + p.derived["roa"], + p.profiles["te(keV)"], + "-", + c=colors[cont], + label=labelAssigned[cont], + lw=lw, + ) + ax = axTi + ax.plot( + p.derived["roa"], p.profiles["ti(keV)"][:, 0], "-", c=colors[cont], lw=lw + ) + if axne is not None: + ax = axne + ax.plot( + p.derived["roa"], + p.profiles["ne(10^19/m^3)"] * 1e-1, + "-", + c=colors[cont], lw=lw, - markersize=ms, - label="$Q^T$" + lab if i == 0 else "", ) - if yerr is not None: - ax.errorbar( - x_var, - Qe, - c=colors[contP] if color is None else color, - yerr=[QeEl, QeEu], - capsize=cp, - capthick=lwc, - fmt="none", + if axnZ is not None: + ax = axnZ + ax.plot( + p.derived["roa"], + p.profiles["ni(10^19/m^3)"][:, self.runWithImpurity] * 1e-1, + "-", + c=colors[cont], + lw=lw, + ) + if axw0 is not None: + ax = axw0 + ax.plot( + p.derived["roa"], + p.profiles["w0(rad/s)"] * 1e-3, + "-", + c=colors[cont], lw=lw, - markersize=ms, - label="$Q$" + lab if i == 0 else "", ) - ax = axTi_f - if lines[0] is not None: + ax = axTe_g + ax.plot( + p.derived["roa"][:ix], + p.derived["aLTe"][:ix], + "-o", + c=colors[cont], + markersize=0, + lw=lw, + ) + ax = axTi_g + ax.plot( + p.derived["roa"][:ix], + p.derived["aLTi"][:ix, 0], + "-o", + c=colors[cont], + markersize=0, + lw=lw, + ) + if axne_g is not None: + ax = axne_g ax.plot( - x_var, - Qi, - lines[0], - c=colors[contP] if color is None else color, - label=f"#{i}", + p.derived["roa"][:ix], + p.derived["aLne"][:ix], + "-o", + c=colors[cont], + markersize=0, lw=lw, - markersize=ms, ) - if lines[1] is not None: + + if axnZ_g is not None: + ax = axnZ_g ax.plot( - x_var, - Qi_tar, - lines[1], - c=colors[contP] if color is None else color, + p.derived["roa"][:ix], + p.derived["aLni"][:ix, self.runWithImpurity], + "-o", + c=colors[cont], + markersize=0, lw=lw, - markersize=ms, ) - if yerr is not None: - ax.errorbar( - x_var, - Qi, - c=colors[contP] if color is None else color, - yerr=[QiEl, QiEu], - capsize=cp, - capthick=lwc, - fmt="none", + if axw0_g is not None: + ax = axw0_g + ax.plot( + p.derived["roa"][:ix], + p.derived["dw0dr"][:ix] * factor_dw0dr, + "-o", + c=colors[cont], + markersize=0, lw=lw, - markersize=ms, ) - if axne_f is not None: - ax = axne_f - if lines[0] is not None: - ax.plot( - x_var, - Ge, - lines[0], - c=colors[contP] if color is None else color, - label=f"#{i}", - lw=lw, - markersize=ms, - ) - if lines[1] is not None: - ax.plot( - x_var, - Ge_tar, - lines[1], - c=colors[contP] if color is None else color, - lw=lw, - markersize=ms, - ) - if yerr is not None: - ax.errorbar( - x_var, - Ge, - c=colors[contP] if color is None else color, - yerr=[GeEl, GeEu], - capsize=cp, - capthick=lwc, - fmt="none", - lw=lw, - markersize=ms, - ) + cont += 1 - if axnZ_f is not None: - ax = axnZ_f - if lines[0] is not None: - ax.plot( - x_var, - GZ, - lines[0], - c=colors[contP] if color is None else color, - label=f"#{i}", - lw=lw, - markersize=ms, - ) - if lines[1] is not None: - ax.plot( - x_var, - GZ_tar, - lines[1], - c=colors[contP] if color is None else color, - lw=lw, - markersize=ms, - ) - if yerr is not None: - ax.errorbar( - x_var, - GZ, - c=colors[contP] if color is None else color, - yerr=[GZEl, GZEu], - capsize=cp, - capthick=lwc, - fmt="none", - lw=lw, - markersize=ms, - ) + # ---- Plot profiles next - if axw0_f is not None: - ax = axw0_f - if lines[0] is not None: - ax.plot( - x_var, - Mt, - lines[0], - c=colors[contP] if color is None else color, - label=f"#{i}", - lw=lw, - markersize=ms, - ) - if lines[1] is not None: - ax.plot( - x_var, - Mt_tar, - lines[1], - c=colors[contP] if color is None else color, - lw=lw, - markersize=ms, - ) - if yerr is not None: - ax.errorbar( - x_var, - Mt, - c=colors[contP] if color is None else color, - yerr=[MtEl, MtEu], - capsize=cp, - capthick=lwc, - fmt="none", - lw=lw, - markersize=ms, - ) + if self.profiles_next is not None: + p = self.profiles_next + roa = self.profiles_next_new.derived["roa"] + dw0dr = self.profiles_next_new.derived["dw0dr"] - if plotResidual: - ax = axTe_r - if lines[0] is not None: - ax.plot( - x_var, - (Qe - Qe_tar), - lines[0], - c=colors[contP] if color is None else color, - label="$Q-Q^T$" + lab if i == 0 else "", - lw=lw, - markersize=ms, - ) - if plotErr[cont]: - ax.errorbar( - x_var, - (Qe - Qe_tar), - c=colors[contP] if color is None else color, - yerr=[QeEl, QeEu], - capsize=cp, - capthick=lwc, - fmt="none", - lw=0.5, - markersize=0, - ) - - ax = axTi_r - if lines[0] is not None: - ax.plot( - x_var, - (Qi - Qi_tar), - lines[0], - c=colors[contP] if color is None else color, - label=f"#{i}", - lw=lw, - markersize=ms, - ) - if plotErr[cont]: - ax.errorbar( - x_var, - (Qi - Qi_tar), - c=colors[contP] if color is None else color, - yerr=[QiEl, QiEu], - capsize=cp, - capthick=lwc, - fmt="none", - lw=0.5, - markersize=0, - ) - - if axne_r is not None: - ax = axne_r - if lines[0] is not None: - ax.plot( - x_var, - (Ge - Ge_tar), - lines[0], - c=colors[contP] if color is None else color, - label=f"#{i}", - lw=lw, - markersize=ms, - ) - if plotErr[cont]: - ax.errorbar( - x_var, - (Ge - Ge_tar), - c=colors[contP] if color is None else color, - yerr=[GeEl, GeEu], - capsize=cp, - capthick=lwc, - fmt="none", - lw=0.5, - markersize=0, - ) - - if axnZ_r is not None: - ax = axnZ_r - if lines[0] is not None: - ax.plot( - x_var, - (GZ - GZ_tar), - lines[0], - c=colors[contP] if color is None else color, - label=f"#{i}", - lw=lw, - markersize=ms, - ) - if plotErr[cont]: - ax.errorbar( - x_var, - (GZ - GZ_tar), - c=colors[contP] if color is None else color, - yerr=[GZEl, GZEu], - capsize=cp, - capthick=lwc, - fmt="none", - lw=0.5, - markersize=0, - ) - if axw0_r is not None: - ax = axw0_r - if lines[0] is not None: - ax.plot( - x_var, - (Mt - Mt_tar), - lines[0], - c=colors[contP] if color is None else color, - label=f"#{i}", - lw=lw, - markersize=ms, - ) - if plotErr[cont]: - ax.errorbar( - x_var, - (Mt - Mt_tar), - c=colors[contP] if color is None else color, - yerr=[MtEl, MtEu], - capsize=cp, - capthick=lwc, - fmt="none", - lw=0.5, - markersize=0, - ) - - return contP - - -def plotExpected( - prfs_model, - folder, - fn, - labelsFluxes={}, - step=-1, - plotPoints=[0], - labelAssigned=["0"], - plotNext=True, - MITIMextra_dict=None, - stds=2, -): - model = prfs_model.steps[step].GP["combined_model"] - - x_train_num = prfs_model.steps[step].train_X.shape[0] - - posZ = prfs_model.mainFunction.PORTALSparameters["ImpurityOfInterest"] - 1 - - # ---- Training - x_train = torch.from_numpy(prfs_model.steps[step].train_X).to(model.train_X) - y_trainreal = torch.from_numpy(prfs_model.steps[step].train_Y).to(model.train_X) - yL_trainreal = torch.from_numpy(prfs_model.steps[step].train_Ystd).to(model.train_X) - yU_trainreal = torch.from_numpy(prfs_model.steps[step].train_Ystd).to(model.train_X) - - y_train, yU_train, yL_train, _ = model.predict(x_train) - - # ---- Next - x_next = y_next = yU_next = yL_next = None - if plotNext: - try: - x_next = prfs_model.steps[step].x_next - y_next, yU_next, yL_next, _ = model.predict(x_next) - except: - pass - - # ---- Get profiles - profiles = [] - if MITIMextra_dict is not None: - print(f"\t- Reading TGYRO and PROFILES from MITIMextra_dict") - for i in plotPoints: - profiles.append(MITIMextra_dict[i]["tgyro"].results["use"].profiles) - else: - for i in plotPoints: - file = f"{folder}/Execution/Evaluation.{i}/model_complete/input.gacode" - p = PROFILEStools.PROFILES_GACODE(file, calculateDerived=False) - profiles.append(p) - - profiles_next = None - if x_next is not None: - try: - file = f"{folder}/Execution/Evaluation.{x_train_num}/model_complete/input.gacode" - profiles_next = PROFILEStools.PROFILES_GACODE(file, calculateDerived=False) - - try: - file = f"{folder}/Execution/Evaluation.{x_train_num}/model_complete/input.gacode.new" - profiles_next_new = PROFILEStools.PROFILES_GACODE( - file, calculateDerived=True - ) - profiles_next_new.printInfo(label="NEXT") - except: - profiles_next_new = profiles_next - profiles_next_new.deriveQuantities() - except: - pass - - # ---- Plot - fig = fn.add_figure(label="PORTALS Expected") - - numprofs = len(prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]) - - if numprofs <= 4: - wspace = 0.3 - else: - wspace = 0.5 - - grid = plt.GridSpec(nrows=4, ncols=numprofs, hspace=0.2, wspace=wspace) - - axTe = fig.add_subplot(grid[0, 0]) - axTe.set_title("Electron Temperature") - axTe_g = fig.add_subplot(grid[1, 0], sharex=axTe) - axTe_f = fig.add_subplot(grid[2, 0], sharex=axTe) - axTe_r = fig.add_subplot(grid[3, 0], sharex=axTe) - - axTi = fig.add_subplot(grid[0, 1], sharex=axTe) - axTi.set_title("Ion Temperature") - axTi_g = fig.add_subplot(grid[1, 1], sharex=axTe) - axTi_f = fig.add_subplot(grid[2, 1], sharex=axTe) - axTi_r = fig.add_subplot(grid[3, 1], sharex=axTe) - - cont = 0 - if "ne" in prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - axne = fig.add_subplot(grid[0, 2 + cont], sharex=axTe) - axne.set_title("Electron Density") - axne_g = fig.add_subplot(grid[1, 2 + cont], sharex=axTe) - axne_f = fig.add_subplot(grid[2, 2 + cont], sharex=axTe) - axne_r = fig.add_subplot(grid[3, 2 + cont], sharex=axTe) - cont += 1 - else: - axne = axne_g = axne_f = axne_r = None - if "nZ" in prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - impos = prfs_model.mainFunction.PORTALSparameters["ImpurityOfInterest"] - labIon = f"Ion {impos} ({profiles[0].Species[impos-1]['N']}{int(profiles[0].Species[impos-1]['Z'])},{int(profiles[0].Species[impos-1]['A'])})" - axnZ = fig.add_subplot(grid[0, 2 + cont], sharex=axTe) - axnZ.set_title(f"{labIon} Density") - axnZ_g = fig.add_subplot(grid[1, 2 + cont], sharex=axTe) - axnZ_f = fig.add_subplot(grid[2, 2 + cont], sharex=axTe) - axnZ_r = fig.add_subplot(grid[3, 2 + cont], sharex=axTe) - cont += 1 - else: - axnZ = axnZ_g = axnZ_f = axnZ_r = None - - if "w0" in prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: - axw0 = fig.add_subplot(grid[0, 2 + cont], sharex=axTe) - axw0.set_title("Rotation") - axw0_g = fig.add_subplot(grid[1, 2 + cont], sharex=axTe) - axw0_f = fig.add_subplot(grid[2, 2 + cont], sharex=axTe) - axw0_r = fig.add_subplot(grid[3, 2 + cont], sharex=axTe) - cont += 1 - else: - axw0 = axw0_g = axw0_f = axw0_r = None - - colorsA = GRAPHICStools.listColors() - colors = [] - coli = -1 - for label in labelAssigned: - if "best" in label: - colors.append("g") - elif "last" in label: - colors.append("m") - elif "base" in label: - colors.append("r") - else: - coli += 1 - while colorsA[coli] in ["g", "m", "r"]: - coli += 1 - colors.append(colorsA[coli]) - - rho = profiles[0].profiles["rho(-)"] - roa = profiles[0].derived["roa"] - rhoVals = prfs_model.mainFunction.TGYROparameters["RhoLocations"] - roaVals = np.interp(rhoVals, rho, roa) - lastX = roaVals[-1] - - # ---- Plot profiles - cont = -1 - for i in range(len(profiles)): - cont += 1 - - p = profiles[i] - - ix = np.argmin(np.abs(p.derived["roa"] - lastX)) + 1 + ix = np.argmin(np.abs(roa - lastX)) + 1 - lw = 1.0 if cont > 0 else 1.5 + lw = 1.5 ax = axTe ax.plot( - p.derived["roa"], + roa, p.profiles["te(keV)"], "-", - c=colors[cont], - label=labelAssigned[cont], + c="k", + label=f"#{x_train_num} (next)", lw=lw, ) ax = axTi - ax.plot( - p.derived["roa"], p.profiles["ti(keV)"][:, 0], "-", c=colors[cont], lw=lw - ) + ax.plot(roa, p.profiles["ti(keV)"][:, 0], "-", c="k", lw=lw) if axne is not None: ax = axne - ax.plot( - p.derived["roa"], - p.profiles["ne(10^19/m^3)"] * 1e-1, - "-", - c=colors[cont], - lw=lw, - ) + ax.plot(roa, p.profiles["ne(10^19/m^3)"] * 1e-1, "-", c="k", lw=lw) + if axnZ is not None: ax = axnZ ax.plot( - p.derived["roa"], - p.profiles["ni(10^19/m^3)"][:, posZ] * 1e-1, + roa, + p.profiles["ni(10^19/m^3)"][:, self.runWithImpurity] * 1e-1, "-", - c=colors[cont], + c="k", lw=lw, ) if axw0 is not None: ax = axw0 - ax.plot( - p.derived["roa"], - p.profiles["w0(rad/s)"] * 1e-3, - "-", - c=colors[cont], - lw=lw, - ) + ax.plot(roa, p.profiles["w0(rad/s)"] * 1e-3, "-", c="k", lw=lw) ax = axTe_g - ax.plot( - p.derived["roa"][:ix], - p.derived["aLTe"][:ix], - "-o", - c=colors[cont], - markersize=0, - lw=lw, - ) - ax = axTi_g - ax.plot( - p.derived["roa"][:ix], - p.derived["aLTi"][:ix, 0], - "-o", - c=colors[cont], - markersize=0, - lw=lw, - ) - if axne_g is not None: - ax = axne_g - ax.plot( - p.derived["roa"][:ix], - p.derived["aLne"][:ix], - "-o", - c=colors[cont], - markersize=0, - lw=lw, - ) - - if axnZ_g is not None: - ax = axnZ_g - ax.plot( - p.derived["roa"][:ix], - p.derived["aLni"][:ix, posZ], - "-o", - c=colors[cont], - markersize=0, - lw=lw, - ) - if axw0_g is not None: - ax = axw0_g - ax.plot( - p.derived["roa"][:ix], - p.derived["dw0dr"][:ix] * factor_dw0dr, - "-o", - c=colors[cont], - markersize=0, - lw=lw, - ) - - cont += 1 - - # ---- Plot profiles next - - if profiles_next is not None: - p = profiles_next - roa = profiles_next_new.derived["roa"] - dw0dr = profiles_next_new.derived["dw0dr"] - - ix = np.argmin(np.abs(roa - lastX)) + 1 - - lw = 1.5 - - ax = axTe - ax.plot( - roa, - p.profiles["te(keV)"], - "-", - c="k", - label=f"#{x_train_num} (next)", - lw=lw, - ) - ax = axTi - ax.plot(roa, p.profiles["ti(keV)"][:, 0], "-", c="k", lw=lw) - if axne is not None: - ax = axne - ax.plot(roa, p.profiles["ne(10^19/m^3)"] * 1e-1, "-", c="k", lw=lw) - - if axnZ is not None: - ax = axnZ - ax.plot(roa, p.profiles["ni(10^19/m^3)"][:, posZ] * 1e-1, "-", c="k", lw=lw) - if axw0 is not None: - ax = axw0 - ax.plot(roa, p.profiles["w0(rad/s)"] * 1e-3, "-", c="k", lw=lw) - - ax = axTe_g - ax.plot(roa[:ix], p.derived["aLTe"][:ix], "o-", c="k", markersize=0, lw=lw) + ax.plot(roa[:ix], p.derived["aLTe"][:ix], "o-", c="k", markersize=0, lw=lw) ax = axTi_g ax.plot(roa[:ix], p.derived["aLTi"][:ix, 0], "o-", c="k", markersize=0, lw=lw) @@ -2080,7 +1309,12 @@ def plotExpected( if axnZ_g is not None: ax = axnZ_g ax.plot( - roa[:ix], p.derived["aLni"][:ix, posZ], "-o", c="k", markersize=0, lw=lw + roa[:ix], + p.derived["aLni"][:ix, self.runWithImpurity], + "-o", + c="k", + markersize=0, + lw=lw, ) if axw0_g is not None: ax = axw0_g @@ -2093,11 +1327,11 @@ def plotExpected( ranges = [-30, 30] - rho = profiles_next_new.profiles["rho(-)"] - rhoVals = prfs_model.mainFunction.TGYROparameters["RhoLocations"] + rho = self.profiles_next_new.profiles["rho(-)"] + rhoVals = self.TGYROparameters["RhoLocations"] roaVals = np.interp(rhoVals, rho, roa) - p0 = profiles[0] + p0 = self.profiles[plotPoints[0]] zVals = [] z = ((p.derived["aLTe"] - p0.derived["aLTe"]) / p0.derived["aLTe"]) * 100.0 for roai in roaVals: @@ -2105,7 +1339,7 @@ def plotExpected( axTe_g_twin.plot(roaVals, zVals, "--s", c=colors[0], lw=0.5, markersize=4) if len(labelAssigned) > 1 and "last" in labelAssigned[1]: - p0 = profiles[1] + p0 = self.profiles[plotPoints[1]] zVals = [] z = ((p.derived["aLTe"] - p0.derived["aLTe"]) / p0.derived["aLTe"]) * 100.0 for roai in roaVals: @@ -2115,7 +1349,7 @@ def plotExpected( axTe_g_twin.set_ylim(ranges) axTe_g_twin.set_ylabel("(%) from last or best", fontsize=8) - p0 = profiles[0] + p0 = self.profiles[plotPoints[0]] zVals = [] z = ( (p.derived["aLTi"][:, 0] - p0.derived["aLTi"][:, 0]) @@ -2126,7 +1360,7 @@ def plotExpected( axTi_g_twin.plot(roaVals, zVals, "--s", c=colors[0], lw=0.5, markersize=4) if len(labelAssigned) > 1 and "last" in labelAssigned[1]: - p0 = profiles[1] + p0 = self.profiles[plotPoints[1]] zVals = [] z = ( (p.derived["aLTi"][:, 0] - p0.derived["aLTi"][:, 0]) @@ -2145,7 +1379,7 @@ def plotExpected( if axne_g is not None: axne_g_twin = axne_g.twinx() - p0 = profiles[0] + p0 = self.profiles[plotPoints[0]] zVals = [] z = ((p.derived["aLne"] - p0.derived["aLne"]) / p0.derived["aLne"]) * 100.0 for roai in roaVals: @@ -2153,7 +1387,7 @@ def plotExpected( axne_g_twin.plot(roaVals, zVals, "--s", c=colors[0], lw=0.5, markersize=4) if len(labelAssigned) > 1 and "last" in labelAssigned[1]: - p0 = profiles[1] + p0 = self.profiles[plotPoints[1]] zVals = [] z = ( (p.derived["aLne"] - p0.derived["aLne"]) / p0.derived["aLne"] @@ -2172,22 +1406,28 @@ def plotExpected( if axnZ_g is not None: axnZ_g_twin = axnZ_g.twinx() - p0 = profiles[0] + p0 = self.profiles[plotPoints[0]] zVals = [] z = ( - (p.derived["aLni"][:, posZ] - p0.derived["aLni"][:, posZ]) - / p0.derived["aLni"][:, posZ] + ( + p.derived["aLni"][:, self.runWithImpurity] + - p0.derived["aLni"][:, self.runWithImpurity] + ) + / p0.derived["aLni"][:, self.runWithImpurity] ) * 100.0 for roai in roaVals: zVals.append(np.interp(roai, roa, z)) axnZ_g_twin.plot(roaVals, zVals, "--s", c=colors[0], lw=0.5, markersize=4) if len(labelAssigned) > 1 and "last" in labelAssigned[1]: - p0 = profiles[1] + p0 = self.profiles[plotPoints[1]] zVals = [] z = ( - (p.derived["aLni"][:, posZ] - p0.derived["aLni"][:, posZ]) - / p0.derived["aLni"][:, posZ] + ( + p.derived["aLni"][:, self.runWithImpurity] + - p0.derived["aLni"][:, self.runWithImpurity] + ) + / p0.derived["aLni"][:, self.runWithImpurity] ) * 100.0 for roai in roaVals: zVals.append(np.interp(roai, roa, z)) @@ -2203,7 +1443,7 @@ def plotExpected( if axw0_g is not None: axw0_g_twin = axw0_g.twinx() - p0 = profiles[0] + p0 = self.profiles[plotPoints[0]] zVals = [] z = ((dw0dr - p0.derived["dw0dr"]) / p0.derived["dw0dr"]) * 100.0 for roai in roaVals: @@ -2211,7 +1451,7 @@ def plotExpected( axw0_g_twin.plot(roaVals, zVals, "--s", c=colors[0], lw=0.5, markersize=4) if len(labelAssigned) > 1 and "last" in labelAssigned[1]: - p0 = profiles[1] + p0 = self.profiles[plotPoints[1]] zVals = [] z = ((dw0dr - p0.derived["dw0dr"]) / p0.derived["dw0dr"]) * 100.0 for roai in roaVals: @@ -2235,7 +1475,7 @@ def plotExpected( # ---- Plot fluxes cont = plotVars( - prfs_model, + self.opt_fun.prfs_model, y_trainreal, [axTe_f, axTi_f, axne_f, axnZ_f, axw0_f], [axTe_r, axTi_r, axne_r, axnZ_r, axw0_r], @@ -2248,7 +1488,7 @@ def plotExpected( colors=colors, ) _ = plotVars( - prfs_model, + self.opt_fun.prfs_model, y_train, [axTe_f, axTi_f, axne_f, axnZ_f, axw0_f], [axTe_r, axTi_r, axne_r, axnZ_r, axw0_r], @@ -2262,7 +1502,7 @@ def plotExpected( if y_next is not None: cont = plotVars( - prfs_model, + self.opt_fun.prfs_model, y_next, [axTe_f, axTi_f, axne_f, axnZ_f, axw0_f], [axTe_r, axTi_r, axne_r, axnZ_r, axw0_r], @@ -2311,12 +1551,7 @@ def plotExpected( ax.set_ylabel("$w_0$ (krad/s)") GRAPHICStools.addDenseAxis(ax, n=n) - roacoarse = ( - prfs_model.mainFunction.surrogate_parameters["powerstate"] - .plasma["roa"][0, 1:] - .cpu() - .numpy() - ) + roacoarse = self.powerstate.plasma["roa"][0, 1:].cpu().numpy() ax = axTe_g ax.set_xlim([0, 1]) @@ -2390,7 +1625,7 @@ def plotExpected( ax = axTe_f ax.set_xlim([0, 1]) - ax.set_ylabel(labelsFluxes["te"]) + ax.set_ylabel(self.labelsFluxes["te"]) ax.set_ylim(bottom=0) # ax.legend(loc='best',prop={'size':6}) # ax.set_xticklabels([]) @@ -2398,7 +1633,7 @@ def plotExpected( ax = axTi_f ax.set_xlim([0, 1]) - ax.set_ylabel(labelsFluxes["ti"]) + ax.set_ylabel(self.labelsFluxes["ti"]) ax.set_ylim(bottom=0) # ax.set_xticklabels([]) GRAPHICStools.addDenseAxis(ax, n=n) @@ -2406,7 +1641,7 @@ def plotExpected( if axne_f is not None: ax = axne_f ax.set_xlim([0, 1]) - ax.set_ylabel(labelsFluxes["ne"]) + ax.set_ylabel(self.labelsFluxes["ne"]) # GRAPHICStools.addDenseAxis(ax,n=n) # ax.set_xticklabels([]) GRAPHICStools.addDenseAxis(ax, n=n) @@ -2414,7 +1649,7 @@ def plotExpected( if axnZ_f is not None: ax = axnZ_f ax.set_xlim([0, 1]) - ax.set_ylabel(labelsFluxes["nZ"]) + ax.set_ylabel(self.labelsFluxes["nZ"]) # GRAPHICStools.addDenseAxis(ax,n=n) # ax.set_xticklabels([]) GRAPHICStools.addDenseAxis(ax, n=n) @@ -2422,64 +1657,1055 @@ def plotExpected( if axw0_f is not None: ax = axw0_f ax.set_xlim([0, 1]) - ax.set_ylabel(labelsFluxes["w0"]) + ax.set_ylabel(self.labelsFluxes["w0"]) # GRAPHICStools.addDenseAxis(ax,n=n) # ax.set_xticklabels([]) GRAPHICStools.addDenseAxis(ax, n=n) - ax = axTe_r - ax.set_xlim([0, 1]) - ax.set_xlabel("$r/a$") - ax.set_ylabel("Residual " + labelsFluxes["te"]) - GRAPHICStools.addDenseAxis(ax, n=n) - ax.axhline(y=0, lw=0.5, ls="--", c="k") + ax = axTe_r + ax.set_xlim([0, 1]) + ax.set_xlabel("$r/a$") + ax.set_ylabel("Residual " + self.labelsFluxes["te"]) + GRAPHICStools.addDenseAxis(ax, n=n) + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + ax = axTi_r + ax.set_xlim([0, 1]) + ax.set_xlabel("$r/a$") + ax.set_ylabel("Residual " + self.labelsFluxes["ti"]) + GRAPHICStools.addDenseAxis(ax, n=n) + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + if axne_r is not None: + ax = axne_r + ax.set_xlim([0, 1]) + ax.set_xlabel("$r/a$") + ax.set_ylabel("Residual " + self.labelsFluxes["ne"]) + GRAPHICStools.addDenseAxis(ax, n=n) + ax.axhline(y=0, lw=0.5, ls="--", c="k") # + + if axnZ_r is not None: + ax = axnZ_r + ax.set_xlim([0, 1]) + ax.set_xlabel("$r/a$") + ax.set_ylabel("Residual " + self.labelsFluxes["nZ"]) + GRAPHICStools.addDenseAxis(ax, n=n) + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + if axw0_r is not None: + ax = axw0_r + ax.set_xlim([0, 1]) + ax.set_xlabel("$r/a$") + ax.set_ylabel("Residual " + self.labelsFluxes["w0"]) + GRAPHICStools.addDenseAxis(ax, n=n) + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + try: + Qe, Qi, Ge, GZ, Mt, Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = varToReal( + y_trainreal[self.opt_fun.prfs_model.BOmetrics["overall"]["indBest"], :] + .detach() + .cpu() + .numpy(), + self.opt_fun.prfs_model, + ) + rangePlotResidual = np.max([np.max(Qe_tar), np.max(Qi_tar), np.max(Ge_tar)]) + for ax in [axTe_r, axTi_r, axne_r]: + ax.set_ylim( + [-rangePlotResidual * 0.5, rangePlotResidual * 0.5] + ) # 50% of max targets + except: + pass + + +def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): + print("- Plotting PORTALS summary of TGYRO and PROFILES classes") + + indecesPlot = [ + self.ibest, + self.i0, + self.iextra, + ] + + # ------------------------------------------------------- + # Plot TGYROs + # ------------------------------------------------------- + + self.tgyros[indecesPlot[1]].plot( + fn=fn, prelabel=f"({indecesPlot[1]}) TGYRO - ", fn_color=fn_color + ) + if indecesPlot[0] < len(self.tgyros): + self.tgyros[indecesPlot[0]].plot( + fn=fn, prelabel=f"({indecesPlot[0]}) TGYRO - ", fn_color=fn_color + ) + + # ------------------------------------------------------- + # Plot PROFILES + # ------------------------------------------------------- + + figs = [ + fn.add_figure(label="PROFILES - Profiles", tab_color=fn_color), + fn.add_figure(label="PROFILES - Powers", tab_color=fn_color), + fn.add_figure(label="PROFILES - Geometry", tab_color=fn_color), + fn.add_figure(label="PROFILES - Gradients", tab_color=fn_color), + fn.add_figure(label="PROFILES - Flows", tab_color=fn_color), + fn.add_figure(label="PROFILES - Other", tab_color=fn_color), + fn.add_figure(label="PROFILES - Impurities", tab_color=fn_color), + ] + + if indecesPlot[0] < len(self.profiles): + _ = PROFILEStools.plotAll( + [ + self.profiles[indecesPlot[1]], + self.profiles[indecesPlot[0]], + ], + figs=figs, + extralabs=[f"{indecesPlot[1]}", f"{indecesPlot[0]}"], + ) + + # ------------------------------------------------------- + # Plot Comparison + # ------------------------------------------------------- + + profile_original = self.mitim_runs[0]["tgyro"].results["tglf_neo"].profiles + profile_best = self.mitim_runs[self.ibest]["tgyro"].results["tglf_neo"].profiles + + profile_original_unCorrected = self.mitim_runs["profiles_original_un"] + profile_original_0 = self.mitim_runs["profiles_original"] + + fig4 = fn.add_figure(label="PROFILES Comparison", tab_color=fn_color) + grid = plt.GridSpec( + 2, + np.max([3, len(self.ProfilesPredicted)]), + hspace=0.3, + wspace=0.3, + ) + axs4 = [ + fig4.add_subplot(grid[0, 0]), + fig4.add_subplot(grid[1, 0]), + fig4.add_subplot(grid[0, 1]), + fig4.add_subplot(grid[1, 1]), + fig4.add_subplot(grid[0, 2]), + fig4.add_subplot(grid[1, 2]), + ] + + cont = 1 + if self.runWithImpurity: + axs4.append(fig4.add_subplot(grid[0, 2 + cont])) + axs4.append(fig4.add_subplot(grid[1, 2 + cont])) + cont += 1 + if self.runWithRotation: + axs4.append(fig4.add_subplot(grid[0, 2 + cont])) + axs4.append(fig4.add_subplot(grid[1, 2 + cont])) + + colors = GRAPHICStools.listColors() + + for i, (profiles, label, alpha) in enumerate( + zip( + [ + profile_original_unCorrected, + profile_original_0, + profile_original, + profile_best, + ], + ["Original", "Corrected", "Initial", "Final"], + [0.2, 1.0, 1.0, 1.0], + ) + ): + profiles.plotGradients( + axs4, + color=colors[i], + label=label, + lastRho=self.TGYROparameters["RhoLocations"][-1], + alpha=alpha, + useRoa=True, + RhoLocationsPlot=self.TGYROparameters["RhoLocations"], + plotImpurity=self.runWithImpurity, + plotRotation=self.runWithRotation, + ) + + axs4[0].legend(loc="best") + + +def PORTALSanalyzer_plotRanges(self, fig=None): + if fig is None: + plt.ion() + fig = plt.figure() + + pps = np.max( + [3, len(self.ProfilesPredicted)] + ) # Because plotGradients require at least Te, Ti, ne + grid = plt.GridSpec(2, pps, hspace=0.3, wspace=0.3) + axsR = [] + for i in range(pps): + axsR.append(fig.add_subplot(grid[0, i])) + axsR.append(fig.add_subplot(grid[1, i])) + + produceInfoRanges( + self.opt_fun.prfs_model.mainFunction, + self.opt_fun.prfs_model.bounds_orig, + axsR=axsR, + color="k", + lw=0.2, + alpha=0.05, + label="original", + ) + produceInfoRanges( + self.opt_fun.prfs_model.mainFunction, + self.opt_fun.prfs_model.bounds, + axsR=axsR, + color="c", + lw=0.2, + alpha=0.05, + label="final", + ) + + p = self.mitim_runs[0]["tgyro"].results["tglf_neo"].profiles + p.plotGradients( + axsR, + color="b", + lastRho=self.TGYROparameters["RhoLocations"][-1], + ms=0, + lw=1.0, + label="#0", + ls="-o" if self.opt_fun.prfs_model.avoidPoints else "--o", + plotImpurity=self.runWithImpurity, + plotRotation=self.runWithRotation, + ) + + for ikey in self.mitim_runs: + if not isinstance(self.mitim_runs[ikey], dict): + break + + p = self.mitim_runs[ikey]["tgyro"].results["tglf_neo"].profiles + p.plotGradients( + axsR, + color="r", + lastRho=self.TGYROparameters["RhoLocations"][-1], + ms=0, + lw=0.3, + ls="-o" if self.opt_fun.prfs_model.avoidPoints else "-.o", + plotImpurity=self.runWithImpurity, + plotRotation=self.runWithRotation, + ) + + p.plotGradients( + axsR, + color="g", + lastRho=self.TGYROparameters["RhoLocations"][-1], + ms=0, + lw=1.0, + label=f"#{self.opt_fun.res.best_absolute_index} (best)", + plotImpurity=self.runWithImpurity, + plotRotation=self.runWithRotation, + ) + + axsR[0].legend(loc="best") + + +def PORTALSanalyzer_plotModelComparison( + self, + fig=None, + axs=None, + UseTGLFfull_x=None, + includeErrors=True, + includeMetric=True, + includeLegAll=True, +): + print("- Plotting PORTALS Simulations - Model comparison") + + if (fig is None) and (axs is None): + plt.ion() + fig = plt.figure(figsize=(15, 6 if len(self.ProfilesPredicted) < 4 else 10)) + + if axs is None: + if len(self.ProfilesPredicted) < 4: + axs = fig.subplots(ncols=3) + else: + axs = fig.subplots(ncols=3, nrows=2) + + plt.subplots_adjust(wspace=0.25, hspace=0.25) + + axs = axs.flatten() + + # te + quantityX = "QeGB_sim_turb" if UseTGLFfull_x is None else "[TGLF]Qe" + quantityX_stds = "QeGB_sim_turb_stds" if UseTGLFfull_x is None else None + quantityY = "QeGB_sim_turb" + quantityY_stds = "QeGB_sim_turb_stds" + plotModelComparison_quantity( + self, + axs[0], + quantityX=quantityX, + quantityX_stds=quantityX_stds, + quantityY=quantityY, + quantityY_stds=quantityY_stds, + quantity_label="$Q_e^{GB}$", + title="Electron energy flux (GB)", + includeErrors=includeErrors, + includeMetric=includeMetric, + includeLeg=True, + ) + + axs[0].set_xscale("log") + axs[0].set_yscale("log") + + # ti + quantityX = "QiGBIons_sim_turb_thr" if UseTGLFfull_x is None else "[TGLF]Qi" + quantityX_stds = "QiGBIons_sim_turb_thr_stds" if UseTGLFfull_x is None else None + quantityY = "QiGBIons_sim_turb_thr" + quantityY_stds = "QiGBIons_sim_turb_thr_stds" + plotModelComparison_quantity( + self, + axs[1], + quantityX=quantityX, + quantityX_stds=quantityX_stds, + quantityY=quantityY, + quantityY_stds=quantityY_stds, + quantity_label="$Q_i^{GB}$", + title="Ion energy flux (GB)", + includeErrors=includeErrors, + includeMetric=includeMetric, + includeLeg=includeLegAll, + ) + + axs[1].set_xscale("log") + axs[1].set_yscale("log") + + # ne + quantityX = "GeGB_sim_turb" if UseTGLFfull_x is None else "[TGLF]Ge" + quantityX_stds = "GeGB_sim_turb_stds" if UseTGLFfull_x is None else None + quantityY = "GeGB_sim_turb" + quantityY_stds = "GeGB_sim_turb_stds" + plotModelComparison_quantity( + self, + axs[2], + quantityX=quantityX, + quantityX_stds=quantityX_stds, + quantityY=quantityY, + quantityY_stds=quantityY_stds, + quantity_label="$\\Gamma_e^{GB}$", + title="Electron particle flux (GB)", + includeErrors=includeErrors, + includeMetric=includeMetric, + includeLeg=includeLegAll, + ) + + if UseTGLFfull_x is None: + val_calc = self.mitim_runs[0]["tgyro"].results["use"].__dict__[quantityX][0, 1:] + else: + val_calc = np.array( + [ + self.tglf_full.results["ev0"]["TGLFout"][j].__dict__[ + quantityX.replace("[TGLF]", "") + ] + for j in range(len(self.rhos)) + ] + ) + + thre = 10 ** round(np.log10(np.abs(val_calc).min())) + axs[2].set_xscale("symlog", linthresh=thre) + axs[2].set_yscale("symlog", linthresh=thre) + axs[2].tick_params(axis="both", which="major", labelsize=8) + + cont = 1 + if "nZ" in self.ProfilesPredicted: + # nZ + quantityX = "GiGB_sim_turb" if UseTGLFfull_x is None else "[TGLF]GiAll" + quantityX_stds = "GiGB_sim_turb_stds" if UseTGLFfull_x is None else None + quantityY = "GiGB_sim_turb" + quantityY_stds = "GiGB_sim_turb_stds" + plotModelComparison_quantity( + self, + axs[2 + cont], + quantityX=quantityX, + quantityX_stds=quantityX_stds, + quantityY=quantityY, + quantityY_stds=quantityY_stds, + quantity_label="$\\Gamma_Z^{GB}$", + title="Impurity particle flux (GB)", + runWithImpurity=self.runWithImpurity, + includeErrors=includeErrors, + includeMetric=includeMetric, + includeLeg=includeLegAll, + ) + + if UseTGLFfull_x is None: + val_calc = ( + self.mitim_runs[0]["tgyro"] + .results["use"] + .__dict__[quantityX][self.runWithImpurity, 0, 1:] + ) + else: + val_calc = np.array( + [ + self.tglf_full.results["ev0"]["TGLFout"][j].__dict__[ + quantityX.replace("[TGLF]", "") + ] + for j in range(len(self.rhos)) + ] + )[self.runWithImpurity] + + thre = 10 ** round(np.log10(np.abs(val_calc).min())) + axs[2 + cont].set_xscale("symlog", linthresh=thre) + axs[2 + cont].set_yscale("symlog", linthresh=thre) + axs[2 + cont].tick_params(axis="both", which="major", labelsize=8) + + cont += 1 + + if "w0" in self.ProfilesPredicted: + if UseTGLFfull_x is not None: + raise Exception("Momentum plot not implemented yet") + # w0 + quantityX = "MtGB_sim_turb" + quantityX_stds = "MtGB_sim_turb_stds" + quantityY = "MtGB_sim_turb" + quantityY_stds = "MtGB_sim_turb_stds" + plotModelComparison_quantity( + self, + axs[2 + cont], + quantityX=quantityX, + quantityX_stds=quantityX_stds, + quantityY=quantityY, + quantityY_stds=quantityY_stds, + quantity_label="$M_T^{GB}$", + title="Momentum Flux (GB)", + includeErrors=includeErrors, + includeMetric=includeMetric, + includeLeg=includeLegAll, + ) + + thre = 10 ** round( + np.log10( + np.abs( + self.mitim_runs[0]["tgyro"] + .results["use"] + .__dict__[quantityX][0, 1:] + ).min() + ) + ) + axs[2 + cont].set_xscale("symlog", linthresh=thre) + axs[2 + cont].set_yscale("symlog", linthresh=thre) + axs[2 + cont].tick_params(axis="both", which="major", labelsize=8) + + cont += 1 + + if self.PORTALSparameters["surrogateForTurbExch"]: + if UseTGLFfull_x is not None: + raise Exception("Turbulent exchange plot not implemented yet") + # Sexch + quantityX = "EXeGB_sim_turb" + quantityX_stds = "EXeGB_sim_turb_stds" + quantityY = "EXeGB_sim_turb" + quantityY_stds = "EXeGB_sim_turb_stds" + plotModelComparison_quantity( + self, + axs[2 + cont], + quantityX=quantityX, + quantityX_stds=quantityX_stds, + quantityY=quantityY, + quantityY_stds=quantityY_stds, + quantity_label="$S_{exch}^{GB}$", + title="Turbulent Exchange (GB)", + includeErrors=includeErrors, + includeMetric=includeMetric, + includeLeg=includeLegAll, + ) + + thre = 10 ** round( + np.log10( + np.abs( + self.mitim_runs[0]["tgyro"] + .results["use"] + .__dict__[quantityX][0, 1:] + ).min() + ) + ) + axs[2 + cont].set_xscale("symlog", linthresh=thre) + axs[2 + cont].set_yscale("symlog", linthresh=thre) + axs[2 + cont].tick_params(axis="both", which="major", labelsize=8) + + cont += 1 + + return axs + + +def plotModelComparison_quantity( + self, + ax, + quantityX="QeGB_sim_turb", + quantityX_stds="QeGB_sim_turb_stds", + quantityY="QeGB_sim_turb", + quantityY_stds="QeGB_sim_turb_stds", + quantity_label="", + title="", + runWithImpurity=None, + includeErrors=True, + includeMetric=True, + includeLeg=True, +): + resultsX = "tglf_neo" + quantity_label_resultsX = "(TGLF)" + + if "cgyro_neo" in self.mitim_runs[0]["tgyro"].results: + resultsY = "cgyro_neo" + quantity_label_resultsY = "(CGYRO)" + else: + resultsY = resultsX + quantity_label_resultsY = quantity_label_resultsX + + X, X_stds = [], [] + Y, Y_stds = [], [] + for i in range(self.ilast + 1): + """ + Read the fluxes to be plotted in Y from the TGYRO results + """ + t = self.mitim_runs[i]["tgyro"].results + Y.append( + t[resultsY].__dict__[quantityY][ + ... if runWithImpurity is None else runWithImpurity, 0, 1: + ] + ) + Y_stds.append( + t[resultsY].__dict__[quantityY_stds][ + ... if runWithImpurity is None else runWithImpurity, 0, 1: + ] + ) + + """ + Read the fluxes to be plotted in X from... + """ + + # ...from the TGLF full results + if "[TGLF]" in quantityX: + X.append( + [ + self.tglf_full.results[f"ev{i}"]["TGLFout"][j].__dict__[ + quantityX.replace("[TGLF]", "") + ] + for j in range(len(self.rhos)) + ] + ) + X_stds.append([np.nan for j in range(len(self.rhos))]) + + # ...from the TGLF results + else: + X.append( + t[resultsX].__dict__[quantityX][ + (... if runWithImpurity is None else runWithImpurity), 0, 1: + ] + ) + X_stds.append( + t[resultsX].__dict__[quantityX_stds][ + ... if runWithImpurity is None else runWithImpurity, 0, 1: + ] + ) + + X = np.array(X) + Y = np.array(Y) + X_stds = np.array(X_stds) + Y_stds = np.array(Y_stds) + + colors = GRAPHICStools.listColors() + + for ir in range(X.shape[1]): + label = f"$r/a={self.roa[ir]:.2f}$" + if includeMetric: + metric, lab_metric = add_metric(None, X[:, ir], Y[:, ir]) + label += f", {lab_metric}: {metric:.2f}" + + ax.errorbar( + X[:, ir], + Y[:, ir], + xerr=X_stds[:, ir] if includeErrors else None, + yerr=Y_stds[:, ir] if includeErrors else None, + c=colors[ir], + markersize=2, + capsize=2, + fmt="s", + elinewidth=1.0, + capthick=1.0, + label=label, + ) + + # ------------------------------------------------------- + # Decorations + # ------------------------------------------------------- + + minFlux = np.min([X.min(), Y.min()]) + maxFlux = np.max([X.max(), Y.max()]) + + minFlux = minFlux - 0.25 * (maxFlux - minFlux) + maxFlux = maxFlux + 0.25 * (maxFlux - minFlux) + + ax.plot([minFlux, maxFlux], [minFlux, maxFlux], "-", color="k", lw=0.5) + + ax.set_xlabel(f"{quantity_label} {quantity_label_resultsX}") + ax.set_ylabel(f"{quantity_label} {quantity_label_resultsY}") + ax.set_title(title) + GRAPHICStools.addDenseAxis(ax) + + sizeLeg = 7 + + if includeLeg: + legend = ax.legend(loc="best", prop={"size": sizeLeg}) + + if includeMetric: + metric, lab_metric = add_metric( + ax if not includeLeg else None, X, Y, fontsize=sizeLeg + ) + if includeLeg: + legend.set_title(f"{lab_metric}: {metric:.2f}") + plt.setp( + legend.get_title(), + bbox=dict( + facecolor="lightgreen", + alpha=0.3, + edgecolor="black", + boxstyle="round,pad=0.2", + ), + ) + legend.get_title().set_fontsize(sizeLeg) + + +# --------------------------------------------------------------------------------------------------------------------- + + +def add_metric(ax, X, Y, typeM="RMSE", fontsize=8): + if typeM == "RMSE": + metric = np.sqrt(np.mean((X - Y) ** 2)) + metric_lab = "RMSE" + if ax is not None: + ax.text( + 0.05, + 0.95, + f"{metric_lab}: {metric:.2f}", + ha="left", + va="top", + transform=ax.transAxes, + bbox=dict( + facecolor="lightgreen", + alpha=0.3, + edgecolor="black", + boxstyle="round,pad=0.2", + ), + fontsize=fontsize, + ) + + return metric, metric_lab + + +def varToReal(y, prfs_model): + of, cal, res = prfs_model.mainFunction.scalarized_objective( + torch.Tensor(y).to(prfs_model.mainFunction.dfT).unsqueeze(0) + ) + + cont = 0 + Qe, Qi, Ge, GZ, Mt = [], [], [], [], [] + Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = [], [], [], [], [] + for prof in prfs_model.mainFunction.TGYROparameters["ProfilesPredicted"]: + for rad in prfs_model.mainFunction.TGYROparameters["RhoLocations"]: + if prof == "te": + Qe.append(of[0, cont]) + Qe_tar.append(cal[0, cont]) + if prof == "ti": + Qi.append(of[0, cont]) + Qi_tar.append(cal[0, cont]) + if prof == "ne": + Ge.append(of[0, cont]) + Ge_tar.append(cal[0, cont]) + if prof == "nZ": + GZ.append(of[0, cont]) + GZ_tar.append(cal[0, cont]) + if prof == "w0": + Mt.append(of[0, cont]) + Mt_tar.append(cal[0, cont]) + + cont += 1 + + Qe, Qi, Ge, GZ, Mt = ( + np.array(Qe), + np.array(Qi), + np.array(Ge), + np.array(GZ), + np.array(Mt), + ) + Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = ( + np.array(Qe_tar), + np.array(Qi_tar), + np.array(Ge_tar), + np.array(GZ_tar), + np.array(Mt_tar), + ) + + return Qe, Qi, Ge, GZ, Mt, Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar + + +def plotVars( + prfs_model, + y, + axs, + axsR, + contP=0, + lines=["-s", "--o"], + yerr=None, + plotPoints=None, + plotResidual=True, + lab="", + color=None, + plotErr=[False] * 10, + colors=GRAPHICStools.listColors(), +): + [axTe_f, axTi_f, axne_f, axnZ_f, axw0_f] = axs + [axTe_r, axTi_r, axne_r, axnZ_r, axw0_r] = axsR + + ms, cp, lwc = 4, 2, 0.5 + + if plotPoints is None: + plotPoints = range(y.shape[0]) + + cont = -1 + for i in plotPoints: + cont += 1 + + lw = 1.5 if i == 0 else 1.0 + + contP += 1 + + x_var = ( + prfs_model.mainFunction.surrogate_parameters["powerstate"] + .plasma["roa"][0, 1:] + .cpu() + .numpy() + ) # prfs_model.mainFunction.TGYROparameters['RhoLocations'] + + try: + Qe, Qi, Ge, GZ, Mt, Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = varToReal( + y[i, :].detach().cpu().numpy(), prfs_model + ) + except: + continue + + if yerr is not None: + ( + QeEl, + QiEl, + GeEl, + GZEl, + MtEl, + Qe_tarEl, + Qi_tarEl, + Ge_tarEl, + GZ_tarEl, + Mt_tarEl, + ) = varToReal(yerr[0][i, :].detach().cpu().numpy(), prfs_model) + ( + QeEu, + QiEu, + GeEu, + GZEu, + MtEu, + Qe_tarEu, + Qi_tarEu, + Ge_tarEu, + GZ_tarEu, + Mt_tarEu, + ) = varToReal(yerr[1][i, :].detach().cpu().numpy(), prfs_model) + + ax = axTe_f + + if lines[0] is not None: + ax.plot( + x_var, + Qe, + lines[0], + c=colors[contP] if color is None else color, + label="$Q$" + lab if i == 0 else "", + lw=lw, + markersize=ms, + ) + if lines[1] is not None: + ax.plot( + x_var, + Qe_tar, + lines[1], + c=colors[contP] if color is None else color, + lw=lw, + markersize=ms, + label="$Q^T$" + lab if i == 0 else "", + ) + if yerr is not None: + ax.errorbar( + x_var, + Qe, + c=colors[contP] if color is None else color, + yerr=[QeEl, QeEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=lw, + markersize=ms, + label="$Q$" + lab if i == 0 else "", + ) + + ax = axTi_f + if lines[0] is not None: + ax.plot( + x_var, + Qi, + lines[0], + c=colors[contP] if color is None else color, + label=f"#{i}", + lw=lw, + markersize=ms, + ) + if lines[1] is not None: + ax.plot( + x_var, + Qi_tar, + lines[1], + c=colors[contP] if color is None else color, + lw=lw, + markersize=ms, + ) + if yerr is not None: + ax.errorbar( + x_var, + Qi, + c=colors[contP] if color is None else color, + yerr=[QiEl, QiEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=lw, + markersize=ms, + ) + + if axne_f is not None: + ax = axne_f + if lines[0] is not None: + ax.plot( + x_var, + Ge, + lines[0], + c=colors[contP] if color is None else color, + label=f"#{i}", + lw=lw, + markersize=ms, + ) + if lines[1] is not None: + ax.plot( + x_var, + Ge_tar, + lines[1], + c=colors[contP] if color is None else color, + lw=lw, + markersize=ms, + ) + if yerr is not None: + ax.errorbar( + x_var, + Ge, + c=colors[contP] if color is None else color, + yerr=[GeEl, GeEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=lw, + markersize=ms, + ) + + if axnZ_f is not None: + ax = axnZ_f + if lines[0] is not None: + ax.plot( + x_var, + GZ, + lines[0], + c=colors[contP] if color is None else color, + label=f"#{i}", + lw=lw, + markersize=ms, + ) + if lines[1] is not None: + ax.plot( + x_var, + GZ_tar, + lines[1], + c=colors[contP] if color is None else color, + lw=lw, + markersize=ms, + ) + if yerr is not None: + ax.errorbar( + x_var, + GZ, + c=colors[contP] if color is None else color, + yerr=[GZEl, GZEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=lw, + markersize=ms, + ) + + if axw0_f is not None: + ax = axw0_f + if lines[0] is not None: + ax.plot( + x_var, + Mt, + lines[0], + c=colors[contP] if color is None else color, + label=f"#{i}", + lw=lw, + markersize=ms, + ) + if lines[1] is not None: + ax.plot( + x_var, + Mt_tar, + lines[1], + c=colors[contP] if color is None else color, + lw=lw, + markersize=ms, + ) + if yerr is not None: + ax.errorbar( + x_var, + Mt, + c=colors[contP] if color is None else color, + yerr=[MtEl, MtEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=lw, + markersize=ms, + ) - ax = axTi_r - ax.set_xlim([0, 1]) - ax.set_xlabel("$r/a$") - ax.set_ylabel("Residual " + labelsFluxes["ti"]) - GRAPHICStools.addDenseAxis(ax, n=n) - ax.axhline(y=0, lw=0.5, ls="--", c="k") + if plotResidual: + ax = axTe_r + if lines[0] is not None: + ax.plot( + x_var, + (Qe - Qe_tar), + lines[0], + c=colors[contP] if color is None else color, + label="$Q-Q^T$" + lab if i == 0 else "", + lw=lw, + markersize=ms, + ) + if plotErr[cont]: + ax.errorbar( + x_var, + (Qe - Qe_tar), + c=colors[contP] if color is None else color, + yerr=[QeEl, QeEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=0.5, + markersize=0, + ) - if axne_r is not None: - ax = axne_r - ax.set_xlim([0, 1]) - ax.set_xlabel("$r/a$") - ax.set_ylabel("Residual " + labelsFluxes["ne"]) - GRAPHICStools.addDenseAxis(ax, n=n) - ax.axhline(y=0, lw=0.5, ls="--", c="k") # + ax = axTi_r + if lines[0] is not None: + ax.plot( + x_var, + (Qi - Qi_tar), + lines[0], + c=colors[contP] if color is None else color, + label=f"#{i}", + lw=lw, + markersize=ms, + ) + if plotErr[cont]: + ax.errorbar( + x_var, + (Qi - Qi_tar), + c=colors[contP] if color is None else color, + yerr=[QiEl, QiEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=0.5, + markersize=0, + ) - if axnZ_r is not None: - ax = axnZ_r - ax.set_xlim([0, 1]) - ax.set_xlabel("$r/a$") - ax.set_ylabel("Residual " + labelsFluxes["nZ"]) - GRAPHICStools.addDenseAxis(ax, n=n) - ax.axhline(y=0, lw=0.5, ls="--", c="k") + if axne_r is not None: + ax = axne_r + if lines[0] is not None: + ax.plot( + x_var, + (Ge - Ge_tar), + lines[0], + c=colors[contP] if color is None else color, + label=f"#{i}", + lw=lw, + markersize=ms, + ) + if plotErr[cont]: + ax.errorbar( + x_var, + (Ge - Ge_tar), + c=colors[contP] if color is None else color, + yerr=[GeEl, GeEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=0.5, + markersize=0, + ) - if axw0_r is not None: - ax = axw0_r - ax.set_xlim([0, 1]) - ax.set_xlabel("$r/a$") - ax.set_ylabel("Residual " + labelsFluxes["w0"]) - GRAPHICStools.addDenseAxis(ax, n=n) - ax.axhline(y=0, lw=0.5, ls="--", c="k") + if axnZ_r is not None: + ax = axnZ_r + if lines[0] is not None: + ax.plot( + x_var, + (GZ - GZ_tar), + lines[0], + c=colors[contP] if color is None else color, + label=f"#{i}", + lw=lw, + markersize=ms, + ) + if plotErr[cont]: + ax.errorbar( + x_var, + (GZ - GZ_tar), + c=colors[contP] if color is None else color, + yerr=[GZEl, GZEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=0.5, + markersize=0, + ) + if axw0_r is not None: + ax = axw0_r + if lines[0] is not None: + ax.plot( + x_var, + (Mt - Mt_tar), + lines[0], + c=colors[contP] if color is None else color, + label=f"#{i}", + lw=lw, + markersize=ms, + ) + if plotErr[cont]: + ax.errorbar( + x_var, + (Mt - Mt_tar), + c=colors[contP] if color is None else color, + yerr=[MtEl, MtEu], + capsize=cp, + capthick=lwc, + fmt="none", + lw=0.5, + markersize=0, + ) - try: - Qe, Qi, Ge, GZ, Mt, Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = varToReal( - y_trainreal[prfs_model.BOmetrics["overall"]["indBest"], :] - .detach() - .cpu() - .numpy(), - prfs_model, - ) - rangePlotResidual = np.max([np.max(Qe_tar), np.max(Qi_tar), np.max(Ge_tar)]) - for ax in [axTe_r, axTi_r, axne_r]: - ax.set_ylim( - [-rangePlotResidual * 0.5, rangePlotResidual * 0.5] - ) # 50% of max targets - except: - pass + return contP def plotFluxComparison( @@ -2490,7 +2716,8 @@ def plotFluxComparison( axne_f, axnZ_f, axw0_f, - posZ=3, + forceZeroParticleFlux=False, + runWithImpurity=3, labZ="Z", includeFirst=True, alpha=1.0, @@ -2618,6 +2845,9 @@ def plotFluxComparison( except: sigma = t.Qe_sim_turb[0][ixF:] * 0.0 + if forceZeroParticleFlux: + Ge_tar = Ge_tar * 0.0 + if axne_f is not None: axne_f.plot( r[0][ixF:], @@ -2650,22 +2880,26 @@ def plotFluxComparison( if axnZ_f is not None: if useConvectiveFluxes: GZ, GZ_tar = ( - t.Ci_sim_turb[posZ, :, :] + t.Ci_sim_neo[posZ, :, :], + t.Ci_sim_turb[runWithImpurity, :, :] + + t.Ci_sim_neo[runWithImpurity, :, :], t.Ge_tar * 0.0, ) try: sigma = ( - t.Ci_sim_turb_stds[posZ, 0][ixF:] + t.Ci_sim_neo_stds[posZ, 0][ixF:] + t.Ci_sim_turb_stds[runWithImpurity, 0][ixF:] + + t.Ci_sim_neo_stds[runWithImpurity, 0][ixF:] ) except: sigma = t.Qe_sim_turb[0][ixF:] * 0.0 else: GZ, GZ_tar = ( - t.Gi_sim_turb[posZ, :, :] + t.Gi_sim_neo[posZ, :, :] + t.Gi_sim_turb[runWithImpurity, :, :] + + t.Gi_sim_neo[runWithImpurity, :, :] ), t.Ge_tar * 0.0 try: sigma = ( - t.Gi_sim_turb_stds[posZ, 0][ixF:] + t.Gi_sim_neo_stds[posZ, 0][ixF:] + t.Gi_sim_turb_stds[runWithImpurity, 0][ixF:] + + t.Gi_sim_neo_stds[runWithImpurity, 0][ixF:] ) except: sigma = t.Qe_sim_turb[0][ixF:] * 0.0 @@ -2891,15 +3125,162 @@ def plotFluxComparison( Q = np.max([np.abs(Qmin), np.abs(Qmax)]) axw0_f.set_ylim([-Q, Q]) - return ( - QeBest_min, - QeBest_max, - QiBest_min, - QiBest_max, - GeBest_min, - GeBest_max, - GZBest_min, - GZBest_max, - MtBest_min, - MtBest_max, + +def produceInfoRanges( + self_complete, bounds, axsR, label="", color="k", lw=0.2, alpha=0.05 +): + rhos = np.append([0], self_complete.TGYROparameters["RhoLocations"]) + aLTe, aLTi, aLne, aLnZ, aLw0 = ( + np.zeros((len(rhos), 2)), + np.zeros((len(rhos), 2)), + np.zeros((len(rhos), 2)), + np.zeros((len(rhos), 2)), + np.zeros((len(rhos), 2)), + ) + for i in range(len(rhos) - 1): + if f"aLte_{i+1}" in bounds: + aLTe[i + 1, :] = bounds[f"aLte_{i+1}"] + if f"aLti_{i+1}" in bounds: + aLTi[i + 1, :] = bounds[f"aLti_{i+1}"] + if f"aLne_{i+1}" in bounds: + aLne[i + 1, :] = bounds[f"aLne_{i+1}"] + if f"aLnZ_{i+1}" in bounds: + aLnZ[i + 1, :] = bounds[f"aLnZ_{i+1}"] + if f"aLw0_{i+1}" in bounds: + aLw0[i + 1, :] = bounds[f"aLw0_{i+1}"] + + X = torch.zeros( + ((len(rhos) - 1) * len(self_complete.TGYROparameters["ProfilesPredicted"]), 2) + ) + l = len(rhos) - 1 + X[0:l, :] = torch.from_numpy(aLTe[1:, :]) + X[l : 2 * l, :] = torch.from_numpy(aLTi[1:, :]) + + cont = 0 + if "ne" in self_complete.TGYROparameters["ProfilesPredicted"]: + X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLne[1:, :]) + cont += 1 + if "nZ" in self_complete.TGYROparameters["ProfilesPredicted"]: + X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLnZ[1:, :]) + cont += 1 + if "w0" in self_complete.TGYROparameters["ProfilesPredicted"]: + X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLw0[1:, :]) + cont += 1 + + X = X.transpose(0, 1) + + powerstate = PORTALStools.constructEvaluationProfiles( + X, self_complete.surrogate_parameters, recalculateTargets=False + ) + + GRAPHICStools.fillGraph( + axsR[0], + powerstate.plasma["rho"][0], + powerstate.plasma["te"][0], + y_up=powerstate.plasma["te"][1], + alpha=alpha, + color=color, + lw=lw, + label=label, + ) + GRAPHICStools.fillGraph( + axsR[1], + rhos, + aLTe[:, 0], + y_up=aLTe[:, 1], + alpha=alpha, + color=color, + label=label, + lw=lw, + ) + + GRAPHICStools.fillGraph( + axsR[2], + powerstate.plasma["rho"][0], + powerstate.plasma["ti"][0], + y_up=powerstate.plasma["ti"][1], + alpha=alpha, + color=color, + label=label, + lw=lw, + ) + GRAPHICStools.fillGraph( + axsR[3], + rhos, + aLTi[:, 0], + y_up=aLTi[:, 1], + alpha=alpha, + color=color, + label=label, + lw=lw, ) + + cont = 0 + if "ne" in self_complete.TGYROparameters["ProfilesPredicted"]: + GRAPHICStools.fillGraph( + axsR[3 + cont + 1], + powerstate.plasma["rho"][0], + powerstate.plasma["ne"][0] * 0.1, + y_up=powerstate.plasma["ne"][1] * 0.1, + alpha=alpha, + color=color, + label=label, + lw=lw, + ) + GRAPHICStools.fillGraph( + axsR[3 + cont + 2], + rhos, + aLne[:, 0], + y_up=aLne[:, 1], + alpha=alpha, + color=color, + label=label, + lw=lw, + ) + cont += 2 + + if "nZ" in self_complete.TGYROparameters["ProfilesPredicted"]: + GRAPHICStools.fillGraph( + axsR[3 + cont + 1], + powerstate.plasma["rho"][0], + powerstate.plasma["nZ"][0] * 0.1, + y_up=powerstate.plasma["nZ"][1] * 0.1, + alpha=alpha, + color=color, + label=label, + lw=lw, + ) + GRAPHICStools.fillGraph( + axsR[3 + cont + 2], + rhos, + aLnZ[:, 0], + y_up=aLnZ[:, 1], + alpha=alpha, + color=color, + label=label, + lw=lw, + ) + cont += 2 + + if "w0" in self_complete.TGYROparameters["ProfilesPredicted"]: + GRAPHICStools.fillGraph( + axsR[3 + cont + 1], + powerstate.plasma["rho"][0], + powerstate.plasma["w0"][0] * 1e-3, + y_up=powerstate.plasma["w0"][1] * 1e-3, + alpha=alpha, + color=color, + label=label, + lw=lw, + ) + GRAPHICStools.fillGraph( + axsR[3 + cont + 2], + rhos, + aLw0[:, 0], + y_up=aLw0[:, 1], + alpha=alpha, + color=color, + label=label, + lw=lw, + ) + cont += 2 diff --git a/src/mitim_modules/portals/exe/check_initialization.py b/src/mitim_modules/portals/exe/check_initialization.py deleted file mode 100644 index 228c8534..00000000 --- a/src/mitim_modules/portals/exe/check_initialization.py +++ /dev/null @@ -1,117 +0,0 @@ -import sys -import matplotlib.pyplot as plt -from mitim_modules.powertorch import STATEtools -from mitim_tools.misc_tools import IOtools, GRAPHICStools -from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.misc_tools.GUItools import FigureNotebook - -""" -To check how the method-6 initializaiton is going... - -E.g. - run ~/MITIM/mitim_opt/mitim/exe/check_initialization.py ~/PRF/project_2022_mitim/jet/cgyro/jet_cgyro_D/ -""" - -folder = IOtools.expandPath(sys.argv[1]) - -# Read powerstates -ps = [] -profs = [] -for i in range(10): - try: - profs.append( - PROFILEStools.PROFILES_GACODE( - f"{folder}/Outputs/ProfilesEvaluated/input.gacode.{i}" - ) - ) - p = STATEtools.read_saved_state( - f"{folder}/Initialization/initialization_simple_relax/portals_{IOtools.reducePathLevel(folder)[1]}_ev{i}/powerstate.pkl" - ) - except: - break - - p.profiles.deriveQuantities() - ps.append(p) - -plt.ioff() -fn = FigureNotebook(0, "PowerState", geometry="1800x900") -figMain = fn.add_figure(label="PowerState") -figG = fn.add_figure(label="Sequence") - -grid = plt.GridSpec(4, 6, hspace=0.3, wspace=0.4) -axs = [ - figMain.add_subplot(grid[0, 1]), - figMain.add_subplot(grid[0, 2]), - figMain.add_subplot(grid[0, 3]), - figMain.add_subplot(grid[0, 4]), - figMain.add_subplot(grid[0, 5]), - figMain.add_subplot(grid[1, 1]), - figMain.add_subplot(grid[1, 2]), - figMain.add_subplot(grid[1, 3]), - figMain.add_subplot(grid[1, 4]), - figMain.add_subplot(grid[1, 5]), - figMain.add_subplot(grid[2, 1]), - figMain.add_subplot(grid[2, 2]), - figMain.add_subplot(grid[2, 3]), - figMain.add_subplot(grid[2, 4]), - figMain.add_subplot(grid[2, 5]), - figMain.add_subplot(grid[3, 1]), - figMain.add_subplot(grid[3, 2]), - figMain.add_subplot(grid[3, 3]), - figMain.add_subplot(grid[3, 4]), - figMain.add_subplot(grid[3, 5]), -] - -axsRes = figMain.add_subplot(grid[:, 0]) - -colors = GRAPHICStools.listColors() - -# POWERPLOT - -for i in range(len(ps)): - ps[i].plot(axs=axs, axsRes=axsRes, c=colors[i], label=f"#{i}") - -axs[0].legend(prop={"size": 8}) - -axsRes.set_xlim([0, i]) - -# GRADIENTS - -grid = plt.GridSpec(2, 5, hspace=0.3, wspace=0.3) -axsGrads = [] -for j in range(5): - for i in range(2): - axsGrads.append(figG.add_subplot(grid[i, j])) -for i, p in enumerate(profs): - p.plotGradients( - axsGrads, - color=colors[i], - plotImpurity=3, - plotRotation=True, - lastRho=ps[0].plasma["rho"][-1, -1].item(), - ) - -axsGrads_extra = [ - axs[0], - axs[5], - axs[1], - axs[6], - axs[2], - axs[7], - axs[3], - axs[8], - axs[4], - axs[9], -] -for i, p in enumerate(profs): - p.plotGradients( - axsGrads_extra, - color=colors[i], - plotImpurity=3, - plotRotation=True, - lastRho=ps[0].plasma["rho"][-1, -1].item(), - lw=0.5, - ms=0, - ) - -fn.show() diff --git a/src/mitim_modules/portals/exe/compareBest.py b/src/mitim_modules/portals/exe/compareBest.py deleted file mode 100644 index 9793d070..00000000 --- a/src/mitim_modules/portals/exe/compareBest.py +++ /dev/null @@ -1,59 +0,0 @@ -""" -This plots final profiles -""" - -import sys, argparse, copy -import matplotlib.pyplot as plt -import dill as pickle_dill -from mitim_tools.misc_tools import GRAPHICStools -from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.opt_tools import STRATEGYtools - - -parser = argparse.ArgumentParser() -parser.add_argument("folders", type=str, nargs="*") -args = parser.parse_args() -folders = args.folders - - -p = [] -for folderWork in folders: - opt_fun = STRATEGYtools.FUNmain(folderWork) - opt_fun.read_optimization_results(analysis_level=4) - - with open(opt_fun.prfs_model.MITIMextra, "rb") as handle: - MITIMextra_dict = pickle_dill.load(handle) - - MITIMextra_dict[opt_fun.res.best_absolute_index]["tgyro"].profiles - p.append(MITIMextra_dict[opt_fun.res.best_absolute_index]["tgyro"].profiles) - -labels = [] -for i in folders: - labels.append(i.split("/")[-1]) - -plt.ion() -fig, axs = plt.subplots(ncols=3, nrows=2, figsize=(15, 9)) - -colors = GRAPHICStools.listColors() - -# porig = MITIMextra_dict[0]['tgyro'].profiles -# porig.plotGradients(axs.flatten(),color='k',label='original',lastRho=MITIMextra_dict[opt_fun.res.best_absolute_index]['tgyro'].rhosToSimulate[-1]) - -for p0, color, label in zip(p, colors, labels): - p0.plotGradients( - axs.flatten(), - color=color, - label=label, - lastRho=MITIMextra_dict[opt_fun.res.best_absolute_index][ - "tgyro" - ].rhosToSimulate[-1], - ) - -for ax in axs.flatten(): - for i in MITIMextra_dict[opt_fun.res.best_absolute_index]["tgyro"].rhosToSimulate: - ax.axvline(x=i, ls="--", c="k", lw=0.2) - -axs[0, 0].legend() - -for ax in axs.flatten(): - ax.set_ylim(bottom=0) diff --git a/src/mitim_modules/portals/exe/comparePORTALS.py b/src/mitim_modules/portals/exe/comparePORTALS.py deleted file mode 100644 index ef2b3bf3..00000000 --- a/src/mitim_modules/portals/exe/comparePORTALS.py +++ /dev/null @@ -1,152 +0,0 @@ -""" -To compare mitim and TGYRO results -""" - -import sys -import matplotlib.pyplot as plt -import numpy as np -from IPython import embed -from mitim_tools.gacode_tools import TGYROtools -from mitim_tools.misc_tools import GRAPHICStools -from mitim_tools.opt_tools import STRATEGYtools -from mitim_tools.opt_tools.aux import BOgraphics -import dill as pickle_dill - -# run ~/MITIM/mitim_opt/mitim/exe/comparemitim.py ~/runs_portals/set6/run38/ ~/runs_portals/set6/run38/TGYRO_test/tgyro1_1/ -# run comparemitim.py ~/PRF/mitim_dev/run1/ ~/PRF/mitim_dev/tgyro1_1/ mfews01.psfc.mit.edu-9224:runs_portals/set6/run38/ - - -def compareSolvers( - foldermitim, foldersTGYRO, labelsTGYRO, foldermitim_remote=None, fig=None -): - plt.ion() - if fig is None: - fig = plt.figure(figsize=(6, 8)) - grid = plt.GridSpec(1, 2, hspace=0.3, wspace=0.3) - ax0 = fig.add_subplot(grid[0, 0]) - ax1 = fig.add_subplot(grid[0, 1], sharex=ax0) - - c = GRAPHICStools.listColors() - - # ---------------------------------------------------------------------------------------------------- - # ----- Plot TGYRO - # ---------------------------------------------------------------------------------------------------- - - resTGYRO = TGYROtools.TGYROoutput(foldersTGYRO[0]) - i = 0 - for folderTGYRO, labelTGYRO in zip(foldersTGYRO, labelsTGYRO): - # ----- Read TGYRO - try: - resTGYRO = TGYROtools.TGYROoutput(folderTGYRO) - except: - continue - - iterations_array = resTGYRO.calls_solver / (resTGYRO.rho.shape[1] - 1) - residual_array = resTGYRO.residual_manual_real - ax0.plot( - iterations_array, - residual_array, - "*-", - color=c[i], - label=labelTGYRO, - lw=0.5, - markersize=3, - ) - i += 1 - - # ---------------------------------------------------------------------------------------------------- - # ----- Plot mitim - # ---------------------------------------------------------------------------------------------------- - - opt_fun = STRATEGYtools.FUNmain(foldermitim) - opt_fun.read_optimization_results( - plotYN=False, folderRemote=foldermitim_remote, analysis_level=4 - ) - self_complete = opt_fun.plot_optimization_results(analysis_level=4, plotYN=False) - with open(self_complete.MITIMextra, "rb") as handle: - MITIMextra_dict = pickle_dill.load(handle) - - iterationsMultiplier = 1.0 # resTGYRO.rho.shape[1]-1 - iterationsOffset = 1 # resTGYRO.rho.shape[1]-1 - - # Residual based on TGYRO definition - resmitim = [] - iterations_arrayP = [] - for i in MITIMextra_dict: - try: - resmitim.append( - MITIMextra_dict[i]["tgyro"] - .results[tuple(MITIMextra_dict[i]["tgyro"].results.keys())[0]] - .residual_manual_real[0] - ) - iterations_arrayP.append(i) - except: - break - resmitim = np.array(resmitim) - iterations_arrayP = np.array(iterations_arrayP) - - ax0.plot( - iterations_arrayP * iterationsMultiplier + iterationsOffset, - resmitim, - "s-", - color="g", - label="MITIM", - lw=1.0, - markersize=5, - ) - - # ---------------------------------------------------------------------------------------------------- - # Decor - # ---------------------------------------------------------------------------------------------------- - - plotPercent = 1e-2 - - ax = ax0 - - ax.legend(loc="upper right", prop={"size": 8}) - - if iterationsMultiplier > 1.0: - ax.set_xlabel("Calls to TGLF model") - else: - ax.set_xlabel("Profile evaluations (Calls to TGLF per radius)") - if plotPercent > 0: - GRAPHICStools.drawLineWithTxt( - ax, - resmitim[0] * plotPercent, - label="$10^{-2}$x Original", - orientation="horizontal", - color="k", - lw=0.5, - ls="--", - alpha=1.0, - fontsize=8, - fromtop=0.8, - fontweight="normal", - verticalalignment="bottom", - horizontalalignment="left", - separation=0, - ) - - ax.set_ylabel("$\\widehat{L_1}$ Residue ($MW/m^2$), TGYRO definition") - - ax.set_yscale("log") - ax.set_xlim(left=0.0) - GRAPHICStools.addDenseAxis(ax) - - return ax0, ax1 - - -if __name__ == "__main__": - foldermitim = sys.argv[1] - folderTGYRO = sys.argv[2] - try: - foldermitim_remote = sys.argv[3] - except: - foldermitim_remote = None - - ax0, ax1 = compareSolvers( - foldermitim, - [folderTGYRO], - labelsTGYRO=["TGYRO"], - foldermitim_remote=foldermitim_remote, - ) diff --git a/src/mitim_modules/portals/exe/readMetrics.py b/src/mitim_modules/portals/exe/readMetrics.py index 541b2f9f..6c250f80 100644 --- a/src/mitim_modules/portals/exe/readMetrics.py +++ b/src/mitim_modules/portals/exe/readMetrics.py @@ -1,9 +1,8 @@ -import sys, argparse -import dill as pickle_dill +import argparse import matplotlib.pyplot as plt from mitim_tools.misc_tools import IOtools -from mitim_tools.opt_tools import STRATEGYtools -from mitim_modules.portals.aux import PORTALSplot +from mitim_modules.portals.aux import PORTALSanalysis +from IPython import embed """ This script is to plot only the convergence figure, not the rest of surrogates that takes long. @@ -11,37 +10,37 @@ """ parser = argparse.ArgumentParser() -parser.add_argument("--folders", required=True, type=str, nargs="*") -parser.add_argument("--remote", type=str, required=False, default=None) +parser.add_argument("folders", type=str, nargs="*") +parser.add_argument("--remote", "-r", type=str, required=False, default=None) parser.add_argument( "--max", type=int, required=False, default=None ) # Define max bounds of fluxes based on this one, like 0, -1 or None(best) parser.add_argument("--index_extra", type=int, required=False, default=None) parser.add_argument( - "--all", type=bool, required=False, default=False + "--all", required=False, default=False, action="store_true" ) # Plot all fluxes? parser.add_argument( "--file", type=str, required=False, default=None ) # File to save .eps +parser.add_argument( + "--complete", "-c", required=False, default=False, action="store_true" +) + args = parser.parse_args() folders = args.folders - -size = 8 -plt.rc("font", family="serif", serif="Times", size=size) -plt.rc("xtick.minor", size=size) -plt.close("all") - +portals_total = [] for folderWork in folders: folderRemote_reduced = args.remote file = args.file indexToMaximize = args.max index_extra = args.index_extra plotAllFluxes = args.all + complete = args.complete folderRemote = ( f"{folderRemote_reduced}/{IOtools.reducePathLevel(folderWork)[-1]}/" @@ -49,37 +48,62 @@ else None ) - # Read standard OPT - opt_fun = STRATEGYtools.FUNmain(folderWork) - opt_fun.read_optimization_results( - analysis_level=4, plotYN=False, folderRemote=folderRemote + # Read PORTALS + portals = PORTALSanalysis.PORTALSanalyzer.from_folder( + folderWork, folderRemote=folderRemote ) - # Analyze mitim results - self_complete = opt_fun.plot_optimization_results(analysis_level=4, plotYN=False) - - # Interpret results - with open(self_complete.MITIMextra, "rb") as handle: - MITIMextra_dict = pickle_dill.load(handle) - portals_plot = PORTALSplot.PORTALSresults( - self_complete.folder, - opt_fun.prfs_model, - opt_fun.res, - MITIMextra_dict=MITIMextra_dict, - indecesPlot=[opt_fun.res.best_absolute_index, 0, index_extra], - ) + portals_total.append(portals) - # Plot - plt.ion() - fig = plt.figure(figsize=(18, 9)) - PORTALSplot.plotConvergencePORTALS( - portals_plot, - fig=fig, - indexToMaximize=indexToMaximize, - plotAllFluxes=plotAllFluxes, - ) +# PLOTTING - # Save +if not complete: + size = 8 + plt.rc("font", family="serif", serif="Times", size=size) + plt.rc("xtick.minor", size=size) +plt.close("all") + +is_any_ini = False +for i in range(len(folders)): + is_any_ini = is_any_ini or isinstance( + portals_total[i], PORTALSanalysis.PORTALSinitializer + ) -if file is not None: - plt.savefig(file, transparent=True, dpi=300) +requiresFN = (len(folders) > 1) or complete or is_any_ini + +if requiresFN: + from mitim_tools.misc_tools.GUItools import FigureNotebook + + fn = FigureNotebook("PORTALS", geometry="1600x1000") +else: + fn = None + +for i in range(len(folders)): + lab = f"{IOtools.reducePathLevel(folders[i])[-1]}" + + portals_total[i].fn = fn + + # Plot metrics + if (not complete) or ( + isinstance(portals_total[i], PORTALSanalysis.PORTALSinitializer) + ): + if isinstance(portals_total[i], PORTALSanalysis.PORTALSinitializer): + fig = None + elif requiresFN: + fig = fn.add_figure(label=lab) + else: + plt.ion() + fig = plt.figure(figsize=(15, 8)) + + portals_total[i].plotMetrics( + fig=fig, + indexToMaximize=indexToMaximize, + plotAllFluxes=plotAllFluxes, + index_extra=index_extra, + file_save=file if len(folders) == 1 else None, + extra_lab=lab, + ) + + # Plot PORTALS + else: + portals_total[i].plotPORTALS() diff --git a/src/mitim_modules/portals/exe/runTGLF.py b/src/mitim_modules/portals/exe/runTGLF.py new file mode 100644 index 00000000..b4436644 --- /dev/null +++ b/src/mitim_modules/portals/exe/runTGLF.py @@ -0,0 +1,81 @@ +import argparse +import numpy as np +from mitim_modules.portals.aux import PORTALSanalysis + +""" +This script is useful to understand why PORTALS may fail at reproducing TGLF fluxes. You can select the iteration +to use as base case, and scan parameters to see how TGLF behaves (understand if it has discontinuities) + e.g. + runTGLF.py --folder run11/ --ev 5 --pos 0 2 --params RLTS_2 RLTS_1 --wf 0.2 1.0 --var 0.05 + +Notes: + - wf runs scan with waveform too (slightly more expensive, as it will require 1 extra sim per run, but cheaper) + - drives will simply run the drives analysis, ignoring the params option +""" + +# --- Inputs + +parser = argparse.ArgumentParser() +parser.add_argument("--folder", required=True, type=str) +parser.add_argument("--ev", type=int, required=False, default=None) +parser.add_argument("--pos", type=int, required=False, default=[0.5], nargs="*") +parser.add_argument("--params", type=str, required=False, default=["RLTS_2"], nargs="*") +parser.add_argument("--wf", type=float, required=False, default=None, nargs="*") +parser.add_argument( + "--var", type=float, required=False, default=0.05 +) # Variation in inputs (10% default) +parser.add_argument( + "--restart", "-r", required=False, default=False, action="store_true" +) +parser.add_argument("--drives", required=False, default=False, action="store_true") + +args = parser.parse_args() +folder = args.folder +ev = args.ev +params = args.params +pos = args.pos +wf = args.wf +var = args.var +restart = args.restart +drives = args.drives + +# --- Workflow + +portals = PORTALSanalysis.PORTALSanalyzer.from_folder(folder) +tglf, TGLFsettings, extraOptions = portals.extractTGLF(positions=pos, evaluation=ev) + +if not drives: + varUpDown = np.linspace(1.0 - var, 1.0 + var, 10) + + labels = [] + for param in params: + tglf.runScan( + subFolderTGLF="scan", + variable=param, + varUpDown=varUpDown, + TGLFsettings=TGLFsettings, + extraOptions=extraOptions, + restart=restart, + runWaveForms=wf, + ) + + tglf.readScan(label=f"scan_{param}", variable=param) + labels.append(f"scan_{param}") + + # --- Extra TGLF plotting + + tglf.plotScan(labels=labels, variableLabel="X", relativeX=True) + +else: + tglf.runScanTurbulenceDrives( + subFolderTGLF="turb", + resolutionPoints=5, + variation=var, + variablesDrives=["RLTS_1", "RLTS_2", "RLNS_1", "XNUE", "TAUS_2", "BETAE"], + TGLFsettings=TGLFsettings, + extraOptions=extraOptions, + restart=restart, + runWaveForms=wf, + ) + + tglf.plotScanTurbulenceDrives(label="turb") diff --git a/src/mitim_modules/portals/exe/runTGLFdrivesfromPORTALS.py b/src/mitim_modules/portals/exe/runTGLFdrivesfromPORTALS.py index 9c7247d3..ec8ebd93 100644 --- a/src/mitim_modules/portals/exe/runTGLFdrivesfromPORTALS.py +++ b/src/mitim_modules/portals/exe/runTGLFdrivesfromPORTALS.py @@ -1,17 +1,11 @@ -import os, argparse, copy -import dill as pickle_dill -import numpy as np -from IPython import embed -from mitim_tools.misc_tools import IOtools -from mitim_tools.gacode_tools import TGLFtools -from mitim_tools.opt_tools import STRATEGYtools +import argparse +from mitim_modules.portals.aux import PORTALSanalysis """ - -This script is useful to understand why MITIM may fail at reproducing TGLF fluxes. You can select the iteration -to use as base case to see how TGLF behaves (understand if it has discontinuities) +This script is useful to understand why surrogates may fail at reproducing TGLF fluxes. +You can select the iteration to use as base case to see how TGLF behaves (if it has discontinuities) e.g. - runTGLFdrivesfrommitim.py --folder run11/ --ev -1 --pos 0 2 --var 0.05 --wf 0.2 1.0 + runTGLFdrivesfrommitim.py --folder run11/ --ev 5 --pos 0 2 --var 0.05 --wf 0.2 1.0 Notes: - wf runs scan with waveform too (slightly more expensive, as it will require 1 extra sim per run, but cheaper) @@ -21,12 +15,13 @@ parser = argparse.ArgumentParser() parser.add_argument("--folder", required=True, type=str) -parser.add_argument("--ev", type=int, required=False, default=-1) -parser.add_argument("--pos", type=int, required=False, default=[0.5], nargs="*") +parser.add_argument("--ev", type=int, required=False, default=None) +parser.add_argument("--pos", type=int, required=False, default=[0], nargs="*") parser.add_argument("--wf", type=float, required=False, default=None, nargs="*") parser.add_argument( "--var", type=float, required=False, default=0.01 ) # Variation in inputs (1% default) +parser.add_argument("-r", required=False, default=False, action="store_true") args = parser.parse_args() folder = args.folder @@ -34,59 +29,22 @@ pos = args.pos wf = args.wf var = args.var +restart = args.r # --- Workflow -workingFolder = f"{folder}/turb_drives/" - -if not os.path.exists(workingFolder): - os.system(f"mkdir {workingFolder}") - -opt_fun = STRATEGYtools.FUNmain(folder) -opt_fun.read_optimization_results(analysis_level=4, plotYN=False) - -rhos = [] -for i in pos: - rhos.append(opt_fun.prfs_model.mainFunction.TGYROparameters["RhoLocations"][i]) - -if ev < 0: - ev = opt_fun.res.best_absolute_index - -inputgacode = f"{folder}/Execution/Evaluation.{ev}/tgyro_complete/input.gacode" - -if not os.path.exists(inputgacode): - print( - "input.gacode could not be read in Execution folder, likely because of a remote run. Resorting to MITIMextra" - ) - self_complete = opt_fun.plot_optimization_results(analysis_level=4, plotYN=False) - with open(self_complete.MITIMextra, "rb") as handle: - MITIMextra_dict = pickle_dill.load(handle) - - inputgacode = f"{workingFolder}/input.gacode.start" - MITIMextra_dict[ev]["tgyro"].profiles.writeCurrentStatus(file=inputgacode) - - -tglf = TGLFtools.TGLF(rhos=rhos) -cdf = tglf.prep(workingFolder, restart=False, inputgacode=inputgacode) - -extraOptions = opt_fun.prfs_model.mainFunction.TGLFparameters["extraOptionsTGLF"] - -extraOptions["NMODES"] = 2 - -name = f"turb" +portals = PORTALSanalysis.PORTALSanalyzer.from_folder(folder) +tglf, TGLFsettings, extraOptions = portals.extractTGLF(positions=pos, evaluation=ev) tglf.runScanTurbulenceDrives( - subFolderTGLF=name, + subFolderTGLF="turb", resolutionPoints=5, variation=var, variablesDrives=["RLTS_1", "RLTS_2", "RLNS_1", "XNUE", "TAUS_2", "BETAE"], - TGLFsettings=opt_fun.prfs_model.mainFunction.TGLFparameters["TGLFsettings"], + TGLFsettings=TGLFsettings, extraOptions=extraOptions, - restart=False, + restart=restart, runWaveForms=wf, ) - -# --- Extra TGLF plotting - -tglf.plotScanTurbulenceDrives(label=name) +tglf.plotScanTurbulenceDrives(label="turb") diff --git a/src/mitim_modules/portals/exe/runTGLFscanfromPORTALS.py b/src/mitim_modules/portals/exe/runTGLFscanfromPORTALS.py deleted file mode 100644 index e9f94a88..00000000 --- a/src/mitim_modules/portals/exe/runTGLFscanfromPORTALS.py +++ /dev/null @@ -1,76 +0,0 @@ -import os, argparse, copy -import numpy as np -from IPython import embed -from mitim_tools.misc_tools import IOtools -from mitim_tools.gacode_tools import TGLFtools -from mitim_tools.opt_tools import STRATEGYtools - -""" - -This script is useful to understand why MITIM may fail at reproducing TGLF fluxes. You can select the iteration -to use as base case, and scan parameters to see how TGLF behaves (understand if it has discontinuities) - e.g. - runTGLFscanfrommitim.py --folder run11/ --ev -1 --pos 0 2 --params RLTS_2 RLTS_1 --wf 0.2 1.0 - -Notes: - - wf runs scan with waveform too (slightly more expensive, as it will require 1 extra sim per run, but cheaper) -""" - -# --- Inputs - -parser = argparse.ArgumentParser() -parser.add_argument("--folder", required=True, type=str) -parser.add_argument("--ev", type=int, required=False, default=-1) -parser.add_argument("--pos", type=int, required=False, default=[0.5], nargs="*") -parser.add_argument("--params", type=str, required=False, default=["RLTS_2"], nargs="*") -parser.add_argument("--wf", type=float, required=False, default=None, nargs="*") - -args = parser.parse_args() -folder = args.folder -ev = args.ev -params = args.params -pos = args.pos -wf = args.wf - -# --- Workflow - -opt_fun = STRATEGYtools.FUNmain(folder) -opt_fun.read_optimization_results(analysis_level=4, plotYN=False) - -rhos = [] -for i in pos: - rhos.append(opt_fun.prfs_model.mainFunction.TGYROparameters["RhoLocations"][i]) - -if ev < 0: - ev = opt_fun.res.best_absolute_index - -inputgacode = f"{folder}/Execution/Evaluation.{ev}/tgyro_complete/input.gacode" -workingFolder = f"{folder}/scans{ev}/" -varUpDown = np.linspace(0.9, 1.1, 10) - -tglf = TGLFtools.TGLF(rhos=rhos) -cdf = tglf.prep(workingFolder, restart=False, inputgacode=inputgacode) - -labels = [] -for param in params: - extraOptions = copy.deepcopy( - opt_fun.prfs_model.mainFunction.TGLFparameters["extraOptionsTGLF"] - ) - - tglf.runScan( - subFolderTGLF=f"scan_{param}/", - variable=param, - varUpDown=varUpDown, - TGLFsettings=opt_fun.prfs_model.mainFunction.TGLFparameters["TGLFsettings"], - extraOptions=extraOptions, - restart=False, - runWaveForms=wf, - ) - - tglf.readScan(label=f"scan_{param}", variable=param) - labels.append(f"scan_{param}") - - -# --- Extra TGLF plotting - -tglf.plotScan(labels=labels, variableLabel="X", relativeX=True) diff --git a/src/mitim_modules/portals/exe/runTGYROfromPORTALS.py b/src/mitim_modules/portals/exe/runTGYRO.py similarity index 76% rename from src/mitim_modules/portals/exe/runTGYROfromPORTALS.py rename to src/mitim_modules/portals/exe/runTGYRO.py index 6eaffc26..78a963bd 100644 --- a/src/mitim_modules/portals/exe/runTGYROfromPORTALS.py +++ b/src/mitim_modules/portals/exe/runTGYRO.py @@ -1,15 +1,16 @@ -import sys, os, argparse, copy, time +import argparse +import copy +import time import numpy as np -from IPython import embed from mitim_tools.misc_tools import IOtools, FARMINGtools -from mitim_tools.opt_tools import STRATEGYtools -from mitim_tools.gacode_tools import TGYROtools, PROFILEStools -from mitim_modules.portals import PORTALSmain +from mitim_modules.portals.aux import PORTALSanalysis +from mitim_modules.portals.exe.comparePORTALS import compareSolvers +from IPython import embed """ This script will run TGYRO using the settings used for mitim. It will do it in a subfolder of the mitim run. -e.g.: runTGYROfrommitim.py --folder run1/ --seeds 5 --methods 1 6 +e.g.: runTGYRO.py --folder run1/ --seeds 5 --methods 1 6 It will run 16 in parallel @@ -41,27 +42,12 @@ # Preparation # ------------------------------------------------------------------------------------------ -# Folder to work on -folder = IOtools.expandPath(folderO + "/tgyro_std_analysis/") -if not os.path.exists(folder): - IOtools.askNewFolder(folder) - - -# # Read mitim results -# opt_fun = STRATEGYtools.FUNmain(folderO) -# opt_fun.read_optimization_results(plotYN=False,analysis_level=1) -# portals = opt_fun.prfs_model.mainFunction - -# # Prepare TGYRO -# readFile = f'{folder}/input.gacode_copy_initialization' -# with open(readFile,'w') as f: f.writelines(portals.file_in_lines_initial_input_gacode) -# profiles = PROFILEStools.PROFILES_GACODE(readFile,calculateDerived=True) +portals = PORTALSanalysis.PORTALSanalyzer.from_folder(folderO) -opt_fun, portals, self_complete, MITIMextra_dict = PORTALSmain.readFullmitim(folderO) -profiles = MITIMextra_dict[0]["tgyro"].profiles - -tgyro = TGYROtools.TGYRO() -tgyro.prep(folder, profilesclass_custom=profiles, restart=restart, forceIfRestart=True) +folder = IOtools.expandPath(folderO + "/tgyro_std_analysis/") +tgyro, rhos, PredictionSet, TGLFsettings, extraOptionsTGLF = portals.extractTGYRO( + folder=folder, restart=restart, evaluation=0 +) # ------------------------------------------------------------------------------------------ # Run TGYRO (function prep) @@ -70,7 +56,7 @@ def run_tgyro_parallel(Params, cont): time.sleep( - cont * 5 + cont * Params["seconds_sleep"] ) # Wait a bit to avoid bandwidth problems of ssh or scp connections # Grab data @@ -79,7 +65,6 @@ def run_tgyro_parallel(Params, cont): method = Params["method"] iterations = Params["iterations"] i = Params["scan"][cont] - seconds_sleep = Params["seconds_sleep"] if param == 1: name = f"m{method}_r{i:.3f}" @@ -97,19 +82,15 @@ def run_tgyro_parallel(Params, cont): iterations=iterations, restart=restartTGYRO, forceIfRestart=True, - special_radii=portals.TGYROparameters["RhoLocations"], - PredictionSet=[ - int("te" in portals.TGYROparameters["ProfilesPredicted"]), - int("ti" in portals.TGYROparameters["ProfilesPredicted"]), - int("ne" in portals.TGYROparameters["ProfilesPredicted"]), - ], + special_radii=rhos, + PredictionSet=PredictionSet, minutesJob=120, launchSlurm=True, - TGLFsettings=portals.TGLFparameters["TGLFsettings"], - extraOptionsTGLF=portals.TGLFparameters["extraOptionsTGLF"], + TGLFsettings=TGLFsettings, + extraOptionsTGLF=extraOptionsTGLF, TGYRO_physics_options={ "TargetType": 3, - "quasineutrality": [1, 2] if profiles.DTplasmaBool else [1], + "quasineutrality": [1, 2] if tgyro_here.profiles.DTplasmaBool else [1], }, ) @@ -122,7 +103,7 @@ def run_tgyro_parallel(Params, cont): for param in params: if method == 6: scan = np.linspace(0.01, 0.2, seeds) - iterations = 100 + iterations = 2 elif method == 1: if param == 1: scan = np.linspace(1, 3, seeds) @@ -200,12 +181,8 @@ def run_tgyro_parallel(Params, cont): # Plot # ------------------------------------------------------------------------------------------ -tgyro.plotRun(labels=names, doNotShow=True) - -from mitim_modules.portals.exe.comparemitim import compareSolvers +tgyro.plot(labels=names, doNotShow=True) compareSolvers( folderO, folders, nice_names, fig=tgyro.fn.add_figure(label="mitim COMPARISON") ) - -tgyro.fn.show() diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 16e51498..d556f516 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -1,19 +1,21 @@ -import copy, torch, datetime +import copy +import torch +import datetime import matplotlib.pyplot as plt import numpy as np -from IPython import embed import dill as pickle from mitim_tools.misc_tools import PLASMAtools, IOtools from mitim_modules.powertorch.aux import TRANSFORMtools, POWERplot from mitim_modules.powertorch.iteration import ITtools from mitim_modules.powertorch.physics import TARGETStools, TRANSPORTtools, CALCtools from mitim_tools.misc_tools.IOtools import printMsg as print +from IPython import embed UseCUDAifAvailable = True def read_saved_state(file): - print(f" - Reading state file {file}") + print(f"\t- Reading state file {IOtools.clipstr(file)}") with open(file, "rb") as handle: state = pickle.load(handle) return state @@ -192,7 +194,6 @@ def insertProfiles( insertPowers=insertPowers, rederive=rederive_profiles, ProfilesPredicted=self.ProfilesPredicted, - impurityPosition=self.impurityPosition, ) if writeFile is not None: @@ -354,17 +355,16 @@ def plot(self, axs=None, axsRes=None, figs=None, c="r", label=""): if axs is None: from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - fn = FigureNotebook(0, "PowerState", geometry="1800x900") - figMain = fn.add_figure(label="PowerState") - - figProf_1 = fn.add_figure(label="Profiles") - figProf_2 = fn.add_figure(label="Powers") - figProf_3 = fn.add_figure(label="Geometry") - figProf_4 = fn.add_figure(label="Gradients") - figFlows = fn.add_figure(label="Flows") - figProf_6 = fn.add_figure(label="Other") - fig7 = fn.add_figure(label="Impurities") + self.fn = FigureNotebook("PowerState", geometry="1800x900") + figMain = self.fn.add_figure(label="PowerState") + + figProf_1 = self.fn.add_figure(label="Profiles") + figProf_2 = self.fn.add_figure(label="Powers") + figProf_3 = self.fn.add_figure(label="Geometry") + figProf_4 = self.fn.add_figure(label="Gradients") + figFlows = self.fn.add_figure(label="Flows") + figProf_6 = self.fn.add_figure(label="Other") + fig7 = self.fn.add_figure(label="Impurities") figs = figProf_1, figProf_2, figProf_3, figProf_4, figFlows, figProf_6, fig7 grid = plt.GridSpec(4, 6, hspace=0.3, wspace=0.3) @@ -394,15 +394,8 @@ def plot(self, axs=None, axsRes=None, figs=None, c="r", label=""): axsRes = figMain.add_subplot(grid[:, 0]) - provided = False - else: - provided = True - POWERplot.plot(self, axs, axsRes, figs, c=c, label=label) - if not provided: - fn.show() - # ------------------------------------------------------------------ # Main tools # ------------------------------------------------------------------ @@ -648,6 +641,7 @@ def calculateProfileFunctions(self, calculateRotationStuff=True, mref_u=2.01355) self.plasma["Ggb"], self.plasma["Pgb"], self.plasma["Sgb"], + self.plasma["Qgb_convection"], ) = PLASMAtools.gyrobohmUnits( self.plasma["te"], self.plasma["ne"] * 1e-1, @@ -856,6 +850,11 @@ def calculateTargets(self): self.plasma["te"], self.plasma["CZ"] ) # MW/m^2 + if ( + "forceZeroParticleFlux" in self.TransportOptions["ModelOptions"] + ) and self.TransportOptions["ModelOptions"]["forceZeroParticleFlux"]: + self.plasma["Ce"] = self.plasma["Ce"] * 0 + """ ************************************************************************************************** GB Normalized diff --git a/src/mitim_modules/powertorch/aux/POWERplot.py b/src/mitim_modules/powertorch/aux/POWERplot.py index 00bc4198..a591b203 100644 --- a/src/mitim_modules/powertorch/aux/POWERplot.py +++ b/src/mitim_modules/powertorch/aux/POWERplot.py @@ -8,7 +8,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label=""): profiles_new = self.insertProfiles(self.profiles, insertPowers=True) if figs is not None: - PROFILEStools.plotAll([self.profiles, profiles_new], figs=figs) + fn = PROFILEStools.plotAll([self.profiles, profiles_new], figs=figs) # plotPlasma(self,self.FluxMatch_plasma_orig,axs,color='b') plotPlasma(self, self.plasma, axs, color=c, label=label) diff --git a/src/mitim_modules/powertorch/aux/TRANSFORMtools.py b/src/mitim_modules/powertorch/aux/TRANSFORMtools.py index 28e026e3..4b5640a1 100644 --- a/src/mitim_modules/powertorch/aux/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/aux/TRANSFORMtools.py @@ -1,16 +1,16 @@ -import copy, torch +import copy +import torch import numpy as np -from IPython import embed from mitim_modules.powertorch.aux import PARAMtools from mitim_modules.powertorch.physics import TARGETStools from mitim_tools.misc_tools import IOtools from mitim_tools.misc_tools.IOtools import printMsg as print from mitim_tools.misc_tools.CONFIGread import read_verbose_level +from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpFunction +from IPython import embed verbose_level = read_verbose_level() -from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpFunction - """ -------------------------------------------------------------------------------------------------------------------------------------------- FUNCTIONS THAT DEAL WITH THE CONNECTION BETWEEN POWERSTATE AND PROFILES_GACODE (READING AND WRITING) @@ -26,7 +26,6 @@ def fromPowerToGacode( insertPowers=True, rederive=True, ProfilesPredicted=["te", "ti", "ne"], - impurityPosition=1, ): """ Notes: @@ -131,7 +130,7 @@ def fromPowerToGacode( # ------------------------------------------------------------------------------------------ if "w0" in ProfilesPredicted: - print(f"\t- Changing w0") + print("\t- Changing w0") factor_mult = 1 / factorMult_w0(self) x, y = self.deparametrizers["w0"]( self.plasma["roa"][PositionInBatch, :], @@ -493,6 +492,20 @@ def defineMovingGradients(self, profiles): ) self.plasma["aLw0"] = aLy_coarse[:-1, 1] + # ---------------------------------------------------------------------------------------------------- + # Check that it's not completely zero + # ---------------------------------------------------------------------------------------------------- + + for key in ["aLte", "aLti", "aLne", "aLnZ", "aLw0"]: + if key[2:] in self.ProfilesPredicted: + if self.plasma[key].sum() == 0.0: + addT = 1e-15 + print( + f"\t- All values of {key} detected to be zero, to avoid NaNs, inserting {addT} at the edge", + typeMsg="w", + ) + self.plasma[key][..., -1] += addT + def factorMult_w0(self): """ diff --git a/src/mitim_modules/powertorch/physics/TRANSPORTtools.py b/src/mitim_modules/powertorch/physics/TRANSPORTtools.py index 990cfa5f..113d7b3e 100644 --- a/src/mitim_modules/powertorch/physics/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/physics/TRANSPORTtools.py @@ -1,12 +1,13 @@ -import copy, os, torch +import copy +import os +import torch import numpy as np -from IPython import embed from mitim_tools.misc_tools import PLASMAtools, IOtools from mitim_tools.gacode_tools import TGYROtools from mitim_modules.portals.aux import PORTALScgyro from mitim_modules.powertorch.aux import TRANSFORMtools - from mitim_tools.misc_tools.IOtools import printMsg as print +from IPython import embed # ------------------------------------------------------------------ # SIMPLE @@ -103,7 +104,6 @@ def tgyro_model( includeFast = ModelOptions["includeFastInQi"] impurityPosition = ModelOptions["impurityPosition"] useConvectiveFluxes = ModelOptions["useConvectiveFluxes"] - ProfilesPredicted = TGYROparameters["ProfilesPredicted"] UseFineGridTargets = ModelOptions["UseFineGridTargets"] launchTGYROviaSlurm = ModelOptions.get("launchTGYROviaSlurm", False) @@ -135,6 +135,8 @@ def tgyro_model( applyCorrections=TGYROparameters["applyCorrections"], ) + # VGEN? + # copy for future modifications self.file_profs_mod = f"{self.file_profs}_modified" os.system(f"cp {self.file_profs} {self.file_profs_mod}") @@ -157,7 +159,7 @@ def tgyro_model( print("\t- Launching TGYRO evaluation as a terminal job") self.tgyro_current.run( - subFolderTGYRO=f"tglf_neo_original/", + subFolderTGYRO="tglf_neo_original/", restart=restart, forceIfRestart=True, special_radii=RadiisToRun, @@ -171,7 +173,7 @@ def tgyro_model( forcedName=name, ) - self.tgyro_current.read(label=f"tglf_neo_original") + self.tgyro_current.read(label="tglf_neo_original") # Copy one with evaluated targets self.file_profs_targets = f"{self.tgyro_current.FolderTGYRO}/input.gacode.new" @@ -189,7 +191,7 @@ def tgyro_model( # Add errors and merge fluxes as we would do if this was a CGYRO run curateTGYROfiles( - self.tgyro_current.results[f"tglf_neo_original"], + self.tgyro_current.results["tglf_neo_original"], f"{FolderEvaluation_TGYRO}/tglf_neo/", percentError, impurityPosition=impurityPosition, @@ -198,7 +200,7 @@ def tgyro_model( # Read again to capture errors self.tgyro_current.read( - label=f"tglf_neo", folder=f"{FolderEvaluation_TGYRO}/tglf_neo/" + label="tglf_neo", folder=f"{FolderEvaluation_TGYRO}/tglf_neo/" ) labels_results.append("tglf_neo") @@ -210,8 +212,6 @@ def tgyro_model( useConvectiveFluxes=useConvectiveFluxes, includeFast=includeFast, impurityPosition=impurityPosition, - ProfilesPredicted=ProfilesPredicted, - provideTurbulentExchange=provideTurbulentExchange, UseFineGridTargets=UseFineGridTargets, OriginalFimp=OriginalFimp, forceZeroParticleFlux=forceZeroParticleFlux, @@ -267,15 +267,13 @@ def tgyro_model( useConvectiveFluxes=useConvectiveFluxes, includeFast=includeFast, impurityPosition=impurityPosition, - ProfilesPredicted=ProfilesPredicted, - provideTurbulentExchange=provideTurbulentExchange, UseFineGridTargets=UseFineGridTargets, OriginalFimp=OriginalFimp, forceZeroParticleFlux=forceZeroParticleFlux, dfT=self.dfT, ) - print(f"\t- Checking model modifications:") + print("\t- Checking model modifications:") for r in ["Qe_turb", "Qi_turb", "Ge_turb", "GZ_turb", "Mt_turb", "PexchTurb"]: print( f"\t\t{r}(tglf) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(portals_variables_orig[r][0][1:],portals_variables_orig[r+'_stds'][0][1:]) ])}" diff --git a/src/mitim_modules/vitals/VITALSmain.py b/src/mitim_modules/vitals/VITALSmain.py index 67365778..f6680e24 100644 --- a/src/mitim_modules/vitals/VITALSmain.py +++ b/src/mitim_modules/vitals/VITALSmain.py @@ -373,7 +373,7 @@ def analyze_results( # Plot if plotYN: - self.tglf_final.plotRun( + self.tglf_final.plot( labels=["Original", "VITALS"], fn=fn, extratitle="TGLF- " ) diff --git a/src/mitim_tools/__init__.py b/src/mitim_tools/__init__.py index 75977e6f..5becc17c 100644 --- a/src/mitim_tools/__init__.py +++ b/src/mitim_tools/__init__.py @@ -1 +1 @@ -__version__ = '1.0.0' \ No newline at end of file +__version__ = "1.0.0" diff --git a/src/mitim_tools/astra_tools/ASTRA_CDFtools.py b/src/mitim_tools/astra_tools/ASTRA_CDFtools.py index ccf890bc..a4349395 100644 --- a/src/mitim_tools/astra_tools/ASTRA_CDFtools.py +++ b/src/mitim_tools/astra_tools/ASTRA_CDFtools.py @@ -347,7 +347,7 @@ def build_notebook(self, time_aims): self.getProfiles() # time_index = time_index(time) name = "ASTRA CDF Viewer" - self.fn = FigureNotebook(0, name) + self.fn = FigureNotebook(name) self.plot_temp(time_aims=time_aims) self.plot_gradients(time_aims=time_aims) diff --git a/src/mitim_tools/experiment_tools/CMODtools.py b/src/mitim_tools/experiment_tools/CMODtools.py index 43143b50..611b5955 100644 --- a/src/mitim_tools/experiment_tools/CMODtools.py +++ b/src/mitim_tools/experiment_tools/CMODtools.py @@ -974,6 +974,7 @@ def getZeff_neo( f.write(str(shotNumber)) Command = f"cd {folder} && idl < idl_in" + # FIX error, result = FARMINGtools.runCommand_remote( Command, machine=socket.gethostname() ) @@ -1008,6 +1009,7 @@ def getMMX(shotNumber, runid, folderWork): f.write(f"CMOD\n{shotNumber}\nA\nA\nQ\nY\nP\nPRF\nW\nQ") Command = f"cd {folderScratch} && scrunch2 < scrunch.in" + # FIX error, result = FARMINGtools.runCommand_remote( Command, machine=socket.gethostname() ) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index ceb88ff7..debe302c 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -1,21 +1,15 @@ -import os, time, copy +import os import numpy as np from IPython import embed import matplotlib.pyplot as plt - from mitim_tools.gacode_tools import TGYROtools from mitim_tools.gacode_tools.aux import GACODEdefaults, GACODErun from mitim_tools.misc_tools import IOtools, GRAPHICStools from mitim_tools.gacode_tools.aux import GACODEplotting from mitim_tools.misc_tools.IOtools import printMsg as print -from pygacode.cgyro.data import cgyrodata from pygacode.cgyro.data_plot import cgyrodata_plot from pygacode import gacodefuncs -""" - -""" - class CGYRO: def __init__( @@ -107,7 +101,6 @@ def runLinear( kys=[0.3], cores_per_ky=32, restart=False, - gacodeCGYRO="gacode", forceIfRestart=False, extra_name="", ): @@ -118,7 +111,6 @@ def runLinear( self.kys = kys numcores = int(len(self.kys) * cores_per_ky) - self.gacodeCGYRO = gacodeCGYRO self.FolderCGYRO = IOtools.expandPath(self.FolderGACODE + subFolderCGYRO + "/") self.FolderCGYRO_tmp = self.FolderCGYRO + "/tmp_standard/" @@ -165,7 +157,6 @@ def runLinear( numcores=numcores, filesToRetrieve=self.ResultsFiles, name=f'cgyro_{self.nameRunid}_{subFolderCGYRO.strip("/")}{extra_name}', - gacode_compilation=self.gacodeCGYRO, ) def read(self, label="cgyro1", folder=None): @@ -208,11 +199,11 @@ def get_flux(self, label="", moment="e", ispec=0, retrieveSpecieFlag=True): cgyro = self.results[label] # Time + usec = cgyro.getflux() cgyro.getnorm("elec") t = cgyro.tnorm # Flux - usec = cgyro.getflux("auto") ys = np.sum(cgyro.ky_flux, axis=(2, 3)) if moment == "n": y = ys[ispec, 0, :] @@ -366,7 +357,7 @@ def plot_fluxes(self, axs=None, label="", c="b", lw=1, plotLegend=True): lw=lw, ls=ls[0], ) - ax.set_title(f"Ion particle fluxes") + ax.set_title("Ion particle fluxes") # Extra ax = axs["C"] @@ -442,15 +433,13 @@ def plot_fluxes(self, axs=None, label="", c="b", lw=1, plotLegend=True): ) def plot(self, labels=[""]): - plt.rcParams["figure.max_open_warning"] = False from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - fn = FigureNotebook(0, "CGYRO Notebook", geometry="1600x1000") + self.fn = FigureNotebook("CGYRO Notebook", geometry="1600x1000") colors = GRAPHICStools.listColors() - fig = fn.add_figure(label="Fluxes Time Traces") + fig = self.fn.add_figure(label="Fluxes Time Traces") axsFluxes_t = fig.subplot_mosaic( """ ABC @@ -466,24 +455,20 @@ def plot(self, labels=[""]): plotLegend=j == len(labels) - 1, ) - fn.show() - def plotLS(self, labels=["cgyro1"], fig=None): colors = GRAPHICStools.listColors() if fig is None: # fig = plt.figure(figsize=(15,9)) - plt.rcParams["figure.max_open_warning"] = False + from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - fn = FigureNotebook( - 0, + self.fnLS = FigureNotebook( f"CGYRO Notebook, run #{self.nameRunid}, time {self.time:3f}s", geometry="1600x1000", ) - fig1 = fn.add_figure(label="Linear Stability") - fig2 = fn.add_figure(label="Ballooning") + fig1 = self.fnLS.add_figure(label="Linear Stability") + fig2 = self.fnLS.add_figure(label="Ballooning") grid = plt.GridSpec(2, 2, hspace=0.3, wspace=0.3) ax00 = fig1.add_subplot(grid[0, 0]) @@ -636,8 +621,6 @@ def plotLS(self, labels=["cgyro1"], fig=None): ax.axvline(x=0, lw=0.5, ls="--", c="k") ax.axhline(y=0, lw=0.5, ls="--", c="k") - fn.show() - def changeANDwrite_CGYRO(rhos, ky, FolderCGYRO, CGYROsettings=1): inputFilesCGYRO = {} diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py new file mode 100644 index 00000000..101001a9 --- /dev/null +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -0,0 +1,86 @@ +import os +from mitim_tools.misc_tools import IOtools +from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.gacode_tools.aux import GACODErun +from mitim_tools.misc_tools.IOtools import printMsg as print +from IPython import embed + + +class NEO: + def __init__(self): + pass + + def prep(self, inputgacode, folder): + self.inputgacode = inputgacode + self.folder = IOtools.expandPath(folder) + + if not os.path.exists(self.folder): + os.system(f"mkdir {self.folder}") + + def run_vgen(self, subfolder="vgen1", vgenOptions={}, restart=False): + while subfolder[-1] == "/": + subfolder = subfolder[:-1] + + self.folder_vgen = f"{self.folder}/{subfolder}/" + + # ---- Default options + + vgenOptions.setdefault("er", 2) + vgenOptions.setdefault("vel", 1) + vgenOptions.setdefault("numspecies", len(self.inputgacode.Species)) + vgenOptions.setdefault("matched_ion", 1) + vgenOptions.setdefault("nth", "17,39") + + # ---- Prepare + + runThisCase = check_if_files_exist( + self.folder_vgen, + [ + "vgen/input.gacode", + "vgen/input.neo.gen", + "out.vgen.neoequil00", + "out.vgen.neoexpnorm00", + "out.vgen.neontheta00", + "vgen.dat", + ], + ) + + if (not runThisCase) and restart: + runThisCase = print( + "\t- Files found in folder, but restart requested. Are you sure?", + typeMsg="q", + ) + + if runThisCase: + IOtools.askNewFolder(self.folder_vgen, force=True) + + self.inputgacode.writeCurrentStatus(file=f"{self.folder_vgen}/input.gacode") + + # ---- Run + + if runThisCase: + file_new = GACODErun.runVGEN( + self.folder_vgen, vgenOptions=vgenOptions, name_run=subfolder + ) + else: + print( + f"\t- Required files found in {subfolder}, not running VGEN", + typeMsg="i", + ) + file_new = f"{self.folder_vgen}/vgen/input.gacode" + + # ---- Postprocess + + self.inputgacode_vgen = PROFILEStools.PROFILES_GACODE( + file_new, calculateDerived=True, mi_ref=self.inputgacode.mi_ref + ) + + +def check_if_files_exist(folder, list_files): + folder = IOtools.expandPath(folder) + + for file in list_files: + if not os.path.exists(f"{folder}/{file}"): + return False + + return True diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 148e95e6..895b4237 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -1,4 +1,5 @@ -import copy, torch +import copy +import torch import numpy as np import matplotlib.pyplot as plt from collections import OrderedDict @@ -7,7 +8,8 @@ from mitim_tools.misc_tools import GRAPHICStools, MATHtools, PLASMAtools, IOtools from mitim_modules.powertorch.physics import GEOMETRYtools, CALCtools from mitim_tools.gs_tools import GEQtools -from mitim_tools.gacode_tools.aux import TRANSPinteraction, PROFILEStoMODELS, GACODErun +from mitim_tools.gacode_tools import NEOtools +from mitim_tools.gacode_tools.aux import TRANSPinteraction, PROFILEStoMODELS from mitim_tools.transp_tools import CDFtools from mitim_tools.im_tools.modules import PEDmodule from mitim_tools.misc_tools.CONFIGread import read_verbose_level @@ -18,7 +20,7 @@ try: from mitim_tools.gacode_tools.aux import PORTALSinteraction -except: +except ImportError: print( "- I could not import PORTALSinteraction, likely a consequence of botorch incompatbility", typeMsg="w", @@ -107,31 +109,55 @@ def runTRANSPfromGACODE(self, folderTRANSP, machine="SPARC"): folderTRANSP, machine=machine ) - def runNEOforEr( + def calculate_Er( self, folder, - vgenOptions={ - "er": 2, - "vel": 1, - "numspecies": None, - "matched_ion": 1, - "nth": "17,39", - }, - name="", + rhos=None, + vgenOptions={}, + name="vgen1", + includeAll=False, + write_new_file=None, + restart=False, ): - if vgenOptions["numspecies"] is None: - vgenOptions["numspecies"] = len(self.Species) + profiles = copy.deepcopy(self) + + # Resolution? + resol_changed = False + if rhos is not None: + profiles.changeResolution(rho_new=rhos) + resol_changed = True + + self.neo = NEOtools.NEO() + self.neo.prep(profiles, folder) + self.neo.run_vgen(subfolder=name, vgenOptions=vgenOptions, restart=restart) + + profiles_new = copy.deepcopy(self.neo.inputgacode_vgen) + if resol_changed: + profiles_new.changeResolution(rho_new=self.profiles["rho(-)"]) + + # Get the information from the NEO run + + variables = ["w0(rad/s)"] + if includeAll: + variables += [ + "vpol(m/s)", + "vtor(m/s)", + "jbs(MA/m^2)", + "jbstor(MA/m^2)", + "johm(MA/m^2)", + ] - print( - f"\t- Running NEO (with {vgenOptions['numspecies']} species) to populate w0(rad/s) in input.gacode file" - ) - print(f"\t\t> Matching ion {vgenOptions['matched_ion']} Vtor") + for ikey in variables: + if ikey in profiles_new.profiles: + print( + f'\t- Inserting {ikey} from NEO run{" (went back to original resolution by interpolation)" if resol_changed else ""}' + ) + self.profiles[ikey] = profiles_new.profiles[ikey] - file_new = GACODErun.runVGEN( - self.file, folder, vgenOptions=vgenOptions, nameChange=name - ) + self.deriveQuantities() - self.__init__(file_new, calculateDerived=True, mi_ref=self.mi_ref) + if write_new_file is not None: + self.writeCurrentStatus(file=write_new_file) # ***************** @@ -142,6 +168,8 @@ def readHeader(self): self.header = self.lines[:istartProfs] def readProfiles(self): + singleLine, title, var = None, None, None # for ruff complaints + # --- found = False self.profiles = OrderedDict() @@ -428,7 +456,7 @@ def deriveQuantities(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry=True) self.profiles["te(keV)"], self.derived["mi_ref"], self.derived["B_unit"] ) - self.derived["q_gb"], self.derived["g_gb"], _, _ = PLASMAtools.gyrobohmUnits( + self.derived["q_gb"], self.derived["g_gb"], _, _, _ = PLASMAtools.gyrobohmUnits( self.profiles["te(keV)"], self.profiles["ne(10^19/m^3)"] * 1e-1, self.derived["mi_ref"], @@ -508,6 +536,8 @@ def deriveQuantities(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry=True) self.derived["ge"] * 1e-20, r, volp ) # Because the units were #/sec/m^3 + self.derived["geIn"] = self.derived["ge_10E20miller"][-1] # 1E20 particles/sec + self.derived["qe_MWm2"] = self.derived["qe_MWmiller"] / (volp) self.derived["qi_MWm2"] = self.derived["qi_MWmiller"] / (volp) self.derived["ge_10E20m2"] = self.derived["ge_10E20miller"] / (volp) @@ -762,6 +792,10 @@ def deriveQuantities(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry=True) self.derived["tauE"] = self.derived["Wthr"] / self.derived["qHeat"] # Seconds + self.derived["tauP"] = self.derived["Ne"] / self.derived["geIn"] # Seconds + + self.derived["tauPotauE"] = self.derived["tauP"] / self.derived["tauE"] + # Dilutions self.derived["fi"] = self.profiles["ni(10^19/m^3)"] / np.atleast_2d( self.profiles["ne(10^19/m^3)"] @@ -910,7 +944,7 @@ def deriveQuantities(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry=True) self.derived["BetaN"] = ( Beta / ( - float(self.profiles["current(MA)"][-1]) + np.abs(float(self.profiles["current(MA)"][-1])) / (self.derived["a"] * self.derived["B0"]) ) * 100.0 @@ -1025,68 +1059,62 @@ def deriveContentByVolumes(self, rhos=[0.5], impurityPosition=3): return We, Wi, Ne, Ni - def printInfo(self, label=""): - ImpurityText = "" - for i in range(len(self.Species)): - ImpurityText += f"{self.Species[i]['N']}({self.Species[i]['Z']:.0f},{self.Species[i]['A']:.0f}) = {self.derived['fi_vol'][i]:.1e}, " - ImpurityText = ImpurityText[:-2] + def printInfo(self, label="", reDeriveIfNotFound=True): + try: + ImpurityText = "" + for i in range(len(self.Species)): + ImpurityText += f"{self.Species[i]['N']}({self.Species[i]['Z']:.0f},{self.Species[i]['A']:.0f}) = {self.derived['fi_vol'][i]:.1e}, " + ImpurityText = ImpurityText[:-2] - print(f"\n***********************{label}****************") - print("Performance:") - print( - "\tQ = {0:.2f} (Pfus = {1:.1f}MW, Pin = {2:.1f}MW)".format( - self.derived["Q"], self.derived["Pfus"], self.derived["qIn"] + print(f"\n***********************{label}****************") + print("Performance:") + print( + "\tQ = {0:.2f} (Pfus = {1:.1f}MW, Pin = {2:.1f}MW)".format( + self.derived["Q"], self.derived["Pfus"], self.derived["qIn"] + ) ) - ) - print( - "\tH98y2 = {0:.2f} (tauE = {1:.3f} s)".format( - self.derived["H98"], self.derived["tauE"] + print( + "\tH98y2 = {0:.2f} (tauE = {1:.3f} s)".format( + self.derived["H98"], self.derived["tauE"] + ) ) - ) - print( - "\tH89p = {0:.2f} (H97L = {1:.2f})".format( - self.derived["H89"], self.derived["H97L"] + print( + "\tH89p = {0:.2f} (H97L = {1:.2f})".format( + self.derived["H89"], self.derived["H97L"] + ) ) - ) - print( - "\tnu_ne = {0:.2f} (nu_eff = {1:.2f})".format( - self.derived["ne_peaking"], self.derived["nu_eff"] + print( + "\tnu_ne = {0:.2f} (nu_eff = {1:.2f})".format( + self.derived["ne_peaking"], self.derived["nu_eff"] + ) ) - ) - print( - "\tnu_ne0.2 = {0:.2f} (nu_eff w/Zeff2 = {1:.2f})".format( - self.derived["ne_peaking0.2"], self.derived["nu_eff2"] + print( + "\tnu_ne0.2 = {0:.2f} (nu_eff w/Zeff2 = {1:.2f})".format( + self.derived["ne_peaking0.2"], self.derived["nu_eff2"] + ) ) - ) - print(f"\tnu_Ti = {self.derived['Ti_peaking']:.2f}") - try: + print(f"\tnu_Ti = {self.derived['Ti_peaking']:.2f}") print( "\tBetaN = {0:.3f} (approx, based on B0 and p_thr)".format( self.derived["BetaN"] ) ) - except: - pass - print( - "\tPrad = {0:.1f}MW ({1:.1f}% of total)".format( - self.derived["Prad"], - self.derived["Prad"] / self.derived["qHeat"] * 100.0, + print( + "\tPrad = {0:.1f}MW ({1:.1f}% of total)".format( + self.derived["Prad"], + self.derived["Prad"] / self.derived["qHeat"] * 100.0, + ) ) - ) - try: print( "\tPsol = {0:.1f}MW (fLH = {1:.2f})".format( self.derived["Psol"], self.derived["LHratio"] ) ) - except: - pass - print( - "Operational point (, = [{0:.2f},{1:.2f}) and species:".format( - self.derived["ne_vol20"], self.derived["Te_vol"] + print( + "Operational point ( [,] = [{0:.2f},{1:.2f}] ) and species:".format( + self.derived["ne_vol20"], self.derived["Te_vol"] + ) ) - ) - try: print( "\t = {0:.2f} keV (/ = {1:.2f}, Ti0/Te0 = {2:.2f})".format( self.derived["Ti_vol"], @@ -1094,37 +1122,41 @@ def printInfo(self, label=""): self.derived["tite"][0], ) ) - except: print(f"\t = {self.derived['Ti_vol']:.2f} keV") - print( - "\tfG = {0:.2f} ( = {1:.2f} * 10^20 m^-3)".format( - self.derived["fG"], self.derived["ne_vol20"] + print( + "\tfG = {0:.2f} ( = {1:.2f} * 10^20 m^-3)".format( + self.derived["fG"], self.derived["ne_vol20"] + ) ) - ) - try: print( - f"\tZeff = {self.derived['Zeff_vol']:.2f} (M_main = {self.derived['mbg_main']:.2f}, f_main = {self.derived['fmain']:.2f}) [QN err = {self.derived['QN_Error']:.4f}]" + f"\tZeff = {self.derived['Zeff_vol']:.2f} (M_main = {self.derived['mbg_main']:.2f}, f_main = {self.derived['fmain']:.2f}) [QN err = {self.derived['QN_Error']:.1e}]" ) - except: - pass - try: print(f"\tMach = {self.derived['MachNum_vol']:.2f} (vol avg)") - except: - pass - print("Content:") - print( - "\tWe = {0:.2f} MJ, Wi_thr = {1:.2f} MJ (W_thr = {2:.2f} MJ)".format( - self.derived["We"], self.derived["Wi_thr"], self.derived["Wthr"] + print("Content:") + print( + "\tWe = {0:.2f} MJ, Wi_thr = {1:.2f} MJ (W_thr = {2:.2f} MJ)".format( + self.derived["We"], self.derived["Wi_thr"], self.derived["Wthr"] + ) ) - ) - print( - "\tNe = {0:.1f}*10^20, Ni_thr = {1:.1f}*10^20 (N_thr = {2:.1f}*10^20)".format( - self.derived["Ne"], self.derived["Ni_thr"], self.derived["Nthr"] + print( + "\tNe = {0:.1f}*10^20, Ni_thr = {1:.1f}*10^20 (N_thr = {2:.1f}*10^20)".format( + self.derived["Ne"], self.derived["Ni_thr"], self.derived["Nthr"] + ) ) - ) - print("Species concentration:") - print(f"\t{ImpurityText}") - print("******************************************************") + print( + f"\ttauE = { self.derived['tauE']:.3f} s, tauP = {self.derived['tauP']:.3f} s (tauP/tauE = {self.derived['tauPotauE']:.2f})" + ) + print("Species concentration:") + print(f"\t{ImpurityText}") + print("******************************************************") + except KeyError: + print( + "\t- When printing info, not all keys found, probably because this input.gacode class came from an old MITIM version", + typeMsg="w", + ) + if reDeriveIfNotFound: + self.deriveQuantities() + self.printInfo(label=label, reDeriveIfNotFound=False) def makeAllThermalIonsHaveSameTemp(self, refIon=0): SpecRef = self.Species[refIon]["N"] @@ -1143,10 +1175,10 @@ def scaleAllThermalDensities(self, scaleFactor=1.0): for sp in range(len(self.Species)): if self.Species[sp]["S"] == "therm": print( - f"\t\t\t- Scaling density of {self.Species[sp]['N']} by an average factor of {np.mean(scaleFactor):.3f}" + f"\t\t\t- Scaling density of {self.Species[sp]['N']} by an average factor of {np.mean(scaleFactor_ions):.3f}" ) ni_orig = self.profiles["ni(10^19/m^3)"][:, sp] - self.profiles["ni(10^19/m^3)"][:, sp] = scaleFactor * ni_orig + self.profiles["ni(10^19/m^3)"][:, sp] = scaleFactor_ions * ni_orig def writeCurrentStatus(self, file=None, limitedNames=False): print("\t- Writting input.gacode file") @@ -1276,7 +1308,7 @@ def changeResolution( self.deriveQuantities(mi_ref=self.derived["mi_ref"]) print( - f"\t- Resolution of profiles changed to {n} points with function {interpFunction}" + f"\t\t- Resolution of profiles changed to {n} points with function {interpFunction}" ) def DTplasma(self): @@ -1791,19 +1823,15 @@ def plot( lsFlows="-", legFlows=True, showtexts=True, + lastRhoGradients=0.89, ): if axs1 is None: if fn is None: - plt.rcParams["figure.max_open_warning"] = False from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - fn = FigureNotebook(0, "PROFILES Notebook", geometry="1600x1000") - wasProvided = False - else: - wasProvided = True + self.fn = FigureNotebook("PROFILES Notebook", geometry="1600x1000") - fig = fn.add_figure(label="Profiles" + fnlab) + fig = self.fn.add_figure(label="Profiles" + fnlab) grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) axs1 = [ fig.add_subplot(grid[0, 0]), @@ -1817,7 +1845,7 @@ def plot( fig.add_subplot(grid[2, 2]), ] - fig2 = fn.add_figure(label="Powers" + fnlab) + fig2 = self.fn.add_figure(label="Powers" + fnlab) grid = plt.GridSpec(3, 2, hspace=0.3, wspace=0.3) axs2 = [ fig2.add_subplot(grid[0, 0]), @@ -1828,7 +1856,7 @@ def plot( fig2.add_subplot(grid[2, 1]), ] - fig3 = fn.add_figure(label="Geometry" + fnlab) + fig3 = self.fn.add_figure(label="Geometry" + fnlab) grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.5) ax00c = fig3.add_subplot(grid[0, 0]) axs3 = [ @@ -1846,7 +1874,7 @@ def plot( fig3.add_subplot(grid[2, 3], sharex=ax00c), ] - fig4 = fn.add_figure(label="Gradients" + fnlab) + fig4 = self.fn.add_figure(label="Gradients" + fnlab) grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) axs4 = [ fig4.add_subplot(grid[0, 0]), @@ -1857,7 +1885,7 @@ def plot( fig4.add_subplot(grid[1, 2]), ] - fig5 = fn.add_figure(label="Flows" + fnlab) + fig5 = self.fn.add_figure(label="Flows" + fnlab) grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) axsFlows = [ @@ -1869,7 +1897,7 @@ def plot( fig5.add_subplot(grid[1, 2]), ] - fig6 = fn.add_figure(label="Other" + fnlab) + fig6 = self.fn.add_figure(label="Other" + fnlab) grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) axs6 = [ fig6.add_subplot(grid[0, 0]), @@ -1881,7 +1909,7 @@ def plot( fig6.add_subplot(grid[1, 3]), ] - fig7 = fn.add_figure(label="Impurities" + fnlab) + fig7 = self.fn.add_figure(label="Impurities" + fnlab) grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) axsImps = [ fig7.add_subplot(grid[0, 0]), @@ -1892,9 +1920,6 @@ def plot( fig7.add_subplot(grid[1, 2]), ] - else: - wasProvided = True - [ax00, ax10, ax20, ax01, ax11, ax21, ax02, ax12, ax22] = axs1 [ax00b, ax01b, ax10b, ax11b, ax20b, ax21b] = axs2 [ @@ -2440,7 +2465,9 @@ def plot( GRAPHICStools.autoscale_y(ax, bottomy=0) # Derived - self.plotGradients(axs4, color=color, lw=lw) + self.plotGradients( + axs4, color=color, lw=lw, lastRho=lastRhoGradients, label=extralab + ) # Others ax = axs6[0] @@ -2658,9 +2685,6 @@ def plot( GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax, bottomy=0) - if not wasProvided: - fn.show() - def plotGradients( self, axs4, @@ -2765,6 +2789,7 @@ def plotGradients( ax.set_ylabel("$T_e$ (keV)") ax.set_xlabel(labelx) GRAPHICStools.autoscale_y(ax, bottomy=0) + ax.legend(loc="best", fontsize=7) ax = axs4[2] ax.set_ylabel("$T_i$ (keV)") ax.set_xlabel(labelx) @@ -3463,7 +3488,6 @@ def plotPeaking( ix = np.argmin(np.abs(self.profiles["rho(-)"] - 0.9)) if debugPlot: - plt.ioff() fig, axq = plt.subplots() ne = self.profiles["ne(10^19/m^3)"] @@ -3526,7 +3550,7 @@ def compareProfiles(profiles_list, fig=None, labs_list=[""] * 10, lws=[3] * 10): ax01 = fig.add_subplot(grid[0, 1]) ax11 = fig.add_subplot(grid[1, 1]) ax02 = fig.add_subplot(grid[0, 2]) - ax12 = fig.add_subplot(grid[1, 2]) + # ax12 = fig.add_subplot(grid[1, 2]) cols = GRAPHICStools.listColors() @@ -3549,15 +3573,14 @@ def compareProfiles(profiles_list, fig=None, labs_list=[""] * 10, lws=[3] * 10): ) -def plotAll(profiles_list, figs=None, extralabs=None): +def plotAll(profiles_list, figs=None, extralabs=None, lastRhoGradients=0.89): if figs is not None: figProf_1, figProf_2, figProf_3, figProf_4, figFlows, figProf_6, fig7 = figs + fn = None else: - plt.rcParams["figure.max_open_warning"] = False from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - fn = FigureNotebook(0, "Profiles", geometry="1800x900") + fn = FigureNotebook("Profiles", geometry="1800x900") figProf_1 = fn.add_figure(label="Profiles") figProf_2 = fn.add_figure(label="Powers") figProf_3 = fn.add_figure(label="Geometry") @@ -3668,10 +3691,10 @@ def plotAll(profiles_list, figs=None, extralabs=None): lsFlows=ls[i], legFlows=i == 0, showtexts=False, + lastRhoGradients=lastRhoGradients, ) - if figs is None: - fn.show() + return fn def readTGYRO_profile_extra(file, varLabel="B_unit (T)"): @@ -3743,7 +3766,7 @@ def ionName(Z, A): # elif A == 69: return 'Ga' -def gradientsMerger(p0, p_true, roa=0.46, blending=0.1, debug=False): +def gradientsMerger(p0, p_true, roa=0.46, blending=0.1): p = copy.deepcopy(p0) aLTe_true = np.interp( @@ -3816,28 +3839,6 @@ def gradientsMerger(p0, p_true, roa=0.46, blending=0.1, debug=False): .numpy()[0] ) - if debug: - fig, axs = plt.subplots(nrows=2) - ax = axs[0] - ax.plot(p.derived["roa"], p.profiles["te(keV)"], c="b") - ax.plot(roa_true, p_true.profiles["te(keV)"], c="r") - ax.plot(p.derived["roa"], Te, c="g") - - ax.axvline(x=roa, c="k", ls="--") - ax.axvline(x=roa - blending, c="k", ls="--") - - ax = axs[1] - ax.plot(p.derived["roa"], p.derived["aLTe"], c="b") - ax.plot(p_true.derived["roa"], p_true.derived["aLTe"], c="r") - ax.plot(p.derived["roa"], aLTe, c="g") - - ax.axvline(x=roa, c="k", ls="--") - ax.axvline(x=roa - blending, c="k", ls="--") - - plt.show() - - embed() - p.profiles["te(keV)"] = Te p.profiles["ti(keV)"][:, 0] = Ti p.profiles["ne(10^19/m^3)"] = ne diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 2fbf514d..6a0e3062 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -1,7 +1,8 @@ -import os, copy, pickle +import os +import copy +import pickle import numpy as np import matplotlib.pyplot as plt -from IPython import embed from mitim_tools.gacode_tools import TGYROtools, PROFILEStools from mitim_tools.misc_tools import ( IOtools, @@ -20,6 +21,8 @@ from mitim_tools.misc_tools.CONFIGread import read_verbose_level +from IPython import embed + verbose_level = read_verbose_level() mi_D = 2.01355 @@ -78,7 +81,7 @@ def __init__( tglf.read(label='run1',folder='~/testTGLF/tglf1/') # Plot - plt.ion(); tglf.plotRun(labels=['run1']) + plt.ion(); tglf.plot(labels=['run1']) ********************************* ***** Example use for scans ***** @@ -162,7 +165,7 @@ def __init__( if self.LocationCDF is not None: _, self.nameRunid = IOtools.getLocInfo(self.LocationCDF) else: - self.nameRunid = "mitim" + self.nameRunid = "0" self.time, self.avTime = time, avTime self.rhos = np.array(rhos) @@ -171,23 +174,41 @@ def __init__( self.scans, self.tgyro, self.ky_single, - self.NormalizationSets, - ) = ({}, {}, None, None, None) + ) = ({}, {}, None, None) + + self.NormalizationSets = { + "TRANSP": None, + "PROFILES": None, + "TGYRO": None, + "EXP": None, + "input_gacode": None, + "SELECTED": None, + } - def save_pkl(self, file): - print(f"> Saving tglf class as pickle file: {IOtools.clipstr(file)}") - try: + def prepare_for_save_TGLF(self): + """ + This is a function that will be called when saving the class as pickle. + It will delete things that are not easily pickled. + """ + + if "fn" in self.__dict__: del self.fn # otherwise it cannot deepcopy - except: - pass + tglf_copy = copy.deepcopy(self) - # Delete things that are not easily pickled del tglf_copy.convolution_fun_fluct for label in tglf_copy.results: if "convolution_fun_fluct" in tglf_copy.results[label]: tglf_copy.results[label]["convolution_fun_fluct"] = None + return tglf_copy + + def save_pkl(self, file): + print(f"> Saving tglf class as pickle file: {IOtools.clipstr(file)}") + + # Prepare + tglf_copy = self.prepare_for_save_TGLF() + # Write with open(file, "wb") as handle: pickle.dump(tglf_copy, handle) @@ -254,7 +275,9 @@ def prep( inp = TGLFinput(fii) exists = exists and not inp.onlyControl else: - print(f"\t\t- Running scans because it does not exist file {fii}") + print( + f"\t\t- Running scans because it does not exist file {IOtools.clipstr(fii)}" + ) exists = False if exists: print( @@ -343,19 +366,37 @@ def prep_from_tglf( # Main folder where things are self.FolderGACODE = IOtools.expandPath(FolderGACODE) - # input_tglf_file - inputclass = TGLFinput(file=input_tglf_file) - - self.rhos = [inputclass.geom["RMIN_LOC"]] - self.inputsTGLF = {self.rhos[0]: inputclass} - print(f"\t- This file correspond to {self.rhos[0]} according to RMIN_LOC") - + # Main folder where things are self.NormalizationSets, _ = NORMtools.normalizations( PROFILEStools.PROFILES_GACODE(input_gacode) if input_gacode is not None else None ) + # input_tglf_file + inputclass = TGLFinput(file=input_tglf_file) + + roa = inputclass.geom["RMIN_LOC"] + print(f"\t- This file correspond to r/a={roa} according to RMIN_LOC") + + if self.NormalizationSets["input_gacode"] is not None: + rho = np.interp( + roa, + self.NormalizationSets["input_gacode"].derived["roa"], + self.NormalizationSets["input_gacode"].profiles["rho(-)"], + ) + print(f"\t\t- rho={rho:.4f}, using input.gacode for conversion") + else: + print( + "\t\t- No input.gacode for conversion, assuming rho=r/a, EXTREME CAUTION PLEASE", + typeMsg="w", + ) + rho = roa + + self.rhos = [rho] + + self.inputsTGLF = {self.rhos[0]: inputclass} + def run( self, subFolderTGLF, # 'tglf1/', @@ -369,7 +410,7 @@ def run( launchSlurm=True, restart=False, forceIfRestart=False, - extra_name="", + extra_name="exe", slurm_setup={ "cores": 4, "minutes": 5, @@ -378,11 +419,10 @@ def run( inputs = copy.deepcopy(self.inputsTGLF) self.FolderTGLF = IOtools.expandPath(self.FolderGACODE + subFolderTGLF + "/") - self.FolderTGLF_tmp = self.FolderTGLF + "/tmp_tglf/" ResultsFiles_new = [] for i in self.ResultsFiles: - if "sbatch.out" not in i: + if "mitim.out" not in i: ResultsFiles_new.append(i) self.ResultsFiles = ResultsFiles_new @@ -405,7 +445,6 @@ def run( if len(rhosEvaluate) == len(self.rhos): IOtools.askNewFolder(self.FolderTGLF, force=forceIfRestart) - IOtools.askNewFolder(self.FolderTGLF_tmp) else: if not exists: print( @@ -474,12 +513,11 @@ def run( GACODErun.runTGLF( rhosEvaluate, self.FolderTGLF, - self.FolderTGLF_tmp, self.latest_inputsFileTGLF, filesToRetrieve=self.ResultsFiles, minutes=slurm_setup["minutes"], cores_tglf=slurm_setup["cores"], - name=f"tglf_{self.nameRunid}_{subFolderTGLF.strip('/')}{extra_name}", + name=f"tglf_{self.nameRunid}{subFolderTGLF.replace('/','_')}{extra_name}", launchSlurm=launchSlurm, ) else: @@ -510,11 +548,11 @@ def run( self.read(label=f"ky{ky_single0}", folder=FolderTGLF_old) - self.FoldersTGLF_WF[f"ky{ky_single0}"] = {} - for i, ir in enumerate(self.rhos): - self.FolderTGLF = f"{FolderTGLF_old}/wf_{ir}_{ky_single0}/" - self.FoldersTGLF_WF[f"ky{ky_single0}"][ir] = self.FolderTGLF + self.FolderTGLF = f"{FolderTGLF_old}/ky{ky_single0}/" + self.FoldersTGLF_WF[f"ky{ky_single0}"] = self.FolderTGLF + ky_singles = [] + for i, ir in enumerate(self.rhos): # -------- Get the closest unstable mode to the one requested if forceClosestUnstableWF: # Only unstable ones @@ -546,52 +584,39 @@ def run( ky_single = closest_ky else: ky_single = ky_single0 - # ------------------------------------------------------------ - - existsWF = not restart - for j in self.ResultsFiles: - existsWF = existsWF and os.path.exists( - f"{self.FolderTGLF}{j}_{ir:.4f}" - ) - - if not existsWF: - IOtools.askNewFolder(self.FolderTGLF, force=forceIfRestart) - os.system( - f"cp {FolderTGLF_old}/input.tglf* {self.FolderTGLF}/" - ) - extraOptions_WF = copy.deepcopy(extraOptions) - - extraOptions_WF["USE_TRANSPORT_MODEL"] = "F" - extraOptions_WF["KY"] = ky_single - extraOptions_WF[ - "VEXB_SHEAR" - ] = 0.0 # See email from G. Staebler on 05/16/2021 - inputFileTGLF, _ = changeANDwrite_TGLF( - [ir], - {ir: inputs[ir]}, - self.FolderTGLF, - TGLFsettings=TGLFsettings, - extraOptions=extraOptions_WF, - multipliers=multipliers, - ApplyCorrections=ApplyCorrections, - Quasineutral=Quasineutral, - ) + ky_singles.append(ky_single) + # ------------------------------------------------------------ - GACODErun.runTGLF( - [ir], - self.FolderTGLF, - self.FolderTGLF_tmp, - inputFileTGLF, - filesToRetrieve=self.ResultsFiles, - minutes=slurm_setup["minutes"], - cores_tglf=slurm_setup["cores"], - name=f"{self.nameRunid}_{subFolderTGLF.strip('/')}", - launchSlurm=launchSlurm, - ) + extraOptions_WF = copy.deepcopy(extraOptions) + + extraOptions_WF["USE_TRANSPORT_MODEL"] = "F" + extraOptions_WF["WRITE_WAVEFUNCTION_FLAG"] = 1 + extraOptions_WF["KY"] = ky_singles + extraOptions_WF[ + "VEXB_SHEAR" + ] = 0.0 # See email from G. Staebler on 05/16/2021 + + tglf_wf = copy.deepcopy(self) + + tglf_wf.run( + f"{subFolderTGLF}/ky{ky_single0}", + TGLFsettings=TGLFsettings, + extraOptions=extraOptions_WF, + multipliers=multipliers, + runWaveForms=[], + ApplyCorrections=ApplyCorrections, + Quasineutral=Quasineutral, + launchSlurm=launchSlurm, + restart=restart, + forceIfRestart=forceIfRestart, + extra_name=extra_name, + slurm_setup=slurm_setup, + ) del self.results[f"ky{ky_single0}"] except FileNotFoundError: + embed() self.ky_single = None print( "> Waveform analysis FAILED (maybe due to TGLF version?)", typeMsg="w" @@ -607,13 +632,13 @@ def read( self, label="tglf1", folder=None, # If None, search in the previously run folder - suffix=None, # If None, search with my standard _0.55 suffixes corresponding to rho + suffix=None, # If None, search with my standard _0.55 suffixes corresponding to rho of this TGLF class d_perp_cm=None, # It can be a dictionary with rhos. If None provided, use the last one employed ): print("> Reading TGLF results") if d_perp_cm is not None: - if type(d_perp_cm) is float: + if isinstance(d_perp_cm, float): self.d_perp_dict = {} for rho in self.rhos: self.d_perp_dict[rho] = d_perp_cm @@ -680,23 +705,21 @@ def read( self.results[label]["wavefunction"] = {} if self.ky_single is not None: for ky_single0 in self.ky_single: - self.results[label]["wavefunction"][f"ky{ky_single0}"] = {} if f"ky{ky_single0}" not in self.FoldersTGLF_WF: continue - for ir in self.rhos: - if ir not in self.FoldersTGLF_WF[f"ky{ky_single0}"]: - continue + self.results[label]["wavefunction"][f"ky{ky_single0}"] = {} + for ir in self.rhos: suffix0 = f"_{ir:.4f}" if suffix is None else suffix self.results[label]["wavefunction"][f"ky{ky_single0}"][ ir ] = GACODEinterpret.Waveform_read( - f"{self.FoldersTGLF_WF[f'ky{ky_single0}'][ir]}/out.tglf.wavefunction{suffix0}", - f"{self.FoldersTGLF_WF[f'ky{ky_single0}'][ir]}/out.tglf.run{suffix0}", + f"{self.FoldersTGLF_WF[f'ky{ky_single0}']}/out.tglf.wavefunction{suffix0}", + f"{self.FoldersTGLF_WF[f'ky{ky_single0}']}/out.tglf.run{suffix0}", ) - def plotRun( + def plot( self, fn=None, labels=["tglf1"], @@ -759,17 +782,15 @@ def plotRun( ) for il in self.results[label]["TGLFout"][irho].fields: - if not (il in max_fields): + if il not in max_fields: max_fields.append(il) - plt.rcParams["figure.max_open_warning"] = False - plt.ioff() if fn is None: - fnsh, self.fn = True, GUItools.FigureNotebook( - 0, "TGLF MITIM Notebook", geometry="1700x900", vertical=True + self.fn = GUItools.FigureNotebook( + "TGLF MITIM Notebook", geometry="1700x900", vertical=True ) else: - fnsh, self.fn = False, fn + self.fn = fn # *** TGLF Figures fig1 = self.fn.add_figure(label=f"{extratitle}Summary") @@ -1818,7 +1839,6 @@ def plotRun( theta = wf["theta"] / np.pi markers = GRAPHICStools.listmarkers() - typeline = GRAPHICStools.listmarkersLS() # ES max0, min0 = GACODEplotting.plotWaveform( @@ -1954,9 +1974,7 @@ def plotRun( ax, size=6, ratio=0.6, title=title_legend ) ax.set_title("Growth Rate") - ax.set_xlim( - [self.ky_single[kycont] * 0.8, self.ky_single[kycont] * 1.2] - ) + ax.set_xlim([self.ky_single[kycont] - 2.0, self.ky_single[kycont] + 2]) # ax.set_yscale('log') ax = ax10 @@ -1966,9 +1984,7 @@ def plotRun( if addLegend: GRAPHICStools.addLegendApart(ax, size=6, ratio=0.6, withleg=False) ax.set_title("Real Frequency") - ax.set_xlim( - [self.ky_single[kycont] * 0.8, self.ky_single[kycont] * 1.2] - ) + ax.set_xlim([self.ky_single[kycont] - 2.0, self.ky_single[kycont] + 2]) ax = ax01 ax.set_xlabel("Poloidal angle $\\theta$ ($\\pi$)") @@ -2029,9 +2045,6 @@ def plotRun( legYN=contLab == 0, ) - if fnsh: - self.fn.show() - # ~~~~~~~~~~~~~~ Scan options def runScan( @@ -2061,7 +2074,7 @@ def runScan( # ------------------------------------- if (1.0 not in varUpDown) and relativeChanges: print( - f"\n* Since variations vector did not include base case, I am adding it", + "\n* Since variations vector did not include base case, I am adding it", typeMsg="i", ) varUpDown_new = [] @@ -2079,7 +2092,7 @@ def runScan( varUpDown_new[i] = round(varUpDown_new[i], 3) print(f"\n- Proceeding to scan {variable}:") - for mult in varUpDown_new: + for cont_mult, mult in enumerate(varUpDown_new): mult = round(mult, 6) if relativeChanges: @@ -2100,20 +2113,25 @@ def runScan( if not relativeChanges: for ikey in multipliers_mod: - extraOptions[ikey] = multipliers_mod[ikey] + kwargs_TGLFrun["extraOptions"][ikey] = multipliers_mod[ikey] multipliers_mod = {} # Force ensure quasineutrality if the if variable in ["AS_3", "AS_4", "AS_5", "AS_6"]: - Quasineutral = True + kwargs_TGLFrun["Quasineutral"] = True + + # Only ask the restart in the first round + kwargs_TGLFrun["forceIfRestart"] = cont_mult > 0 or ( + "forceIfRestart" in kwargs_TGLFrun and kwargs_TGLFrun["forceIfRestart"] + ) self.run( subFolderTGLF=f"{self.subFolderTGLF_scan}_{name}", - multipliers=multipliers_mod, # forceIfRestart=True, + multipliers=multipliers_mod, **kwargs_TGLFrun, ) - self.read(label=f"{subFolderTGLF}{name}") + self.read(label=f"{self.subFolderTGLF_scan}_{name}") def readScan( self, label="scan1", subFolderTGLF=None, variable="RLTS_1", positionIon=2 @@ -2121,8 +2139,8 @@ def readScan( if subFolderTGLF is None: subFolderTGLF = self.subFolderTGLF_scan - if subFolderTGLF[-1] != "/": - subFolderTGLF += "/" + while self.subFolderTGLF_scan[-1] == "/": + self.subFolderTGLF_scan = self.subFolderTGLF_scan[:-1] self.scans[label] = {} self.scans[label]["variable"] = variable @@ -2149,7 +2167,12 @@ def readScan( etalow_g, etalow_f, etalow_k = [], [], [] cont = 0 for ikey in self.results: - if ikey.rpartition(variable)[0] == subFolderTGLF: + isThisTheRightReadResults = (subFolderTGLF in ikey) and ( + variable + == "_".join(ikey.split("_")[:-1]).split(subFolderTGLF + "_")[-1] + ) + + if isThisTheRightReadResults: x0, Qe0, Qi0, Ge0, Gi0, ky0, g0, f0, eta10, eta20, itg0, tem0, etg0 = ( [], [], @@ -2311,10 +2334,8 @@ def plotScan( ) if figs is None: - plt.rcParams["figure.max_open_warning"] = False - plt.ioff() self.fn = GUItools.FigureNotebook( - 0, "TGLF Scan MITIM Notebook", geometry="1500x900", vertical=True + "TGLF Scan MITIM Notebook", geometry="1500x900", vertical=True ) if unnormalization_successful: fig1 = self.fn.add_figure(label="Fluxes") @@ -2753,41 +2774,38 @@ def plotScan( # -------------------------------------------------------- total_plots = 0 - for ikey in self.scans: + for ikey in labels: total_plots += self.scans[ikey]["Qe_gb"].size if plotTGLFs and ( total_plots < 10 or ( print( - f">> TGLFscan module will plot {total_plots} individual TGLF, you can choose not to plot them invidually", + f">> TGLFscan module wants to *also* plot {total_plots} individual TGLF, are you sure you want to do this?", typeMsg="q", ) ) ): - for contLabel, ikey in enumerate(self.scans): - for contrho, rho in enumerate(self.rhos): - labelsExtraPlot, labels_legend = [], [] - values_tot = ( - self.scans[ikey]["xV"][0] - / self.scans[ikey]["xV"][0, self.scans[ikey]["positionBase"]] - ) - for contVal, val in enumerate(values_tot): - fullres = ( - f"{ikey}/{self.scans[ikey]['variable']}_{round(val,7)}" - ) - labelsExtraPlot.append(fullres) - labels_legend.append(f"{round(val,7)}") + for contLabel, ikey in enumerate(labels): + labelsExtraPlot, labels_legend = [], [] + values_tot = ( + self.scans[ikey]["xV"][0] + / self.scans[ikey]["xV"][0, self.scans[ikey]["positionBase"]] + ) + for val in values_tot: + fullres = f"{ikey}_{round(val,7)}" + labelsExtraPlot.append(fullres) + labels_legend.append(f"{round(val,7)}") - colorsC = np.transpose( - np.array(colorsChosen[contLabel]), axes=(1, 0, 2) - ) - colorsC_tuple = [] - for i in range(colorsC.shape[0]): - for j in range(colorsC.shape[1]): - colorsC_tuple.append(tuple(colorsC[i, j, :])) + colorsC = np.transpose( + np.array(colorsChosen[contLabel]), axes=(1, 0, 2) + ) + colorsC_tuple = [] + for i in range(colorsC.shape[0]): + for j in range(colorsC.shape[1]): + colorsC_tuple.append(tuple(colorsC[i, j, :])) - self.plotRun( + self.plot( extratitle=f"{ikey} - ", labels=labelsExtraPlot, labels_legend=labels_legend, @@ -2798,9 +2816,6 @@ def plotScan( plotGACODE=False, ) - if figs is None: - self.fn.show() - # ~~~~~~~~~~~~~~ Extra complete analysis options def runScanTurbulenceDrives( @@ -2818,23 +2833,55 @@ def runScanTurbulenceDrives( varUpDown = np.linspace(1 - variation, 1 + variation, resolutionPoints) - for variable in self.variablesDrives: - name = subFolderTGLF + "_" + variable + for cont, variable in enumerate(self.variablesDrives): + # Only ask the restart in the first round + kwargs_TGLFrun["forceIfRestart"] = cont > 0 or ( + "forceIfRestart" in kwargs_TGLFrun and kwargs_TGLFrun["forceIfRestart"] + ) + self.runScan( - subFolderTGLF=name + "/", + subFolderTGLF=subFolderTGLF, variable=variable, varUpDown=varUpDown, **kwargs_TGLFrun, ) - self.readScan(label=name, variable=variable) + self.readScan(label=f"{subFolderTGLF}_{variable}", variable=variable) - def plotScanTurbulenceDrives(self, label="scan1", figs=None): + def plotScanTurbulenceDrives(self, label="scan1", figs=None, **kwargs_TGLFscanPlot): labels = [] for variable in self.variablesDrives: labels.append(f"{label}_{variable}") - self.plotScan(labels=labels, figs=figs, variableLabel="X", relativeX=True) + if figs is None: + self.fn = GUItools.FigureNotebook( + "TGLF Drives MITIM Notebook", geometry="1500x900", vertical=True + ) + fig1 = self.fn.add_figure(label="Fluxes - Relative") + fig2 = self.fn.add_figure(label="Fluxes (GB) - Relative") + fig3 = self.fn.add_figure(label="Linear Stability - Relative") + figs1 = [fig1, fig2, fig3] + fig1 = self.fn.add_figure(label="Fluxes") + fig2 = self.fn.add_figure(label="Fluxes (GB)") + fig3 = self.fn.add_figure(label="Linear Stability") + figs2 = [fig1, fig2, fig3] + else: + figs1, figs2 = None, None + + kwargs_TGLFscanPlot.pop("figs", None) + + self.plotScan( + labels=labels, + figs=figs1, + variableLabel="X", + relativeX=True, + **kwargs_TGLFscanPlot, + ) + + kwargs_TGLFscanPlot["plotTGLFs"] = False + self.plotScan( + labels=labels, figs=figs2, variableLabel="X", **kwargs_TGLFscanPlot + ) def runAnalysis( self, @@ -3015,10 +3062,8 @@ def runAnalysis( def plotAnalysis(self, labels=["analysis1"], analysisType="chi_e", figs=None): if figs is None: - plt.rcParams["figure.max_open_warning"] = False - plt.ioff() self.fn = GUItools.FigureNotebook( - 0, "TGLF Analysis MITIM Notebook", geometry="1500x900" + "TGLF Analysis MITIM Notebook", geometry="1500x900" ) fig1 = self.fn.add_figure(label="Analysis") fig2 = self.fn.add_figure(label="Fluxes") @@ -3319,16 +3364,14 @@ def plotAnalysis(self, labels=["analysis1"], analysisType="chi_e", figs=None): ax.set_title(f"Integrated profile using BC={BC}") GRAPHICStools.addDenseAxis(ax) - if figs is None: - self.fn.show() - def updateConvolution(self): self.DRMAJDX_LOC = {} if "latest_inputsFileTGLFDict" not in self.__dict__: for i in self.rhos: self.DRMAJDX_LOC[i] = 0.0 print( - " ~~ Using DRMAJDX_LOC=0 because no input file was stored", typeMsg="w" + "\t- [convolution] Using DRMAJDX_LOC =0 because no input file was stored", + typeMsg="i", ) else: for i in self.latest_inputsFileTGLFDict: @@ -3413,7 +3456,7 @@ def changeANDwrite_TGLF( modInputTGLF = {} ns_max = [] - for rho in rhos: + for i, rho in enumerate(rhos): print(f"\t- Changing input file for rho={rho:.4f}") NS = inputs[rho].plasma["NS"] inputTGLF_rho = modifyInputToTGLF( @@ -3422,6 +3465,7 @@ def changeANDwrite_TGLF( extraOptions=extraOptions, multipliers=multipliers, NS=NS, + position_change=i, ) newfile = f"{FolderTGLF}/input.tglf_{rho:.4f}" @@ -4032,7 +4076,12 @@ def identifySpecie(dict_species, dict_find): def modifyInputToTGLF( - inputTGLF, TGLFsettings=None, extraOptions={}, multipliers={}, NS=2 + inputTGLF, + TGLFsettings=None, + extraOptions={}, + multipliers={}, + NS=2, + position_change=0, ): if TGLFsettings is not None: _, TGLFoptions, label = GACODEdefaults.addTGLFcontrol( @@ -4043,7 +4092,6 @@ def modifyInputToTGLF( print( f" \t- Using presets TGLFsettings = {TGLFsettings} ({label})", typeMsg="i" ) - TGLFoptions_orig = copy.deepcopy(inputTGLF.controls) inputTGLF.controls = TGLFoptions else: @@ -4058,6 +4106,11 @@ def modifyInputToTGLF( if len(extraOptions) > 0: print("\t- External options:") for ikey in extraOptions: + if isinstance(extraOptions[ikey], (list, np.ndarray)): + value_to_change_to = extraOptions[ikey][position_change] + else: + value_to_change_to = extraOptions[ikey] + # is a specie one? try: isspecie = ikey.split("_")[0] in inputTGLF.species[1] @@ -4068,20 +4121,20 @@ def modifyInputToTGLF( specie = int(ikey.split("_")[-1]) varK = "_".join(ikey.split("_")[:-1]) var_orig = inputTGLF.species[specie][varK] - var_new = extraOptions[ikey] + var_new = value_to_change_to inputTGLF.species[specie][varK] = var_new else: if ikey in inputTGLF.controls: var_orig = inputTGLF.controls[ikey] - var_new = extraOptions[ikey] + var_new = value_to_change_to inputTGLF.controls[ikey] = var_new elif ikey in inputTGLF.geom: var_orig = inputTGLF.geom[ikey] - var_new = extraOptions[ikey] + var_new = value_to_change_to inputTGLF.geom[ikey] = var_new elif ikey in inputTGLF.plasma: var_orig = inputTGLF.plasma[ikey] - var_new = extraOptions[ikey] + var_new = value_to_change_to inputTGLF.plasma[ikey] = var_new else: # If the variable in extraOptions wasn't in there, consider it a control param @@ -4090,7 +4143,7 @@ def modifyInputToTGLF( typeMsg="i", ) var_orig = None - var_new = extraOptions[ikey] + var_new = value_to_change_to inputTGLF.controls[ikey] = var_new print( @@ -4215,7 +4268,7 @@ def __init__(self, FolderGACODE, suffix=""): f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} with suffix {suffix}" ) - self.inputclass = TGLFinput(file=self.FolderGACODE + "input.tglf" + self.suffix) + self.inputclass = TGLFinput(file=f"{self.FolderGACODE}/input.tglf{self.suffix}") self.roa = self.inputclass.geom["RMIN_LOC"] self.read() diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index b903917d..e2dec886 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -1,13 +1,11 @@ -import os, copy, pdb +import os +import copy import numpy as np import matplotlib.pyplot as plt -from collections import OrderedDict from IPython import embed from mitim_tools.misc_tools import ( IOtools, - FARMINGtools, GRAPHICStools, - MATHtools, PLASMAtools, ) from mitim_tools.gacode_tools import TGLFtools, PROFILEStools @@ -55,7 +53,7 @@ def __init__(self, cdf=None, time=100.0, avTime=0.0): if cdf is not None: _, self.nameRunid = IOtools.getLocInfo(self.LocationCDF) else: - self.nameRunid = "12345A01" + self.nameRunid = "mitim" self.time, self.avTime = time, avTime @@ -127,9 +125,7 @@ def prep( if (not os.path.exists(self.FolderGACODE)) or restart: print( - "\t- Folder {0} does not exist, or restart has been requested... creating folder to store".format( - FolderGACODE - ) + f"\t- Folder {FolderGACODE} does not exist, or restart has been requested... creating folder to store" ) IOtools.askNewFolder( self.FolderGACODE, force=forceIfRestart or (not restart) @@ -214,6 +210,17 @@ def run( ) # limitSpecies is not to consider the last ions. In there are 6 ions in input.gacode, and limitSpecies=4, the last 2 won't be considered quasineutrality = TGYRO_physics_options.get("quasineutrality", []) + if ( + (iterations == 0) + and ("tgyro_method" in TGYRO_solver_options) + and (TGYRO_solver_options["tgyro_method"] != 1) + ): + print( + f"\t- Zero iteration must be run with method 1, changing it from {TGYRO_solver_options['tgyro_method']} to 1", + typeMsg="w", + ) + TGYRO_solver_options["tgyro_method"] = 1 + # ------------------------------------------------------------------------ print("\n> Run TGYRO") @@ -231,10 +238,14 @@ def run( ] if vectorRange[2] == 1: - print(" -> TGYRO cannot run with only one point, adding 0.9", typeMsg="w") - vectorRange[2], vectorRange[1] = 2, 0.9 + extra_add = vectorRange[0] + 1e-3 + print( + f" -> TGYRO cannot run with only one point ({vectorRange[0]}), adding {extra_add}", + typeMsg="w", + ) + vectorRange[2], vectorRange[1] = 2, extra_add if special_radii_mod is not None: - special_radii_mod = np.append(special_radii_mod, [0.9]) + special_radii_mod = np.append(special_radii_mod, [extra_add]) self.rhosToSimulate = ( special_radii_mod @@ -320,7 +331,7 @@ def run( # Do the required files exist? exists = not restart - txt_nonexist = '' + txt_nonexist = "" if exists: for j in self.outputFiles: file = f"{self.FolderTGYRO}{j}" @@ -334,7 +345,7 @@ def run( "\t- Some of the required output files did not exist, running TGYRO", typeMsg="i", ) - #print(txt_nonexist, typeMsg="w") + # print(txt_nonexist, typeMsg="w") # ---------------------------------------------------------------- # ---------------------------------------------------------------- @@ -503,8 +514,8 @@ def run( inputgacode_new.writeCurrentStatus() # ------------------------------------------------------------------------------------------------------------------------ - # Copy those files that I'm interested in into the main folder - for file in self.outputFiles: + # Copy those files that I'm interested in, plus the extra file, into the main folder + for file in self.outputFiles + ["input.tgyro"]: os.system(f"cp {self.FolderTGYRO_tmp}/{file} {self.FolderTGYRO}/{file}") # Rename the input.tglf.news to the actual rho they where at @@ -674,7 +685,7 @@ def runTGLFsensitivities(self, fromlabel="tgyro1", rho=0.5, restart=False): rho: TGLFtools.TGLFinput(file=f"{self.FolderTGYRO}/input.tglf_{rho:.4f}") } - self.tglf[fromlabel] = TGLFtools.TGLF(rhos=rhos) + self.tglf[fromlabel] = TGLFtools.TGLF(rhos=[rho]) self.tglf[fromlabel].prep( "tglf_runs/", specificInputs=inputsTGLF, @@ -759,16 +770,12 @@ def run_tglf_scan( return res - def plotRun(self, fn=None, labels=["tgyro1"], doNotShow=False): + def plot(self, fn=None, labels=["tgyro1"], doNotShow=False, fn_color=None): if fn is None: - plt.rcParams["figure.max_open_warning"] = False from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - self.fn = FigureNotebook(0, "TGYRO Output Notebook", geometry="1800x900") - fnProvided = False + self.fn = FigureNotebook("TGYRO Output Notebook", geometry="1800x900") else: - fnProvided = True self.fn = fn # **** Plot the TGYRO output for all the labels @@ -783,7 +790,7 @@ def plotRun(self, fn=None, labels=["tgyro1"], doNotShow=False): # **** Final real - fig1 = self.fn.add_figure(label="Final Comp.") + fig1 = self.fn.add_figure(tab_color=fn_color, label="Final Comp.") grid = plt.GridSpec(3, 4, hspace=0.4, wspace=0.3) ax00 = fig1.add_subplot(grid[0, 0]) ax01 = fig1.add_subplot(grid[0, 1]) @@ -1058,12 +1065,12 @@ def plotRun(self, fn=None, labels=["tgyro1"], doNotShow=False): GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax, bottomy=0.0) - figProf_1 = self.fn.add_figure(label="GACODE-Prof.") - figProf_2 = self.fn.add_figure(label="GACODE-Power") - figProf_3 = self.fn.add_figure(label="GACODE-Geom.") - figProf_4 = self.fn.add_figure(label="GACODE-Grad.") - figProf_6 = self.fn.add_figure(label="GACODE-Other") - figProf_7 = self.fn.add_figure(label="GACODE-Impurities") + figProf_1 = self.fn.add_figure(tab_color=fn_color, label="GACODE-Prof.") + figProf_2 = self.fn.add_figure(tab_color=fn_color, label="GACODE-Power") + figProf_3 = self.fn.add_figure(tab_color=fn_color, label="GACODE-Geom.") + figProf_4 = self.fn.add_figure(tab_color=fn_color, label="GACODE-Grad.") + figProf_6 = self.fn.add_figure(tab_color=fn_color, label="GACODE-Other") + figProf_7 = self.fn.add_figure(tab_color=fn_color, label="GACODE-Impurities") grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) axsProf_1 = [ @@ -1170,14 +1177,16 @@ def plotRun(self, fn=None, labels=["tgyro1"], doNotShow=False): print("Could not plot profiles_final", typeMsg="w") try: - figFlows = self.fn.add_figure(label="GACODE-FlowsFin") + figFlows = self.fn.add_figure( + tab_color=fn_color, label="GACODE-FlowsFin" + ) self.results[label].plotBalance(fig=figFlows) except: print("Could not plot final flows", typeMsg="w") for label in labels: if label in self.tglf: - self.tglf[label].plotRun( + self.tglf[label].plot( fn=self.fn, labels=[f"{self.nameRuns_default}_tglf1"], fontsizeLeg=5, @@ -1186,32 +1195,6 @@ def plotRun(self, fn=None, labels=["tgyro1"], doNotShow=False): plotGACODE=False, ) - if (not fnProvided) and (not doNotShow): - self.fn.show() - - def analysisTGLF(self, label): - # TO FIX - - inputgacode = self.results[label].profiles_final.file - - # rhos = [0.5] # rho locations - # TGLFsettings = 4 # settings to run tglf - # restartTGLF = False # Restart the TGLF run - - # # --- Workflow - - # workingFolder = '~/PRF/REGS/scratchs/tglf_case6/' - # if not os.path.exists(workingFolder): IOtools.askNewFolder(workingFolder) - # inputgacode_new = '{0}/input.gacode'.format(workingFolder) - # os.system('cp {0} {1}'.format(inputgacode,inputgacode_new)) - - # tglf = TGLFtools.TGLF(rhos=rhos) - # cdf = tglf.prep(workingFolder,restart=False) - # tglf.run(subFolderTGLF='run1/',TGLFsettings=TGLFsettings,restart=restartTGLF) - # tglf.read(label='run1') - - # tglf.plotRun(labels=['run1']) - class TGYROoutput: def __init__(self, FolderTGYRO, profiles=None): @@ -1811,6 +1794,12 @@ def readprocess(self): ev_profs, ev_extras = nprof, nr elif self.inputs["TGYRO_ITERATION_METHOD"] == "6": ev_profs, ev_extras = 1, 0 + else: + print( + f"\t- TGYRO_ITERATION_METHOD={self.inputs['TGYRO_ITERATION_METHOD']} logic not implemented yet, assuming same as 1", + typeMsg="w", + ) + ev_profs, ev_extras = nprof, nr ev_calls = ev_profs * nr + ev_extras @@ -2141,10 +2130,7 @@ def derived(self): ) def useFineGridTargets(self, impurityPosition=1): - print( - "\t* It has been requested that I calculate targets based on the fine grid of input.gacode.new", - typeMsg="i", - ) + print("\t\t\t* Recalculating targets on the fine grid of input.gacode.new") if self.profiles_final is None: print( @@ -2189,23 +2175,22 @@ def useFineGridTargets(self, impurityPosition=1): def TGYROmodeledVariables(self, **kwargs): return PORTALSinteraction.TGYROmodeledVariables(self, **kwargs) - def plot(self, fn=None, label="", prelabel=""): + def plot(self, fn=None, label="", prelabel="", fn_color=None): if fn is None: - plt.rcParams["figure.max_open_warning"] = False from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - fn = FigureNotebook(0, "TGYRO Output Notebook", geometry="1800x900") - fnprovided = False + self.fn = FigureNotebook("TGYRO Output Notebook", geometry="1800x900") else: - fnprovided = True + self.fn = fn # ------------------------------------------------------------------------------ # Summary 1 # ------------------------------------------------------------------------------ - fig1 = fn.add_figure(label=prelabel + "Overview" + label) + fig1 = self.fn.add_figure( + tab_color=fn_color, label=prelabel + "Overview" + label + ) grid = plt.GridSpec(3, 4, hspace=0.45, wspace=0.3) ax00 = fig1.add_subplot(grid[:, 0]) @@ -2305,10 +2290,6 @@ def plot(self, fn=None, label="", prelabel=""): ax.set_xlabel("$r/a$") ax.set_xlim([0, 1]) ax.set_ylabel("$a/L_T$") - min0 = np.min(self.aLte[0]) - min1 = np.min(self.aLte[-1]) - max0 = np.max(self.aLte[0]) - max1 = np.max(self.aLte[-1]) ax.set_title("Electron Temperature Gradient") GRAPHICStools.addDenseAxis(ax) @@ -2435,10 +2416,6 @@ def plot(self, fn=None, label="", prelabel=""): ax.set_xlabel("$r/a$") ax.set_xlim([0, 1]) ax.set_ylabel("$a/L_T$") - min0 = np.min(self.aLti[0]) - min1 = np.min(self.aLti[-1]) - max0 = np.max(self.aLti[0]) - max1 = np.max(self.aLti[-1]) ax.set_title("Ion Temperature Gradient") GRAPHICStools.addDenseAxis(ax) @@ -2568,10 +2545,6 @@ def plot(self, fn=None, label="", prelabel=""): ax.set_xlabel("$r/a$") ax.set_xlim([0, 1]) ax.set_ylabel("$a/L_n$") - min0 = np.min(self.aLne[0]) - min1 = np.min(self.aLne[-1]) - max0 = np.max(self.aLne[0]) - max1 = np.max(self.aLne[-1]) ax.set_title("Density Gradient") GRAPHICStools.addDenseAxis(ax) @@ -2625,7 +2598,7 @@ def plot(self, fn=None, label="", prelabel=""): # Summary 2 # ------------------------------------------------------------------------------ - fig1 = fn.add_figure(label=prelabel + "Match" + label) + fig1 = self.fn.add_figure(tab_color=fn_color, label=prelabel + "Match" + label) grid = plt.GridSpec(4, 4, hspace=0.45, wspace=0.3) ax00 = fig1.add_subplot(grid[0, 0]) @@ -3216,7 +3189,9 @@ def plot(self, fn=None, label="", prelabel=""): # Convergence # ------------------------------------------------------------------------------ - fig1 = fn.add_figure(label=prelabel + "Convergence" + label) + fig1 = self.fn.add_figure( + tab_color=fn_color, label=prelabel + "Convergence" + label + ) try: self.plotConvergence(fig1=fig1) @@ -3230,7 +3205,7 @@ def plot(self, fn=None, label="", prelabel=""): # Fluxes # ------------------------------------------------------------------------------ - fig1 = fn.add_figure(label=prelabel + "Fluxes" + label) + fig1 = self.fn.add_figure(tab_color=fn_color, label=prelabel + "Fluxes" + label) grid = plt.GridSpec(4, self.Qi_sim_turb.shape[0] + 3, hspace=0.3, wspace=0.3) @@ -3328,7 +3303,7 @@ def plot(self, fn=None, label="", prelabel=""): # ax.set_xlabel('$r/a$') ax.axhline(y=0, lw=0.5, c="k", ls="--") ax.legend(prop={"size": 6}, loc="best") - ax.set_title("Ions".format(i + 1)) + ax.set_title("Ions") GRAPHICStools.addDenseAxis(ax) @@ -3633,7 +3608,7 @@ def plot(self, fn=None, label="", prelabel=""): # Powers # ------------------------------------------------------------------------------ - fig1 = fn.add_figure(label=prelabel + "Powers" + label) + fig1 = self.fn.add_figure(tab_color=fn_color, label=prelabel + "Powers" + label) grid = plt.GridSpec(2, 5, hspace=0.2, wspace=0.4) ax00 = fig1.add_subplot(grid[0, 0]) @@ -3823,11 +3798,14 @@ def plot(self, fn=None, label="", prelabel=""): "--o", c="blue", label="exch", - markersize=1, - lw=0.5, + markersize=3, ) ax.plot( - self.roa[-1], self.Qe_tarMW_exch[-1], "--o", c="red", markersize=1, lw=0.5 + self.roa[-1], + self.Qe_tarMW_exch[-1], + "--o", + c="red", + markersize=3, ) ax.plot( self.roa[0], @@ -3835,11 +3813,14 @@ def plot(self, fn=None, label="", prelabel=""): "-.o", c="blue", label="expwd", - markersize=1, - lw=0.5, + markersize=3, ) ax.plot( - self.roa[-1], self.Qe_tarMW_expwd[-1], "-.o", c="red", markersize=1, lw=0.5 + self.roa[-1], + self.Qe_tarMW_expwd[-1], + "-.o", + c="red", + markersize=3, ) ax.plot( @@ -4143,7 +4124,9 @@ def plot(self, fn=None, label="", prelabel=""): # Metrics # ------------------------------------------------------------------------------ - fig1 = fn.add_figure(label=prelabel + "Perform." + label) + fig1 = self.fn.add_figure( + tab_color=fn_color, label=prelabel + "Perform." + label + ) grid = plt.GridSpec(2, 2, hspace=0.45, wspace=0.3) # Fusion Gain @@ -4344,12 +4327,11 @@ def plot(self, fn=None, label="", prelabel=""): # Final # ------------------------------------------------------------------------------ - fig1 = fn.add_figure(label=prelabel + "Flows" + label) + fig1 = self.fn.add_figure( + tab_color=fn_color, label=prelabel + "Flows" + label + ) self.plotBalance(fig=fig1) - if not fnprovided: - fn.show() - """ Note that input.gacode and TGYRO may differ in Pfus, Prad, Exch, etc because TGYRO has internal calculations. I believe those are passed to the output of tgyro, even though the tgyro option is static. @@ -4435,7 +4417,7 @@ def plotConvergence(self, fig1=None): ax.set_ylabel("Residual (GB)") ax.set_yscale("log") whichticks = ax.get_xticks() - ax2 = GRAPHICStools.addXaxis( + _ = GRAPHICStools.addXaxis( ax, self.iterations, self.calls_solver, @@ -4527,7 +4509,7 @@ def plotConvergence(self, fig1=None): GRAPHICStools.addDenseAxis(ax) # GRAPHICStools.autoscale_y(ax) - ax2 = GRAPHICStools.addXaxis( + _ = GRAPHICStools.addXaxis( ax, self.iterations, self.calls_solver, @@ -4577,16 +4559,9 @@ def plotConvergence(self, fig1=None): def plotAll(TGYROoutputs, labels=None, fn=None): if fn is None: - # fig = plt.figure(figsize=(15,9)) - plt.rcParams["figure.max_open_warning"] = False from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - fn = FigureNotebook(0, "TGYRO Output Notebook", geometry="1800x900") - fnprovided = False - - else: - fnprovided = True + fn = FigureNotebook("TGYRO Output Notebook", geometry="1800x900") if labels is None: labels = [f" {i}" for i in np.arange(1, len(TGYROoutputs) + 1, 1)] @@ -4594,9 +4569,6 @@ def plotAll(TGYROoutputs, labels=None, fn=None): for i, TGYROoutput in enumerate(TGYROoutputs): TGYROoutput.plot(fn=fn, label=labels[i]) - if not fnprovided: - fn.show() - class TGYROinput: def __init__(self, input_profiles, file=None, onlyThermal=False, limitSpecies=100): @@ -4738,9 +4710,11 @@ def produceInputs_TGYROworkflow( print("\t- Testing... do TGYRO files already exist?") if os.path.exists(f"{finalFolder}/{file_to_look}"): ProfilesGenerated, StateGenerated = True, True - print(f"\t\t+++++++ {file_to_look} already generated") + print(f"\t\t+++++++ {IOtools.clipstr(file_to_look)} already generated") else: - print(f"\t\t+++++++ {finalFolder}/{file_to_look} file not found") + print( + f"\t\t+++++++ {IOtools.clipstr(f'{finalFolder}/{file_to_look}')} file not found" + ) ProfilesGenerated = False if os.path.exists(finalFolder + "/10001.cdf"): StateGenerated = True diff --git a/src/mitim_tools/gacode_tools/aux/GACODEdefaults.py b/src/mitim_tools/gacode_tools/aux/GACODEdefaults.py index a2538f24..720b11dc 100644 --- a/src/mitim_tools/gacode_tools/aux/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/aux/GACODEdefaults.py @@ -1,152 +1,12 @@ +import json import numpy as np +from mitim_tools.misc_tools import IOtools + from IPython import embed -from mitim_tools.misc_tools import FARMINGtools, IOtools from mitim_tools.misc_tools.IOtools import printMsg as print -def constructStandardTGLF(NS=2): - """ - Consistent with localdump of TGLF - """ - - TGLFoptions = { - "USE_TRANSPORT_MODEL": True, - "GEOMETRY_FLAG": 1, # 0: s-a, 1: Miller, 2: Fourier, 3: ELITE - "ADIABATIC_ELEC": False, # T: Use adiabatic electrons. - "SIGN_IT": 1, # Sign of 𝐼𝑇 with repsect to CCW toroidal direction from top. - "SIGN_BT": 1, # Sign of 𝐡𝑇 with repsect to CCW toroidal direction from top. - "IFLUX": True, # Compute quasilinear weights and mode amplitudes. - "THETA_TRAPPED": 0.7, # Parameter to adjust trapped fraction model. - "NN_MAX_ERROR": -1.0, # Threshold for TGLF-NN execution versus full TGLF calculation - } - - """ - Wavenumber grid - ----------------------- - """ - TGLFoptions["KYGRID_MODEL"] = 1 # 1: Standard grid - TGLFoptions["NKY"] = 19 # Number of poloidal modes in the high-k spectrum - TGLFoptions["KY"] = 0.3 # ky for single-mode call to TGLF. - - """ - Hermite basis functions - ----------------------- - """ - TGLFoptions[ - "FIND_WIDTH" - ] = True # T: Find the width that maximizes the growth rate, F: Use WIDTH - TGLFoptions[ - "USE_BISECTION" - ] = True # T: Use bisection search method to find width that maximizes growth rate - TGLFoptions[ - "WIDTH" - ] = 1.65 # Max. width to search for Hermite polynomial basis (1.65 default, <1.5 recommended in Staebler PoP05) - TGLFoptions[ - "WIDTH_MIN" - ] = 0.3 # Min. width to search for Hermite polynomial basis (0.3 default) - TGLFoptions["NWIDTH"] = 21 # Number of search points - TGLFoptions["NBASIS_MIN"] = 2 # Minimum number of parallel basis functions - TGLFoptions["NBASIS_MAX"] = 4 # Maximum number of parallel basis functions - TGLFoptions["NXGRID"] = 16 # Number of nodes in Gauss-Hermite quadrature. - - """ - Modes options - ----------------------- - """ - TGLFoptions[ - "IBRANCH" - ] = ( - -1 - ) # 0 = find two most unstable modes one for each sign of frequency, electron drift direction (1), ion drift direction (2), - # -1 = sort the unstable modes by growthrate in rank order - TGLFoptions["NMODES"] = int( - NS + 2 - ) # For IBRANCH=-1, number of modes to store. ASTRA uses NS+2 - - """ - Saturation model - ----------------------- - """ - TGLFoptions["SAT_RULE"] = 0 - TGLFoptions["ETG_FACTOR"] = 1.25 - - """ - Physics included - ----------------------- - """ - TGLFoptions[ - "USE_BPER" - ] = False # Include transverse magnetic fluctuations, 𝛿𝐴‖. -> PERPENDICULAR FLUCTUATIONS - TGLFoptions[ - "USE_BPAR" - ] = False # Include compressional magnetic fluctuations, \delta B_{\lVert }} -> COMPRESSIONAL EFFECTS - TGLFoptions[ - "USE_MHD_RULE" - ] = False # Ignore pressure gradient contribution to curvature drift. - TGLFoptions["XNU_MODEL"] = 2 # Collision model (2=new) - TGLFoptions["VPAR_MODEL"] = 0 # 0=low-Mach-number limit (DEPRECATED?) - TGLFoptions["VPAR_SHEAR_MODEL"] = 1 - - """ - ExB shear model - ----------------------- - """ - TGLFoptions[ - "ALPHA_QUENCH" - ] = 0.0 # 1.0 = use quench rule, 0.0 = use new spectral shift model (0.0 recommeded by Gary, 05/11/2020) - TGLFoptions[ - "ALPHA_E" - ] = 1.0 # Multiplies ExB velocity shear for spectral shift model (1.0 ecommeded by Gary, 05/11/2020) - - TGLFoptions["ALPHA_MACH"] = 0.0 - TGLFoptions["ALPHA_ZF"] = 1.0 - - """ - Multipliers - ----------------------- - """ - - TGLFoptions["DEBYE_FACTOR"] = 1.0 # Multiplies the debye length - TGLFoptions[ - "XNU_FACTOR" - ] = 1.0 # Multiplies the trapped/passing boundary electron-ion collision terms - TGLFoptions["ALPHA_P"] = 1.0 # Multiplies parallel velocity shear for all species - TGLFoptions["PARK"] = 1.0 # Multiplies the parallel gradient term. - TGLFoptions["GHAT"] = 1.0 # Multiplies the curvature drift closure terms. - TGLFoptions["GCHAT"] = 1.0 # Multiplies the curvature drift irreducible terms. - - # WHY? - TGLFoptions["DAMP_PSI"] = 0.0 # Damping factor for psi - TGLFoptions["DAMP_SIG"] = 0.0 # Damping factor for sig - TGLFoptions["LINSKER_FACTOR"] = 0.0 # Multiplies the Linsker terms - TGLFoptions["GRADB_FACTOR"] = 0.0 # Multiplies the gradB terms - - TGLFoptions[ - "WD_ZERO" - ] = 0.1 # Cutoff for curvature drift eigenvalues to prevent zero. - TGLFoptions[ - "FILTER" - ] = 2.0 # Sets threshold for frequency/drift frequency to filter out non-driftwave instabilities. - TGLFoptions["WDIA_TRAPPED"] = 0.0 - - """ - Other Options - ----------------------- - """ - - TGLFoptions[ - "USE_AVE_ION_GRID" - ] = False # T:Use weighted average charge of ions for the gyroradius reference, F: Use first ion - TGLFoptions[ - "USE_INBOARD_DETRAPPED" - ] = False # Set trapped fraction to zero if eigenmode is inward ballooning. - - TGLFoptions["RLNP_CUTOFF"] = 18.0 # Limits SAT2 factor of R/Lp < RLNP_CUTOFF - - return TGLFoptions - - def addTGLFcontrol(TGLFsettings=1, NS=2, minimal=False): """ ******************************************************************************** @@ -171,7 +31,10 @@ def addTGLFcontrol(TGLFsettings=1, NS=2, minimal=False): # Define every flag else: - TGLFoptions = constructStandardTGLF(NS=NS) + TGLFoptions = IOtools.generateMITIMNamelist( + "$MITIM_PATH/templates/input.tglf.controls", caseInsensitive=False + ) + TGLFoptions["NMODES"] = NS + 2 """ ******************************************************************************** @@ -180,81 +43,22 @@ def addTGLFcontrol(TGLFsettings=1, NS=2, minimal=False): ******************************************************************************** """ - label = "unspecified" - - # SAT1 ----- Old SAT1 standard to recover previous results - if TGLFsettings == 1: - TGLFoptions["SAT_RULE"] = 1 - TGLFoptions["UNITS"] = "GYRO" - - label = "SAT1" - - # SAT0 ----- Old SAT0 standard to recover previous results - if TGLFsettings == 2: - TGLFoptions["SAT_RULE"] = 0 - TGLFoptions[ - "UNITS" - ] = "GYRO" # In SAT0, this is the only option, see tglf.startup.f90 - TGLFoptions["ETG_FACTOR"] = 1.25 - - label = "SAT0" - - # SAT1geo ----- SAT1 standard (CGYRO) - if TGLFsettings == 3: - TGLFoptions["SAT_RULE"] = 1 - TGLFoptions["UNITS"] = "CGYRO" - - label = "SAT1geo" - - # SAT2 ----- SAT2 standard - if TGLFsettings == 4: - label = "SAT2" - - # SAT2 - TGLFoptions["SAT_RULE"] = 2 - TGLFoptions[ - "UNITS" - ] = "CGYRO" # In SAT2, CGYRO/GENE are the only options, see tglf.startup.f90 - TGLFoptions["XNU_MODEL"] = 3 # This is forced anyway, see tglf.startup.f90 - TGLFoptions["WDIA_TRAPPED"] = 1.0 # This is forced anyway, see tglf.startup.f90 - - # SAT2em - if TGLFsettings == 5: - label = "SAT2em" - - # SAT2 - TGLFoptions["SAT_RULE"] = 2 - TGLFoptions[ - "UNITS" - ] = "CGYRO" # In SAT2, CGYRO/GENE are the only options, see tglf.startup.f90 - TGLFoptions["XNU_MODEL"] = 3 # This is forced anyway, see tglf.startup.f90 - TGLFoptions["WDIA_TRAPPED"] = 1.0 # This is forced anyway, see tglf.startup.f90 - - # EM - TGLFoptions["USE_BPER"] = True - - # ------------------------ - # PRF's Experiments - # ------------------------ - - # SAT2em with higher resolution - if TGLFsettings == 101: - # SAT2 - TGLFoptions["SAT_RULE"] = 2 - TGLFoptions[ - "UNITS" - ] = "CGYRO" # In SAT2, CGYRO/GENE are the only options, see tglf.startup.f90 - TGLFoptions["XNU_MODEL"] = 3 # This is forced anyway, see tglf.startup.f90 - TGLFoptions["WDIA_TRAPPED"] = 1.0 # This is forced anyway, see tglf.startup.f90 - - # EM - TGLFoptions["USE_BPER"] = True - - # Extra - TGLFoptions["KYGRID_MODEL"] = 4 - TGLFoptions["NBASIS_MAX"] = 6 - - label = "SAT2em basis" + with open( + IOtools.expandPath("$MITIM_PATH/templates/input.tglf.models.json"), "r" + ) as f: + settings = json.load(f) + + if str(TGLFsettings) in settings: + sett = settings[str(TGLFsettings)] + label = sett["label"] + for ikey in sett["controls"]: + TGLFoptions[ikey] = sett["controls"][ikey] + else: + print( + "\t- TGLFsettings not found in input.tglf.models.json, using defaults", + typeMsg="w", + ) + label = "unspecified" # -------------------------------- # From dictionary to text diff --git a/src/mitim_tools/gacode_tools/aux/GACODEinterpret.py b/src/mitim_tools/gacode_tools/aux/GACODEinterpret.py index 170cc56d..c2f5d552 100644 --- a/src/mitim_tools/gacode_tools/aux/GACODEinterpret.py +++ b/src/mitim_tools/gacode_tools/aux/GACODEinterpret.py @@ -68,12 +68,12 @@ def Waveform_read(file, fileOut): aux = f.readlines() for ir in range(len(aux)): - if "ky:" in aux[ir]: break - + if "ky:" in aux[ir]: + break ky = float(aux[ir].split()[-1]) g, f = [], [] - for i in range(len(aux[ir+2:])): + for i in range(len(aux[ir + 2 :])): a = aux[ir + 2 + i].split() f.append(float(a[1])) g.append(float(a[2])) diff --git a/src/mitim_tools/gacode_tools/aux/GACODErun.py b/src/mitim_tools/gacode_tools/aux/GACODErun.py index 1babe626..b331c69e 100644 --- a/src/mitim_tools/gacode_tools/aux/GACODErun.py +++ b/src/mitim_tools/gacode_tools/aux/GACODErun.py @@ -1,17 +1,15 @@ -import os, copy, time +import os +import time import numpy as np import matplotlib.pyplot as plt from scipy.interpolate import interp1d -from IPython import embed from mitim_tools.transp_tools.tools import PLASMASTATEtools from mitim_tools.misc_tools import FARMINGtools, IOtools, MATHtools, GRAPHICStools -from mitim_tools.gacode_tools.aux import GACODEdefaults -from mitim_tools.misc_tools import CONFIGread - from mitim_tools.misc_tools.IOtools import printMsg as print - from mitim_tools.misc_tools.CONFIGread import read_verbose_level +from IPython import embed + verbose_level = read_verbose_level() @@ -32,14 +30,20 @@ def runTGYRO( launchSlurm = False -> Launch locally as a bash script """ - gacode_compilation = "gacode" - # This routine assumes that folderWork contains input.profiles and input.tgyro already - machineSettings = CONFIGread.machineSettings( - code="tgyro", - gacode_compilation=gacode_compilation, - nameScratch=f"mitim_tmp_{nameRunid}/", + tgyro_job = FARMINGtools.mitim_job(folderWork) + + tgyro_job.define_machine( + "tgyro", + f"mitim_{nameRunid}/", + launchSlurm=launchSlurm, + slurm_settings={ + "minutes": minutes, + "ntasks": 1, + "name": nameJob, + "cpuspertask": nparallel, + }, ) # ------ Run TGYRO @@ -52,7 +56,7 @@ def runTGYRO( # Execution command # --------------- - folderExecution = machineSettings["folderWork"] + folderExecution = tgyro_job.machineSettings["folderWork"] TGYROcommand = f"tgyro -e . -n {nparallel} -p {folderExecution}" @@ -73,29 +77,20 @@ def runTGYRO( "done", ] - # --------------- + # --------------------------------------------- # Execute - # --------------- - - ntasks = 1 - cores_tgyro = nparallel + # --------------------------------------------- - # machineSettings['clear'] = False - FARMINGtools.SLURMcomplete( + tgyro_job.prep( TGYROcommand, - folderWork, - inputFiles, - outputFiles, - minutes, - ntasks, - nameJob, - machineSettings, - cpuspertask=cores_tgyro, + input_files=inputFiles, + output_files=outputFiles, shellPreCommands=shellPreCommands, shellPostCommands=shellPostCommands, - launchSlurm=launchSlurm, ) + tgyro_job.run() + def findNamelist(LocationCDF, folderWork=None, nameRunid="10000", ForceFirst=True): # ----------------------------------------------------------- @@ -151,18 +146,6 @@ def prepareTGYRO( ): nameWork = "10001" - machineSettingsTRXPL = CONFIGread.machineSettings( - code="trxpl", nameScratch=f"mitim_tmp_{nameRunid}/" - ) - - gacode_compilation = "gacode" - - machineSettingsPROF = CONFIGread.machineSettings( - code="profiles_gen", - gacode_compilation=gacode_compilation, - nameScratch=f"mitim_tmp_{nameRunid}/", - ) - if not StateGenerated: print("\t- Running TRXPL to extract g-file and plasmastate") CDFtoTRXPLoutput( @@ -173,7 +156,6 @@ def prepareTGYRO( BtIp_dirs=BtIp_dirs, nameOutputs=nameWork, folderWork=folderWork, - machineSettings=machineSettingsTRXPL, grids=gridsTRXPL, sendState=sendState, ) @@ -181,7 +163,6 @@ def prepareTGYRO( print("\t- Running PROFILES_GEN to generate input.profiles/input.gacode files") runPROFILES_GEN( folderWork, - machineSettings=machineSettingsPROF, nameFiles=nameWork, UsePRFmodification=fixPlasmaState, includeGEQ=includeGEQ, @@ -196,7 +177,6 @@ def CDFtoTRXPLoutput( BtIp_dirs=[0, 0], nameOutputs="10001", folderWork="~/scratchFolder/", - machineSettings={"machine": "local", "folderWork": "~/"}, grids=[151, 101, 101], sendState=True, ): @@ -215,7 +195,6 @@ def CDFtoTRXPLoutput( IpDir=BtIp_dirs[1], avTime=avTime, nameFiles=nameFiles, - machineSettings=machineSettings, sendState=sendState, nameOutputs=nameOutputs, grids=grids, @@ -236,7 +215,6 @@ def CDFtoTRXPLoutput( IpDir=BtIp_dirs[1], avTime=avTime, nameFiles=nameFiles, - machineSettings=machineSettings, sendState=sendState, nameOutputs=nameOutputs, grids=grids, @@ -247,20 +225,25 @@ def executeCGYRO( FolderCGYRO, linesCGYRO, fileProfiles, - gacode_compilation="gacode", outputFiles=["out.cgyro.run"], name="", numcores=32, ): - machineSettings = CONFIGread.machineSettings( - code="cgyro", - nameScratch=f"mitim_tmp_{name}/", - gacode_compilation=gacode_compilation, - ) - if not os.path.exists(FolderCGYRO): os.system(f"mkdir {FolderCGYRO}") + cgyro_job = FARMINGtools.mitim_job(FolderCGYRO) + + cgyro_job.define_machine( + "cgyro", + f"mitim_cgyro_{name}/", + slurm_settings={ + "minutes": 60, + "ntasks": numcores, + "name": name, + }, + ) + # --------------- # Prepare files # --------------- @@ -273,7 +256,7 @@ def executeCGYRO( # Execution command # --------------- - folderExecution = machineSettings["folderWork"] + folderExecution = cgyro_job.machineSettings["folderWork"] CGYROcommand = f"cgyro -e . -n {numcores} -p {folderExecution}" shellPreCommands = [] @@ -282,23 +265,15 @@ def executeCGYRO( # Execute # --------------- - inputFiles = [fileProfiles, fileCGYRO] - minutes = 60 # NTH - nparallel = numcores - nameJob = name - - FARMINGtools.SLURMcomplete( + cgyro_job.prep( CGYROcommand, - FolderCGYRO, - inputFiles, - outputFiles, - minutes, - nparallel, - nameJob, - machineSettings, + input_files=[fileProfiles, fileCGYRO], + output_files=outputFiles, shellPreCommands=shellPreCommands, ) + cgyro_job.run() + def runTRXPL( FolderTRXPL, @@ -308,7 +283,6 @@ def runTRXPL( avTime=0.0, nameFiles="10000", nameOutputs="10001", - machineSettings={"machine": "local", "folderWork": "~/"}, sendState=True, grids=[151, 101, 101], ): @@ -352,21 +326,26 @@ def runTRXPL( ), typeMsg="i", ) - command = f"cd {machineSettings['folderWork']} && trxpl < trxpl.in" - FARMINGtools.runCommand( + + trxpl_job = FARMINGtools.mitim_job(FolderTRXPL) + trxpl_job.define_machine( + "trxpl", + f"mitim_trxpl_{nameOutputs}/", + ) + + command = "trxpl < trxpl.in" + + trxpl_job.prep( command, - inputFiles, - outputFiles=[f"{nameOutputs}.cdf", f"{nameOutputs}.geq"], - whereOutput=FolderTRXPL, - machineSettings=machineSettings, - pauseBeforeCommand=not sendState, + input_files=inputFiles, + output_files=[f"{nameOutputs}.cdf", f"{nameOutputs}.geq"], ) + trxpl_job.run() def runPROFILES_GEN( FolderTGLF, nameFiles="10001", - machineSettings={"machine": "local", "folderWork": "~/"}, UsePRFmodification=False, includeGEQ=True, ): @@ -386,8 +365,6 @@ def runPROFILES_GEN( if includeGEQ: inputFiles.append(FolderTGLF + f"{nameFiles}.geq") - command = f"cd {machineSettings['folderWork']} && bash profiles_gen.sh" - # **** Write command txt = f"profiles_gen -i {nameFiles}.cdf" if includeGEQ: @@ -399,13 +376,19 @@ def runPROFILES_GEN( # ****************** print(f"\t\t- Proceeding to run PROFILES_GEN with: {txt}") - FARMINGtools.runCommand( - command, - inputFiles, - outputFiles=["input.gacode"], - whereOutput=FolderTGLF, - machineSettings=machineSettings, + + pgen_job = FARMINGtools.mitim_job(FolderTGLF) + pgen_job.define_machine( + "profiles_gen", + f"mitim_profiles_gen_{nameFiles}/", + ) + + pgen_job.prep( + "bash profiles_gen.sh", + input_files=inputFiles, + output_files=["input.gacode"], ) + pgen_job.run() if ( runWithoutEqIfFail @@ -424,28 +407,15 @@ def runPROFILES_GEN( # ****************** print(f"\t\t- Proceeding to run PROFILES_GEN with: {txt}") - FARMINGtools.runCommand( - command, - inputFiles, - outputFiles=["input.gacode"], - whereOutput=FolderTGLF, - machineSettings=machineSettings, - ) + pgen_job.run() def runVGEN( - inputgacode_file, workingFolder, numcores=32, - minutes=15, - vgenOptions={ - "er": 2, - "vel": 1, - "numspecies": None, - "matched_ion": 1, - "nth": "17,39", - }, - nameChange="", + minutes=60, + vgenOptions={}, + name_run="vgen1", ): """ Driver for the vgen (velocity-generation) capability of NEO. @@ -466,40 +436,52 @@ def runVGEN( -nth: Minimum and maximum theta resolutions (e.g. 17,39) """ + vgen_job = FARMINGtools.mitim_job(workingFolder) + + vgen_job.define_machine( + "profiles_gen", + f"mitim_vgen_{name_run}/", + slurm_settings={ + "minutes": minutes, + "ntasks": numcores, + "name": f"neo_vgen_{name_run}", + }, + ) + + print( + f"\t- Running NEO (with {vgenOptions['numspecies']} species) to populate w0(rad/s) in input.gacode file" + ) + print(f"\t\t> Matching ion {vgenOptions['matched_ion']} Vtor") + options = f"-er {vgenOptions['er']} -vel {vgenOptions['vel']} -in {vgenOptions['numspecies']} -ix {vgenOptions['matched_ion']} -nth {vgenOptions['nth']}" # *********************************** print( - f"\t\t- Proceeding to generate Er from NEO run using profiles_gen ({options})" + f"\t\t- Proceeding to generate Er from NEO run using profiles_gen -vgen ({options})" ) + inputgacode_file = f"{workingFolder}/input.gacode" + _, nameFile = IOtools.reducePathLevel(inputgacode_file, level=1, isItFile=True) - machineSettings = CONFIGread.machineSettings( - code="profiles_gen", nameScratch=f"mitim_tmp_vgen_{nameChange}_{nameFile}/" - ) - command = f"cd {machineSettings['folderWork']} && bash profiles_vgen.sh" + command = f"cd {vgen_job.machineSettings['folderWork']} && bash profiles_vgen.sh" with open(f"{workingFolder}/profiles_vgen.sh", "w") as f: f.write(f"profiles_gen -vgen -i {nameFile} {options} -n {numcores}") - FARMINGtools.SLURMcomplete( + # --------------- + # Execute + # --------------- + + vgen_job.prep( command, - workingFolder, - [inputgacode_file, f"{workingFolder}/profiles_vgen.sh"], - ["slurm_output.dat", "slurm_error.dat"], - minutes, - numcores, - f"neo_vgen_{nameChange}_{nameFile}", - machineSettings, - outputFolders=["vgen/"], + input_files=[inputgacode_file, f"{workingFolder}/profiles_vgen.sh"], + output_files=["slurm_output.dat", "slurm_error.dat"], ) - if len(nameChange) > 0: - os.system(f"mv {workingFolder}/vgen {workingFolder}/vgen_{nameChange}") - file_new = f"{workingFolder}/vgen_{nameChange}/input.gacode" - else: - file_new = f"{workingFolder}/vgen/input.gacode" + vgen_job.run() + + file_new = f"{workingFolder}/vgen/input.gacode" return file_new @@ -520,7 +502,7 @@ def buildDictFromInput(inputFile): parsed[splits[0].split()[0]] = splits[1].split()[0] for i in parsed: - if type(parsed[i]) == str: + if isinstance(parsed[i], str): if ( parsed[i].lower() == "t" or parsed[i].lower() == "true" @@ -823,7 +805,6 @@ def runCGYRO( extraFlag="", filesToRetrieve=["cgyro/out.cgyro.info"], name="", - gacode_compilation="gacode", ): MaxRunsWithoutWait, timewaitSendRun = 4, 10 @@ -836,7 +817,6 @@ def runCGYRO( tmpFolder, inputFilesCGYRO[rho], inputGacode, - gacode_compilation=gacode_compilation, numcores=numcores, outputFiles=filesToRetrieve, name=name, @@ -865,14 +845,12 @@ def runCGYRO( def runTGLF( rhos, finalFolder, - tmpFolder, inputFilesTGLF, minutes=5, cores_tglf=4, extraFlag="", filesToRetrieve=["out.tglf.gbflux"], name="", - gacode_compilation="gacode", launchSlurm=True, ): """ @@ -880,22 +858,15 @@ def runTGLF( launchSlurm = False -> Launch locally as a bash script """ - machineSettings = CONFIGread.machineSettings( - code="tglf", - nameScratch=f"mitim_tmp_{name}/", - gacode_compilation=gacode_compilation, - ) - - if not os.path.exists(tmpFolder): - os.system(f"mkdir {tmpFolder}") + tmpFolder = finalFolder + "/tmp_tglf/" + IOtools.askNewFolder(tmpFolder) - if not launchSlurm: - machineSettings["machine"], machineSettings["folderWork"] = ( - "local", - tmpFolder + "tmp_run/", - ) + tglf_job = FARMINGtools.mitim_job(tmpFolder) - folderExecution = machineSettings["folderWork"] + tglf_job.define_machine_quick( + "tglf", + f"mitim_{name}/", + ) # --------------------------------------------- # Prepare files and folders @@ -922,7 +893,7 @@ def runTGLF( total_tglf_cores = int(cores_tglf * len(rhos)) - if launchSlurm and (machineSettings["partition"] is not None): + if launchSlurm and ("partition" in tglf_job.machineSettings["slurm"]): typeRun = "job" if total_tglf_cores <= 32 else "array" else: typeRun = "bash" @@ -930,9 +901,7 @@ def runTGLF( if typeRun in ["bash", "job"]: TGLFcommand = "" for i, rho in enumerate(rhos): - TGLFcommand += ( - f"tglf -e rho_{rho:.4f}/ -n {cores_tglf} -p {folderExecution}/ &\n" - ) + TGLFcommand += f"tglf -e rho_{rho:.4f}/ -n {cores_tglf} -p {tglf_job.folderExecution}/ &\n" TGLFcommand += ( "\nwait" # This is needed so that the script doesn't end before each job @@ -949,7 +918,7 @@ def runTGLF( ) rho_array = ",".join([f"{int(rho*1E4)}" for rho in rhos]) - TGLFcommand = f'tglf -e rho_0."$SLURM_ARRAY_TASK_ID"/ -n {cores_tglf} -p {folderExecution}/ 1> {folderExecution}/rho_0."$SLURM_ARRAY_TASK_ID"/slurm_output.dat 2> {folderExecution}/rho_0."$SLURM_ARRAY_TASK_ID"/slurm_error.dat\n' + TGLFcommand = f'tglf -e rho_0."$SLURM_ARRAY_TASK_ID"/ -n {cores_tglf} -p {tglf_job.folderExecution}/ 1> {tglf_job.folderExecution}/rho_0."$SLURM_ARRAY_TASK_ID"/slurm_output.dat 2> {tglf_job.folderExecution}/rho_0."$SLURM_ARRAY_TASK_ID"/slurm_error.dat\n' ntasks = 1 cpuspertask = cores_tglf @@ -958,25 +927,28 @@ def runTGLF( # Execute # --------------------------------------------- - # machineSettings["clear"] = False - FARMINGtools.SLURMcomplete( - TGLFcommand, - tmpFolder, - [], - [], - minutes, - ntasks, - name, - machineSettings, - cpuspertask=cpuspertask, - inputFolders=folders, - outputFolders=folders_red, - shellPreCommands=[], + tglf_job.define_machine( + "tglf", + f"mitim_{name}/", launchSlurm=launchSlurm, - job_array=rho_array, - nodes=1, + slurm_settings={ + "minutes": minutes, + "ntasks": ntasks, + "name": name, + "cpuspertask": cpuspertask, + "job_array": rho_array, + "nodes": 1, + }, + ) + + tglf_job.prep( + TGLFcommand, + input_folders=folders, + output_folders=folders_red, ) + tglf_job.run() + # --------------------------------------------- # Organize # --------------------------------------------- @@ -1001,13 +973,6 @@ def runTGLF( typeMsg="w", verbose=verbose_level, ) - # else: - # print( - # f"\t\t~ {file} successfully retrieved, converted into {original_file}", - # verbose=verbose_level, - # ) - - # os.system(f"mv {fileTGLF} {finalFolder}/rho_{rho:.4f}/input.tglf_{rho:.4f}") if fineall: print("\t\t- All files were successfully retrieved") diff --git a/src/mitim_tools/gacode_tools/aux/NORMtools.py b/src/mitim_tools/gacode_tools/aux/NORMtools.py index af1fabae..65204e7f 100644 --- a/src/mitim_tools/gacode_tools/aux/NORMtools.py +++ b/src/mitim_tools/gacode_tools/aux/NORMtools.py @@ -68,7 +68,7 @@ def normalizations( "exp_Ge_gb": Norms[norm_select]["exp_Ge"] / Norms[norm_select]["g_gb"], } else: - print("\t\t\t- Normalization set is empty", typeMsg="w") + print(f"\t\t- Normalization set ({norm_select}) is empty", typeMsg="w") return Norms, cdf @@ -127,7 +127,7 @@ def normalizations_profiles(profiles): return Set_norm, Set_norm["rho"], Set_norm["roa"], Set_norm["mi_ref"] else: - print("\t\t\t- Cannot read normalization", typeMsg="w") + print("\t\t- Cannot read normalization", typeMsg="w") return None, None, None, None diff --git a/src/mitim_tools/gacode_tools/aux/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/aux/PORTALSinteraction.py index 2e9d7030..ea21d361 100644 --- a/src/mitim_tools/gacode_tools/aux/PORTALSinteraction.py +++ b/src/mitim_tools/gacode_tools/aux/PORTALSinteraction.py @@ -135,12 +135,13 @@ def TGYROmodeledVariables( forceZeroParticleFlux=False, includeFast=False, impurityPosition=1, - ProfilesPredicted=["te", "ti" "ne"], - provideTurbulentExchange=True, UseFineGridTargets=False, OriginalFimp=1.0, dfT=torch.Tensor(), ): + """ + impurityPosition will be substracted one + """ if "tgyro_stds" not in self.__dict__: self.tgyro_stds = False @@ -305,8 +306,6 @@ def TGYROmodeledVariables( self.EXe_sim_turb_stds[:, :] if self.tgyro_stds else None ) - # portals_variables['PexchTurbFlux'] = self.Qe_tarMW_expwd/self.dvoldr # MW/m^2 - if forceZeroParticleFlux: portals_variables["Ge"] = self.Ge_tar[:, :] * 0.0 @@ -467,7 +466,7 @@ def calculatePseudos(var_dict, PORTALSparameters, TGYROparameters, powerstate): # Source term is (TARGET - TRANSPORT) source = cal - of - # Residual is defined as the negative (bc it's maximization) normalized (1/N) norm of radial & channel residuals -> L2 or L1 + # Residual is defined as the negative (bc it's maximization) normalized (1/N) norm of radial & channel residuals -> L2 res = -1 / source.shape[-1] * torch.norm(source, p=2, dim=-1) return of, cal, source, res diff --git a/src/mitim_tools/gacode_tools/exe/plot_cgyro.py b/src/mitim_tools/gacode_tools/exe/read_cgyro.py similarity index 68% rename from src/mitim_tools/gacode_tools/exe/plot_cgyro.py rename to src/mitim_tools/gacode_tools/exe/read_cgyro.py index 9fd2c7c7..0c64d097 100644 --- a/src/mitim_tools/gacode_tools/exe/plot_cgyro.py +++ b/src/mitim_tools/gacode_tools/exe/read_cgyro.py @@ -1,4 +1,5 @@ import argparse +from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import CGYROtools """ @@ -6,7 +7,7 @@ """ parser = argparse.ArgumentParser() -parser.add_argument("--folders", required=True, type=str, nargs="*") +parser.add_argument("folders", type=str, nargs="*") args = parser.parse_args() folders = args.folders @@ -16,7 +17,7 @@ labels = [] for i, folder in enumerate(folders): - labels.append(f"cgyro{i+1}") + labels.append(f"{IOtools.reducePathLevel(folder)[-1]}") c.read(label=labels[-1], folder=folder) c.plot(labels=labels) diff --git a/src/mitim_tools/gacode_tools/exe/read_gacodes.py b/src/mitim_tools/gacode_tools/exe/read_gacode.py similarity index 54% rename from src/mitim_tools/gacode_tools/exe/read_gacodes.py rename to src/mitim_tools/gacode_tools/exe/read_gacode.py index af504a54..682e4f67 100644 --- a/src/mitim_tools/gacode_tools/exe/read_gacodes.py +++ b/src/mitim_tools/gacode_tools/exe/read_gacode.py @@ -4,14 +4,18 @@ """ Quick way to plot several input.gacode files together e.g. - read_gacodes.py --files input.gacode1 input.gacode2 + read_gacodes.py input.gacode1 input.gacode2 [--rho 0.9] """ parser = argparse.ArgumentParser() -parser.add_argument("--files", required=True, type=str, nargs="*") +parser.add_argument("files", type=str, nargs="*") +parser.add_argument( + "--rho", type=float, required=False, default=0.89 +) # Last rho for gradients plot args = parser.parse_args() files = args.files +rho = args.rho # Read profs = [] @@ -23,4 +27,4 @@ # Plot -PROFILEStools.plotAll(profs) +fn = PROFILEStools.plotAll(profs, lastRhoGradients=rho) diff --git a/src/mitim_tools/gacode_tools/exe/read_tglf.py b/src/mitim_tools/gacode_tools/exe/read_tglf.py index 6f54546a..a67e90e0 100644 --- a/src/mitim_tools/gacode_tools/exe/read_tglf.py +++ b/src/mitim_tools/gacode_tools/exe/read_tglf.py @@ -1,7 +1,7 @@ """ This example reads TGLF from an already existing folder (no normalizations if no input_gacode file provided) - read_tglf.py --folder run0/ [--suffix _0.55] [--gacode input.gacode] + read_tglf.py run0/ [--suffixes _0.55] [--gacode input.gacode] """ @@ -10,21 +10,28 @@ from mitim_tools.gacode_tools import TGLFtools parser = argparse.ArgumentParser() -parser.add_argument("--folder", required=True, type=str) -parser.add_argument("--suffix", required=False, type=str, default=None) +parser.add_argument("folders", type=str, nargs="*") +parser.add_argument("--suffixes", required=False, type=str, nargs="*", default=None) parser.add_argument("--gacode", required=False, type=str, default=None) args = parser.parse_args() -folder = IOtools.expandPath(args.folder) +folders = [IOtools.expandPath(folder) for folder in args.folders] input_gacode = IOtools.expandPath(args.gacode) if args.gacode is not None else None -suffix = args.suffix +suffixes = args.suffixes -if suffix is None: - suffix = "" +if suffixes is None: + suffixes = ["" for _ in range(len(folders))] + +for i in range(len(suffixes)): + if suffixes[i] == "_": + suffixes[i] = "" tglf = TGLFtools.TGLF() -tglf.prep_from_tglf(folder, f"{folder}/input.tglf{suffix}", input_gacode=input_gacode) -tglf.read(folder=f"{folder}/", suffix=suffix, label="run1") +tglf.prep_from_tglf( + folders[0], f"{folders[0]}/input.tglf{suffixes[0]}", input_gacode=input_gacode +) +for i, folder in enumerate(folders): + tglf.read(folder=f"{folder}/", suffix=suffixes[i], label=f"run{i}") -tglf.plotRun(labels=["run1"]) +tglf.plot(labels=[f"run{i}" for i in range(len(folders))]) diff --git a/src/mitim_tools/gacode_tools/exe/read_tgyros.py b/src/mitim_tools/gacode_tools/exe/read_tgyro.py similarity index 89% rename from src/mitim_tools/gacode_tools/exe/read_tgyros.py rename to src/mitim_tools/gacode_tools/exe/read_tgyro.py index 26c6805e..71b9a1b7 100644 --- a/src/mitim_tools/gacode_tools/exe/read_tgyros.py +++ b/src/mitim_tools/gacode_tools/exe/read_tgyro.py @@ -8,7 +8,7 @@ """ parser = argparse.ArgumentParser() -parser.add_argument("--folders", required=True, type=str, nargs="*") +parser.add_argument("folders", type=str, nargs="*") args = parser.parse_args() folders = args.folders diff --git a/src/mitim_tools/gacode_tools/exe/run_profilesgen.py b/src/mitim_tools/gacode_tools/exe/run_profilesgen.py index 260a40e8..17fbb721 100644 --- a/src/mitim_tools/gacode_tools/exe/run_profilesgen.py +++ b/src/mitim_tools/gacode_tools/exe/run_profilesgen.py @@ -1,6 +1,6 @@ -import sys, os +import sys +import os from mitim_tools.misc_tools import IOtools -from mitim_tools.misc_tools import CONFIGread from mitim_tools.gacode_tools.aux import GACODErun """ @@ -13,13 +13,8 @@ nameWork = s2.split(".cdf")[0] -machineSettings = CONFIGread.machineSettings( - code="profiles_gen", nameScratch=f"mitim_tmp_profiles_gen_{nameWork}/" -) - GACODErun.runPROFILES_GEN( folderWork, - machineSettings=machineSettings, nameFiles=nameWork, UsePRFmodification=True, includeGEQ=True, diff --git a/src/mitim_tools/gacode_tools/exe/run_tglf.py b/src/mitim_tools/gacode_tools/exe/run_tglf.py index 436a1af2..336bbfeb 100644 --- a/src/mitim_tools/gacode_tools/exe/run_tglf.py +++ b/src/mitim_tools/gacode_tools/exe/run_tglf.py @@ -1,7 +1,7 @@ """ This example runs TGLF from an already existing file (no normalizations if no input_gacode file provided) - run_tglf.py --folder run0/ --tglf input.tglf [--gacode input.gacode] [--scan RLTS_2] [--drives True] + run_tglf.py --folder run0/ --tglf input.tglf [--gacode input.gacode] [--scan RLTS_2] [--drives] [--restart] Sequence: - If drives: do drives analysis @@ -20,11 +20,15 @@ # ------------------------------------------------------------------------------ parser = argparse.ArgumentParser() -parser.add_argument("--folder", required=True, type=str) -parser.add_argument("--tglf", required=True, type=str) +parser.add_argument("folder", type=str) +parser.add_argument("tglf", type=str) parser.add_argument("--gacode", required=False, type=str, default=None) parser.add_argument("--scan", required=False, type=str, default=None) -parser.add_argument("--drives", required=False, type=bool, default=False) +parser.add_argument("--drives", required=False, default=False, action="store_true") +parser.add_argument( + "--restart", "-r", required=False, default=False, action="store_true" +) + args = parser.parse_args() @@ -33,6 +37,7 @@ input_gacode = IOtools.expandPath(args.gacode) if args.gacode is not None else None scan = args.scan drives = args.drives +restart = args.restart # ------------------------------------------------------------------------------ # Preparation @@ -55,11 +60,12 @@ variable=scan, varUpDown=np.linspace(0.2, 2.0, 5), TGLFsettings=None, + restart=restart, ) tglf.readScan(label="scan1", variable=scan) tglf.plotScan(labels=["scan1"], variableLabel=scan) else: - tglf.run(subFolderTGLF="run1/", TGLFsettings=None) + tglf.run(subFolderTGLF="run1/", TGLFsettings=None, restart=restart) tglf.read(label="run1") - tglf.plotRun(labels=["run1"]) + tglf.plot(labels=["run1"]) diff --git a/src/mitim_tools/gs_tools/FREEGStools.py b/src/mitim_tools/gs_tools/FREEGStools.py index b1c2b2ec..edf8787e 100644 --- a/src/mitim_tools/gs_tools/FREEGStools.py +++ b/src/mitim_tools/gs_tools/FREEGStools.py @@ -186,7 +186,7 @@ def evaluator( # -------------------------------------------------------- if plot and prfs is not None: - GSplotting.plotResult( + fn = GSplotting.plotResult( prfs, metrics, Constraints, diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 73f4bf65..07848f56 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -245,24 +245,23 @@ def write(self, filename=None): def plot(self, fn=None, extraLabel=""): if fn is None: wasProvided = False - plt.rcParams["figure.max_open_warning"] = False + from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - fn = FigureNotebook(0, "GEQDSK Notebook", geometry="1600x1000") + self.fn = FigureNotebook("GEQDSK Notebook", geometry="1600x1000") else: wasProvided = True - + self.fn = fn # ----------------------------------------------------------------------------- # OMFIT Summary # ----------------------------------------------------------------------------- - # fig = fn.add_figure(label=extraLabel+'OMFIT Summ') + # fig = self.fn.add_figure(label=extraLabel+'OMFIT Summ') # self.g.plot() # ----------------------------------------------------------------------------- # Flux # ----------------------------------------------------------------------------- - fig = fn.add_figure(label=extraLabel + "Surfaces") + fig = self.fn.add_figure(label=extraLabel + "Surfaces") grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) ax1 = fig.add_subplot(grid[:, 0]) ax2 = fig.add_subplot(grid[:, 1]) @@ -274,7 +273,7 @@ def plot(self, fn=None, extraLabel=""): # ----------------------------------------------------------------------------- # Currents # ----------------------------------------------------------------------------- - fig = fn.add_figure(label=extraLabel + "Currents") + fig = self.fn.add_figure(label=extraLabel + "Currents") grid = plt.GridSpec(3, 5, hspace=0.3, wspace=0.3) ax1 = fig.add_subplot(grid[2, 0]) ax2 = fig.add_subplot(grid[:2, 0]) @@ -294,7 +293,7 @@ def plot(self, fn=None, extraLabel=""): # ----------------------------------------------------------------------------- # Fields # ----------------------------------------------------------------------------- - fig = fn.add_figure(label=extraLabel + "Fields") + fig = self.fn.add_figure(label=extraLabel + "Fields") grid = plt.GridSpec(3, 5, hspace=0.3, wspace=0.3) ax1 = fig.add_subplot(grid[2, 0]) ax2 = fig.add_subplot(grid[:2, 0]) @@ -314,7 +313,7 @@ def plot(self, fn=None, extraLabel=""): # ----------------------------------------------------------------------------- # Checks # ----------------------------------------------------------------------------- - fig = fn.add_figure(label=extraLabel + "GS Quality") + fig = self.fn.add_figure(label=extraLabel + "GS Quality") grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) ax1 = fig.add_subplot(grid[0, 0]) ax1E = ax1.twinx() @@ -328,7 +327,7 @@ def plot(self, fn=None, extraLabel=""): # ----------------------------------------------------------------------------- # Parameterization # ----------------------------------------------------------------------------- - fig = fn.add_figure(label=extraLabel + "Parameteriz.") + fig = self.fn.add_figure(label=extraLabel + "Parameteriz.") grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.3) ax1 = fig.add_subplot(grid[:, 0]) ax2 = fig.add_subplot(grid[0, 1]) @@ -341,7 +340,7 @@ def plot(self, fn=None, extraLabel=""): # ----------------------------------------------------------------------------- # Plasma # ----------------------------------------------------------------------------- - fig = fn.add_figure(label=extraLabel + "Plasma") + fig = self.fn.add_figure(label=extraLabel + "Plasma") grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) ax_plasma = [ @@ -359,7 +358,7 @@ def plot(self, fn=None, extraLabel=""): # ----------------------------------------------------------------------------- # Geometry # ----------------------------------------------------------------------------- - fig = fn.add_figure(label=extraLabel + "Geometry") + fig = self.fn.add_figure(label=extraLabel + "Geometry") grid = plt.GridSpec(2, 2, hspace=0.3, wspace=0.3) ax1 = fig.add_subplot(grid[0, 0]) ax2 = fig.add_subplot(grid[0, 1]) @@ -368,9 +367,6 @@ def plot(self, fn=None, extraLabel=""): self.plotGeometry(axs=[ax1, ax2, ax3, ax4]) - if not wasProvided: - fn.show() - return ax_plasma def plotFS(self, axs=None, color="b", label=""): @@ -1136,11 +1132,10 @@ def plotSurfaces( def compareGeqdsk(geqdsks, fn=None, extraLabel="", plotAll=True, labelsGs=None): if fn is None: wasProvided = False - plt.rcParams["figure.max_open_warning"] = False + from mitim_tools.misc_tools.GUItools import FigureNotebook - plt.ioff() - fn = FigureNotebook(0, "GEQDSK Notebook", geometry="1600x1000") + fn = FigureNotebook("GEQDSK Notebook", geometry="1600x1000") else: wasProvided = True @@ -1203,10 +1198,7 @@ def compareGeqdsk(geqdsks, fn=None, extraLabel="", plotAll=True, labelsGs=None): label=f"{labelsGs[i]} ", ) - if not wasProvided: - fn.show() - - return ax_plasma + return ax_plasma, fn def plotEnclosed(Rmajor, a, Zmajor, kappaU, kappaL, deltaU, deltaL, ax=None, c="k"): diff --git a/src/mitim_tools/gs_tools/aux/GSplotting.py b/src/mitim_tools/gs_tools/aux/GSplotting.py index 97693a4b..f0803961 100644 --- a/src/mitim_tools/gs_tools/aux/GSplotting.py +++ b/src/mitim_tools/gs_tools/aux/GSplotting.py @@ -52,7 +52,7 @@ def writeResults( gs = write_gfiles(prfs, whereOutput) if plotGs: - GEQtools.compareGeqdsk(gs, fn=fn, plotAll=True, labelsGs=None) + axs, fn = GEQtools.compareGeqdsk(gs, fn=fn, plotAll=True, labelsGs=None) # ------------- Pickles file = f"{whereOutput}/{namePkl}.pkl" @@ -168,12 +168,10 @@ def plotResult( cols = cols[colorOffset:] - plt.ioff() if figs is None: - fnProvided = False from mitim_tools.misc_tools.GUItools import FigureNotebook - fn = FigureNotebook(0, "FreeGS Notebook", geometry="1600x1000") + fn = FigureNotebook("FreeGS Notebook", geometry="1600x1000") fig1 = fn.add_figure(label="Eq. & Coils") fig2 = fn.add_figure(label="Metrics") @@ -185,7 +183,6 @@ def plotResult( figRes = fn.add_figure(label="Summary") else: - fnProvided = True fig1 = figs[0] fig2 = figs[1] fig3 = figs[2] @@ -514,8 +511,7 @@ def plotResult( except: pass - if not fnProvided: - fn.show() + return fn def plotMetricsPRFS( diff --git a/src/mitim_tools/gs_tools/exe/read_eqs.py b/src/mitim_tools/gs_tools/exe/read_eq.py similarity index 66% rename from src/mitim_tools/gs_tools/exe/read_eqs.py rename to src/mitim_tools/gs_tools/exe/read_eq.py index 2660dc78..dd1941a7 100644 --- a/src/mitim_tools/gs_tools/exe/read_eqs.py +++ b/src/mitim_tools/gs_tools/exe/read_eq.py @@ -1,8 +1,8 @@ -import sys, argparse +import argparse from mitim_tools.gs_tools import GEQtools parser = argparse.ArgumentParser() -parser.add_argument("--files", required=True, type=str, nargs="*") +parser.add_argument("files", type=str, nargs="*") args = parser.parse_args() files = args.files @@ -14,4 +14,4 @@ if len(gs) == 1: gs[0].plot() else: - GEQtools.compareGeqdsk(gs) + axs, fn = GEQtools.compareGeqdsk(gs) diff --git a/src/mitim_tools/im_tools/modules/EQmodule.py b/src/mitim_tools/im_tools/modules/EQmodule.py index 1407c0e6..32f5e07e 100644 --- a/src/mitim_tools/im_tools/modules/EQmodule.py +++ b/src/mitim_tools/im_tools/modules/EQmodule.py @@ -204,7 +204,6 @@ def writeBoundaryFiles( # ~~~ Plot boundary if saveFigure: - plt.ioff() fig, ax = plt.subplots() ax.plot(rs, zs, "-o", markersize=0.5, lw=0.5, c="m", label="LCFS") else: @@ -270,18 +269,22 @@ def generateMRY( f">> Running scruncher with {momentsScruncher} moments to create MRY file from boundary files..." ) - machineSettings = CONFIGread.machineSettings( - code="scruncher", nameScratch=f"tmp_scruncher_{name}/" + scruncher_job = FARMINGtools.mitim_job(FolderEquilibrium) + + scruncher_job.define_machine( + "scruncher", + f"tmp_scruncher_{name}/", + launchSlurm=False, ) - FARMINGtools.runCommand( - f"cd {machineSettings['folderWork']} && scruncher < scrunch_in", - filesInput, - outputFiles=["M123456.MRY"], - whereOutput=FolderEquilibrium, - machineSettings=machineSettings, + scruncher_job.prep( + "scruncher < scrunch_in", + output_files=["M123456.MRY"], + input_files=filesInput, ) + scruncher_job.run() + fileUF = f"{FolderMRY}/PRF{nameBaseShot}.MRY" os.system(f"cp {FolderEquilibrium}/M123456.MRY {fileUF}") @@ -1025,7 +1028,6 @@ def extractBoundaries_TSC( # ~~~ Plot boundary if saveFigure and not onlyExtractGfiles: - plt.ioff() fig, ax = plt.subplots() ax.plot(rs, zs, lw=4, c="m", label="LCFS") for ikey in rsT: @@ -1070,7 +1072,6 @@ def extractBoundaries_Sweep(gfileNames, times, folderOutput, saveFigure=True): # ~~~ Plot boundary if saveFigure: - plt.ioff() fig, ax = plt.subplots() ax.plot(rs, zs, lw=4, c="m", label="LCFS") ax.plot(rs1, zs1, lw=3, c="b", label="e1") diff --git a/src/mitim_tools/im_tools/modules/PEDmodule.py b/src/mitim_tools/im_tools/modules/PEDmodule.py index 3e12b595..72351b2a 100644 --- a/src/mitim_tools/im_tools/modules/PEDmodule.py +++ b/src/mitim_tools/im_tools/modules/PEDmodule.py @@ -755,15 +755,19 @@ def fit_pedestal_mtanh( """ - machineSettings = CONFIGread.machineSettings( - code="idl", nameScratch=f"mitim_tmp_idl_{nameRunid}/" + pedestal_job = FARMINGtools.mitim_job(folderWork) + + pedestal_job.define_machine( + "idl", + f"mitim_idl_{nameRunid}/", + launchSlurm=False, ) - path = machineSettings["folderWork"] + path = pedestal_job.folderExecution plasmastate_path = path + IOtools.reducePathLevel(plasmastate, isItFile=True)[-1] with open(folderWork + "/idl_in", "w") as f: - f.write(f".r /home/nthoward/SPARC_mtanh/make_pedestal_profiles_portals.pro\n") + f.write(".r /home/nthoward/SPARC_mtanh/make_pedestal_profiles_portals.pro\n") f.write( f"make_profiles,[{width_top},{netop},{p1},{ptop}],'{plasmastate_path}','{path}'\n\n" ) @@ -774,17 +778,17 @@ def fit_pedestal_mtanh( command = f"cd {path} && {start} && idl < idl_in" inputFiles = [folderWork + "/idl_in", plasmastate] - outputFiles = ["/mtanh_fits"] + outputFiles = ["mtanh_fits"] print(f"\t\t- Proceeding to run idl pedestal fitter (psi_pol = {width_top:.3f})") - FARMINGtools.runCommand( + + pedestal_job.prep( command, - inputFiles, - outputFiles=outputFiles, - whereOutput=folderWork, - machineSettings=machineSettings, - timeoutSecs=30, - ) # For some reason, if I don't time out, it just doesn't come back... + output_files=outputFiles, + input_files=inputFiles, + ) + + pedestal_job.run(timeoutSecs=30) x, ne, Te, Ti = read_mtanh(folderWork + "/mtanh_fits") diff --git a/src/mitim_tools/im_tools/outputs/database.py b/src/mitim_tools/im_tools/outputs/database.py index f19a438b..4a020fae 100644 --- a/src/mitim_tools/im_tools/outputs/database.py +++ b/src/mitim_tools/im_tools/outputs/database.py @@ -13,7 +13,6 @@ def plotDB(cdfs, runIDs=None, fn=None, extralab="", timePlot=None, avTime=0.0): runIDs = np.arange(0, len(cdfs)) if fn is None: - plt.ioff() figsize = (20, 10) fig1 = plt.figure(figsize=figsize) fig4 = plt.figure(figsize=figsize) @@ -810,6 +809,5 @@ def addlab(ax, label, typeA=1): except: print("problem " + str(i)) - plt.ioff() - fn = GUItools.FigureNotebook(0, "SPARC scenarios", geometry="1700x900") + fn = GUItools.FigureNotebook("SPARC scenarios", geometry="1700x900") figs = plotDB(cdfs, runIDs=runIDs, fn=fn) diff --git a/src/mitim_tools/im_tools/outputs/plotALL.py b/src/mitim_tools/im_tools/outputs/plotALL.py index 1086bb43..2caee0fe 100644 --- a/src/mitim_tools/im_tools/outputs/plotALL.py +++ b/src/mitim_tools/im_tools/outputs/plotALL.py @@ -78,8 +78,8 @@ # ------ Plot Summary of all runs, and convergence summary -plt.ioff() -fn = GUItools.FigureNotebook(0, "ALL", geometry="1700x900") + +fn = GUItools.FigureNotebook("ALL", geometry="1700x900") fig1 = fn.add_figure(label="Data") axSummary = [] @@ -116,7 +116,7 @@ cdfs_p_clean.append(c) -# fn = GUItools.FigureNotebook(0,'SPARC scenarios',geometry='1700x900') +# fn = GUItools.FigureNotebook('SPARC scenarios',geometry='1700x900') figs = plotDB( cdfs_clean, @@ -134,5 +134,3 @@ timePlot=timePlot, avTime=avTime, ) - -fn.show() diff --git a/src/mitim_tools/im_tools/outputs/plotCompare.py b/src/mitim_tools/im_tools/outputs/plotCompare.py index 9085ee6b..aac4209f 100644 --- a/src/mitim_tools/im_tools/outputs/plotCompare.py +++ b/src/mitim_tools/im_tools/outputs/plotCompare.py @@ -29,8 +29,8 @@ # ------ Plot Summary of all runs, and convergence summary -plt.ioff() -fn = GUItools.FigureNotebook(0, "ALL", geometry="1700x900") + +fn = GUItools.FigureNotebook("ALL", geometry="1700x900") fig1 = fn.add_figure(label="Data") axSummary = [] @@ -70,7 +70,7 @@ # from analysis_tools.SPARCdivertor.dataPlot import plotDB -# fn = GUItools.FigureNotebook(0,'SPARC scenarios',geometry='1700x900') +# fn = GUItools.FigureNotebook('SPARC scenarios',geometry='1700x900') figs = plotDB( cdfs_clean, diff --git a/src/mitim_tools/misc_tools/CONFIGread.py b/src/mitim_tools/misc_tools/CONFIGread.py index 6f18473b..3798e67e 100644 --- a/src/mitim_tools/misc_tools/CONFIGread.py +++ b/src/mitim_tools/misc_tools/CONFIGread.py @@ -1,4 +1,10 @@ -import os, json, socket, warnings, logging, getpass, numpy +import os +import json +import socket +import warnings +import logging +import getpass +from mitim_tools.misc_tools import IOtools from IPython import embed # PRF Note: Do not load IOtools, otherwise circularity problem @@ -67,7 +73,6 @@ def isThisEngaging(): def machineSettings( code="tgyro", - gacode_compilation="gacode", nameScratch="mitim_tmp/", forceUsername=None, ): @@ -82,9 +87,9 @@ def machineSettings( machine = s["preferences"][code] """ - Set-up per code and machine - ------------------------------------------------- - """ + Set-up per code and machine + ------------------------------------------------- + """ if forceUsername is not None: username = forceUsername @@ -96,24 +101,28 @@ def machineSettings( machineSettings = { "machine": s[machine]["machine"], "user": username, - "clear": True, "tunnel": None, "port": None, "identity": None, - "partition": None, - "modules": "source $MITIM_PATH/config/mitim.bashrc", + "modules": "source $MITIM_PATH/config/mitim.bashrc", "folderWork": scratch, - "exclude": s[machine]["exclude"] if "exclude" in s[machine] else None, + "slurm": {}, "isTunnelSameMachine": bool(s[machine]["isTunnelSameMachine"]) if "isTunnelSameMachine" in s[machine] else False, } # I can give extra things to load in the config file - if "modules" in s[machine] and s[machine]["modules"] is not None and s[machine]["modules"] != "": - machineSettings["modules"] = f'{machineSettings["modules"]}\n{s[machine]["modules"]}' + if ( + "modules" in s[machine] + and s[machine]["modules"] is not None + and s[machine]["modules"] != "" + ): + machineSettings[ + "modules" + ] = f'{machineSettings["modules"]}\n{s[machine]["modules"]}' - checkers = ["identity", "partition", "tunnel", "port"] + checkers = ["slurm", "identity", "tunnel", "port"] for i in checkers: if i in s[machine]: machineSettings[i] = s[machine][i] @@ -138,6 +147,11 @@ def machineSettings( # ************************************************************************************************************************ + if machineSettings["machine"] == "local": + machineSettings["folderWork"] = IOtools.expandPath( + machineSettings["folderWork"] + ) + if forceUsername is not None: machineSettings["identity"] = f"~/.ssh/id_rsa_{forceUsername}" diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 5bd1b761..48ad6a4e 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -1,16 +1,25 @@ """ - - Set of tools to farm out simulations to run in either remote clusters or locally, serially or parallel - +Set of tools to farm out simulations to run in either remote clusters or locally, serially or parallel """ -import os, sys, subprocess, socket, signal, pickle, datetime, torch, copy +from tqdm import tqdm +import os +import time +import sys +import subprocess +import socket +import signal +import datetime +import torch +import copy +import paramiko +import tarfile import numpy as np -from IPython import embed from contextlib import contextmanager -from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools import IOtools, CONFIGread from mitim_tools.misc_tools.IOtools import printMsg as print from mitim_tools.misc_tools.CONFIGread import read_verbose_level +from IPython import embed verbose_level = read_verbose_level() @@ -21,797 +30,705 @@ UseCUDAifAvailable = True -""" - Timeout function - - I just need to add a function: - with timeout(secs): - do things - - I don't recommend that "do things" includes a context manager like with Popen or it won't work... not sure why """ +New handling of jobs in remote or local clusters. Example use: + + folderWork = path_to_local_folder + + # Define job + job = FARMINGtools.mitim_job(folderWork) + + # Define machine + job.define_machine( + code_name, # must be defined in config_user.json + job_name, # name you want to give to the job + slurm_settings={ + 'minutes':minutes, + 'ntasks':ntasks, + 'name':job_name, + }, + ) + # Prepare job (remember that you can get job.folderExecution, which is where the job will be executed remotely) + job.prep( + Command_to_execute_as_string, # e.g. f'cd {job.folderExecution} && python3 job.py' + output_files=outputFiles, + input_files=inputFiles, + input_folders=inputFolders, + output_folders=outputFolders, + ) -def raise_timeout(signum, frame): - raise TimeoutError + # Run job + job.run() +""" -@contextmanager -def timeout(time, proc=None): - time = int(time) - if time < 1e6: - print( - f'\t\t* Note: this process will be killed if time exceeds {time}sec of execution ({datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")})', - typeMsg="i", - ) - # Register a function to raise a TimeoutError on the signal. - signal.signal(signal.SIGALRM, raise_timeout) - # Schedule the signal to be sent after ``time``. - signal.alarm(time) +class mitim_job: + def __init__(self, folder_local): + self.folder_local = folder_local + self.jobid = None + + def define_machine( + self, + code, + nameScratch, + launchSlurm=True, + slurm_settings={}, + ): + # Separated in case I need to quickly grab the machine settings + self.define_machine_quick(code, nameScratch, slurm_settings=slurm_settings) + + self.launchSlurm = launchSlurm + + if self.launchSlurm and (len(self.machineSettings["slurm"]) == 0): + self.launchSlurm = False + print( + "\t- slurm requested but no slurm setup to this machine in config... not doing slurm", + typeMsg="w", + ) - try: - yield - except TimeoutError: - print( - f'\t\t\t* Killing process! ({datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")})', - typeMsg="w", + # Print Slurm info + if self.launchSlurm: + print("\t- Slurm Settings:") + print("\t\t- Job settings:") + for key in self.slurm_settings: + if self.slurm_settings[key] is not None: + print(f"\t\t\t- {key}: {self.slurm_settings[key]}") + print("\t\t- Partition settings:") + print(f'\t\t\t- machine: {self.machineSettings["machine"]}') + print(f'\t\t\t- username: {self.machineSettings["user"]}') + for key in self.machineSettings["slurm"]: + print(f'\t\t\t- {key}: {self.machineSettings["slurm"][key]}') + + def define_machine_quick(self, code, nameScratch, slurm_settings={}): + self.slurm_settings = slurm_settings + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Defaults for slurm + self.slurm_settings.setdefault("name", "mitim_job") + self.slurm_settings.setdefault("minutes", 10) + self.slurm_settings.setdefault("cpuspertask", 1) + self.slurm_settings.setdefault("ntasks", 1) + self.slurm_settings.setdefault("nodes", None) + self.slurm_settings.setdefault("job_array", None) + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self.machineSettings = CONFIGread.machineSettings( + code=code, + nameScratch=nameScratch, ) - if proc is not None: - proc.kill() - outs, errs = proc.communicate() - finally: - # Unregister the signal so it won't be triggered if the timeout is not reached. - signal.signal(signal.SIGALRM, signal.SIG_IGN) + self.folderExecution = self.machineSettings["folderWork"] + def prep( + self, + command, + input_files=[], + input_folders=[], + output_files=[], + output_folders=[], + shellPreCommands=[], + shellPostCommands=[], + label_log_files="", + ): + """ + command: + Option 1: string with commands to execute separated by && + Option 2: list of strings with commands to execute + """ + + # Pass to class + self.command = command + self.input_files = input_files + self.input_folders = input_folders + self.output_files = output_files + self.output_folders = output_folders + + self.shellPreCommands = shellPreCommands + self.shellPostCommands = shellPostCommands + self.label_log_files = label_log_files + + def run(self, waitYN=True, timeoutSecs=1e6, removeScratchFolders=True): + if not waitYN: + removeScratchFolders = False + + # ****** Prepare SLURM job ***************************** + comm, fileSBTACH, fileSHELL = create_slurm_execution_files( + self.command, + self.machineSettings["folderWork"], + self.machineSettings["modules"], + job_array=self.slurm_settings["job_array"], + folder_local=self.folder_local, + shellPreCommands=self.shellPreCommands, + shellPostCommands=self.shellPostCommands, + nameJob=self.slurm_settings["name"], + minutes=self.slurm_settings["minutes"], + nodes=self.slurm_settings["nodes"], + ntasks=self.slurm_settings["ntasks"], + cpuspertask=self.slurm_settings["cpuspertask"], + slurm=self.machineSettings["slurm"], + launchSlurm=self.launchSlurm, + label_log_files=self.label_log_files, + wait_until_sbatch=waitYN, + ) + # ****************************************************** + + if fileSBTACH not in self.input_files: + self.input_files.append(fileSBTACH) + if fileSHELL not in self.input_files: + self.input_files.append(fileSHELL) + + self.output_files = curateOutFiles(self.output_files) + + # Relative paths TO FIX + self.input_files = [ + os.path.relpath(path, self.folder_local) for path in self.input_files + ] + self.input_folders = [ + os.path.relpath(path, self.folder_local) for path in self.input_folders + ] + + # Process + self.full_process( + comm, removeScratchFolders=removeScratchFolders, timeoutSecs=timeoutSecs + ) -# ------------------------------------- + # Get jobid + if self.launchSlurm: + with open(self.folder_local + "/mitim.out", "r") as f: + aux = f.readlines() + self.jobid = aux[0].split()[-1] + else: + self.jobid = None + + # -------------------------------------------------------------------- + # SSH executions + # -------------------------------------------------------------------- + + def full_process(self, comm, timeoutSecs=1e6, removeScratchFolders=True): + """ + My philosophy is to always wait for the execution of all commands. If I need + to not wait, that's handled by a slurm submission without --wait, but I still + want to execute that. + """ + wait_for_all_commands = True + + if timeoutSecs < 1e6: + timeOut_txt = f", will timeout execution in {timeoutSecs}s" + else: + timeOut_txt = "" + time_init = datetime.datetime.now() -def runFunction_Complete( - FunctionToRun, - Params_in, - WhereIsFunction="", - machineSettings={"machine": "local", "folderWork": "~/", "clear": False}, - scratchFolder="~/scratch/", -): - """ - Runs function externally. - Important note... this will use the function as is at the time of execution - """ - - txt = f""" -import sys,pickle -{WhereIsFunction} + print( + f"\n\t-------------- Running process ({time_init.strftime('%Y-%m-%d %H:%M:%S')}{timeOut_txt}) --------------" + ) -paramF = sys.argv[1] -Params_in = pickle.load(open(paramF,'rb')) -print(' ~~ Parameters successfully read') + self.connect(log_file=f"{self.folder_local}/paramiko.log") + if removeScratchFolders: + self.remove_scratch_folder() + self.create_scratch_folder() + self.send() + + output, error = self.execute( + comm, + make_relative=True, + wait_for_all_commands=wait_for_all_commands, + printYN=True, + timeoutSecs=timeoutSecs if timeoutSecs < 1e6 else None, + ) -Params_for_function = Params_in[:-2] -folder = Params_in[-2] + for _ in range(2): + self.retrieve() + received = self.check_all_received() + if received: + print("\t\t- All correct") + break + print("\t\t- Not all received, trying again", typeMsg="w") + time.sleep(10) + + if received: + if wait_for_all_commands and removeScratchFolders: + self.remove_scratch_folder() + else: + cont = print( + "\t* Not all expected files received, not removing scratch folder", + typeMsg="q", + ) + if not cont: + print( + "[mitim] Stopped with embed(), you can look at output and error", + typeMsg="w", + ) + embed() -Params_out = {FunctionToRun}(*Params_for_function) + self.close() -with open(folder+'/Params_out.pkl','wb') as handle: pickle.dump(Params_out,handle,protocol=2) -""" + print( + f"\t-------------- Finished process (took {IOtools.getTimeDifference(time_init)}) --------------\n" + ) - IOtools.askNewFolder(scratchFolder, force=True) - fileExe = scratchFolder + "/exe.py" - with open(fileExe, "w") as f: - f.write(txt) + def connect(self, *args, **kwargs): + if self.machineSettings["machine"] != "local": + return self.connect_ssh(*args, **kwargs) + else: + self.jump_client, self.ssh, self.sftp = None, None, None - ScriptToRun = f'python3 {machineSettings["folderWork"]}/exe.py' + def connect_ssh(self, log_file=None): + self.jump_host = self.machineSettings["tunnel"] + self.jump_user = self.machineSettings["user"] - Params_out = runFunction( - ScriptToRun, - Params_in, - InputFiles=[fileExe], - machineSettings=machineSettings, - scratchFolder=scratchFolder, - ) + self.target_host = self.machineSettings["machine"] + self.target_user = self.machineSettings["user"] - return Params_out + print("\t* Connecting to remote server:") + print( + f'\t\t{self.target_user}@{self.target_host}{f", via tunnel {self.jump_user}@" +self.jump_host if self.jump_host is not None else ""}{":" + str(self.machineSettings["port"]) if self.machineSettings["port"] is not None else ""}{" with identity: " + self.machineSettings["identity"] if self.machineSettings["identity"] is not None else ""}' + ) + if log_file is not None: + paramiko.util.log_to_file(log_file) -def runFunction( - ScriptToRun, - Params_in, - InputFiles=[], - machineSettings={"machine": "local", "folderWork": "~/", "clear": False}, - scratchFolder="~/scratch/", -): - """ - Params_in is tuple - """ + try: + self.define_jump() + self.define_server() + except paramiko.ssh_exception.AuthenticationException: + # If it fails, try to disable rsa-sha2-512 and rsa-sha2-256 (e.g. for iris.gat.com) + self.define_jump() + self.define_server( + disabled_algorithms={"pubkeys": ["rsa-sha2-512", "rsa-sha2-256"]} + ) - scratchFolder = IOtools.expandPath(scratchFolder) - - if not os.path.exists(scratchFolder): - IOtools.askNewFolder(scratchFolder, force=True) - - # First, put parameters into a pickle - Params_in += (machineSettings["folderWork"],) - pickF = "parameters_in.pkl" - with open(scratchFolder + "/" + pickF, "wb") as handle: - pickle.dump(Params_in, handle, protocol=1) - - # Then, run function that takes the location of Params.pkl and creates Params_out.pkl - commandMain = f'{ScriptToRun} {machineSettings["folderWork"]}/{pickF}' - InputFiles.append(scratchFolder + "/" + pickF) - runCommand( - commandMain, - InputFiles, - outputFiles=["Params_out.pkl"], - whereOutput=scratchFolder, - machineSettings=machineSettings, - ) + def define_server(self, disabled_algorithms=None): + # Create a new SSH client for the target machine + self.ssh = paramiko.SSHClient() + self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + # Connect to the host + self.ssh.connect( + self.target_host, + username=self.target_user, + disabled_algorithms=disabled_algorithms, + key_filename=self.key_filename, + port=self.port, + sock=self.sock, + allow_agent=True, + ) - # Read output - Params_out = pickle.load(open(scratchFolder + "/Params_out.pkl", "rb")) + try: + self.sftp = self.ssh.open_sftp() + except paramiko.sftp.SFTPError: + raise Exception( + "[mitim] SFTPError: Your bashrc on the server likely contains print statements" + ) - return Params_out + def define_jump(self): + if self.jump_host is not None: + # Create an SSH client instance for the jump host + self.jump_client = paramiko.SSHClient() + self.jump_client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + # Connect to the jump host + self.jump_client.connect( + self.jump_host, + username=self.jump_user, + port=self.machineSettings["port"] + if self.machineSettings["port"] is not None + else 22, + key_filename=self.machineSettings["identity"], + allow_agent=True, + ) + # Use the existing transport of self.jump_client for tunneling + transport = self.jump_client.get_transport() -def run_subprocess(commandExecute, shell=False, timeoutSecs=None, localRun=False): - """ - Note that before I had a context such as "with Popen() as p:" but that failed to catch time outs! - So, even though I don't know why... I'm doing this directly, with opening and closing it + # Create a channel to the target through the tunnel + create_port_in_tunnel = 22 + channel = transport.open_channel( + "direct-tcpip", + (self.target_host, create_port_in_tunnel), + (self.target_host, 0), + ) - For local runs, I had originally: - error=None; result=None; - os.system(Command) - Now, it uses subprocess with shell. This is because I couldn't load "source" because is a shell command, with simple os.system() - New solution is not the safest but it works. - """ + # to pass to self.ssh + self.port = create_port_in_tunnel + self.sock = channel + self.key_filename = None - if localRun: - print("\t\t\t++ Running local process") - shell = True + else: + self.jump_client = None + self.port = ( + self.machineSettings["port"] + if self.machineSettings["port"] is not None + else 22 + ) + self.sock = None + self.key_filename = self.machineSettings["identity"] - executable = ( - "/bin/bash" if shell else None - ) # If I'm simply running on the shell, then it's better to indicate bash + def create_scratch_folder(self): + print(f'\t* Creating{" remote" if self.ssh is not None else ""} folder:') + print(f'\t\t{self.machineSettings["folderWork"]}') - p = subprocess.Popen( - commandExecute, - shell=shell, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - executable=executable, - ) + command = "mkdir -p " + self.machineSettings["folderWork"] - result, error = None, None - if timeoutSecs is not None: - with timeout(timeoutSecs, proc=p): - result, error = p.communicate() - p.stdout.close() - p.stderr.close() - else: - result, error = p.communicate() - p.stdout.close() - p.stderr.close() + output, error = self.execute(command) - return error, result + return output, error + def send(self): + print( + f'\t* Sending files{" to remove server" if self.ssh is not None else ""}:' + ) -def runCommand_remote( - command, - machine="local", - user=None, - tunnel=None, - userTunnel=None, - port=None, - identity=None, - shell=False, - timeoutSecs=None, -): - # Building command for port and identity options -------------- + # Create a tarball of the local directory + print("\t\t- Tarballing") + with tarfile.open( + os.path.join(self.folder_local, "mitim_send.tar.gz"), "w:gz" + ) as tar: + for file in self.input_files + self.input_folders: + tar.add(os.path.join(self.folder_local, file), arcname=file) + + # Send it + print("\t\t- Sending") + if self.ssh is not None: + with TqdmUpTo( + unit="B", + unit_scale=True, + miniters=1, + desc="mitim_send.tar.gz", + bar_format=" " * 20 + + "{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{rate_fmt}{postfix}]", + ) as t: + self.sftp.put( + os.path.join(self.folder_local, "mitim_send.tar.gz"), + os.path.join( + self.machineSettings["folderWork"], "mitim_send.tar.gz" + ), + callback=lambda sent, total_size: t.update_to(sent, total_size), + ) + else: + os.system( + "cp " + + os.path.join(self.folder_local, "mitim_send.tar.gz") + + " " + + os.path.join(self.machineSettings["folderWork"], "mitim_send.tar.gz") + ) - sshCommand = ["ssh"] + # Extract it + print("\t\t- Extracting tarball") + self.execute( + "tar -xzf " + + os.path.join(self.machineSettings["folderWork"], "mitim_send.tar.gz") + + " -C " + + self.machineSettings["folderWork"] + ) - if port is not None: - portCommand = f"-p {port}" - sshCommand.append(portCommand) - else: - portCommand = "" - if identity is not None: - sshCommand.append("-i") - sshCommand.append(identity) + # Remove tarballs + print("\t\t- Removing tarballs") + os.remove(os.path.join(self.folder_local, "mitim_send.tar.gz")) + self.execute( + "rm " + + os.path.join(self.machineSettings["folderWork"], "mitim_send.tar.gz") + ) - if user is None: - userCommand = "" - else: - userCommand = user + "@" + def execute(self, command_str, make_relative=False, **kwargs): + # Always start by going to the folder + if make_relative: + command_str_mod = f"cd {self.folderExecution} && {command_str}" + else: + command_str_mod = command_str - if userTunnel is None: - userTunnelCommand = "" - else: - userTunnelCommand = userTunnel + "@" + if self.ssh is not None: + return self.execute_remote(command_str_mod, **kwargs) + else: + return self.execute_local(command_str_mod, **kwargs) + + def execute_remote( + self, + command_str, + printYN=False, + timeoutSecs=None, + wait_for_all_commands=True, + **kwargs, + ): + if printYN: + print("\t* Executing (remote):", typeMsg="i") + print(f"\t\t{command_str}") + + output = None + error = None - # -------------------------------------------------------------- + try: + stdin, stdout, stderr = self.ssh.exec_command( + command_str, timeout=timeoutSecs + ) + # Wait for the command to complete and read the output + if wait_for_all_commands: + stdin.close() + output = stdout.read() + error = stderr.read() + except socket.timeout: + print("Command timed out") + + return output, error + + def execute_local(self, command_str, printYN=False, timeoutSecs=None, **kwargs): + if printYN: + print("\t* Executing (local):") + print(f"\t\t{command_str}") + + output, error = run_subprocess( + [command_str], timeoutSecs=timeoutSecs, localRun=True + ) - if machine == "local" or machine == socket.gethostname(): - error, result = run_subprocess(command, localRun=True) - else: - if tunnel is None or len(tunnel) == 0: - commandExecute = [f"{userCommand}{machine}", command] - else: - commandExecute = [ - f"{userTunnelCommand}{tunnel}", - f'ssh {user}@{machine} "{command}"', - ] - - sshCommand.extend(commandExecute) - error, result = run_subprocess(sshCommand, shell=shell, timeoutSecs=timeoutSecs) - - return error, result - - -def sendCommand_remote( - file, - folderWork, - machine="local", - user=None, - tunnel=None, - userTunnel=None, - folderWorkTunnel=None, - port=None, - identity=None, - isItFolders=False, -): - # Building command for port and identity options -------------- + return output, error - sshCommand = ["ssh"] + def retrieve(self): + print( + f'\t* Retrieving files{" from remote server" if self.ssh is not None else ""}:' + ) - if port is not None: - portCommand = f"-p {port}" - sshCommand.append(portCommand) - else: - portCommand = "" - if identity is not None: - identityCommand = f"-i {identity}" - sshCommand.append(identityCommand) - else: - identityCommand = "" + # Create a tarball of the output files and folders on the remote machine + print("\t\t- Tarballing") + self.execute( + "tar -czf " + + os.path.join(self.machineSettings["folderWork"], "mitim_receive.tar.gz") + + " -C " + + self.machineSettings["folderWork"] + + " " + + " ".join(self.output_files + self.output_folders) + ) - if user is None: - userCommand = "" - else: - userCommand = user + "@" + # Download the tarball + print("\t\t- Downloading") + if self.ssh is not None: + with TqdmUpTo( + unit="B", + unit_scale=True, + miniters=1, + desc="mitim_receive.tar.gz", + bar_format=" " * 20 + + "{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{rate_fmt}{postfix}]", + ) as t: + self.sftp.get( + os.path.join( + self.machineSettings["folderWork"], "mitim_receive.tar.gz" + ), + os.path.join(self.folder_local, "mitim_receive.tar.gz"), + callback=lambda sent, total_size: t.update_to(sent, total_size), + ) + else: + os.system( + "cp " + + os.path.join( + self.machineSettings["folderWork"], "mitim_receive.tar.gz" + ) + + " " + + os.path.join(self.folder_local, "mitim_receive.tar.gz") + ) - if userTunnel is None: - userTunnelCommand = "" - else: - userTunnelCommand = userTunnel + "@" + # Extract the tarball locally + print("\t\t- Extracting tarball") + with tarfile.open( + os.path.join(self.folder_local, "mitim_receive.tar.gz"), "r:gz" + ) as tar: + tar.extractall(path=self.folder_local) + + # Remove tarballs + print("\t\t- Removing tarballs") + os.remove(os.path.join(self.folder_local, "mitim_receive.tar.gz")) + self.execute( + "rm " + + os.path.join(self.machineSettings["folderWork"], "mitim_receive.tar.gz") + ) - if folderWorkTunnel is None: - folderWorkTunnel = folderWork + def remove_scratch_folder(self): + print(f'\t* Removing{" remote" if self.ssh is not None else ""} folder') - if machine == "local" and not isItFolders: - inFiles = file[1:-1].split(",") - else: - inFiles = file.split(" ") + output, error = self.execute("rm -rf " + self.machineSettings["folderWork"]) - # -------------------------------------------------------------- + return output, error - if machine == "local" or machine == socket.gethostname(): - if not isItFolders: - for file0 in inFiles: - if len(file0) > 0: - os.system(f"cp {file0} {folderWork}/.") - else: - for file0 in inFiles: - os.system(f"cp -r {file0} {folderWork}/.") - elif tunnel is None: - if not isItFolders: - commai = f"scp {quiet_tag}{portCommand.upper()} {identityCommand} {file} {userCommand}{machine}:{folderWork}/." - else: - commai = f"scp {quiet_tag}{portCommand.upper()} {identityCommand} -r {file} {userCommand}{machine}:{folderWork}/." - # run_subprocess(commai,localRun=True) - os.system(commai) - else: - # Send files to tunnel - sendCommand_remote( - file, - folderWorkTunnel, - machine=tunnel, - user=userTunnel, - tunnel=None, - port=port, - isItFolders=isItFolders, - ) + def close(self, *args, **kwargs): + if self.machineSettings["machine"] != "local": + return self.close_ssh(*args, **kwargs) - # Send files from tunnel to remote (this ssh connections is better to do file by file) - - for file0 in inFiles: - if not isItFolders: - if len(file0.split("/")) > 1: - file1 = IOtools.reducePathLevel(file0, level=1, isItFile=True)[1] - else: - file1 = file0 - commandExecute = [ - f"{userTunnelCommand}{tunnel}", - f"scp {quiet_tag}{folderWorkTunnel}/{file1} {userCommand}{machine}:{folderWork}/.", - ] - else: - file1 = IOtools.reducePathLevel(file0, level=1, isItFile=False)[1] - commandExecute = [ - f"{userTunnelCommand}{tunnel}", - f"scp -r {quiet_tag}{folderWorkTunnel}/{file1} {userCommand}{machine}:{folderWork}/.", - ] - sshCommand.extend(commandExecute) - - with subprocess.Popen( - sshCommand, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE - ) as p: - result, error = p.communicate() - - -def receiveCommand_remote( - file, - folderWork, - f2, - machine="local", - user=None, - tunnel=None, - userTunnel=None, - port=None, - identity=None, - whereFiles=None, - folderWorkTunnel=None, - isItFolders=False, - scpCommand="scp -TO", -): - if folderWorkTunnel is None: - folderWorkTunnel = folderWork + def close_ssh(self): + print("\t* Closing connection") - if whereFiles is None: - folderWork_extra = folderWork - else: - folderWork_extra = whereFiles + self.sftp.close() + self.ssh.close() - # Building command for port and identity options -------------- + if self.jump_client is not None: + self.jump_client.close() - sshCommand = ["ssh"] + # -------------------------------------------------------------------- - if port is not None: - portCommand = f"-p {port}" - sshCommand.append(portCommand) - else: - portCommand = "" - if identity is not None: - identityCommand = f"-i {identity}" - sshCommand.append(identityCommand) - else: - identityCommand = "" + def check(self): + """ + Check job status slurm - # ---- - if not isItFolders: - outputFiles = file[1:-1].split(",") - else: - outputFiles = [file] - for i in range(len(outputFiles)): - outputFiles[i] = folderWork + outputFiles[i] + - If the job was launched with run(waitYN=False), then the script will not + wait for the full slurm execution, but will launch it and retrieve the jobid, + which the check command uses to check the status of the job. + - If the class was initiated but not run, it will not have the jobid, so it will + try to find it from the job_name, which must match the submitted one. + """ - if len(outputFiles) < 2: - file_mod = outputFiles[0] - else: - # file_mod = '{'+','.join(outputFiles)+'}' - file_mod = '"' + " ".join(outputFiles) + '"' + if self.jobid is not None: + txt_look = f"-j {self.jobid}" + else: + txt_look = f"-n {self.slurm_settings['name']}" - # -------------------------------------------------------------- + command = ( + f'squeue {txt_look} -o "%.15i %.24P %.18j %.10u %.10T %.10M %.10l %.5D %R" ' + ) - if user is None: - userCommand = "" - else: - userCommand = user + "@" + self.connect() + self.output_files.append("slurm_output.dat") + output, error = self.execute(command, printYN=True) + self.retrieve() # Retrieve self.output_files too + self.close() - if userTunnel is None: - userTunnelCommand = "" - else: - userTunnelCommand = userTunnel + "@" - - if machine == "local": - if verbose_level in [4, 5]: - print("\t\t\t++ Copying local file") - for file0 in outputFiles: - if isItFolders: - os.system(f"cp -r {file0} {f2}/.") - else: - os.system(f"cp {file0} {f2}/.") - elif tunnel is None: - if verbose_level in [4, 5]: - print(f"\t\t\t- Secure ssh copy from {machine}:{folderWork} to local") - port_iden = f"{portCommand.upper()} {identityCommand}" - if isItFolders: - commandExecute = f"{scpCommand} -r {quiet_tag}{port_iden} {userCommand}{machine}:{file_mod} {f2}/." # https://unix.stackexchange.com/questions/708517/scp-multiple-files-with-single-command + output = str(output)[3:].split("\\n") + if len(output[1].split()) == 1: + self.infoSLURM = None else: - commandExecute = f"{scpCommand} {quiet_tag}{port_iden} {userCommand}{machine}:{file_mod} {f2}/." - # if verbose_level in [4, 5]: - # print(f'\t\t\t{commandExecute}') - os.system(commandExecute) - else: - # Execute command to retrieve files from remote to tunnel (this ssh connections is better to do file by file) - for file_mod_extra in outputFiles: - if verbose_level in [4, 5]: - print( - f"\t\t\t- In {tunnel}: Secure ssh copy from {machine}:{folderWork_extra} to {tunnel}:{folderWorkTunnel}" - ) - if isItFolders: - comm = f"scp -r {quiet_tag}{userCommand}{machine}:{file_mod_extra} {folderWorkTunnel}/." - else: - comm = f"scp {quiet_tag}{userCommand}{machine}:{file_mod_extra} {folderWorkTunnel}/." - - commandExecute = [f"{userTunnelCommand}{tunnel}", comm] - sshCommand.extend(commandExecute) - - with subprocess.Popen( - sshCommand, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE - ) as p: - result, error = p.communicate() - - # Execute command to retrieve files from tunnel to local - receiveCommand_remote( - file, - folderWorkTunnel, - f2, - machine=tunnel, - user=userTunnel, - tunnel=None, - port=port, - isItFolders=isItFolders, - scpCommand=scpCommand, - ) + self.infoSLURM = {} + for i in range(len(output[0].split())): + self.infoSLURM[output[0].split()[i]] = output[1].split()[i] + with open(self.folder_local + "/slurm_output.dat", "r") as f: + self.log_file = f.readlines() -def runCommand( - commandMain, - inputFiles, - outputFiles=[], - inputFolders=[], - outputFolders=[], - whereOutput="~/.", - machineSettings={"machine": "local", "folderWork": "~/", "clear": False}, - timeoutSecs=1e6, - whereFiles=None, - pauseBeforeCommand=False, -): - timeOut_txt = ( - f", will timeout execution in {timeoutSecs}s" if timeoutSecs < 1e6 else "" - ) + # If jobid was given as None, I retrieved the info from the job_name, but now provide here the actual id + if self.infoSLURM is not None: + self.jobid_found = self.infoSLURM["JOBID"] + else: + self.jobid_found = None + + # Print info + txt = "\t- Job was checked" + if (self.jobid is None) and (self.jobid_found is not None): + txt += f' (jobid {self.jobid_found}, found from name "{self.slurm_settings["name"]}")' + elif self.jobid is not None: + txt += f" (jobid {self.jobid})" + if self.infoSLURM is not None: + txt += f', is {self.infoSLURM["STATE"]} (job.infoSLURM)' + else: + txt += ", is not on the grid (job.infoSLURM is None)" + if self.log_file is not None: + txt += f". Log file (job.log_file) was retrieved, and has {len(self.log_file)} lines" + print(txt) + + def check_all_received(self): + print("\t* Checking if all expected files and folders were received") + received = True + for file in self.output_files: + if not os.path.exists(os.path.join(self.folder_local, file)): + print(f"\t\t- File {file} not received", typeMsg="w") + received = False + for folder in self.output_folders: + if not os.path.exists(os.path.join(self.folder_local, folder)): + print(f"\t\t- Folder {folder} not received", typeMsg="w") + received = False + + return received + + +class TqdmUpTo(tqdm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.initialized = False + + def update_to(self, sent, total_size): + if not self.initialized: + self.total = total_size + self.initialized = True + self.update(sent - self.n) # will also set self.n = sent - machine = machineSettings["machine"] - folderWork = machineSettings["folderWork"] - clearYN = machineSettings.setdefault("clear", False) - user = machineSettings.setdefault("user", None) - tunnel = machineSettings.setdefault("tunnel", None) - port = machineSettings.setdefault("port", None) - identity = machineSettings.setdefault("identity", None) - userTunnel = machineSettings.setdefault("userTunnel", user) - folderWorkTunnel = machineSettings.setdefault("folderWorkTunnel", folderWork) +""" + Timeout function + - I just need to add a function: + with timeout(secs): + do things + - I don't recommend that "do things" includes a context manager like with Popen or it won't work... not sure why +""" - isTunnelSameMachine = machineSettings.setdefault("isTunnelSameMachine", 0) - # --- If I am at the tunnel, do not tunnel! - if (tunnel is not None) and (socket.gethostname() in tunnel): - tunnel = None - # ---- +def raise_timeout(signum, frame): + raise TimeoutError - if (tunnel is not None) and isTunnelSameMachine: + +@contextmanager +def timeout(time, proc=None): + time = int(time) + if time < 1e6: print( - "\t- Tunnel and machine share the same file system, do not tunnel for file handling", + f'\t\t* Note: this process will be killed if time exceeds {time}sec of execution ({datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")})', typeMsg="i", ) - tunnelF = None - machineF = tunnel - else: - tunnelF = tunnel - machineF = machine - currentTime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - - print( - f"\n\t------------------------ Running process ({currentTime}{timeOut_txt}) ------------------------" - ) - if verbose_level in [4, 5]: - print(f"\t * process: {commandMain}") - machine_info = ( - f"\t * in machine {machine}" - if machine != "local" - else "\t * locally in this machine" - ) - if tunnel is not None: - tunnel_info = f" through {tunnel}" - if identity is not None: - tunnel_info += f" with identity {identity}" - if port is not None: - tunnel_info += f" with port {port}" - machine_info += tunnel_info - - print(machine_info) - - print("\t * following these steps:") - - # ------------ Create temporal folders if they do not exist - - contStep = 1 - - print(f"\t\t{contStep}. Creating temporal folders") - - if tunnelF is not None: - # -------- Sometimes cybele doesn't have scratch... not sure why - folderScratch = IOtools.reducePathLevel( - folderWorkTunnel, level=1, isItFile=False - )[0] - command = f"mkdir {folderScratch}" - error, result = runCommand_remote( - command, - machine=tunnel, - user=userTunnel, - userTunnel=None, - tunnel=None, - port=port, - identity=identity, - ) - # --------------------------- - - command = f"mkdir {folderWorkTunnel}" - error, result = runCommand_remote( - command, - machine=tunnel, - user=userTunnel, - userTunnel=None, - tunnel=None, - port=port, - identity=identity, - ) - - command = f"mkdir {folderWork}" - error, result = runCommand_remote( - command, - machine=machineF, - user=user, - tunnel=tunnelF, - userTunnel=userTunnel, - port=port, - identity=identity, - ) - - # ------------ Send files to machine - - if len(inputFiles) > 0: - contStep += 1 - print(f"\t\t{contStep}. Sending required files") - if machine == "local": - combined_files = "{" + ",".join(inputFiles) + ",}" - else: - combined_files = " ".join(inputFiles) - sendCommand_remote( - combined_files, - folderWork, - folderWorkTunnel=folderWorkTunnel, - machine=machineF, - user=user, - tunnel=tunnelF, - userTunnel=userTunnel, - port=port, - identity=identity, - isItFolders=False, - ) + # Register a function to raise a TimeoutError on the signal. + signal.signal(signal.SIGALRM, raise_timeout) + # Schedule the signal to be sent after ``time``. + signal.alarm(time) - # ------------ Send folders to machine - - if len(inputFolders) > 0: - contStep += 1 - print(f"\t\t{contStep}. Sending required folders") - combined_folders = " ".join(inputFolders) - sendCommand_remote( - combined_folders, - folderWork, - folderWorkTunnel=folderWorkTunnel, - machine=machineF, - user=user, - tunnel=tunnelF, - userTunnel=userTunnel, - port=port, - identity=identity, - isItFolders=True, + try: + yield + except TimeoutError: + print( + f'\t\t\t* Killing process! ({datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")})', + typeMsg="w", ) + if proc is not None: + proc.kill() + outs, errs = proc.communicate() + finally: + # Unregister the signal so it won't be triggered if the timeout is not reached. + signal.signal(signal.SIGALRM, signal.SIG_IGN) - # ------------ Run command - - if pauseBeforeCommand: - embed() - - error, result = "", "" - contStep += 1 - print(f"\t\t{contStep}. Running command") - if verbose_level in [4, 5]: - print(f"\t\t\t{commandMain}") - error, result = runCommand_remote( - commandMain, - machine=machine, - user=user, - userTunnel=userTunnel, - tunnel=tunnel, - port=port, - identity=identity, - timeoutSecs=timeoutSecs, - ) - # --------------------------------------------------------------- - # ------------ Bring files from machine - # --------------------------------------------------------------- - - scpCommand = "scp -TO" - allFilesCorrect = bringFiles( - outputFiles, - outputFolders, - folderWork, - whereOutput, - folderWorkTunnel, - machineF, - user, - tunnelF, - userTunnel, - port, - identity, - whereFiles, - scpCommand, - contStep=contStep, - ) +def run_subprocess(commandExecute, timeoutSecs=None, localRun=False): + """ + Note that before I had a context such as "with Popen() as p:" but that failed to catch time outs! + So, even though I don't know why... I'm doing this directly, with opening and closing it - if not allFilesCorrect: - print( - "* Exploring SCP option -T instead of -TO because retrieval failed", - typeMsg="w", - ) - scpCommand = "scp -T" - allFilesCorrect = bringFiles( - outputFiles, - outputFolders, - folderWork, - whereOutput, - folderWorkTunnel, - machineF, - user, - tunnelF, - userTunnel, - port, - identity, - whereFiles, - scpCommand, - contStep=contStep, - ) + For local runs, I had originally: + error=None; result=None; + os.system(Command) + Now, it uses subprocess with shell. This is because I couldn't load "source" because is a shell command, with simple os.system() + New solution is not the safest but it works. + """ - if not allFilesCorrect: - print( - "\t- Not all output files received from remote computer, printing process error:", - typeMsg="w", - ) - try: - print("".join(error)) - except: - print(error) - print("\t- Printing process result:") - try: - print("".join(result)) - except: - print(result) - print( - "\t- For easy debugging, I am not clearing the remote folder", typeMsg="w" - ) - clearYN = False - - # Clear stuff - if clearYN: - command = f"rm -r {folderWork}" - - error, result = runCommand_remote( - command, - machine=machine, - user=user, - userTunnel=userTunnel, - tunnel=tunnel, - port=port, - identity=identity, - ) - if tunnelF is not None: - command = f"rm -r {folderWorkTunnel}" - error, result = runCommand_remote( - command, - machine=tunnel, - user=user, - tunnel=None, - port=port, - identity=identity, - ) + if localRun: + shell = True + executable = "/bin/bash" + else: + shell = False + executable = None - print( - f"\t------------------------ Finished process ({datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}) ------------------------" + p = subprocess.Popen( + commandExecute, + shell=shell, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + executable=executable, ) - print("\n") - - -def bringFiles( - outputFiles, - outputFolders, - folderWork, - whereOutput, - folderWorkTunnel, - machine, - user, - tunnel, - userTunnel, - port, - identity, - whereFiles, - scpCommand, - contStep=0, -): - if len(outputFiles) > 0: - contStep += 1 - print(f"\t\t{contStep}. Receiving files") - combined_files = "{" + ",".join(outputFiles) + "}" - receiveCommand_remote( - combined_files, - folderWork, - whereOutput, - isItFolders=False, - folderWorkTunnel=folderWorkTunnel, - machine=machine, - user=user, - tunnel=tunnel, - userTunnel=userTunnel, - port=port, - identity=identity, - whereFiles=whereFiles, - scpCommand=scpCommand, - ) - if len(outputFolders) > 0: - contStep += 1 - print(f"\t\t{contStep}. Receiving folders") - for outputFolder in outputFolders: - # First remove current folders - upo = 0 - tester = f"{whereOutput}/{outputFolder}_old{upo}" - while os.path.exists(f"{whereOutput}/{outputFolder}"): - if os.path.exists(tester): - upo += 1 - tester = f"{whereOutput}/{outputFolder}_old{upo}" - else: - os.system(f"mv {whereOutput}/{outputFolder} {tester}") - - receiveCommand_remote( - outputFolder, - folderWork, - whereOutput, - isItFolders=True, - folderWorkTunnel=folderWorkTunnel, - machine=machine, - user=user, - tunnel=tunnel, - userTunnel=userTunnel, - port=port, - identity=identity, - whereFiles=whereFiles, - scpCommand=scpCommand, - ) - - allFilesCorrect = True - for file in outputFiles: - fileF = os.path.exists(f"{whereOutput}/{file}") - allFilesCorrect = allFilesCorrect and fileF - if not fileF: - print(f"\t\t\t- File {file} not found", typeMsg="w") - for folder in outputFolders: - folderF = os.path.exists(f"{whereOutput}/{folder}") - allFilesCorrect = allFilesCorrect and folderF - if not folderF: - print(f"\t\t\t- Folder {folder} not found", typeMsg="w") + result, error = None, None + if timeoutSecs is not None: + with timeout(timeoutSecs, proc=p): + result, error = p.communicate() + p.stdout.close() + p.stderr.close() + else: + result, error = p.communicate() + p.stdout.close() + p.stderr.close() - return allFilesCorrect + return result, error # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -833,9 +750,9 @@ def FunctionToParallelize(Params,cont): """ -def init(l): +def init(l_lock): global lock - lock = l + lock = l_lock class PRF_ParallelClass_reduced(object): @@ -865,8 +782,8 @@ def ParallelProcedure( so that for instance not two at the same time open and write the same file. """ - l = multiprocessing.Lock() - pool = multiprocessing.Pool(initializer=init, initargs=(l,), processes=parallel) + l0 = multiprocessing.Lock() + pool = multiprocessing.Pool(initializer=init, initargs=(l0,), processes=parallel) if array: print( @@ -899,42 +816,46 @@ def SerialProcedure(Function, Params, howmany): return y, yE -def SLURM( +def create_slurm_execution_files( command, folder_remote, modules_remote, + slurm={}, folder_local=None, shellPreCommands=[], shellPostCommands=[], launchSlurm=True, nameJob="test", - partition="sched_mit_psfc", minutes=5, ntasks=1, cpuspertask=4, job_array=None, nodes=None, - email="noemail", - waitYN=True, - extranamelogs="", - exclude=None, - default_exclusive=False, + label_log_files="", + wait_until_sbatch=True, ): if folder_local is None: folder_local = folder_remote - if type(command) == str: + if isinstance(command, str): command = [command] folderExecution = IOtools.expandPath(folder_remote) - fileSBTACH = f"{folder_local}/bash.src" - fileSHELL = f"{folder_local}/mitim.sh" - fileSBTACH_remote = f"{folder_remote}/bash.src" + fileSBTACH = f"{folder_local}/mitim_bash.src" + fileSHELL = f"{folder_local}/mitim_shell_executor.sh" + fileSBTACH_remote = f"{folder_remote}/mitim_bash.src" minutes = int(minutes) + partition = slurm.setdefault("partition", None) + email = slurm.setdefault("email", None) + exclude = slurm.setdefault("exclude", None) + account = slurm.setdefault("account", None) + constraint = slurm.setdefault("constraint", None) + memory_req = slurm.setdefault("mem", None) + """ ******************************************************************************************** - Write bash.src file to execute + Write mitim_bash.src file to execute ******************************************************************************************** - Contains sourcing of mitim.bashrc, so that it's done at node level """ @@ -952,21 +873,29 @@ def SLURM( commandSBATCH.append("#!/bin/bash -l") commandSBATCH.append(f"#SBATCH --job-name {nameJob}") commandSBATCH.append( - f"#SBATCH --output {folderExecution}/slurm_output{extranamelogs}.dat" + f"#SBATCH --output {folderExecution}/slurm_output{label_log_files}.dat" ) commandSBATCH.append( - f"#SBATCH --error {folderExecution}/slurm_error{extranamelogs}.dat" + f"#SBATCH --error {folderExecution}/slurm_error{label_log_files}.dat" ) - if email != "noemail": + if email is not None: commandSBATCH.append("#SBATCH --mail-user=" + email) # ******* Partition / Billing commandSBATCH.append(f"#SBATCH --partition {partition}") + + if account is not None: + commandSBATCH.append(f"#SBATCH --account {account}") + if constraint is not None: + commandSBATCH.append(f"#SBATCH --constraint {constraint}") + + if memory_req is not None: + commandSBATCH.append(f"#SBATCH --mem {memory_req}") + commandSBATCH.append(f"#SBATCH --time {time_com}") if job_array is None: - if default_exclusive: - commandSBATCH.append("#SBATCH --exclusive") + commandSBATCH.append("#SBATCH --exclusive") else: commandSBATCH.append(f"#SBATCH --array={job_array}") @@ -1001,7 +930,7 @@ def SLURM( commandSBATCH.append("") - wait_txt = " --wait" if waitYN else "" + wait_txt = " --wait" if wait_until_sbatch else "" if launchSlurm: comm, launch = commandSBATCH, "sbatch" + wait_txt else: @@ -1014,7 +943,7 @@ def SLURM( """ ******************************************************************************************** - Write mitim.sh file that handles the execution of the bash.src with pre and post commands + Write mitim_shell_executor.sh file that handles the execution of the mitim_bash.src with pre and post commands ******************************************************************************************** - Contains sourcing of mitim.bashrc, so that it's done at machine level """ @@ -1028,9 +957,9 @@ def SLURM( for i in range(len(shellPostCommands)): commandSHELL.append(shellPostCommands[i]) # Evaluate Job performance - commandSHELL.append( - "python3 $MITIM_PATH/src/mitim_tools/misc_tools/FARMINGtools.py sbatch.out" - ) + # commandSHELL.append( + # "python3 $MITIM_PATH/src/mitim_tools/misc_tools/FARMINGtools.py mitim.out" + # ) if os.path.exists(fileSHELL): os.system(f"rm {fileSHELL}") @@ -1043,102 +972,16 @@ def SLURM( ******************************************************************************************** """ - comm = f"cd {folderExecution} && bash mitim.sh > sbatch.out" + comm = "bash mitim_shell_executor.sh > mitim.out" return comm, fileSBTACH, fileSHELL -def SLURMcomplete( - command, - folderLocal, - inputFiles, - outputFiles, - minutes, - ntasks, - nameJob, - machineSettings, - job_array=None, - nodes=None, - launchSlurm=True, - cpuspertask=1, - inputFolders=[], - outputFolders=[], - shellPreCommands=[], - shellPostCommands=[], - email="noemail", - waitYN=True, - extranamelogs="", - default_exclusive=False, -): - folderWork, modules, partition = ( - machineSettings["folderWork"], - machineSettings["modules"], - machineSettings["partition"], - ) - - if not waitYN: - # If I'm not waiting, make sure i don't clear the folder - machineSettings["clear"] = False - - if launchSlurm and (partition is None): - launchSlurm = False - print( - "\t- slurm requested but no partition assigned to this machine in config... not doing slurm", - typeMsg="w", - ) - - # ****** Prepare SLURM job ***************************** - comm, fileSBTACH, fileSHELL = SLURM( - command, - folderWork, - modules, - job_array=job_array, - folder_local=folderLocal, - shellPreCommands=shellPreCommands, - shellPostCommands=shellPostCommands, - nameJob=nameJob, - minutes=minutes, - nodes=nodes, - ntasks=ntasks, - cpuspertask=cpuspertask, - partition=partition, - launchSlurm=launchSlurm, - email=email, - waitYN=waitYN, - extranamelogs=extranamelogs, - exclude=machineSettings["exclude"], - default_exclusive=default_exclusive, - ) - # ****************************************************** - - inputFiles.append(fileSBTACH) - inputFiles.append(fileSHELL) - - outputFiles = curateOutFiles(outputFiles) - - runCommand( - comm, - inputFiles, - inputFolders=inputFolders, - outputFiles=outputFiles, - outputFolders=outputFolders, - whereOutput=folderLocal, - machineSettings=machineSettings, - ) - - with open(folderLocal + "/sbatch.out", "r") as f: - aux = f.readlines() - - jobid = aux[0].split()[-1] - - return jobid - - def curateOutFiles(outputFiles): # Avoid repetitions, otherwise, e.g., they will fail to rename - if "sbatch.out" not in outputFiles: - outputFiles.append("sbatch.out") + if "mitim.out" not in outputFiles: + outputFiles.append("mitim.out") outputFiles_new = [] for file in outputFiles: @@ -1148,50 +991,9 @@ def curateOutFiles(outputFiles): return outputFiles_new -def getSLURMstatus(FolderSLURM, machineSettings, jobid=None, name=None, grablog=None): - """ - Search by jobid or by name - """ - - folderWork = machineSettings["folderWork"] - - if jobid is not None: - txt_look = f"-j {jobid}" - else: - txt_look = f"-n {name}" - - command = f'cd {folderWork} && squeue {txt_look} -o "%.15i %.24P %.18j %.10u %.10T %.10M %.10l %.5D %R" > squeue.out' - - outputFiles = ["squeue.out"] - - if grablog is not None: - command += f" && cp {grablog} {folderWork}/." - outputFiles.append(IOtools.reducePathLevel(grablog)[-1]) - - runCommand( - command, - [], - outputFiles=outputFiles, - machineSettings=machineSettings, - whereOutput=FolderSLURM, - ) - - with open(f"{FolderSLURM}/squeue.out", "r") as f: - aux = f.readlines() - - if len(aux) > 1: - info = {} - for i in range(len(aux[0].split())): - info[aux[0].split()[i]] = aux[1].split()[i] - else: - info = None - - return info - - def printEfficiencySLURM(out_file): """ - It reads jobid from sbatch.out or slurm_output.dat + It reads jobid from mitim.out or slurm_output.dat """ with open(out_file, "r") as f: diff --git a/src/mitim_tools/misc_tools/GRAPHICStools.py b/src/mitim_tools/misc_tools/GRAPHICStools.py index 678b4d2c..c74068e2 100644 --- a/src/mitim_tools/misc_tools/GRAPHICStools.py +++ b/src/mitim_tools/misc_tools/GRAPHICStools.py @@ -1,6 +1,7 @@ -import pdb, sys, socket +import socket import numpy as np import matplotlib.pyplot as plt +from collections import OrderedDict from mpl_toolkits.axes_grid1 import make_axes_locatable import matplotlib.colors as mcol import matplotlib.cm as cm @@ -17,6 +18,52 @@ def aroundZeroLims(zlims): return zlims +def convert_to_hex_soft(color): + """ + Color can be an integer, in that case it'll grab that position in the color_hex dictionary + """ + + # chatGPT created this dictionary to convert to hex colors + color_hex = OrderedDict( + { + "b": "#CCCCFF", # very light blue + "m": "#FFCCFF", # very light magenta + "r": "#FFCCCC", # very light red + "t": "#99CCCC", # very light teal + "y": "#FFFFCC", # very light yellow + "g": "#CCFFCC", # very light green instead of green + "c": "#CCFFFF", # very light cyan instead of cyan + "n": "#666699", # lighter navy + "o": "#CCCC99", # very light olive instead of olive + "p": "#FFCCFF", # very light purple instead of purple + "s": "#E0E0E0", # lighter silver + "a": "#FFA07A", # light salmon + "h": "#ADFF2F", # green yellow + "i": "#FFB6C1", # light pink + "j": "#FFD700", # gold + "q": "#20B2AA", # light sea green + "u": "#87CEFA", # light sky blue + "v": "#778899", # light slate gray + "x": "#B0C4DE", # light steel blue + "z": "#F08080", # light coral + "k": "#666666", # medium gray instead of black + "w": "#FAFAFA", # very light gray instead of white + } + ) + + if isinstance(color, int): + # If color is an integer, get the color at that index + color_keys = list(color_hex.keys()) + if color < len(color_keys): + return color_hex[color_keys[color]] + else: + return None + elif (color is not None) and (color in color_hex): + return color_hex[color] + else: + return None + + def plotRange( t, x, diff --git a/src/mitim_tools/misc_tools/GUItools.py b/src/mitim_tools/misc_tools/GUItools.py index f4037aa2..2953f996 100644 --- a/src/mitim_tools/misc_tools/GUItools.py +++ b/src/mitim_tools/misc_tools/GUItools.py @@ -4,7 +4,6 @@ """ import sys -from IPython import embed from mitim_tools.misc_tools import IOtools, GRAPHICStools from mitim_tools.misc_tools.IOtools import printMsg as print @@ -14,11 +13,11 @@ from matplotlib.backends.backend_qtagg import ( NavigationToolbar2QT as NavigationToolbar, ) - from PyQt6 import QtWidgets, QtCore + from PyQt6 import QtWidgets, QtCore, QtGui from PyQt6.QtWidgets import QTabWidget, QTabBar # ----------------------------- -except: +except ImportError: print( " > PyQt6 module or backends could not be loaded by MITIM, notebooks will not work but I let you continue", typeMsg="w", @@ -33,6 +32,7 @@ class QTabBar: import matplotlib.pyplot as plt from mitim_tools.misc_tools.CONFIGread import read_dpi +from IPython import embed plt.rcParams["figure.max_open_warning"] = False @@ -40,9 +40,9 @@ class QTabBar: class FigureNotebook: - def __init__( - self, dummy, windowtitle, parent=None, geometry="1800x900", vertical=True - ): + def __init__(self, windowtitle, parent=None, geometry="1800x900", vertical=True): + plt.ioff() + self.app = QtWidgets.QApplication.instance() if self.app is None: self.app = QtWidgets.QApplication(sys.argv) @@ -64,9 +64,9 @@ def __init__( self.MainWindow.resize(int(geometry.split("x")[0]), int(geometry.split("x")[1])) self.MainWindow.show() - def add_figure(self, label=""): + def add_figure(self, label="", tab_color=None): figure = plt.figure(dpi=dpi_notebook) - self.addPlot(label, figure) + self.addPlot(label, figure, tab_color=tab_color) return figure @@ -77,7 +77,11 @@ def subplots(self, ncols=1, nrows=1, sharey=False, sharex=False, label=""): return fig, ax - def addPlot(self, title, figure): + def addPlot(self, title, figure, tab_color=None, tab_alpha=0.2): + """ + tab_color can be a color name or an integer to grab colors in order + """ + new_tab = QtWidgets.QWidget() layout = QtWidgets.QVBoxLayout() new_tab.setLayout(layout) @@ -98,6 +102,13 @@ def addPlot(self, title, figure): self.figure_handles.append(figure) self.tab_handles.append(new_tab) + # Set the color for the tab if specified + tab_color_hex = GRAPHICStools.convert_to_hex_soft(tab_color) + if tab_color_hex: + tab_color_hex = QtGui.QColor(tab_color_hex) + tab_color_hex.setAlphaF(tab_alpha) + self.tabs.tabBar().setTabColor(self.tabs.count() - 1, tab_color_hex) + def show(self): print(f"\n> MITIM Notebook open, titled: {self.windowtitle}", typeMsg="i") print("\t- Close the notebook to continue") @@ -123,9 +134,10 @@ def __init__(self, parent=None, vertical=False, xextend=1600): super().__init__(parent) self.vertical = vertical + self.tab_colors = {} if self.vertical: - self.setFixedSize(xextend, 150) + self.setFixedSize(xextend, 170) else: self.setFixedSize(xextend, 30) @@ -135,17 +147,21 @@ def __init__(self, parent=None, vertical=False, xextend=1600): font-size: 9pt; } QTabBar::tab:selected { - background: #90EE90; + background: #00FF00; color: #191970; font: bold; } QTabBar::tab:hover { - background: #ADD8E6; - color: #000000; + background: #90EE90; + color: #191970; } """ ) + def setTabColor(self, index, color): + self.tab_colors[index] = color + self.update() + def tabSizeHint(self, i): if self.vertical: tw = int(self.width() / (self.count())) @@ -161,6 +177,11 @@ def paintEvent(self, event): for i in range(self.count()): self.initStyleOption(opt, i) + if i in self.tab_colors: + opt.palette.setColor( + QtGui.QPalette.ColorRole.Button, + QtGui.QColor(self.tab_colors[i]), + ) painter.drawControl( QtWidgets.QStyle.ControlElement.CE_TabBarTabShape, opt ) diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index aaa70249..ed3539a0 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -1,23 +1,30 @@ -import re, os, shutil, sys, time, datetime, socket, random, zipfile, cProfile, termios, tty, h5py +import re +import os +import shutil +import sys +import time +import datetime +import socket +import random +import zipfile +import cProfile +import termios +import tty +import h5py from collections import OrderedDict import numpy as np import matplotlib.pyplot as plt try: import pandas -except: +except ImportError: pass try: from IPython import embed -except: +except ImportError: pass -# Running here globally because most mitim use at some point IOtools -from mitim_tools.misc_tools.CONFIGread import read_verbose_level - -verbose_level = read_verbose_level() - import urllib.request as urlREQ # urllibR import urllib.error as urlERR # urllibE @@ -42,7 +49,7 @@ def __exit__(self, *args): ) -def clipstr(txt, chars=30): +def clipstr(txt, chars=40): return f"{'...' if len(txt) > chars else ''}{txt[-chars:]}" @@ -643,6 +650,9 @@ def findFileByExtension( ) fileReturn = None else: + from mitim_tools.misc_tools.CONFIGread import read_verbose_level + + verbose_level = read_verbose_level() printMsg( f"\t\t\t~ Folder ...{folder[np.max([-40,-len(folder)]):]} does not exist, returning None", verbose=verbose_level, @@ -1286,11 +1296,14 @@ def ArrayToString(ll): def expandPath(txt, fixSpaces=False, ensurePathValid=False): + while txt[-2:] == "//": + txt = txt[:-1] + if txt[0] == ".": - if len(txt) > 1: - txt = os.path.realpath(txt) - else: + if (len(txt) == 1) or (len(txt) == 2 and txt[-1] == "/"): txt = os.path.realpath(txt) + "/" + else: + txt = os.path.realpath(txt) if ensurePathValid and (txt[0] not in ["~", "/"]): txt = os.path.realpath("./" + txt + "/") diff --git a/src/mitim_tools/misc_tools/MATHtools.py b/src/mitim_tools/misc_tools/MATHtools.py index 6b05cc68..a608b1ef 100644 --- a/src/mitim_tools/misc_tools/MATHtools.py +++ b/src/mitim_tools/misc_tools/MATHtools.py @@ -625,10 +625,6 @@ def drawContours(Rold, Yold, Zold, resol=5e3, psiN_boundary=0.99999): [Rg, Yg] = np.meshgrid(R, Y) - try: - plt.ioff() - except: - pass cs = plt.contour(Rg, Yg, Z, resol, levels=[psiN_boundary]) Rpsi, Ypsi = [], [] diff --git a/src/mitim_tools/misc_tools/PLASMAtools.py b/src/mitim_tools/misc_tools/PLASMAtools.py index 8f03c18a..7d50026a 100644 --- a/src/mitim_tools/misc_tools/PLASMAtools.py +++ b/src/mitim_tools/misc_tools/PLASMAtools.py @@ -2,7 +2,8 @@ Plasma physics models/info/approximations """ -import scipy, torch +import scipy +import torch import numpy as np import matplotlib.pyplot as plt from IPython import embed @@ -41,6 +42,9 @@ me_u = 5.4488741e-04 # as in input.gacode +factor_convection = 3 / 2 # IMPORTANT + + def magneticshear(q, rmin, R0): """ [Rice PRL 2013] @@ -189,13 +193,13 @@ def nminfactor(nmin, n): return np.array(nminfact) -def convective_flux(Te, Gamma_e, factor=3 / 2): +def convective_flux(Te, Gamma_e): # keV and 1E20 m^-3/s (or /m^2) Te_J = Te * 1e3 * e_J Gamma = Gamma_e * 1e20 - Qe_conv = factor * Te_J * Gamma * 1e-6 # MW (or /m^2) + Qe_conv = factor_convection * Te_J * Gamma * 1e-6 # MW (or /m^2) return Qe_conv @@ -379,7 +383,10 @@ def gyrobohmUnits(Te_keV, ne_20, mref_u, Bunit, a): # Exchange source Sgb = Qgb / a - return Qgb, Ggb, Pgb, Sgb + # Particle flux when using convective flux + Qgb_convection = Qgb * factor_convection + + return Qgb, Ggb, Pgb, Sgb, Qgb_convection def conduction(n19, TkeV, chi, aLT, a): @@ -544,8 +551,8 @@ def calculateCoulombLogarithm(Te, ne, Z=None, ni=None, Ti=None): # TO FIX """ # Freidberg p. 194 - L = 4.9e7 * Te**1.5 / (ne * 1e-20) ** 0.5 - logL = np.log(L) + # L = 4.9e7 * Te**1.5 / (ne * 1e-20) ** 0.5 + # logL = np.log(L) # Lee = 30.0 - np.log( ne**0.5 * (Te*1E-3)**-1.5 ) # Lc = 24.0 - np.log( np.sqrt( ne/Te ) ) diff --git a/src/mitim_tools/opt_tools/BOTORCHtools.py b/src/mitim_tools/opt_tools/BOTORCHtools.py index 2c2bef12..5d574883 100644 --- a/src/mitim_tools/opt_tools/BOTORCHtools.py +++ b/src/mitim_tools/opt_tools/BOTORCHtools.py @@ -4,7 +4,9 @@ ************************************************************************************************************** """ -import torch, botorch, gpytorch +import torch +import botorch +import gpytorch from IPython import embed from mitim_tools.misc_tools.IOtools import printMsg as print from mitim_tools.misc_tools.CONFIGread import read_verbose_level diff --git a/src/mitim_tools/opt_tools/STEPtools.py b/src/mitim_tools/opt_tools/STEPtools.py index 6d64aea8..9822420a 100644 --- a/src/mitim_tools/opt_tools/STEPtools.py +++ b/src/mitim_tools/opt_tools/STEPtools.py @@ -367,7 +367,7 @@ def fit_step(self, avoidPoints=[], fitWithTrainingDataIfContains=None): self.GP["individual_models"], self.GP["combined_model"] ) print("--> Launching tests evaluate accuracy on training set (absolute units)") - self.GP["combined_model"].testTraining(printYN=False) + self.GP["combined_model"].testTraining() txt_time = IOtools.getTimeDifference(time1) diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index f65a716c..aa67bd0f 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -1,4 +1,9 @@ -import os, copy, datetime, array, traceback, torch +import os +import copy +import datetime +import array +import traceback +import torch from collections import OrderedDict from IPython import embed import dill as pickle_dill @@ -281,9 +286,6 @@ def plot_optimization_results( f"\t- Perform extra analysis for this sub-module (analysis level {analysis_level})" ) - if plotYN: - plt.ioff() - self.read_optimization_results( plotYN=plotYN and (analysis_level >= 0), folderRemote=folderRemote, @@ -295,11 +297,10 @@ def plot_optimization_results( self_complete = None if analysis_level > 1: - plt.ioff() """ - If the analyze_results exists, I'm in a child class, so just proceed to analyze. - Otherwise, let's grab the method from the pickled - """ + If the analyze_results exists, I'm in a child class, so just proceed to analyze. + Otherwise, let's grab the method from the pickled + """ if hasattr(self, "analyze_results"): self_complete = self.analyze_results( plotYN=plotYN, fn=self.fn, analysis_level=analysis_level @@ -333,13 +334,11 @@ def plot_optimization_results( ) if plotYN and (analysis_level >= 0): - print(f"- Plotting took {IOtools.getTimeDifference(time1)}", typeMsg="i") + print(f"\n- Plotting took {IOtools.getTimeDifference(time1)}") if save_folder is not None: self.fn.save(save_folder) - self.fn.show() - return self_complete @@ -848,6 +847,12 @@ def prepare_for_save_PRFBO(self, copyClass): if "evaluators" in copyClass.steps[i].__dict__: del copyClass.steps[i].evaluators + # ------------------------------------------------------------------------------------------------- + # Add time stamp + # ------------------------------------------------------------------------------------------------- + + copyClass.timeStamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + return copyClass def prepare_for_read_PRFBO(self, copyClass): @@ -866,7 +871,7 @@ def prepare_for_read_PRFBO(self, copyClass): return copyClass - def save(self, storeClass=True, name="MITIMstate.pkl"): + def save(self, name="MITIMstate.pkl"): print("* Proceeding to save new MITIM state pickle file") stateFile = f"{self.folderOutputs}/{name}" stateFile_tmp = f"{self.folderOutputs}/{name}_tmp" @@ -929,12 +934,11 @@ def read( step = aux.steps[iteration] print( - f"\t* Read {stateFile[-30:]}{'...' if len(stateFile) > 30 else ''} state file, grabbed step #{iteration}", + f"\t* Read {IOtools.clipstr(stateFile)} state file, grabbed step #{iteration}", typeMsg="f", ) - except: - print(f"\n\nCould not read {stateFile} state file, because:", typeMsg="w") - print(traceback.format_exc()) + except FileNotFoundError: + print(f"\t- State file {stateFile} not found", typeMsg="w") step, aux = None, None return aux if provideFullClass else step @@ -1475,17 +1479,16 @@ def plot( GPs = self.steps - if fn is None: + if doNotShow: plt.ioff() + + if fn is None: from mitim_tools.misc_tools.GUItools import FigureNotebook geometry = ( "1200x1000" if len(GPs[0].GP["individual_models"]) == 1 else "1700x1000" ) - fn = FigureNotebook(0, "MITIM BO Strategy", geometry=geometry) - fnprov = False - else: - fnprov = True + fn = FigureNotebook("MITIM BO Strategy", geometry=geometry) """ **************************************************************** @@ -1508,18 +1511,29 @@ def plot( f"- Plotting MITIM step #{k} information ({ck+1}/{len(rangePlot)})... {len(self.outputs)} GP models need to be plotted ({pointsEvaluateEachGPdimension} points/dim)..." ) + tab_color = ck + 5 + figsFund, figs, figsFundTrain = [], [], [] for i in range(Tabs_needed): figsFund.append( - fn.add_figure(label=f"#{k} Fundamental Surr. ({i+1}/{Tabs_needed})") + fn.add_figure( + label=f"#{k} Fundamental Surr. ({i+1}/{Tabs_needed})", + tab_color=tab_color, + ) ) for i in range(Tabs_needed): figsFundTrain.append( - fn.add_figure(label=f"#{k} Fundamental Train ({i+1}/{Tabs_needed})") + fn.add_figure( + label=f"#{k} Fundamental Train ({i+1}/{Tabs_needed})", + tab_color=tab_color, + ) ) for i in range(Tabs_needed): figs.append( - fn.add_figure(label=f"#{k} Surrogate ({i+1}/{Tabs_needed})") + fn.add_figure( + label=f"#{k} Surrogate ({i+1}/{Tabs_needed})", + tab_color=tab_color, + ) ) grid = plt.GridSpec( @@ -1629,7 +1643,7 @@ def plot( ) # Plot model specifics from last model - self.plotModelStatus(boStep=k, fn=fn, stds=stds) + self.plotModelStatus(boStep=k, fn=fn, stds=stds, tab_color=tab_color) print("- Finished plotting of step models") @@ -1639,12 +1653,14 @@ def plot( **************************************************************** """ + tab_color = ck + 5 + 1 + # ---- Trust region ---------------------------------------------------------- - figTR = fn.add_figure(label="Trust Region") + figTR = fn.add_figure(label="Trust Region", tab_color=tab_color) try: SBOcorrections.plotTrustRegionInformation(self, fig=figTR) except: - print("Problem plotting trust region") + print("\t- Problem plotting trust region", typeMsg="w") # ---- ResultsOptimization --------------------------------------------------- if plotResultsOptimization: @@ -1654,20 +1670,21 @@ def plot( logFile = self.logFile else: logFile = None - self.ResultsOptimization.plot(fn=fn, doNotShow=True, log=logFile) - - if (not fnprov) and (not doNotShow): - fn.show() + self.ResultsOptimization.plot( + fn=fn, doNotShow=True, log=logFile, tab_color=tab_color + ) return fn - def plotModelStatus(self, fn=None, boStep=-1, plotsPerFigure=20, stds=2): + def plotModelStatus( + self, fn=None, boStep=-1, plotsPerFigure=20, stds=2, tab_color=None + ): step = self.steps[boStep] GP = step.GP["combined_model"] # ---- Jacobian ------------------------------------------------------- - fig = fn.add_figure(label=f"#{boStep}: Jacobian") + fig = fn.add_figure(label=f"#{boStep}: Jacobian", tab_color=tab_color) maxPoints = 1 # 4 xExplore = [] if "x_next" in step.__dict__.keys(): @@ -1702,7 +1719,9 @@ def plotModelStatus(self, fn=None, boStep=-1, plotsPerFigure=20, stds=2): figsQuality = [] for i in range(numfigs): figsQuality.append( - fn.add_figure(label=f"#{boStep}: Quality {i+1}/{numfigs}") + fn.add_figure( + label=f"#{boStep}: Quality {i+1}/{numfigs}", tab_color=tab_color + ) ) axs = GP.testTraining( @@ -1725,8 +1744,8 @@ def plotModelStatus(self, fn=None, boStep=-1, plotsPerFigure=20, stds=2): if "InfoOptimization" not in step.__dict__.keys(): return - figOPT1 = fn.add_figure(label=f"#{boStep}: Optim Perfom.") - figOPT2 = fn.add_figure(label=f"#{boStep}: Optim Ranges") + figOPT1 = fn.add_figure(label=f"#{boStep}: Optim Perfom.", tab_color=tab_color) + figOPT2 = fn.add_figure(label=f"#{boStep}: Optim Ranges", tab_color=tab_color) self.plotSurrogateOptimization(fig1=figOPT1, fig2=figOPT2, boStep=boStep) # --------------------------------------------------------------------- @@ -1750,16 +1769,11 @@ def plotSurrogateOptimization(self, fig1=None, fig2=None, boStep=-1): colors = GRAPHICStools.listColors() if fig1 is None: - plt.ioff() from mitim_tools.misc_tools.GUItools import FigureNotebook - fn = FigureNotebook(0, "PRF BO Strategy", geometry="1700x1000") + fn = FigureNotebook("PRF BO Strategy", geometry="1700x1000") fig2 = fn.add_figure(label=f"#{boStep}: Optim Ranges") fig1 = fn.add_figure(label=f"#{boStep}: Optim Perfom.") - fnprov = False - - else: - fnprov = True grid = plt.GridSpec(nrows=2, ncols=2, hspace=0.2, wspace=0.2) ax0_r = fig2.add_subplot(grid[:, 0]) @@ -1887,9 +1901,6 @@ def plotSurrogateOptimization(self, fig1=None, fig2=None, boStep=-1): for i in range(len(axs)): GRAPHICStools.addDenseAxis(axs[i]) - if not fnprov: - fn.show() - def read_from_scratch(file): """ diff --git a/src/mitim_tools/opt_tools/SURROGATEtools.py b/src/mitim_tools/opt_tools/SURROGATEtools.py index 5903e0d3..4d5abbde 100644 --- a/src/mitim_tools/opt_tools/SURROGATEtools.py +++ b/src/mitim_tools/opt_tools/SURROGATEtools.py @@ -1,13 +1,17 @@ -import torch, gpytorch, botorch, copy, datetime, contextlib +import torch +import gpytorch +import botorch +import copy +import contextlib import numpy as np import matplotlib.pyplot as plt import dill as pickle_dill -from IPython import embed -from mitim_tools.misc_tools import MATHtools, GRAPHICStools, IOtools -from mitim_tools.opt_tools import BOTORCHtools, OPTtools +from mitim_tools.misc_tools import GRAPHICStools, IOtools +from mitim_tools.opt_tools import BOTORCHtools from mitim_tools.opt_tools.aux import BOgraphics from mitim_tools.misc_tools.IOtools import printMsg as print from mitim_tools.misc_tools.CONFIGread import read_verbose_level +from IPython import embed verbose_level = read_verbose_level() @@ -171,6 +175,13 @@ def __init__( self.train_Y_added = torch.empty((0, dimTransformedDV_y)) self.train_Yvar_added = torch.empty((0, dimTransformedDV_y)) + # -------------------------------------------------------------------------------------- + # Make sure that very small variations are not captured + # -------------------------------------------------------------------------------------- + + if self.train_X_added.shape[0] > 0: + self.ensureMinimalVariationSuppressed(input_transform_physics) + # -------------------------------------------------------------------------------------- # Make sure at least 2 points # -------------------------------------------------------------------------------------- @@ -474,6 +485,34 @@ def writeFileTraining(self, input_transform_physics, outcome_transform_physics): with open(self.fileTraining, "rb") as f: data_dict = pickle_dill.load(f) + if self.train_X_added_full.shape[-1] < train_X_Complete.shape[-1]: + print( + "\t\t- Points from file have less input dimensions, extending with NaNs for writing new file", + typeMsg="w", + ) + self.train_X_added_full = torch.cat( + ( + self.train_X_added_full, + torch.full( + ( + self.train_X_added_full.shape[0], + train_X_Complete.shape[-1] + - self.train_X_added_full.shape[-1], + ), + torch.nan, + ), + ), + axis=-1, + ) + elif self.train_X_added_full.shape[-1] > train_X_Complete.shape[-1]: + print( + "\t\t- Points from file have more input dimensions, removing last dimensions for writing new file", + typeMsg="w", + ) + self.train_X_added_full = self.train_X_added_full[ + :, : train_X_Complete.shape[-1] + ] + x = torch.cat((self.train_X_added_full, train_X_Complete), axis=0) y = torch.cat((self.train_Y_added, train_Y), axis=0) yvar = torch.cat((self.train_Yvar_added, train_Yvar), axis=0) @@ -515,7 +554,6 @@ def testTraining( x_next=None, y_next=None, ystd_next=None, - printYN=False, axs=None, plotsPerFigure=20, ylabels=None, @@ -532,7 +570,6 @@ def testTraining( yPredicted, yU, yL, _ = self.predict(xT) - x = xT.cpu().numpy() y = y.cpu().numpy() yPredicted = yPredicted.detach().cpu().numpy() yL = yL.detach().cpu().numpy() @@ -667,6 +704,36 @@ def testTraining( return axs + def ensureMinimalVariationSuppressed(self, input_transform_physics, thr=1e-6): + """ + In some cases, the added data from file might have extremely small variations in some of the fixed + inputs, as compared to the trained data of this run. In such a case, modify this variation + """ + + # Do dimensions of the non-added points change? + x_transform = input_transform_physics(self.train_X) + indecesUnchanged = torch.where( + (x_transform.max(axis=0)[0] - x_transform.min(axis=0)[0]) + / x_transform.mean(axis=0)[0] + < thr + )[0] + + HasThisBeenApplied = 0 + + for i in indecesUnchanged: + if ( + (self.train_X_added[:, i] - x_transform[0, i]) / x_transform[0, i] + ).max() < thr: + HasThisBeenApplied += 1 + for j in range(self.train_X_added.shape[0]): + self.train_X_added[j, i] = x_transform[0, i] + + if HasThisBeenApplied > 0: + print( + f"\t- Supression of small variations {thr:.1e} in added data applied to {HasThisBeenApplied} dims", + typeMsg="w", + ) + def ensureMinimumNoise(self): if ("MinimumRelativeNoise" in self.surrogateOptions) and ( self.surrogateOptions["MinimumRelativeNoise"] is not None @@ -720,11 +787,11 @@ def assess_optimization(self, track_fval): print("\t- Fitting summary:", verbose=verbose_level) if verbose_level in [4, 5]: - print(f"\t\t* Model raw parameters:") + print("\t\t* Model raw parameters:") for param_name, param in self.gpmodel.named_parameters(): BOgraphics.printParam(param_name, param, extralab="\t\t\t") - print(f"\t\t* Model constraints:") + print("\t\t* Model constraints:") dictParam = {} for constraint_name, constraint in self.gpmodel.named_constraints(): BOgraphics.printConstraint(constraint_name, constraint, extralab="\t\t") @@ -734,7 +801,7 @@ def assess_optimization(self, track_fval): This is an "inconvenient" way to calculate the actual parameters https://docs.gpytorch.ai/en/stable/examples/00_Basic_Usage/Hyperparameters.html?highlight=constraints#How-do-constraints-work? but I like it. """ - print(f"\t\t* Model actual parameters:") + print("\t\t* Model actual parameters:") for param_name, param in self.gpmodel.named_parameters(): if param_name in dictParam: param = dictParam[param_name].transform(param) @@ -803,7 +870,7 @@ def extendPoints(file, output): y = torch.Tensor() yvar = torch.Tensor() for i in data: - list_values, list_valuesE = list(data[i].items()), list(dataE[i].items()) + list_values = list(data[i].items()) x_new = torch.Tensor() for j in range(len(list_values)): @@ -863,8 +930,8 @@ def writeTabulars( TabularData.data[iC][f"x_{j}"] = TabularDataStds.data[iC][ f"x_{j}" ] = round(X[i, j], 16) - TabularData.data[iC][f"y"] = round(Y[i, 0], 16) - TabularDataStds.data[iC][f"y"] = round(Yvar[i, 0] ** 0.5, 16) + TabularData.data[iC]["y"] = round(Y[i, 0], 16) + TabularDataStds.data[iC]["y"] = round(Yvar[i, 0] ** 0.5, 16) iC += 1 outputs.append(output) @@ -899,7 +966,7 @@ def simpleModel( dtype=torch.float64 ) if yvar is not None: - if type(yvar) is float: + if isinstance(yvar, float): yvar = [yvar] yvar = torch.from_numpy(yvar).to(x) diff --git a/src/mitim_tools/opt_tools/aux/BOgraphics.py b/src/mitim_tools/opt_tools/aux/BOgraphics.py index f16db061..32e0a679 100644 --- a/src/mitim_tools/opt_tools/aux/BOgraphics.py +++ b/src/mitim_tools/opt_tools/aux/BOgraphics.py @@ -1,16 +1,17 @@ -import os, pyDOE, copy, collections, pdb, multiprocessing, datetime, torch, sys +import os +import copy +import torch +import sys import dill as pickle_dill import numpy as np import matplotlib as mpl -import pandas as pd -from mpl_toolkits.axes_grid1 import make_axes_locatable from matplotlib import pyplot as plt from scipy.interpolate import griddata from collections import OrderedDict from IPython import embed -from mitim_tools.misc_tools import IOtools, FARMINGtools, GRAPHICStools, GUItools -from mitim_tools.opt_tools import SURROGATEtools, STRATEGYtools +from mitim_tools.misc_tools import IOtools, GRAPHICStools, GUItools, MATHtools +from mitim_tools.opt_tools import STRATEGYtools from mitim_tools.opt_tools.aux import TESTtools from mitim_tools.misc_tools.IOtools import printMsg as print @@ -583,7 +584,7 @@ def plotSensitivities_surrogate_model( def plotTraining_surrogate_model( - self, axs=None, relative_to=-1, figIndex_inner=0, stds=2.0 + self, axs=None, relative_to=-1, figIndex_inner=0, stds=2.0, legYN=True ): colors = GRAPHICStools.listColors() @@ -636,7 +637,7 @@ def plotTraining_surrogate_model( label=newLabels[j], ) - ax2.plot(x, trainYtr[:, 0], "-s", markersize=3, color="b") + ax2.plot(x, trainYtr[:, 0], "-s", markersize=3, color="b", label="train") ax2.errorbar( x, trainYtr[:, 0], @@ -648,7 +649,7 @@ def plotTraining_surrogate_model( mean, upper, lower, _ = self.predict(train_X, produceFundamental=True) mean = mean[:, 0].detach().cpu().numpy() - ax2.plot(x, mean, "-s", color="r", lw=0.5, markersize=3) + ax2.plot(x, mean, "-s", color="r", lw=0.5, markersize=3, label="model") ax2.errorbar( x, mean, @@ -682,6 +683,9 @@ def plotTraining_surrogate_model( GRAPHICStools.addDenseAxis(ax2) ax2.set_xlim(left=0) + if legYN: + ax2.legend(loc="best", prop={"size": 4}) + def localBehavior_surrogate_model( self, x, outputs=None, plotYN=True, ax=None, prefix="", rangeZeroth=1e-6 @@ -825,30 +829,10 @@ def retrieveResults( print(" - Grabbing remote") if not os.path.exists(folderWork): os.system(f"mkdir {folderWork}") - if not os.path.exists(folderWork + "/Outputs/"): - os.system(f"mkdir {folderWork}/Outputs/") - os.system( - f"scp {port} {username}@{machine}:{folderRemote0}/Outputs/ResultsOptimization.out {folderWork}/Outputs/." - ) - os.system( - f"scp {port} {username}@{machine}:{folderRemote0}/Outputs/MITIM.log {folderWork}/Outputs/." - ) + os.system( - f"scp {port} {username}@{machine}:{folderRemote0}/Outputs/MITIMstate* {folderWork}/Outputs/." + f"scp -TO -r {port} {username}@{machine}:{folderRemote0}/Outputs {folderWork}" ) - if analysis_level >= 0: - os.system( - f"scp {port} {username}@{machine}:{folderRemote0}/Outputs/TabularData.dat {folderWork}/Outputs/." - ) - os.system( - f"scp {port} {username}@{machine}:{folderRemote0}/Outputs/TabularDataStds.dat {folderWork}/Outputs/." - ) - if analysis_level > 0: - os.system( - f"scp {port} {username}@{machine}:{folderRemote0}/Outputs/MITIMextra* {folderWork}/Outputs/." - ) - - resFile = f"{folderWork}/ResultsOptimization.out" # ---------------------------------------------------------------------------------------------------------------- # Viewing workflow @@ -857,6 +841,11 @@ def retrieveResults( print("\t\t--> Opening MITIMstate.pkl") prfs_model = STRATEGYtools.read_from_scratch(f"{folderWork}/Outputs/MITIMstate.pkl") + if "timeStamp" in prfs_model.__dict__: + print(f"\t\t\t- Time stamp of MITIMstate.pkl: {prfs_model.timeStamp}") + else: + print("\t\t\t- Time stamp of MITIMstate.pkl not found") + # ---------------- Read ResultsOptimization fileOutputs = folderWork + "/Outputs/ResultsOptimization.out" res = ResultsOptimization(file=fileOutputs) @@ -1339,9 +1328,7 @@ def save(self): print("\t* ResultsOptimization updated", verbose=verbose_level) def read(self): - print( - f"\t\t--> Opening {self.file[-30:]}{'...' if len(self.file) > 30 else ''}" - ) + print(f"\t\t--> Opening {IOtools.clipstr(self.file)}") with open(self.file, "r") as f: lines = f.readlines() @@ -1702,26 +1689,26 @@ def gatherOptima(self, basedOnBest=True): except: break - def plot(self, fn=None, doNotShow=True, separateOFs=False, log=None): + def plot( + self, fn=None, doNotShow=True, separateOFs=False, log=None, tab_color=None + ): if fn is None: - plt.ioff() from mitim_tools.misc_tools.GUItools import FigureNotebook - fn = FigureNotebook(0, "Calibration", geometry="1600x1000") - fnprov = False + self.fn = FigureNotebook("Calibration", geometry="1600x1000") else: - fnprov = True - - fig1 = fn.add_figure(label="Complete") - fig1e = fn.add_figure(label="Complete (rel.)") - fig2 = fn.add_figure(label="Metrics") - fig3 = fn.add_figure(label="Deviations") - fig3b = fn.add_figure(label="Separate") - fig3c = fn.add_figure(label="Together") - fig3cE = fn.add_figure(label="Together All") - fig4 = fn.add_figure(label="Improvement") + self.fn = fn + + fig1 = self.fn.add_figure(label="Complete", tab_color=tab_color) + fig1e = self.fn.add_figure(label="Complete (rel.)", tab_color=tab_color) + fig2 = self.fn.add_figure(label="Metrics", tab_color=tab_color) + fig3 = self.fn.add_figure(label="Deviations", tab_color=tab_color) + fig3b = self.fn.add_figure(label="Separate", tab_color=tab_color) + fig3c = self.fn.add_figure(label="Together", tab_color=tab_color) + fig3cE = self.fn.add_figure(label="Together All", tab_color=tab_color) + fig4 = self.fn.add_figure(label="Improvement", tab_color=tab_color) if log is not None: - figTimes = fn.add_figure(label="Times") + figTimes = self.fn.add_figure(label="Times", tab_color=tab_color) grid = plt.GridSpec(1, 2, hspace=0.3, wspace=0.3) axsTimes = [figTimes.add_subplot(grid[0]), figTimes.add_subplot(grid[1])] @@ -1742,7 +1729,7 @@ def plot(self, fn=None, doNotShow=True, separateOFs=False, log=None): onlyFinals=False, ) self.plotMetrics(fig2) - self.plotCalibrations(figs=[fig3, fig3b, fig3c, fig3cE]) + self.plotCalibrations(figs=[fig3, fig3b, fig3c, fig3cE], tab_color=tab_color) grid = plt.GridSpec(1, 3, hspace=0.3, wspace=0.3) ax0 = fig4.add_subplot(grid[0, 0]) @@ -1756,10 +1743,7 @@ def plot(self, fn=None, doNotShow=True, separateOFs=False, log=None): if log is not None: log.plot(axs=[axsTimes[0], axsTimes[1]]) - if (not fnprov) and (not doNotShow): - fn.show() - - return fn + return self.fn def plotDVs( self, @@ -2219,8 +2203,6 @@ def plotGoodness(self, ax=None, axDiff=None): axDiff.axhline(y=1, ls="--", c="k", lw=1) def plotMetrics(self, fig): - plt.ion() - grid = plt.GridSpec(nrows=2, ncols=1, hspace=0.4, wspace=0.4) ax1 = fig.add_subplot(grid[0]) @@ -2230,20 +2212,17 @@ def plotMetrics(self, fig): return ax1 - def plotCalibrations(self, figs=None): + def plotCalibrations(self, figs=None, tab_color=None): if figs is None: - plt.ioff() from mitim_tools.misc_tools.GUItools import FigureNotebook - fn = FigureNotebook(0, "Calibration", geometry="1600x1000") - fnprov = False - fig3 = fn.add_figure(label="Deviations") - fig3b = fn.add_figure(label="Separate") - fig3c = fn.add_figure(label="Together") - fig3c = fn.add_figure(label="Together All") + self.fnCals = FigureNotebook("Calibration", geometry="1600x1000") + fig3 = self.fnCals.add_figure(label="Deviations", tab_color=tab_color) + fig3b = self.fnCals.add_figure(label="Separate", tab_color=tab_color) + fig3c = self.fnCals.add_figure(label="Together", tab_color=tab_color) + fig3c = self.fnCals.add_figure(label="Together All", tab_color=tab_color) else: [fig3, fig3b, fig3c, fig3cE] = figs - fnprov = True # ---------------- Plot stuff @@ -3286,7 +3265,7 @@ def plotGA_results( plot_colors = GRAPHICStools.listColors() if fn is None: - fn = GUItools.FigureNotebook(0, "MITIM GA Notebook", geometry="1500x1000") + fn = GUItools.FigureNotebook("MITIM GA Notebook", geometry="1500x1000") fig = fn.add_figure(label="Pareto Front" + subname) @@ -3354,17 +3333,18 @@ def plotGA_results( ax = figAnalysis.add_subplot(grid[1, 0]) plotGA_fitness(info, ax=ax) + return fn + def plotGA_essential(GAOF, fn=None, NumGenerations=5, plotAllmembers=False, subname=""): - plt.ioff() if fn is None: - fn = GUItools.FigureNotebook(0, "MITIM GA Notebook", geometry="1500x1000") + fn = GUItools.FigureNotebook("MITIM GA Notebook", geometry="1500x1000") if plotAllmembers: members = GAOF["All_x"] else: members = None - plotGA_results( + fn = plotGA_results( GAOF["Paretos_x"], GAOF["fitness"], GAOF["toolbox"], @@ -3376,6 +3356,8 @@ def plotGA_essential(GAOF, fn=None, NumGenerations=5, plotAllmembers=False, subn trained=GAOF["trained"], ) + return fn + def printConstraint(constraint_name, constraint, extralab=""): print(f"{extralab} * {constraint_name:55} = {constraint}") diff --git a/src/mitim_tools/opt_tools/exe/evaluate_model.py b/src/mitim_tools/opt_tools/exe/evaluate_model.py index 7e31e5e1..0f0ea9bf 100644 --- a/src/mitim_tools/opt_tools/exe/evaluate_model.py +++ b/src/mitim_tools/opt_tools/exe/evaluate_model.py @@ -1,4 +1,5 @@ -import sys, copy, torch, datetime, cProfile, argparse +import torch +import argparse import numpy as np import matplotlib.pyplot as plt from mitim_tools.misc_tools import IOtools diff --git a/src/mitim_tools/opt_tools/exe/evaluate_optimizer_botorch.py b/src/mitim_tools/opt_tools/exe/evaluate_optimizer_botorch.py index f8d3893f..a13170bb 100644 --- a/src/mitim_tools/opt_tools/exe/evaluate_optimizer_botorch.py +++ b/src/mitim_tools/opt_tools/exe/evaluate_optimizer_botorch.py @@ -1,15 +1,18 @@ -import torch, datetime, argparse, botorch, copy +import datetime +import argparse +import botorch import dill as pickle_dill import numpy as np -from IPython import embed import matplotlib.pyplot as plt from mitim_tools.misc_tools import IOtools, GRAPHICStools from mitim_tools.opt_tools import STRATEGYtools, BOTORCHtools, OPTtools +from IPython import embed + """ This script performs a series of tests using BOTORCH optimizer for the last step of the MITIM folder. e.g. - optimizer_tester.py --folder run1/ --test 1 --save True --seeds 10 --var 9 + optimizer_tester.py --folder run1/ --test 1 --save True --seeds 10 --var 9 """ # *************************************************************************************************** @@ -19,7 +22,7 @@ parser = argparse.ArgumentParser() parser.add_argument("--folder", required=True, type=str) parser.add_argument("--test", required=False, type=int, default=1) -parser.add_argument("--save", required=False, type=bool, default=False) +parser.add_argument("--save", required=False, default=False, action="store_true") parser.add_argument("--seeds", required=False, type=int, default=1) parser.add_argument("--var", required=False, type=int, default=9) args = parser.parse_args() diff --git a/src/mitim_tools/opt_tools/exe/read.py b/src/mitim_tools/opt_tools/exe/read.py index 0af12f78..cfe15ce5 100644 --- a/src/mitim_tools/opt_tools/exe/read.py +++ b/src/mitim_tools/opt_tools/exe/read.py @@ -1,14 +1,15 @@ -import sys, argparse, copy, torch +import argparse +import copy import matplotlib.pyplot as plt import numpy as np from mitim_tools.opt_tools.aux import BOgraphics from mitim_tools.misc_tools import IOtools, GRAPHICStools from mitim_tools.opt_tools import STRATEGYtools -from IPython import embed - from mitim_tools.misc_tools.IOtools import printMsg as print from mitim_tools.misc_tools.CONFIGread import read_verbose_level +from IPython import embed + verbose_level = read_verbose_level() """ @@ -43,8 +44,8 @@ parser.add_argument( "--type", type=int, required=False, default=-1 ) # 0: Only ResultsOpt plotting, 1: Also pickle, 2: Also final analysis, 3: Others -parser.add_argument("--folders", required=True, type=str, nargs="*") -parser.add_argument("--remote", type=str, required=False, default=None) +parser.add_argument("folders", type=str, nargs="*") +parser.add_argument("--remote", "-r", type=str, required=False, default=None) parser.add_argument("--seeds", type=int, required=False, default=None) parser.add_argument("--resolution", type=int, required=False, default=50) parser.add_argument("--save", type=str, required=False, default=None) @@ -251,7 +252,7 @@ def plotCompare(folders, plotMeanMax=[True, False]): else: retrieval_level = analysis_level -txt = f"***************************************************************************\n" +txt = "***************************************************************************\n" for folderWork in foldersWork: txt += f"* Reading results in {folderWork}\n" @@ -319,7 +320,7 @@ def plotCompare(folders, plotMeanMax=[True, False]): print(f"Plotting Violin with {xf.shape[0]} points") GRAPHICStools.plotViolin([xf], labels=["run"], ax=ax, colors=["b"]) - ax.set_xlabel(f"Number of evaluations to converge") + ax.set_xlabel("Number of evaluations to converge") # ax.set_title(f'Residual reduced by x{1/percent:.0f}') ax.set_xlim([0, 50]) diff --git a/src/mitim_tools/opt_tools/exe/save_model_figures.py b/src/mitim_tools/opt_tools/exe/save_model_figures.py index ba07fa2a..eb10912a 100644 --- a/src/mitim_tools/opt_tools/exe/save_model_figures.py +++ b/src/mitim_tools/opt_tools/exe/save_model_figures.py @@ -34,7 +34,7 @@ if not os.path.exists(save): os.system(f"mkdir {save}") -plt.ioff() + for i in range(len(gps)): name_save = f"{save}/model{i+1}" diff --git a/src/mitim_tools/opt_tools/exe/slurm.py b/src/mitim_tools/opt_tools/exe/slurm.py index ddb1871e..371cc7fa 100644 --- a/src/mitim_tools/opt_tools/exe/slurm.py +++ b/src/mitim_tools/opt_tools/exe/slurm.py @@ -1,6 +1,5 @@ -import sys, os, argparse -import numpy as np -from IPython import embed +import os +import argparse from mitim_tools.misc_tools import FARMINGtools, IOtools """ @@ -33,7 +32,7 @@ def commander( f"python3 {script} {folderWork} --seed {seed}", ] - comm, fileSBTACH, fileSHELL = FARMINGtools.SLURM( + _, fileSBTACH, _ = FARMINGtools.create_slurm_execution_files( command, folderWork, None, @@ -50,17 +49,18 @@ def commander( def run_slurm( script, - num, + folder, partition, venv, seeds=1, - MainFolder="./", hours=8, n=32, seed_specific=0, ): script = IOtools.expandPath(script) - folderWork = IOtools.expandPath(f"{MainFolder}/{num}") + folderWork = IOtools.expandPath(folder) + + num = folder.split("/")[-1] if seeds > 1: for i in range(seeds): @@ -92,25 +92,29 @@ def run_slurm( if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("script", type=str) - parser.add_argument("name", type=str) + parser.add_argument("--folder", type=str, required=False, default="run1/") parser.add_argument( - "--partition", type=str, required=False, default="sched_mit_psfc" + "--partition", + type=str, + required=False, + default=IOtools.expandPath("$MITIM_PARTITION"), ) parser.add_argument("--seeds", type=int, required=False, default=1) + parser.add_argument( + "--env", type=str, required=False, default=IOtools.expandPath("~/env/mitim-env") + ) parser.add_argument("--hours", type=int, required=False, default=8) parser.add_argument("--n", type=int, required=False, default=64) parser.add_argument("--seed_specific", type=int, required=False, default=0) args = parser.parse_args() - venv = IOtools.expandPath("~/.env/mitim-env") - # Run run_slurm( args.script, - args.name, + args.folder, args.partition, - venv, + args.env, seeds=args.seeds, hours=args.hours, n=args.n, diff --git a/src/mitim_tools/popcon_tools/FunctionalForms.py b/src/mitim_tools/popcon_tools/FunctionalForms.py index eaeff2b1..d8a1e6ba 100644 --- a/src/mitim_tools/popcon_tools/FunctionalForms.py +++ b/src/mitim_tools/popcon_tools/FunctionalForms.py @@ -1,4 +1,3 @@ -import torch, pdb import numpy as np import matplotlib.pyplot as plt from IPython import embed diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index fb10d6b2..02aa94ef 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -1,12 +1,15 @@ -import sys, os, pickle, copy, datetime, netCDF4, warnings +import os +import pickle +import copy +import datetime +import netCDF4 +import warnings import numpy as np import matplotlib.pyplot as plt import matplotlib.patches as patches from collections import OrderedDict -from IPython import embed from mitim_tools.misc_tools import ( IOtools, - FARMINGtools, MATHtools, PLASMAtools, GRAPHICStools, @@ -25,6 +28,7 @@ from mitim_tools.misc_tools.GUItools import FigureNotebook from mitim_tools.misc_tools.IOtools import printMsg as print from mitim_tools.misc_tools.CONFIGread import read_verbose_level +from IPython import embed verbose_level = read_verbose_level() @@ -34,7 +38,6 @@ def __init__( self, netCDFfile, ssh=None, - printCheckPoints=False, ZerothTime=False, readFBM=False, readTGLF=False, @@ -62,10 +65,6 @@ def __init__( np.seterr(under="ignore") - if printCheckPoints: - timeInitial = datetime.datetime.now() - itcheck = 1 - netCDFfile = IOtools.expandPath(netCDFfile, ensurePathValid=True) # ~~~~~~~~ Open CDF file ~~~~~~~~ @@ -204,17 +203,6 @@ def __init__( if self.fbm_He4_gc is not None: print("\t\t- Gathered He4 FBM files") - # D beam - if self.nbD_avol.max() > thr: - self.fbm_Dbeam_gc, self.fbm_Dbeam_po = FBMtools.getFBMprocess( - self.folderWork, - self.nameRunid, - datanum=datanum, - FBMparticle="D_NBI", - ) - if self.fbm_Dbeam_gc is not None: - print("\t\t- Gathered D beam FBM files") - # T (D+D) if self.nfusT_avol.max() > thr: self.fbm_T_gc, self.fbm_T_po = FBMtools.getFBMprocess( @@ -226,6 +214,17 @@ def __init__( if self.fbm_T_gc is not None: print("\t\t- Gathered T FBM files") + # D beam + if self.nbD_avol.max() > thr: + self.fbm_Dbeam_gc, self.fbm_Dbeam_po = FBMtools.getFBMprocess( + self.folderWork, + self.nameRunid, + datanum=datanum, + FBMparticle="D_NBI", + ) + if self.fbm_Dbeam_gc is not None: + print("\t\t- Gathered D beam FBM files") + # ~~~~~~~~ TORIC ~~~~~~~~ readTORIC = readTORIC and np.sum(self.PichT) > 0.0 + self.eps00 * ( @@ -879,7 +878,7 @@ def getTGLFparameters(self, mi_u=None, correctGrid=True): self.TGLF_a = constant_radius(self.a, lenX=len(self.x_lw)) - self.TGLF_Qgb, self.TGLF_Ggb, _, _ = PLASMAtools.gyrobohmUnits( + self.TGLF_Qgb, self.TGLF_Ggb, _, _, _ = PLASMAtools.gyrobohmUnits( self.TGLF_Te, self.TGLF_ne, self.TGLF_mref, self.TGLF_Bunit_x, self.TGLF_a ) @@ -1332,19 +1331,25 @@ def getFusionIons(self): self.nfusHe3_avol = ( volumeAverage(self.f, "FDENS_3") * 1e6 * 1e-20 ) # in 10^20m^-3 + except KeyError: + self.nfusHe3 = self.nD * 0.0 + self.eps00 + self.nfusHe3_avol = self.nD_avol * 0.0 + self.eps00 + + try: self.nfusT = self.f["FDENS_T"][:] * 1e6 * 1e-20 # in 10^20m^-3 self.nfusT_avol = ( volumeAverage(self.f, "FDENS_T") * 1e6 * 1e-20 ) # in 10^20m^-3 + except KeyError: + self.nfusT = self.nD * 0.0 + self.eps00 + self.nfusT_avol = self.nD_avol * 0.0 + self.eps00 + + try: self.nfusH = self.f["FDENS_P"][:] * 1e6 * 1e-20 # in 10^20m^-3 self.nfusH_avol = ( volumeAverage(self.f, "FDENS_P") * 1e6 * 1e-20 ) # in 10^20m^-3 - except: - self.nfusHe3 = self.nD * 0.0 + self.eps00 - self.nfusHe3_avol = self.nD_avol * 0.0 + self.eps00 - self.nfusT = self.nD * 0.0 + self.eps00 - self.nfusT_avol = self.nD_avol * 0.0 + self.eps00 + except KeyError: self.nfusH = self.nD * 0.0 + self.eps00 self.nfusH_avol = self.nD_avol * 0.0 + self.eps00 @@ -5669,17 +5674,15 @@ def plotAllGeos(self, ax, time): def plotISOLVER(self, fn=None, time=None): if fn is None: - std = True - plt.ioff() - fn = FigureNotebook(0, f"ISOLVER Notebook, run #{self.nameRunid}") + self.fnIsolver = FigureNotebook(f"ISOLVER Notebook, run #{self.nameRunid}") else: - std = False + self.fnIsolver = fn - fig1 = fn.add_figure(label="ISOLVER_eq") - fig2 = fn.add_figure(label="ISOLVER_coils 1") - fig2_e = fn.add_figure(label="ISOLVER_coils 2") - fig3 = fn.add_figure(label="ISOLVER_sum") - fig4 = fn.add_figure(label="ISOLVER_lcfs") + fig1 = self.fnIsolver.add_figure(label="ISOLVER_eq") + fig2 = self.fnIsolver.add_figure(label="ISOLVER_coils 1") + fig2_e = self.fnIsolver.add_figure(label="ISOLVER_coils 2") + fig3 = self.fnIsolver.add_figure(label="ISOLVER_sum") + fig4 = self.fnIsolver.add_figure(label="ISOLVER_lcfs") if time is None: time = self.t[self.ind_saw] @@ -5755,9 +5758,6 @@ def plotISOLVER(self, fn=None, time=None): self.isolver.plotEquilibria(time, fig=fig1) - if std: - fn.show() - def plotTGLF(self, figGR=None, figFL=None): ax1, ax2, ax3 = self.TGLF.plotComplete_GR(fig=figGR) ax1, ax2, ax3 = self.TGLF.plotComplete_FL(fig=figFL) @@ -6892,10 +6892,10 @@ def plotSawtooth(self, fig=None, time=None): # Detail if self.calcualtePorcelli: self.plotPorcelliInternals(ax=ax1s) + GRAPHICStools.addLegendApart(ax1s, ratio=0.7) ax1s.set_title("Porcelli Parameters") ax1s.set_xlabel("Time (s)") GRAPHICStools.addDenseAxis(ax1s) - GRAPHICStools.addLegendApart(ax1s, ratio=0.7) ax2s.plot(self.t, self.porcelli_s1, c="b", label="$s_1$") ax2s.legend(loc="upper left", prop={"size": self.mainLegendSize}) @@ -8148,7 +8148,7 @@ def plotSeparateSystems(self, fig=None): # ax1.set_ylim([0,20]) # ax1.set_ylabel('Error (%)') - ax.legend(loc="best", prop={"size": self.mainLegendSize}) + ax.legend(loc="best", prop={"size": self.mainLegendSize}) ax.set_title("ICRF") ax.set_ylabel("Power Antenna ($MW$)") ax.set_ylim(bottom=0) @@ -8159,7 +8159,7 @@ def plotSeparateSystems(self, fig=None): if np.sum(self.PichT) > 1.0e-5: for i in range(len(self.FichT_ant)): ax.plot(self.t, self.FichT_ant[i], lw=2, label=f"{i + 1}") - ax.legend(loc="best", prop={"size": self.mainLegendSize}) + ax.legend(loc="best", prop={"size": self.mainLegendSize}) ax.set_ylabel("Frequency Antenna ($MHz$)") ax.set_xlabel("Time (s)") ax.set_ylim(bottom=0) @@ -10118,19 +10118,18 @@ def plotImpurities(self, fig=None, time=None, maxStates=10): ax.set_ylim([0, self.fZAVE_Z[it, 0] + 5]) - def diagramSpecies(self, axs=None, time=None, fn=None, label=""): + def diagramSpecies(self, time=None, fn=None, label="", fn_color=None): if time is None: it = self.ind_saw else: it = np.argmin(np.abs(self.t - time)) - if axs is None: - if fn is None: - fig, axs = plt.subplots( - ncols=1, sharex=True, sharey=False, figsize=(14, 9) - ) - else: - fig, axs = fn.subplots(ncols=1, sharex=True, sharey=False, label=label) + if fn is None: + plt.ion() + _, axs = plt.subplots(ncols=1, sharex=True, sharey=False, figsize=(14, 9)) + else: + fig = fn.add_figure(label=label, tab_color=fn_color) + axs = fig.subplots(ncols=1, sharex=True, sharey=False) ax = axs ax.axis("off") @@ -10528,7 +10527,7 @@ def diagramSpecies(self, axs=None, time=None, fn=None, label=""): transform=ax.transAxes, ) - def diagramFlows(self, axs=None, time=None, fn=None, label=""): + def diagramFlows(self, axs=None, time=None, fn=None, label="", fn_color=None): if time is None: it = self.ind_saw else: @@ -10540,7 +10539,8 @@ def diagramFlows(self, axs=None, time=None, fn=None, label=""): ncols=1, sharex=True, sharey=False, figsize=(14, 9) ) else: - fig, axs = fn.subplots(ncols=1, sharex=True, sharey=False, label=label) + fig = fn.add_figure(label=label, tab_color=fn_color) + axs = fig.subplots(ncols=1, sharex=True, sharey=False) ax = axs ax.axis("off") @@ -13126,248 +13126,251 @@ def plotPerformance(self, fig=None, time=None): GRAPHICStools.addLegendApart(ax, ratio=0.7, withleg=False) - def plotRun(self, time=None, timesAv=None, plot_analysis=True): - plt.ioff() - + def plot(self, fn=None, time=None, timesAv=None, plot_analysis=False, counter=0): if time is None: time = self.t[self.ind_saw] name = f"MITIM Notebook, run #{self.nameRunid}, profiles at time t={time:.3f}s" - fn = FigureNotebook(0, name) + fn_color = counter if counter > 0 else None + + if fn is None: + self.fn = FigureNotebook(name) + else: + self.fn = fn # Machine - fig = fn.add_figure(label="Machine") + fig = self.fn.add_figure(tab_color=fn_color, label="Machine") self.plotMachine(fig=fig, time=time) # Equil - fig = fn.add_figure(label="Equilibrium") + fig = self.fn.add_figure(tab_color=fn_color, label="Equilibrium") self.plotEquilParams(fig=fig, time=time) # GS - fig = fn.add_figure(label="Grad-Shafranov") + fig = self.fn.add_figure(tab_color=fn_color, label="Grad-Shafranov") self.plotGS(fig=fig, time=time) # Geometry - fig = fn.add_figure(label="Geometry") + fig = self.fn.add_figure(tab_color=fn_color, label="Geometry") self.plotGEO(fig=fig, time=time) # Profiles - fig = fn.add_figure(label="Profiles") + fig = self.fn.add_figure(tab_color=fn_color, label="Profiles") self.plotProfiles(fig=fig, time=time) # PRESSURES - fig = fn.add_figure(label="Pressure") + fig = self.fn.add_figure(tab_color=fn_color, label="Pressure") self.plotPressures(fig=fig, time=time) # Systems - fig = fn.add_figure(label="Power (Auxiliary)") + fig = self.fn.add_figure(tab_color=fn_color, label="Power (Auxiliary)") self.plotSeparateSystems(fig=fig) # Heating - fig = fn.add_figure(label="Power (Total)") + fig = self.fn.add_figure(tab_color=fn_color, label="Power (Total)") self.plotHeating(fig=fig) # Radial Powe3 - fig = fn.add_figure(label="Power (Radial)") - fig2 = fn.add_figure(label="Power (Cumul.)") + fig = self.fn.add_figure(tab_color=fn_color, label="Power (Radial)") + fig2 = self.fn.add_figure(tab_color=fn_color, label="Power (Cumul.)") self.plotRadialPower(time=time, fig=fig, figCum=fig2) # ICRF if np.sum(self.PichT) > 0.0 + self.eps00 * (len(self.t) + 1): - fig = fn.add_figure(label="ICRF (Total)") + fig = self.fn.add_figure(tab_color=fn_color, label="ICRF (Total)") self.plotICRF_t(fig=fig) - fig = fn.add_figure(label="ICRF (Radial)") + fig = self.fn.add_figure(tab_color=fn_color, label="ICRF (Radial)") self.plotICRF(fig=fig, time=time) # ECRF if np.sum(self.PechT) > 0.0 + self.eps00 * (len(self.t) + 1): - fig = fn.add_figure(label="ECRF") + fig = self.fn.add_figure(tab_color=fn_color, label="ECRF") self.plotECRF(fig=fig, time=time) # NBI if np.sum(self.PnbiT) > 0.0 + self.eps00 * (len(self.t) + 1): - fig = fn.add_figure(label="NBI") + fig = self.fn.add_figure(tab_color=fn_color, label="NBI") self.plotNBI(fig=fig, time=time) # LH if np.sum(self.PlhT) > 0.0 + 2 * self.eps00 * (len(self.t) + 1): - fig = fn.add_figure(label="LowerHyb") + fig = self.fn.add_figure(tab_color=fn_color, label="LowerHyb") self.plotLowerHybrid(fig=fig, time=time) # Transport - fig = fn.add_figure(label="Transport") + fig = self.fn.add_figure(tab_color=fn_color, label="Transport") self.plotTransport(fig=fig, time=time) # Derivatives - fig = fn.add_figure(label="Gradients") + fig = self.fn.add_figure(tab_color=fn_color, label="Gradients") self.plotDerivatives(fig=fig, time=time) # Porcelli - fig = fn.add_figure(label="Sawtooth Trigger") + fig = self.fn.add_figure(tab_color=fn_color, label="Sawtooth Trigger") self.plotSawtooth(fig=fig, time=time) # Around sawtooth # try: - fig = fn.add_figure(label="Sawtooth Effect") + fig = self.fn.add_figure(tab_color=fn_color, label="Sawtooth Effect") try: self.plotAroundSawtoothQuantities(fig=fig) except: print("Could not plot plotAroundSawtoothQuantities", typeMsg="w") - fig = fn.add_figure(label="Sawtooth Mixing") + fig = self.fn.add_figure(tab_color=fn_color, label="Sawtooth Mixing") self.plotSawtoothMixing(fig=fig) # Electric Field - fig = fn.add_figure(label="Current Diffusion") + fig = self.fn.add_figure(tab_color=fn_color, label="Current Diffusion") self.plotEM(fig=fig, time=time) - fig = fn.add_figure(label="Poynting") + fig = self.fn.add_figure(tab_color=fn_color, label="Poynting") self.plotUmag(fig=fig, time=time) # Fundamental - fig = fn.add_figure(label="Fundamental Plasma") + fig = self.fn.add_figure(tab_color=fn_color, label="Fundamental Plasma") self.plotFundamental(fig=fig, time=time) # Stability - fig = fn.add_figure(label="MHD Stability") + fig = self.fn.add_figure(tab_color=fn_color, label="MHD Stability") self.plotStability(fig=fig, time=time) # Electric Field - fig = fn.add_figure(label="Electric Field") + fig = self.fn.add_figure(tab_color=fn_color, label="Electric Field") self.plotElectricField(fig=fig, time=time) # Rotation - fig = fn.add_figure(label="Rotation") + fig = self.fn.add_figure(tab_color=fn_color, label="Rotation") self.plotRotation(fig=fig, time=time) # Time scales try: - fig = fn.add_figure(label="Time Scales") + fig = self.fn.add_figure(tab_color=fn_color, label="Time Scales") self.plotTimeScales(fig=fig, time=time) except: pass - fig = fn.add_figure(label="Averaging") + fig = self.fn.add_figure(tab_color=fn_color, label="Averaging") self.plotTimeAverages(fig=fig, times=timesAv) # Impurities - fig = fn.add_figure(label="Impurities") + fig = self.fn.add_figure(tab_color=fn_color, label="Impurities") self.plotImpurities(fig=fig, time=time) # Radiation - fig = fn.add_figure(label="Radiation") + fig = self.fn.add_figure(tab_color=fn_color, label="Radiation") self.plotRadiation(fig=fig, time=time) # Convergence - fig = fn.add_figure(label="Flux Matching") + fig = self.fn.add_figure(tab_color=fn_color, label="Flux Matching") self.checkRun(fig=fig, time=time, printYN=False) # LH - fig = fn.add_figure(label="LH Transition") + fig = self.fn.add_figure(tab_color=fn_color, label="LH Transition") self.plotLH(fig=fig, time=time) # Particle Balance - fig = fn.add_figure(label="Particle Balance") + fig = self.fn.add_figure(tab_color=fn_color, label="Particle Balance") self.plotParticleBalance(fig=fig, time=time) - fig = fn.add_figure(label="Ions Balance") + fig = self.fn.add_figure(tab_color=fn_color, label="Ions Balance") self.plotIonsBalance(fig=fig, time=time) # SLow down if self.neutrons_thrDT[-1] > self.eps00: - fig = fn.add_figure(label="Slow Down") + fig = self.fn.add_figure(tab_color=fn_color, label="Slow Down") self.plotSlowDown(fig=fig, time=time) # Fast - fig = fn.add_figure(label="Fast (Radial)") + fig = self.fn.add_figure(tab_color=fn_color, label="Fast (Radial)") self.plotFast(fig=fig, time=time) # Fast - fig = fn.add_figure(label="Fast (Stabilization)") + fig = self.fn.add_figure(tab_color=fn_color, label="Fast (Stabilization)") self.plotFast2(fig=fig, time=time) # SLow down if self.neutrons_thrDT[-1] > self.eps00: - fig = fn.add_figure(label="Fast (Transport)") + fig = self.fn.add_figure(tab_color=fn_color, label="Fast (Transport)") self.plotFastTransport(fig=fig, time=time) # Neutrals - fig = fn.add_figure(label="Neutrals") + fig = self.fn.add_figure(tab_color=fn_color, label="Neutrals") self.plotNeutrals(fig=fig, time=time) # Performance - fig = fn.add_figure(label="Performance") + fig = self.fn.add_figure(tab_color=fn_color, label="Performance") self.plotPerformance(fig=fig, time=time) # Nuclear - fig = fn.add_figure(label="Neutrons") + fig = self.fn.add_figure(tab_color=fn_color, label="Neutrons") self.plotNuclear(fig=fig, time=time) # Boundary - fig = fn.add_figure(label="Boundary") + fig = self.fn.add_figure(tab_color=fn_color, label="Boundary") self.plotDivertor(fig=fig, time=time) # ----------- DIAGRAMS # Species - self.diagramSpecies(time=time, fn=fn, label="Species") + self.diagramSpecies(time=time, fn=self.fn, label="Species", fn_color=fn_color) # Flows - self.diagramFlows(time=time, fn=fn, label="Flows") + self.diagramFlows(time=time, fn=self.fn, label="Flows", fn_color=fn_color) # CPU - fig = fn.add_figure(label="CPU usage") + fig = self.fn.add_figure(tab_color=fn_color, label="CPU usage") self.plotCPUperformance(fig=fig, time=time) # ----------- EXTRA # TORIC for i, toric in enumerate(self.torics): - fig = fn.add_figure(label=f"TORIC #{i+1}") + fig = self.fn.add_figure(tab_color=fn_color, label=f"TORIC #{i+1}") self.plotTORIC(fig=fig, position=i) # FBM if self.fbm_He4_gc is not None: - fig = fn.add_figure(label="FBM He4 GC") + fig = self.fn.add_figure(tab_color=fn_color, label="FBM He4 GC") self.plotFBM( fig=fig, particleFile=self.fbm_He4_gc, finalProfile=self.nfusHe4 ) - fig = fn.add_figure(label="FBM He4 PO") + fig = self.fn.add_figure(tab_color=fn_color, label="FBM He4 PO") self.plotFBM( fig=fig, particleFile=self.fbm_He4_po, finalProfile=self.nfusHe4 ) if self.fbm_Dbeam_gc is not None: - fig = fn.add_figure(label="FBM Dbeam GC") + fig = self.fn.add_figure(tab_color=fn_color, label="FBM Dbeam GC") self.plotFBM(fig=fig, particleFile=self.fbm_Dbeam_gc, finalProfile=self.nbD) - fig = fn.add_figure(label="FBM Dbeam PO") + fig = self.fn.add_figure(tab_color=fn_color, label="FBM Dbeam PO") self.plotFBM(fig=fig, particleFile=self.fbm_Dbeam_po, finalProfile=self.nbD) if self.fbm_Dbeam_gc.birth is not None: - fig = fn.add_figure(label="BIRTH D") + fig = self.fn.add_figure(tab_color=fn_color, label="BIRTH D") self.plotBirth(particleFile=self.fbm_Dbeam_gc, fig=fig) if self.fbm_T_gc is not None: - fig = fn.add_figure(label="FBM T GC") + fig = self.fn.add_figure(tab_color=fn_color, label="FBM T GC") self.plotFBM(fig=fig, particleFile=self.fbm_T_gc, finalProfile=self.nfusT) - fig = fn.add_figure(label="FBM T PO") + fig = self.fn.add_figure(tab_color=fn_color, label="FBM T PO") self.plotFBM(fig=fig, particleFile=self.fbm_T_po, finalProfile=self.nfusT) # TGLF if hasattr(self, "TGLF") and self.TGLF is not None: - figGR = fn.add_figure(label="TGLF1") - figFL = fn.add_figure(label="TGLF2") + figGR = self.fn.add_figure(tab_color=fn_color, label="TGLF1") + figFL = self.fn.add_figure(tab_color=fn_color, label="TGLF2") self.plotTGLF(figGR=figGR, figFL=figFL) # ~~~~~~~~~~~~~ Comparisons if hasattr(self, "exp") and self.exp is not None: - fig = fn.add_figure(label="EXP") + fig = self.fn.add_figure(tab_color=fn_color, label="EXP") self.plotComparison(fig=fig) # ~~~~~~~~~~~~~ Comparisons if hasattr(self, "isolver") and self.isolver is not None: - self.plotISOLVER(fn=fn, time=time) + self.plotISOLVER(fn=self.fn, time=time) # ~~~~~~~~~~~~~ Final g-file # Here I make the exception of reading it at plotting, because I may have generated it since loading the class @@ -13382,10 +13385,12 @@ def plotRun(self, time=None, timesAv=None, plot_analysis=True): if self.gfile_out is not None: if self.gfile_in is None: - ax_plasma = self.gfile_out.plot(fn=fn, extraLabel="G_out - ") + ax_plasma = self.gfile_out.plot(fn=self.fn, extraLabel="G_out - ") else: - ax_plasma = GEQtools.compareGeqdsk( - [self.gfile_in, self.gfile_out], fn=fn, labelsGs=["G_in", "G_out"] + ax_plasma, fnG = GEQtools.compareGeqdsk( + [self.gfile_in, self.gfile_out], + fn=self.fn, + labelsGs=["G_in", "G_out"], ) it = np.argmin(np.abs(self.t - time)) @@ -13453,20 +13458,22 @@ def plotRun(self, time=None, timesAv=None, plot_analysis=True): if plot_analysis: # Pulse try: - fig = fn.add_figure(label="ANALYSIS - Heat Pulse") + fig = self.fn.add_figure( + tab_color=fn_color, label="ANALYSIS - Heat Pulse" + ) self.plotPulse(fig=fig) except: pass - fig = fn.add_figure(label="ANALYSIS - initial") + fig = self.fn.add_figure(tab_color=fn_color, label="ANALYSIS - initial") self.analyze_initial(fig=fig) if len(self.tlastsawU) > 1: - fig = fn.add_figure(label="ANALYSIS - sawtooth") + fig = self.fn.add_figure( + tab_color=fn_color, label="ANALYSIS - sawtooth" + ) self.analyze_sawtooth(fig=fig) - fn.show() - # -------------------------------------- # Additional analysis # -------------------------------------- @@ -13845,22 +13852,19 @@ def runTGLFstandalone( # ----------------------------------------------------------------------------------- if plotCompare: - plt.ioff() - fn = FigureNotebook( - 0, "TGLF-TRANSP Notebook", geometry="1500x900", vertical=True + self.fn_std = FigureNotebook( + "TGLF-TRANSP Notebook", geometry="1500x900", vertical=True ) - self.TGLFstd[nameF].plotRun(labels=[labelTGLF], fn=fn) + self.TGLFstd[nameF].plot(labels=[labelTGLF], fn=self.fn_std) - fig1 = fn.add_figure(label="Comparison Flux") + fig1 = self.fn_std.add_figure(label="Comparison Flux") self.plotStdTRANSP(fig=fig1, tglfRun=labelTGLF, time=time) - fig2 = fn.add_figure(label="Comparison GR") + fig2 = self.fn_std.add_figure(label="Comparison GR") self.plotGRTRANSP(fig=fig2, tglfRun=labelTGLF, time=time) - fig3 = fn.add_figure(label="Comparison FL") + fig3 = self.fn_std.add_figure(label="Comparison FL") self.plotFLTRANSP(fig=fig3, tglfRun=labelTGLF, time=time) - fn.show() - return self.TGLFstd[nameF] def transportAnalysis( @@ -14464,12 +14468,10 @@ def compareChiPert( # Plot # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if plotYN: - plt.ioff() - - fn = FigureNotebook(0, f"mitim Notebook, run #{self.nameRunid}") - fig1 = fn.add_figure(label="Compare") - fig2 = fn.add_figure(label="Pulse") - fig3 = fn.add_figure(label="TGLF") + self.fnChiPert = FigureNotebook(f"mitim Notebook, run #{self.nameRunid}") + fig1 = self.fnChiPert.add_figure(label="Compare") + fig2 = self.fnChiPert.add_figure(label="Pulse") + fig3 = self.fnChiPert.add_figure(label="TGLF") grid = plt.GridSpec(2, 2, hspace=0.6, wspace=0.2) ax00 = fig1.add_subplot(grid[0, 0]) @@ -14543,8 +14545,6 @@ def compareChiPert( # Plot TGLF calculation details self.ChiPert_tglf.plotAnalysis(labels=["chi_pert"], fig=fig3) - fn.show() - # -------- def getGFILE(self): diff --git a/src/mitim_tools/transp_tools/exe/read_transp.py b/src/mitim_tools/transp_tools/exe/read_transp.py index e5e81a75..6583e945 100644 --- a/src/mitim_tools/transp_tools/exe/read_transp.py +++ b/src/mitim_tools/transp_tools/exe/read_transp.py @@ -1,29 +1,38 @@ -import argparse, socket -import numpy as np -import matplotlib.pyplot as plt +import argparse from mitim_tools.transp_tools import CDFtools -from mitim_tools.misc_tools.CONFIGread import read_verbose_level - -verbose_level = read_verbose_level() - parser = argparse.ArgumentParser() -parser.add_argument("--files", required=True, type=str, nargs="*") +parser.add_argument("files", type=str, nargs="*") +parser.add_argument( + "--full", "-f", required=False, default=True, action="store_true" # Full read +) +parser.add_argument( + "--read", "-r", required=False, default=False, action="store_true" # Only read +) args = parser.parse_args() expl = args.files +plotYN = not args.read +fullYN = args.full cdfs = [] -readFBM = True -readTGLF = True -readTORIC = True -readGFILE = True -readStructures = True + ZerothTime = False -printCheckPoints = False -readGEQDSK = True +if fullYN: + readFBM = True + readTGLF = True + readTORIC = True + readGFILE = True + readStructures = True + readGEQDSK = True +else: + readFBM = False + readTGLF = False + readTORIC = False + readGFILE = False + readStructures = False + readGEQDSK = False -plt.ioff() for i in expl: if ":" in i: @@ -31,7 +40,6 @@ CDFtools.CDFreactor( i.split(":")[1], ssh=i.split(":")[0], - printCheckPoints=printCheckPoints, readFBM=readFBM, readTGLF=readTGLF, readTORIC=readTORIC, @@ -45,7 +53,6 @@ cdfs.append( CDFtools.CDFreactor( i, - printCheckPoints=printCheckPoints, readFBM=readFBM, readTGLF=readTGLF, readTORIC=readTORIC, @@ -55,3 +62,10 @@ ZerothTime=ZerothTime, ) ) + +if plotYN: + from mitim_tools.misc_tools.GUItools import FigureNotebook + + fn = FigureNotebook("TRANSP run") + for i in range(len(cdfs)): + cdfs[i].plot(fn=fn, counter=i) diff --git a/src/mitim_tools/transp_tools/exe/run_clean.py b/src/mitim_tools/transp_tools/exe/run_clean.py index adfd6a50..94127f75 100644 --- a/src/mitim_tools/transp_tools/exe/run_clean.py +++ b/src/mitim_tools/transp_tools/exe/run_clean.py @@ -5,10 +5,9 @@ run_clean.py 88664P CMOD --numbers 1 --to 5 """ -import sys, os, argparse +import argparse import numpy as np -from mitim_tools.misc_tools import CONFIGread -from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools import CONFIGread, IOtools from mitim_tools.transp_tools import TRANSPtools # User inputs diff --git a/src/mitim_tools/transp_tools/routines/compareNML.py b/src/mitim_tools/transp_tools/routines/compareNML.py index f29db9a9..5002566f 100644 --- a/src/mitim_tools/transp_tools/routines/compareNML.py +++ b/src/mitim_tools/transp_tools/routines/compareNML.py @@ -1,7 +1,9 @@ -import sys, copy +import sys import numpy as np -from IPython import embed from mitim_tools.misc_tools import IOtools +from IPython import embed + +from mitim_tools.misc_tools.IOtools import printMsg as print """ This is used to commpare namelists values @@ -38,7 +40,7 @@ def separateArrays(d): # Separate values in commas dnew = IOtools.CaseInsensitiveDict() for ikey in d: - if type(d[ikey]) == str and "," in d[ikey]: + if isinstance(d[ikey], str) and "," in d[ikey]: arr = d[ikey].split(",") for cont, i in enumerate(arr): try: @@ -51,13 +53,13 @@ def separateArrays(d): # Quotes for ikey in dnew: - if type(dnew[ikey]) == str and "'" in dnew[ikey]: + if isinstance(dnew[ikey], str) and "'" in dnew[ikey]: dnew[ikey] = dnew[ikey].replace("'", '"') return dnew -def cleanDifferences(d, tol_rel=0.0001): +def cleanDifferences(d, tol_rel=1e-7): d_new = {} for key in d: if key not in ["inputdir"]: @@ -96,32 +98,30 @@ def compareDictionaries(d1, d2): return different -def printTable(diff): +def printTable(diff, warning_percent=1e-1): print(f"{'':>15}{file1.split('/')[-1]:>25}{file2.split('/')[-1]:>25}") for key in diff: if diff[key][0] is not None: if diff[key][1] is not None: if diff[key][0] != 0.0: try: - perc = round( - 100 * np.abs((diff[key][0] - diff[key][1]) / diff[key][0]), - 2, + perc = 100 * np.abs( + (diff[key][0] - diff[key][1]) / diff[key][0] ) except: perc = np.nan else: perc = np.nan print( - "{0:>15}{1:>25}{2:>25} ({3}%)".format( - key, str(diff[key][0]), str(diff[key][1]), str(perc) - ) + f"{key:>15}{str(diff[key][0]):>25}{str(diff[key][1]):>25} (~{perc:.0e}%)", + typeMsg="w" if perc > warning_percent else "", ) else: print(f"{key:>15}{str(diff[key][0]):>25}{'':>25}") else: print(f"{key:>15}{'':>25}{str(diff[key][1]):>25}") print( - "--------------------------------------------------------------------------------------------------" + "--------------------------------------------------------------------------------" ) diff --git a/src/mitim_tools/transp_tools/src/TRANSPglobus.py b/src/mitim_tools/transp_tools/src/TRANSPglobus.py index 5ee2c1ce..018f7ca8 100644 --- a/src/mitim_tools/transp_tools/src/TRANSPglobus.py +++ b/src/mitim_tools/transp_tools/src/TRANSPglobus.py @@ -35,8 +35,8 @@ class TRANSPglobus(TRANSPmain.TRANSPgeneric): def __init__(self, FolderTRANSP, tokamak): super().__init__(FolderTRANSP, tokamak) - def defineRunParameters(self, *args, **kargs): - super().defineRunParameters(*args, **kargs) + def defineRunParameters(self, *args, **kwargs): + super().defineRunParameters(*args, **kwargs) s = CONFIGread.load_settings() self.pppluser = s["globus"]["username"] @@ -52,7 +52,7 @@ def defineRunParameters(self, *args, **kargs): ------------------------------------------------------------------------------------------------------ """ - def run(self, version="pshare", **kargs): + def run(self, version="pshare", **kwargs): # Make sure that the MPIs are set up properly self.mpisettings = TRANSPmain.ensureMPIcompatibility( self.nml_file, self.nml_file_ptsolver, self.mpisettings @@ -64,7 +64,7 @@ def run(self, version="pshare", **kargs): # tr_send self.send() - def check(self, permitSomeTime=False, NBImissing_isStopped=True, **kargs): + def check(self, permitSomeTime=False, NBImissing_isStopped=True, **kwargs): permitSomeTime = False # False has been HARD CODED for tests # ------------------------------------------------------------------------------------------------------------------------------ @@ -109,7 +109,7 @@ def get( retrieveAC=False, fullRequest=True, checkForActive=True, - **kargs, + **kwargs, ): """ This is for unfinished runs (tr_look), grab only the CDF for checks @@ -380,7 +380,7 @@ def get( self.cdfs[label] = Reactor - def fetch(self, label="run1", retrieveAC=False, **kargs): + def fetch(self, label="run1", retrieveAC=False, **kwargs): """ This is for finished runs (tr_fetch) """ @@ -398,7 +398,7 @@ def fetch(self, label="run1", retrieveAC=False, **kargs): return self.cdfs[label] - def delete(self, howManyCancel=1, MinWaitDeletion=0, **kargs): + def delete(self, howManyCancel=1, MinWaitDeletion=0, **kwargs): TRANSPglobus_ntcc.tr_cancel( self.runid, self.FolderTRANSP, diff --git a/src/mitim_tools/transp_tools/src/TRANSPglobus_ntcc.py b/src/mitim_tools/transp_tools/src/TRANSPglobus_ntcc.py index e9bc448c..2ff11cd8 100644 --- a/src/mitim_tools/transp_tools/src/TRANSPglobus_ntcc.py +++ b/src/mitim_tools/transp_tools/src/TRANSPglobus_ntcc.py @@ -26,18 +26,19 @@ def tr_start( if FolderTRANSP[-1] != "/": FolderTRANSP += "/" - subfolder = IOtools.reducePathLevel(FolderTRANSP, level=1, isItFile=False)[ - 1 - ] #'/FolderTRANSP/' - # -------------------------------- # Where to run # -------------------------------- - machineSettings = CONFIGread.machineSettings( - code="ntcc", nameScratch=f"mitim_tmp_{runid}/" + transp_job = FARMINGtools.mitim_job(FolderTRANSP) + + transp_job.define_machine( + "ntcc", + f"mitim_{runid}/", + launchSlurm=False, ) - scratchArea = f"{machineSettings['folderWork']}/{subfolder}/" + + scratchArea = f"{transp_job.folderExecution}/" # ------------------------------------- # Prepare tr_start IDL script responses @@ -75,7 +76,6 @@ def tr_start( # Command # -------------------------------- - inputFiles = [] inputFolders = [FolderTRANSP] outputFiles = [ f"{runid}_{tok}_tmp.tar.gz", @@ -109,16 +109,14 @@ def tr_start( timeoutSecs = 120 # THis is needed, to avoid crazy high file sizes!! - FARMINGtools.runCommand( + transp_job.prep( Command, - inputFiles, - outputFiles=outputFiles, - inputFolders=inputFolders, - machineSettings=machineSettings, - whereOutput=FolderTRANSP, - timeoutSecs=timeoutSecs, + output_files=outputFiles, + input_folders=inputFolders, ) + transp_job.run(timeoutSecs=timeoutSecs) + # -------------------------------- # Checker # -------------------------------- @@ -136,22 +134,26 @@ def tr_dat(runid, tok, FolderTRANSP): typeMsg="i", ) - machineSettings = CONFIGread.machineSettings( - code="ntcc", nameScratch=f"mitim_tmp_{runid}/" - ) + transp_job = FARMINGtools.mitim_job(FolderTRANSP) subfolder = IOtools.reducePathLevel(FolderTRANSP, level=1, isItFile=False)[ 1 ] #'/FolderTRANSP/' - scratchArea = f"{machineSettings['folderWork']}/{subfolder}/" + + transp_job.define_machine( + "ntcc", + f"mitim_{runid}/", + launchSlurm=False, + ) + + scratchArea = f"{transp_job.folderExecution}" # -------------------------------- # Command # -------------------------------- - inputFiles = [] inputFolders = [FolderTRANSP] - outputFiles = [f"{subfolder}/outputtrdat.txt"] + outputFiles = ["outputtrdat.txt"] os.system(f"rm {FolderTRANSP}/outputtrdat.txt") Command = f"cd {scratchArea} && trdat {tok} {runid} Q >> outputtrdat.txt" @@ -162,16 +164,14 @@ def tr_dat(runid, tok, FolderTRANSP): timeoutSecs = 120 # THis is needed, to avoid crazy high file sizes!! - FARMINGtools.runCommand( + transp_job.prep( Command, - inputFiles, - outputFiles=outputFiles, - inputFolders=inputFolders, - machineSettings=machineSettings, - whereOutput=FolderTRANSP, - timeoutSecs=timeoutSecs, + output_files=outputFiles, + input_folders=inputFolders, ) + transp_job.run(timeoutSecs=timeoutSecs) + # -------------------------------- # Interpret # -------------------------------- @@ -179,7 +179,7 @@ def tr_dat(runid, tok, FolderTRANSP): NMLtools.interpret_trdat(f"{FolderTRANSP}/outputtrdat.txt") -def tr_send(FolderTRANSP, runid, tok, waitseconds=1 * 60): +def tr_send(FolderTRANSP, runid, tok): """ Kill look command after waitseconds """ @@ -191,10 +191,15 @@ def tr_send(FolderTRANSP, runid, tok, waitseconds=1 * 60): # Where to run # -------------------------------- - machineSettings = CONFIGread.machineSettings( - code="ntcc", nameScratch=f"mitim_tmp_{runid}/" + transp_job = FARMINGtools.mitim_job(FolderTRANSP) + + transp_job.define_machine( + "ntcc", + f"mitim_{runid}/", + launchSlurm=False, ) - scratchArea = machineSettings["folderWork"] + + scratchArea = transp_job.folderExecution # -------------------------------- # Command @@ -218,15 +223,14 @@ def tr_send(FolderTRANSP, runid, tok, waitseconds=1 * 60): # Run # -------------------------------- - FARMINGtools.runCommand( + transp_job.prep( Command, - inputFiles, - outputFiles=outputFiles, - machineSettings=machineSettings, - whereOutput=FolderTRANSP, - timeoutSecs=waitseconds, + output_files=outputFiles, + input_files=inputFiles, ) + transp_job.run(waitYN=False) # A send command gets out as soon as it's sent + def tr_look(FolderTRANSP, runid, tok, waitseconds=60): """ @@ -240,10 +244,15 @@ def tr_look(FolderTRANSP, runid, tok, waitseconds=60): # Where to run # -------------------------------- - machineSettings = CONFIGread.machineSettings( - code="ntcc", nameScratch=f"mitim_tmp_{runid}/" + transp_job = FARMINGtools.mitim_job(FolderTRANSP) + + transp_job.define_machine( + "ntcc", + f"mitim_{runid}/", + launchSlurm=False, ) - scratchArea = machineSettings["folderWork"] + + scratchArea = transp_job.folderExecution """ -------------------------------- @@ -264,14 +273,12 @@ def tr_look(FolderTRANSP, runid, tok, waitseconds=60): # Run # -------------------------------- - FARMINGtools.runCommand( + transp_job.prep( Command, - inputFiles, - outputFiles=outputFiles, - machineSettings=machineSettings, - whereOutput=FolderTRANSP, - timeoutSecs=waitseconds, + input_files=inputFiles, + output_files=outputFiles, ) + transp_job.run(timeoutSecs=waitseconds) def tr_get(file, server, runid, FolderTRANSP, tok, remove_previous_before=False): @@ -288,18 +295,21 @@ def tr_get(file, server, runid, FolderTRANSP, tok, remove_previous_before=False) # Where to run # -------------------------------- - machineSettings = CONFIGread.machineSettings( - code="ntcc", nameScratch=f"mitim_tmp_{runid}/" + transp_job = FARMINGtools.mitim_job(FolderTRANSP) + + transp_job.define_machine( + "ntcc", + f"mitim_{runid}/", + launchSlurm=False, ) - scratchArea = machineSettings["folderWork"] + + scratchArea = transp_job.folderExecution # -------------------------------- # Command # -------------------------------- - Command = "cd {0} && globus-url-copy -nodcau gsiftp://{1} file://{0}/{2} >> outputtrfetch.txt".format( - scratchArea, file_in_server, file - ) + Command = f"cd {scratchArea} && globus-url-copy -nodcau gsiftp://{file_in_server} file://{scratchArea}/{file} >> outputtrfetch.txt" inputFiles = [] outputFiles = [file] @@ -311,18 +321,18 @@ def tr_get(file, server, runid, FolderTRANSP, tok, remove_previous_before=False) # -------------------------------- if remove_previous_before: - error, result = FARMINGtools.run_subprocess( + ouptut, error = FARMINGtools.run_subprocess( f"cd {FolderTRANSP} && rm {file}", localRun=True ) - FARMINGtools.runCommand( + transp_job.prep( Command, - inputFiles, - outputFiles=outputFiles, - machineSettings=machineSettings, - whereOutput=FolderTRANSP, + output_files=outputFiles, + input_files=inputFiles, ) + transp_job.run() + def tr_cancel(runid, FolderTRANSP, tok, howManyCancel=1, MinWaitDeletion=2): if FolderTRANSP[-1] != "/": @@ -332,10 +342,15 @@ def tr_cancel(runid, FolderTRANSP, tok, howManyCancel=1, MinWaitDeletion=2): # Where to run # -------------------------------- - machineSettings = CONFIGread.machineSettings( - code="ntcc", nameScratch=f"mitim_tmp_{runid}/" + transp_job = FARMINGtools.mitim_job(FolderTRANSP) + + transp_job.define_machine( + "ntcc", + f"mitim_{runid}/", + launchSlurm=False, ) - scratchArea = machineSettings["folderWork"] + + scratchArea = transp_job.folderExecution # -------------------------------- # Command @@ -356,19 +371,16 @@ def tr_cancel(runid, FolderTRANSP, tok, howManyCancel=1, MinWaitDeletion=2): print(f">> Deleting {runid} from the grid") for k in range(howManyCancel): - FARMINGtools.runCommand( + transp_job.prep( Command, - inputFiles, - outputFiles=outputFiles, - timeoutSecs=5, - machineSettings=machineSettings, - whereOutput=FolderTRANSP, + output_files=outputFiles, + input_files=inputFiles, ) + transp_job.run(waitYN=False) + # Leave some more time for deletion print( - ">> Cancel request has been submitted, but waiting now {0}min to allow for deletion to happen properly".format( - MinWaitDeletion - ) + f">> Cancel request has been submitted, but waiting now {MinWaitDeletion}min to allow for deletion to happen properly" ) time.sleep(MinWaitDeletion * 60.0) diff --git a/src/mitim_tools/transp_tools/src/TRANSPmain.py b/src/mitim_tools/transp_tools/src/TRANSPmain.py index a74a70c5..99695678 100644 --- a/src/mitim_tools/transp_tools/src/TRANSPmain.py +++ b/src/mitim_tools/transp_tools/src/TRANSPmain.py @@ -61,19 +61,19 @@ def __init__(self, FolderTRANSP, tokamak): ------------------------------------------------------------------------------------------------------ """ - def run(self, *args, **kargs): + def run(self, *args, **kwargs): pass - def check(self, *args, **kargs): + def check(self, *args, **kwargs): pass - def get(self, *args, **kargs): + def get(self, *args, **kwargs): pass - def fetch(self, *args, **kargs): + def fetch(self, *args, **kwargs): pass - def delete(self, *args, **kargs): + def delete(self, *args, **kwargs): pass def automatic(*args, **kwargs): @@ -138,9 +138,7 @@ def convergence( datetime.datetime.now() + datetime.timedelta(minutes=minWait) ).strftime("%H:%M") print( - ">> Waiting for run {0} to converge... (check every {1}min, at {2})".format( - self.runid, minWait, tt - ) + f">> Waiting for run {self.runid} to converge... (check every {minWait}min, at {tt})" ) try: @@ -151,7 +149,6 @@ def convergence( time.sleep(60 * minWait) self.get(label="r1", retrieveAC=retrieveAC, checkForActive=checkForActive) - statusStop = self.statusStop # Try to read the hours it has been static (not progressed a single ms) try: @@ -163,11 +160,9 @@ def convergence( if hoursStatic > maxhoursStatic: print( - " >> {0} run has not progressed in {1:.1f}h (> {2:.1f}h), assume there is a problem with it".format( - self.runid, hoursStatic, maxhoursStatic - ) + f" >> {self.runid} run has not progressed in {hoursStatic:.1f}h (> {maxhoursStatic:.1f}h), assume there is a problem with it" ) - statusStop = 1 # 0 + self.statusStop = 1 # 0 except: pass @@ -176,9 +171,7 @@ def convergence( try: if lastTime > self.cdfs["r1"].t[-1]: print( - " >> {0} run has come back in time! Please check what happened by comparing CDF and CDF_prev".format( - self.runid - ), + f" >> {self.runid} run has come back in time! Please check what happened by comparing CDF and CDF_prev", typeMsg="q" if not automaticProcess else "qa", ) else: @@ -194,7 +187,7 @@ def convergence( # ~~~~~~~~~~~~~~ If the run is correct, evaluate if it has converged # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if statusStop == 0: + if self.statusStop == 0: ConvergedRun = self.cdfs["r1"].fullMITIMconvergence( minWait, timeDifference, @@ -205,11 +198,11 @@ def convergence( else: break - return ConvergedRun, statusStop + return ConvergedRun, self.statusStop def plot(self, label="run1", time=None): if self.cdfs[label] is not None: - self.cdfs[label].plotRun(time=time) + self.cdfs[label].plot(time=time) def checkUntilFinished( self, label="run1", checkMin=5, grabIntermediateEachMin=300.0, retrieveAC=False @@ -230,9 +223,7 @@ def checkUntilFinished( if first: print( - ">> Simulation just submitted, will check status in {0}min (at {1})".format( - checkMin, tt - ) + f">> Simulation just submitted, will check status in {checkMin}min (at {tt})" ) else: print( diff --git a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py index 614f6b12..95e16e5c 100644 --- a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py +++ b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py @@ -1,11 +1,13 @@ -import os, copy, datetime, time +import os +import copy +import datetime +import time import numpy as np from IPython import embed from mitim_tools.transp_tools.src import TRANSPmain from mitim_tools.misc_tools import IOtools, FARMINGtools from mitim_tools.misc_tools import CONFIGread from mitim_tools.transp_tools.tools import NMLtools -from mitim_tools.transp_tools import CDFtools from mitim_tools.misc_tools.IOtools import printMsg as print @@ -16,8 +18,8 @@ def __init__(self, FolderTRANSP, tokamak): self.job_id, self.job_name = None, None - def defineRunParameters(self, *args, **kargs): - super().defineRunParameters(*args, **kargs) + def defineRunParameters(self, *args, minutesAllocation=60 * 10, **kwargs): + super().defineRunParameters(*args, **kwargs) self.job_name = f"transp_{self.tok}_{self.runid}" @@ -27,35 +29,50 @@ def defineRunParameters(self, *args, **kargs): ) self.folderExecution = machineSettings["folderWork"] - """ - ------------------------------------------------------------------------------------------------------ - Main routines - ------------------------------------------------------------------------------------------------------ - """ - - def run(self, restartFromPrevious=False, minutesAllocation=60 * 10, **kargs): # Make sure that the MPIs are set up properly self.mpisettings = TRANSPmain.ensureMPIcompatibility( self.nml_file, self.nml_file_ptsolver, self.mpisettings ) - self.job_id = runSINGULARITY( - self.FolderTRANSP, + # --------------------------------------------------------------------------------------------------------------------------------------- + # Number of cores (must be inside 1 node) + # --------------------------------------------------------------------------------------------------------------------------------------- + nparallel = 1 + for j in self.mpisettings: + nparallel = int(np.max([nparallel, self.mpisettings[j]])) + if self.mpisettings[j] == 1: + self.mpisettings[j] = 0 # definition used for the transp-source + + self.job = FARMINGtools.mitim_job(self.FolderTRANSP) + + self.job.define_machine( + "transp", + f"transp_{self.tok}_{self.runid}/", + slurm_settings={ + "minutes": minutesAllocation, + "ntasks": nparallel, + "name": self.job_name, + }, + ) + + def run(self, restartFromPrevious=False, **kwargs): + runSINGULARITY( + self.job, self.runid, self.shotnumber, self.tok, self.mpisettings, - nameJob=self.job_name, - minutes=minutesAllocation, restartFromPrevious=restartFromPrevious, ) - def check(self, **kargs): - infoSLURM, self.job_id, self.log_file = getRunInformation( - self.job_id, self.FolderTRANSP, self.job_name, self.folderExecution - ) + self.jobid = self.job.jobid + + def check(self, **kwargs): + self.job.check() - info, status = interpretRun(infoSLURM, self.log_file) + info, status = interpretRun(self.job.infoSLURM, self.job.log_file) + + self.jobid = self.job.jobid_found return info, status, None @@ -63,15 +80,14 @@ def get( self, label="run1", retrieveAC=False, - fullRequest=True, minutesAllocation=60, - **kargs, + **kwargs, ): runSINGULARITY_look( self.FolderTRANSP, self.runid, self.tok, - self.folderExecution, + self.job_name, minutes=minutesAllocation, ) @@ -79,37 +95,61 @@ def get( self.FolderTRANSP, self.runid, retrieveAC=retrieveAC ) - def fetch(self, label="run1", retrieveAC=False, minutesAllocation=60, **kargs): + dictInfo, _, _ = self.check() + + # THIS NEEDS MORE WORK, LIKE IN GLOBUS # TO FIX + if dictInfo["info"]["status"] == "finished": + self.statusStop = -2 + else: + self.statusStop = 0 + + def fetch(self, label="run1", retrieveAC=False, minutesAllocation=60, **kwargs): runSINGULARITY_finish( - self.FolderTRANSP, self.runid, self.tok, minutes=minutesAllocation + self.FolderTRANSP, + self.runid, + self.tok, + self.job_name, + minutes=minutesAllocation, + ) + + # Get reactor to call for ACs as well + self.cdfs[label] = TRANSPmain.storeCDF( + self.FolderTRANSP, self.runid, retrieveAC=False ) # ------------------ # Organize AC files # ------------------ + if retrieveAC: - print("Checker AC, work on it") - embed() - ICRF, TORBEAM, NUBEAM = self.determineACs() + ICRF, TORBEAM, NUBEAM = self.determineACs(self.cdfs[label]) organizeACfiles( self.runid, self.FolderTRANSP, ICRF=ICRF, TORBEAM=TORBEAM, NUBEAM=NUBEAM ) - self.cdfs[label] = TRANSPmain.storeCDF( - self.FolderTRANSP, self.runid, retrieveAC=retrieveAC - ) + # Re-Read again + self.cdfs[label] = TRANSPmain.storeCDF( + self.FolderTRANSP, self.runid, retrieveAC=retrieveAC + ) return self.cdfs[label] - def delete(self, howManyCancel=1, MinWaitDeletion=0, **kargs): - machineSettings = CONFIGread.machineSettings( - code="transp", nameScratch=f"transp_{self.tok}_{self.runid}/" + def delete(self, howManyCancel=1, MinWaitDeletion=0, **kwargs): + transp_job = FARMINGtools.mitim_job(self.FolderTRANSP) + + transp_job.define_machine( + "transp", + self.job_name, + launchSlurm=False, + ) + + transp_job.prep( + f"scancel {self.job_id}", + label_log_files="_finish", ) - machineSettings["clear"] = True + for i in range(howManyCancel): - FARMINGtools.runCommand( - f"scancel {self.job_id}", [], machineSettings=machineSettings - ) + transp_job.run() time.sleep(MinWaitDeletion * 60.0) @@ -118,42 +158,39 @@ def automatic( convCriteria, minWait=60, timeStartPrediction=0, - FolderOutputs=None, phasetxt="", automaticProcess=False, retrieveAC=False, - **kargs, + **kwargs, ): # Launch run self.run(restartFromPrevious=False) - statusStop = -1 + self.statusStop = -1 # If run is not found on the grid (-1: not found, 0: running, 1: stopped, -2: success) - while statusStop == -1: + while self.statusStop == -1: # ~~~~~ Check status of run before sending look (to avoid problem with OUTTIMES) if retrieveAC: - dictInfo, _ = self.check() + dictInfo, _, _ = self.check() infoCheck = dictInfo["info"]["status"] while infoCheck != "finished": mins = 10 currentTime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") print( - " >>>>>>>>>>> {0}, run not finished yet, but wait for AC generation (wait {1} min for next check)".format( - currentTime, mins - ) + f" >>>>>>>>>>> {currentTime}, run not finished yet, but wait for AC generation (wait {mins} min for next check)" ) time.sleep(60.0 * mins) - dictInfo, _ = self.check() + dictInfo, _, _ = self.check() infoCheck = dictInfo["info"]["status"] - statusStop = -2 + self.statusStop = -2 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # ~~~~~ Standard convergence test # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ else: - ConvergedRun, statusStop = self.convergence( + ConvergedRun, self.statusStop = self.convergence( convCriteria, minWait=minWait, timeStartPrediction=timeStartPrediction, @@ -173,27 +210,23 @@ def automatic( # --------------------------------------------------------------------------- # If run has stopped - if statusStop == 1: + if self.statusStop == 1: print(f" >>>>>>>>>>> Run {self.runid} has STOPPED") HasItFailed = True # If run has finished running - elif statusStop == -2: + elif self.statusStop == -2: print( - " >>>>>>>>>>> Run {0} has finished in the grid, assume converged".format( - self.runid - ) + f" >>>>>>>>>>> Run {self.runid} has finished in the grid, assume converged" ) HasItFailed = False self.fetch(label="run1", retrieveAC=retrieveAC) # If run is for some reason stuck and does not admit looks, repeat process - elif statusStop == 10: + elif self.statusStop == 10: print( - " >>>>>>>>>>> Run {0} does not admit looks, removing and running loop again".format( - self.runid - ) + f" >>>>>>>>>>> Run {self.runid} does not admit looks, removing and running loop again" ) self.delete(howManyCancel=2, MinWaitDeletion=2) @@ -226,31 +259,17 @@ def automatic( def runSINGULARITY( - folderWork, + transp_job, runid, shotnumber, tok, mpis, - minutes=60, - nameJob="transp", restartFromPrevious=False, ): - machineSettings = CONFIGread.machineSettings( - code="transp", nameScratch=f"transp_{tok}_{runid}/" - ) + folderWork = transp_job.folder_local + nparallel = transp_job.slurm_settings["ntasks"] - folderExecution = machineSettings["folderWork"] - - # --------------------------------------------------------------------------------------------------------------------------------------- - # Number of cores (must be inside 1 node) - # --------------------------------------------------------------------------------------------------------------------------------------- - nparallel = 1 - for j in mpis: - nparallel = int(np.max([nparallel, mpis[j]])) - if mpis[j] == 1: - mpis[j] = 0 # definition used for the transp-source - - NMLtools.adaptNML(folderWork, runid, shotnumber, folderExecution) + NMLtools.adaptNML(folderWork, runid, shotnumber, transp_job.folderExecution) # ---------------------------------------------------------------------------------------------------------------------------------------- # Common things @@ -258,10 +277,10 @@ def runSINGULARITY( inputFolders, inputFiles, shellPreCommands = [], [], [] - if "nobackup1" in folderExecution: - txt_bind = ( - "--bind /nobackup1 " # As Jai suggestion to solve problem with nobackup1 - ) + start_folder = transp_job.folderExecution.split("/")[1] # e.g. pool001, nobackup1 + + if start_folder not in ["home", "Users"]: + txt_bind = f"--bind /{start_folder} " # As Jai suggestion to solve problem with nobackup1 else: txt_bind = "" @@ -298,13 +317,7 @@ def runSINGULARITY( inputFiles.append(file) with open(file, "w") as f: f.write( - "{0}/results/\n{0}/data/\n{0}/tmp/\ny\ny\n{1}\n{2}\n{3}\n{4}\n0\n0".format( - folderExecution, - nparallel, - mpis["trmpi"], - mpis["toricmpi"], - mpis["ptrmpi"], - ) + f"{transp_job.folderExecution}/results/\n{transp_job.folderExecution}/data/\n{transp_job.folderExecution}/tmp/\ny\ny\n{nparallel}\n{mpis['trmpi']}\n{mpis['toricmpi']}\n{mpis['ptrmpi']}\n0\n0" ) # Direct env (dirty fix until Jai fixes the env app) @@ -312,12 +325,12 @@ def runSINGULARITY( inputFiles.append(file) ENVcommand = f""" export WORKDIR=./ -export RESULTDIR={folderExecution}/results/ -mkdir -p {folderExecution}/results/ -export DATADIR={folderExecution}/data/ -mkdir -p {folderExecution}/data/ -export TMPDIR_TR={folderExecution}/tmp/ -mkdir -p {folderExecution}/tmp/ +export RESULTDIR={transp_job.folderExecution}/results/ +mkdir -p {transp_job.folderExecution}/results/ +export DATADIR={transp_job.folderExecution}/data/ +mkdir -p {transp_job.folderExecution}/data/ +export TMPDIR_TR={transp_job.folderExecution}/tmp/ +mkdir -p {transp_job.folderExecution}/tmp/ export NPROCS={nparallel} export NBI_NPROCS={mpis["trmpi"]} export NTOR_NPROCS={mpis["toricmpi"]} @@ -340,17 +353,17 @@ def runSINGULARITY( # --------------- TRANSPcommand_prep = f""" -#singularity run --app environ $TRANSP_SINGULARITY < {folderExecution}/env_prf -singularity run {txt_bind}--app pretr $TRANSP_SINGULARITY {tok}{txt} {runid} < {folderExecution}/pre_prf +#singularity run --app environ $TRANSP_SINGULARITY < {transp_job.folderExecution}/env_prf +singularity run {txt_bind}--app pretr $TRANSP_SINGULARITY {tok}{txt} {runid} < {transp_job.folderExecution}/pre_prf singularity run {txt_bind}--app trdat $TRANSP_SINGULARITY {tok} {runid} w q |& tee {runid}tr_dat.log """ TRANSPcommand = f""" -#singularity run --app environ $TRANSP_SINGULARITY < {folderExecution}/env_prf -singularity run {txt_bind}--app pretr $TRANSP_SINGULARITY {tok}{txt} {runid} < {folderExecution}/pre_prf +#singularity run --app environ $TRANSP_SINGULARITY < {transp_job.folderExecution}/env_prf +singularity run {txt_bind}--app pretr $TRANSP_SINGULARITY {tok}{txt} {runid} < {transp_job.folderExecution}/pre_prf singularity run {txt_bind}--app trdat $TRANSP_SINGULARITY {tok} {runid} w q |& tee {runid}tr_dat.log singularity run {txt_bind}--app link $TRANSP_SINGULARITY {runid} -singularity run {txt_bind}--cleanenv --app transp $TRANSP_SINGULARITY {runid} |& tee {folderExecution}/{runid}tr.log +singularity run {txt_bind}--cleanenv --app transp $TRANSP_SINGULARITY {runid} |& tee {transp_job.folderExecution}/{runid}tr.log """ # ********** Start from previous @@ -360,11 +373,9 @@ def runSINGULARITY( TRANSPcommand_prep = None - TRANSPcommand = """ -singularity run {4}--cleanenv --app transp $TRANSP_SINGULARITY {3} R |& tee {0}/{3}tr.log -""".format( - folderExecution, tok, txt, runid, txt_bind - ) + TRANSPcommand = f""" +singularity run {txt_bind}--cleanenv --app transp $TRANSP_SINGULARITY {runid} R |& tee {transp_job.folderExecution}/{runid}tr.log +""" # ------------------ # Execute pre-checks @@ -374,78 +385,47 @@ def runSINGULARITY( os.system(f"rm {folderWork}/{runid}tr_dat.log") # Run first the prep (with tr_dat) - os.system(f"rm {folderWork}/tmp_inputs/bash.src") - os.system(f"rm {folderWork}/tmp_inputs/mitim.sh") + os.system(f"rm {folderWork}/tmp_inputs/mitim_bash.src") + os.system(f"rm {folderWork}/tmp_inputs/mitim_shell_executor.sh") - jobid = FARMINGtools.SLURMcomplete( + transp_job.prep( TRANSPcommand_prep, - folderWork, - inputFiles, - [f"{runid}tr_dat.log"], - 5, - nparallel, - nameJob, - machineSettings, + input_files=inputFiles, + input_folders=inputFolders, + output_files=[f"{runid}tr_dat.log"], shellPreCommands=shellPreCommands, - inputFolders=inputFolders, - waitYN=True, ) + # tr_dat doesn't need slurm + lS = copy.deepcopy(transp_job.launchSlurm) + transp_job.launchSlurm = False + + transp_job.run() + + transp_job.launchSlurm = lS # Back to original + # Interpret NMLtools.interpret_trdat(f"{folderWork}/{runid}tr_dat.log") # inputFiles = inputFiles[:-2] # Because in SLURMcomplete they are added - os.system(f"rm {folderWork}/tmp_inputs/bash.src") - os.system(f"rm {folderWork}/tmp_inputs/mitim.sh") + os.system(f"rm {folderWork}/tmp_inputs/mitim_bash.src") + os.system(f"rm {folderWork}/tmp_inputs/mitim_shell_executor.sh") # --------------- # Execute Full # --------------- - jobid = FARMINGtools.SLURMcomplete( + transp_job.prep( TRANSPcommand, - folderWork, - inputFiles, - [], - minutes, - nparallel, - nameJob, - machineSettings, + input_files=inputFiles, + input_folders=inputFolders, shellPreCommands=shellPreCommands, - inputFolders=inputFolders, - waitYN=False, ) - os.system(f"rm -r {folderWork}/tmp_inputs") - - return jobid - - -def getRunInformation(jobid, folder, job_name, folderExecution): - print('* Submitting a "check" request to the cluster', typeMsg="i") + transp_job.run(waitYN=False) - machineSettings = CONFIGread.machineSettings( - code="transp", nameScratch=f"transp_{jobid}_check/" - ) - - # Grab slurm state and log file from TRANSP - infoSLURM = FARMINGtools.getSLURMstatus( - folder, - machineSettings, - jobid=jobid, - grablog=f"{folderExecution}/slurm_output.dat", - name=job_name, - ) - - with open(folder + "/slurm_output.dat", "r") as f: - log_file = f.readlines() - - # If jobid was given as None, I retrieved the info from the job_name, but now provide here the actual id - if infoSLURM is not None: - jobid = infoSLURM["JOBID"] - - return infoSLURM, jobid, log_file + os.system(f"rm -r {folderWork}/tmp_inputs") def interpretRun(infoSLURM, log_file): @@ -480,10 +460,9 @@ def interpretRun(infoSLURM, log_file): info["info"]["status"] = "finished" else: print( - "Not identified status... assume for the moment that it has finished?", + "Not identified status... assume for the moment that it has finished", typeMsg="w", ) - embed() pringLogTail(log_file) status = 1 info["info"]["status"] = "finished" @@ -508,28 +487,31 @@ def pringLogTail(log_file, howmanylines=50): print(txt, typeMsg="w") -def runSINGULARITY_finish(folderWork, runid, tok, minutes=60): - nameJob = f"transp_{tok}_{runid}_finish" +def runSINGULARITY_finish(folderWork, runid, tok, job_name, minutes=60): + transp_job = FARMINGtools.mitim_job(folderWork) - machineSettings = CONFIGread.machineSettings( - code="transp_look", nameScratch=f"transp_{tok}_{runid}/" + transp_job.define_machine( + "transp", + job_name, + launchSlurm=False, ) # --------------- # Execution command # --------------- - if "nobackup1" in machineSettings["folderWork"]: - txt_bind = ( - "--bind /nobackup1 " # As Jai suggestion to solve problem with nobackup1 - ) + start_folder = transp_job.machineSettings["folderWork"].split("/")[ + 1 + ] # e.g. pool001, nobackup1 + + if start_folder not in ["home", "Users"]: + txt_bind = f"--bind /{start_folder} " # As Jai suggestion to solve problem with nobackup1 else: txt_bind = "" TRANSPcommand = f""" -cd {machineSettings['folderWork']} && singularity run {txt_bind}--app trlook $TRANSP_SINGULARITY {tok} {runid} -cd {machineSettings['folderWork']} && singularity run {txt_bind}--app finishup $TRANSP_SINGULARITY {runid} -cd {machineSettings['folderWork']} && tar -czvf TRANSPresults.tar results/{tok}.00 +cd {transp_job.machineSettings['folderWork']} && singularity run {txt_bind}--app trlook $TRANSP_SINGULARITY {tok} {runid} +cd {transp_job.machineSettings['folderWork']} && singularity run {txt_bind}--app finishup $TRANSP_SINGULARITY {runid} """ # --------------- @@ -538,46 +520,43 @@ def runSINGULARITY_finish(folderWork, runid, tok, minutes=60): print('* Submitting a "finish" request to the cluster', typeMsg="i") - machineSettings["clear"] = False - jobid = FARMINGtools.SLURMcomplete( + transp_job.prep( TRANSPcommand, - folderWork, - [], - ["TRANSPresults.tar"], - minutes, - 1, - nameJob, - machineSettings, - extranamelogs="_finish", + output_folders=["results/"], + label_log_files="_finish", ) - os.system( - f"cd {folderWork} && tar -xzvf TRANSPresults.tar && cp -r results/{tok}.00/* ." - ) - os.system(f"rm -r {folderWork}/TRANSPresults.tar {folderWork}/results/") + transp_job.run( + removeScratchFolders=False + ) # Because it needs to read what it was there from run() + + os.system(f"cd {folderWork}&& cp -r results/{tok}.00/* .") -def runSINGULARITY_look(folderWork, runid, tok, folderExecution, minutes=60): - nameJob = f"transp_{tok}_{runid}_look" +def runSINGULARITY_look(folderWork, runid, tok, job_name, minutes=60): + transp_job = FARMINGtools.mitim_job(folderWork) - machineSettings = CONFIGread.machineSettings( - code="transp_look", nameScratch=f"{nameJob}/" + transp_job.define_machine( + "transp", + job_name, + launchSlurm=False, ) # --------------- # Execution command # --------------- - if "nobackup1" in machineSettings["folderWork"]: - txt_bind = ( - "--bind /nobackup1 " # As Jai suggestion to solve problem with nobackup1 - ) + start_folder = transp_job.machineSettings["folderWork"].split("/")[ + 1 + ] # e.g. pool001, nobackup1 + + if start_folder not in ["home", "Users"]: + txt_bind = f"--bind /{start_folder} " # As Jai suggestion to solve problem with nobackup1 else: txt_bind = "" - # I have to do the look in another folder (I think it fails if I simply grab, so I copy things) TRANSPcommand = f""" -cd {machineSettings['folderWork']} && cp {folderExecution}/*PLN {folderExecution}/transp-bashrc {machineSettings['folderWork']}/. && singularity run {txt_bind}--app plotcon $TRANSP_SINGULARITY {runid} +cd {transp_job.folderExecution} && singularity run {txt_bind}--app plotcon $TRANSP_SINGULARITY {runid} """ # --------------- @@ -588,10 +567,15 @@ def runSINGULARITY_look(folderWork, runid, tok, folderExecution, minutes=60): outputFiles = [f"{runid}.CDF"] - jobid = FARMINGtools.SLURMcomplete( - TRANSPcommand, folderWork, [], outputFiles, minutes, 1, nameJob, machineSettings + transp_job.prep( + TRANSPcommand, + output_files=outputFiles, ) + transp_job.run( + removeScratchFolders=False + ) # Because it needs to read what it was there from run() + def organizeACfiles( runid, FolderTRANSP, nummax=1, ICRF=False, NUBEAM=False, TORBEAM=False @@ -604,22 +588,22 @@ def organizeACfiles( if NUBEAM: for i in range(nummax): name = runid + f".DATA{i + 1}" - os.system("mv {0}/{1} {0}/NUBEAM_folder".format(FolderTRANSP, name)) + os.system(f"mv {FolderTRANSP}/{name} {FolderTRANSP}/NUBEAM_folder") name = runid + f"_birth.cdf{i + 1}" - os.system("mv {0}/{1} {0}/NUBEAM_folder".format(FolderTRANSP, name)) + os.system(f"mv {FolderTRANSP}/{name} {FolderTRANSP}/NUBEAM_folder") if ICRF: for i in range(nummax): name = runid + f"_ICRF_TAR.GZ{i + 1}" - os.system("mv {0}/{1} {0}/TORIC_folder".format(FolderTRANSP, name)) + os.system(f"mv {FolderTRANSP}/{name} {FolderTRANSP}/TORIC_folder") name = runid + f"_FI_TAR.GZ{i + 1}" - os.system("mv {0}/{1} {0}/FI_folder".format(FolderTRANSP, name)) + os.system(f"mv {FolderTRANSP}/{name} {FolderTRANSP}/FI_folder") name = runid + "FPPRF.DATA" - os.system("mv {0}/{1} {0}/NUBEAM_folder".format(FolderTRANSP, name)) + os.system(f"mv {FolderTRANSP}/{name} {FolderTRANSP}/NUBEAM_folder") if TORBEAM: for i in range(nummax): name = runid + f"_TOR_TAR.GZ{i + 1}" - os.system("mv {0}/{1} {0}/TORBEAM_folder".format(FolderTRANSP, name)) + os.system(f"mv {FolderTRANSP}/{name} {FolderTRANSP}/TORBEAM_folder") diff --git a/src/mitim_tools/transp_tools/tools/FBMtools.py b/src/mitim_tools/transp_tools/tools/FBMtools.py index da11ab73..f4602079 100644 --- a/src/mitim_tools/transp_tools/tools/FBMtools.py +++ b/src/mitim_tools/transp_tools/tools/FBMtools.py @@ -315,34 +315,36 @@ def runGetFBM( MaxSeconds=60 * 1, commandOrder=None, ): - machineSettings = CONFIGread.machineSettings( - code="get_fbm", nameScratch=f"tmp_fbm_{name}/" - ) - folder, fileonly = IOtools.reducePathLevel(file, level=1) - print(f"\t\t\t- Running get_fbm command") + fbm_job = FARMINGtools.mitim_job(folder) + + fbm_job.define_machine( + "get_fbm", + f"tmp_fbm_{name}/", + launchSlurm=False, + ) + + print("\t\t\t- Running get_fbm command") if commandOrder is not None: - print(f"\t\t\t- [First running AC corrector]") - FARMINGtools.runCommand( - f"cd {machineSettings['folderWork']} && {commandOrder} && mv {fileonly} {fileonly}_converted", - [file], - outputFiles=[f"{fileonly}_converted"], - whereOutput=folder, - machineSettings=machineSettings, - timeoutSecs=MaxSeconds, + print("\t\t\t- [First running AC corrector]") + + fbm_job.prep( + f"{commandOrder} && mv {fileonly} {fileonly}_converted", + output_files=[f"{fileonly}_converted"], + input_files=[file], ) + fbm_job.run(timeoutSecs=MaxSeconds) + os.system(f"mv {file} {file}_original") os.system(f"mv {file}_converted {file}") - FARMINGtools.runCommand( - f"cd {machineSettings['folderWork']} && {commandMain}", - [file], - outputFiles=[finFile2], - whereOutput=folder, - machineSettings=machineSettings, - timeoutSecs=MaxSeconds, + fbm_job.prep( + commandMain, + output_files=[finFile2], + input_files=[file], ) + fbm_job.run(timeoutSecs=MaxSeconds) class birthCDF: diff --git a/src/mitim_tools/transp_tools/tools/NMLtools.py b/src/mitim_tools/transp_tools/tools/NMLtools.py index 44ba35c5..e078365c 100644 --- a/src/mitim_tools/transp_tools/tools/NMLtools.py +++ b/src/mitim_tools/transp_tools/tools/NMLtools.py @@ -1,7 +1,7 @@ import os -from mitim_tools.misc_tools import IOtools, FARMINGtools +import numpy as np +from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools.aux import GACODEdefaults -from mitim_tools.misc_tools import CONFIGread from mitim_tools.misc_tools.IOtools import printMsg as print from IPython import embed @@ -76,8 +76,6 @@ def interpret_trdat(file): errors = [pos for pos, char in enumerate(file_plain) if char == "?"] - errors_clean = [] - truly_error = False for cont, i in enumerate(errors): if (i == 0 or errors[cont - 1] < i - 1) and file_plain[ @@ -86,9 +84,19 @@ def interpret_trdat(file): truly_error = True if truly_error: - print('\t- Detected "?" in TRDAT output, printing tr_dat:', typeMsg="w") - print("".join(aux), typeMsg="w") - print("Do you wish to continue? It will likely fail! (c)", typeMsg="q") + print( + '\t- Detected "?" in TRDAT output, printing tr_dat around errors:', + typeMsg="w", + ) + print("-------------------------------", typeMsg="w") + for i in range(len(aux)): + if ("?" in aux[i]) and ("????" not in aux[i]): + print("".join(aux[np.max(i - 5, 0) : i + 2]), typeMsg="w") + print("-------------------------------", typeMsg="w") + if not print( + "Do you wish to continue? It will likely fail! (c)", typeMsg="q" + ): + embed() else: print("\t- TRDAT output did not show any error", typeMsg="i") diff --git a/src/mitim_tools/transp_tools/tools/toric_tools_JCW.py b/src/mitim_tools/transp_tools/tools/toric_tools_JCW.py index 237819e1..bf8d88bc 100644 --- a/src/mitim_tools/transp_tools/tools/toric_tools_JCW.py +++ b/src/mitim_tools/transp_tools/tools/toric_tools_JCW.py @@ -638,7 +638,7 @@ def plot_2Dfield( # ax.text(-0.2,-0.2,sublabel,transform = ax.transAxes) # print ("interactive off while plotting") - # plt.ioff() + # if logl > 0: title = "log10 " + title diff --git a/templates/input.tglf.controls b/templates/input.tglf.controls new file mode 100644 index 00000000..cc8d84b8 --- /dev/null +++ b/templates/input.tglf.controls @@ -0,0 +1,71 @@ +#------------------------------------------------------------------------- +# Template input.tglf file (controls-only) +#------------------------------------------------------------------------- + +# ~~~ Main +USE_TRANSPORT_MODEL = True +WRITE_WAVEFUNCTION_FLAG = 0 +IFLUX = True # Compute quasilinear weights and mode amplitudes +GEOMETRY_FLAG = 1 # 0: s-a, 1: Miller, 2: Fourier, 3: ELITE +SIGN_IT = 1.0 # Sign of IT with repsect to CCW toroidal direction from top +SIGN_BT = 1.0 # Sign of BT with repsect to CCW toroidal direction from top +NN_MAX_ERROR = -1.0 # Threshold for TGLF-NN execution versus full TGLF calculation + +# ~~~ Saturation model +SAT_RULE = 0 +UNITS = GYRO +ETG_FACTOR = 1.25 + +# ~~~ Physics included +ADIABATIC_ELEC = False # T: Use adiabatic electrons +THETA_TRAPPED = 0.7 # Parameter to adjust trapped fraction model +USE_BPER = False # Include transverse magnetic fluctuations, 𝛿𝐴‖ -> PERPENDICULAR FLUCTUATIONS +USE_BPAR = False # Include compressional magnetic fluctuations, \delta B_{\lVert }} -> COMPRESSIONAL EFFECTS +USE_MHD_RULE = False # Ignore pressure gradient contribution to curvature drift +XNU_MODEL = 2 # Collision model (2=new) +VPAR_MODEL = 0 # 0=low-Mach-number limit (DEPRECATED?) +VPAR_SHEAR_MODEL = 1 +USE_AVE_ION_GRID = False # T:Use weighted average charge of ions for the gyroradius reference, F: Use first ion +USE_INBOARD_DETRAPPED = False # Set trapped fraction to zero if eigenmode is inward ballooning. +RLNP_CUTOFF = 18.0 # Limits SAT2 factor of R/Lp < RLNP_CUTOFF +NEW_EIKONAL = True + +# ~~~ Wavenumber grid +KYGRID_MODEL = 1 # 1: Standard grid +NKY = 19 # Number of poloidal modes in the high-k spectrum +KY = 0.3 # ky for single-mode call to TGLF + +# ~~~ Hermite basis functions +FIND_WIDTH = True # T: Find the width that maximizes the growth rate, F: Use WIDTH +USE_BISECTION = True # T: Use bisection search method to find width that maximizes growth rate +WIDTH = 1.65 # Max. width to search for Hermite polynomial basis (1.65 default, <1.5 recommended in Staebler PoP05) +WIDTH_MIN = 0.3 # Min. width to search for Hermite polynomial basis (0.3 default) +NWIDTH = 21 # Number of search points +NBASIS_MIN = 2 # Minimum number of parallel basis functions +NBASIS_MAX = 4 # Maximum number of parallel basis functions +NXGRID = 16 # Number of nodes in Gauss-Hermite quadrature + +# ~~~ Modes options +IBRANCH = -1 # 0 = find two most unstable modes one for each sign of frequency, electron drift direction (1), ion drift direction (2), -1 = sort the unstable modes by growthrate in rank order +NMODES = 5 # For IBRANCH=-1, number of modes to store. ASTRA uses NS+2 + +# ~~~ ExB shear model +ALPHA_QUENCH = 0.0 # 1.0 = use quench rule, 0.0 = use new spectral shift model (0.0 recommeded by Gary, 05/11/2020) +ALPHA_E = 1.0 # Multiplies ExB velocity shear for spectral shift model (1.0 ecommeded by Gary, 05/11/2020) +ALPHA_MACH = 0.0 +ALPHA_ZF = 1.0 + +# ~~~ Multipliers +DEBYE_FACTOR = 1.0 # Multiplies the debye length +XNU_FACTOR = 1.0 # Multiplies the trapped/passing boundary electron-ion collision terms +ALPHA_P = 1.0 # Multiplies parallel velocity shear for all species +PARK = 1.0 # Multiplies the parallel gradient term +GHAT = 1.0 # Multiplies the curvature drift closure terms +GCHAT = 1.0 # Multiplies the curvature drift irreducible terms +DAMP_PSI = 0.0 +DAMP_SIG = 0.0 +LINSKER_FACTOR = 0.0 +GRADB_FACTOR = 0.0 +WD_ZERO = 0.1 # Cutoff for curvature drift eigenvalues to prevent zero +FILTER = 2.0 # Sets threshold for frequency/drift frequency to filter out non-driftwave instabilities +WDIA_TRAPPED = 0.0 diff --git a/templates/input.tglf.models.json b/templates/input.tglf.models.json new file mode 100644 index 00000000..d4ae2591 --- /dev/null +++ b/templates/input.tglf.models.json @@ -0,0 +1,64 @@ +{ + "1": { + "label": "SAT1", + "controls": { + "SAT_RULE": 1, + "UNITS": "GYRO" + } + }, + "2": { + "label": "SAT0", + "controls": { + "SAT_RULE": 0, + "UNITS": "GYRO", + "ETG_FACTOR": 1.25 + } + }, + "3": { + "label": "SAT1geo", + "controls": { + "SAT_RULE": 1, + "UNITS": "CGYRO" + } + }, + "4": { + "label": "SAT2", + "controls": { + "SAT_RULE": 2, + "UNITS": "CGYRO", + "XNU_MODEL": 3, + "WDIA_TRAPPED": 1.0 + } + }, + "5": { + "label": "SAT2em", + "controls": { + "SAT_RULE": 2, + "UNITS": "CGYRO", + "XNU_MODEL": 3, + "WDIA_TRAPPED": 1.0, + "USE_BPER": true + } + }, + "6": { + "label": "SAT3", + "controls": { + "SAT_RULE": 3, + "UNITS": "CGYRO", + "XNU_MODEL": 3, + "WDIA_TRAPPED": 1.0 + } + }, + "101": { + "label": "[Experimentation] SAT3em basis", + "controls": { + "SAT_RULE": 2, + "UNITS": "CGYRO", + "XNU_MODEL": 3, + "WDIA_TRAPPED": 1.0, + "USE_BPER": true, + "KYGRID_MODEL": 4, + "NBASIS_MAX": 6 + } + } +} \ No newline at end of file diff --git a/tests/TGLF_workflow.py b/tests/TGLF_workflow.py index f7a981fe..4e1329ba 100644 --- a/tests/TGLF_workflow.py +++ b/tests/TGLF_workflow.py @@ -25,4 +25,4 @@ ) tglf.read(label="run1") -tglf.plotRun(labels=["run1"]) +tglf.plot(labels=["run1"]) diff --git a/tests/TGLFfull_workflow.py b/tests/TGLFfull_workflow.py index 7aa7980c..0e9d2317 100644 --- a/tests/TGLFfull_workflow.py +++ b/tests/TGLFfull_workflow.py @@ -24,7 +24,7 @@ tglf.run( subFolderTGLF="runBase/", TGLFsettings=5, - runWaveForms=[0.3], + runWaveForms=[0.1,0.3], restart=restart, forceIfRestart=True, ) @@ -33,7 +33,7 @@ tglf.run( subFolderTGLF="runSAT0/", TGLFsettings=2, - runWaveForms=[0.3], + runWaveForms=[0.1,0.3], restart=restart, forceIfRestart=True, ) @@ -45,4 +45,4 @@ tglf.NormalizationSets["EXP"]["exp_Qe_error"] = 0.005 tglf.NormalizationSets["EXP"]["exp_Qi_error"] = 0.005 -tglf.plotRun(labels=["runBase", "runSAT0"]) +tglf.plot(labels=["runBase", "runSAT0"]) diff --git a/tests/TGYRO_workflow.py b/tests/TGYRO_workflow.py index fcf2697d..5c590e90 100644 --- a/tests/TGYRO_workflow.py +++ b/tests/TGYRO_workflow.py @@ -35,7 +35,7 @@ tgyro.run( subFolderTGYRO="run1/", - iterations=5, + iterations=3, restart=True, forceIfRestart=True, special_radii=rhos, @@ -45,4 +45,4 @@ TGYRO_physics_options=physics_options, ) tgyro.read(label="run1") -tgyro.plotRun(labels=["run1"]) +tgyro.plot(labels=["run1"]) diff --git a/tests/TRANSP_workflow.py b/tests/TRANSP_workflow.py index b56ff2ad..9d3dc47a 100644 --- a/tests/TRANSP_workflow.py +++ b/tests/TRANSP_workflow.py @@ -12,7 +12,7 @@ In engaging, with 32 cores, should take ~1h20min """ -import os, sys, random +import os from mitim_tools.misc_tools import IOtools from mitim_tools.transp_tools import TRANSPtools from mitim_tools.misc_tools import CONFIGread @@ -39,17 +39,7 @@ # Randomize number of regression, not to collide with other people launching the same run, but there's a chance anyway... s = CONFIGread.load_settings() -runid = s["globus"]["username"][0].upper() + str(random.randint(1, 99)).zfill( - 2 -) # e.g. 'S01' - -# Define TRANSP class and where it is run -t = TRANSPtools.TRANSP(folder, "CMOD") - -# Define user and run parameters -t.defineRunParameters( - "12345" + runid, "12345", mpisettings={"trmpi": 32, "toricmpi": 32, "ptrmpi": 32} -) +runid = 'P24' # ---- Prepare NML and UFILES if os.path.exists(folder): @@ -58,16 +48,24 @@ os.system(f"mv {folder}/12345X01TR.DAT {folder}/12345{runid}TR.DAT") # --------------------------- + +# Define TRANSP class and where it is run +t = TRANSPtools.TRANSP(folder, "CMOD") + +# Define user and run parameters +t.defineRunParameters( + "12345" + runid, "12345", + mpisettings={"trmpi": 32, "toricmpi": 32, "ptrmpi": 32}, + minutesAllocation = 10 +) + # Submit run t.run() # Check c = t.checkUntilFinished( - label="run1", checkMin=5, grabIntermediateEachMin=20, retrieveAC=True + label="run1", checkMin=2, grabIntermediateEachMin=20, retrieveAC=True ) -# Write outputs of TRANSP (optional) -c.writeOutput(time=1e4, avTime=0.001) - # Plot t.plot(label="run1") diff --git a/tests/data/FolderTRANSP/12345X01TR.DAT b/tests/data/FolderTRANSP/12345X01TR.DAT index 3cff7977..c0e16f0b 100644 --- a/tests/data/FolderTRANSP/12345X01TR.DAT +++ b/tests/data/FolderTRANSP/12345X01TR.DAT @@ -314,7 +314,7 @@ c_sawtooth(29) = 0.4 ! chfast in Eq 13 (default porcelli = 0.4) !----- General Model -nalpha = 1 ! Fusion products model (0=MC for alphas, 1= fast model) +nalpha = 0 ! Fusion products model (0=MC for alphas, 1= fast model) nptclf = 10000 ! Number of monte carlo particles for fusion product calcs nlfatom = T ! Include atomic physics effects on products (e.g. CX) nl_ignore_mini_conflicts = T ! Ignore issue with He3 product coinciden with ICRF minority @@ -324,7 +324,7 @@ nl_ignore_mini_conflicts = T ! Ignore issue with He3 product coinciden with IC nlfhe4 = False ! Turn on MC slowing-down of He4 from D+T reactions plfhe4 = 1e2 ! Source power threshold to run MC (in W) -nlfst = F ! Turn on MC slowing-down of T from D+D reactions (D+D=T+p) +nlfst = True ! Turn on MC slowing-down of T from D+D reactions (D+D=T+p) plfst = 1e0 ! Source power threshold to run MC NLFST nlfsp = F ! Turn on MC slowing-down of protons from D+D reactions (D+D=T+p) @@ -532,7 +532,7 @@ nribol = -4 ! Added by PORTALS nrivp2 = -4 ! Added by PORTALS extsaw = 'SAW' ! Added by PORTALS presaw = 'PRF' ! Added by PORTALS -inputdir='/home/pablorf/scratch//mitim_tmp_12345X01//FolderTRANSP/' +inputdir='/home/pablorf/scratch//mitim_12345X01//FolderTRANSP/' MTHDAVG = 2 ! Added by MITIM AVGTIM = 0.001 ! Added by MITIM OUTTIM = 0.802 ! Added by MITIM diff --git a/tutorials/TGLF_tutorial.py b/tutorials/TGLF_tutorial.py index 4186acf9..5805bb01 100644 --- a/tutorials/TGLF_tutorial.py +++ b/tutorials/TGLF_tutorial.py @@ -1,3 +1,4 @@ +import numpy as np from mitim_tools.gacode_tools import TGLFtools from mitim_tools.misc_tools import IOtools @@ -11,8 +12,19 @@ # Prepare the TGLF class cdf = tglf.prep(folder, inputgacode=inputgacode_file, restart=False) +''' +*************************************************************************** +Run standandalone TGLF +*************************************************************************** +''' + # Run TGLF in subfolder -tglf.run(subFolderTGLF="yes_em_folder/", TGLFsettings=5, extraOptions={}, restart=True) +tglf.run( + subFolderTGLF="yes_em_folder/", + TGLFsettings=5, + extraOptions={}, + restart=False + ) # Read results of previous run and store them in results dictionary tglf.read(label="yes_em") @@ -22,11 +34,79 @@ subFolderTGLF="no_em_folder/", TGLFsettings=5, extraOptions={"USE_BPER": False}, - restart=True, + restart=False, ) # Read results of previous run and store them in results dictionary tglf.read(label="no_em") # Plot the two cases together -tglf.plotRun(labels=["yes_em", "no_em"]) +tglf.plot(labels=["yes_em", "no_em"]) + +''' +*************************************************************************** +Run TGLF scan +*************************************************************************** +''' + +tglf.runScan( subFolderTGLF = 'scan1/', + TGLFsettings = 5, + restart = False, + variable = 'RLTS_1', + varUpDown = np.linspace(0.5,1.5,3)) +tglf.readScan(label='scan1',variable = 'RLTS_1') + + +tglf.runScan( subFolderTGLF = 'scan2/', + TGLFsettings = 5, + restart = False, + variable = 'RLTS_2', + varUpDown = np.linspace(0.5,1.5,3)) +tglf.readScan(label='scan2',variable = 'RLTS_2') + + +tglf.plotScan(labels=['scan1','scan2']) + +''' +*************************************************************************** +Automatic scan of turbulence drives +*************************************************************************** +''' + +tglf.runScanTurbulenceDrives( + subFolderTGLF = 'turb_drives/', + TGLFsettings = 5, + restart = False) + +tglf.plotScanTurbulenceDrives(label='turb_drives') + +''' +*************************************************************************** +Automatic scan of turbulence drives +*************************************************************************** +''' + +tglf.runAnalysis( + subFolderTGLF = 'chi_e/', + analysisType = 'chi_e', + TGLFsettings = 5, + restart = False, + label = 'chi_eu') + +tglf.plotAnalysis(labels=['chi_eu'],analysisType='chi_e') + +''' +*************************************************************************** +Explore all available MITIM settings for TGLF (with waveforms) +*************************************************************************** +''' + +for i in[1,2,3,4,5,101]: + tglf.run( + subFolderTGLF = f'settings{i}/', + runWaveForms = [0.67], + TGLFsettings = i, + restart = False) + tglf.read(label=f'settings{i}') + +tglf.plot(labels=[f'settings{i}' for i in range(1,6)])