From 56f417fbcfa8ad29d23e03bd6db92e06d45be0f8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 20 Apr 2025 11:23:42 -0400 Subject: [PATCH 001/385] By default, CGYRO file always includes Mt and GZ --- .../portals/utils/PORTALScgyro.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/src/mitim_modules/portals/utils/PORTALScgyro.py b/src/mitim_modules/portals/utils/PORTALScgyro.py index d2752041..46f4e09a 100644 --- a/src/mitim_modules/portals/utils/PORTALScgyro.py +++ b/src/mitim_modules/portals/utils/PORTALScgyro.py @@ -34,17 +34,15 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod trick_harcoded_f is the name of the file until the iteration number. E.g. 'example_run/Outputs/cgyro_results/iter_rmp_75_' e.g.: - includeMtAndGz_hardcoded, train_sep,start_num,last_one,trick_hardcoded_f = True, 1, 0,100, 'example_run/Outputs/cgyro_results/d3d_5chan_it_' + train_sep,start_num,last_one,trick_hardcoded_f = 1, 0,100, 'example_run/Outputs/cgyro_results/d3d_5chan_it_' """ - includeMtAndGz_hardcoded = PORTALSparameters["hardCodedCGYRO"]["includeMtAndGz_hardcoded"] train_sep = PORTALSparameters["hardCodedCGYRO"]["train_sep"] start_num = PORTALSparameters["hardCodedCGYRO"]["start_num"] last_one = PORTALSparameters["hardCodedCGYRO"]["last_one"] trick_hardcoded_f = PORTALSparameters["hardCodedCGYRO"]["trick_hardcoded_f"] else: - includeMtAndGz_hardcoded = None train_sep = None start_num = None last_one = None @@ -68,7 +66,7 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod OriginalFimp = PORTALSparameters["fImp_orig"] cgyroing_file = ( - lambda file_cgyro, numPORTALS_this=0, includeMtAndGz=False: cgyroing( + lambda file_cgyro, numPORTALS_this=0: cgyroing( FolderEvaluation, unmodified_profiles, numPORTALS, @@ -81,7 +79,6 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod evaluationsInFile=f"{numPORTALS_this}", impurityPosition=impurityPosition, file=file_cgyro, - includeMtAndGz=includeMtAndGz, ) ) print(f"\t- Suggested function call for mitim evaluation {numPORTALS} (lambda for cgyroing):",typeMsg="i") @@ -99,13 +96,11 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod cgyroing_file( f"{trick_hardcoded_f}{start_num}.txt", numPORTALS_this=numPORTALS, - includeMtAndGz=includeMtAndGz_hardcoded, ) else: cgyroing_file( f"{trick_hardcoded_f}{int(numPORTALS)-train_sep+1+start_num}.txt", numPORTALS_this=0, - includeMtAndGz=includeMtAndGz_hardcoded, ) @@ -122,7 +117,6 @@ def cgyroing( file=None, evaluationsInFile=0, impurityPosition=3, - includeMtAndGz=False, ): """ Variables need to have dimensions of (evaluation,rho) @@ -150,7 +144,7 @@ def cgyroing( PexchE, _, _, - ) = readCGYROresults(file, radii, includeMtAndGz=includeMtAndGz) + ) = readCGYROresults(file, radii) cont = 0 for i in evaluations: @@ -219,7 +213,7 @@ def wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro, ErrorRaised=0.005): ) -def readlineNTH(line, full_file=False, unnormalize=True): +def readlineNTH(line, full_file=True, unnormalize=True): s = line.split() i = 2 @@ -357,7 +351,7 @@ def readlineNTH(line, full_file=False, unnormalize=True): ) -def readCGYROresults(file, radii, includeMtAndGz=False, unnormalize=True): +def readCGYROresults(file, radii, unnormalize=True): """ Arrays are in (batch,radii) MW/m^2 and 1E20 @@ -428,7 +422,7 @@ def readCGYROresults(file, radii, includeMtAndGz=False, unnormalize=True): Pexch_std_read, tstart_read, tend_read, - ) = readlineNTH(lines[i], full_file=includeMtAndGz, unnormalize=unnormalize) + ) = readlineNTH(lines[i], unnormalize=unnormalize) # -------------------------------------------------------- # Radial location position From 5282c2cf2c596636327ee15c2d6684377a714ac8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 11:23:31 -0700 Subject: [PATCH 002/385] Better naming for surrogate selector --- src/mitim_modules/portals/PORTALSmain.py | 2 +- src/mitim_modules/portals/PORTALStools.py | 2 +- src/mitim_modules/portals/utils/PORTALSoptimization.py | 2 +- src/mitim_tools/opt_tools/STEPtools.py | 9 ++------- 4 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 6fb6db59..5256ebbd 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -64,7 +64,7 @@ def default_namelist(optimization_options, CGYROrun=False): optimization_options['acquisition_options']['relative_improvement_for_stopping'] = 1e-2 # Surrogate - optimization_options["surrogate_options"]["selectSurrogate"] = partial(PORTALStools.selectSurrogate, CGYROrun=CGYROrun) + optimization_options["surrogate_options"]["surrogate_selection"] = partial(PORTALStools.surrogate_selection_portals, CGYROrun=CGYROrun) if CGYROrun: # CGYRO runs should prioritize accuracy diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 310cac4e..190c2c34 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -8,7 +8,7 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -def selectSurrogate(output, surrogate_options, CGYROrun=False): +def surrogate_selection_portals(output, surrogate_options, CGYROrun=False): print(f'\t- Selecting surrogate options for "{output}" to be run') diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index dd9268e7..22eadce3 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -30,7 +30,7 @@ def initialization_simple_relax(self): MainFolder.mkdir(parents=True, exist_ok=True) a, b = IOtools.reducePathLevel(self.folderExecution, level=1) - namingConvention = f"portals_sr_{b}_ev" + namingConvention = "portals_sr_ev" if self.seed is not None and self.seed != 0: random.seed(self.seed) diff --git a/src/mitim_tools/opt_tools/STEPtools.py b/src/mitim_tools/opt_tools/STEPtools.py index a13553e8..41f86102 100644 --- a/src/mitim_tools/opt_tools/STEPtools.py +++ b/src/mitim_tools/opt_tools/STEPtools.py @@ -184,13 +184,8 @@ def fit_step(self, avoidPoints=None, fitWithTrainingDataIfContains=None): surrogate_options = copy.deepcopy(self.surrogate_options) # Then, depending on application (e.g. targets in mitim are fitted differently) - if ( - "selectSurrogate" in surrogate_options - and surrogate_options["selectSurrogate"] is not None - ): - surrogate_options = surrogate_options["selectSurrogate"]( - outi, surrogate_options - ) + if "surrogate_selection" in surrogate_options and surrogate_options["surrogate_selection"] is not None: + surrogate_options = surrogate_options["surrogate_selection"](outi, surrogate_options) # --------------------------------------------------------------------------------------------------- # To avoid problems with fixed values (e.g. calibration terms that are fixed) From 0b3d71a668074db9ae5cce5d09adc8882dfdb9d0 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 11:39:20 -0700 Subject: [PATCH 003/385] Changed naming of transforms. Allows backcompatibility for now --- src/mitim_modules/freegsu/utils/FREEGSUparams.py | 4 ++-- src/mitim_modules/portals/PORTALStools.py | 15 +++++++++++---- src/mitim_modules/portals/utils/PORTALSinit.py | 4 ++-- .../powertorch/utils/TRANSFORMtools.py | 9 +-------- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/mitim_modules/freegsu/utils/FREEGSUparams.py b/src/mitim_modules/freegsu/utils/FREEGSUparams.py index bd84cee9..f8c13f8a 100644 --- a/src/mitim_modules/freegsu/utils/FREEGSUparams.py +++ b/src/mitim_modules/freegsu/utils/FREEGSUparams.py @@ -108,7 +108,7 @@ def createProblemParameters( dvs_min.extend(dvs_min2) dvs_max.extend(dvs_max2) - transformation = produceNewInputs + transformation = input_transform_freegs else: transformation = None @@ -376,7 +376,7 @@ def extractCont(x, cont): return v -def produceNewInputs(X, output, bounds, ParamProfile): +def input_transform_freegs(X, output, bounds, ParamProfile): """ X will be a tensor (with or without gradients) batch*dim, unnormalized """ diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 190c2c34..b1749b6b 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -94,7 +94,7 @@ def default_portals_transformation_variables(additional_params = []): return portals_transformation_variables, portals_transformation_variables_trace -def produceNewInputs(Xorig, output, surrogate_parameters, surrogate_transformation_variables): +def input_transform_portals(Xorig, output, surrogate_parameters, surrogate_transformation_variables): """ - Xorig will be a tensor (batch1...N,dim) unnormalized (with or without gradients). @@ -146,13 +146,11 @@ def produceNewInputs(Xorig, output, surrogate_parameters, surrogate_transformati return xFit, parameters_combined - # ---------------------------------------------------------------------- # Transformation of Outputs # ---------------------------------------------------------------------- - -def transformPORTALS(X, surrogate_parameters, output): +def output_transform_portals(X, surrogate_parameters, output): """ 1. Make sure all batches are squeezed into a single dimension ------------------------------------------------------------------ @@ -406,3 +404,12 @@ def stopping_criteria_portals(mitim_bo, parameters = {}): else: print("\t- No convergence yet, providing as iteration values the scalarized objective") return False, yvals + +# TODO: Remove in the future, this is just to enable back compatibility with the old code +def selectSurrogate(*args,**kwargs): + return surrogate_selection_portals(*args,**kwargs) +def produceNewInputs(*args,**kwargs): + return input_transform_portals(*args,**kwargs) +def transformPORTALS(*args,**kwargs): + return output_transform_portals(*args,**kwargs) +# ----- diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 30f4772e..d82396f9 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -339,8 +339,8 @@ def initializeProblem( Variables[ikey] = prepportals_transformation_variables(portals_fun, ikey) portals_fun.surrogate_parameters = { - "transformationInputs": PORTALStools.produceNewInputs, - "transformationOutputs": PORTALStools.transformPORTALS, + "transformationInputs": PORTALStools.input_transform_portals, + "transformationOutputs": PORTALStools.output_transform_portals, "powerstate": portals_fun.powerstate, "applyImpurityGammaTrick": portals_fun.PORTALSparameters["applyImpurityGammaTrick"], "useFluxRatios": portals_fun.PORTALSparameters["useFluxRatios"], diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 405add08..84fef4f9 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -416,7 +416,6 @@ def parameterize_curve( x_coarse_tensor, parameterize_in_aLx=True, multiplier_quantity=1.0, - PreventNegative=False, ): """ Notes: @@ -450,9 +449,6 @@ def parameterize_curve( x_coarse = x_coarse_tensor[1:].cpu().numpy() - # Clip to zero if I want to prevent negative values - ygrad_coord = ygrad_coord.clip(0) if PreventNegative else ygrad_coord - """ Define region to get control points from ------------------------------------------------------------ @@ -504,10 +500,7 @@ def deparametrizer_coarse(x, y, multiplier=multiplier_quantity): I need to do in a finer grid so that it is consistent with TGYRO. x, y must be (batch, radii), y_bc must be (1) """ - return ( - x, - integrator_function(x, y, y_bc_real) / multiplier, - ) + return x, integrator_function(x, y, y_bc_real) / multiplier def deparametrizer_coarse_middle(x, y, multiplier=multiplier_quantity): """ From 45373f8c27dbd412b50f18ca98f3a66c0da4f240 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 12:07:37 -0700 Subject: [PATCH 004/385] Towards a generalization of input object in powerstate --- .../portals/utils/PORTALSinit.py | 2 +- src/mitim_modules/powertorch/STATEtools.py | 62 +++++-------------- .../powertorch/physics/TRANSPORTtools.py | 2 +- .../powertorch/scripts/calculateTargets.py | 2 +- .../powertorch/utils/POWERplot.py | 2 +- .../powertorch/utils/TRANSFORMtools.py | 52 +++++++++++++++- src/mitim_tools/popcon_tools/RAPIDStools.py | 2 +- 7 files changed, 69 insertions(+), 55 deletions(-) diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index d82396f9..423c18e1 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -203,7 +203,7 @@ def initializeProblem( ) # Write this updated profiles class (with parameterized profiles) - _ = portals_fun.powerstate.to_gacode( + _ = portals_fun.powerstate.from_powerstate( write_input_gacode=FolderInitialization / "input.gacode", postprocess_input_gacode=portals_fun.MODELparameters["applyCorrections"], ) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index dda5a86d..927cd43f 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -41,7 +41,7 @@ def __init__( ): ''' Inputs: - - profiles: PROFILES_GACODE object + - profiles: Object for PROFILES_GACODE or others - EvolutionOptions: - rhoPredicted: radial grid (MUST NOT CONTAIN ZERO, it will be added internally) - ProfilesPredicted: list of profiles to predict @@ -117,13 +117,20 @@ def _ensure_ne_before_nz(lst): self.labelsFM.append([f'aL{profile}', list(self.profile_map[profile])[0], list(self.profile_map[profile])[1]]) # ------------------------------------------------------------------------------------- - # input.gacode + # Object type (e.g. input.gacode) # ------------------------------------------------------------------------------------- - # Use a copy because I'm deriving, it may be expensive and I don't want to carry that out outside of this class - self.profiles = copy.deepcopy(profiles) - if "derived" not in self.profiles.__dict__: - self.profiles.deriveQuantities() + if isinstance(profiles, PROFILEStools.PROFILES_GACODE): + self.to_powerstate = TRANSFORMtools.gacode_to_powerstate + self.from_powerstate = TRANSFORMtools.to_gacode + + # Use a copy because I'm deriving, it may be expensive and I don't want to carry that out outside of this class + self.profiles = copy.deepcopy(profiles) + if "derived" not in self.profiles.__dict__: + self.profiles.deriveQuantities() + + else: + raise ValueError("[MITIM] The input profile object is not recognized, please use PROFILES_GACODE") # ------------------------------------------------------------------------------------- # Fine targets (need to do it here so that it's only once per definition of powerstate) @@ -138,12 +145,8 @@ def _ensure_ne_before_nz(lst): # Standard creation of plasma dictionary # ------------------------------------------------------------------------------------- - # Resolution of input.gacode - if increase_profile_resol: - TRANSFORMtools.improve_resolution_profiles(self.profiles, rho_vec) - # Convert to powerstate - TRANSFORMtools.gacode_to_powerstate(self, self.profiles, self.plasma["rho"]) + self.to_powerstate(self,increase_profile_resol=increase_profile_resol) # Convert into a batch so that always the quantities are (batch,dimX) self.batch_size = 0 @@ -195,44 +198,7 @@ def _fine_grid(self): # Revert plasma back self.plasma = plasma_copy - def to_gacode( - self, - write_input_gacode=None, - position_in_powerstate_batch=0, - postprocess_input_gacode={}, - insert_highres_powers=False, - rederive_profiles=True, - ): - ''' - Notes: - - insert_highres_powers: whether to insert high resolution powers (will calculate them with powerstate targets object, not other custom ones) - ''' - print(">> Inserting powerstate into input.gacode") - - profiles = TRANSFORMtools.powerstate_to_gacode( - self, - position_in_powerstate_batch=position_in_powerstate_batch, - postprocess_input_gacode=postprocess_input_gacode, - insert_highres_powers=insert_highres_powers, - rederive=rederive_profiles, - ) - # Write input.gacode - if write_input_gacode is not None: - write_input_gacode = Path(write_input_gacode) - print(f"\t- Writing input.gacode file: {IOtools.clipstr(write_input_gacode)}") - write_input_gacode.parent.mkdir(parents=True, exist_ok=True) - profiles.writeCurrentStatus(file=write_input_gacode) - - # If corrections modify the ions set... it's better to re-read, otherwise powerstate will be confused - if rederive_profiles: - TRANSFORMtools.defineIons(self, profiles, self.plasma["rho"][position_in_powerstate_batch, :], self.dfT) - # Repeat, that's how it's done earlier - self._repeat_tensors(batch_size=self.plasma["rho"].shape[0], - specific_keys=["ni","ions_set_mi","ions_set_Zi","ions_set_Dion","ions_set_Tion","ions_set_c_rad"], - positionToUnrepeat=None) - - return profiles # ------------------------------------------------------------------ # Storing and combining diff --git a/src/mitim_modules/powertorch/physics/TRANSPORTtools.py b/src/mitim_modules/powertorch/physics/TRANSPORTtools.py index 8955116c..8cd92763 100644 --- a/src/mitim_modules/powertorch/physics/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/physics/TRANSPORTtools.py @@ -82,7 +82,7 @@ def _produce_profiles(self,deriveQuantities=True): powerstate_detached = self.powerstate.copy_state() - self.powerstate.profiles = powerstate_detached.to_gacode( + self.powerstate.profiles = powerstate_detached.from_powerstate( write_input_gacode=self.file_profs, postprocess_input_gacode=self.applyCorrections, rederive_profiles = deriveQuantities, # Derive quantities so that it's ready for analysis and plotting later diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 9b478dc0..b4e343de 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -112,7 +112,7 @@ def calculator( p.profiles.deriveQuantities() - p.to_gacode( + p.from_powerstate( write_input_gacode=folder / "input.gacode.new.powerstate", position_in_powerstate_batch=0, postprocess_input_gacode={ diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index 51bbe449..8e9da478 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -12,7 +12,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co if figs is not None: # Insert profiles with the latest powerstate - profiles_new = self.to_gacode(insert_highres_powers=True) + profiles_new = self.from_powerstate(insert_highres_powers=True) # Plot the inserted profiles together with the original ones _ = PROFILEStools.plotAll([self.profiles, profiles_new], figs=figs) diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 84fef4f9..304e8649 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -1,9 +1,10 @@ import copy import torch +from pathlib import Path import numpy as np import pandas as pd from mitim_modules.powertorch.physics import CALCtools -from mitim_tools.misc_tools import LOGtools +from mitim_tools.misc_tools import LOGtools, IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.powertorch.physics import TARGETStools from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -13,7 +14,7 @@ # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function -def gacode_to_powerstate(self, input_gacode, rho_vec): +def gacode_to_powerstate(self, increase_profile_resol=False): """ This function converts from the fine input.gacode grid to a powertorch object and grid. Notes: @@ -30,6 +31,13 @@ def gacode_to_powerstate(self, input_gacode, rho_vec): print("\t- Producing powerstate object from input.gacode") + input_gacode = self.profiles + rho_vec = self.plasma["rho"] + + # Resolution of input.gacode + if increase_profile_resol: + improve_resolution_profiles(input_gacode, rho_vec) + # ********************************************************************************************* # Radial grid # ********************************************************************************************* @@ -162,6 +170,7 @@ def gacode_to_powerstate(self, input_gacode, rho_vec): # Define deparametrizer functions for the varying profiles and gradients from here # ********************************************************************************************* + # [quantiy in powerstate, quantity in input.gacode, index of the ion, multiplier, parameterize_in_aLx] cases_to_parameterize = [ ["te", "te(keV)", None, 1.0, True], ["ti", "ti(keV)", 0, 1.0, True], @@ -200,6 +209,45 @@ def gacode_to_powerstate(self, input_gacode, rho_vec): print(f"\t- All values of {key[0]} detected to be zero, to avoid NaNs, inserting {addT} at the edge",typeMsg="w") self.plasma[f"aL{key[0]}"][..., -1] += addT + def to_gacode( + self, + write_input_gacode=None, + position_in_powerstate_batch=0, + postprocess_input_gacode={}, + insert_highres_powers=False, + rederive_profiles=True, + ): + ''' + Notes: + - insert_highres_powers: whether to insert high resolution powers (will calculate them with powerstate targets object, not other custom ones) + ''' + print(">> Inserting powerstate into input.gacode") + + profiles = powerstate_to_gacode( + self, + position_in_powerstate_batch=position_in_powerstate_batch, + postprocess_input_gacode=postprocess_input_gacode, + insert_highres_powers=insert_highres_powers, + rederive=rederive_profiles, + ) + + # Write input.gacode + if write_input_gacode is not None: + write_input_gacode = Path(write_input_gacode) + print(f"\t- Writing input.gacode file: {IOtools.clipstr(write_input_gacode)}") + write_input_gacode.parent.mkdir(parents=True, exist_ok=True) + profiles.writeCurrentStatus(file=write_input_gacode) + + # If corrections modify the ions set... it's better to re-read, otherwise powerstate will be confused + if rederive_profiles: + defineIons(self, profiles, self.plasma["rho"][position_in_powerstate_batch, :], self.dfT) + # Repeat, that's how it's done earlier + self._repeat_tensors(batch_size=self.plasma["rho"].shape[0], + specific_keys=["ni","ions_set_mi","ions_set_Zi","ions_set_Dion","ions_set_Tion","ions_set_c_rad"], + positionToUnrepeat=None) + + return profiles + def powerstate_to_gacode( self, postprocess_input_gacode={}, diff --git a/src/mitim_tools/popcon_tools/RAPIDStools.py b/src/mitim_tools/popcon_tools/RAPIDStools.py index 3d6ad27e..77f0a6e1 100644 --- a/src/mitim_tools/popcon_tools/RAPIDStools.py +++ b/src/mitim_tools/popcon_tools/RAPIDStools.py @@ -138,7 +138,7 @@ def pedestal(p): power = STATEtools.powerstate(p,EvolutionOptions={"rhoPredicted": np.linspace(0.0, 0.9, 50)[1:]}) power.calculate(None, folder='~/scratch/power/') - profiles_new = power.to_gacode(insert_highres_powers=True) + profiles_new = power.from_powerstate(insert_highres_powers=True) return ptop_kPa,profiles_new, eped_evaluation From 49987d71747723679b9a0b605a5d5c218bfaacb5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 14:11:18 -0700 Subject: [PATCH 005/385] Towards deparametrizer renaming --- src/mitim_modules/powertorch/STATEtools.py | 11 +++-------- src/mitim_modules/powertorch/physics/TARGETStools.py | 2 +- src/mitim_modules/powertorch/utils/TRANSFORMtools.py | 2 +- 3 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 927cd43f..010c6f26 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -2,7 +2,6 @@ import torch import datetime import shutil -from pathlib import Path import matplotlib.pyplot as plt import dill as pickle from mitim_tools.misc_tools import PLASMAtools, IOtools @@ -529,7 +528,7 @@ def _cpu_tensors(self): if hasattr(self, 'profiles'): self.profiles.toNumpyArrays() - def update_var(self, name, var=None, specific_deparametrizer=None): + def update_var(self, name, var=None, specific_profile_constructor=None): """ This inserts gradients and updates coarse profiles @@ -543,17 +542,13 @@ def update_var(self, name, var=None, specific_deparametrizer=None): # General function to update a variable # ------------------------------------------------------------------------------------- - deparametrizers_choice = ( - self.deparametrizers_coarse - if specific_deparametrizer is None - else specific_deparametrizer - ) + profile_constructor_choice = self.deparametrizers_coarse if specific_profile_constructor is None else specific_profile_constructor def _update_plasma_var(var_key, clamp_min=None, clamp_max=None): if var is not None: self.plasma[f"aL{var_key}"][: var.shape[0], :] = var[:, :] aLT_withZero = self.plasma[f"aL{var_key}"] - _, varN = deparametrizers_choice[var_key]( + _, varN = profile_constructor_choice[var_key]( self.plasma["roa"], aLT_withZero) self.plasma[var_key] = varN.clamp(min=clamp_min, max=clamp_max) if ( (clamp_min is not None) or (clamp_max is not None) ) else varN self.plasma[f"aL{var_key}"] = torch.cat( diff --git a/src/mitim_modules/powertorch/physics/TARGETStools.py b/src/mitim_modules/powertorch/physics/TARGETStools.py index a13bd90c..f75fded6 100644 --- a/src/mitim_modules/powertorch/physics/TARGETStools.py +++ b/src/mitim_modules/powertorch/physics/TARGETStools.py @@ -82,7 +82,7 @@ def fine_grid(self): # Integrate through fine de-parameterization # ---------------------------------------------------- for i in self.powerstate.ProfilesPredicted: - _ = self.powerstate.update_var(i,specific_deparametrizer=self.powerstate.deparametrizers_coarse_middle) + _ = self.powerstate.update_var(i,specific_profile_constructor=self.powerstate.deparametrizers_coarse_middle) def flux_integrate(self): """ diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 304e8649..6965a215 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -167,7 +167,7 @@ def gacode_to_powerstate(self, increase_profile_resol=False): self.plasma["kradcm"] = 1e-5 / self.plasma["a"] # ********************************************************************************************* - # Define deparametrizer functions for the varying profiles and gradients from here + # Define profile_constructor functions for the varying profiles and gradients from here # ********************************************************************************************* # [quantiy in powerstate, quantity in input.gacode, index of the ion, multiplier, parameterize_in_aLx] From 39c5c5e8e1b56b756f879b0f76888a1b415da5e1 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 14:20:46 -0700 Subject: [PATCH 006/385] Towards a deparametrizer renaming (2) --- .../powertorch/utils/TRANSFORMtools.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 6965a215..e1cf3d91 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -192,7 +192,7 @@ def gacode_to_powerstate(self, increase_profile_resol=False): self.deparametrizers_fine[key[0]], self.deparametrizers_coarse[key[0]], self.deparametrizers_coarse_middle[key[0]], - ) = parameterize_curve( + ) = parameterize_profile( input_gacode.derived["roa"], quant, self.plasma["roa"], @@ -292,7 +292,7 @@ def powerstate_to_gacode( print(f"\t- Inserting {key[0]} into input.gacode profiles") # ********************************************************************************************* - # From a/Lx to x via fine deparametrizer + # From a/Lx to x via fine profile_constructor # ********************************************************************************************* x, y = self.deparametrizers_fine[key[0]]( self.plasma["roa"][position_in_powerstate_batch, :], @@ -458,7 +458,7 @@ def defineIons(self, input_gacode, rho_vec, dfT): self.plasma["ions_set_Tion"] = Tion self.plasma["ions_set_c_rad"] = c_rad -def parameterize_curve( +def parameterize_profile( x_coord, y_coord_raw, x_coarse_tensor, @@ -536,10 +536,10 @@ def parameterize_curve( y_bc_real = torch.from_numpy(interpolation_function([x_coarse[-2]], x_coord, y_coord.cpu().numpy())).to(ygrad_coord) # ********************************************************************************************************** - # Define deparametrizer functions + # Define profile_constructor functions # ********************************************************************************************************** - def deparametrizer_coarse(x, y, multiplier=multiplier_quantity): + def profile_constructor_coarse(x, y, multiplier=multiplier_quantity): """ Construct curve in a coarse grid ---------------------------------------------------------------------------------------------------- @@ -550,7 +550,7 @@ def deparametrizer_coarse(x, y, multiplier=multiplier_quantity): """ return x, integrator_function(x, y, y_bc_real) / multiplier - def deparametrizer_coarse_middle(x, y, multiplier=multiplier_quantity): + def profile_constructor_middle(x, y, multiplier=multiplier_quantity): """ Deparamterizes a finer profile based on the values in the coarse. Reason why something like this is not used for the full profile is because derivative of this will not be as original, @@ -559,7 +559,7 @@ def deparametrizer_coarse_middle(x, y, multiplier=multiplier_quantity): yCPs = CALCtools.Interp1d()(aLy_coarse[:, 0][:-1].repeat((y.shape[0], 1)), y, x) return x, integrator_function(x, yCPs, y_bc_real) / multiplier - def deparametrizer_fine(x, y, multiplier=multiplier_quantity): + def profile_constructor_fine(x, y, multiplier=multiplier_quantity): """ Notes: - x is a 1D array, but y can be a 2D array for a batch of individuals: (batch,x) @@ -619,9 +619,9 @@ def deparametrizer_fine(x, y, multiplier=multiplier_quantity): return ( aLy_coarse, - deparametrizer_fine, - deparametrizer_coarse, - deparametrizer_coarse_middle, + profile_constructor_fine, + profile_constructor_coarse, + profile_constructor_coarse_middle, ) def improve_resolution_profiles(profiles, rhoMODEL): From 30aa3404f6effae5bfbbd96cfc6bbef083106102 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 14:35:31 -0700 Subject: [PATCH 007/385] misc --- src/mitim_modules/powertorch/utils/TRANSFORMtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index e1cf3d91..34a99f2f 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -621,7 +621,7 @@ def profile_constructor_fine(x, y, multiplier=multiplier_quantity): aLy_coarse, profile_constructor_fine, profile_constructor_coarse, - profile_constructor_coarse_middle, + profile_constructor_middle, ) def improve_resolution_profiles(profiles, rhoMODEL): From b2370e613e5f660bda5411ada76f754a4ff8e060 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 15:29:54 -0700 Subject: [PATCH 008/385] Full code refactoring (NO BACK-COMPATIBILITY) to separate transport, targets and deparametrizers (profile_constructors) --- src/mitim_modules/maestro/utils/EPEDbeat.py | 2 +- .../maestro/utils/PORTALSbeat.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 11 +- src/mitim_modules/portals/PORTALStools.py | 9 - .../portals/utils/PORTALSinit.py | 2 +- .../portals/utils/PORTALSoptimization.py | 5 +- src/mitim_modules/powertorch/STATEtools.py | 18 +- .../powertorch/physics/__init__.py | 0 .../physics_models/parameterizers.py | 175 ++++++++++++ .../radiation_chebyshev.csv | 0 .../targets_analytic.py} | 196 +------------ .../physics_models/transport_analytic.py | 90 ++++++ .../transport_gacode.py} | 229 +-------------- .../powertorch/scripts/calculateTargets.py | 8 +- .../{physics => utils}/CALCtools.py | 0 .../{physics => utils}/GEOMETRYtools.py | 0 .../powertorch/utils/TARGETStools.py | 193 +++++++++++++ .../powertorch/utils/TRANSFORMtools.py | 267 ++++-------------- .../powertorch/utils/TRANSPORTtools.py | 144 ++++++++++ src/mitim_tools/gacode_tools/PROFILEStools.py | 6 +- src/mitim_tools/gacode_tools/TGLFtools.py | 2 +- src/mitim_tools/misc_tools/IOtools.py | 2 +- src/mitim_tools/misc_tools/PLASMAtools.py | 2 +- .../popcon_tools/FunctionalForms.py | 2 +- .../popcon_tools/scripts/test_functionals.py | 2 +- .../popcon_tools/utils/FUNCTIONALScalc.py | 2 +- 26 files changed, 693 insertions(+), 676 deletions(-) delete mode 100644 src/mitim_modules/powertorch/physics/__init__.py create mode 100644 src/mitim_modules/powertorch/physics_models/parameterizers.py rename src/mitim_modules/powertorch/{physics => physics_models}/radiation_chebyshev.csv (100%) rename src/mitim_modules/powertorch/{physics/TARGETStools.py => physics_models/targets_analytic.py} (52%) create mode 100644 src/mitim_modules/powertorch/physics_models/transport_analytic.py rename src/mitim_modules/powertorch/{physics/TRANSPORTtools.py => physics_models/transport_gacode.py} (73%) rename src/mitim_modules/powertorch/{physics => utils}/CALCtools.py (100%) rename src/mitim_modules/powertorch/{physics => utils}/GEOMETRYtools.py (100%) create mode 100644 src/mitim_modules/powertorch/utils/TARGETStools.py create mode 100644 src/mitim_modules/powertorch/utils/TRANSPORTtools.py diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index 6167ec0a..0b9303e8 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -10,7 +10,7 @@ from mitim_tools.popcon_tools import FunctionalForms from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_modules.maestro.utils.MAESTRObeat import beat -from mitim_modules.powertorch.physics import CALCtools +from mitim_modules.powertorch.utils import CALCtools from IPython import embed # <> Function to interpolate a curve <> diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index b8bf1552..6dc4aaf2 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -44,7 +44,7 @@ def prepare(self, profiles = self.profiles_current for i in range(len(profiles.Species)): - data_df = pd.read_csv(__mitimroot__ / "src" / "mitim_modules" / "powertorch" / "physics" / "radiation_chebyshev.csv") + data_df = pd.read_csv(__mitimroot__ / "src" / "mitim_modules" / "powertorch" / "physics_models" / "radiation_chebyshev.csv") if not (data_df['Ion'].str.lower()==profiles.Species[i]["N"].lower()).any(): print(f"\t\t- {profiles.Species[i]['N']} not found in radiation table, looking for closest Z (+- 5) USING THE Z SPECIFIED IN THE INPUT.GACODE (fully stripped assumption)",typeMsg='w') # Find closest Z diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 5256ebbd..cfcc5831 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -14,11 +14,12 @@ PORTALSoptimization, PORTALSanalysis, ) -from mitim_modules.powertorch.physics import TRANSPORTtools, TARGETStools +from mitim_modules.powertorch.physics_models import targets_analytic, transport_gacode from mitim_tools.opt_tools import STRATEGYtools from mitim_tools.opt_tools.utils import BOgraphics from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +from tqdm import tnrange """ @@ -193,8 +194,8 @@ def __init__( """ # Selection of model - transport_evaluator = TRANSPORTtools.tgyro_model - targets_evaluator = TARGETStools.analytical_model + transport_evaluator = transport_gacode.tgyro_model + targets_evaluator = targets_analytic.analytical_model self.PORTALSparameters = { "percentError": [5,10,1], # (%) Error (std, in percent) of model evaluation [TGLF (treated as minimum if scan trick), NEO, TARGET] @@ -468,7 +469,7 @@ def check_flags(self): print("\t- Requested fineTargetsResolution, so running powerstate target calculations",typeMsg="w") self.PORTALSparameters["TargetCalc"] = "powerstate" - if not issubclass(self.PORTALSparameters["transport_evaluator"], TRANSPORTtools.tgyro_model) and (self.PORTALSparameters["TargetCalc"] == "tgyro"): + if not issubclass(self.PORTALSparameters["transport_evaluator"], transport_gacode.tgyro_model) and (self.PORTALSparameters["TargetCalc"] == "tgyro"): print("\t- Requested TGYRO targets, but transport evaluator is not tgyro, so changing to powerstate",typeMsg="w") self.PORTALSparameters["TargetCalc"] = "powerstate" @@ -555,7 +556,7 @@ def reuseTrainingTabular( self_copy.powerstate.TransportOptions["transport_evaluator"] = None self_copy.powerstate.TargetOptions["ModelOptions"]["TypeTarget"] = "powerstate" else: - self_copy.powerstate.TransportOptions["transport_evaluator"] = TRANSPORTtools.tgyro_model + self_copy.powerstate.TransportOptions["transport_evaluator"] = transport_gacode.tgyro_model _, dictOFs = runModelEvaluator( self_copy, diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index b1749b6b..f466058c 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -404,12 +404,3 @@ def stopping_criteria_portals(mitim_bo, parameters = {}): else: print("\t- No convergence yet, providing as iteration values the scalarized objective") return False, yvals - -# TODO: Remove in the future, this is just to enable back compatibility with the old code -def selectSurrogate(*args,**kwargs): - return surrogate_selection_portals(*args,**kwargs) -def produceNewInputs(*args,**kwargs): - return input_transform_portals(*args,**kwargs) -def transformPORTALS(*args,**kwargs): - return output_transform_portals(*args,**kwargs) -# ----- diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 423c18e1..96bb1bb3 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -104,7 +104,7 @@ def initializeProblem( # Check if I will be able to calculate radiation speciesNotFound = [] for i in range(len(profiles.Species)): - data_df = pd.read_csv(__mitimroot__ / "src" / "mitim_modules" / "powertorch" / "physics" / "radiation_chebyshev.csv") + data_df = pd.read_csv(__mitimroot__ / "src" / "mitim_modules" / "powertorch" / "physics_models" / "radiation_chebyshev.csv") if not (data_df['Ion'].str.lower()==profiles.Species[i]["N"].lower()).any(): speciesNotFound.append(profiles.Species[i]["N"]) diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index 22eadce3..ba2a9d7a 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -1,9 +1,10 @@ import copy +from mitim_modules.powertorch.physics_models import transport_analytic import torch import shutil import random from functools import partial -from mitim_modules.powertorch.physics import TRANSPORTtools +from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_tools.misc_tools import IOtools from mitim_modules.powertorch import STATEtools from mitim_tools.opt_tools.utils import BOgraphics @@ -128,7 +129,7 @@ def flux_match_surrogate(step,profiles, plot_results=False, fn = None, file_writ TransportOptions = copy.deepcopy(step.surrogate_parameters["powerstate"].TransportOptions) # Define transport calculation function as a surrogate model - TransportOptions['transport_evaluator'] = TRANSPORTtools.surrogate_model + TransportOptions['transport_evaluator'] = transport_analytic.surrogate TransportOptions['ModelOptions'] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} # Create powerstate with the same options as the original portals but with the new profiles diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 010c6f26..b25eec5e 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -2,13 +2,15 @@ import torch import datetime import shutil +from types import MethodType import matplotlib.pyplot as plt import dill as pickle from mitim_tools.misc_tools import PLASMAtools, IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.powertorch.utils import TRANSFORMtools, POWERplot from mitim_tools.opt_tools.optimizers import optim -from mitim_modules.powertorch.physics import TARGETStools, CALCtools, TRANSPORTtools +from mitim_modules.powertorch.utils import TARGETStools, CALCtools, TRANSPORTtools +from mitim_modules.powertorch.physics_models import targets_analytic from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -27,7 +29,7 @@ def __init__( "ModelOptions": {} }, TargetOptions={ - "targets_evaluator": TARGETStools.analytical_model, + "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": 3, "TargetCalc": "powerstate" @@ -121,7 +123,7 @@ def _ensure_ne_before_nz(lst): if isinstance(profiles, PROFILEStools.PROFILES_GACODE): self.to_powerstate = TRANSFORMtools.gacode_to_powerstate - self.from_powerstate = TRANSFORMtools.to_gacode + self.from_powerstate = MethodType(TRANSFORMtools.to_gacode, self) # Use a copy because I'm deriving, it may be expensive and I don't want to carry that out outside of this class self.profiles = copy.deepcopy(profiles) @@ -144,8 +146,12 @@ def _ensure_ne_before_nz(lst): # Standard creation of plasma dictionary # ------------------------------------------------------------------------------------- + # Resolution of input.gacode + if increase_profile_resol: + TRANSFORMtools.improve_resolution_profiles(self.profiles, rho_vec) + # Convert to powerstate - self.to_powerstate(self,increase_profile_resol=increase_profile_resol) + self.to_powerstate(self) # Convert into a batch so that always the quantities are (batch,dimX) self.batch_size = 0 @@ -191,7 +197,7 @@ def _fine_grid(self): ) # Recalculate with higher resolution - TRANSFORMtools.gacode_to_powerstate(self, self.profiles, rho_new) + TRANSFORMtools.gacode_to_powerstate(self, rho_vec = rho_new) self.plasma_fine = copy.deepcopy(self.plasma) # Revert plasma back @@ -542,7 +548,7 @@ def update_var(self, name, var=None, specific_profile_constructor=None): # General function to update a variable # ------------------------------------------------------------------------------------- - profile_constructor_choice = self.deparametrizers_coarse if specific_profile_constructor is None else specific_profile_constructor + profile_constructor_choice = self.profile_constructors_coarse if specific_profile_constructor is None else specific_profile_constructor def _update_plasma_var(var_key, clamp_min=None, clamp_max=None): if var is not None: diff --git a/src/mitim_modules/powertorch/physics/__init__.py b/src/mitim_modules/powertorch/physics/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/mitim_modules/powertorch/physics_models/parameterizers.py b/src/mitim_modules/powertorch/physics_models/parameterizers.py new file mode 100644 index 00000000..36bfe27b --- /dev/null +++ b/src/mitim_modules/powertorch/physics_models/parameterizers.py @@ -0,0 +1,175 @@ +import copy +import torch +import numpy as np +from mitim_modules.powertorch.utils import CALCtools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +# <> Function to interpolate a curve <> +from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function + +def piecewise_linear( + x_coord, + y_coord_raw, + x_coarse_tensor, + parameterize_in_aLx=True, + multiplier_quantity=1.0, + ): + """ + Notes: + - x_coarse_tensor must be torch + """ + + # ********************************************************************************************************** + # Define the integrator and derivator functions (based on whether I want to parameterize in aLx or in gradX) + # ********************************************************************************************************** + + if parameterize_in_aLx: + # 1/Lx = -1/X*dX/dr + integrator_function, derivator_function = ( + CALCtools.integrateGradient, + CALCtools.produceGradient, + ) + else: + # -dX/dr + integrator_function, derivator_function = ( + CALCtools.integrateGradient_lin, + CALCtools.produceGradient_lin, + ) + + y_coord = torch.from_numpy(y_coord_raw).to(x_coarse_tensor) * multiplier_quantity + + ygrad_coord = derivator_function( torch.from_numpy(x_coord).to(x_coarse_tensor), y_coord ) + + # ********************************************************************************************************** + # Get control points + # ********************************************************************************************************** + + x_coarse = x_coarse_tensor[1:].cpu().numpy() + + """ + Define region to get control points from + ------------------------------------------------------------ + Trick: Addition of extra point + This is important because if I don't, when I combine the trailing edge and the new + modified profile, there's going to be a discontinuity in the gradient. + """ + + ir_end = np.argmin(np.abs(x_coord - x_coarse[-1])) + + if ir_end < len(x_coord) - 1: + ir = ir_end + 2 # To prevent that TGYRO does a 2nd order derivative + x_coarse = np.append(x_coarse, [x_coord[ir]]) + else: + ir = ir_end + + # Definition of trailing edge. Any point after, and including, the extra point + x_trail = torch.from_numpy(x_coord[ir:]).to(x_coarse_tensor) + y_trail = y_coord[ir:] + x_notrail = torch.from_numpy(x_coord[: ir + 1]).to(x_coarse_tensor) + + # Produce control points, including a zero at the beginning + aLy_coarse = [[0.0, 0.0]] + for cont, i in enumerate(x_coarse): + yValue = ygrad_coord[np.argmin(np.abs(x_coord - i))] + aLy_coarse.append([i, yValue.cpu().item()]) + + aLy_coarse = torch.from_numpy(np.array(aLy_coarse)).to(ygrad_coord) + + # Since the last one is an extra point very close, I'm making it the same + aLy_coarse[-1, 1] = aLy_coarse[-2, 1] + + # Boundary condition at point moved by gridPointsAllowed + y_bc = torch.from_numpy(interpolation_function([x_coarse[-1]], x_coord, y_coord.cpu().numpy())).to(ygrad_coord) + + # Boundary condition at point (ACTUAL THAT I WANT to keep fixed, i.e. rho=0.8) + y_bc_real = torch.from_numpy(interpolation_function([x_coarse[-2]], x_coord, y_coord.cpu().numpy())).to(ygrad_coord) + + # ********************************************************************************************************** + # Define profile_constructor functions + # ********************************************************************************************************** + + def profile_constructor_coarse(x, y, multiplier=multiplier_quantity): + """ + Construct curve in a coarse grid + ---------------------------------------------------------------------------------------------------- + This constructs a curve in any grid, with any batch given in y=y. + Useful for surrogate evaluations. Fast in a coarse grid. For HF evaluations, + I need to do in a finer grid so that it is consistent with TGYRO. + x, y must be (batch, radii), y_bc must be (1) + """ + return x, integrator_function(x, y, y_bc_real) / multiplier + + def profile_constructor_middle(x, y, multiplier=multiplier_quantity): + """ + Deparamterizes a finer profile based on the values in the coarse. + Reason why something like this is not used for the full profile is because derivative of this will not be as original, + which is needed to match TGYRO + """ + yCPs = CALCtools.Interp1d()(aLy_coarse[:, 0][:-1].repeat((y.shape[0], 1)), y, x) + return x, integrator_function(x, yCPs, y_bc_real) / multiplier + + def profile_constructor_fine(x, y, multiplier=multiplier_quantity): + """ + Notes: + - x is a 1D array, but y can be a 2D array for a batch of individuals: (batch,x) + - I am assuming it is 1/LT for parameterization, but gives T + """ + + y = torch.atleast_2d(y) + x = x[0, :] if x.dim() == 2 else x + + # Add the extra trick point + x = torch.cat((x, aLy_coarse[-1][0].repeat((1)))) + y = torch.cat((y, aLy_coarse[-1][-1].repeat((y.shape[0], 1))), dim=1) + + # Model curve (basically, what happens in between points) + yBS = CALCtools.Interp1d()(x.repeat(y.shape[0], 1), y, x_notrail.repeat(y.shape[0], 1)) + + """ + --------------------------------------------------------------------------------------------------------- + Trick 1: smoothAroundCoarsing + TGYRO will use a 2nd order scheme to obtain gradients out of the profile, so a piecewise linear + will simply not give the right derivatives. + Here, this rough trick is to modify the points in gradient space around the coarse grid with the + same value of gradient, so in principle it doesn't matter the order of the derivative. + """ + num_around = 1 + for i in range(x.shape[0] - 2): + ir = torch.argmin(torch.abs(x[i + 1] - x_notrail)) + for k in range(-num_around, num_around + 1, 1): + yBS[:, ir + k] = yBS[:, ir] + # -------------------------------------------------------------------------------------------------------- + + yBS = integrator_function(x_notrail.repeat(yBS.shape[0], 1), yBS.clone(), y_bc) + + """ + Trick 2: Correct y_bc + The y_bc for the profile integration started at gridPointsAllowed, but that's not the real + y_bc. I want the temperature fixed at my first point that I actually care for. + Here, I multiply the profile to get that. + Multiplication works because: + 1/LT = 1/T * dT/dr + 1/LT' = 1/(T*m) * d(T*m)/dr = 1/T * dT/dr = 1/LT + Same logarithmic gradient, but with the right boundary condition + + """ + ir = torch.argmin(torch.abs(x_notrail - x[-2])) + yBS = yBS * torch.transpose((y_bc_real / yBS[:, ir]).repeat(yBS.shape[1], 1), 0, 1) + + # Add trailing edge + y_trailnew = copy.deepcopy(y_trail).repeat(yBS.shape[0], 1) + + x_notrail_t = torch.cat((x_notrail[:-1], x_trail), dim=0) + yBS = torch.cat((yBS[:, :-1], y_trailnew), dim=1) + + return x_notrail_t, yBS / multiplier + + # ********************************************************************************************************** + + return ( + aLy_coarse, + profile_constructor_fine, + profile_constructor_coarse, + profile_constructor_middle, + ) \ No newline at end of file diff --git a/src/mitim_modules/powertorch/physics/radiation_chebyshev.csv b/src/mitim_modules/powertorch/physics_models/radiation_chebyshev.csv similarity index 100% rename from src/mitim_modules/powertorch/physics/radiation_chebyshev.csv rename to src/mitim_modules/powertorch/physics_models/radiation_chebyshev.csv diff --git a/src/mitim_modules/powertorch/physics/TARGETStools.py b/src/mitim_modules/powertorch/physics_models/targets_analytic.py similarity index 52% rename from src/mitim_modules/powertorch/physics/TARGETStools.py rename to src/mitim_modules/powertorch/physics_models/targets_analytic.py index f75fded6..b36dd63d 100644 --- a/src/mitim_modules/powertorch/physics/TARGETStools.py +++ b/src/mitim_modules/powertorch/physics_models/targets_analytic.py @@ -1,201 +1,9 @@ import torch from mitim_tools.misc_tools import PLASMAtools +from mitim_modules.powertorch.utils import TARGETStools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -# ------------------------------------------------------------------ -# Main classes -# ------------------------------------------------------------------ - -class power_targets: - ''' - Default class for power target models, change "evaluate" method to implement a new model - ''' - - def evaluate(self): - print("No model implemented for power targets", typeMsg="w") - - def __init__(self,powerstate): - self.powerstate = powerstate - - # Make sub-targets equal to zero - variables_to_zero = ["qfuse", "qfusi", "qie", "qrad", "qrad_bremms", "qrad_line", "qrad_sync"] - for i in variables_to_zero: - self.powerstate.plasma[i] = self.powerstate.plasma["te"] * 0.0 - - # ---------------------------------------------------- - # Fixed Targets (targets without a model) - # ---------------------------------------------------- - - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 1: - self.Pe_orig, self.Pi_orig = ( - self.powerstate.plasma["Pe_orig_fusradexch"], - self.powerstate.plasma["Pi_orig_fusradexch"], - ) # Original integrated from input.gacode - elif self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 2: - self.Pe_orig, self.Pi_orig = ( - self.powerstate.plasma["Pe_orig_fusrad"], - self.powerstate.plasma["Pi_orig_fusrad"], - ) - elif self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: - self.Pe_orig, self.Pi_orig = self.powerstate.plasma["te"] * 0.0, self.powerstate.plasma["te"] * 0.0 - - # For the moment, I don't have a model for these, so I just grab the original from input.gacode - self.CextraE = self.powerstate.plasma["Gaux_e"] # 1E20/s/m^2 - self.CextraZ = self.powerstate.plasma["Gaux_Z"] # 1E20/s/m^2 - self.Mextra = self.powerstate.plasma["Maux"] # J/m^2 - - def fine_grid(self): - - """ - Make all quantities needed on the fine resolution - ------------------------------------------------- - In the powerstate creation, the plasma variables are stored in two different resolutions, one for the coarse grid and one for the fine grid, - if the option is activated. - - Here, at calculation stage I use some precalculated quantities in the fine grid and then integrate the gradients into that resolution - - Note that the set ['te','ti','ne','nZ','w0','ni'] will automatically be substituted during the update_var() that comes next, so - it's ok that I lose the torch leaf here. However, I must do this copy here because if any of those variables are not updated in - update_var() then it would fail. But first store them for later use. - """ - - self.plasma_original = {} - - # Bring to fine grid - variables_to_fine = ["B_unit", "B_ref", "volp", "rmin", "roa", "rho", "ni"] - for variable in variables_to_fine: - self.plasma_original[variable] = self.powerstate.plasma[variable].clone() - self.powerstate.plasma[variable] = self.powerstate.plasma_fine[variable] - - # Bring also the gradients and kinetic variables - for variable in self.powerstate.profile_map.keys(): - - # Kinetic variables (te,ti,ne,nZ,w0,ni) - self.plasma_original[variable] = self.powerstate.plasma[variable].clone() - self.powerstate.plasma[variable] = self.powerstate.plasma_fine[variable] - - # Bring also the gradients that are part of the torch trees, so that the derivative is not lost - self.plasma_original[f'aL{variable}'] = self.powerstate.plasma[f'aL{variable}'].clone() - - # ---------------------------------------------------- - # Integrate through fine de-parameterization - # ---------------------------------------------------- - for i in self.powerstate.ProfilesPredicted: - _ = self.powerstate.update_var(i,specific_profile_constructor=self.powerstate.deparametrizers_coarse_middle) - - def flux_integrate(self): - """ - ************************************************************************************************** - Calculate integral of all targets, and then sum aux. - Reason why I do it this convoluted way is to make it faster in mitim, not to run integrateQuadPoly all the time. - Run once for all the batch and also for electrons and ions - (in MW/m^2) - ************************************************************************************************** - """ - - qe = self.powerstate.plasma["te"]*0.0 - qi = self.powerstate.plasma["te"]*0.0 - - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] >= 2: - qe += -self.powerstate.plasma["qie"] - qi += self.powerstate.plasma["qie"] - - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: - qe += self.powerstate.plasma["qfuse"] - self.powerstate.plasma["qrad"] - qi += self.powerstate.plasma["qfusi"] - - q = torch.cat((qe, qi)).to(qe) - self.P = self.powerstate.volume_integrate(q, force_dim=q.shape[0]) - - def coarse_grid(self): - - # ************************************************************************************************** - # Come back to original grid for targets - # ************************************************************************************************** - - # Interpolate results from fine to coarse (i.e. whole point is that it is better than integrate interpolated values) - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] >= 2: - for i in ["qie"]: - self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] - - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: - for i in [ - "qfuse", - "qfusi", - "qrad", - "qrad_bremms", - "qrad_line", - "qrad_sync", - ]: - self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] - - self.P = self.P[:, self.powerstate.positions_targets] - - # Recover variables calculated prior to the fine-targets method - for i in self.plasma_original: - self.powerstate.plasma[i] = self.plasma_original[i] - - def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, assumedPercentError=1.0): - - # ************************************************************************************************** - # Plug-in Targets - # ************************************************************************************************** - - self.powerstate.plasma["Pe"] = ( - self.powerstate.plasma["Paux_e"] + self.P[: self.P.shape[0]//2, :] + self.Pe_orig - ) # MW/m^2 - self.powerstate.plasma["Pi"] = ( - self.powerstate.plasma["Paux_i"] + self.P[self.P.shape[0]//2 :, :] + self.Pi_orig - ) # MW/m^2 - self.powerstate.plasma["Ce_raw"] = self.CextraE - self.powerstate.plasma["CZ_raw"] = self.CextraZ - self.powerstate.plasma["Mt"] = self.Mextra - - # Merge convective fluxes - - if useConvectiveFluxes: - self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux( - self.powerstate.plasma["te"], self.powerstate.plasma["Ce_raw"] - ) # MW/m^2 - self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux( - self.powerstate.plasma["te"], self.powerstate.plasma["CZ_raw"] - ) # MW/m^2 - else: - self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce_raw"] - self.powerstate.plasma["CZ"] = self.powerstate.plasma["CZ_raw"] - - if forceZeroParticleFlux: - self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce"] * 0 - self.powerstate.plasma["Ce_raw"] = self.powerstate.plasma["Ce_raw"] * 0 - - # ************************************************************************************************** - # Error - # ************************************************************************************************** - - variables_to_error = ["Pe", "Pi", "Ce", "CZ", "Mt", "Ce_raw", "CZ_raw"] - - for i in variables_to_error: - self.powerstate.plasma[i + "_stds"] = abs(self.powerstate.plasma[i]) * assumedPercentError / 100 - - """ - ************************************************************************************************** - GB Normalized - ************************************************************************************************** - Note: This is useful for mitim surrogate variables of targets - """ - - gb_mapping = { - "Pe": "Qgb", - "Pi": "Qgb", - "Ce": "Qgb" if useConvectiveFluxes else "Ggb", - "CZ": "Qgb" if useConvectiveFluxes else "Ggb", - "Mt": "Pgb", - } - - for i in gb_mapping.keys(): - self.powerstate.plasma[f"{i}GB"] = self.powerstate.plasma[i] / self.powerstate.plasma[gb_mapping[i]] - # ---------------------------------------------------------------------------------------------------- # Full analytical models taken from TGYRO # ---------------------------------------------------------------------------------------------------- @@ -214,7 +22,7 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, c4, c5, c6, c7 = 4.60643e-3, 1.3500e-2, -1.06750e-4, 1.36600e-5 bg, er = 34.3827, 1.124656e6 -class analytical_model(power_targets): +class analytical_model(TARGETStools.power_targets): def __init__(self,powerstate, **kwargs): super().__init__(powerstate, **kwargs) diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py new file mode 100644 index 00000000..87839749 --- /dev/null +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -0,0 +1,90 @@ +import torch +from mitim_tools.misc_tools import PLASMAtools +from mitim_modules.powertorch.utils import TRANSPORTtools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +# ------------------------------------------------------------------ +# SIMPLE Diffusion (#TODO: implement with particle flux and the raw) +# ------------------------------------------------------------------ + +class diffusion_model(TRANSPORTtools.power_transport): + def __init__(self, powerstate, **kwargs): + super().__init__(powerstate, **kwargs) + + # Ensure that the provided diffusivities include the zero location + self.chi_e = self.powerstate.TransportOptions["ModelOptions"]["chi_e"] + self.chi_i = self.powerstate.TransportOptions["ModelOptions"]["chi_i"] + + if self.chi_e.shape[0] < self.powerstate.plasma['rho'].shape[-1]: + self.chi_e = torch.cat((torch.zeros(1), self.chi_e)) + + if self.chi_i.shape[0] < self.powerstate.plasma['rho'].shape[-1]: + self.chi_i = torch.cat((torch.zeros(1), self.chi_i)) + + def produce_profiles(self): + pass + + def evaluate(self): + + # Make sure the chis are applied to all the points in the batch + Pe_tr = PLASMAtools.conduction( + self.powerstate.plasma["ne"], + self.powerstate.plasma["te"], + self.chi_e.repeat(self.powerstate.plasma['rho'].shape[0],1), + self.powerstate.plasma["aLte"], + self.powerstate.plasma["a"].unsqueeze(-1), + ) + Pi_tr = PLASMAtools.conduction( + self.powerstate.plasma["ni"].sum(axis=-1), + self.powerstate.plasma["ti"], + self.chi_i.repeat(self.powerstate.plasma['rho'].shape[0],1), + self.powerstate.plasma["aLti"], + self.powerstate.plasma["a"].unsqueeze(-1), + ) + + self.powerstate.plasma["Pe_tr_turb"] = Pe_tr * 2 / 3 + self.powerstate.plasma["Pi_tr_turb"] = Pi_tr * 2 / 3 + + self.powerstate.plasma["Pe_tr_neo"] = Pe_tr * 1 / 3 + self.powerstate.plasma["Pi_tr_neo"] = Pi_tr * 1 / 3 + + self.powerstate.plasma["Pe_tr"] = self.powerstate.plasma["Pe_tr_turb"] + self.powerstate.plasma["Pe_tr_neo"] + self.powerstate.plasma["Pi_tr"] = self.powerstate.plasma["Pi_tr_turb"] + self.powerstate.plasma["Pi_tr_neo"] + +# ------------------------------------------------------------------ +# SURROGATE +# ------------------------------------------------------------------ + +class surrogate(TRANSPORTtools.power_transport): + def __init__(self, powerstate, **kwargs): + super().__init__(powerstate, **kwargs) + + def produce_profiles(self): + pass + + def evaluate(self): + + """ + flux_fun as given in ModelOptions must produce Q and Qtargets in order of te,ti,ne + """ + + X = torch.Tensor() + for prof in self.powerstate.ProfilesPredicted: + X = torch.cat((X,self.powerstate.plasma['aL'+prof][:,1:]),axis=1) + + _, Q, _, _ = self.powerstate.TransportOptions["ModelOptions"]["flux_fun"](X) + + numeach = self.powerstate.plasma["rho"].shape[1] - 1 + + quantities = { + "te": "Pe", + "ti": "Pi", + "ne": "Ce", + "nZ": "CZ", + "w0": "Mt", + } + + for c, i in enumerate(self.powerstate.ProfilesPredicted): + self.powerstate.plasma[f"{quantities[i]}_tr"] = torch.cat((torch.tensor([[0.0]]),Q[:, numeach * c : numeach * (c + 1)]),dim=1) + diff --git a/src/mitim_modules/powertorch/physics/TRANSPORTtools.py b/src/mitim_modules/powertorch/physics_models/transport_gacode.py similarity index 73% rename from src/mitim_modules/powertorch/physics/TRANSPORTtools.py rename to src/mitim_modules/powertorch/physics_models/transport_gacode.py index 8cd92763..57166e5c 100644 --- a/src/mitim_modules/powertorch/physics/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/physics_models/transport_gacode.py @@ -2,154 +2,18 @@ import shutil import torch import numpy as np -from mitim_tools.misc_tools import PLASMAtools, IOtools -from mitim_tools.gacode_tools import TGYROtools, PROFILEStools +from mitim_tools.misc_tools import IOtools +from mitim_tools.gacode_tools import TGYROtools from mitim_modules.portals.utils import PORTALScgyro +from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -class power_transport: - ''' - Default class for power transport models, change "evaluate" method to implement a new model and produce_profiles if the model requires written input.gacode written - - Notes: - - After evaluation, the self.model_results attribute will contain the results of the model, which can be used for plotting and analysis - - model results can have .plot() method that can grab kwargs or be similar to TGYRO plot - - ''' - def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_number = 0): - - self.name = name - self.folder = IOtools.expandPath(folder) - self.evaluation_number = evaluation_number - self.powerstate = powerstate - - # Allowed fluxes in powerstate so far - self.quantities = ['Pe', 'Pi', 'Ce', 'CZ', 'Mt'] - - # Each flux has a turbulent and neoclassical component - self.variables = [f'{i}_tr_turb' for i in self.quantities] + [f'{i}_tr_neo' for i in self.quantities] - - # Each flux component has a standard deviation - self.variables += [f'{i}_stds' for i in self.variables] - - # There is also turbulent exchange - self.variables += ['PexchTurb', 'PexchTurb_stds'] - - # And total transport flux - self.variables += [f'{i}_tr' for i in self.quantities] - - # Model results is None by default, but can be assigned in evaluate - self.model_results = None - - # Assign zeros to transport ones if not evaluated - for i in self.variables: - self.powerstate.plasma[i] = self.powerstate.plasma["te"] * 0.0 - - # There is also target components - self.variables += [f'{i}' for i in self.quantities] + [f'{i}_stds' for i in self.quantities] - - # ---------------------------------------------------------------------------------------- - # labels for plotting - # ---------------------------------------------------------------------------------------- - - self.powerstate.labelsFluxes = { - "te": "$Q_e$ ($MW/m^2$)", - "ti": "$Q_i$ ($MW/m^2$)", - "ne": ( - "$Q_{conv}$ ($MW/m^2$)" - if self.powerstate.TransportOptions["ModelOptions"].get("useConvectiveFluxes", True) - else "$\\Gamma_e$ ($10^{20}/s/m^2$)" - ), - "nZ": ( - "$Q_{conv}$ $\\cdot f_{Z,0}$ ($MW/m^2$)" - if self.powerstate.TransportOptions["ModelOptions"].get("useConvectiveFluxes", True) - else "$\\Gamma_Z$ $\\cdot f_{Z,0}$ ($10^{20}/s/m^2$)" - ), - "w0": "$M_T$ ($J/m^2$)", - } - - def produce_profiles(self): - # Only add self._produce_profiles() if it's needed (e.g. full TGLF), otherwise this is somewhat expensive (e.g. for flux matching) - pass - - def _produce_profiles(self,deriveQuantities=True): - - self.applyCorrections = self.powerstate.TransportOptions["ModelOptions"].get("MODELparameters", {}).get("applyCorrections", {}) - - # Write this updated profiles class (with parameterized profiles and target powers) - self.file_profs = self.folder / "input.gacode" - - powerstate_detached = self.powerstate.copy_state() - - self.powerstate.profiles = powerstate_detached.from_powerstate( - write_input_gacode=self.file_profs, - postprocess_input_gacode=self.applyCorrections, - rederive_profiles = deriveQuantities, # Derive quantities so that it's ready for analysis and plotting later - insert_highres_powers = deriveQuantities, # Insert powers so that Q, Pfus and all that it's consistent when read later - ) - - self.profiles_transport = copy.deepcopy(self.powerstate.profiles) - - self._modify_profiles() - - def _modify_profiles(self): - ''' - Modify the profiles (e.g. lumping) before running the transport model - ''' - - # After producing the profiles, copy for future modifications - self.file_profs_unmod = self.file_profs.parent / f"{self.file_profs.name}_unmodified" - shutil.copy2(self.file_profs, self.file_profs_unmod) - - profiles_postprocessing_fun = self.powerstate.TransportOptions["ModelOptions"].get("profiles_postprocessing_fun", None) - - if profiles_postprocessing_fun is not None: - print(f"\t- Modifying input.gacode to run transport calculations based on {profiles_postprocessing_fun}",typeMsg="i") - self.profiles_transport = profiles_postprocessing_fun(self.file_profs) - - # Position of impurity ion may have changed - p_old = PROFILEStools.PROFILES_GACODE(self.file_profs_unmod) - p_new = PROFILEStools.PROFILES_GACODE(self.file_profs) - - impurity_of_interest = p_old.Species[self.powerstate.impurityPosition] - - try: - impurityPosition_new = p_new.Species.index(impurity_of_interest) - - except ValueError: - print(f"\t- Impurity {impurity_of_interest} not found in new profiles, keeping position {self.powerstate.impurityPosition}",typeMsg="w") - impurityPosition_new = self.powerstate.impurityPosition - - if impurityPosition_new != self.powerstate.impurityPosition: - print(f"\t- Impurity position has changed from {self.powerstate.impurityPosition} to {impurityPosition_new}",typeMsg="w") - self.powerstate.impurityPosition_transport = p_new.Species.index(impurity_of_interest) - - # ---------------------------------------------------------------------------------------------------- - # EVALUATE (custom part) - # ---------------------------------------------------------------------------------------------------- - def evaluate(self): - ''' - This needs to populate the following in self.powerstate.plasma - - Pe, Pe_tr, Pe_tr_turb, Pe_tr_neo -> MW/m^2 - - Pi, Pi_tr, Pi_tr_turb, Pi_tr_neo -> MW/m^2 - - Ce, Ce_tr, Ce_tr_turb, Ce_tr_neo -> MW/m^2 - * Ce_raw, Ce_raw_tr, Ce_raw_tr_turb, Ce_raw_tr_neo -> 10^20/s/m^2 - - CZ, CZ_tr, CZ_tr_turb, CZ_tr_neo -> MW/m^2 (but modified as needed, for example dividing by fZ0) - * CZ_raw, CZ_raw_tr, CZ_raw_tr_turb, CZ_raw_tr_neo -> 10^20/s/m^2 (NOT modified) - - Mt, Mt_tr, Mt_tr_turb, Mt_tr_neo -> J/m^2 - - PexchTurb -> MW/m^3 - and their respective standard deviations - ''' - - print(">> No transport fluxes to evaluate", typeMsg="w") - pass - # ---------------------------------------------------------------------------------------------------- # FULL TGYRO # ---------------------------------------------------------------------------------------------------- -class tgyro_model(power_transport): +class tgyro_model(TRANSPORTtools.power_transport): def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) @@ -514,90 +378,6 @@ def calculate_mean_std(Q): return Qe_point, Qi_point, Ge_point, GZ_point, Mt_point, Pexch_point, Qe_std, Qi_std, Ge_std, GZ_std, Mt_std, Pexch_std -# ------------------------------------------------------------------ -# SIMPLE Diffusion (#TODO: implement with particle flux and the raw) -# ------------------------------------------------------------------ - -class diffusion_model(power_transport): - def __init__(self, powerstate, **kwargs): - super().__init__(powerstate, **kwargs) - - # Ensure that the provided diffusivities include the zero location - self.chi_e = self.powerstate.TransportOptions["ModelOptions"]["chi_e"] - self.chi_i = self.powerstate.TransportOptions["ModelOptions"]["chi_i"] - - if self.chi_e.shape[0] < self.powerstate.plasma['rho'].shape[-1]: - self.chi_e = torch.cat((torch.zeros(1), self.chi_e)) - - if self.chi_i.shape[0] < self.powerstate.plasma['rho'].shape[-1]: - self.chi_i = torch.cat((torch.zeros(1), self.chi_i)) - - def produce_profiles(self): - pass - - def evaluate(self): - - # Make sure the chis are applied to all the points in the batch - Pe_tr = PLASMAtools.conduction( - self.powerstate.plasma["ne"], - self.powerstate.plasma["te"], - self.chi_e.repeat(self.powerstate.plasma['rho'].shape[0],1), - self.powerstate.plasma["aLte"], - self.powerstate.plasma["a"].unsqueeze(-1), - ) - Pi_tr = PLASMAtools.conduction( - self.powerstate.plasma["ni"].sum(axis=-1), - self.powerstate.plasma["ti"], - self.chi_i.repeat(self.powerstate.plasma['rho'].shape[0],1), - self.powerstate.plasma["aLti"], - self.powerstate.plasma["a"].unsqueeze(-1), - ) - - self.powerstate.plasma["Pe_tr_turb"] = Pe_tr * 2 / 3 - self.powerstate.plasma["Pi_tr_turb"] = Pi_tr * 2 / 3 - - self.powerstate.plasma["Pe_tr_neo"] = Pe_tr * 1 / 3 - self.powerstate.plasma["Pi_tr_neo"] = Pi_tr * 1 / 3 - - self.powerstate.plasma["Pe_tr"] = self.powerstate.plasma["Pe_tr_turb"] + self.powerstate.plasma["Pe_tr_neo"] - self.powerstate.plasma["Pi_tr"] = self.powerstate.plasma["Pi_tr_turb"] + self.powerstate.plasma["Pi_tr_neo"] - -# ------------------------------------------------------------------ -# SURROGATE -# ------------------------------------------------------------------ - -class surrogate_model(power_transport): - def __init__(self, powerstate, **kwargs): - super().__init__(powerstate, **kwargs) - - def produce_profiles(self): - pass - - def evaluate(self): - - """ - flux_fun as given in ModelOptions must produce Q and Qtargets in order of te,ti,ne - """ - - X = torch.Tensor() - for prof in self.powerstate.ProfilesPredicted: - X = torch.cat((X,self.powerstate.plasma['aL'+prof][:,1:]),axis=1) - - _, Q, _, _ = self.powerstate.TransportOptions["ModelOptions"]["flux_fun"](X) - - numeach = self.powerstate.plasma["rho"].shape[1] - 1 - - quantities = { - "te": "Pe", - "ti": "Pi", - "ne": "Ce", - "nZ": "CZ", - "w0": "Mt", - } - - for c, i in enumerate(self.powerstate.ProfilesPredicted): - self.powerstate.plasma[f"{quantities[i]}_tr"] = torch.cat((torch.tensor([[0.0]]),Q[:, numeach * c : numeach * (c + 1)]),dim=1) - # ************************************************************************************************** # Functions # ************************************************************************************************** @@ -879,3 +659,4 @@ def dummyCDF(GeneralFolder, FolderEvaluation): cdf = FolderEvaluation / f"{subname}_ev{name}.CDF" return cdf + diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index b4e343de..ae7cb7ed 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -7,7 +7,7 @@ from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.powertorch import STATEtools -from mitim_modules.powertorch.physics import TRANSPORTtools,TARGETStools +from mitim_modules.powertorch.physics_models import targets_analytic, transport_gacode from IPython import embed def calculator( @@ -33,13 +33,13 @@ def calculator( 'fineTargetsResolution': fineTargetsResolution, }, TargetOptions={ - "targets_evaluator": TARGETStools.analytical_model, + "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": TypeTarget, "TargetCalc": "tgyro"}, }, TransportOptions={ - "transport_evaluator": TRANSPORTtools.tgyro_model, + "transport_evaluator": transport_gacode.tgyro_model, "ModelOptions": { "cold_start": cold_start, "launchSlurm": True, @@ -75,7 +75,7 @@ def calculator( 'fineTargetsResolution': fineTargetsResolution, }, TargetOptions={ - "targets_evaluator": TARGETStools.analytical_model, + "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": TypeTarget, "TargetCalc": "powerstate"}, diff --git a/src/mitim_modules/powertorch/physics/CALCtools.py b/src/mitim_modules/powertorch/utils/CALCtools.py similarity index 100% rename from src/mitim_modules/powertorch/physics/CALCtools.py rename to src/mitim_modules/powertorch/utils/CALCtools.py diff --git a/src/mitim_modules/powertorch/physics/GEOMETRYtools.py b/src/mitim_modules/powertorch/utils/GEOMETRYtools.py similarity index 100% rename from src/mitim_modules/powertorch/physics/GEOMETRYtools.py rename to src/mitim_modules/powertorch/utils/GEOMETRYtools.py diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py new file mode 100644 index 00000000..526562f6 --- /dev/null +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -0,0 +1,193 @@ +import torch +from mitim_tools.misc_tools import PLASMAtools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +class power_targets: + ''' + Default class for power target models, change "evaluate" method to implement a new model + ''' + + def evaluate(self): + print("No model implemented for power targets", typeMsg="w") + + def __init__(self,powerstate): + self.powerstate = powerstate + + # Make sub-targets equal to zero + variables_to_zero = ["qfuse", "qfusi", "qie", "qrad", "qrad_bremms", "qrad_line", "qrad_sync"] + for i in variables_to_zero: + self.powerstate.plasma[i] = self.powerstate.plasma["te"] * 0.0 + + # ---------------------------------------------------- + # Fixed Targets (targets without a model) + # ---------------------------------------------------- + + if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 1: + self.Pe_orig, self.Pi_orig = ( + self.powerstate.plasma["Pe_orig_fusradexch"], + self.powerstate.plasma["Pi_orig_fusradexch"], + ) # Original integrated from input.gacode + elif self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 2: + self.Pe_orig, self.Pi_orig = ( + self.powerstate.plasma["Pe_orig_fusrad"], + self.powerstate.plasma["Pi_orig_fusrad"], + ) + elif self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: + self.Pe_orig, self.Pi_orig = self.powerstate.plasma["te"] * 0.0, self.powerstate.plasma["te"] * 0.0 + + # For the moment, I don't have a model for these, so I just grab the original from input.gacode + self.CextraE = self.powerstate.plasma["Gaux_e"] # 1E20/s/m^2 + self.CextraZ = self.powerstate.plasma["Gaux_Z"] # 1E20/s/m^2 + self.Mextra = self.powerstate.plasma["Maux"] # J/m^2 + + def fine_grid(self): + + """ + Make all quantities needed on the fine resolution + ------------------------------------------------- + In the powerstate creation, the plasma variables are stored in two different resolutions, one for the coarse grid and one for the fine grid, + if the option is activated. + + Here, at calculation stage I use some precalculated quantities in the fine grid and then integrate the gradients into that resolution + + Note that the set ['te','ti','ne','nZ','w0','ni'] will automatically be substituted during the update_var() that comes next, so + it's ok that I lose the torch leaf here. However, I must do this copy here because if any of those variables are not updated in + update_var() then it would fail. But first store them for later use. + """ + + self.plasma_original = {} + + # Bring to fine grid + variables_to_fine = ["B_unit", "B_ref", "volp", "rmin", "roa", "rho", "ni"] + for variable in variables_to_fine: + self.plasma_original[variable] = self.powerstate.plasma[variable].clone() + self.powerstate.plasma[variable] = self.powerstate.plasma_fine[variable] + + # Bring also the gradients and kinetic variables + for variable in self.powerstate.profile_map.keys(): + + # Kinetic variables (te,ti,ne,nZ,w0,ni) + self.plasma_original[variable] = self.powerstate.plasma[variable].clone() + self.powerstate.plasma[variable] = self.powerstate.plasma_fine[variable] + + # Bring also the gradients that are part of the torch trees, so that the derivative is not lost + self.plasma_original[f'aL{variable}'] = self.powerstate.plasma[f'aL{variable}'].clone() + + # ---------------------------------------------------- + # Integrate through fine de-parameterization + # ---------------------------------------------------- + for i in self.powerstate.ProfilesPredicted: + _ = self.powerstate.update_var(i,specific_profile_constructor=self.powerstate.profile_constructors_coarse_middle) + + def flux_integrate(self): + """ + ************************************************************************************************** + Calculate integral of all targets, and then sum aux. + Reason why I do it this convoluted way is to make it faster in mitim, not to run integrateQuadPoly all the time. + Run once for all the batch and also for electrons and ions + (in MW/m^2) + ************************************************************************************************** + """ + + qe = self.powerstate.plasma["te"]*0.0 + qi = self.powerstate.plasma["te"]*0.0 + + if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] >= 2: + qe += -self.powerstate.plasma["qie"] + qi += self.powerstate.plasma["qie"] + + if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: + qe += self.powerstate.plasma["qfuse"] - self.powerstate.plasma["qrad"] + qi += self.powerstate.plasma["qfusi"] + + q = torch.cat((qe, qi)).to(qe) + self.P = self.powerstate.volume_integrate(q, force_dim=q.shape[0]) + + def coarse_grid(self): + + # ************************************************************************************************** + # Come back to original grid for targets + # ************************************************************************************************** + + # Interpolate results from fine to coarse (i.e. whole point is that it is better than integrate interpolated values) + if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] >= 2: + for i in ["qie"]: + self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] + + if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: + for i in [ + "qfuse", + "qfusi", + "qrad", + "qrad_bremms", + "qrad_line", + "qrad_sync", + ]: + self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] + + self.P = self.P[:, self.powerstate.positions_targets] + + # Recover variables calculated prior to the fine-targets method + for i in self.plasma_original: + self.powerstate.plasma[i] = self.plasma_original[i] + + def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, assumedPercentError=1.0): + + # ************************************************************************************************** + # Plug-in Targets + # ************************************************************************************************** + + self.powerstate.plasma["Pe"] = ( + self.powerstate.plasma["Paux_e"] + self.P[: self.P.shape[0]//2, :] + self.Pe_orig + ) # MW/m^2 + self.powerstate.plasma["Pi"] = ( + self.powerstate.plasma["Paux_i"] + self.P[self.P.shape[0]//2 :, :] + self.Pi_orig + ) # MW/m^2 + self.powerstate.plasma["Ce_raw"] = self.CextraE + self.powerstate.plasma["CZ_raw"] = self.CextraZ + self.powerstate.plasma["Mt"] = self.Mextra + + # Merge convective fluxes + + if useConvectiveFluxes: + self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux( + self.powerstate.plasma["te"], self.powerstate.plasma["Ce_raw"] + ) # MW/m^2 + self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux( + self.powerstate.plasma["te"], self.powerstate.plasma["CZ_raw"] + ) # MW/m^2 + else: + self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce_raw"] + self.powerstate.plasma["CZ"] = self.powerstate.plasma["CZ_raw"] + + if forceZeroParticleFlux: + self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce"] * 0 + self.powerstate.plasma["Ce_raw"] = self.powerstate.plasma["Ce_raw"] * 0 + + # ************************************************************************************************** + # Error + # ************************************************************************************************** + + variables_to_error = ["Pe", "Pi", "Ce", "CZ", "Mt", "Ce_raw", "CZ_raw"] + + for i in variables_to_error: + self.powerstate.plasma[i + "_stds"] = abs(self.powerstate.plasma[i]) * assumedPercentError / 100 + + """ + ************************************************************************************************** + GB Normalized + ************************************************************************************************** + Note: This is useful for mitim surrogate variables of targets + """ + + gb_mapping = { + "Pe": "Qgb", + "Pi": "Qgb", + "Ce": "Qgb" if useConvectiveFluxes else "Ggb", + "CZ": "Qgb" if useConvectiveFluxes else "Ggb", + "Mt": "Pgb", + } + + for i in gb_mapping.keys(): + self.powerstate.plasma[f"{i}GB"] = self.powerstate.plasma[i] / self.powerstate.plasma[gb_mapping[i]] diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 34a99f2f..ff357f58 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -3,10 +3,9 @@ from pathlib import Path import numpy as np import pandas as pd -from mitim_modules.powertorch.physics import CALCtools from mitim_tools.misc_tools import LOGtools, IOtools from mitim_tools.gacode_tools import PROFILEStools -from mitim_modules.powertorch.physics import TARGETStools +from mitim_modules.powertorch.physics_models import targets_analytic, parameterizers from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __mitimroot__ from IPython import embed @@ -14,7 +13,7 @@ # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function -def gacode_to_powerstate(self, increase_profile_resol=False): +def gacode_to_powerstate(self, rho_vec=None): """ This function converts from the fine input.gacode grid to a powertorch object and grid. Notes: @@ -31,12 +30,9 @@ def gacode_to_powerstate(self, increase_profile_resol=False): print("\t- Producing powerstate object from input.gacode") - input_gacode = self.profiles - rho_vec = self.plasma["rho"] - - # Resolution of input.gacode - if increase_profile_resol: - improve_resolution_profiles(input_gacode, rho_vec) + input_gacode = self.profiles + if rho_vec is None: + rho_vec = self.plasma["rho"] # ********************************************************************************************* # Radial grid @@ -183,16 +179,16 @@ def gacode_to_powerstate(self, increase_profile_resol=False): for i in range(input_gacode.profiles['ni(10^19/m^3)'].shape[1]): cases_to_parameterize.append([f"ni{i}", "ni(10^19/m^3)", i, 1.0, True]) - self.deparametrizers_fine, self.deparametrizers_coarse, self.deparametrizers_coarse_middle = {}, {}, {} + self.profile_constructors_fine, self.profile_constructors_coarse, self.profile_constructors_coarse_middle = {}, {}, {} for key in cases_to_parameterize: quant = input_gacode.profiles[key[1]] if key[2] is None else input_gacode.profiles[key[1]][:, key[2]] ( aLy_coarse, - self.deparametrizers_fine[key[0]], - self.deparametrizers_coarse[key[0]], - self.deparametrizers_coarse_middle[key[0]], - ) = parameterize_profile( + self.profile_constructors_fine[key[0]], + self.profile_constructors_coarse[key[0]], + self.profile_constructors_coarse_middle[key[0]], + ) = parameterizers.piecewise_linear( input_gacode.derived["roa"], quant, self.plasma["roa"], @@ -209,44 +205,44 @@ def gacode_to_powerstate(self, increase_profile_resol=False): print(f"\t- All values of {key[0]} detected to be zero, to avoid NaNs, inserting {addT} at the edge",typeMsg="w") self.plasma[f"aL{key[0]}"][..., -1] += addT - def to_gacode( - self, - write_input_gacode=None, - position_in_powerstate_batch=0, - postprocess_input_gacode={}, - insert_highres_powers=False, - rederive_profiles=True, - ): - ''' - Notes: - - insert_highres_powers: whether to insert high resolution powers (will calculate them with powerstate targets object, not other custom ones) - ''' - print(">> Inserting powerstate into input.gacode") - - profiles = powerstate_to_gacode( - self, - position_in_powerstate_batch=position_in_powerstate_batch, - postprocess_input_gacode=postprocess_input_gacode, - insert_highres_powers=insert_highres_powers, - rederive=rederive_profiles, - ) +def to_gacode( + self, + write_input_gacode=None, + position_in_powerstate_batch=0, + postprocess_input_gacode={}, + insert_highres_powers=False, + rederive_profiles=True, +): + ''' + Notes: + - insert_highres_powers: whether to insert high resolution powers (will calculate them with powerstate targets object, not other custom ones) + ''' + print(">> Inserting powerstate into input.gacode") - # Write input.gacode - if write_input_gacode is not None: - write_input_gacode = Path(write_input_gacode) - print(f"\t- Writing input.gacode file: {IOtools.clipstr(write_input_gacode)}") - write_input_gacode.parent.mkdir(parents=True, exist_ok=True) - profiles.writeCurrentStatus(file=write_input_gacode) + profiles = powerstate_to_gacode( + self, + position_in_powerstate_batch=position_in_powerstate_batch, + postprocess_input_gacode=postprocess_input_gacode, + insert_highres_powers=insert_highres_powers, + rederive=rederive_profiles, + ) - # If corrections modify the ions set... it's better to re-read, otherwise powerstate will be confused - if rederive_profiles: - defineIons(self, profiles, self.plasma["rho"][position_in_powerstate_batch, :], self.dfT) - # Repeat, that's how it's done earlier - self._repeat_tensors(batch_size=self.plasma["rho"].shape[0], - specific_keys=["ni","ions_set_mi","ions_set_Zi","ions_set_Dion","ions_set_Tion","ions_set_c_rad"], - positionToUnrepeat=None) + # Write input.gacode + if write_input_gacode is not None: + write_input_gacode = Path(write_input_gacode) + print(f"\t- Writing input.gacode file: {IOtools.clipstr(write_input_gacode)}") + write_input_gacode.parent.mkdir(parents=True, exist_ok=True) + profiles.writeCurrentStatus(file=write_input_gacode) + + # If corrections modify the ions set... it's better to re-read, otherwise powerstate will be confused + if rederive_profiles: + defineIons(self, profiles, self.plasma["rho"][position_in_powerstate_batch, :], self.dfT) + # Repeat, that's how it's done earlier + self._repeat_tensors(batch_size=self.plasma["rho"].shape[0], + specific_keys=["ni","ions_set_mi","ions_set_Zi","ions_set_Dion","ions_set_Tion","ions_set_c_rad"], + positionToUnrepeat=None) - return profiles + return profiles def powerstate_to_gacode( self, @@ -294,7 +290,7 @@ def powerstate_to_gacode( # ********************************************************************************************* # From a/Lx to x via fine profile_constructor # ********************************************************************************************* - x, y = self.deparametrizers_fine[key[0]]( + x, y = self.profile_constructors_fine[key[0]]( self.plasma["roa"][position_in_powerstate_batch, :], self.plasma[f"aL{key[0]}"][position_in_powerstate_batch, :], ) @@ -375,7 +371,7 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): profiles, EvolutionOptions={"rhoPredicted": rhoy}, TargetOptions={ - "targets_evaluator": TARGETStools.analytical_model, + "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": self.TargetOptions["ModelOptions"]["TypeTarget"], # Important to keep the same as in the original "TargetCalc": "powerstate", @@ -433,7 +429,7 @@ def defineIons(self, input_gacode, rho_vec, dfT): Zi.append(input_gacode.profiles["z"][i]) # Grab chebyshev coefficients from file - data_df = pd.read_csv(__mitimroot__ / "src" / "mitim_modules" / "powertorch" / "physics" / "radiation_chebyshev.csv") + data_df = pd.read_csv(__mitimroot__ / "src" / "mitim_modules" / "powertorch" / "physics_models" / "radiation_chebyshev.csv") try: c = data_df[data_df['Ion'].str.lower()==input_gacode.profiles["name"][i].lower()].to_numpy()[0,2:].astype(float) except IndexError: @@ -458,172 +454,6 @@ def defineIons(self, input_gacode, rho_vec, dfT): self.plasma["ions_set_Tion"] = Tion self.plasma["ions_set_c_rad"] = c_rad -def parameterize_profile( - x_coord, - y_coord_raw, - x_coarse_tensor, - parameterize_in_aLx=True, - multiplier_quantity=1.0, - ): - """ - Notes: - - x_coarse_tensor must be torch - """ - - # ********************************************************************************************************** - # Define the integrator and derivator functions (based on whether I want to parameterize in aLx or in gradX) - # ********************************************************************************************************** - - if parameterize_in_aLx: - # 1/Lx = -1/X*dX/dr - integrator_function, derivator_function = ( - CALCtools.integrateGradient, - CALCtools.produceGradient, - ) - else: - # -dX/dr - integrator_function, derivator_function = ( - CALCtools.integrateGradient_lin, - CALCtools.produceGradient_lin, - ) - - y_coord = torch.from_numpy(y_coord_raw).to(x_coarse_tensor) * multiplier_quantity - - ygrad_coord = derivator_function( torch.from_numpy(x_coord).to(x_coarse_tensor), y_coord ) - - # ********************************************************************************************************** - # Get control points - # ********************************************************************************************************** - - x_coarse = x_coarse_tensor[1:].cpu().numpy() - - """ - Define region to get control points from - ------------------------------------------------------------ - Trick: Addition of extra point - This is important because if I don't, when I combine the trailing edge and the new - modified profile, there's going to be a discontinuity in the gradient. - """ - - ir_end = np.argmin(np.abs(x_coord - x_coarse[-1])) - - if ir_end < len(x_coord) - 1: - ir = ir_end + 2 # To prevent that TGYRO does a 2nd order derivative - x_coarse = np.append(x_coarse, [x_coord[ir]]) - else: - ir = ir_end - - # Definition of trailing edge. Any point after, and including, the extra point - x_trail = torch.from_numpy(x_coord[ir:]).to(x_coarse_tensor) - y_trail = y_coord[ir:] - x_notrail = torch.from_numpy(x_coord[: ir + 1]).to(x_coarse_tensor) - - # Produce control points, including a zero at the beginning - aLy_coarse = [[0.0, 0.0]] - for cont, i in enumerate(x_coarse): - yValue = ygrad_coord[np.argmin(np.abs(x_coord - i))] - aLy_coarse.append([i, yValue.cpu().item()]) - - aLy_coarse = torch.from_numpy(np.array(aLy_coarse)).to(ygrad_coord) - - # Since the last one is an extra point very close, I'm making it the same - aLy_coarse[-1, 1] = aLy_coarse[-2, 1] - - # Boundary condition at point moved by gridPointsAllowed - y_bc = torch.from_numpy(interpolation_function([x_coarse[-1]], x_coord, y_coord.cpu().numpy())).to(ygrad_coord) - - # Boundary condition at point (ACTUAL THAT I WANT to keep fixed, i.e. rho=0.8) - y_bc_real = torch.from_numpy(interpolation_function([x_coarse[-2]], x_coord, y_coord.cpu().numpy())).to(ygrad_coord) - - # ********************************************************************************************************** - # Define profile_constructor functions - # ********************************************************************************************************** - - def profile_constructor_coarse(x, y, multiplier=multiplier_quantity): - """ - Construct curve in a coarse grid - ---------------------------------------------------------------------------------------------------- - This constructs a curve in any grid, with any batch given in y=y. - Useful for surrogate evaluations. Fast in a coarse grid. For HF evaluations, - I need to do in a finer grid so that it is consistent with TGYRO. - x, y must be (batch, radii), y_bc must be (1) - """ - return x, integrator_function(x, y, y_bc_real) / multiplier - - def profile_constructor_middle(x, y, multiplier=multiplier_quantity): - """ - Deparamterizes a finer profile based on the values in the coarse. - Reason why something like this is not used for the full profile is because derivative of this will not be as original, - which is needed to match TGYRO - """ - yCPs = CALCtools.Interp1d()(aLy_coarse[:, 0][:-1].repeat((y.shape[0], 1)), y, x) - return x, integrator_function(x, yCPs, y_bc_real) / multiplier - - def profile_constructor_fine(x, y, multiplier=multiplier_quantity): - """ - Notes: - - x is a 1D array, but y can be a 2D array for a batch of individuals: (batch,x) - - I am assuming it is 1/LT for parameterization, but gives T - """ - - y = torch.atleast_2d(y) - x = x[0, :] if x.dim() == 2 else x - - # Add the extra trick point - x = torch.cat((x, aLy_coarse[-1][0].repeat((1)))) - y = torch.cat((y, aLy_coarse[-1][-1].repeat((y.shape[0], 1))), dim=1) - - # Model curve (basically, what happens in between points) - yBS = CALCtools.Interp1d()(x.repeat(y.shape[0], 1), y, x_notrail.repeat(y.shape[0], 1)) - - """ - --------------------------------------------------------------------------------------------------------- - Trick 1: smoothAroundCoarsing - TGYRO will use a 2nd order scheme to obtain gradients out of the profile, so a piecewise linear - will simply not give the right derivatives. - Here, this rough trick is to modify the points in gradient space around the coarse grid with the - same value of gradient, so in principle it doesn't matter the order of the derivative. - """ - num_around = 1 - for i in range(x.shape[0] - 2): - ir = torch.argmin(torch.abs(x[i + 1] - x_notrail)) - for k in range(-num_around, num_around + 1, 1): - yBS[:, ir + k] = yBS[:, ir] - # -------------------------------------------------------------------------------------------------------- - - yBS = integrator_function(x_notrail.repeat(yBS.shape[0], 1), yBS.clone(), y_bc) - - """ - Trick 2: Correct y_bc - The y_bc for the profile integration started at gridPointsAllowed, but that's not the real - y_bc. I want the temperature fixed at my first point that I actually care for. - Here, I multiply the profile to get that. - Multiplication works because: - 1/LT = 1/T * dT/dr - 1/LT' = 1/(T*m) * d(T*m)/dr = 1/T * dT/dr = 1/LT - Same logarithmic gradient, but with the right boundary condition - - """ - ir = torch.argmin(torch.abs(x_notrail - x[-2])) - yBS = yBS * torch.transpose((y_bc_real / yBS[:, ir]).repeat(yBS.shape[1], 1), 0, 1) - - # Add trailing edge - y_trailnew = copy.deepcopy(y_trail).repeat(yBS.shape[0], 1) - - x_notrail_t = torch.cat((x_notrail[:-1], x_trail), dim=0) - yBS = torch.cat((yBS[:, :-1], y_trailnew), dim=1) - - return x_notrail_t, yBS / multiplier - - # ********************************************************************************************************** - - return ( - aLy_coarse, - profile_constructor_fine, - profile_constructor_coarse, - profile_constructor_middle, - ) - def improve_resolution_profiles(profiles, rhoMODEL): """ Resolution of input.gacode @@ -674,7 +504,6 @@ def improve_resolution_profiles(profiles, rhoMODEL): # ---------------------------------------------------------------------------------- profiles.changeResolution(rho_new=rho_new) - def debug_transformation(p, p_new, s): rho = s.plasma['rho'][0][1:] diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py new file mode 100644 index 00000000..11e8e0b6 --- /dev/null +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -0,0 +1,144 @@ +import copy +import shutil +from mitim_tools.misc_tools import IOtools +from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +class power_transport: + ''' + Default class for power transport models, change "evaluate" method to implement a new model and produce_profiles if the model requires written input.gacode written + + Notes: + - After evaluation, the self.model_results attribute will contain the results of the model, which can be used for plotting and analysis + - model results can have .plot() method that can grab kwargs or be similar to TGYRO plot + + ''' + def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_number = 0): + + self.name = name + self.folder = IOtools.expandPath(folder) + self.evaluation_number = evaluation_number + self.powerstate = powerstate + + # Allowed fluxes in powerstate so far + self.quantities = ['Pe', 'Pi', 'Ce', 'CZ', 'Mt'] + + # Each flux has a turbulent and neoclassical component + self.variables = [f'{i}_tr_turb' for i in self.quantities] + [f'{i}_tr_neo' for i in self.quantities] + + # Each flux component has a standard deviation + self.variables += [f'{i}_stds' for i in self.variables] + + # There is also turbulent exchange + self.variables += ['PexchTurb', 'PexchTurb_stds'] + + # And total transport flux + self.variables += [f'{i}_tr' for i in self.quantities] + + # Model results is None by default, but can be assigned in evaluate + self.model_results = None + + # Assign zeros to transport ones if not evaluated + for i in self.variables: + self.powerstate.plasma[i] = self.powerstate.plasma["te"] * 0.0 + + # There is also target components + self.variables += [f'{i}' for i in self.quantities] + [f'{i}_stds' for i in self.quantities] + + # ---------------------------------------------------------------------------------------- + # labels for plotting + # ---------------------------------------------------------------------------------------- + + self.powerstate.labelsFluxes = { + "te": "$Q_e$ ($MW/m^2$)", + "ti": "$Q_i$ ($MW/m^2$)", + "ne": ( + "$Q_{conv}$ ($MW/m^2$)" + if self.powerstate.TransportOptions["ModelOptions"].get("useConvectiveFluxes", True) + else "$\\Gamma_e$ ($10^{20}/s/m^2$)" + ), + "nZ": ( + "$Q_{conv}$ $\\cdot f_{Z,0}$ ($MW/m^2$)" + if self.powerstate.TransportOptions["ModelOptions"].get("useConvectiveFluxes", True) + else "$\\Gamma_Z$ $\\cdot f_{Z,0}$ ($10^{20}/s/m^2$)" + ), + "w0": "$M_T$ ($J/m^2$)", + } + + def produce_profiles(self): + # Only add self._produce_profiles() if it's needed (e.g. full TGLF), otherwise this is somewhat expensive (e.g. for flux matching) + pass + + def _produce_profiles(self,deriveQuantities=True): + + self.applyCorrections = self.powerstate.TransportOptions["ModelOptions"].get("MODELparameters", {}).get("applyCorrections", {}) + + # Write this updated profiles class (with parameterized profiles and target powers) + self.file_profs = self.folder / "input.gacode" + + powerstate_detached = self.powerstate.copy_state() + + self.powerstate.profiles = powerstate_detached.from_powerstate( + write_input_gacode=self.file_profs, + postprocess_input_gacode=self.applyCorrections, + rederive_profiles = deriveQuantities, # Derive quantities so that it's ready for analysis and plotting later + insert_highres_powers = deriveQuantities, # Insert powers so that Q, Pfus and all that it's consistent when read later + ) + + self.profiles_transport = copy.deepcopy(self.powerstate.profiles) + + self._modify_profiles() + + def _modify_profiles(self): + ''' + Modify the profiles (e.g. lumping) before running the transport model + ''' + + # After producing the profiles, copy for future modifications + self.file_profs_unmod = self.file_profs.parent / f"{self.file_profs.name}_unmodified" + shutil.copy2(self.file_profs, self.file_profs_unmod) + + profiles_postprocessing_fun = self.powerstate.TransportOptions["ModelOptions"].get("profiles_postprocessing_fun", None) + + if profiles_postprocessing_fun is not None: + print(f"\t- Modifying input.gacode to run transport calculations based on {profiles_postprocessing_fun}",typeMsg="i") + self.profiles_transport = profiles_postprocessing_fun(self.file_profs) + + # Position of impurity ion may have changed + p_old = PROFILEStools.PROFILES_GACODE(self.file_profs_unmod) + p_new = PROFILEStools.PROFILES_GACODE(self.file_profs) + + impurity_of_interest = p_old.Species[self.powerstate.impurityPosition] + + try: + impurityPosition_new = p_new.Species.index(impurity_of_interest) + + except ValueError: + print(f"\t- Impurity {impurity_of_interest} not found in new profiles, keeping position {self.powerstate.impurityPosition}",typeMsg="w") + impurityPosition_new = self.powerstate.impurityPosition + + if impurityPosition_new != self.powerstate.impurityPosition: + print(f"\t- Impurity position has changed from {self.powerstate.impurityPosition} to {impurityPosition_new}",typeMsg="w") + self.powerstate.impurityPosition_transport = p_new.Species.index(impurity_of_interest) + + # ---------------------------------------------------------------------------------------------------- + # EVALUATE (custom part) + # ---------------------------------------------------------------------------------------------------- + def evaluate(self): + ''' + This needs to populate the following in self.powerstate.plasma + - Pe, Pe_tr, Pe_tr_turb, Pe_tr_neo -> MW/m^2 + - Pi, Pi_tr, Pi_tr_turb, Pi_tr_neo -> MW/m^2 + - Ce, Ce_tr, Ce_tr_turb, Ce_tr_neo -> MW/m^2 + * Ce_raw, Ce_raw_tr, Ce_raw_tr_turb, Ce_raw_tr_neo -> 10^20/s/m^2 + - CZ, CZ_tr, CZ_tr_turb, CZ_tr_neo -> MW/m^2 (but modified as needed, for example dividing by fZ0) + * CZ_raw, CZ_raw_tr, CZ_raw_tr_turb, CZ_raw_tr_neo -> 10^20/s/m^2 (NOT modified) + - Mt, Mt_tr, Mt_tr_turb, Mt_tr_neo -> J/m^2 + - PexchTurb -> MW/m^3 + and their respective standard deviations + ''' + + print(">> No transport fluxes to evaluate", typeMsg="w") + pass + diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index c5aea661..592bf70d 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -5,7 +5,7 @@ import matplotlib.pyplot as plt from collections import OrderedDict from mitim_tools.misc_tools import GRAPHICStools, MATHtools, PLASMAtools, IOtools -from mitim_modules.powertorch.physics import GEOMETRYtools, CALCtools +from mitim_modules.powertorch.utils import GEOMETRYtools, CALCtools from mitim_tools.gs_tools import GEQtools from mitim_tools.gacode_tools import NEOtools from mitim_tools.gacode_tools.utils import GACODEdefaults @@ -1539,9 +1539,7 @@ def writeMiminalKinetic(self, file): valt = f"{val:.7e}".rjust(15) f.write(f"{pos}{valt}\n") - def changeResolution( - self, n=100, rho_new=None, interpolation_function=MATHtools.extrapolateCubicSpline - ): + def changeResolution(self, n=100, rho_new=None, interpolation_function=MATHtools.extrapolateCubicSpline): rho = copy.deepcopy(self.profiles["rho(-)"]) if rho_new is None: diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index ee6c47aa..44f2e9c4 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -3654,7 +3654,7 @@ def plotAnalysis(self, labels=["analysis1"], analysisType="chi_e", figs=None): rho_mod = np.append([0], rho) aLn = np.append([0], y) import torch - from mitim_modules.powertorch.physics import CALCtools + from mitim_modules.powertorch.utils import CALCtools BC = 1.0 T = CALCtools.integrateGradient( diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index ae1b5a6c..ac74c0ac 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -750,7 +750,7 @@ def findFileByExtension( f"\t\t\t~ Folder ...{fstr} does not exist, returning None", ) - # TODO: We really should not change return type + #TODO: We really should not change return type #retval = None #if retpath is not None: # if not provide_full_path: diff --git a/src/mitim_tools/misc_tools/PLASMAtools.py b/src/mitim_tools/misc_tools/PLASMAtools.py index 55c64ed9..e044737a 100644 --- a/src/mitim_tools/misc_tools/PLASMAtools.py +++ b/src/mitim_tools/misc_tools/PLASMAtools.py @@ -8,7 +8,7 @@ import matplotlib.pyplot as plt from IPython import embed from mitim_tools.misc_tools import MATHtools -from mitim_modules.powertorch.physics import CALCtools +from mitim_modules.powertorch.utils import CALCtools from mitim_tools.popcon_tools import FunctionalForms from mitim_tools.misc_tools.LOGtools import printMsg as print diff --git a/src/mitim_tools/popcon_tools/FunctionalForms.py b/src/mitim_tools/popcon_tools/FunctionalForms.py index 45b60748..f2d9cf7f 100644 --- a/src/mitim_tools/popcon_tools/FunctionalForms.py +++ b/src/mitim_tools/popcon_tools/FunctionalForms.py @@ -4,7 +4,7 @@ from mitim_tools.popcon_tools.utils import PRFfunctionals, FUNCTIONALScalc from mitim_tools.misc_tools import MATHtools, GRAPHICStools from mitim_tools.misc_tools.LOGtools import printMsg as print -from mitim_modules.powertorch.physics import CALCtools +from mitim_modules.powertorch.utils import CALCtools from IPython import embed diff --git a/src/mitim_tools/popcon_tools/scripts/test_functionals.py b/src/mitim_tools/popcon_tools/scripts/test_functionals.py index 949af0b4..7e8ceb21 100644 --- a/src/mitim_tools/popcon_tools/scripts/test_functionals.py +++ b/src/mitim_tools/popcon_tools/scripts/test_functionals.py @@ -1,6 +1,6 @@ import torch, datetime import matplotlib.pyplot as plt -from mitim_modules.powertorch.physics import CALCtools +from mitim_modules.powertorch.utils import CALCtools from mitim_tools.misc_tools import GRAPHICStools, IOtools, PLASMAtools from mitim_tools.popcon_tools.FunctionalForms import ( diff --git a/src/mitim_tools/popcon_tools/utils/FUNCTIONALScalc.py b/src/mitim_tools/popcon_tools/utils/FUNCTIONALScalc.py index 4758f47a..db8ec4f8 100644 --- a/src/mitim_tools/popcon_tools/utils/FUNCTIONALScalc.py +++ b/src/mitim_tools/popcon_tools/utils/FUNCTIONALScalc.py @@ -1,7 +1,7 @@ import torch import numpy as np from IPython import embed -from mitim_modules.powertorch.physics import CALCtools +from mitim_modules.powertorch.utils import CALCtools from mitim_tools.misc_tools.LOGtools import printMsg as print From 10e536aefa2dbad82ca35d04a0a9f726d3f6c5ba Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 17:21:16 -0700 Subject: [PATCH 009/385] tutorial fixes --- src/mitim_modules/portals/utils/PORTALSanalysis.py | 2 +- tutorials/PORTALS_tutorial.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 8e33cdd8..5d3af5af 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -961,7 +961,7 @@ def __init__(self, folder): for i in range(100): try: p = STATEtools.read_saved_state( - self.folder / "Initialization" / "initialization_simple_relax" / f"portals_sr_{IOtools.reducePathLevel(self.folder)[1]}_ev_{i}" / "powerstate.pkl" + self.folder / "Initialization" / "initialization_simple_relax" / f"portals_sr_ev_{i}" / "powerstate.pkl" ) except FileNotFoundError: break diff --git a/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index ea36678c..3e7983c1 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -8,7 +8,7 @@ # Starting input.gacode file inputgacode = __mitimroot__ / "tests" / "data" / "input.gacode" -folder = __mitimroot__ / "tests" / "scratch" / "portals_tut" +folder = __mitimroot__ / "tests" / "scratch" / "portals_tutorial" # Initialize PORTALS class portals_fun = PORTALSmain.portals(folder) @@ -20,8 +20,8 @@ portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti", "ne"] # Codes to use -from mitim_modules.powertorch.physics import TRANSPORTtools -portals_fun.PORTALSparameters["transport_evaluator"] = TRANSPORTtools.tgyro_model +from mitim_modules.powertorch.physics_models.transport_gacode import tgyro_model +portals_fun.PORTALSparameters["transport_evaluator"] = tgyro_model # TGLF specifications portals_fun.MODELparameters["transport_model"] = { From b401206a064a8be46ebd1af5facadf13b6912ebe Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 18:02:41 -0700 Subject: [PATCH 010/385] Separation of CGYRO model from TGYRO --- src/mitim_modules/portals/PORTALSmain.py | 12 +- src/mitim_modules/powertorch/STATEtools.py | 67 +++-- .../physics_models/transport_cgyro.py | 141 ++++++++++ ...transport_gacode.py => transport_tgyro.py} | 247 +++++------------- .../powertorch/scripts/calculateTargets.py | 4 +- tutorials/PORTALS_tutorial.py | 2 +- 6 files changed, 263 insertions(+), 210 deletions(-) create mode 100644 src/mitim_modules/powertorch/physics_models/transport_cgyro.py rename src/mitim_modules/powertorch/physics_models/{transport_gacode.py => transport_tgyro.py} (74%) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index cfcc5831..b05542ea 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -14,12 +14,11 @@ PORTALSoptimization, PORTALSanalysis, ) -from mitim_modules.powertorch.physics_models import targets_analytic, transport_gacode +from mitim_modules.powertorch.physics_models import targets_analytic, transport_tgyro, transport_cgyro from mitim_tools.opt_tools import STRATEGYtools from mitim_tools.opt_tools.utils import BOgraphics from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -from tqdm import tnrange """ @@ -194,7 +193,10 @@ def __init__( """ # Selection of model - transport_evaluator = transport_gacode.tgyro_model + if CGYROrun: + transport_evaluator = transport_cgyro.cgyro_model + else: + transport_evaluator = transport_tgyro.tgyro_model targets_evaluator = targets_analytic.analytical_model self.PORTALSparameters = { @@ -469,7 +471,7 @@ def check_flags(self): print("\t- Requested fineTargetsResolution, so running powerstate target calculations",typeMsg="w") self.PORTALSparameters["TargetCalc"] = "powerstate" - if not issubclass(self.PORTALSparameters["transport_evaluator"], transport_gacode.tgyro_model) and (self.PORTALSparameters["TargetCalc"] == "tgyro"): + if not issubclass(self.PORTALSparameters["transport_evaluator"], transport_tgyro.tgyro_model) and (self.PORTALSparameters["TargetCalc"] == "tgyro"): print("\t- Requested TGYRO targets, but transport evaluator is not tgyro, so changing to powerstate",typeMsg="w") self.PORTALSparameters["TargetCalc"] = "powerstate" @@ -556,7 +558,7 @@ def reuseTrainingTabular( self_copy.powerstate.TransportOptions["transport_evaluator"] = None self_copy.powerstate.TargetOptions["ModelOptions"]["TypeTarget"] = "powerstate" else: - self_copy.powerstate.TransportOptions["transport_evaluator"] = transport_gacode.tgyro_model + self_copy.powerstate.TransportOptions["transport_evaluator"] = transport_tgyro.tgyro_model _, dictOFs = runModelEvaluator( self_copy, diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index b25eec5e..ced8dd95 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -21,28 +21,16 @@ class powerstate: def __init__( self, - profiles, + profiles_object, increase_profile_resol=True, - EvolutionOptions={}, - TransportOptions={ - "transport_evaluator": None, - "ModelOptions": {} - }, - TargetOptions={ - "targets_evaluator": targets_analytic.analytical_model, - "ModelOptions": { - "TypeTarget": 3, - "TargetCalc": "powerstate" - }, - }, - tensor_opts = { - "dtype": torch.double, - "device": torch.device("cpu"), - } + EvolutionOptions=None, + TransportOptions=None, + TargetOptions=None, + tensor_opts=None, ): ''' Inputs: - - profiles: Object for PROFILES_GACODE or others + - profiles_object: Object for PROFILES_GACODE or others - EvolutionOptions: - rhoPredicted: radial grid (MUST NOT CONTAIN ZERO, it will be added internally) - ProfilesPredicted: list of profiles to predict @@ -53,6 +41,31 @@ def __init__( - TargetOptions: dictionary with targets_evaluator and ModelOptions ''' + if EvolutionOptions is None: + EvolutionOptions = {} + if TransportOptions is None: + TransportOptions = { + "transport_evaluator": None, + "ModelOptions": {} + } + if TargetOptions is None: + TargetOptions = { + "targets_evaluator": targets_analytic.analytical_model, + "ModelOptions": { + "TypeTarget": 3, + "TargetCalc": "powerstate" + }, + } + if tensor_opts is None: + tensor_opts = { + "dtype": torch.double, + "device": torch.device("cpu"), + } + + # ------------------------------------------------------------------------------------- + # Check inputs + # ------------------------------------------------------------------------------------- + print('>> Creating powerstate object...') self.TransportOptions = TransportOptions @@ -121,12 +134,12 @@ def _ensure_ne_before_nz(lst): # Object type (e.g. input.gacode) # ------------------------------------------------------------------------------------- - if isinstance(profiles, PROFILEStools.PROFILES_GACODE): + if isinstance(profiles_object, PROFILEStools.PROFILES_GACODE): self.to_powerstate = TRANSFORMtools.gacode_to_powerstate self.from_powerstate = MethodType(TRANSFORMtools.to_gacode, self) # Use a copy because I'm deriving, it may be expensive and I don't want to carry that out outside of this class - self.profiles = copy.deepcopy(profiles) + self.profiles = copy.deepcopy(profiles_object) if "derived" not in self.profiles.__dict__: self.profiles.deriveQuantities() @@ -203,8 +216,6 @@ def _fine_grid(self): # Revert plasma back self.plasma = plasma_copy - - # ------------------------------------------------------------------ # Storing and combining # ------------------------------------------------------------------ @@ -385,9 +396,12 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): Xpass = X[best_candidate, :].detach() # Store values - if y_history is not None: y_history.append(yRes) - if x_history is not None: x_history.append(Xpass) - if metric_history is not None: metric_history.append(yMetric) + if y_history is not None: + y_history.append(yRes) + if x_history is not None: + x_history.append(Xpass) + if metric_history is not None: + metric_history.append(yMetric) return QTransport, QTarget, yMetric @@ -403,8 +417,7 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): _,Yopt, Xopt, metric_history = solver_fun(evaluator,x0, bounds=self.bounds_current,solver_options=solver_options) # For simplicity, return the trajectory of only the best candidate - self.FluxMatch_Yopt = Yopt - self.FluxMatch_Xopt = Xopt + self.FluxMatch_Yopt, self.FluxMatch_Xopt = Yopt, Xopt print("**********************************************************************************************") print(f"\t- Flux matching of powerstate finished, and took {IOtools.getTimeDifference(timeBeginning)}\n") diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py new file mode 100644 index 00000000..07d93afa --- /dev/null +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -0,0 +1,141 @@ +import copy +import shutil +import torch +from mitim_tools.misc_tools import IOtools +from mitim_modules.portals.utils import PORTALScgyro +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed +from mitim_modules.powertorch.physics_models import transport_tgyro + +class cgyro_model(transport_tgyro.tgyro_model): + def __init__(self, powerstate, **kwargs): + super().__init__(powerstate, **kwargs) + + def produce_profiles(self): + super().produce_profiles() + + def evaluate(self): + + # --------------------------------------------------------- + # Run original evaluator + # --------------------------------------------------------- + + tgyro = self._evaluate_tglf_neo() + + # --------------------------------------------------------- + # Run cgyro_trick + # --------------------------------------------------------- + + powerstate_orig = self._evaluate_cgyro(tgyro) + + # --------------------------------------------------------- + # Process results + # --------------------------------------------------------- + + self._postprocess_results(tgyro, "cgyro_neo", powerstate_orig) + + def _evaluate_cgyro(self, tgyro): + + print("\t- Checking whether cgyro_neo folder exists and it was written correctly via cgyro_trick...") + + correctly_run = (self.folder / "cgyro_neo").exists() + if correctly_run: + print("\t\t- Folder exists, but was cgyro_trick run?") + with open(self.folder / "cgyro_neo" / "mitim_flag", "r") as f: + correctly_run = bool(float(f.readline())) + + if correctly_run: + print("\t\t\t* Yes, it was", typeMsg="w") + else: + print("\t\t\t* No, it was not, repating process", typeMsg="i") + + # Remove cgyro_neo folder + if (self.folder / "cgyro_neo").exists(): + IOtools.shutil_rmtree(self.folder / "cgyro_neo") + + # Copy tglf_neo results + shutil.copytree(self.folder / "tglf_neo", self.folder / "cgyro_neo") + + # CGYRO writter + cgyro_trick(self,self.folder / "cgyro_neo") + + # Read TGYRO files and construct portals variables + + tgyro.read(label="cgyro_neo", folder=self.folder / "cgyro_neo") + + powerstate_orig = copy.deepcopy(self.powerstate) + + return powerstate_orig + + def _postprocess_results(self, tgyro, label, powerstate_orig): + super()._postprocess_results(tgyro, label) + + # ------------------------------------------------------------ + # Some checks + # ------------------------------------------------------------ + + print("\t- Checking model modifications:") + for r in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: #, "PexchTurb"]: #TODO: FIX + print(f"\t\t{r}(tglf) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(powerstate_orig.plasma[r][0][1:],powerstate_orig.plasma[r+'_stds'][0][1:]) ])}") + print(f"\t\t{r}(cgyro) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(self.powerstate.plasma[r][0][1:],self.powerstate.plasma[r+'_stds'][0][1:]) ])}") + +def cgyro_trick(self,FolderEvaluation_TGYRO): + + with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: + f.write("0") + + # ************************************************************************************************************************** + # Print Information + # ************************************************************************************************************************** + + txt = "\nFluxes to be matched by CGYRO ( TARGETS - NEO ):" + + for var, varn in zip( + ["r/a ", "rho ", "a/LTe", "a/LTi", "a/Lne", "a/LnZ", "a/Lw0"], + ["roa", "rho", "aLte", "aLti", "aLne", "aLnZ", "aLw0"], + ): + txt += f"\n{var} = " + for j in range(self.powerstate.plasma["rho"].shape[1] - 1): + txt += f"{self.powerstate.plasma[varn][0,j+1]:.6f} " + + for var, varn in zip( + ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ce (MW/m^2)", "CZ (MW/m^2)", "Mt (J/m^2) "], + ["Pe", "Pi", "Ce", "CZ", "Mt"], + ): + txt += f"\n{var} = " + for j in range(self.powerstate.plasma["rho"].shape[1] - 1): + txt += f"{self.powerstate.plasma[varn][0,j+1]-self.powerstate.plasma[f'{varn}_tr_neo'][0,j+1]:.4e} " + + print(txt) + + # Copy profiles so that later it is easy to grab all the input.gacodes that were evaluated + self._profiles_to_store() + + # ************************************************************************************************************************** + # Evaluate CGYRO + # ************************************************************************************************************************** + + PORTALScgyro.evaluateCGYRO( + self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["PORTALSparameters"], + self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"], + self.evaluation_number, + FolderEvaluation_TGYRO, + self.file_profs, + self.powerstate.plasma["roa"][0,1:], + self.powerstate.ProfilesPredicted, + ) + + # ************************************************************************************************************************** + # EXTRA + # ************************************************************************************************************************** + + # Make tensors + for i in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: + try: + self.powerstate.plasma[i] = torch.from_numpy(self.powerstate.plasma[i]).to(self.powerstate.dfT).unsqueeze(0) + except: + pass + + # Write a flag indicating this was performed, to avoid an issue that... the script crashes when it has copied tglf_neo, without cgyro_trick modification + with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: + f.write("1") diff --git a/src/mitim_modules/powertorch/physics_models/transport_gacode.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py similarity index 74% rename from src/mitim_modules/powertorch/physics_models/transport_gacode.py rename to src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 57166e5c..2cba83a1 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_gacode.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -1,6 +1,5 @@ import copy import shutil -import torch import numpy as np from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import TGYROtools @@ -22,6 +21,75 @@ def produce_profiles(self): def evaluate(self): + tgyro = self._evaluate_tglf_neo() + + self._postprocess_results(tgyro, "tglf_neo") + + # ************************************************************************************ + # Private functions for TGLF and NEO evaluations + # ************************************************************************************ + + def _profiles_to_store(self): + + if "extra_params" in self.powerstate.TransportOptions["ModelOptions"] and "folder" in self.powerstate.TransportOptions["ModelOptions"]["extra_params"]: + whereFolder = IOtools.expandPath(self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") + if not whereFolder.exists(): + IOtools.askNewFolder(whereFolder) + + fil = whereFolder / f"input.gacode.{self.evaluation_number}" + shutil.copy2(self.file_profs, fil) + shutil.copy2(self.file_profs_unmod, fil.parent / f"{fil.name}_unmodified") + shutil.copy2(self.file_profs_targets, fil.parent / f"{fil.name}.new") + print(f"\t- Copied profiles to {IOtools.clipstr(fil)}") + else: + print("\t- Could not move files", typeMsg="w") + + + def _postprocess_results(self, tgyro, label): + + ModelOptions = self.powerstate.TransportOptions["ModelOptions"] + + includeFast = ModelOptions.get("includeFastInQi",False) + useConvectiveFluxes = ModelOptions.get("useConvectiveFluxes", True) + UseFineGridTargets = ModelOptions.get("UseFineGridTargets", False) + provideTurbulentExchange = ModelOptions.get("TurbulentExchange", False) + OriginalFimp = ModelOptions.get("OriginalFimp", 1.0) + forceZeroParticleFlux = ModelOptions.get("forceZeroParticleFlux", False) + + # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) + impurityPosition = self.powerstate.impurityPosition_transport #ModelOptions.get("impurityPosition", 1) + + # Produce right quantities (TGYRO -> powerstate.plasma) + self.powerstate = tgyro.results[label].TGYROmodeledVariables( + self.powerstate, + useConvectiveFluxes=useConvectiveFluxes, + includeFast=includeFast, + impurityPosition=impurityPosition, + UseFineGridTargets=UseFineGridTargets, + OriginalFimp=OriginalFimp, + forceZeroParticleFlux=forceZeroParticleFlux, + provideTurbulentExchange=provideTurbulentExchange, + provideTargets=self.powerstate.TargetOptions['ModelOptions']['TargetCalc'] == "tgyro", + ) + + tgyro.results["use"] = tgyro.results[label] + + # Copy profiles to share + self._profiles_to_store() + + # ------------------------------------------------------------------------------------------------------------------------ + # Results class that can be used for further plotting and analysis in PORTALS + # ------------------------------------------------------------------------------------------------------------------------ + + self.model_results = copy.deepcopy(tgyro.results["use"]) # Pass the TGYRO results class that should be use for plotting and analysis + + self.model_results.extra_analysis = {} + for ikey in tgyro.results: + if ikey != "use": + self.model_results.extra_analysis[ikey] = tgyro.results[ikey] + + def _evaluate_tglf_neo(self): + # ------------------------------------------------------------------------------------------------------------------------ # Model Options # ------------------------------------------------------------------------------------------------------------------------ @@ -30,13 +98,9 @@ def evaluate(self): MODELparameters = ModelOptions.get("MODELparameters",None) includeFast = ModelOptions.get("includeFastInQi",False) - useConvectiveFluxes = ModelOptions.get("useConvectiveFluxes", True) - UseFineGridTargets = ModelOptions.get("UseFineGridTargets", False) launchMODELviaSlurm = ModelOptions.get("launchMODELviaSlurm", False) cold_start = ModelOptions.get("cold_start", False) provideTurbulentExchange = ModelOptions.get("TurbulentExchange", False) - OriginalFimp = ModelOptions.get("OriginalFimp", 1.0) - forceZeroParticleFlux = ModelOptions.get("forceZeroParticleFlux", False) percentError = ModelOptions.get("percentError", [5, 1, 0.5]) use_tglf_scan_trick = ModelOptions.get("use_tglf_scan_trick", None) cores_per_tglf_instance = ModelOptions.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) @@ -45,7 +109,7 @@ def evaluate(self): impurityPosition = self.powerstate.impurityPosition_transport #ModelOptions.get("impurityPosition", 1) # ------------------------------------------------------------------------------------------------------------------------ - # 1. tglf_neo_original: Run TGYRO workflow - TGLF + NEO in subfolder tglf_neo_original (original as in... without stds or merging) + # tglf_neo_original: Run TGYRO workflow - TGLF + NEO in subfolder tglf_neo_original (original as in... without stds or merging) # ------------------------------------------------------------------------------------------------------------------------ RadiisToRun = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] @@ -83,7 +147,7 @@ def evaluate(self): self.file_profs_targets = tgyro.FolderTGYRO / "input.gacode.new" # ------------------------------------------------------------------------------------------------------------------------ - # 2. tglf_neo: Write TGLF, NEO and TARGET errors in tgyro files as well + # tglf_neo: Write TGLF, NEO and TARGET errors in tgyro files as well # ------------------------------------------------------------------------------------------------------------------------ # Copy original TGYRO folder @@ -138,95 +202,7 @@ def evaluate(self): # $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ # $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ - # ------------------------------------------------------------------------------------------------------------------------ - # 3. tglf_neo: Populate powerstate with the TGYRO results - # ------------------------------------------------------------------------------------------------------------------------ - - # Produce right quantities (TGYRO -> powerstate.plasma) - self.powerstate = tgyro.results["tglf_neo"].TGYROmodeledVariables( - self.powerstate, - useConvectiveFluxes=useConvectiveFluxes, - includeFast=includeFast, - impurityPosition=impurityPosition, - UseFineGridTargets=UseFineGridTargets, - OriginalFimp=OriginalFimp, - forceZeroParticleFlux=forceZeroParticleFlux, - provideTurbulentExchange=provideTurbulentExchange, - provideTargets=self.powerstate.TargetOptions['ModelOptions']['TargetCalc'] == "tgyro", - ) - - # ------------------------------------------------------------------------------------------------------------------------ - # 4. cgyro_neo: Trick to fake a tgyro output to reflect CGYRO - # ------------------------------------------------------------------------------------------------------------------------ - - if MODELparameters['transport_model']['turbulence'] == 'CGYRO': - - print("\t- Checking whether cgyro_neo folder exists and it was written correctly via cgyro_trick...") - - correctly_run = (self.folder / "cgyro_neo").exists() - if correctly_run: - print("\t\t- Folder exists, but was cgyro_trick run?") - with open(self.folder / "cgyro_neo" / "mitim_flag", "r") as f: - correctly_run = bool(float(f.readline())) - - if correctly_run: - print("\t\t\t* Yes, it was", typeMsg="w") - else: - print("\t\t\t* No, it was not, repating process", typeMsg="i") - - # Remove cgyro_neo folder - if (self.folder / "cgyro_neo").exists(): - IOtools.shutil_rmtree(self.folder / "cgyro_neo") - - # Copy tglf_neo results - shutil.copytree(self.folder / "tglf_neo", self.folder / "cgyro_neo") - - # CGYRO writter - cgyro_trick(self,self.folder / "cgyro_neo") - - # Read TGYRO files and construct portals variables - - tgyro.read(label="cgyro_neo", folder=self.folder / "cgyro_neo") - - powerstate_orig = copy.deepcopy(self.powerstate) - - self.powerstate = tgyro.results["cgyro_neo"].TGYROmodeledVariables( - self.powerstate, - useConvectiveFluxes=useConvectiveFluxes, - includeFast=includeFast, - impurityPosition=impurityPosition, - UseFineGridTargets=UseFineGridTargets, - OriginalFimp=OriginalFimp, - forceZeroParticleFlux=forceZeroParticleFlux, - provideTurbulentExchange=provideTurbulentExchange, - provideTargets=self.powerstate.TargetOptions['ModelOptions']['TargetCalc'] == "tgyro", - ) - - print("\t- Checking model modifications:") - for r in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: #, "PexchTurb"]: #TODO: FIX - print(f"\t\t{r}(tglf) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(powerstate_orig.plasma[r][0][1:],powerstate_orig.plasma[r+'_stds'][0][1:]) ])}") - print(f"\t\t{r}(cgyro) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(self.powerstate.plasma[r][0][1:],self.powerstate.plasma[r+'_stds'][0][1:]) ])}") - - # ** - tgyro.results["use"] = tgyro.results["cgyro_neo"] - - else: - # copy profiles too! - profilesToShare(self) - - # ** - tgyro.results["use"] = tgyro.results["tglf_neo"] - - # ------------------------------------------------------------------------------------------------------------------------ - # Results class that can be used for further plotting and analysis in PORTALS - # ------------------------------------------------------------------------------------------------------------------------ - - self.model_results = copy.deepcopy(tgyro.results["use"]) # Pass the TGYRO results class that should be use for plotting and analysis - - self.model_results.extra_analysis = {} - for ikey in tgyro.results: - if ikey != "use": - self.model_results.extra_analysis[ikey] = tgyro.results[ikey] + return tgyro def tglf_scan_trick( fluxesTGYRO, @@ -377,7 +353,6 @@ def calculate_mean_std(Q): return Qe_point, Qi_point, Ge_point, GZ_point, Mt_point, Pexch_point, Qe_std, Qi_std, Ge_std, GZ_std, Mt_std, Pexch_std - # ************************************************************************************************** # Functions # ************************************************************************************************** @@ -558,83 +533,6 @@ def curateTGYROfiles( special_label="_stds", ) - -def profilesToShare(self): - if "extra_params" in self.powerstate.TransportOptions["ModelOptions"] and "folder" in self.powerstate.TransportOptions["ModelOptions"]["extra_params"]: - whereFolder = IOtools.expandPath(self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") - if not whereFolder.exists(): - IOtools.askNewFolder(whereFolder) - - fil = whereFolder / f"input.gacode.{self.evaluation_number}" - shutil.copy2(self.file_profs, fil) - shutil.copy2(self.file_profs_unmod, fil.parent / f"{fil.name}_unmodified") - shutil.copy2(self.file_profs_targets, fil.parent / f"{fil.name}.new") - print(f"\t- Copied profiles to {IOtools.clipstr(fil)}") - else: - print("\t- Could not move files", typeMsg="w") - - -def cgyro_trick(self,FolderEvaluation_TGYRO): - - with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: - f.write("0") - - # ************************************************************************************************************************** - # Print Information - # ************************************************************************************************************************** - - txt = "\nFluxes to be matched by CGYRO ( TARGETS - NEO ):" - - for var, varn in zip( - ["r/a ", "rho ", "a/LTe", "a/LTi", "a/Lne", "a/LnZ", "a/Lw0"], - ["roa", "rho", "aLte", "aLti", "aLne", "aLnZ", "aLw0"], - ): - txt += f"\n{var} = " - for j in range(self.powerstate.plasma["rho"].shape[1] - 1): - txt += f"{self.powerstate.plasma[varn][0,j+1]:.6f} " - - for var, varn in zip( - ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ce (MW/m^2)", "CZ (MW/m^2)", "Mt (J/m^2) "], - ["Pe", "Pi", "Ce", "CZ", "Mt"], - ): - txt += f"\n{var} = " - for j in range(self.powerstate.plasma["rho"].shape[1] - 1): - txt += f"{self.powerstate.plasma[varn][0,j+1]-self.powerstate.plasma[f'{varn}_tr_neo'][0,j+1]:.4e} " - - print(txt) - - # Copy profiles so that later it is easy to grab all the input.gacodes that were evaluated - profilesToShare(self) - - # ************************************************************************************************************************** - # Evaluate CGYRO - # ************************************************************************************************************************** - - PORTALScgyro.evaluateCGYRO( - self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["PORTALSparameters"], - self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"], - self.evaluation_number, - FolderEvaluation_TGYRO, - self.file_profs, - self.powerstate.plasma["roa"][0,1:], - self.powerstate.ProfilesPredicted, - ) - - # ************************************************************************************************************************** - # EXTRA - # ************************************************************************************************************************** - - # Make tensors - for i in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: - try: - self.powerstate.plasma[i] = torch.from_numpy(self.powerstate.plasma[i]).to(self.powerstate.dfT).unsqueeze(0) - except: - pass - - # Write a flag indicating this was performed, to avoid an issue that... the script crashes when it has copied tglf_neo, without cgyro_trick modification - with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: - f.write("1") - def dummyCDF(GeneralFolder, FolderEvaluation): """ This routine creates path to a dummy CDF file in FolderEvaluation, with the name "simulation_evaluation.CDF" @@ -659,4 +557,3 @@ def dummyCDF(GeneralFolder, FolderEvaluation): cdf = FolderEvaluation / f"{subname}_ev{name}.CDF" return cdf - diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index ae7cb7ed..2889ff87 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -7,7 +7,7 @@ from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.powertorch import STATEtools -from mitim_modules.powertorch.physics_models import targets_analytic, transport_gacode +from mitim_modules.powertorch.physics_models import targets_analytic, transport_tgyro from IPython import embed def calculator( @@ -39,7 +39,7 @@ def calculator( "TargetCalc": "tgyro"}, }, TransportOptions={ - "transport_evaluator": transport_gacode.tgyro_model, + "transport_evaluator": transport_tgyro.tgyro_model, "ModelOptions": { "cold_start": cold_start, "launchSlurm": True, diff --git a/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index 3e7983c1..b23ccfb1 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -20,7 +20,7 @@ portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti", "ne"] # Codes to use -from mitim_modules.powertorch.physics_models.transport_gacode import tgyro_model +from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model portals_fun.PORTALSparameters["transport_evaluator"] = tgyro_model # TGLF specifications From 7b8b21b1a51d40e3e10765275995837c9983f776 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 18:05:53 -0700 Subject: [PATCH 011/385] removed references to "turbulence" --- src/mitim_modules/maestro/tmp_tests/maestro_test1.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 3 ++- src/mitim_modules/powertorch/physics_models/transport_cgyro.py | 2 +- src/mitim_modules/powertorch/scripts/calculateTargets.py | 2 +- templates/maestro_namelist.json | 1 - tutorials/PORTALS_tutorial.py | 1 - 6 files changed, 5 insertions(+), 6 deletions(-) diff --git a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py b/src/mitim_modules/maestro/tmp_tests/maestro_test1.py index f155818a..e1018d28 100644 --- a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py +++ b/src/mitim_modules/maestro/tmp_tests/maestro_test1.py @@ -46,7 +46,7 @@ "MODELparameters": { "RoaLocations": [0.35,0.55,0.75,0.875,0.9], "ProfilesPredicted": ["te", "ti", "ne"], "Physics_options": {"TypeTarget": 3}, - "transport_model": {"turbulence":'TGLF',"TGLFsettings": 6, "extraOptionsTGLF": {'USE_BPER':True}}}, + "transport_model": {"TGLFsettings": 6, "extraOptionsTGLF": {'USE_BPER':True}}}, "INITparameters": {"FastIsThermal": True, "removeIons": [5,6], "quasineutrality": True}, "optimization_options": { "convergence_options": { diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index b05542ea..2427258a 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -171,7 +171,7 @@ def __init__( "Tfast_ratio": False, # Keep the ratio of Tfast/Te constant throughout the Te evolution "ensureMachNumber": None, # Change w0 to match this Mach number when Ti varies }, - "transport_model": {"turbulence":'TGLF',"TGLFsettings": 6, "extraOptionsTGLF": {}} + "transport_model": {"TGLFsettings": 6, "extraOptionsTGLF": {}} } for key in self.MODELparameters.keys(): @@ -197,6 +197,7 @@ def __init__( transport_evaluator = transport_cgyro.cgyro_model else: transport_evaluator = transport_tgyro.tgyro_model + targets_evaluator = targets_analytic.analytical_model self.PORTALSparameters = { diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 07d93afa..cef77d06 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -3,9 +3,9 @@ import torch from mitim_tools.misc_tools import IOtools from mitim_modules.portals.utils import PORTALScgyro +from mitim_modules.powertorch.physics_models import transport_tgyro from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -from mitim_modules.powertorch.physics_models import transport_tgyro class cgyro_model(transport_tgyro.tgyro_model): def __init__(self, powerstate, **kwargs): diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 2889ff87..271fc2c9 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -59,7 +59,7 @@ def calculator( "ni_thermals": True, "recompute_ptot": False, }, - "transport_model": {"turbulence": 'TGLF',"TGLFsettings": 5, "extraOptionsTGLF": {}}, + "transport_model": {"TGLFsettings": 5, "extraOptionsTGLF": {}}, }, "includeFastInQi": False, }, diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index 70dcb052..ca8f857f 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -70,7 +70,6 @@ "RoaLocations": [0.35,0.45,0.55,0.65,0.75,0.875,0.9], "Physics_options": {"TypeTarget": 3}, "transport_model": { - "turbulence": "TGLF", "TGLFsettings": 100, "extraOptionsTGLF": {"USE_BPER": true} } diff --git a/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index b23ccfb1..9d4e0f82 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -25,7 +25,6 @@ # TGLF specifications portals_fun.MODELparameters["transport_model"] = { - "turbulence":'TGLF', "TGLFsettings": 6, # Check out templates/input.tglf.models.json for more options "extraOptionsTGLF": {"USE_BPER": False} # Turn off BPER } From edb79f04cd9c9d40e0fc973fef5c334b4348ba87 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 21:22:35 -0700 Subject: [PATCH 012/385] CGYRO model better organization --- .../physics_models/transport_cgyro.py | 37 ++---- .../physics_models/transport_tgyro.py | 121 +++++++++--------- 2 files changed, 71 insertions(+), 87 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index cef77d06..c3a1421e 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -16,25 +16,26 @@ def produce_profiles(self): def evaluate(self): - # --------------------------------------------------------- # Run original evaluator - # --------------------------------------------------------- - tgyro = self._evaluate_tglf_neo() - # --------------------------------------------------------- - # Run cgyro_trick - # --------------------------------------------------------- - - powerstate_orig = self._evaluate_cgyro(tgyro) + # Run CGYRO trick + powerstate_orig = self._trick_cgyro(tgyro) - # --------------------------------------------------------- # Process results - # --------------------------------------------------------- + self._postprocess_results(tgyro, "cgyro_neo") + + # Some checks + print("\t- Checking model modifications:") + for r in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: #, "PexchTurb"]: #TODO: FIX + print(f"\t\t{r}(tglf) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(powerstate_orig.plasma[r][0][1:],powerstate_orig.plasma[r+'_stds'][0][1:]) ])}") + print(f"\t\t{r}(cgyro) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(self.powerstate.plasma[r][0][1:],self.powerstate.plasma[r+'_stds'][0][1:]) ])}") - self._postprocess_results(tgyro, "cgyro_neo", powerstate_orig) + # ************************************************************************************ + # Private functions for CGYRO evaluation + # ************************************************************************************ - def _evaluate_cgyro(self, tgyro): + def _trick_cgyro(self, tgyro): print("\t- Checking whether cgyro_neo folder exists and it was written correctly via cgyro_trick...") @@ -67,18 +68,6 @@ def _evaluate_cgyro(self, tgyro): return powerstate_orig - def _postprocess_results(self, tgyro, label, powerstate_orig): - super()._postprocess_results(tgyro, label) - - # ------------------------------------------------------------ - # Some checks - # ------------------------------------------------------------ - - print("\t- Checking model modifications:") - for r in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: #, "PexchTurb"]: #TODO: FIX - print(f"\t\t{r}(tglf) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(powerstate_orig.plasma[r][0][1:],powerstate_orig.plasma[r+'_stds'][0][1:]) ])}") - print(f"\t\t{r}(cgyro) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(self.powerstate.plasma[r][0][1:],self.powerstate.plasma[r+'_stds'][0][1:]) ])}") - def cgyro_trick(self,FolderEvaluation_TGYRO): with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 2cba83a1..2a8d3cf9 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -29,71 +29,8 @@ def evaluate(self): # Private functions for TGLF and NEO evaluations # ************************************************************************************ - def _profiles_to_store(self): - - if "extra_params" in self.powerstate.TransportOptions["ModelOptions"] and "folder" in self.powerstate.TransportOptions["ModelOptions"]["extra_params"]: - whereFolder = IOtools.expandPath(self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") - if not whereFolder.exists(): - IOtools.askNewFolder(whereFolder) - - fil = whereFolder / f"input.gacode.{self.evaluation_number}" - shutil.copy2(self.file_profs, fil) - shutil.copy2(self.file_profs_unmod, fil.parent / f"{fil.name}_unmodified") - shutil.copy2(self.file_profs_targets, fil.parent / f"{fil.name}.new") - print(f"\t- Copied profiles to {IOtools.clipstr(fil)}") - else: - print("\t- Could not move files", typeMsg="w") - - - def _postprocess_results(self, tgyro, label): - - ModelOptions = self.powerstate.TransportOptions["ModelOptions"] - - includeFast = ModelOptions.get("includeFastInQi",False) - useConvectiveFluxes = ModelOptions.get("useConvectiveFluxes", True) - UseFineGridTargets = ModelOptions.get("UseFineGridTargets", False) - provideTurbulentExchange = ModelOptions.get("TurbulentExchange", False) - OriginalFimp = ModelOptions.get("OriginalFimp", 1.0) - forceZeroParticleFlux = ModelOptions.get("forceZeroParticleFlux", False) - - # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) - impurityPosition = self.powerstate.impurityPosition_transport #ModelOptions.get("impurityPosition", 1) - - # Produce right quantities (TGYRO -> powerstate.plasma) - self.powerstate = tgyro.results[label].TGYROmodeledVariables( - self.powerstate, - useConvectiveFluxes=useConvectiveFluxes, - includeFast=includeFast, - impurityPosition=impurityPosition, - UseFineGridTargets=UseFineGridTargets, - OriginalFimp=OriginalFimp, - forceZeroParticleFlux=forceZeroParticleFlux, - provideTurbulentExchange=provideTurbulentExchange, - provideTargets=self.powerstate.TargetOptions['ModelOptions']['TargetCalc'] == "tgyro", - ) - - tgyro.results["use"] = tgyro.results[label] - - # Copy profiles to share - self._profiles_to_store() - - # ------------------------------------------------------------------------------------------------------------------------ - # Results class that can be used for further plotting and analysis in PORTALS - # ------------------------------------------------------------------------------------------------------------------------ - - self.model_results = copy.deepcopy(tgyro.results["use"]) # Pass the TGYRO results class that should be use for plotting and analysis - - self.model_results.extra_analysis = {} - for ikey in tgyro.results: - if ikey != "use": - self.model_results.extra_analysis[ikey] = tgyro.results[ikey] - def _evaluate_tglf_neo(self): - # ------------------------------------------------------------------------------------------------------------------------ - # Model Options - # ------------------------------------------------------------------------------------------------------------------------ - ModelOptions = self.powerstate.TransportOptions["ModelOptions"] MODELparameters = ModelOptions.get("MODELparameters",None) @@ -204,6 +141,64 @@ def _evaluate_tglf_neo(self): return tgyro + def _postprocess_results(self, tgyro, label): + + ModelOptions = self.powerstate.TransportOptions["ModelOptions"] + + includeFast = ModelOptions.get("includeFastInQi",False) + useConvectiveFluxes = ModelOptions.get("useConvectiveFluxes", True) + UseFineGridTargets = ModelOptions.get("UseFineGridTargets", False) + provideTurbulentExchange = ModelOptions.get("TurbulentExchange", False) + OriginalFimp = ModelOptions.get("OriginalFimp", 1.0) + forceZeroParticleFlux = ModelOptions.get("forceZeroParticleFlux", False) + + # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) + impurityPosition = self.powerstate.impurityPosition_transport #ModelOptions.get("impurityPosition", 1) + + # Produce right quantities (TGYRO -> powerstate.plasma) + self.powerstate = tgyro.results[label].TGYROmodeledVariables( + self.powerstate, + useConvectiveFluxes=useConvectiveFluxes, + includeFast=includeFast, + impurityPosition=impurityPosition, + UseFineGridTargets=UseFineGridTargets, + OriginalFimp=OriginalFimp, + forceZeroParticleFlux=forceZeroParticleFlux, + provideTurbulentExchange=provideTurbulentExchange, + provideTargets=self.powerstate.TargetOptions['ModelOptions']['TargetCalc'] == "tgyro", + ) + + tgyro.results["use"] = tgyro.results[label] + + # Copy profiles to share + self._profiles_to_store() + + # ------------------------------------------------------------------------------------------------------------------------ + # Results class that can be used for further plotting and analysis in PORTALS + # ------------------------------------------------------------------------------------------------------------------------ + + self.model_results = copy.deepcopy(tgyro.results["use"]) # Pass the TGYRO results class that should be use for plotting and analysis + + self.model_results.extra_analysis = {} + for ikey in tgyro.results: + if ikey != "use": + self.model_results.extra_analysis[ikey] = tgyro.results[ikey] + + def _profiles_to_store(self): + + if "extra_params" in self.powerstate.TransportOptions["ModelOptions"] and "folder" in self.powerstate.TransportOptions["ModelOptions"]["extra_params"]: + whereFolder = IOtools.expandPath(self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") + if not whereFolder.exists(): + IOtools.askNewFolder(whereFolder) + + fil = whereFolder / f"input.gacode.{self.evaluation_number}" + shutil.copy2(self.file_profs, fil) + shutil.copy2(self.file_profs_unmod, fil.parent / f"{fil.name}_unmodified") + shutil.copy2(self.file_profs_targets, fil.parent / f"{fil.name}.new") + print(f"\t- Copied profiles to {IOtools.clipstr(fil)}") + else: + print("\t- Could not move files", typeMsg="w") + def tglf_scan_trick( fluxesTGYRO, tgyro, From 767676b19007cd0afa50e4fbf7e6e6f5ebdd0131 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Apr 2025 22:54:13 -0700 Subject: [PATCH 013/385] misc appearance --- .../portals/utils/PORTALScgyro.py | 135 ++---------------- 1 file changed, 11 insertions(+), 124 deletions(-) diff --git a/src/mitim_modules/portals/utils/PORTALScgyro.py b/src/mitim_modules/portals/utils/PORTALScgyro.py index 46f4e09a..4dc32b48 100644 --- a/src/mitim_modules/portals/utils/PORTALScgyro.py +++ b/src/mitim_modules/portals/utils/PORTALScgyro.py @@ -17,7 +17,6 @@ The CGYRO file must use particle flux. Convective transformation occurs later """ - def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmodified_profiles, radii, ProfilesPredicted): print("\n ** CGYRO evaluation of fluxes has been requested before passing information to the STRATEGY module **",typeMsg="i",) @@ -125,41 +124,17 @@ def cgyroing( evaluations = np.array([int(i) for i in evaluations.split(",")]) evaluationsInFile = np.array([int(i) for i in evaluationsInFile.split(",")]) - ( - aLTe, - aLTi, - aLne, - Q_gb, - Qe, - Qi, - Ge, - GZ, - Mt, - Pexch, - QeE, - QiE, - GeE, - GZE, - MtE, - PexchE, - _, - _, - ) = readCGYROresults(file, radii) + aLTe,aLTi,aLne,Q_gb,Qe,Qi,Ge,GZ,Mt,Pexch,QeE,QiE,GeE,GZE,MtE,PexchE,_,_ = readCGYROresults(file, radii) cont = 0 - for i in evaluations: + for _ in evaluations: k = evaluationsInFile[cont] cont += 1 - print( - f"\t- Modifying {IOtools.clipstr(FolderEvaluation)} with position {k} in CGYRO results file {IOtools.clipstr(file)}" - ) + print(f"\t- Modifying {IOtools.clipstr(FolderEvaluation)} with position {k} in CGYRO results file {IOtools.clipstr(file)}") # Get TGYRO - tgyro = TGYROtools.TGYROoutput( - FolderEvaluation, - profiles=PROFILEStools.PROFILES_GACODE(unmodified_profiles), - ) + tgyro = TGYROtools.TGYROoutput(FolderEvaluation,profiles=PROFILEStools.PROFILES_GACODE(unmodified_profiles)) # Quick checker of correct file wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro) @@ -306,49 +281,9 @@ def readlineNTH(line, full_file=True, unnormalize=True): PexchReal = Pexch PexchReal_std = Pexch_std - return ( - roa, - aLTe, - aLTi, - aLne, - Q_gb, - QeReal, - QiReal, - GeReal, - GZReal, - MtReal, - PexchReal, - QeReal_std, - QiReal_std, - GeReal_std, - GZReal_std, - MtReal_std, - PexchReal_std, - tstart, - tend, - ) + return roa,aLTe,aLTi,aLne,Q_gb,QeReal,QiReal,GeReal,GZReal,MtReal,PexchReal,QeReal_std,QiReal_std,GeReal_std,GZReal_std,MtReal_std,PexchReal_std,tstart,tend else: - return ( - roa, - aLTe, - aLTi, - aLne, - Q_gb, - QeReal, - QiReal, - GeReal, - 0.0, - 0.0, - 0.0, - QeReal_std, - QiReal_std, - GeReal_std, - 0.0, - 0.0, - 0.0, - tstart, - tend, - ) + return roa,aLTe,aLTi,aLne,Q_gb,QeReal,QiReal,GeReal,0.0,0.0,0.0,QeReal_std,QiReal_std,GeReal_std,0.0,0.0,0.0,tstart,tend def readCGYROresults(file, radii, unnormalize=True): @@ -434,47 +369,7 @@ def readCGYROresults(file, radii, unnormalize=True): # Assign to that radial location # -------------------------------------------------------- - ( - roa[p[r], r], - aLTe[p[r], r], - aLTi[p[r], r], - aLne[p[r], r], - Q_gb[p[r], r], - Qe[p[r], r], - Qi[p[r], r], - Ge[p[r], r], - GZ[p[r], r], - Mt[p[r], r], - Pexch[p[r], r], - Qe_std[p[r], r], - Qi_std[p[r], r], - Ge_std[p[r], r], - GZ_std[p[r], r], - Mt_std[p[r], r], - Pexch_std[p[r], r], - tstart[p[r], r], - tend[p[r], r], - ) = ( - roa_read, - aLTe_read, - aLTi_read, - aLne_read, - Q_gb_read, - Qe_read, - Qi_read, - Ge_read, - GZ_read, - Mt_read, - Pexch_read, - Qe_std_read, - Qi_std_read, - Ge_std_read, - GZ_std_read, - Mt_std_read, - Pexch_std_read, - tstart_read, - tend_read, - ) + roa[p[r], r],aLTe[p[r], r],aLTi[p[r], r],aLne[p[r], r],Q_gb[p[r], r],Qe[p[r], r],Qi[p[r], r],Ge[p[r], r],GZ[p[r], r],Mt[p[r], r],Pexch[p[r], r],Qe_std[p[r], r],Qi_std[p[r], r],Ge_std[p[r], r],GZ_std[p[r], r],Mt_std[p[r], r],Pexch_std[p[r], r],tstart[p[r], r],tend[p[r], r] = roa_read,aLTe_read,aLTi_read,aLne_read,Q_gb_read,Qe_read,Qi_read,Ge_read,GZ_read,Mt_read,Pexch_read,Qe_std_read,Qi_std_read,Ge_std_read,GZ_std_read,Mt_std_read,Pexch_std_read,tstart_read,tend_read p[r] += 1 @@ -671,18 +566,10 @@ def modifyEVO( GZTGB = GZT / tgyro.Gamma_GB[-1, 1:] MtTGB = MtT / tgyro.Pi_GB[-1, 1:] - modTGYROfile( - folder / "out.tgyro.evo_te", QeTGB, pos=positionMod, fileN_suffix=special_label - ) - modTGYROfile( - folder / "out.tgyro.evo_ti", QiTGB, pos=positionMod, fileN_suffix=special_label - ) - modTGYROfile( - folder / "out.tgyro.evo_ne", GeTGB, pos=positionMod, fileN_suffix=special_label - ) - modTGYROfile( - folder / "out.tgyro.evo_er", MtTGB, pos=positionMod, fileN_suffix=special_label - ) + modTGYROfile(folder / "out.tgyro.evo_te", QeTGB, pos=positionMod, fileN_suffix=special_label) + modTGYROfile(folder / "out.tgyro.evo_ti", QiTGB, pos=positionMod, fileN_suffix=special_label) + modTGYROfile(folder / "out.tgyro.evo_ne", GeTGB, pos=positionMod, fileN_suffix=special_label) + modTGYROfile(folder / "out.tgyro.evo_er", MtTGB, pos=positionMod, fileN_suffix=special_label) for i in range(tgyro.Qi_sim_turb.shape[0]): if i == impurityPosition: From d5beb87f8411879b5e784bccf11f29e3ea434c4f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 08:37:35 -0700 Subject: [PATCH 014/385] Moved CGYRO handlers from PORTALS to POWERTORCH --- .../portals/utils/PORTALScgyro.py | 720 ------------------ .../physics_models/transport_cgyro.py | 715 ++++++++++++++++- .../physics_models/transport_tgyro.py | 10 +- 3 files changed, 717 insertions(+), 728 deletions(-) delete mode 100644 src/mitim_modules/portals/utils/PORTALScgyro.py diff --git a/src/mitim_modules/portals/utils/PORTALScgyro.py b/src/mitim_modules/portals/utils/PORTALScgyro.py deleted file mode 100644 index 4dc32b48..00000000 --- a/src/mitim_modules/portals/utils/PORTALScgyro.py +++ /dev/null @@ -1,720 +0,0 @@ -import shutil -import copy -import numpy as np -from IPython import embed -from mitim_tools.misc_tools import IOtools, PLASMAtools -from mitim_tools.gacode_tools import PROFILEStools, TGYROtools -from mitim_tools.misc_tools.LOGtools import printMsg as print - -""" -__________________ -To run standalone: - run ~/MITIM/mitim_opt/mitim/utils/PORTALScgyro.py ./run5/ ~/PRF/mitim_cgyro/sparc_results.txt 0,1,2,3,4 -or - run ~/MITIM/mitim_opt/mitim/utils/PORTALScgyro.py ./run5/ ~/PRF/mitim_cgyro/sparc_results.txt 0[Evaluation.X] 0[position_in_txt] -__________________ -The CGYRO file must contain GB units, and the gb unit is MW/m^2, 1E19m^2/s -The CGYRO file must use particle flux. Convective transformation occurs later -""" - -def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmodified_profiles, radii, ProfilesPredicted): - print("\n ** CGYRO evaluation of fluxes has been requested before passing information to the STRATEGY module **",typeMsg="i",) - - if isinstance(numPORTALS, int): - numPORTALS = str(numPORTALS) - - # ------------------------------------------------------------------------------------------------ - # Harcoded - # ------------------------------------------------------------------------------------------------ - if PORTALSparameters['hardCodedCGYRO'] is not 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. 'example_run/Outputs/cgyro_results/iter_rmp_75_' - - e.g.: - train_sep,start_num,last_one,trick_hardcoded_f = 1, 0,100, 'example_run/Outputs/cgyro_results/d3d_5chan_it_' - - """ - - train_sep = PORTALSparameters["hardCodedCGYRO"]["train_sep"] - start_num = PORTALSparameters["hardCodedCGYRO"]["start_num"] - last_one = PORTALSparameters["hardCodedCGYRO"]["last_one"] - trick_hardcoded_f = PORTALSparameters["hardCodedCGYRO"]["trick_hardcoded_f"] - else: - train_sep = None - start_num = None - last_one = None - trick_hardcoded_f = None - # ------------------------------------------------------------------------------------------------ - - minErrorPercent = PORTALSparameters["percentError_stable"] - Qi_criterion_stable = PORTALSparameters["Qi_criterion_stable"] - percentNeo = PORTALSparameters["percentError"][1] - useConvectiveFluxes = PORTALSparameters["useConvectiveFluxes"] - - try: - impurityPosition = PROFILEStools.impurity_location(PROFILEStools.PROFILES_GACODE(unmodified_profiles), PORTALSparameters["ImpurityOfInterest"]) - except ValueError: - if 'nZ' in ProfilesPredicted: - raise ValueError(f"Impurity {PORTALSparameters['ImpurityOfInterest']} not found in the profiles and needed for CGYRO evaluation") - else: - impurityPosition = 0 - print(f'\t- Impurity location not found. Using hardcoded value of {impurityPosition}') - - OriginalFimp = PORTALSparameters["fImp_orig"] - - cgyroing_file = ( - lambda file_cgyro, numPORTALS_this=0: cgyroing( - FolderEvaluation, - unmodified_profiles, - numPORTALS, - minErrorPercent, - Qi_criterion_stable, - useConvectiveFluxes, - percentNeo, - radii, - OriginalFimp=OriginalFimp, - evaluationsInFile=f"{numPORTALS_this}", - impurityPosition=impurityPosition, - file=file_cgyro, - ) - ) - print(f"\t- Suggested function call for mitim evaluation {numPORTALS} (lambda for cgyroing):",typeMsg="i") - cgyropath = IOtools.expandPath(folder, ensurePathValid=True) / 'Outputs' / 'cgyro_results' / f'cgyro_it_{numPORTALS}.txt' - print(f"\tcgyroing_file('{cgyropath}')") - - print('\t- Then insert "exit" and RETURN', typeMsg="i") - if (trick_hardcoded_f is None) or (int(numPORTALS) > last_one): - embed() - else: - # ------------------------------------------------------------------ - # Hard-coded stuff for quick modifications - # ------------------------------------------------------------------ - if int(numPORTALS) < train_sep: - cgyroing_file( - f"{trick_hardcoded_f}{start_num}.txt", - numPORTALS_this=numPORTALS, - ) - else: - cgyroing_file( - f"{trick_hardcoded_f}{int(numPORTALS)-train_sep+1+start_num}.txt", - numPORTALS_this=0, - ) - - -def cgyroing( - FolderEvaluation, - unmodified_profiles, - evaluations, - minErrorPercent, - Qi_criterion_stable, - useConvectiveFluxes, - percentNeo, - radii, - OriginalFimp=1.0, - file=None, - evaluationsInFile=0, - impurityPosition=3, -): - """ - Variables need to have dimensions of (evaluation,rho) - """ - - evaluations = np.array([int(i) for i in evaluations.split(",")]) - evaluationsInFile = np.array([int(i) for i in evaluationsInFile.split(",")]) - - aLTe,aLTi,aLne,Q_gb,Qe,Qi,Ge,GZ,Mt,Pexch,QeE,QiE,GeE,GZE,MtE,PexchE,_,_ = readCGYROresults(file, radii) - - cont = 0 - for _ in evaluations: - k = evaluationsInFile[cont] - cont += 1 - - print(f"\t- Modifying {IOtools.clipstr(FolderEvaluation)} with position {k} in CGYRO results file {IOtools.clipstr(file)}") - - # Get TGYRO - tgyro = TGYROtools.TGYROoutput(FolderEvaluation,profiles=PROFILEStools.PROFILES_GACODE(unmodified_profiles)) - - # Quick checker of correct file - wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro) - - modifyResults( - Qe[k, :], - Qi[k, :], - Ge[k, :], - GZ[k, :], - Mt[k, :], - Pexch[k, :], - QeE[k, :], - QiE[k, :], - GeE[k, :], - GZE[k, :], - MtE[k, :], - PexchE[k, :], - tgyro, - FolderEvaluation, - minErrorPercent=minErrorPercent, - useConvectiveFluxes=useConvectiveFluxes, - Qi_criterion_stable=Qi_criterion_stable, - percentNeo=percentNeo, - impurityPosition=impurityPosition, - OriginalFimp=OriginalFimp, - ) - - -def wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro, ErrorRaised=0.005): - print("\t- Checking that this was the correct run...") - - tgyro_new = copy.deepcopy(tgyro) - tgyro_new.aLti = tgyro_new.aLti[:, 0, :] - - variables = [ - [aLTe, tgyro_new.aLte, "aLTe"], - [aLTi, tgyro_new.aLti, "aLTi"], - [aLne, tgyro_new.aLne, "aLne"], - [Q_gb, tgyro_new.Q_GB, "Qgb"], - ] - - for var in variables: - [c, t, n] = var - - for pos in range(c.shape[0]): - for i in range(c.shape[1]): - error = np.max(abs((t[pos, i + 1] - c[pos, i]) / t[pos, i + 1])) - print( - f"\t\t* Error in {n}[{i}] was {error*100.0:.2f}% (TGYRO {t[pos,i+1]:.3f} vs. CGYRO {c[pos,i]:.3f})", - typeMsg="w" if error > ErrorRaised else "", - ) - - -def readlineNTH(line, full_file=True, unnormalize=True): - s = line.split() - - i = 2 - roa = float(s[i]) - i += 3 - aLne = float(s[i]) - i += 3 - aLTi = float(s[i]) - i += 3 - aLTe = float(s[i]) - i += 3 - - Qi = float(s[i]) - i += 3 - Qi_std = float(s[i]) - i += 3 - Qe = float(s[i]) - i += 3 - Qe_std = float(s[i]) - i += 3 - Ge = float(s[i]) - i += 3 - Ge_std = float(s[i]) - i += 3 - - if full_file: - GZ = float(s[i]) - i += 3 - GZ_std = float(s[i]) - i += 3 - - Mt = float(s[i]) - i += 3 - Mt_std = float(s[i]) - i += 3 - - Pexch = float(s[i]) - i += 3 - Pexch_std = float(s[i]) - i += 3 - - Q_gb = float(s[i]) - i += 3 - G_gb = float(s[i]) * 1e-1 - i += 3 # From 1E19 to 1E20 - - if full_file: - Mt_gb = float(s[i]) - i += 3 - Pexch_gb = float(s[i]) - i += 3 - - tstart = float(s[i]) - i += 3 - tend = float(s[i]) - i += 3 - - if unnormalize: - QiReal = Qi * Q_gb - QiReal_std = Qi_std * Q_gb - QeReal = Qe * Q_gb - QeReal_std = Qe_std * Q_gb - GeReal = Ge * G_gb - GeReal_std = Ge_std * G_gb - else: - QiReal = Qi - QiReal_std = Qi_std - QeReal = Qe - QeReal_std = Qe_std - GeReal = Ge - GeReal_std = Ge_std - - if full_file: - if unnormalize: - GZReal = GZ * G_gb - GZReal_std = GZ_std * G_gb - - MtReal = Mt * Mt_gb - MtReal_std = Mt_std * Mt_gb - - PexchReal = Pexch * Pexch_gb - PexchReal_std = Pexch_std * Pexch_gb - else: - GZReal = GZ - GZReal_std = GZ_std - - MtReal = Mt - MtReal_std = Mt_std - - PexchReal = Pexch - PexchReal_std = Pexch_std - - return roa,aLTe,aLTi,aLne,Q_gb,QeReal,QiReal,GeReal,GZReal,MtReal,PexchReal,QeReal_std,QiReal_std,GeReal_std,GZReal_std,MtReal_std,PexchReal_std,tstart,tend - else: - return roa,aLTe,aLTi,aLne,Q_gb,QeReal,QiReal,GeReal,0.0,0.0,0.0,QeReal_std,QiReal_std,GeReal_std,0.0,0.0,0.0,tstart,tend - - -def readCGYROresults(file, radii, unnormalize=True): - """ - Arrays are in (batch,radii) - MW/m^2 and 1E20 - """ - - with open(file, "r") as f: - lines = f.readlines() - - rad = len(radii) - num = len(lines) // rad - - roa = np.zeros((num, rad)) - aLTe = np.zeros((num, rad)) - aLTi = np.zeros((num, rad)) - aLne = np.zeros((num, rad)) - Q_gb = np.zeros((num, rad)) - - Qe = np.zeros((num, rad)) - Qe_std = np.zeros((num, rad)) - Qi = np.zeros((num, rad)) - Qi_std = np.zeros((num, rad)) - Ge = np.zeros((num, rad)) - Ge_std = np.zeros((num, rad)) - - GZ = np.zeros((num, rad)) - GZ_std = np.zeros((num, rad)) - - Mt = np.zeros((num, rad)) - Mt_std = np.zeros((num, rad)) - - Pexch = np.zeros((num, rad)) - Pexch_std = np.zeros((num, rad)) - - tstart = np.zeros((num, rad)) - tend = np.zeros((num, rad)) - - p = {} - for r in range(len(radii)): - p[r] = 0 - for i in range(len(lines)): - - # -------------------------------------------------------- - # Line not empty - # -------------------------------------------------------- - if len(lines[i].split()) < 10: - continue - - # -------------------------------------------------------- - # Read line - # -------------------------------------------------------- - ( - roa_read, - aLTe_read, - aLTi_read, - aLne_read, - Q_gb_read, - Qe_read, - Qi_read, - Ge_read, - GZ_read, - Mt_read, - Pexch_read, - Qe_std_read, - Qi_std_read, - Ge_std_read, - GZ_std_read, - Mt_std_read, - Pexch_std_read, - tstart_read, - tend_read, - ) = readlineNTH(lines[i], unnormalize=unnormalize) - - # -------------------------------------------------------- - # Radial location position - # -------------------------------------------------------- - threshold_radii = 1E-4 - r = np.where(np.abs(radii-roa_read) last_one): + embed() + else: + # ------------------------------------------------------------------ + # Hard-coded stuff for quick modifications + # ------------------------------------------------------------------ + if int(numPORTALS) < train_sep: + cgyroing_file( + f"{trick_hardcoded_f}{start_num}.txt", + numPORTALS_this=numPORTALS, + ) + else: + cgyroing_file( + f"{trick_hardcoded_f}{int(numPORTALS)-train_sep+1+start_num}.txt", + numPORTALS_this=0, + ) + + +def cgyroing( + FolderEvaluation, + unmodified_profiles, + evaluations, + minErrorPercent, + Qi_criterion_stable, + useConvectiveFluxes, + percentNeo, + radii, + OriginalFimp=1.0, + file=None, + evaluationsInFile=0, + impurityPosition=3, +): + """ + Variables need to have dimensions of (evaluation,rho) + """ + + evaluations = np.array([int(i) for i in evaluations.split(",")]) + evaluationsInFile = np.array([int(i) for i in evaluationsInFile.split(",")]) + + aLTe,aLTi,aLne,Q_gb,Qe,Qi,Ge,GZ,Mt,Pexch,QeE,QiE,GeE,GZE,MtE,PexchE,_,_ = readCGYROresults(file, radii) + + cont = 0 + for _ in evaluations: + k = evaluationsInFile[cont] + cont += 1 + + print(f"\t- Modifying {IOtools.clipstr(FolderEvaluation)} with position {k} in CGYRO results file {IOtools.clipstr(file)}") + + # Get TGYRO + tgyro = TGYROtools.TGYROoutput(FolderEvaluation,profiles=PROFILEStools.PROFILES_GACODE(unmodified_profiles)) + + # Quick checker of correct file + wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro) + + modifyResults( + Qe[k, :], + Qi[k, :], + Ge[k, :], + GZ[k, :], + Mt[k, :], + Pexch[k, :], + QeE[k, :], + QiE[k, :], + GeE[k, :], + GZE[k, :], + MtE[k, :], + PexchE[k, :], + tgyro, + FolderEvaluation, + minErrorPercent=minErrorPercent, + useConvectiveFluxes=useConvectiveFluxes, + Qi_criterion_stable=Qi_criterion_stable, + percentNeo=percentNeo, + impurityPosition=impurityPosition, + OriginalFimp=OriginalFimp, + ) + + +def wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro, ErrorRaised=0.005): + print("\t- Checking that this was the correct run...") + + tgyro_new = copy.deepcopy(tgyro) + tgyro_new.aLti = tgyro_new.aLti[:, 0, :] + + variables = [ + [aLTe, tgyro_new.aLte, "aLTe"], + [aLTi, tgyro_new.aLti, "aLTi"], + [aLne, tgyro_new.aLne, "aLne"], + [Q_gb, tgyro_new.Q_GB, "Qgb"], + ] + + for var in variables: + [c, t, n] = var + + for pos in range(c.shape[0]): + for i in range(c.shape[1]): + error = np.max(abs((t[pos, i + 1] - c[pos, i]) / t[pos, i + 1])) + print( + f"\t\t* Error in {n}[{i}] was {error*100.0:.2f}% (TGYRO {t[pos,i+1]:.3f} vs. CGYRO {c[pos,i]:.3f})", + typeMsg="w" if error > ErrorRaised else "", + ) + + +def readlineNTH(line, full_file=True, unnormalize=True): + s = line.split() + + i = 2 + roa = float(s[i]) + i += 3 + aLne = float(s[i]) + i += 3 + aLTi = float(s[i]) + i += 3 + aLTe = float(s[i]) + i += 3 + + Qi = float(s[i]) + i += 3 + Qi_std = float(s[i]) + i += 3 + Qe = float(s[i]) + i += 3 + Qe_std = float(s[i]) + i += 3 + Ge = float(s[i]) + i += 3 + Ge_std = float(s[i]) + i += 3 + + if full_file: + GZ = float(s[i]) + i += 3 + GZ_std = float(s[i]) + i += 3 + + Mt = float(s[i]) + i += 3 + Mt_std = float(s[i]) + i += 3 + + Pexch = float(s[i]) + i += 3 + Pexch_std = float(s[i]) + i += 3 + + Q_gb = float(s[i]) + i += 3 + G_gb = float(s[i]) * 1e-1 + i += 3 # From 1E19 to 1E20 + + if full_file: + Mt_gb = float(s[i]) + i += 3 + Pexch_gb = float(s[i]) + i += 3 + + tstart = float(s[i]) + i += 3 + tend = float(s[i]) + i += 3 + + if unnormalize: + QiReal = Qi * Q_gb + QiReal_std = Qi_std * Q_gb + QeReal = Qe * Q_gb + QeReal_std = Qe_std * Q_gb + GeReal = Ge * G_gb + GeReal_std = Ge_std * G_gb + else: + QiReal = Qi + QiReal_std = Qi_std + QeReal = Qe + QeReal_std = Qe_std + GeReal = Ge + GeReal_std = Ge_std + + if full_file: + if unnormalize: + GZReal = GZ * G_gb + GZReal_std = GZ_std * G_gb + + MtReal = Mt * Mt_gb + MtReal_std = Mt_std * Mt_gb + + PexchReal = Pexch * Pexch_gb + PexchReal_std = Pexch_std * Pexch_gb + else: + GZReal = GZ + GZReal_std = GZ_std + + MtReal = Mt + MtReal_std = Mt_std + + PexchReal = Pexch + PexchReal_std = Pexch_std + + return roa,aLTe,aLTi,aLne,Q_gb,QeReal,QiReal,GeReal,GZReal,MtReal,PexchReal,QeReal_std,QiReal_std,GeReal_std,GZReal_std,MtReal_std,PexchReal_std,tstart,tend + else: + return roa,aLTe,aLTi,aLne,Q_gb,QeReal,QiReal,GeReal,0.0,0.0,0.0,QeReal_std,QiReal_std,GeReal_std,0.0,0.0,0.0,tstart,tend + + +def readCGYROresults(file, radii, unnormalize=True): + """ + Arrays are in (batch,radii) + MW/m^2 and 1E20 + """ + + with open(file, "r") as f: + lines = f.readlines() + + rad = len(radii) + num = len(lines) // rad + + roa = np.zeros((num, rad)) + aLTe = np.zeros((num, rad)) + aLTi = np.zeros((num, rad)) + aLne = np.zeros((num, rad)) + Q_gb = np.zeros((num, rad)) + + Qe = np.zeros((num, rad)) + Qe_std = np.zeros((num, rad)) + Qi = np.zeros((num, rad)) + Qi_std = np.zeros((num, rad)) + Ge = np.zeros((num, rad)) + Ge_std = np.zeros((num, rad)) + + GZ = np.zeros((num, rad)) + GZ_std = np.zeros((num, rad)) + + Mt = np.zeros((num, rad)) + Mt_std = np.zeros((num, rad)) + + Pexch = np.zeros((num, rad)) + Pexch_std = np.zeros((num, rad)) + + tstart = np.zeros((num, rad)) + tend = np.zeros((num, rad)) + + p = {} + for r in range(len(radii)): + p[r] = 0 + for i in range(len(lines)): + + # -------------------------------------------------------- + # Line not empty + # -------------------------------------------------------- + if len(lines[i].split()) < 10: + continue + + # -------------------------------------------------------- + # Read line + # -------------------------------------------------------- + ( + roa_read, + aLTe_read, + aLTi_read, + aLne_read, + Q_gb_read, + Qe_read, + Qi_read, + Ge_read, + GZ_read, + Mt_read, + Pexch_read, + Qe_std_read, + Qi_std_read, + Ge_std_read, + GZ_std_read, + Mt_std_read, + Pexch_std_read, + tstart_read, + tend_read, + ) = readlineNTH(lines[i], unnormalize=unnormalize) + + # -------------------------------------------------------- + # Radial location position + # -------------------------------------------------------- + threshold_radii = 1E-4 + r = np.where(np.abs(radii-roa_read) Date: Thu, 24 Apr 2025 08:53:20 -0700 Subject: [PATCH 015/385] Cleaned up CGYRO trick (1) --- .../physics_models/transport_cgyro.py | 146 +++++++++--------- .../physics_models/transport_tgyro.py | 4 - 2 files changed, 70 insertions(+), 80 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 50d282c9..928d4aad 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -12,9 +12,6 @@ class cgyro_model(transport_tgyro.tgyro_model): def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) - def produce_profiles(self): - super().produce_profiles() - def evaluate(self): # Run original evaluator @@ -38,97 +35,100 @@ def evaluate(self): def _trick_cgyro(self, tgyro): + FolderEvaluation_TGYRO = self.folder / "cgyro_neo" + print("\t- Checking whether cgyro_neo folder exists and it was written correctly via cgyro_trick...") - correctly_run = (self.folder / "cgyro_neo").exists() + correctly_run = FolderEvaluation_TGYRO.exists() if correctly_run: print("\t\t- Folder exists, but was cgyro_trick run?") - with open(self.folder / "cgyro_neo" / "mitim_flag", "r") as f: + with open(FolderEvaluation_TGYRO / "mitim_flag", "r") as f: correctly_run = bool(float(f.readline())) if correctly_run: print("\t\t\t* Yes, it was", typeMsg="w") else: - print("\t\t\t* No, it was not, repating process", typeMsg="i") + print("\t\t\t* No, it was run, repating process", typeMsg="i") # Remove cgyro_neo folder - if (self.folder / "cgyro_neo").exists(): - IOtools.shutil_rmtree(self.folder / "cgyro_neo") + if FolderEvaluation_TGYRO.exists(): + IOtools.shutil_rmtree(FolderEvaluation_TGYRO) # Copy tglf_neo results - shutil.copytree(self.folder / "tglf_neo", self.folder / "cgyro_neo") + shutil.copytree(self.folder / "tglf_neo", FolderEvaluation_TGYRO) + # ********************************************************** # CGYRO writter - cgyro_trick(self,self.folder / "cgyro_neo") + # ********************************************************** + + # Write a flag indicating this was not performed + with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: + f.write("0") + + self._cgyro_trick(FolderEvaluation_TGYRO) + + # Write a flag indicating this was performed, to avoid an issue that... the script crashes when it has copied tglf_neo, without cgyro_trick modification + with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: + f.write("1") # Read TGYRO files and construct portals variables - tgyro.read(label="cgyro_neo", folder=self.folder / "cgyro_neo") + tgyro.read(label="cgyro_neo", folder=FolderEvaluation_TGYRO) powerstate_orig = copy.deepcopy(self.powerstate) return powerstate_orig -def cgyro_trick(self,FolderEvaluation_TGYRO): - - with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: - f.write("0") - - # ************************************************************************************************************************** - # Print Information - # ************************************************************************************************************************** - - txt = "\nFluxes to be matched by CGYRO ( TARGETS - NEO ):" - - for var, varn in zip( - ["r/a ", "rho ", "a/LTe", "a/LTi", "a/Lne", "a/LnZ", "a/Lw0"], - ["roa", "rho", "aLte", "aLti", "aLne", "aLnZ", "aLw0"], - ): - txt += f"\n{var} = " - for j in range(self.powerstate.plasma["rho"].shape[1] - 1): - txt += f"{self.powerstate.plasma[varn][0,j+1]:.6f} " - - for var, varn in zip( - ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ce (MW/m^2)", "CZ (MW/m^2)", "Mt (J/m^2) "], - ["Pe", "Pi", "Ce", "CZ", "Mt"], - ): - txt += f"\n{var} = " - for j in range(self.powerstate.plasma["rho"].shape[1] - 1): - txt += f"{self.powerstate.plasma[varn][0,j+1]-self.powerstate.plasma[f'{varn}_tr_neo'][0,j+1]:.4e} " - - print(txt) - - # Copy profiles so that later it is easy to grab all the input.gacodes that were evaluated - self._profiles_to_store() - - # ************************************************************************************************************************** - # Evaluate CGYRO - # ************************************************************************************************************************** - - evaluateCGYRO( - self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["PORTALSparameters"], - self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"], - self.evaluation_number, - FolderEvaluation_TGYRO, - self.file_profs, - self.powerstate.plasma["roa"][0,1:], - self.powerstate.ProfilesPredicted, - ) + def _cgyro_trick(self,FolderEvaluation_TGYRO): + + # Print Information + print(self._print_info()) + + # Copy profiles so that later it is easy to grab all the input.gacodes that were evaluated + self._profiles_to_store() + + # ************************************************************************************************************************** + # Evaluate CGYRO + # ************************************************************************************************************************** + + evaluateCGYRO( + self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["PORTALSparameters"], + self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"], + self.evaluation_number, + FolderEvaluation_TGYRO, + self.file_profs, + self.powerstate.plasma["roa"][0,1:], + self.powerstate.ProfilesPredicted, + ) + + # Make tensors + for i in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: + try: + self.powerstate.plasma[i] = torch.from_numpy(self.powerstate.plasma[i]).to(self.powerstate.dfT).unsqueeze(0) + except: + pass + + def _print_info(self): + + txt = "\nFluxes to be matched by CGYRO ( TARGETS - NEO ):" - # ************************************************************************************************************************** - # EXTRA - # ************************************************************************************************************************** + for var, varn in zip( + ["r/a ", "rho ", "a/LTe", "a/LTi", "a/Lne", "a/LnZ", "a/Lw0"], + ["roa", "rho", "aLte", "aLti", "aLne", "aLnZ", "aLw0"], + ): + txt += f"\n{var} = " + for j in range(self.powerstate.plasma["rho"].shape[1] - 1): + txt += f"{self.powerstate.plasma[varn][0,j+1]:.6f} " - # Make tensors - for i in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: - try: - self.powerstate.plasma[i] = torch.from_numpy(self.powerstate.plasma[i]).to(self.powerstate.dfT).unsqueeze(0) - except: - pass + for var, varn in zip( + ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ce (MW/m^2)", "CZ (MW/m^2)", "Mt (J/m^2) "], + ["Pe", "Pi", "Ce", "CZ", "Mt"], + ): + txt += f"\n{var} = " + for j in range(self.powerstate.plasma["rho"].shape[1] - 1): + txt += f"{self.powerstate.plasma[varn][0,j+1]-self.powerstate.plasma[f'{varn}_tr_neo'][0,j+1]:.4e} " - # Write a flag indicating this was performed, to avoid an issue that... the script crashes when it has copied tglf_neo, without cgyro_trick modification - with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: - f.write("1") + return txt """ The CGYRO file must contain GB units, and the gb unit is MW/m^2, 1E19m^2/s @@ -210,15 +210,9 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod # Hard-coded stuff for quick modifications # ------------------------------------------------------------------ if int(numPORTALS) < train_sep: - cgyroing_file( - f"{trick_hardcoded_f}{start_num}.txt", - numPORTALS_this=numPORTALS, - ) + cgyroing_file(f"{trick_hardcoded_f}{start_num}.txt",numPORTALS_this=numPORTALS) else: - cgyroing_file( - f"{trick_hardcoded_f}{int(numPORTALS)-train_sep+1+start_num}.txt", - numPORTALS_this=0, - ) + cgyroing_file(f"{trick_hardcoded_f}{int(numPORTALS)-train_sep+1+start_num}.txt",numPORTALS_this=0) def cgyroing( diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index b65820f0..0e6871db 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -8,10 +8,6 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -# ---------------------------------------------------------------------------------------------------- -# FULL TGYRO -# ---------------------------------------------------------------------------------------------------- - class tgyro_model(TRANSPORTtools.power_transport): def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) From 50d91481fc28ece8f1f16fde110e2fd4065b253d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 09:09:57 -0700 Subject: [PATCH 016/385] Towards renaming of powertorch fluxes from, e.g., Pe to QeMWm2 for much more clarity --- src/mitim_modules/portals/PORTALSmain.py | 4 +- src/mitim_modules/powertorch/STATEtools.py | 6 +- .../physics_models/transport_analytic.py | 16 +- .../physics_models/transport_cgyro.py | 333 +---------------- .../physics_models/transport_tgyro.py | 335 +++++++++++++++++- .../powertorch/scripts/calculateTargets.py | 2 +- .../scripts/compareRadialResolution.py | 2 +- .../powertorch/scripts/compareWithTGYRO.py | 8 +- .../powertorch/utils/TARGETStools.py | 18 +- .../powertorch/utils/TRANSFORMtools.py | 8 +- .../gacode_tools/utils/PORTALSinteraction.py | 60 ++-- 11 files changed, 396 insertions(+), 396 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 2427258a..661dcc18 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -657,9 +657,9 @@ def map_powerstate_to_portals(powerstate, dictOFs): # Write in OFs for i in range(powerstate.plasma["rho"].shape[1] - 1): # Ignore position 0, which is rho=0 if var == "te": - var0, var1 = "Qe", "Pe" + var0, var1 = "Qe", "QeMWm2" elif var == "ti": - var0, var1 = "Qi", "Pi" + var0, var1 = "Qi", "QiMWm2" elif var == "ne": var0, var1 = "Ge", "Ce" elif var == "nZ": diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index ced8dd95..023ab17b 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -104,8 +104,8 @@ def _ensure_ne_before_nz(lst): tensors will be the same as in self.ProfilesPredicted ''' self.profile_map = { - "te": ("Pe", "Pe_tr"), - "ti": ("Pi", "Pi_tr"), + "te": ("QeMWm2", "QeMWm2_tr"), + "ti": ("QiMWm2", "QiMWm2_tr"), "ne": ("Ce", "Ce_tr"), "nZ": ("CZ", "CZ_tr"), "w0": ("Mt", "Mt_tr") @@ -738,7 +738,7 @@ def _concatenate_flux(plasma, profile_key, flux_key): plasma["P"] = torch.cat((plasma["P"], plasma[profile_key][:, 1:]), dim=1).to(plasma["P"].device) plasma["P_tr"] = torch.cat((plasma["P_tr"], plasma[flux_key][:, 1:]), dim=1).to(plasma["P"].device) - self.plasma["P"], self.plasma["P_tr"] = torch.Tensor().to(self.plasma["Pe"]), torch.Tensor().to(self.plasma["Pe"]) + self.plasma["P"], self.plasma["P_tr"] = torch.Tensor().to(self.plasma["QeMWm2"]), torch.Tensor().to(self.plasma["QeMWm2"]) for profile in self.ProfilesPredicted: _concatenate_flux(self.plasma, *self.profile_map[profile]) diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py index 87839749..07622713 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -43,14 +43,14 @@ def evaluate(self): self.powerstate.plasma["a"].unsqueeze(-1), ) - self.powerstate.plasma["Pe_tr_turb"] = Pe_tr * 2 / 3 - self.powerstate.plasma["Pi_tr_turb"] = Pi_tr * 2 / 3 + self.powerstate.plasma["QeMWm2_tr_turb"] = Pe_tr * 2 / 3 + self.powerstate.plasma["QiMWm2_tr_turb"] = Pi_tr * 2 / 3 - self.powerstate.plasma["Pe_tr_neo"] = Pe_tr * 1 / 3 - self.powerstate.plasma["Pi_tr_neo"] = Pi_tr * 1 / 3 + self.powerstate.plasma["QeMWm2_tr_neo"] = Pe_tr * 1 / 3 + self.powerstate.plasma["QiMWm2_tr_neo"] = Pi_tr * 1 / 3 - self.powerstate.plasma["Pe_tr"] = self.powerstate.plasma["Pe_tr_turb"] + self.powerstate.plasma["Pe_tr_neo"] - self.powerstate.plasma["Pi_tr"] = self.powerstate.plasma["Pi_tr_turb"] + self.powerstate.plasma["Pi_tr_neo"] + self.powerstate.plasma["QeMWm2_tr"] = self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["QeMWm2_tr_neo"] + self.powerstate.plasma["QiMWm2_tr"] = self.powerstate.plasma["QiMWm2_tr_turb"] + self.powerstate.plasma["QiMWm2_tr_neo"] # ------------------------------------------------------------------ # SURROGATE @@ -78,8 +78,8 @@ def evaluate(self): numeach = self.powerstate.plasma["rho"].shape[1] - 1 quantities = { - "te": "Pe", - "ti": "Pi", + "te": "QeMWm2", + "ti": "QiMWm2", "ne": "Ce", "nZ": "CZ", "w0": "Mt", diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 928d4aad..0a475b74 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -2,7 +2,7 @@ import shutil import torch import numpy as np -from mitim_tools.misc_tools import IOtools, PLASMAtools +from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools, TGYROtools from mitim_modules.powertorch.physics_models import transport_tgyro from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -25,7 +25,7 @@ def evaluate(self): # Some checks print("\t- Checking model modifications:") - for r in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: #, "PexchTurb"]: #TODO: FIX + for r in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: #, "PexchTurb"]: #TODO: FIX print(f"\t\t{r}(tglf) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(powerstate_orig.plasma[r][0][1:],powerstate_orig.plasma[r+'_stds'][0][1:]) ])}") print(f"\t\t{r}(cgyro) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(self.powerstate.plasma[r][0][1:],self.powerstate.plasma[r+'_stds'][0][1:]) ])}") @@ -102,7 +102,7 @@ def _cgyro_trick(self,FolderEvaluation_TGYRO): ) # Make tensors - for i in ["Pe_tr_turb", "Pi_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: + for i in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: try: self.powerstate.plasma[i] = torch.from_numpy(self.powerstate.plasma[i]).to(self.powerstate.dfT).unsqueeze(0) except: @@ -122,7 +122,7 @@ def _print_info(self): for var, varn in zip( ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ce (MW/m^2)", "CZ (MW/m^2)", "Mt (J/m^2) "], - ["Pe", "Pi", "Ce", "CZ", "Mt"], + ["QeMWm2", "Pi", "Ce", "CZ", "Mt"], ): txt += f"\n{var} = " for j in range(self.powerstate.plasma["rho"].shape[1] - 1): @@ -251,7 +251,7 @@ def cgyroing( # Quick checker of correct file wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro) - modifyResults( + transport_tgyro.modifyResults( Qe[k, :], Qi[k, :], Ge[k, :], @@ -507,327 +507,4 @@ def readCGYROresults(file, radii, unnormalize=True): ) -def defineReferenceFluxes( - tgyro, factor_tauptauE=5, useConvectiveFluxes=False, impurityPosition=3 -): - Qe_target = abs(tgyro.Qe_tar[0, 1:]) - Qi_target = abs(tgyro.Qi_tar[0, 1:]) - Mt_target = abs(tgyro.Mt_tar[0, 1:]) - - # For particle fluxes, since the targets are often zero... it's more complicated - QeMW_target = abs(tgyro.Qe_tarMW[0, 1:]) - QiMW_target = abs(tgyro.Qi_tarMW[0, 1:]) - We, Wi, Ne, NZ = tgyro.profiles.deriveContentByVolumes( - rhos=tgyro.rho[0, 1:], impurityPosition=impurityPosition - ) - - tau_special = ( - (We + Wi) / (QeMW_target + QiMW_target) * factor_tauptauE - ) # tau_p in seconds - Ge_target_special = (Ne / tau_special) / tgyro.dvoldr[0, 1:] # (1E20/seconds/m^2) - - if useConvectiveFluxes: - Ge_target_special = PLASMAtools.convective_flux( - tgyro.Te[0, 1:], Ge_target_special - ) # (1E20/seconds/m^2) - - GZ_target_special = Ge_target_special * NZ / Ne - - return Qe_target, Qi_target, Ge_target_special, GZ_target_special, Mt_target - - -def modifyResults( - Qe, - Qi, - Ge, - GZ, - Mt, - Pexch, - QeE, - QiE, - GeE, - GZE, - MtE, - PexchE, - tgyro, - folder_tgyro, - minErrorPercent=5.0, - percentNeo=2.0, - useConvectiveFluxes=False, - Qi_criterion_stable=0.0025, - impurityPosition=3, - OriginalFimp=1.0, -): - """ - All in real units, with dimensions of (rho) from axis to edge - """ - - # If a plasma is very close to stable... do something about error - if minErrorPercent is not None: - ( - Qe_target, - Qi_target, - Ge_target_special, - GZ_target_special, - Mt_target, - ) = defineReferenceFluxes( - tgyro, - useConvectiveFluxes=useConvectiveFluxes, - impurityPosition=impurityPosition, - ) - - Qe_min = Qe_target * (minErrorPercent / 100.0) - Qi_min = Qi_target * (minErrorPercent / 100.0) - Ge_min = Ge_target_special * (minErrorPercent / 100.0) - GZ_min = GZ_target_special * (minErrorPercent / 100.0) - Mt_min = Mt_target * (minErrorPercent / 100.0) - - for i in range(Qe.shape[0]): - if Qi[i] < Qi_criterion_stable: - print( - f"\t- Based on 'Qi_criterion_stable', plasma considered stable (Qi = {Qi[i]:.2e} < {Qi_criterion_stable:.2e} MW/m2) at position #{i}, using minimum errors of {minErrorPercent}% of targets", - typeMsg="w", - ) - QeE[i] = Qe_min[i] - print(f"\t\t* QeE = {QeE[i]}") - QiE[i] = Qi_min[i] - print(f"\t\t* QiE = {QiE[i]}") - GeE[i] = Ge_min[i] - print(f"\t\t* GeE = {GeE[i]}") - GZE[i] = GZ_min[i] - print(f"\t\t* GZE = {GZE[i]}") - MtE[i] = Mt_min[i] - print(f"\t\t* MtE = {MtE[i]}") - - # Heat fluxes - QeTot = Qe + tgyro.Qe_sim_neo[0, 1:] - QiTot = Qi + tgyro.QiIons_sim_neo_thr[0, 1:] - - # Particle fluxes - PeTot = Ge + tgyro.Ge_sim_neo[0, 1:] - PZTot = GZ + tgyro.Gi_sim_neo[impurityPosition, 0, 1:] - - # Momentum fluxes - MtTot = Mt + tgyro.Mt_sim_neo[0, 1:] - - # ************************************************************************************ - # **** Modify complete folder (Division of ion fluxes will be wrong, since I put everything in first ion) - # ************************************************************************************ - - # 1. Modify out.tgyro.evo files (which contain turb+neo summed together) - - print(f"\t- Modifying TGYRO out.tgyro.evo files in {IOtools.clipstr(folder_tgyro)}") - modifyEVO( - tgyro, - folder_tgyro, - QeTot, - QiTot, - PeTot, - PZTot, - MtTot, - impurityPosition=impurityPosition, - ) - - # 2. Modify out.tgyro.flux files (which contain turb and neo separated) - - print(f"\t- Modifying TGYRO out.tgyro.flux files in {folder_tgyro}") - modifyFLUX( - tgyro, - folder_tgyro, - Qe, - Qi, - Ge, - GZ, - Mt, - Pexch, - impurityPosition=impurityPosition, - ) - - # 3. Modify files for errors - - print(f"\t- Modifying TGYRO out.tgyro.flux_stds in {folder_tgyro}") - modifyFLUX( - tgyro, - folder_tgyro, - QeE, - QiE, - GeE, - GZE, - MtE, - PexchE, - impurityPosition=impurityPosition, - special_label="_stds", - ) - - -def modifyEVO( - tgyro, - folder, - QeT, - QiT, - GeT, - GZT, - MtT, - impurityPosition=3, - positionMod=1, - special_label=None, -): - QeTGB = QeT / tgyro.Q_GB[-1, 1:] - QiTGB = QiT / tgyro.Q_GB[-1, 1:] - GeTGB = GeT / tgyro.Gamma_GB[-1, 1:] - GZTGB = GZT / tgyro.Gamma_GB[-1, 1:] - MtTGB = MtT / tgyro.Pi_GB[-1, 1:] - - modTGYROfile(folder / "out.tgyro.evo_te", QeTGB, pos=positionMod, fileN_suffix=special_label) - modTGYROfile(folder / "out.tgyro.evo_ti", QiTGB, pos=positionMod, fileN_suffix=special_label) - modTGYROfile(folder / "out.tgyro.evo_ne", GeTGB, pos=positionMod, fileN_suffix=special_label) - modTGYROfile(folder / "out.tgyro.evo_er", MtTGB, pos=positionMod, fileN_suffix=special_label) - - for i in range(tgyro.Qi_sim_turb.shape[0]): - if i == impurityPosition: - var = GZTGB - else: - var = GZTGB * 0.0 - modTGYROfile( - folder / f"out.tgyro.evo_n{i+1}", - var, - pos=positionMod, - fileN_suffix=special_label, - ) - - -def modifyFLUX( - tgyro, - folder, - Qe, - Qi, - Ge, - GZ, - Mt, - S, - QeNeo=None, - QiNeo=None, - GeNeo=None, - GZNeo=None, - MtNeo=None, - impurityPosition=3, - special_label=None, -): - folder = IOtools.expandPath(folder) - - QeGB = Qe / tgyro.Q_GB[-1, 1:] - QiGB = Qi / tgyro.Q_GB[-1, 1:] - GeGB = Ge / tgyro.Gamma_GB[-1, 1:] - GZGB = GZ / tgyro.Gamma_GB[-1, 1:] - MtGB = Mt / tgyro.Pi_GB[-1, 1:] - SGB = S / tgyro.S_GB[-1, 1:] - - # ****************************************************************************************** - # Electrons - # ****************************************************************************************** - - # Particle flux: Update - - modTGYROfile(folder / "out.tgyro.flux_e", GeGB, pos=2, fileN_suffix=special_label) - if GeNeo is not None: - GeGB_neo = GeNeo / tgyro.Gamma_GB[-1, 1:] - modTGYROfile(folder / "out.tgyro.flux_e", GeGB_neo, pos=1, fileN_suffix=special_label) - - # Energy flux: Update - - modTGYROfile(folder / "out.tgyro.flux_e", QeGB, pos=4, fileN_suffix=special_label) - if QeNeo is not None: - QeGB_neo = QeNeo / tgyro.Q_GB[-1, 1:] - modTGYROfile(folder / "out.tgyro.flux_e", QeGB_neo, pos=3, fileN_suffix=special_label) - - # Rotation: Remove (it will be sum to the first ion) - - modTGYROfile(folder / "out.tgyro.flux_e", GeGB * 0.0, pos=6, fileN_suffix=special_label) - modTGYROfile(folder / "out.tgyro.flux_e", GeGB * 0.0, pos=5, fileN_suffix=special_label) - - # Energy exchange - - modTGYROfile(folder / "out.tgyro.flux_e", SGB, pos=7, fileN_suffix=special_label) - - # SMW = S # S is MW/m^3 - # modTGYROfile(f'{folder}/out.tgyro.power_e',SMW,pos=8,fileN_suffix=special_label) - # print('\t\t- Modified turbulent energy exchange in out.tgyro.power_e') - - # ****************************************************************************************** - # Ions - # ****************************************************************************************** - - # Energy flux: Update - - modTGYROfile(folder / "out.tgyro.flux_i1", QiGB, pos=4, fileN_suffix=special_label) - - if QiNeo is not None: - QiGB_neo = QiNeo / tgyro.Q_GB[-1, 1:] - modTGYROfile(folder / "out.tgyro.flux_i1", QiGB_neo, pos=3, fileN_suffix=special_label) - - # Particle flux: Make ion particle fluxes zero, because I don't want to mistake TGLF with CGYRO when looking at tgyro results - - for i in range(tgyro.Qi_sim_turb.shape[0]): - if tgyro.profiles.Species[i]["S"] == "therm": - var = QiGB * 0.0 - modTGYROfile(folder / f"out.tgyro.flux_i{i+1}",var,pos=2,fileN_suffix=special_label,) # Gi_turb - modTGYROfile(folder / f"out.tgyro.evo_n{i+1}", var, pos=1, fileN_suffix=special_label) # Gi (Gi_sim) - - if i != impurityPosition: - modTGYROfile(folder / f"out.tgyro.flux_i{i+1}",var,pos=1,fileN_suffix=special_label) # Gi_neo - - # Rotation: Update - - modTGYROfile(folder / "out.tgyro.flux_i1", MtGB, pos=6, fileN_suffix=special_label) - - if MtNeo is not None: - MtGB_neo = MtNeo / tgyro.Pi_GB[-1, 1:] - modTGYROfile(folder / "out.tgyro.flux_i1", MtGB_neo, pos=5, fileN_suffix=special_label) - - # Energy exchange: Remove (it will be the electrons one) - - modTGYROfile(folder / "out.tgyro.flux_i1", SGB * 0.0, pos=7, fileN_suffix=special_label) - - # ****************************************************************************************** - # Impurities - # ****************************************************************************************** - - # Remove everything from all the rest of non-first ions (except the particles for the impurity chosen) - - for i in range(tgyro.Qi_sim_turb.shape[0] - 1): - if tgyro.profiles.Species[i + 1]["S"] == "therm": - var = QiGB * 0.0 - for pos in [3, 4, 5, 6, 7]: - modTGYROfile(folder / f"out.tgyro.flux_i{i+2}",var,pos=pos,fileN_suffix=special_label) - for pos in [1, 2]: - if i + 2 != impurityPosition: - modTGYROfile(folder / f"out.tgyro.flux_i{i+2}",var,pos=pos,fileN_suffix=special_label) - - modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB,pos=2,fileN_suffix=special_label) - if GZNeo is not None: - GZGB_neo = GZNeo / tgyro.Gamma_GB[-1, 1:] - modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB_neo,pos=1,fileN_suffix=special_label) - - -def modTGYROfile(file, var, pos=0, fileN_suffix=None): - fileN = file if fileN_suffix is None else file.parent / f"{file.name}{fileN_suffix}" - - if not fileN.exists(): - shutil.copy2(file, fileN) - - with open(fileN, "r") as f: - lines = f.readlines() - - with open(fileN, "w") as f: - f.write(lines[0]) - f.write(lines[1]) - f.write(lines[2]) - for i in range(var.shape[0]): - new_s = [float(k) for k in lines[3 + i].split()] - new_s[pos] = var[i] - - line_new = " " - for k in range(len(new_s)): - line_new += f'{"" if k==0 else " "}{new_s[k]:.6e}' - f.write(line_new + "\n") diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 0e6871db..7f48e5b8 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -1,9 +1,8 @@ import copy import shutil import numpy as np -from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools import IOtools, PLASMAtools from mitim_tools.gacode_tools import TGYROtools -from mitim_modules.powertorch.physics_models import transport_cgyro from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -454,7 +453,7 @@ def curateTGYROfiles( # Merge - transport_cgyro.modifyFLUX( + modifyFLUX( tgyro, folder, Qe, @@ -471,7 +470,7 @@ def curateTGYROfiles( impurityPosition=impurityPosition, ) - transport_cgyro.modifyFLUX( + modifyFLUX( tgyro, folder, QeE, @@ -499,7 +498,7 @@ def curateTGYROfiles( GZTargetE = GeTargetE * 0.0 MtTargetE = abs(tgyro.Mt_tar[0, 1:]) * relativeErrorTAR - transport_cgyro.modifyEVO( + modifyEVO( tgyro, folder, QeTargetE * 0.0, @@ -511,7 +510,7 @@ def curateTGYROfiles( positionMod=1, special_label="_stds", ) - transport_cgyro.modifyEVO( + modifyEVO( tgyro, folder, QeTargetE, @@ -548,3 +547,327 @@ def dummyCDF(GeneralFolder, FolderEvaluation): cdf = FolderEvaluation / f"{subname}_ev{name}.CDF" return cdf + +def modifyResults( + Qe, + Qi, + Ge, + GZ, + Mt, + Pexch, + QeE, + QiE, + GeE, + GZE, + MtE, + PexchE, + tgyro, + folder_tgyro, + minErrorPercent=5.0, + percentNeo=2.0, + useConvectiveFluxes=False, + Qi_criterion_stable=0.0025, + impurityPosition=3, + OriginalFimp=1.0, +): + """ + All in real units, with dimensions of (rho) from axis to edge + """ + + # If a plasma is very close to stable... do something about error + if minErrorPercent is not None: + ( + Qe_target, + Qi_target, + Ge_target_special, + GZ_target_special, + Mt_target, + ) = defineReferenceFluxes( + tgyro, + useConvectiveFluxes=useConvectiveFluxes, + impurityPosition=impurityPosition, + ) + + Qe_min = Qe_target * (minErrorPercent / 100.0) + Qi_min = Qi_target * (minErrorPercent / 100.0) + Ge_min = Ge_target_special * (minErrorPercent / 100.0) + GZ_min = GZ_target_special * (minErrorPercent / 100.0) + Mt_min = Mt_target * (minErrorPercent / 100.0) + + for i in range(Qe.shape[0]): + if Qi[i] < Qi_criterion_stable: + print( + f"\t- Based on 'Qi_criterion_stable', plasma considered stable (Qi = {Qi[i]:.2e} < {Qi_criterion_stable:.2e} MW/m2) at position #{i}, using minimum errors of {minErrorPercent}% of targets", + typeMsg="w", + ) + QeE[i] = Qe_min[i] + print(f"\t\t* QeE = {QeE[i]}") + QiE[i] = Qi_min[i] + print(f"\t\t* QiE = {QiE[i]}") + GeE[i] = Ge_min[i] + print(f"\t\t* GeE = {GeE[i]}") + GZE[i] = GZ_min[i] + print(f"\t\t* GZE = {GZE[i]}") + MtE[i] = Mt_min[i] + print(f"\t\t* MtE = {MtE[i]}") + + # Heat fluxes + QeTot = Qe + tgyro.Qe_sim_neo[0, 1:] + QiTot = Qi + tgyro.QiIons_sim_neo_thr[0, 1:] + + # Particle fluxes + PeTot = Ge + tgyro.Ge_sim_neo[0, 1:] + PZTot = GZ + tgyro.Gi_sim_neo[impurityPosition, 0, 1:] + + # Momentum fluxes + MtTot = Mt + tgyro.Mt_sim_neo[0, 1:] + + # ************************************************************************************ + # **** Modify complete folder (Division of ion fluxes will be wrong, since I put everything in first ion) + # ************************************************************************************ + + # 1. Modify out.tgyro.evo files (which contain turb+neo summed together) + + print(f"\t- Modifying TGYRO out.tgyro.evo files in {IOtools.clipstr(folder_tgyro)}") + modifyEVO( + tgyro, + folder_tgyro, + QeTot, + QiTot, + PeTot, + PZTot, + MtTot, + impurityPosition=impurityPosition, + ) + + # 2. Modify out.tgyro.flux files (which contain turb and neo separated) + + print(f"\t- Modifying TGYRO out.tgyro.flux files in {folder_tgyro}") + modifyFLUX( + tgyro, + folder_tgyro, + Qe, + Qi, + Ge, + GZ, + Mt, + Pexch, + impurityPosition=impurityPosition, + ) + + # 3. Modify files for errors + + print(f"\t- Modifying TGYRO out.tgyro.flux_stds in {folder_tgyro}") + modifyFLUX( + tgyro, + folder_tgyro, + QeE, + QiE, + GeE, + GZE, + MtE, + PexchE, + impurityPosition=impurityPosition, + special_label="_stds", + ) + + +def modifyEVO( + tgyro, + folder, + QeT, + QiT, + GeT, + GZT, + MtT, + impurityPosition=3, + positionMod=1, + special_label=None, +): + QeTGB = QeT / tgyro.Q_GB[-1, 1:] + QiTGB = QiT / tgyro.Q_GB[-1, 1:] + GeTGB = GeT / tgyro.Gamma_GB[-1, 1:] + GZTGB = GZT / tgyro.Gamma_GB[-1, 1:] + MtTGB = MtT / tgyro.Pi_GB[-1, 1:] + + modTGYROfile(folder / "out.tgyro.evo_te", QeTGB, pos=positionMod, fileN_suffix=special_label) + modTGYROfile(folder / "out.tgyro.evo_ti", QiTGB, pos=positionMod, fileN_suffix=special_label) + modTGYROfile(folder / "out.tgyro.evo_ne", GeTGB, pos=positionMod, fileN_suffix=special_label) + modTGYROfile(folder / "out.tgyro.evo_er", MtTGB, pos=positionMod, fileN_suffix=special_label) + + for i in range(tgyro.Qi_sim_turb.shape[0]): + if i == impurityPosition: + var = GZTGB + else: + var = GZTGB * 0.0 + modTGYROfile( + folder / f"out.tgyro.evo_n{i+1}", + var, + pos=positionMod, + fileN_suffix=special_label, + ) + + +def modifyFLUX( + tgyro, + folder, + Qe, + Qi, + Ge, + GZ, + Mt, + S, + QeNeo=None, + QiNeo=None, + GeNeo=None, + GZNeo=None, + MtNeo=None, + impurityPosition=3, + special_label=None, +): + folder = IOtools.expandPath(folder) + + QeGB = Qe / tgyro.Q_GB[-1, 1:] + QiGB = Qi / tgyro.Q_GB[-1, 1:] + GeGB = Ge / tgyro.Gamma_GB[-1, 1:] + GZGB = GZ / tgyro.Gamma_GB[-1, 1:] + MtGB = Mt / tgyro.Pi_GB[-1, 1:] + SGB = S / tgyro.S_GB[-1, 1:] + + # ****************************************************************************************** + # Electrons + # ****************************************************************************************** + + # Particle flux: Update + + modTGYROfile(folder / "out.tgyro.flux_e", GeGB, pos=2, fileN_suffix=special_label) + if GeNeo is not None: + GeGB_neo = GeNeo / tgyro.Gamma_GB[-1, 1:] + modTGYROfile(folder / "out.tgyro.flux_e", GeGB_neo, pos=1, fileN_suffix=special_label) + + # Energy flux: Update + + modTGYROfile(folder / "out.tgyro.flux_e", QeGB, pos=4, fileN_suffix=special_label) + if QeNeo is not None: + QeGB_neo = QeNeo / tgyro.Q_GB[-1, 1:] + modTGYROfile(folder / "out.tgyro.flux_e", QeGB_neo, pos=3, fileN_suffix=special_label) + + # Rotation: Remove (it will be sum to the first ion) + + modTGYROfile(folder / "out.tgyro.flux_e", GeGB * 0.0, pos=6, fileN_suffix=special_label) + modTGYROfile(folder / "out.tgyro.flux_e", GeGB * 0.0, pos=5, fileN_suffix=special_label) + + # Energy exchange + + modTGYROfile(folder / "out.tgyro.flux_e", SGB, pos=7, fileN_suffix=special_label) + + # SMW = S # S is MW/m^3 + # modTGYROfile(f'{folder}/out.tgyro.power_e',SMW,pos=8,fileN_suffix=special_label) + # print('\t\t- Modified turbulent energy exchange in out.tgyro.power_e') + + # ****************************************************************************************** + # Ions + # ****************************************************************************************** + + # Energy flux: Update + + modTGYROfile(folder / "out.tgyro.flux_i1", QiGB, pos=4, fileN_suffix=special_label) + + if QiNeo is not None: + QiGB_neo = QiNeo / tgyro.Q_GB[-1, 1:] + modTGYROfile(folder / "out.tgyro.flux_i1", QiGB_neo, pos=3, fileN_suffix=special_label) + + # Particle flux: Make ion particle fluxes zero, because I don't want to mistake TGLF with CGYRO when looking at tgyro results + + for i in range(tgyro.Qi_sim_turb.shape[0]): + if tgyro.profiles.Species[i]["S"] == "therm": + var = QiGB * 0.0 + modTGYROfile(folder / f"out.tgyro.flux_i{i+1}",var,pos=2,fileN_suffix=special_label,) # Gi_turb + modTGYROfile(folder / f"out.tgyro.evo_n{i+1}", var, pos=1, fileN_suffix=special_label) # Gi (Gi_sim) + + if i != impurityPosition: + modTGYROfile(folder / f"out.tgyro.flux_i{i+1}",var,pos=1,fileN_suffix=special_label) # Gi_neo + + # Rotation: Update + + modTGYROfile(folder / "out.tgyro.flux_i1", MtGB, pos=6, fileN_suffix=special_label) + + if MtNeo is not None: + MtGB_neo = MtNeo / tgyro.Pi_GB[-1, 1:] + modTGYROfile(folder / "out.tgyro.flux_i1", MtGB_neo, pos=5, fileN_suffix=special_label) + + # Energy exchange: Remove (it will be the electrons one) + + modTGYROfile(folder / "out.tgyro.flux_i1", SGB * 0.0, pos=7, fileN_suffix=special_label) + + # ****************************************************************************************** + # Impurities + # ****************************************************************************************** + + # Remove everything from all the rest of non-first ions (except the particles for the impurity chosen) + + for i in range(tgyro.Qi_sim_turb.shape[0] - 1): + if tgyro.profiles.Species[i + 1]["S"] == "therm": + var = QiGB * 0.0 + for pos in [3, 4, 5, 6, 7]: + modTGYROfile(folder / f"out.tgyro.flux_i{i+2}",var,pos=pos,fileN_suffix=special_label) + for pos in [1, 2]: + if i + 2 != impurityPosition: + modTGYROfile(folder / f"out.tgyro.flux_i{i+2}",var,pos=pos,fileN_suffix=special_label) + + modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB,pos=2,fileN_suffix=special_label) + if GZNeo is not None: + GZGB_neo = GZNeo / tgyro.Gamma_GB[-1, 1:] + modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB_neo,pos=1,fileN_suffix=special_label) + + +def modTGYROfile(file, var, pos=0, fileN_suffix=None): + fileN = file if fileN_suffix is None else file.parent / f"{file.name}{fileN_suffix}" + + if not fileN.exists(): + shutil.copy2(file, fileN) + + with open(fileN, "r") as f: + lines = f.readlines() + + with open(fileN, "w") as f: + f.write(lines[0]) + f.write(lines[1]) + f.write(lines[2]) + for i in range(var.shape[0]): + new_s = [float(k) for k in lines[3 + i].split()] + new_s[pos] = var[i] + + line_new = " " + for k in range(len(new_s)): + line_new += f'{"" if k==0 else " "}{new_s[k]:.6e}' + f.write(line_new + "\n") + +def defineReferenceFluxes( + tgyro, factor_tauptauE=5, useConvectiveFluxes=False, impurityPosition=3 +): + Qe_target = abs(tgyro.Qe_tar[0, 1:]) + Qi_target = abs(tgyro.Qi_tar[0, 1:]) + Mt_target = abs(tgyro.Mt_tar[0, 1:]) + + # For particle fluxes, since the targets are often zero... it's more complicated + QeMW_target = abs(tgyro.Qe_tarMW[0, 1:]) + QiMW_target = abs(tgyro.Qi_tarMW[0, 1:]) + We, Wi, Ne, NZ = tgyro.profiles.deriveContentByVolumes( + rhos=tgyro.rho[0, 1:], impurityPosition=impurityPosition + ) + + tau_special = ( + (We + Wi) / (QeMW_target + QiMW_target) * factor_tauptauE + ) # tau_p in seconds + Ge_target_special = (Ne / tau_special) / tgyro.dvoldr[0, 1:] # (1E20/seconds/m^2) + + if useConvectiveFluxes: + Ge_target_special = PLASMAtools.convective_flux( + tgyro.Te[0, 1:], Ge_target_special + ) # (1E20/seconds/m^2) + + GZ_target_special = Ge_target_special * NZ / Ne + + return Qe_target, Qi_target, Ge_target_special, GZ_target_special, Mt_target + diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 271fc2c9..26dcf732 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -126,7 +126,7 @@ def calculator( rederive_profiles=False, ) - p.plasma["Pin"] = ( + p.plasma["QiMWm2n"] = ( (p.plasma["Paux_e"] + p.plasma["Paux_i"]) * p.plasma["volp"] )[..., -1] p.plasma["Q"] = p.plasma["Pfus"] / p.plasma["Pin"] diff --git a/src/mitim_modules/powertorch/scripts/compareRadialResolution.py b/src/mitim_modules/powertorch/scripts/compareRadialResolution.py index db600e89..bbdeea21 100644 --- a/src/mitim_modules/powertorch/scripts/compareRadialResolution.py +++ b/src/mitim_modules/powertorch/scripts/compareRadialResolution.py @@ -151,7 +151,7 @@ ax.legend() ax = axs[0, 1] -varsS = ["Pe", "Pi"] +varsS = ["QeMWm2", "QiMWm2"] s, lab = sF, "Fine " for var in varsS: diff --git a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py index 41155bab..fe103393 100644 --- a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py +++ b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py @@ -154,13 +154,13 @@ ax = axs[0, 2] ax.plot(t.rho[0], t.Qe_tar[0], "s-", lw=0.5, label="TGYRO Pe", markersize=markersize) -P = s.plasma["Pe"] +P = s.plasma["QeMWm2"] ax.plot(s.plasma["rho"][0], P[0], ls, lw=0.5, label="STATE Pe", markersize=markersize) MaxError = np.nanmax(np.abs(t.Qe_tarMW[0] - P[0].cpu().numpy()) / t.Qe_tarMW[0] * 100.0) print(f"{MaxError = :.3f} %") ax.plot(t.rho[0], t.Qi_tar[0], "s-", lw=0.5, label="TGYRO Pi", markersize=markersize) -P = s.plasma["Pi"] +P = s.plasma["QiMWm2"] ax.plot(s.plasma["rho"][0], P[0], ls, lw=0.5, label="STATE Pi", markersize=markersize) MaxError = np.nanmax(np.abs(t.Qi_tarMW[0] - P[0].cpu().numpy()) / t.Qi_tarMW[0] * 100.0) print(f"{MaxError = :.3f} %") @@ -178,13 +178,13 @@ ax = axs[1, 2] ax.plot(t.rho[0], t.Qe_tarMW[0], "s-", lw=0.5, label="TGYRO Pe", markersize=markersize) -P = s.plasma["Pe"] * s.plasma["volp"] +P = s.plasma["QeMWm2"] * s.plasma["volp"] ax.plot(s.plasma["rho"][0], P[0], ls, lw=0.5, label="STATE Pe", markersize=markersize) MaxError = np.nanmax(np.abs(t.Qe_tarMW[0] - P[0].cpu().numpy()) / t.Qe_tarMW[0] * 100.0) print(f"{MaxError = :.3f} %") ax.plot(t.rho[0], t.Qi_tarMW[0], "s-", lw=0.5, label="TGYRO Pi", markersize=markersize) -P = s.plasma["Pi"] * s.plasma["volp"] +P = s.plasma["QiMWm2"] * s.plasma["volp"] ax.plot(s.plasma["rho"][0], P[0], ls, lw=0.5, label="STATE Pi", markersize=markersize) MaxError = np.nanmax(np.abs(t.Qi_tarMW[0] - P[0].cpu().numpy()) / t.Qi_tarMW[0] * 100.0) print(f"{MaxError = :.3f} %") diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 526562f6..8e9680b0 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -25,13 +25,13 @@ def __init__(self,powerstate): if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 1: self.Pe_orig, self.Pi_orig = ( - self.powerstate.plasma["Pe_orig_fusradexch"], - self.powerstate.plasma["Pi_orig_fusradexch"], + self.powerstate.plasma["QeMWm2_orig_fusradexch"], + self.powerstate.plasma["QiMWm2_orig_fusradexch"], ) # Original integrated from input.gacode elif self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 2: self.Pe_orig, self.Pi_orig = ( - self.powerstate.plasma["Pe_orig_fusrad"], - self.powerstate.plasma["Pi_orig_fusrad"], + self.powerstate.plasma["QeMWm2_orig_fusrad"], + self.powerstate.plasma["QiMWm2_orig_fusrad"], ) elif self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: self.Pe_orig, self.Pi_orig = self.powerstate.plasma["te"] * 0.0, self.powerstate.plasma["te"] * 0.0 @@ -138,10 +138,10 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, # Plug-in Targets # ************************************************************************************************** - self.powerstate.plasma["Pe"] = ( + self.powerstate.plasma["QeMWm2"] = ( self.powerstate.plasma["Paux_e"] + self.P[: self.P.shape[0]//2, :] + self.Pe_orig ) # MW/m^2 - self.powerstate.plasma["Pi"] = ( + self.powerstate.plasma["QiMWm2"] = ( self.powerstate.plasma["Paux_i"] + self.P[self.P.shape[0]//2 :, :] + self.Pi_orig ) # MW/m^2 self.powerstate.plasma["Ce_raw"] = self.CextraE @@ -169,7 +169,7 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, # Error # ************************************************************************************************** - variables_to_error = ["Pe", "Pi", "Ce", "CZ", "Mt", "Ce_raw", "CZ_raw"] + variables_to_error = ["QeMWm2", "QiMWm2", "Ce", "CZ", "Mt", "Ce_raw", "CZ_raw"] for i in variables_to_error: self.powerstate.plasma[i + "_stds"] = abs(self.powerstate.plasma[i]) * assumedPercentError / 100 @@ -182,8 +182,8 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, """ gb_mapping = { - "Pe": "Qgb", - "Pi": "Qgb", + "QeMWm2": "Qgb", + "QiMWm2": "Qgb", "Ce": "Qgb" if useConvectiveFluxes else "Ggb", "CZ": "Qgb" if useConvectiveFluxes else "Ggb", "Mt": "Pgb", diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index ff357f58..89ea572f 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -129,10 +129,10 @@ def gacode_to_powerstate(self, rho_vec=None): self.plasma["Gaux_Z"] = self.plasma["Gaux_e"] * 0.0 quantitites = {} - quantitites["Pe_orig_fusrad"] = input_gacode.derived["qe_fus_MWmiller"] - input_gacode.derived["qrad_MWmiller"] - quantitites["Pi_orig_fusrad"] = input_gacode.derived["qi_fus_MWmiller"] - quantitites["Pe_orig_fusradexch"] = quantitites["Pe_orig_fusrad"] - input_gacode.derived["qe_exc_MWmiller"] - quantitites["Pi_orig_fusradexch"] = quantitites["Pi_orig_fusrad"] + input_gacode.derived["qe_exc_MWmiller"] + quantitites["QeMWm2_orig_fusrad"] = input_gacode.derived["qe_fus_MWmiller"] - input_gacode.derived["qrad_MWmiller"] + quantitites["QiMWm2_orig_fusrad"] = input_gacode.derived["qi_fus_MWmiller"] + quantitites["QeMWm2_orig_fusradexch"] = quantitites["QeMWm2_orig_fusrad"] - input_gacode.derived["qe_exc_MWmiller"] + quantitites["QiMWm2_orig_fusradexch"] = quantitites["QiMWm2_orig_fusrad"] + input_gacode.derived["qe_exc_MWmiller"] for key in quantitites: diff --git a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py index b4e7d657..bf4c187d 100644 --- a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py +++ b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py @@ -157,15 +157,15 @@ def TGYROmodeledVariables(TGYROresults, # *********** Electron Energy Fluxes # ********************************** - powerstate.plasma["Pe_tr_turb"] = torch.Tensor(TGYROresults.Qe_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Pe_tr_neo"] = torch.Tensor(TGYROresults.Qe_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QeMWm2_tr_turb"] = torch.Tensor(TGYROresults.Qe_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QeMWm2_tr_neo"] = torch.Tensor(TGYROresults.Qe_sim_neo[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Pe_tr_turb_stds"] = torch.Tensor(TGYROresults.Qe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Pe_tr_neo_stds"] = torch.Tensor(TGYROresults.Qe_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QeMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Qe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QeMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Qe_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: - powerstate.plasma["Pe"] = torch.Tensor(TGYROresults.Qe_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Pe_stds"] = torch.Tensor(TGYROresults.Qe_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QeMWm2"] = torch.Tensor(TGYROresults.Qe_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QeMWm2_stds"] = torch.Tensor(TGYROresults.Qe_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None # ********************************** # *********** Ion Energy Fluxes @@ -173,23 +173,23 @@ def TGYROmodeledVariables(TGYROresults, if includeFast: - powerstate.plasma["Pi_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Pi_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Pi_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Pi_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QiMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None else: - powerstate.plasma["Pi_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Pi_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Pi_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Pi_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QiMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: - powerstate.plasma["Pi"] = torch.Tensor(TGYROresults.Qi_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Pi_stds"] = torch.Tensor(TGYROresults.Qi_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QiMWm2"] = torch.Tensor(TGYROresults.Qi_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_stds"] = torch.Tensor(TGYROresults.Qi_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None # ********************************** # *********** Momentum Fluxes @@ -291,8 +291,8 @@ def TGYROmodeledVariables(TGYROresults, powerstate.plasma["PexchTurb"] = torch.Tensor(TGYROresults.EXe_sim_turb[:, :nr]).to(powerstate.dfT) powerstate.plasma["PexchTurb_stds"] = torch.Tensor(TGYROresults.EXe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None else: - powerstate.plasma["PexchTurb"] = powerstate.plasma["Pe_tr_turb"] * 0.0 - powerstate.plasma["PexchTurb_stds"] = powerstate.plasma["Pe_tr_turb"] * 0.0 + powerstate.plasma["PexchTurb"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + powerstate.plasma["PexchTurb_stds"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 # ********************************** # *********** Traget extra @@ -329,18 +329,18 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): var_dict = {} mapper = { - "QeTurb": "Pe_tr_turb", - "QiTurb": "Pi_tr_turb", + "QeTurb": "QeMWm2_tr_turb", + "QiTurb": "QiMWm2_tr_turb", "GeTurb": "Ce_tr_turb", "GZTurb": "CZ_tr_turb", "MtTurb": "Mt_tr_turb", - "QeNeo": "Pe_tr_neo", - "QiNeo": "Pi_tr_neo", + "QeNeo": "QeMWm2_tr_neo", + "QiNeo": "QiMWm2_tr_neo", "GeNeo": "Ce_tr_neo", "GZNeo": "CZ_tr_neo", "MtNeo": "Mt_tr_neo", - "QeTar": "Pe", - "QiTar": "Pi", + "QeTar": "QeMWm2", + "QiTar": "QiMWm2", "GeTar": "Ce", "GZTar": "CZ", "MtTar": "Mt", @@ -463,18 +463,18 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): # Prepare dictionary from powerstate (for use in Analysis) mapper = { - "QeTurb": "Pe_tr_turb", - "QiTurb": "Pi_tr_turb", + "QeTurb": "QeMWm2_tr_turb", + "QiTurb": "QiMWm2_tr_turb", "GeTurb": "Ce_tr_turb", "GZTurb": "CZ_tr_turb", "MtTurb": "Mt_tr_turb", - "QeNeo": "Pe_tr_neo", - "QiNeo": "Pi_tr_neo", + "QeNeo": "QeMWm2_tr_neo", + "QiNeo": "QiMWm2_tr_neo", "GeNeo": "Ce_tr_neo", "GZNeo": "CZ_tr_neo", "MtNeo": "Mt_tr_neo", - "QeTar": "Pe", - "QiTar": "Pi", + "QeTar": "QeMWm2", + "QiTar": "QiMWm2", "GeTar": "Ce", "GZTar": "CZ", "MtTar": "Mt", From b9f33bdeced014f508afbb370ddcc1a989e79f6a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 09:15:51 -0700 Subject: [PATCH 017/385] Do not ask the user if the folder Outputs is empty --- .../physics_models/transport_cgyro.py | 24 +------------------ src/mitim_tools/opt_tools/STRATEGYtools.py | 14 +++++------ 2 files changed, 8 insertions(+), 30 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 0a475b74..5ae7f214 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -485,26 +485,4 @@ def readCGYROresults(file, radii, unnormalize=True): p[r] += 1 - return ( - aLTe, - aLTi, - aLne, - Q_gb, - Qe, - Qi, - Ge, - GZ, - Mt, - Pexch, - Qe_std, - Qi_std, - Ge_std, - GZ_std, - Mt_std, - Pexch_std, - tstart, - tend, - ) - - - + return aLTe,aLTi,aLne,Q_gb,Qe,Qi,Ge,GZ,Mt,Pexch,Qe_std,Qi_std,Ge_std,GZ_std,Mt_std,Pexch_std,tstart,tend diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index ad1807ab..0e72c686 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -362,13 +362,6 @@ def __init__( self.seed = seed self.avoidPoints = [] - if (not self.cold_start) and askQuestions: - if not print( - f"\t* Because {cold_start = }, MITIM will try to read existing results from folder", - typeMsg="q", - ): - raise Exception("[MITIM] - User requested to stop") - if self.optimization_object.name_objectives is None: self.optimization_object.name_objectives = "y" @@ -381,6 +374,13 @@ def __init__( self.folderOutputs = self.folderExecution / "Outputs" + if (not self.cold_start) and askQuestions: + + # Check if Outputs folder is empty (if it's empty, do not ask the user, just continue) + if self.folderOutputs.exists() and (len(list(self.folderOutputs.iterdir())) > 0): + if not print(f"\t* Because {cold_start = }, MITIM will try to read existing results from folder",typeMsg="q"): + raise Exception("[MITIM] - User requested to stop") + if optimization_object.optimization_options is not None: if not self.folderOutputs.exists(): IOtools.askNewFolder(self.folderOutputs, force=True) From 6ab0d0eb7bb68eb872708d1b2760236f01d69352 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 09:26:18 -0700 Subject: [PATCH 018/385] Further Pe->QeMWm2 transition --- .../portals/utils/PORTALSplot.py | 36 +++++++++---------- .../powertorch/utils/POWERplot.py | 4 +-- .../powertorch/utils/TRANSPORTtools.py | 2 +- .../gacode_tools/utils/PORTALSinteraction.py | 2 +- src/mitim_tools/opt_tools/STRATEGYtools.py | 33 +++++------------ 5 files changed, 30 insertions(+), 47 deletions(-) diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index 22ec22f2..e0541935 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -217,23 +217,23 @@ def PORTALSanalyzer_plotMetrics( if axTe_f is not None: axTe_f.plot( rho, - power.plasma['Pe_tr_turb'].cpu().numpy() + power.plasma['Pe_tr_neo'].cpu().numpy(), + power.plasma['QeMWm2_tr_turb'].cpu().numpy() + power.plasma['QeMWm2_tr_neo'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph, ) - axTe_f.plot(rho, power.plasma['Pe'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) + axTe_f.plot(rho, power.plasma['QeMWm2'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) if axTi_f is not None: axTi_f.plot( rho, - power.plasma['Pi_tr_turb'].cpu().numpy() + power.plasma['Pi_tr_neo'].cpu().numpy(), + power.plasma['QiMWm2_tr_turb'].cpu().numpy() + power.plasma['QiMWm2_tr_neo'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph, ) - axTi_f.plot(rho, power.plasma['Pi'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) + axTi_f.plot(rho, power.plasma['QiMWm2'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) if axne_f is not None: @@ -2882,7 +2882,7 @@ def plotFluxComparison( if axTe_f is not None: axTe_f.plot( r[0][ixF:], - power.plasma['Pe_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Pe_tr_neo'].cpu().numpy()[0][ixF:], + power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:], "-s", c=col, lw=2, @@ -2891,10 +2891,10 @@ def plotFluxComparison( alpha=alpha, ) - sigma = power.plasma['Pe_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['Pe_tr_neo_stds'].cpu().numpy()[0][ixF:] + sigma = power.plasma['QeMWm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo_stds'].cpu().numpy()[0][ixF:] - m_Qe, M_Qe = (power.plasma['Pe_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Pe_tr_neo'].cpu().numpy()[0][ixF:]) - stds * sigma, ( - power.plasma['Pe_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Pe_tr_neo'].cpu().numpy()[0][ixF:] + m_Qe, M_Qe = (power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:]) - stds * sigma, ( + power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:] ) + stds * sigma axTe_f.fill_between(r[0][ixF:], m_Qe, M_Qe, facecolor=col, alpha=alpha / 3) @@ -2905,7 +2905,7 @@ def plotFluxComparison( if axTi_f is not None: axTi_f.plot( r[0][ixF:], - power.plasma['Pi_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Pi_tr_neo'].cpu().numpy()[0][ixF:], + power.plasma['QiMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neo'].cpu().numpy()[0][ixF:], "-s", markersize=msFlux, c=col, @@ -2915,13 +2915,13 @@ def plotFluxComparison( ) sigma = ( - power.plasma['Pi_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['Pi_tr_neo_stds'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neo_stds'].cpu().numpy()[0][ixF:] ) m_Qi, M_Qi = ( - power.plasma['Pi_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Pi_tr_neo'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neo'].cpu().numpy()[0][ixF:] ) - stds * sigma, ( - power.plasma['Pi_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Pi_tr_neo'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neo'].cpu().numpy()[0][ixF:] ) + stds * sigma axTi_f.fill_between(r[0][ixF:], m_Qi, M_Qi, facecolor=col, alpha=alpha / 3) @@ -3005,8 +3005,8 @@ def plotFluxComparison( # Retrieve targets ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - Qe_tar = power.plasma['Pe'].cpu().numpy()[0][ixF:] - Qi_tar = power.plasma['Pi'].cpu().numpy()[0][ixF:] + Qe_tar = power.plasma['QeMWm2'].cpu().numpy()[0][ixF:] + Qi_tar = power.plasma['QiMWm2'].cpu().numpy()[0][ixF:] Ge_tar = power.plasma['Ce_raw'].cpu().numpy()[0][ixF:] * (1-int(forceZeroParticleFlux)) GZ_tar = power.plasma['CZ_raw'].cpu().numpy()[0][ixF:] Mt_tar = power.plasma['Mt'].cpu().numpy()[0][ixF:] @@ -3126,7 +3126,7 @@ def plotFluxComparison( if axTe_f is not None: (l1,) = axTe_f.plot( r[0][ixF:], - power.plasma['Pe_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Pe_tr_neo'].cpu().numpy()[0][ixF:], + power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:], "-", c="k", lw=2, @@ -3134,12 +3134,12 @@ def plotFluxComparison( label="Transport", ) (l2,) = axTe_f.plot( - r[0][ixF:], power.plasma['Pe'].cpu().numpy()[0][ixF:], "--*", c="k", lw=2, markersize=0, label="Target" + r[0][ixF:], power.plasma['QeMWm2'].cpu().numpy()[0][ixF:], "--*", c="k", lw=2, markersize=0, label="Target" ) l3 = axTe_f.fill_between( r[0][ixF:], - (power.plasma['Pe_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Pe_tr_neo'].cpu().numpy()[0][ixF:]) - stds, - (power.plasma['Pe_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Pe_tr_neo'].cpu().numpy()[0][ixF:]) + stds, + (power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:]) - stds, + (power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:]) + stds, facecolor="k", alpha=0.3, ) diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index 8e9da478..f51a355d 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -25,12 +25,12 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co if "te" in self.ProfilesPredicted: set_plots.append( - [ 'te', 'aLte', 'Pe_tr', 'Pe', + [ 'te', 'aLte', 'QeMWm2_tr', 'QeMWm2', 'Electron Temperature','$T_e$ (keV)','$a/LT_e$','$Q_e$ (GB)','$Q_e$ ($MW/m^2$)', 1.0,"Qgb"]) if "ti" in self.ProfilesPredicted: set_plots.append( - [ 'ti', 'aLti', 'Pi_tr', 'Pi', + [ 'ti', 'aLti', 'QiMWm2_tr', 'QiMWm2', 'Ion Temperature','$T_i$ (keV)','$a/LT_i$','$Q_i$ (GB)','$Q_i$ ($MW/m^2$)', 1.0,"Qgb"]) if "ne" in self.ProfilesPredicted: diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 11e8e0b6..42d75109 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -22,7 +22,7 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ self.powerstate = powerstate # Allowed fluxes in powerstate so far - self.quantities = ['Pe', 'Pi', 'Ce', 'CZ', 'Mt'] + self.quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'Mt'] # Each flux has a turbulent and neoclassical component self.variables = [f'{i}_tr_turb' for i in self.quantities] + [f'{i}_tr_neo' for i in self.quantities] diff --git a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py index bf4c187d..4a96ece2 100644 --- a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py +++ b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py @@ -306,7 +306,7 @@ def TGYROmodeledVariables(TGYROresults, # Sum here turbulence and neoclassical, after modifications # ------------------------------------------------------------------------------------------------------------------------ - quantities = ['Pe', 'Pi', 'Ce', 'CZ', 'Mt', 'Ce_raw', 'CZ_raw'] + quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'Mt', 'Ce_raw', 'CZ_raw'] for ikey in quantities: powerstate.plasma[ikey+"_tr"] = powerstate.plasma[ikey+"_tr_turb"] + powerstate.plasma[ikey+"_tr_neo"] diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 0e72c686..49c0d40b 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -690,22 +690,14 @@ def run(self): current_step = self.read() if current_step is None: - print( - "\t* Because reading pkl step had problems, disabling cold_starting-from-previous from this point on", - typeMsg="w", - ) - print( - "\t* Are you aware of the consequences of continuing?", - typeMsg="q", - ) + print("\t* Because reading pkl step had problems, disabling cold_starting-from-previous from this point on",typeMsg="w") + print("\t* Are you aware of the consequences of continuing?",typeMsg="q") self.cold_start = True if not self.cold_start: # Read next from Tabular - self.x_next, _, _ = self.optimization_data.extract_points( - points=np.arange(len(self.train_X), len(self.train_X) + self.best_points) - ) + self.x_next, _, _ = self.optimization_data.extract_points(points=np.arange(len(self.train_X), len(self.train_X) + self.best_points)) self.x_next = torch.from_numpy(self.x_next).to(self.dfT) # Re-write x_next from the pkl... reason for this is that if optimization is heuristic, I may prefer what was in Tabular @@ -1010,12 +1002,8 @@ def updateSet( """ print("\n~~~~~~~~~~~~~~~ Entering bounds upgrade module ~~~~~~~~~~~~~~~~~~~") print("(if extrapolations were allowed during optimization)") - self.bounds = SBOcorrections.upgradeBounds( - self.bounds, self.train_X, self.avoidPoints_outside - ) - print( - "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n" - ) + self.bounds = SBOcorrections.upgradeBounds(self.bounds, self.train_X, self.avoidPoints_outside) + print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n") # ~~~~~~~~~~~~~~~~~~ # Possible corrections to modeled & optimization region @@ -1110,7 +1098,7 @@ def initializeOptimization(self): if (not self.cold_start) and (self.optimization_data is not None): self.type_initialization = 3 print( - "--> Since cold_start from a previous MITIM has been requested, forcing initialization type to 3 (read from optimization_data)", + "--> Since restart from a previous MITIM has been requested, forcing initialization type to 3 (read from optimization_data)", typeMsg="i", ) @@ -1119,19 +1107,14 @@ def initializeOptimization(self): try: tabExists = len(self.optimization_data.data) >= self.initial_training - print( - f"\t- optimization_data file has {len(self.optimization_data.data)} elements, and initial_training were {self.initial_training}" - ) + print(f"\t- optimization_data file has {len(self.optimization_data.data)} elements, and initial_training were {self.initial_training}") except: tabExists = False print("\n\nCould not read Tabular, because:", typeMsg="w") print(traceback.format_exc()) if not tabExists: - print( - "--> type_initialization 3 requires optimization_data but something failed. Assigning type_initialization=1 and cold_starting from scratch", - typeMsg="i", - ) + print("--> type_initialization 3 requires optimization_data but something failed. Assigning type_initialization=1 and cold_starting from scratch",typeMsg="i",) if self.askQuestions: flagger = print("Are you sure?", typeMsg="q") if not flagger: From d2b989339d029e02dd8d848e5b730e27ef2908f7 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 09:37:36 -0700 Subject: [PATCH 019/385] Towards Pe -> Qe transition (2) --- src/mitim_modules/portals/utils/PORTALSinit.py | 4 ++-- .../powertorch/utils/TARGETStools.py | 15 +++++---------- .../gacode_tools/utils/PORTALSinteraction.py | 2 -- 3 files changed, 7 insertions(+), 14 deletions(-) diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 96bb1bb3..265a86b7 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -460,9 +460,9 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue Variables[output].append(ikey) elif typ in ["QeTar"]: - Variables[output] = ["PeGB"] + Variables[output] = ["QeGB"] elif typ in ["QiTar"]: - Variables[output] = ["PiGB"] + Variables[output] = ["QiGB"] elif typ in ["GeTar"]: Variables[output] = ["CeGB"] elif typ in ["GZTar"]: diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 8e9680b0..910d0d50 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -181,13 +181,8 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, Note: This is useful for mitim surrogate variables of targets """ - gb_mapping = { - "QeMWm2": "Qgb", - "QiMWm2": "Qgb", - "Ce": "Qgb" if useConvectiveFluxes else "Ggb", - "CZ": "Qgb" if useConvectiveFluxes else "Ggb", - "Mt": "Pgb", - } - - for i in gb_mapping.keys(): - self.powerstate.plasma[f"{i}GB"] = self.powerstate.plasma[i] / self.powerstate.plasma[gb_mapping[i]] + self.powerstate.plasma["QeGB"] = self.powerstate.plasma["QeMWm2"] / self.powerstate.plasma["Qgb"] + self.powerstate.plasma["QiGB"] = self.powerstate.plasma["QiMWm2"] / self.powerstate.plasma["Qgb"] + self.powerstate.plasma["CeGB"] = self.powerstate.plasma["Ce"] / self.powerstate.plasma["Qgb" if useConvectiveFluxes else "Ggb"] + self.powerstate.plasma["CZGB"] = self.powerstate.plasma["CZ"] / self.powerstate.plasma["Qgb" if useConvectiveFluxes else "Ggb"] + self.powerstate.plasma["MtGB"] = self.powerstate.plasma["Mt"] / self.powerstate.plasma["Pgb"] diff --git a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py index 4a96ece2..f5b89a8e 100644 --- a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py +++ b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py @@ -147,12 +147,10 @@ def TGYROmodeledVariables(TGYROresults, if UseFineGridTargets: TGYROresults.useFineGridTargets(impurityPosition=impurityPosition) - nr = powerstate.plasma['rho'].shape[-1] if powerstate.plasma['rho'].shape[-1] != TGYROresults.rho.shape[-1]: print('\t- TGYRO was run with an extra point in the grid, treating it carefully now') - # ********************************** # *********** Electron Energy Fluxes # ********************************** From a1a4ecdfaf83465e656426ceff3dd47c9cb49583 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 10:05:25 -0700 Subject: [PATCH 020/385] Turb and Neo naming to _tr_turb and _tr_neo --- src/mitim_modules/portals/PORTALSmain.py | 36 +++------ .../portals/utils/PORTALSanalysis.py | 2 +- .../portals/utils/PORTALSinit.py | 34 ++++----- .../physics_models/transport_tgyro.py | 76 +++++++++---------- .../gacode_tools/utils/PORTALSinteraction.py | 52 ++++++------- .../opt_tools/scripts/evaluate_model.py | 6 +- 6 files changed, 92 insertions(+), 114 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 661dcc18..8b8ae44c 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -420,7 +420,7 @@ def scalarized_objective(self, Y): ------------------------------------------------------------------------- Prepare transport dictionary ------------------------------------------------------------------------- - Note: var_dict['QeTurb'] must have shape (dim1...N, num_radii) + Note: var_dict['Qe_tr_turb'] must have shape (dim1...N, num_radii) """ var_dict = {} @@ -650,8 +650,6 @@ def runModelEvaluator( return powerstate, dictOFs def map_powerstate_to_portals(powerstate, dictOFs): - """ - """ for var in powerstate.ProfilesPredicted: # Write in OFs @@ -672,19 +670,11 @@ def map_powerstate_to_portals(powerstate, dictOFs): --------------------- """ - dictOFs[f"{var0}Turb_{i+1}"]["value"] = powerstate.plasma[ - f"{var1}_tr_turb" - ][0, i+1] - dictOFs[f"{var0}Turb_{i+1}"]["error"] = powerstate.plasma[ - f"{var1}_tr_turb_stds" - ][0, i+1] + dictOFs[f"{var0}_tr_turb_{i+1}"]["value"] = powerstate.plasma[f"{var1}_tr_turb"][0, i+1] + dictOFs[f"{var0}_tr_turb_{i+1}"]["error"] = powerstate.plasma[f"{var1}_tr_turb_stds"][0, i+1] - dictOFs[f"{var0}Neo_{i+1}"]["value"] = powerstate.plasma[ - f"{var1}_tr_neo" - ][0, i+1] - dictOFs[f"{var0}Neo_{i+1}"]["error"] = powerstate.plasma[ - f"{var1}_tr_neo_stds" - ][0, i+1] + dictOFs[f"{var0}_tr_neo_{i+1}"]["value"] = powerstate.plasma[f"{var1}_tr_neo"][0, i+1] + dictOFs[f"{var0}_tr_neo_{i+1}"]["error"] = powerstate.plasma[f"{var1}_tr_neo_stds"][0, i+1] """ TARGET calculation @@ -692,12 +682,8 @@ def map_powerstate_to_portals(powerstate, dictOFs): If that radius & profile position has target, evaluate """ - dictOFs[f"{var0}Tar_{i+1}"]["value"] = powerstate.plasma[f"{var1}"][ - 0, i+1 - ] - dictOFs[f"{var0}Tar_{i+1}"]["error"] = powerstate.plasma[ - f"{var1}_stds" - ][0, i+1] + dictOFs[f"{var0}Tar_{i+1}"]["value"] = powerstate.plasma[f"{var1}"][0, i+1] + dictOFs[f"{var0}Tar_{i+1}"]["error"] = powerstate.plasma[f"{var1}_stds"][0, i+1] """ Turbulent Exchange @@ -705,12 +691,8 @@ def map_powerstate_to_portals(powerstate, dictOFs): """ if 'PexchTurb_1' in dictOFs: for i in range(powerstate.plasma["rho"].shape[1] - 1): - dictOFs[f"PexchTurb_{i+1}"]["value"] = powerstate.plasma["PexchTurb"][ - 0, i+1 - ] - dictOFs[f"PexchTurb_{i+1}"]["error"] = powerstate.plasma[ - "PexchTurb_stds" - ][0, i+1] + dictOFs[f"PexchTurb_{i+1}"]["value"] = powerstate.plasma["PexchTurb"][0, i+1] + dictOFs[f"PexchTurb_{i+1}"]["error"] = powerstate.plasma["PexchTurb_stds"][0, i+1] return dictOFs diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 5d3af5af..202641ba 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -463,7 +463,7 @@ def extractModels(self, step=-1): 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 = models['Qe_tr_turb_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 diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 265a86b7..d140d773 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -293,8 +293,8 @@ def initializeProblem( var = "Mt" for i in range(len(portals_fun.MODELparameters["RhoLocations"])): - ofs.append(f"{var}Turb_{i+1}") - ofs.append(f"{var}Neo_{i+1}") + ofs.append(f"{var}_tr_turb_{i+1}") + ofs.append(f"{var}_tr_neo_{i+1}") ofs.append(f"{var}Tar_{i+1}") @@ -353,35 +353,31 @@ def initializeProblem( def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValues=False): allOuts = portals_fun.optimization_options["problem_options"]["ofs"] portals_transformation_variables = portals_fun.PORTALSparameters["portals_transformation_variables"][ikey] - portals_transformation_variables_trace = portals_fun.PORTALSparameters[ - "portals_transformation_variables_trace" - ][ikey] - additional_params_in_surrogate = portals_fun.PORTALSparameters[ - "additional_params_in_surrogate" - ] + portals_transformation_variables_trace = portals_fun.PORTALSparameters["portals_transformation_variables_trace"][ikey] + additional_params_in_surrogate = portals_fun.PORTALSparameters["additional_params_in_surrogate"] Variables = {} for output in allOuts: if IOtools.isfloat(output): continue - typ, num = output.split("_") - pos = int(num) + typ = '_'.join(output.split("_")[:-1]) + pos = int(output.split("_")[-1]) if typ in [ "Qe", - "QeTurb", - "QeNeo", + "Qe_tr_turb", + "Qe_tr_neo", "Qi", - "QiTurb", - "QiNeo", + "Qi_tr_turb", + "Qi_tr_neo", "Ge", - "GeTurb", - "GeNeo", + "Ge_tr_turb", + "Ge_tr_neo", "PexchTurb", "Mt", - "MtTurb", - "MtNeo", + "Mt_tr_turb", + "Mt_tr_neo", ]: if doNotFitOnFixedValues: isAbsValFixed = pos == ( @@ -419,7 +415,7 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue if useThisOne: Variables[output].append(ikey) - elif typ in ["GZ", "GZTurb", "GZNeo"]: + elif typ in ["GZ", "GZ_tr_turb", "GZ_tr_neo"]: if doNotFitOnFixedValues: isAbsValFixed = pos == ( portals_fun.powerstate.plasma["rho"].shape[-1] - 1 diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 7f48e5b8..55d0ccba 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -433,23 +433,23 @@ def curateTGYROfiles( # Neo # ************************************************************************************************************************** - QeNeo = tgyro.Qe_sim_neo[0, 1:] + Qe_tr_neo = tgyro.Qe_sim_neo[0, 1:] if includeFast: - QiNeo = tgyro.QiIons_sim_neo[0, 1:] + Qi_tr_neo = tgyro.QiIons_sim_neo[0, 1:] else: - QiNeo = tgyro.QiIons_sim_neo_thr[0, 1:] - GeNeo = tgyro.Ge_sim_neo[0, 1:] - GZNeo = tgyro.Gi_sim_neo[impurityPosition, 0, 1:] - MtNeo = tgyro.Mt_sim_neo[0, 1:] + Qi_tr_neo = tgyro.QiIons_sim_neo_thr[0, 1:] + Ge_tr_neo = tgyro.Ge_sim_neo[0, 1:] + GZ_tr_neo = tgyro.Gi_sim_neo[impurityPosition, 0, 1:] + Mt_tr_neo = tgyro.Mt_sim_neo[0, 1:] - QeNeoE = abs(tgyro.Qe_sim_neo[0, 1:]) * relativeErrorNEO + Qe_tr_neoE = abs(tgyro.Qe_sim_neo[0, 1:]) * relativeErrorNEO if includeFast: - QiNeoE = abs(tgyro.QiIons_sim_neo[0, 1:]) * relativeErrorNEO + Qi_tr_neoE = abs(tgyro.QiIons_sim_neo[0, 1:]) * relativeErrorNEO else: - QiNeoE = abs(tgyro.QiIons_sim_neo_thr[0, 1:]) * relativeErrorNEO - GeNeoE = abs(tgyro.Ge_sim_neo[0, 1:]) * relativeErrorNEO - GZNeoE = abs(tgyro.Gi_sim_neo[impurityPosition, 0, 1:]) * relativeErrorNEO - MtNeoE = abs(tgyro.Mt_sim_neo[0, 1:]) * relativeErrorNEO + Qi_tr_neoE = abs(tgyro.QiIons_sim_neo_thr[0, 1:]) * relativeErrorNEO + Ge_tr_neoE = abs(tgyro.Ge_sim_neo[0, 1:]) * relativeErrorNEO + GZ_tr_neoE = abs(tgyro.Gi_sim_neo[impurityPosition, 0, 1:]) * relativeErrorNEO + Mt_tr_neoE = abs(tgyro.Mt_sim_neo[0, 1:]) * relativeErrorNEO # Merge @@ -462,11 +462,11 @@ def curateTGYROfiles( GZ, Mt, Pexch, - QeNeo=QeNeo, - QiNeo=QiNeo, - GeNeo=GeNeo, - GZNeo=GZNeo, - MtNeo=MtNeo, + Qe_tr_neo=Qe_tr_neo, + Qi_tr_neo=Qi_tr_neo, + Ge_tr_neo=Ge_tr_neo, + GZ_tr_neo=GZ_tr_neo, + Mt_tr_neo=Mt_tr_neo, impurityPosition=impurityPosition, ) @@ -479,11 +479,11 @@ def curateTGYROfiles( GZE, MtE, PexchE, - QeNeo=QeNeoE, - QiNeo=QiNeoE, - GeNeo=GeNeoE, - GZNeo=GZNeoE, - MtNeo=MtNeoE, + Qe_tr_neo=Qe_tr_neoE, + Qi_tr_neo=Qi_tr_neoE, + Ge_tr_neo=Ge_tr_neoE, + GZ_tr_neo=GZ_tr_neoE, + Mt_tr_neo=Mt_tr_neoE, impurityPosition=impurityPosition, special_label="_stds", ) @@ -564,7 +564,7 @@ def modifyResults( tgyro, folder_tgyro, minErrorPercent=5.0, - percentNeo=2.0, + percent_tr_neo=2.0, useConvectiveFluxes=False, Qi_criterion_stable=0.0025, impurityPosition=3, @@ -717,11 +717,11 @@ def modifyFLUX( GZ, Mt, S, - QeNeo=None, - QiNeo=None, - GeNeo=None, - GZNeo=None, - MtNeo=None, + Qe_tr_neo=None, + Qi_tr_neo=None, + Ge_tr_neo=None, + GZ_tr_neo=None, + Mt_tr_neo=None, impurityPosition=3, special_label=None, ): @@ -741,15 +741,15 @@ def modifyFLUX( # Particle flux: Update modTGYROfile(folder / "out.tgyro.flux_e", GeGB, pos=2, fileN_suffix=special_label) - if GeNeo is not None: - GeGB_neo = GeNeo / tgyro.Gamma_GB[-1, 1:] + if Ge_tr_neo is not None: + GeGB_neo = Ge_tr_neo / tgyro.Gamma_GB[-1, 1:] modTGYROfile(folder / "out.tgyro.flux_e", GeGB_neo, pos=1, fileN_suffix=special_label) # Energy flux: Update modTGYROfile(folder / "out.tgyro.flux_e", QeGB, pos=4, fileN_suffix=special_label) - if QeNeo is not None: - QeGB_neo = QeNeo / tgyro.Q_GB[-1, 1:] + if Qe_tr_neo is not None: + QeGB_neo = Qe_tr_neo / tgyro.Q_GB[-1, 1:] modTGYROfile(folder / "out.tgyro.flux_e", QeGB_neo, pos=3, fileN_suffix=special_label) # Rotation: Remove (it will be sum to the first ion) @@ -773,8 +773,8 @@ def modifyFLUX( modTGYROfile(folder / "out.tgyro.flux_i1", QiGB, pos=4, fileN_suffix=special_label) - if QiNeo is not None: - QiGB_neo = QiNeo / tgyro.Q_GB[-1, 1:] + if Qi_tr_neo is not None: + QiGB_neo = Qi_tr_neo / tgyro.Q_GB[-1, 1:] modTGYROfile(folder / "out.tgyro.flux_i1", QiGB_neo, pos=3, fileN_suffix=special_label) # Particle flux: Make ion particle fluxes zero, because I don't want to mistake TGLF with CGYRO when looking at tgyro results @@ -792,8 +792,8 @@ def modifyFLUX( modTGYROfile(folder / "out.tgyro.flux_i1", MtGB, pos=6, fileN_suffix=special_label) - if MtNeo is not None: - MtGB_neo = MtNeo / tgyro.Pi_GB[-1, 1:] + if Mt_tr_neo is not None: + MtGB_neo = Mt_tr_neo / tgyro.Pi_GB[-1, 1:] modTGYROfile(folder / "out.tgyro.flux_i1", MtGB_neo, pos=5, fileN_suffix=special_label) # Energy exchange: Remove (it will be the electrons one) @@ -816,8 +816,8 @@ def modifyFLUX( modTGYROfile(folder / f"out.tgyro.flux_i{i+2}",var,pos=pos,fileN_suffix=special_label) modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB,pos=2,fileN_suffix=special_label) - if GZNeo is not None: - GZGB_neo = GZNeo / tgyro.Gamma_GB[-1, 1:] + if GZ_tr_neo is not None: + GZGB_neo = GZ_tr_neo / tgyro.Gamma_GB[-1, 1:] modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB_neo,pos=1,fileN_suffix=special_label) diff --git a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py index f5b89a8e..80b88a91 100644 --- a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py +++ b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py @@ -327,16 +327,16 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): var_dict = {} mapper = { - "QeTurb": "QeMWm2_tr_turb", - "QiTurb": "QiMWm2_tr_turb", - "GeTurb": "Ce_tr_turb", - "GZTurb": "CZ_tr_turb", - "MtTurb": "Mt_tr_turb", - "QeNeo": "QeMWm2_tr_neo", - "QiNeo": "QiMWm2_tr_neo", - "GeNeo": "Ce_tr_neo", - "GZNeo": "CZ_tr_neo", - "MtNeo": "Mt_tr_neo", + "Qe_tr_turb": "QeMWm2_tr_turb", + "Qi_tr_turb": "QiMWm2_tr_turb", + "Ge_tr_turb": "Ce_tr_turb", + "GZ_tr_turb": "CZ_tr_turb", + "Mt_tr_turb": "Mt_tr_turb", + "Qe_tr_neo": "QeMWm2_tr_neo", + "Qi_tr_neo": "QiMWm2_tr_neo", + "Ge_tr_neo": "Ce_tr_neo", + "GZ_tr_neo": "CZ_tr_neo", + "Mt_tr_neo": "Mt_tr_neo", "QeTar": "QeMWm2", "QiTar": "QiMWm2", "GeTar": "Ce", @@ -388,10 +388,10 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): """ ----------------------------------------------------------------------------------- - Transport (Turb+Neo) + Transport (_tr_turb+_tr_neo) ----------------------------------------------------------------------------------- """ - of0 = var_dict[f"{var}Turb"] + var_dict[f"{var}Neo"] + of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neo"] """ ----------------------------------------------------------------------------------- @@ -461,16 +461,16 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): # Prepare dictionary from powerstate (for use in Analysis) mapper = { - "QeTurb": "QeMWm2_tr_turb", - "QiTurb": "QiMWm2_tr_turb", - "GeTurb": "Ce_tr_turb", - "GZTurb": "CZ_tr_turb", - "MtTurb": "Mt_tr_turb", - "QeNeo": "QeMWm2_tr_neo", - "QiNeo": "QiMWm2_tr_neo", - "GeNeo": "Ce_tr_neo", - "GZNeo": "CZ_tr_neo", - "MtNeo": "Mt_tr_neo", + "Qe_tr_turb": "QeMWm2_tr_turb", + "Qi_tr_turb": "QiMWm2_tr_turb", + "Ge_tr_turb": "Ce_tr_turb", + "GZ_tr_turb": "CZ_tr_turb", + "Mt_tr_turb": "Mt_tr_turb", + "Qe_tr_neo": "QeMWm2_tr_neo", + "Qi_tr_neo": "QiMWm2_tr_neo", + "Ge_tr_neo": "Ce_tr_neo", + "GZ_tr_neo": "CZ_tr_neo", + "Mt_tr_neo": "Mt_tr_neo", "QeTar": "QeMWm2", "QiTar": "QiMWm2", "GeTar": "Ce", @@ -487,7 +487,7 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): else: var_dict[ikey + "_stds"] = None - dfT = var_dict["QeTurb"] # as a reference for sizes + dfT = var_dict["Qe_tr_turb"] # as a reference for sizes # ------------------------------------------------------------------------- # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added @@ -524,12 +524,12 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): """ ----------------------------------------------------------------------------------- - Transport (Turb+Neo) + Transport (_tr_turb+_tr_neo) ----------------------------------------------------------------------------------- """ - of0 = var_dict[f"{var}Turb"] + var_dict[f"{var}Neo"] + of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neo"] of0E = ( - var_dict[f"{var}Turb_stds"] ** 2 + var_dict[f"{var}Neo_stds"] ** 2 + var_dict[f"{var}_tr_turb_stds"] ** 2 + var_dict[f"{var}_tr_neo_stds"] ** 2 ) ** 0.5 """ diff --git a/src/mitim_tools/opt_tools/scripts/evaluate_model.py b/src/mitim_tools/opt_tools/scripts/evaluate_model.py index 4b5ac2fa..da2b417d 100644 --- a/src/mitim_tools/opt_tools/scripts/evaluate_model.py +++ b/src/mitim_tools/opt_tools/scripts/evaluate_model.py @@ -10,8 +10,8 @@ This way, you can try plot, re-ft, find best parameters, etc. It calculates speed, and generates profile file to look at bottlenecks e.g. - evaluate_model.py --folder run1/ --output QiTurb_5 --input aLti_5 --around -3 - evaluate_model.py --folder run1/ --step -1 --output QiTurb_5 --file figure.eps + evaluate_model.py --folder run1/ --output Qi_tr_turb_5 --input aLti_5 --around -3 + evaluate_model.py --folder run1/ --step -1 --output Qi_tr_turb_5 --file figure.eps """ # ***************** Inputs @@ -19,7 +19,7 @@ parser = argparse.ArgumentParser() parser.add_argument("--folder", required=True, type=str) parser.add_argument("--step", type=int, required=False, default=-1) -parser.add_argument("--output", required=False, type=str, default="QiTurb_1") +parser.add_argument("--output", required=False, type=str, default="Qi_tr_turb_1") parser.add_argument("--input", required=False, type=str, default="aLti_1") parser.add_argument("--around", type=int, required=False, default=-1) parser.add_argument("--xrange", type=float, required=False, default=0.5) From 1090b3ee22018d933da3cc87461de0546d1eb976 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 10:26:34 -0700 Subject: [PATCH 021/385] Tar to _tar --- .../maestro/utils/PORTALSbeat.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 17 ++++----- src/mitim_modules/portals/PORTALStools.py | 12 +++--- .../portals/utils/PORTALSanalysis.py | 2 +- .../portals/utils/PORTALSinit.py | 16 ++++---- .../gacode_tools/utils/PORTALSinteraction.py | 38 +++++++++---------- 6 files changed, 42 insertions(+), 45 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 6dc4aaf2..2b2c9f5b 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -344,7 +344,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr # In the situation where the last radial location moves, I cannot reuse that surrogate data if last_radial_location_moved and reusing_surrogate_data: print('\t\t- Last radial location was moved, so surrogate data will not be reused for that specific location') - self.optimization_options['surrogate_options']["extrapointsModelsAvoidContent"] = ['Tar',f'_{len(self.MODELparameters[strKeys])}'] + self.optimization_options['surrogate_options']["extrapointsModelsAvoidContent"] = ['_tar',f'_{len(self.MODELparameters[strKeys])}'] self.try_flux_match_only_for_first_point = False def _inform_save(self): diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 8b8ae44c..2b977bbf 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -345,7 +345,7 @@ def _define_reuse_models(self): ''' The user can define a list of strings to avoid reusing surrogates. e.g. - 'Tar' to avoid reusing targets + '_tar' to avoid reusing targets '_5' to avoid reusing position 5 ''' @@ -353,7 +353,7 @@ def _define_reuse_models(self): # Define avoiders if self.optimization_options['surrogate_options']['extrapointsModelsAvoidContent'] is None: - self.optimization_options['surrogate_options']['extrapointsModelsAvoidContent'] = ['Tar'] + self.optimization_options['surrogate_options']['extrapointsModelsAvoidContent'] = ['_tar'] # Define extrapointsModels for key in self.surrogate_parameters['surrogate_transformation_variables_lasttime'].keys(): @@ -425,12 +425,11 @@ def scalarized_objective(self, Y): var_dict = {} for of in ofs_ordered_names: - var, _ = of.split("_") + + var = '_'.join(of.split("_")[:-1]) if var not in var_dict: var_dict[var] = torch.Tensor().to(Y) - var_dict[var] = torch.cat( - (var_dict[var], Y[..., ofs_ordered_names == of]), dim=-1 - ) + var_dict[var] = torch.cat((var_dict[var], Y[..., ofs_ordered_names == of]), dim=-1) """ ------------------------------------------------------------------------- @@ -575,7 +574,7 @@ def reuseTrainingTabular( # ------------------------------------------------------------------------------------ for i in dictOFs: - if "Tar" in i: + if "_tar" in i: print(f"Changing {i} in file") optimization_data.data[i].iloc[numPORTALS] = dictOFs[i]["value"].cpu().numpy().item() @@ -682,8 +681,8 @@ def map_powerstate_to_portals(powerstate, dictOFs): If that radius & profile position has target, evaluate """ - dictOFs[f"{var0}Tar_{i+1}"]["value"] = powerstate.plasma[f"{var1}"][0, i+1] - dictOFs[f"{var0}Tar_{i+1}"]["error"] = powerstate.plasma[f"{var1}_stds"][0, i+1] + dictOFs[f"{var0}_tar_{i+1}"]["value"] = powerstate.plasma[f"{var1}"][0, i+1] + dictOFs[f"{var0}_tar_{i+1}"]["error"] = powerstate.plasma[f"{var1}_stds"][0, i+1] """ Turbulent Exchange diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index f466058c..8a333847 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -2,9 +2,9 @@ import gpytorch import copy import numpy as np +from collections import OrderedDict from mitim_tools.opt_tools import STRATEGYtools from mitim_tools.misc_tools import PLASMAtools -from collections import OrderedDict from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -14,7 +14,7 @@ def surrogate_selection_portals(output, surrogate_options, CGYROrun=False): if output is not None: # If it's a target, just linear - if output[2:5] == "Tar": + if output[3:6] == "tar": surrogate_options["TypeMean"] = 1 surrogate_options["TypeKernel"] = 2 # Constant kernel # If it's not, stndard @@ -125,7 +125,7 @@ def input_transform_portals(Xorig, output, surrogate_parameters, surrogate_trans initialize it with a larger batch """ - _, num = output.split("_") + num = output.split("_")[-1] index = powerstate.indexes_simulation[int(num)] # num=1 -> pos=1, so that it takes the second value in vectors xFit = torch.Tensor().to(X) @@ -224,8 +224,8 @@ def computeTurbExchangeIndividual(PexchTurb, 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]) + varFull = '_'.join(output.split("_")[:-1]) + pos = int(output.split("_")[-1]) # Select GB unit if varFull[:2] == "Qe": @@ -251,7 +251,7 @@ def ImpurityGammaTrick(x, surrogate_parameters, output, powerstate): Trick to make GZ a function of a/Lnz only (flux as GammaZ_hat = GammaZ /nZ ) """ - pos = int(output.split("_")[1]) + pos = int(output.split("_")[-1]) if ("GZ" in output) and surrogate_parameters["applyImpurityGammaTrick"]: factor = powerstate.plasma["ni"][: x.shape[0],powerstate.indexes_simulation[pos],powerstate.impurityPosition].unsqueeze(-1) diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 202641ba..5078e5ae 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -723,7 +723,7 @@ def __init__(self, gpdict): self._training_outputs = {} if isinstance(gpdict, dict): for key in gpdict: - if 'Tar' in key: + if '_tar' in key: self._targets[key] = gpdict[key] else: self._models[key] = gpdict[key] diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index d140d773..9d8ea742 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -296,7 +296,7 @@ def initializeProblem( ofs.append(f"{var}_tr_turb_{i+1}") ofs.append(f"{var}_tr_neo_{i+1}") - ofs.append(f"{var}Tar_{i+1}") + ofs.append(f"{var}_tar_{i+1}") name_objectives.append(f"{var}Res_{i+1}") @@ -417,9 +417,7 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue elif typ in ["GZ", "GZ_tr_turb", "GZ_tr_neo"]: if doNotFitOnFixedValues: - isAbsValFixed = pos == ( - portals_fun.powerstate.plasma["rho"].shape[-1] - 1 - ) + isAbsValFixed = pos == portals_fun.powerstate.plasma["rho"].shape[-1] - 1 else: isAbsValFixed = False @@ -455,15 +453,15 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue if useThisOne: Variables[output].append(ikey) - elif typ in ["QeTar"]: + elif typ in ["Qe_tar"]: Variables[output] = ["QeGB"] - elif typ in ["QiTar"]: + elif typ in ["Qi_tar"]: Variables[output] = ["QiGB"] - elif typ in ["GeTar"]: + elif typ in ["Ge_tar"]: Variables[output] = ["CeGB"] - elif typ in ["GZTar"]: + elif typ in ["GZ_tar"]: Variables[output] = ["CZGB"] - elif typ in ["MtTar"]: + elif typ in ["Mt_tar"]: Variables[output] = ["MtGB"] return Variables diff --git a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py index 80b88a91..707d22da 100644 --- a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py +++ b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py @@ -337,11 +337,11 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): "Ge_tr_neo": "Ce_tr_neo", "GZ_tr_neo": "CZ_tr_neo", "Mt_tr_neo": "Mt_tr_neo", - "QeTar": "QeMWm2", - "QiTar": "QiMWm2", - "GeTar": "Ce", - "GZTar": "CZ", - "MtTar": "Mt", + "Qe_tar": "QeMWm2", + "Qi_tar": "QiMWm2", + "Ge_tar": "Ce", + "GZ_tar": "CZ", + "Mt_tar": "Mt", "PexchTurb": "PexchTurb" } @@ -399,11 +399,11 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): ----------------------------------------------------------------------------------- """ if var == "Qe": - cal0 = var_dict[f"{var}Tar"] + PexchTurb_integrated + cal0 = var_dict[f"{var}_tar"] + PexchTurb_integrated elif var == "Qi": - cal0 = var_dict[f"{var}Tar"] - PexchTurb_integrated + cal0 = var_dict[f"{var}_tar"] - PexchTurb_integrated else: - cal0 = var_dict[f"{var}Tar"] + cal0 = var_dict[f"{var}_tar"] """ ----------------------------------------------------------------------------------- @@ -471,11 +471,11 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): "Ge_tr_neo": "Ce_tr_neo", "GZ_tr_neo": "CZ_tr_neo", "Mt_tr_neo": "Mt_tr_neo", - "QeTar": "QeMWm2", - "QiTar": "QiMWm2", - "GeTar": "Ce", - "GZTar": "CZ", - "MtTar": "Mt", + "Qe_tar": "QeMWm2", + "Qi_tar": "QiMWm2", + "Ge_tar": "Ce", + "GZ_tar": "CZ", + "Mt_tar": "Mt", "PexchTurb": "PexchTurb" } @@ -538,18 +538,18 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): ----------------------------------------------------------------------------------- """ if var == "Qe": - cal0 = var_dict[f"{var}Tar"] + PexchTurb_integrated + cal0 = var_dict[f"{var}_tar"] + PexchTurb_integrated cal0E = ( - var_dict[f"{var}Tar_stds"] ** 2 + PexchTurb_integrated_stds**2 + var_dict[f"{var}_tar_stds"] ** 2 + PexchTurb_integrated_stds**2 ) ** 0.5 elif var == "Qi": - cal0 = var_dict[f"{var}Tar"] - PexchTurb_integrated + cal0 = var_dict[f"{var}_tar"] - PexchTurb_integrated cal0E = ( - var_dict[f"{var}Tar_stds"] ** 2 + PexchTurb_integrated_stds**2 + var_dict[f"{var}_tar_stds"] ** 2 + PexchTurb_integrated_stds**2 ) ** 0.5 else: - cal0 = var_dict[f"{var}Tar"] - cal0E = var_dict[f"{var}Tar_stds"] + cal0 = var_dict[f"{var}_tar"] + cal0E = var_dict[f"{var}_tar_stds"] of, cal = torch.cat((of, of0), dim=-1), torch.cat((cal, cal0), dim=-1) ofE, calE = torch.cat((ofE, of0E), dim=-1), torch.cat((calE, cal0E), dim=-1) From 407c960699a35244478674914fffcb55055c7f0a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 10:27:09 -0700 Subject: [PATCH 022/385] removed deprecated TargetType warning --- src/mitim_modules/portals/PORTALSmain.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 2b977bbf..b3c02f17 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -483,9 +483,6 @@ def check_flags(self): print("\t- In PORTALS TGYRO evaluations, we need to not recompute gradients (GradientsType=0)",typeMsg="i") self.MODELparameters["Physics_options"]["GradientsType"] = 0 - if 'TargetType' in self.MODELparameters["Physics_options"]: - raise Exception("\t- TargetType is not used in PORTALS anymore") - if self.PORTALSparameters["TargetCalc"] == "tgyro" and self.PORTALSparameters['profiles_postprocessing_fun'] is not None: print("\t- Requested custom modification of postprocessing function but targets from tgyro... are you sure?",typeMsg="q") From 73c8a7bb0dd468db18e8a9be8da8a46e44e5baa2 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 11:56:37 -0700 Subject: [PATCH 023/385] Cleaned up way to define fixed targets --- .../powertorch/utils/TARGETStools.py | 45 ++++++------------- .../powertorch/utils/TRANSFORMtools.py | 44 ++++++++---------- 2 files changed, 32 insertions(+), 57 deletions(-) diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 910d0d50..75b32bce 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -23,23 +23,12 @@ def __init__(self,powerstate): # Fixed Targets (targets without a model) # ---------------------------------------------------- - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 1: - self.Pe_orig, self.Pi_orig = ( - self.powerstate.plasma["QeMWm2_orig_fusradexch"], - self.powerstate.plasma["QiMWm2_orig_fusradexch"], - ) # Original integrated from input.gacode - elif self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 2: - self.Pe_orig, self.Pi_orig = ( - self.powerstate.plasma["QeMWm2_orig_fusrad"], - self.powerstate.plasma["QiMWm2_orig_fusrad"], - ) - elif self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: - self.Pe_orig, self.Pi_orig = self.powerstate.plasma["te"] * 0.0, self.powerstate.plasma["te"] * 0.0 - - # For the moment, I don't have a model for these, so I just grab the original from input.gacode - self.CextraE = self.powerstate.plasma["Gaux_e"] # 1E20/s/m^2 - self.CextraZ = self.powerstate.plasma["Gaux_Z"] # 1E20/s/m^2 - self.Mextra = self.powerstate.plasma["Maux"] # J/m^2 + self.Qe_fixedtargets = self.powerstate.plasma["QeMWm2_fixedtargets"] + self.Qi_fixedtargets = self.powerstate.plasma["QiMWm2_fixedtargets"] + + self.Ce_fixedtargets = self.powerstate.plasma["Ge_fixedtargets"] # 1E20/s/m^2 + self.CZ_fixedtargets = self.powerstate.plasma["GZ_fixedtargets"] # 1E20/s/m^2 + self.Mt_fixedtargets = self.powerstate.plasma["Mt_fixedtargets"] # J/m^2 def fine_grid(self): @@ -138,25 +127,17 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, # Plug-in Targets # ************************************************************************************************** - self.powerstate.plasma["QeMWm2"] = ( - self.powerstate.plasma["Paux_e"] + self.P[: self.P.shape[0]//2, :] + self.Pe_orig - ) # MW/m^2 - self.powerstate.plasma["QiMWm2"] = ( - self.powerstate.plasma["Paux_i"] + self.P[self.P.shape[0]//2 :, :] + self.Pi_orig - ) # MW/m^2 - self.powerstate.plasma["Ce_raw"] = self.CextraE - self.powerstate.plasma["CZ_raw"] = self.CextraZ - self.powerstate.plasma["Mt"] = self.Mextra + self.powerstate.plasma["QeMWm2"] = self.Qe_fixedtargets + self.P[: self.P.shape[0]//2, :] # MW/m^2 + self.powerstate.plasma["QiMWm2"] = self.Qi_fixedtargets + self.P[self.P.shape[0]//2 :, :] # MW/m^2 + self.powerstate.plasma["Ce_raw"] = self.Ce_fixedtargets # 1E20/s/m^2 + self.powerstate.plasma["CZ_raw"] = self.CZ_fixedtargets # 1E20/s/m^2 + self.powerstate.plasma["Mt"] = self.Mt_fixedtargets # J/m^2 # Merge convective fluxes if useConvectiveFluxes: - self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux( - self.powerstate.plasma["te"], self.powerstate.plasma["Ce_raw"] - ) # MW/m^2 - self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux( - self.powerstate.plasma["te"], self.powerstate.plasma["CZ_raw"] - ) # MW/m^2 + self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["Ce_raw"]) # MW/m^2 + self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["CZ_raw"]) # MW/m^2 else: self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce_raw"] self.powerstate.plasma["CZ"] = self.powerstate.plasma["CZ_raw"] diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 89ea572f..ee7024ab 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -109,39 +109,33 @@ def gacode_to_powerstate(self, rho_vec=None): ).to(rho_vec) # ********************************************************************************************* - quantities_to_interpolate_and_volp = [ - ["Paux_e", "qe_aux_MWmiller"], - ["Paux_i", "qi_aux_MWmiller"], - ["Gaux_e", "ge_10E20miller"], - ["Maux", "mt_Jmiller"], - ] - - for key in quantities_to_interpolate_and_volp: - - # ********************************************************************************************* - # Extract the quantity via interpolation and tensorization - # ********************************************************************************************* - self.plasma[key[0]] = torch.from_numpy( - interpolation_function(rho_vec.cpu(), rho_use, input_gacode.derived[key[1]]) - ).to(rho_vec) / self.plasma["volp"] - # ********************************************************************************************* - - self.plasma["Gaux_Z"] = self.plasma["Gaux_e"] * 0.0 + # ********************************************************************************************* + # Fixed targets + # ********************************************************************************************* quantitites = {} - quantitites["QeMWm2_orig_fusrad"] = input_gacode.derived["qe_fus_MWmiller"] - input_gacode.derived["qrad_MWmiller"] - quantitites["QiMWm2_orig_fusrad"] = input_gacode.derived["qi_fus_MWmiller"] - quantitites["QeMWm2_orig_fusradexch"] = quantitites["QeMWm2_orig_fusrad"] - input_gacode.derived["qe_exc_MWmiller"] - quantitites["QiMWm2_orig_fusradexch"] = quantitites["QiMWm2_orig_fusrad"] + input_gacode.derived["qe_exc_MWmiller"] + quantitites["QeMWm2_fixedtargets"] = input_gacode.derived["qe_aux_MWmiller"] + quantitites["QiMWm2_fixedtargets"] = input_gacode.derived["qi_aux_MWmiller"] + quantitites["Ge_fixedtargets"] = input_gacode.derived["ge_10E20miller"] + quantitites["GZ_fixedtargets"] = input_gacode.derived["ge_10E20miller"] * 0.0 + quantitites["Mt_fixedtargets"] = input_gacode.derived["mt_Jmiller"] + + if self.TargetOptions["ModelOptions"]["TypeTarget"] < 3: + # Fusion and radiation fixed if 1,2 + quantitites["QeMWm2_fixedtargets"] += input_gacode.derived["qe_fus_MWmiller"] - input_gacode.derived["qrad_MWmiller"] + quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qi_fus_MWmiller"] + + if self.TargetOptions["ModelOptions"]["TypeTarget"] < 2: + # Exchange fixed if 1 + quantitites["QeMWm2_fixedtargets"] -= input_gacode.derived["qe_exc_MWmiller"] + quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qe_exc_MWmiller"] for key in quantitites: # ********************************************************************************************* # Extract the quantity via interpolation and tensorization # ********************************************************************************************* - self.plasma[key] = torch.from_numpy( - interpolation_function(rho_vec.cpu(), rho_use, quantitites[key]) - ).to(rho_vec) / self.plasma["volp"] + self.plasma[key] = torch.from_numpy(interpolation_function(rho_vec.cpu(), rho_use, quantitites[key])).to(rho_vec) / self.plasma["volp"] # ********************************************************************************************* # ********************************************************************************************* From c94dec02ebbcdf5b0733d1f99f4427bac02524b9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 11:59:12 -0700 Subject: [PATCH 024/385] misc --- .../powertorch/utils/TARGETStools.py | 34 ++++++------------- 1 file changed, 10 insertions(+), 24 deletions(-) diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 75b32bce..8ed71285 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -19,17 +19,6 @@ def __init__(self,powerstate): for i in variables_to_zero: self.powerstate.plasma[i] = self.powerstate.plasma["te"] * 0.0 - # ---------------------------------------------------- - # Fixed Targets (targets without a model) - # ---------------------------------------------------- - - self.Qe_fixedtargets = self.powerstate.plasma["QeMWm2_fixedtargets"] - self.Qi_fixedtargets = self.powerstate.plasma["QiMWm2_fixedtargets"] - - self.Ce_fixedtargets = self.powerstate.plasma["Ge_fixedtargets"] # 1E20/s/m^2 - self.CZ_fixedtargets = self.powerstate.plasma["GZ_fixedtargets"] # 1E20/s/m^2 - self.Mt_fixedtargets = self.powerstate.plasma["Mt_fixedtargets"] # J/m^2 - def fine_grid(self): """ @@ -124,14 +113,14 @@ def coarse_grid(self): def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, assumedPercentError=1.0): # ************************************************************************************************** - # Plug-in Targets + # Plug-in targets that were fixed # ************************************************************************************************** - self.powerstate.plasma["QeMWm2"] = self.Qe_fixedtargets + self.P[: self.P.shape[0]//2, :] # MW/m^2 - self.powerstate.plasma["QiMWm2"] = self.Qi_fixedtargets + self.P[self.P.shape[0]//2 :, :] # MW/m^2 - self.powerstate.plasma["Ce_raw"] = self.Ce_fixedtargets # 1E20/s/m^2 - self.powerstate.plasma["CZ_raw"] = self.CZ_fixedtargets # 1E20/s/m^2 - self.powerstate.plasma["Mt"] = self.Mt_fixedtargets # J/m^2 + self.powerstate.plasma["QeMWm2"] = self.powerstate.plasma["QeMWm2_fixedtargets"] + self.P[: self.P.shape[0]//2, :] # MW/m^2 + self.powerstate.plasma["QiMWm2"] = self.powerstate.plasma["QiMWm2_fixedtargets"] + self.P[self.P.shape[0]//2 :, :] # MW/m^2 + self.powerstate.plasma["Ce_raw"] = self.powerstate.plasma["Ge_fixedtargets"] # 1E20/s/m^2 + self.powerstate.plasma["CZ_raw"] = self.powerstate.plasma["GZ_fixedtargets"] # 1E20/s/m^2 + self.powerstate.plasma["Mt"] = self.powerstate.plasma["Mt_fixedtargets"] # J/m^2 # Merge convective fluxes @@ -143,7 +132,7 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, self.powerstate.plasma["CZ"] = self.powerstate.plasma["CZ_raw"] if forceZeroParticleFlux: - self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce"] * 0 + self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce"] * 0 self.powerstate.plasma["Ce_raw"] = self.powerstate.plasma["Ce_raw"] * 0 # ************************************************************************************************** @@ -155,12 +144,9 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, for i in variables_to_error: self.powerstate.plasma[i + "_stds"] = abs(self.powerstate.plasma[i]) * assumedPercentError / 100 - """ - ************************************************************************************************** - GB Normalized - ************************************************************************************************** - Note: This is useful for mitim surrogate variables of targets - """ + # ************************************************************************************************** + # GB Normalized (Note: This is useful for mitim surrogate variables of targets) + # ************************************************************************************************** self.powerstate.plasma["QeGB"] = self.powerstate.plasma["QeMWm2"] / self.powerstate.plasma["Qgb"] self.powerstate.plasma["QiGB"] = self.powerstate.plasma["QiMWm2"] / self.powerstate.plasma["Qgb"] From 8c1054a202e99421d21ce584d57854d09dc6fef9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 16:24:59 -0700 Subject: [PATCH 025/385] Transitioned momentum to new naming convention --- src/mitim_modules/portals/PORTALSmain.py | 2 +- .../portals/utils/PORTALSplot.py | 26 ++++----- src/mitim_modules/powertorch/STATEtools.py | 10 ++-- .../physics_models/transport_analytic.py | 2 +- .../physics_models/transport_cgyro.py | 8 +-- .../powertorch/utils/POWERplot.py | 6 +-- .../powertorch/utils/TARGETStools.py | 26 +++++---- .../powertorch/utils/TRANSFORMtools.py | 2 +- .../powertorch/utils/TRANSPORTtools.py | 4 +- .../gacode_tools/utils/PORTALSinteraction.py | 54 +++++++++---------- src/mitim_tools/opt_tools/BOTORCHtools.py | 11 +--- tests/PORTALS_workflow.py | 16 +++--- 12 files changed, 79 insertions(+), 88 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index b3c02f17..0ecbca15 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -659,7 +659,7 @@ def map_powerstate_to_portals(powerstate, dictOFs): elif var == "nZ": var0, var1 = "GZ", "CZ" elif var == "w0": - var0, var1 = "Mt", "Mt" + var0, var1 = "Mt", "MtJm2" """ TRANSPORT calculation diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index e0541935..fcae089e 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -241,11 +241,11 @@ def PORTALSanalyzer_plotMetrics( axne_f.plot( rho, - power.plasma['Ce_raw_tr_turb'].cpu().numpy()+power.plasma['Ce_raw_tr_neo'].cpu().numpy(), + power.plasma['Ge1E20sm2_tr_turb'].cpu().numpy()+power.plasma['Ge1E20sm2_tr_neo'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph) axne_f.plot( rho, - power.plasma['Ce_raw'].cpu().numpy() * (1 - int(self.forceZeroParticleFlux)), + power.plasma['Ge1E20sm2'].cpu().numpy() * (1 - int(self.forceZeroParticleFlux)), "--", c=col, lw=lwt, @@ -260,13 +260,13 @@ def PORTALSanalyzer_plotMetrics( if axw0_f is not None: axw0_f.plot( rho, - power.plasma['Mt_tr_turb'].cpu().numpy() + power.plasma['Mt_tr_neo'].cpu().numpy(), + power.plasma['MtJm2_tr_turb'].cpu().numpy() + power.plasma['MtJm2_tr_neo'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph, ) - axw0_f.plot(rho, power.plasma['Mt'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) + axw0_f.plot(rho, power.plasma['MtJm2'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) # --------------------------------------------------------------------------------------------------------- @@ -2168,7 +2168,7 @@ def PORTALSanalyzer_plotModelComparison( quantityX_stds = "MtGB_sim_turb_stds" quantityY = "MtGB_sim_turb" quantityY_stds = "MtGB_sim_turb_stds" - metrics["Mt"] = plotModelComparison_quantity( + metrics["MtJm2"] = plotModelComparison_quantity( self, axs[cont], quantityX=quantityX, @@ -2931,7 +2931,7 @@ def plotFluxComparison( if axne_f is not None: - Ge = power.plasma['Ce_raw_tr_turb'].cpu().numpy() + power.plasma['Ce_raw_tr_neo'].cpu().numpy() + Ge = power.plasma['Ge1E20sm2_tr_turb'].cpu().numpy() + power.plasma['Ge1E20sm2_tr_neo'].cpu().numpy() axne_f.plot( r[0][ixF:], @@ -2944,7 +2944,7 @@ def plotFluxComparison( alpha=alpha, ) - sigma = power.plasma['Ce_raw_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['Ce_raw_tr_neo_stds'].cpu().numpy()[0][ixF:] + sigma = power.plasma['Ge1E20sm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['Ge1E20sm2_tr_neo_stds'].cpu().numpy()[0][ixF:] m_Ge, M_Ge = Ge[0][ixF:] - stds * sigma, Ge[0][ixF:] + stds * sigma @@ -2983,7 +2983,7 @@ def plotFluxComparison( if axw0_f is not None: axw0_f.plot( r[0][ixF:], - power.plasma['Mt_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Mt_tr_neo'].cpu().numpy()[0][ixF:], + power.plasma['MtJm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neo'].cpu().numpy()[0][ixF:], "-s", markersize=msFlux, c=col, @@ -2992,10 +2992,10 @@ def plotFluxComparison( alpha=alpha, ) - sigma = power.plasma['Mt_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['Mt_tr_neo_stds'].cpu().numpy()[0][ixF:] + sigma = power.plasma['MtJm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neo_stds'].cpu().numpy()[0][ixF:] - m_Mt, M_Mt = (power.plasma['Mt_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Mt_tr_neo'].cpu().numpy()[0][ixF:]) - stds * sigma, ( - power.plasma['Mt_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['Mt_tr_neo'].cpu().numpy()[0][ixF:] + m_Mt, M_Mt = (power.plasma['MtJm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neo'].cpu().numpy()[0][ixF:]) - stds * sigma, ( + power.plasma['MtJm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neo'].cpu().numpy()[0][ixF:] ) + stds * sigma axw0_f.fill_between(r[0][ixF:], m_Mt, M_Mt, facecolor=col, alpha=alpha / 3) @@ -3007,9 +3007,9 @@ def plotFluxComparison( Qe_tar = power.plasma['QeMWm2'].cpu().numpy()[0][ixF:] Qi_tar = power.plasma['QiMWm2'].cpu().numpy()[0][ixF:] - Ge_tar = power.plasma['Ce_raw'].cpu().numpy()[0][ixF:] * (1-int(forceZeroParticleFlux)) + Ge_tar = power.plasma['Ge1E20sm2'].cpu().numpy()[0][ixF:] * (1-int(forceZeroParticleFlux)) GZ_tar = power.plasma['CZ_raw'].cpu().numpy()[0][ixF:] - Mt_tar = power.plasma['Mt'].cpu().numpy()[0][ixF:] + Mt_tar = power.plasma['MtJm2'].cpu().numpy()[0][ixF:] # Plot ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 023ab17b..703d0500 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -108,7 +108,7 @@ def _ensure_ne_before_nz(lst): "ti": ("QiMWm2", "QiMWm2_tr"), "ne": ("Ce", "Ce_tr"), "nZ": ("CZ", "CZ_tr"), - "w0": ("Mt", "Mt_tr") + "w0": ("MtJm2", "MtJm2_tr") } # ------------------------------------------------------------------------------------- @@ -293,8 +293,8 @@ def calculate( self.calculateProfileFunctions() # 3. Sources and sinks (populates components and Pe,Pi,...) - assumedPercentError = self.TransportOptions["ModelOptions"].get("percentError", [5, 1, 0.5])[-1] - self.calculateTargets(assumedPercentError=assumedPercentError) # Calculate targets based on powerstate functions (it may be overwritten in next step, if chosen) + relative_error_assumed = self.TransportOptions["ModelOptions"].get("percentError", [5, 1, 0.5])[-1] + self.calculateTargets(relative_error_assumed=relative_error_assumed) # Calculate targets based on powerstate functions (it may be overwritten in next step, if chosen) # 4. Turbulent and neoclassical transport (populates components and Pe_tr,Pi_tr,...) self.calculateTransport( @@ -674,7 +674,7 @@ def calculateProfileFunctions(self, calculateRotationQuantities=True, mref=2.013 self.plasma["w0_n"] = self.plasma["w0"] / self.plasma["c_s"] self.plasma["aLw0_n"] = (self.plasma["aLw0"] * self.plasma["w0"] / self.plasma["c_s"]) # aLw0 * w0 = -a*dw0/dr; then aLw0_n = -dw0/dr * a/c_s - def calculateTargets(self, assumedPercentError=1.0): + def calculateTargets(self, relative_error_assumed=1.0): """ Update the targets of the current state """ @@ -701,7 +701,7 @@ def calculateTargets(self, assumedPercentError=1.0): # Merge targets, calculate errors and normalize targets.postprocessing( - assumedPercentError=assumedPercentError, + relative_error_assumed=relative_error_assumed, useConvectiveFluxes=self.useConvectiveFluxes, forceZeroParticleFlux=self.TransportOptions["ModelOptions"].get("forceZeroParticleFlux", False)) diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py index 07622713..d9002af3 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -82,7 +82,7 @@ def evaluate(self): "ti": "QiMWm2", "ne": "Ce", "nZ": "CZ", - "w0": "Mt", + "w0": "MtJm2", } for c, i in enumerate(self.powerstate.ProfilesPredicted): diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 5ae7f214..031ed0a0 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -25,7 +25,7 @@ def evaluate(self): # Some checks print("\t- Checking model modifications:") - for r in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: #, "PexchTurb"]: #TODO: FIX + for r in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "MtJm2_tr_turb"]: #, "PexchTurb"]: #TODO: FIX print(f"\t\t{r}(tglf) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(powerstate_orig.plasma[r][0][1:],powerstate_orig.plasma[r+'_stds'][0][1:]) ])}") print(f"\t\t{r}(cgyro) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(self.powerstate.plasma[r][0][1:],self.powerstate.plasma[r+'_stds'][0][1:]) ])}") @@ -102,7 +102,7 @@ def _cgyro_trick(self,FolderEvaluation_TGYRO): ) # Make tensors - for i in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "Mt_tr_turb"]: + for i in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "MtJm2_tr_turb"]: try: self.powerstate.plasma[i] = torch.from_numpy(self.powerstate.plasma[i]).to(self.powerstate.dfT).unsqueeze(0) except: @@ -121,8 +121,8 @@ def _print_info(self): txt += f"{self.powerstate.plasma[varn][0,j+1]:.6f} " for var, varn in zip( - ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ce (MW/m^2)", "CZ (MW/m^2)", "Mt (J/m^2) "], - ["QeMWm2", "Pi", "Ce", "CZ", "Mt"], + ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ce (MW/m^2)", "CZ (MW/m^2)", "MtJm2 (J/m^2) "], + ["QeMWm2", "Pi", "Ce", "CZ", "MtJm2"], ): txt += f"\n{var} = " for j in range(self.powerstate.plasma["rho"].shape[1] - 1): diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index f51a355d..9d17a272 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -36,9 +36,9 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co if "ne" in self.ProfilesPredicted: # If this model provides the raw particle flux, go for it - if 'Ce_raw_tr' in self.plasma: + if 'Ge1E20sm2_tr' in self.plasma: set_plots.append( - [ 'ne', 'aLne', 'Ce_raw_tr', 'Ce_raw', + [ 'ne', 'aLne', 'Ge1E20sm2_tr', 'Ge1E20sm2', 'Electron Density','$n_e$ ($10^{20}m^{-3}$)','$a/Ln_e$','$\\Gamma_e$ (GB)','$\\Gamma_e$ ($10^{20}m^{-3}/s$)', 1E-1,"Ggb"]) else: @@ -75,7 +75,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co if "w0" in self.ProfilesPredicted: set_plots.append( - [ 'w0', 'aLw0', 'Mt_tr', 'Mt', + [ 'w0', 'aLw0', 'MtJm2_tr', 'MtJm2', 'Rotation','$\\omega_0$ ($krad/s$)','$-d\\omega_0/dr$ ($krad/s/cm$)','$\\Pi$ (GB)','$\\Pi$ ($J/m^2$)', 1E-3,"Pgb"]) diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 8ed71285..1ae63c80 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -53,7 +53,7 @@ def fine_grid(self): self.plasma_original[f'aL{variable}'] = self.powerstate.plasma[f'aL{variable}'].clone() # ---------------------------------------------------- - # Integrate through fine de-parameterization + # Integrate through fine profile constructors # ---------------------------------------------------- for i in self.powerstate.ProfilesPredicted: _ = self.powerstate.update_var(i,specific_profile_constructor=self.powerstate.profile_constructors_coarse_middle) @@ -110,7 +110,7 @@ def coarse_grid(self): for i in self.plasma_original: self.powerstate.plasma[i] = self.plasma_original[i] - def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, assumedPercentError=1.0): + def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, relative_error_assumed=1.0): # ************************************************************************************************** # Plug-in targets that were fixed @@ -118,31 +118,29 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, self.powerstate.plasma["QeMWm2"] = self.powerstate.plasma["QeMWm2_fixedtargets"] + self.P[: self.P.shape[0]//2, :] # MW/m^2 self.powerstate.plasma["QiMWm2"] = self.powerstate.plasma["QiMWm2_fixedtargets"] + self.P[self.P.shape[0]//2 :, :] # MW/m^2 - self.powerstate.plasma["Ce_raw"] = self.powerstate.plasma["Ge_fixedtargets"] # 1E20/s/m^2 + self.powerstate.plasma["Ge1E20sm2"] = self.powerstate.plasma["Ge_fixedtargets"] # 1E20/s/m^2 self.powerstate.plasma["CZ_raw"] = self.powerstate.plasma["GZ_fixedtargets"] # 1E20/s/m^2 - self.powerstate.plasma["Mt"] = self.powerstate.plasma["Mt_fixedtargets"] # J/m^2 + self.powerstate.plasma["MtJm2"] = self.powerstate.plasma["MtJm2_fixedtargets"] # J/m^2 - # Merge convective fluxes + if forceZeroParticleFlux: + self.powerstate.plasma["Ge1E20sm2"] = self.powerstate.plasma["Ge1E20sm2"] * 0 + # Convective fluxes? if useConvectiveFluxes: - self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["Ce_raw"]) # MW/m^2 + self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["Ge1E20sm2"]) # MW/m^2 self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["CZ_raw"]) # MW/m^2 else: - self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce_raw"] + self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ge1E20sm2"] self.powerstate.plasma["CZ"] = self.powerstate.plasma["CZ_raw"] - if forceZeroParticleFlux: - self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce"] * 0 - self.powerstate.plasma["Ce_raw"] = self.powerstate.plasma["Ce_raw"] * 0 - # ************************************************************************************************** # Error # ************************************************************************************************** - variables_to_error = ["QeMWm2", "QiMWm2", "Ce", "CZ", "Mt", "Ce_raw", "CZ_raw"] + variables_to_error = ["QeMWm2", "QiMWm2", "Ce", "CZ", "MtJm2", "Ge1E20sm2", "CZ_raw"] for i in variables_to_error: - self.powerstate.plasma[i + "_stds"] = abs(self.powerstate.plasma[i]) * assumedPercentError / 100 + self.powerstate.plasma[i + "_stds"] = abs(self.powerstate.plasma[i]) * relative_error_assumed / 100 # ************************************************************************************************** # GB Normalized (Note: This is useful for mitim surrogate variables of targets) @@ -152,4 +150,4 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, self.powerstate.plasma["QiGB"] = self.powerstate.plasma["QiMWm2"] / self.powerstate.plasma["Qgb"] self.powerstate.plasma["CeGB"] = self.powerstate.plasma["Ce"] / self.powerstate.plasma["Qgb" if useConvectiveFluxes else "Ggb"] self.powerstate.plasma["CZGB"] = self.powerstate.plasma["CZ"] / self.powerstate.plasma["Qgb" if useConvectiveFluxes else "Ggb"] - self.powerstate.plasma["MtGB"] = self.powerstate.plasma["Mt"] / self.powerstate.plasma["Pgb"] + self.powerstate.plasma["MtGB"] = self.powerstate.plasma["MtJm2"] / self.powerstate.plasma["Pgb"] diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index ee7024ab..68c283a8 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -118,7 +118,7 @@ def gacode_to_powerstate(self, rho_vec=None): quantitites["QiMWm2_fixedtargets"] = input_gacode.derived["qi_aux_MWmiller"] quantitites["Ge_fixedtargets"] = input_gacode.derived["ge_10E20miller"] quantitites["GZ_fixedtargets"] = input_gacode.derived["ge_10E20miller"] * 0.0 - quantitites["Mt_fixedtargets"] = input_gacode.derived["mt_Jmiller"] + quantitites["MtJm2_fixedtargets"] = input_gacode.derived["mt_Jmiller"] if self.TargetOptions["ModelOptions"]["TypeTarget"] < 3: # Fusion and radiation fixed if 1,2 diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 42d75109..b1c8fc6c 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -22,7 +22,7 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ self.powerstate = powerstate # Allowed fluxes in powerstate so far - self.quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'Mt'] + self.quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2'] # Each flux has a turbulent and neoclassical component self.variables = [f'{i}_tr_turb' for i in self.quantities] + [f'{i}_tr_neo' for i in self.quantities] @@ -131,7 +131,7 @@ def evaluate(self): - Pe, Pe_tr, Pe_tr_turb, Pe_tr_neo -> MW/m^2 - Pi, Pi_tr, Pi_tr_turb, Pi_tr_neo -> MW/m^2 - Ce, Ce_tr, Ce_tr_turb, Ce_tr_neo -> MW/m^2 - * Ce_raw, Ce_raw_tr, Ce_raw_tr_turb, Ce_raw_tr_neo -> 10^20/s/m^2 + * Ge1E20sm2, Ge1E20sm2_tr, Ge1E20sm2_tr_turb, Ge1E20sm2_tr_neo -> 10^20/s/m^2 - CZ, CZ_tr, CZ_tr_turb, CZ_tr_neo -> MW/m^2 (but modified as needed, for example dividing by fZ0) * CZ_raw, CZ_raw_tr, CZ_raw_tr_turb, CZ_raw_tr_neo -> 10^20/s/m^2 (NOT modified) - Mt, Mt_tr, Mt_tr_turb, Mt_tr_neo -> J/m^2 diff --git a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py index 707d22da..29fdfd33 100644 --- a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py +++ b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py @@ -193,42 +193,42 @@ def TGYROmodeledVariables(TGYROresults, # *********** Momentum Fluxes # ********************************** - powerstate.plasma["Mt_tr_turb"] = torch.Tensor(TGYROresults.Mt_sim_turb[:, :nr]).to(powerstate.dfT) # So far, let's include fast in momentum - powerstate.plasma["Mt_tr_neo"] = torch.Tensor(TGYROresults.Mt_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["MtJm2_tr_turb"] = torch.Tensor(TGYROresults.Mt_sim_turb[:, :nr]).to(powerstate.dfT) # So far, let's include fast in momentum + powerstate.plasma["MtJm2_tr_neo"] = torch.Tensor(TGYROresults.Mt_sim_neo[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Mt_tr_turb_stds"] = torch.Tensor(TGYROresults.Mt_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Mt_tr_neo_stds"] = torch.Tensor(TGYROresults.Mt_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["MtJm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Mt_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["MtJm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Mt_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: - powerstate.plasma["Mt"] = torch.Tensor(TGYROresults.Mt_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Mt_stds"] = torch.Tensor(TGYROresults.Mt_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["MtJm2"] = torch.Tensor(TGYROresults.Mt_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["MtJm2_stds"] = torch.Tensor(TGYROresults.Mt_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None # ********************************** # *********** Particle Fluxes # ********************************** # Store raw fluxes for better plotting later - powerstate.plasma["Ce_raw_tr_turb"] = torch.Tensor(TGYROresults.Ge_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_raw_tr_neo"] = torch.Tensor(TGYROresults.Ge_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ge1E20sm2_tr_turb"] = torch.Tensor(TGYROresults.Ge_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ge1E20sm2_tr_neo"] = torch.Tensor(TGYROresults.Ge_sim_neo[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_raw_tr_turb_stds"] = torch.Tensor(TGYROresults.Ge_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Ce_raw_tr_neo_stds"] = torch.Tensor(TGYROresults.Ge_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ge1E20sm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Ge_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ge1E20sm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Ge_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: - powerstate.plasma["Ce_raw"] = torch.Tensor(TGYROresults.Ge_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_raw_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ge1E20sm2"] = torch.Tensor(TGYROresults.Ge_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ge1E20sm2_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if not useConvectiveFluxes: - powerstate.plasma["Ce_tr_turb"] = powerstate.plasma["Ce_raw_tr_turb"] - powerstate.plasma["Ce_tr_neo"] = powerstate.plasma["Ce_raw_tr_neo"] + powerstate.plasma["Ce_tr_turb"] = powerstate.plasma["Ge1E20sm2_tr_turb"] + powerstate.plasma["Ce_tr_neo"] = powerstate.plasma["Ge1E20sm2_tr_neo"] - powerstate.plasma["Ce_tr_turb_stds"] = powerstate.plasma["Ce_raw_tr_turb_stds"] - powerstate.plasma["Ce_tr_neo_stds"] = powerstate.plasma["Ce_raw_tr_neo_stds"] + powerstate.plasma["Ce_tr_turb_stds"] = powerstate.plasma["Ge1E20sm2_tr_turb_stds"] + powerstate.plasma["Ce_tr_neo_stds"] = powerstate.plasma["Ge1E20sm2_tr_neo_stds"] if provideTargets: - powerstate.plasma["Ce"] = powerstate.plasma["Ce_raw"] - powerstate.plasma["Ce_stds"] = powerstate.plasma["Ce_raw_stds"] + powerstate.plasma["Ce"] = powerstate.plasma["Ge1E20sm2"] + powerstate.plasma["Ce_stds"] = powerstate.plasma["Ge1E20sm2_stds"] else: @@ -304,7 +304,7 @@ def TGYROmodeledVariables(TGYROresults, # Sum here turbulence and neoclassical, after modifications # ------------------------------------------------------------------------------------------------------------------------ - quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'Mt', 'Ce_raw', 'CZ_raw'] + quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2', 'Ge1E20sm2', 'CZ_raw'] for ikey in quantities: powerstate.plasma[ikey+"_tr"] = powerstate.plasma[ikey+"_tr_turb"] + powerstate.plasma[ikey+"_tr_neo"] @@ -331,17 +331,17 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): "Qi_tr_turb": "QiMWm2_tr_turb", "Ge_tr_turb": "Ce_tr_turb", "GZ_tr_turb": "CZ_tr_turb", - "Mt_tr_turb": "Mt_tr_turb", + "Mt_tr_turb": "MtJm2_tr_turb", "Qe_tr_neo": "QeMWm2_tr_neo", "Qi_tr_neo": "QiMWm2_tr_neo", "Ge_tr_neo": "Ce_tr_neo", "GZ_tr_neo": "CZ_tr_neo", - "Mt_tr_neo": "Mt_tr_neo", + "Mt_tr_neo": "MtJm2_tr_neo", "Qe_tar": "QeMWm2", "Qi_tar": "QiMWm2", "Ge_tar": "Ce", "GZ_tar": "CZ", - "Mt_tar": "Mt", + "Mt_tar": "MtJm2", "PexchTurb": "PexchTurb" } @@ -431,7 +431,7 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): of0 * PORTALSparameters["Pseudo_multipliers"][3], cal0 * PORTALSparameters["Pseudo_multipliers"][3], ) - elif var == "Mt": + elif var == "MtJm2": of0, cal0 = ( of0 * PORTALSparameters["Pseudo_multipliers"][4], cal0 * PORTALSparameters["Pseudo_multipliers"][4], @@ -465,17 +465,17 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): "Qi_tr_turb": "QiMWm2_tr_turb", "Ge_tr_turb": "Ce_tr_turb", "GZ_tr_turb": "CZ_tr_turb", - "Mt_tr_turb": "Mt_tr_turb", + "Mt_tr_turb": "MtJm2_tr_turb", "Qe_tr_neo": "QeMWm2_tr_neo", "Qi_tr_neo": "QiMWm2_tr_neo", "Ge_tr_neo": "Ce_tr_neo", "GZ_tr_neo": "CZ_tr_neo", - "Mt_tr_neo": "Mt_tr_neo", + "Mt_tr_neo": "MtJm2_tr_neo", "Qe_tar": "QeMWm2", "Qi_tar": "QiMWm2", "Ge_tar": "Ce", "GZ_tar": "CZ", - "Mt_tar": "Mt", + "Mt_tar": "MtJm2", "PexchTurb": "PexchTurb" } @@ -520,7 +520,7 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): elif prof == "nZ": var = "GZ" elif prof == "w0": - var = "Mt" + var = "MtJm2" """ ----------------------------------------------------------------------------------- diff --git a/src/mitim_tools/opt_tools/BOTORCHtools.py b/src/mitim_tools/opt_tools/BOTORCHtools.py index d832cb95..f6b07cb6 100644 --- a/src/mitim_tools/opt_tools/BOTORCHtools.py +++ b/src/mitim_tools/opt_tools/BOTORCHtools.py @@ -704,15 +704,8 @@ def __init__( self.indeces_grad = tuple(grad_vector) # ---------------------------------------------------------------- - self.register_parameter( - name="weights_lin", - parameter=torch.nn.Parameter( - torch.randn(*batch_shape, len(self.indeces_grad), 1) - ), - ) - self.register_parameter( - name="bias", parameter=torch.nn.Parameter(torch.randn(*batch_shape, 1)) - ) + self.register_parameter(name="weights_lin",parameter=torch.nn.Parameter(torch.randn(*batch_shape, len(self.indeces_grad), 1)),) + self.register_parameter(name="bias", parameter=torch.nn.Parameter(torch.randn(*batch_shape, 1))) # set the parameter constraint to be [0,1], when nothing is specified diffusion_constraint = gpytorch.constraints.constraints.Positive() diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index 427a935d..421d3251 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -26,8 +26,8 @@ # Initialize class portals_fun = PORTALSmain.portals(folderWork) -portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 -portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 +portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 +portals_fun.optimization_options["initialization_options"]["initial_training"] = 2 portals_fun.MODELparameters["RhoLocations"] = [0.25, 0.45, 0.65, 0.85] portals_fun.MODELparameters['ProfilesPredicted'] = ["te", "ti", "ne", "nZ", 'w0'] portals_fun.PORTALSparameters['ImpurityOfInterest'] = 'N' @@ -52,12 +52,12 @@ portals_fun.plot_optimization_results(analysis_level=4) # For fun and to show capabilities, let's do a flux match of the current surrogates and plot in the same notebook -# PORTALSoptimization.flux_match_surrogate( -# mitim_bo.steps[-1],PROFILEStools.PROFILES_GACODE(inputgacode), -# fn = portals_fun.fn, -# plot_results = True, -# keep_within_bounds = False -# ) +PORTALSoptimization.flux_match_surrogate( + mitim_bo.steps[-1],PROFILEStools.PROFILES_GACODE(inputgacode), + fn = portals_fun.fn, + plot_results = True, + keep_within_bounds = False + ) # Required if running in non-interactive mode portals_fun.fn.show() From 754e30dd27c2fc8f087e40760f6c5a6eebe58b4e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 16:48:11 -0700 Subject: [PATCH 026/385] Corrected ions position --- .../powertorch/physics_models/transport_tgyro.py | 3 ++- src/mitim_tools/gacode_tools/TGLFtools.py | 5 +---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 55d0ccba..56de66b3 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -200,7 +200,8 @@ def tglf_scan_trick( label, RadiisToRun, ProfilesPredicted, - impurityPosition=1, includeFast=False, + impurityPosition=1, + includeFast=False, delta=0.02, cold_start=False, check_coincidence_thr=1E-2, diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 44f2e9c4..4e63342f 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -3126,7 +3126,6 @@ def runScanTurbulenceDrives( resolutionPoints=5, variation=0.5, add_baseline_to = 'none', # 'all' or 'first' or 'none' - add_also_baseline_to_first = True, variablesDrives=["RLTS_1", "RLTS_2", "RLNS_1", "XNUE", "TAUS_2"], positionIon=2, **kwargs_TGLFrun, @@ -3151,9 +3150,7 @@ def runScanTurbulenceDrives( tglf_executor, tglf_executor_full, folders = {}, {}, [] for cont, variable in enumerate(self.variablesDrives): # Only ask the cold_start in the first round - kwargs_TGLFrun["forceIfcold_start"] = cont > 0 or ( - "forceIfcold_start" in kwargs_TGLFrun and kwargs_TGLFrun["forceIfcold_start"] - ) + kwargs_TGLFrun["forceIfcold_start"] = cont > 0 or ("forceIfcold_start" in kwargs_TGLFrun and kwargs_TGLFrun["forceIfcold_start"]) scan_name = f"{subFolderTGLF}_{variable}" # e.g. turbDrives_RLTS_1 From eb4a2db10d20bd08d31ee6109d7597ae5aff7859 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 17:04:55 -0700 Subject: [PATCH 027/385] renaming to tgyro_to_powerstate --- .../powertorch/physics_models/transport_tgyro.py | 4 ++-- src/mitim_tools/gacode_tools/TGYROtools.py | 4 ++-- src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 56de66b3..01f7fa85 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -148,10 +148,10 @@ def _postprocess_results(self, tgyro, label): forceZeroParticleFlux = ModelOptions.get("forceZeroParticleFlux", False) # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) - impurityPosition = self.powerstate.impurityPosition_transport #ModelOptions.get("impurityPosition", 1) + impurityPosition = self.powerstate.impurityPosition_transport # Produce right quantities (TGYRO -> powerstate.plasma) - self.powerstate = tgyro.results[label].TGYROmodeledVariables( + self.powerstate = tgyro.results[label].tgyro_to_powerstate( self.powerstate, useConvectiveFluxes=useConvectiveFluxes, includeFast=includeFast, diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 139ec674..fb7cfcc8 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -2185,8 +2185,8 @@ def useFineGridTargets(self, impurityPosition=1): self.Ge_tarMW = self.Ge_tar * self.dvoldr self.Ce_tarMW = self.Ce_tar * self.dvoldr - def TGYROmodeledVariables(self, *args, **kwargs): - return PORTALSinteraction.TGYROmodeledVariables(self, *args, **kwargs) + def tgyro_to_powerstate(self, *args, **kwargs): + return PORTALSinteraction.tgyro_to_powerstate(self, *args, **kwargs) def plot(self, fn=None, label="", prelabel="", fn_color=None): if fn is None: diff --git a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py index 29fdfd33..45cd1af5 100644 --- a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py +++ b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py @@ -126,7 +126,7 @@ def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5): # This is where the definitions for the summation variables happen for mitim and PORTALSplot # ------------------------------------------------------------------------------------------------------------------------------------------------------ -def TGYROmodeledVariables(TGYROresults, +def tgyro_to_powerstate(TGYROresults, powerstate, useConvectiveFluxes=False, forceZeroParticleFlux=False, From 81b306396f0823f4c74bb0775f421dc322ccf172 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Apr 2025 17:07:14 -0700 Subject: [PATCH 028/385] misc --- .../gacode_tools/utils/PORTALSinteraction.py | 43 +++++-------------- 1 file changed, 11 insertions(+), 32 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py index 45cd1af5..3ae6c89a 100644 --- a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py +++ b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py @@ -52,47 +52,30 @@ def changeRFpower(self, PrfMW=25.0): def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0.2): ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) - self.profiles["te(keV)"] = ( - self.profiles["te(keV)"] * TkeV / self.profiles["te(keV)"][ix] - ) + self.profiles["te(keV)"] = self.profiles["te(keV)"] * TkeV / self.profiles["te(keV)"][ix] - print( - f"- Producing {typeEdge} boundary condition @ rho = {rho}, T = {TkeV} keV", - typeMsg="i", - ) + print(f"- Producing {typeEdge} boundary condition @ rho = {rho}, T = {TkeV} keV",typeMsg="i",) for sp in range(len(self.Species)): if self.Species[sp]["S"] == "therm": - self.profiles["ti(keV)"][:, sp] = ( - self.profiles["ti(keV)"][:, sp] - * TkeV - / self.profiles["ti(keV)"][ix, sp] - ) + self.profiles["ti(keV)"][:, sp] = self.profiles["ti(keV)"][:, sp] * TkeV / self.profiles["ti(keV)"][ix, sp] if typeEdge == "linear": - self.profiles["te(keV)"][ix:] = np.linspace( - TkeV, Tesep, len(self.profiles["rho(-)"][ix:]) - ) + self.profiles["te(keV)"][ix:] = np.linspace(TkeV, Tesep, len(self.profiles["rho(-)"][ix:])) for sp in range(len(self.Species)): if self.Species[sp]["S"] == "therm": - self.profiles["ti(keV)"][ix:, sp] = np.linspace( - TkeV, Tisep, len(self.profiles["rho(-)"][ix:]) - ) + self.profiles["ti(keV)"][ix:, sp] = np.linspace(TkeV, Tisep, len(self.profiles["rho(-)"][ix:])) elif typeEdge == "same": pass else: raise Exception("no edge") - def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5): ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) - print( - f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3", - typeMsg="i", - ) + print(f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") factor = n20 / self.derived["ne_vol20"] @@ -109,13 +92,11 @@ def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5): / self.profiles["ne(10^19/m^3)"][ix:] ) - self.profiles["ne(10^19/m^3)"][ix:] = ( - self.profiles["ne(10^19/m^3)"][ix:] * factor_x - ) + self.profiles["ne(10^19/m^3)"][ix:] = self.profiles["ne(10^19/m^3)"][ix:] * factor_x + for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): - self.profiles["ni(10^19/m^3)"][ix:, i] = ( - self.profiles["ni(10^19/m^3)"][ix:, i] * factor_x - ) + self.profiles["ni(10^19/m^3)"][ix:, i] = self.profiles["ni(10^19/m^3)"][ix:, i] * factor_x + elif typeEdge == "same": pass else: @@ -528,9 +509,7 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): ----------------------------------------------------------------------------------- """ of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neo"] - of0E = ( - var_dict[f"{var}_tr_turb_stds"] ** 2 + var_dict[f"{var}_tr_neo_stds"] ** 2 - ) ** 0.5 + of0E = (var_dict[f"{var}_tr_turb_stds"] ** 2 + var_dict[f"{var}_tr_neo_stds"] ** 2) ** 0.5 """ ----------------------------------------------------------------------------------- From 37bba6115660d88a53b91e2931df9937637c155c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Apr 2025 08:59:28 -0700 Subject: [PATCH 029/385] Removed unnecessary PORTALSinteraction scripts --- src/mitim_modules/portals/PORTALSmain.py | 3 +- src/mitim_modules/portals/PORTALStools.py | 245 ++++++++ .../portals/utils/PORTALSanalysis.py | 8 +- .../physics_models/transport_tgyro.py | 195 ++++++- src/mitim_tools/gacode_tools/PROFILEStools.py | 100 +++- src/mitim_tools/gacode_tools/TGYROtools.py | 25 +- .../gacode_tools/utils/PORTALSinteraction.py | 536 ------------------ 7 files changed, 539 insertions(+), 573 deletions(-) delete mode 100644 src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 0ecbca15..a75d9351 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -7,7 +7,6 @@ from collections import OrderedDict from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.gacode_tools.utils import PORTALSinteraction from mitim_modules.portals import PORTALStools from mitim_modules.portals.utils import ( PORTALSinit, @@ -439,7 +438,7 @@ def scalarized_objective(self, Y): res must have shape (dim1...N) """ - of, cal, _, res = PORTALSinteraction.calculate_residuals(self.powerstate, self.PORTALSparameters,specific_vars=var_dict) + of, cal, _, res = PORTALStools.calculate_residuals(self.powerstate, self.PORTALSparameters,specific_vars=var_dict) return of, cal, res diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 8a333847..1b46f72c 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -404,3 +404,248 @@ def stopping_criteria_portals(mitim_bo, parameters = {}): else: print("\t- No convergence yet, providing as iteration values the scalarized objective") return False, yvals + + + +def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): + """ + Notes + ----- + - Works with tensors + - It should be independent on how many dimensions it has, except that the last dimension is the multi-ofs + """ + + # Case where I have already constructed the dictionary (i.e. in scalarized objective) + if specific_vars is not None: + var_dict = specific_vars + # Prepare dictionary from powerstate (for use in Analysis) + else: + var_dict = {} + + mapper = { + "Qe_tr_turb": "QeMWm2_tr_turb", + "Qi_tr_turb": "QiMWm2_tr_turb", + "Ge_tr_turb": "Ce_tr_turb", + "GZ_tr_turb": "CZ_tr_turb", + "Mt_tr_turb": "MtJm2_tr_turb", + "Qe_tr_neo": "QeMWm2_tr_neo", + "Qi_tr_neo": "QiMWm2_tr_neo", + "Ge_tr_neo": "Ce_tr_neo", + "GZ_tr_neo": "CZ_tr_neo", + "Mt_tr_neo": "MtJm2_tr_neo", + "Qe_tar": "QeMWm2", + "Qi_tar": "QiMWm2", + "Ge_tar": "Ce", + "GZ_tar": "CZ", + "Mt_tar": "MtJm2", + "PexchTurb": "PexchTurb" + } + + for ikey in mapper: + var_dict[ikey] = powerstate.plasma[mapper[ikey]][..., 1:] + if mapper[ikey] + "_stds" in powerstate.plasma: + var_dict[ikey + "_stds"] = powerstate.plasma[mapper[ikey] + "_stds"][..., 1:] + else: + var_dict[ikey + "_stds"] = None + + dfT = list(var_dict.values())[0] # as a reference for sizes + + # ------------------------------------------------------------------------- + # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added + # ------------------------------------------------------------------------- + + if PORTALSparameters["surrogateForTurbExch"]: + PexchTurb_integrated = computeTurbExchangeIndividual( + var_dict["PexchTurb"], powerstate + ) + else: + PexchTurb_integrated = torch.zeros(dfT.shape).to(dfT) + + # ------------------------------------------------------------------------ + # Go through each profile that needs to be predicted, calculate components + # ------------------------------------------------------------------------ + + of, cal, res = ( + torch.Tensor().to(dfT), + torch.Tensor().to(dfT), + torch.Tensor().to(dfT), + ) + for prof in powerstate.ProfilesPredicted: + if prof == "te": + var = "Qe" + elif prof == "ti": + var = "Qi" + elif prof == "ne": + var = "Ge" + elif prof == "nZ": + var = "GZ" + elif prof == "w0": + var = "Mt" + + """ + ----------------------------------------------------------------------------------- + Transport (_tr_turb+_tr_neo) + ----------------------------------------------------------------------------------- + """ + of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neo"] + + """ + ----------------------------------------------------------------------------------- + Target (Sum here the turbulent exchange power) + ----------------------------------------------------------------------------------- + """ + if var == "Qe": + cal0 = var_dict[f"{var}_tar"] + PexchTurb_integrated + elif var == "Qi": + cal0 = var_dict[f"{var}_tar"] - PexchTurb_integrated + else: + cal0 = var_dict[f"{var}_tar"] + + """ + ----------------------------------------------------------------------------------- + Ad-hoc modifications for different weighting + ----------------------------------------------------------------------------------- + """ + + if var == "Qe": + of0, cal0 = ( + of0 * PORTALSparameters["Pseudo_multipliers"][0], + cal0 * PORTALSparameters["Pseudo_multipliers"][0], + ) + elif var == "Qi": + of0, cal0 = ( + of0 * PORTALSparameters["Pseudo_multipliers"][1], + cal0 * PORTALSparameters["Pseudo_multipliers"][1], + ) + elif var == "Ge": + of0, cal0 = ( + of0 * PORTALSparameters["Pseudo_multipliers"][2], + cal0 * PORTALSparameters["Pseudo_multipliers"][2], + ) + elif var == "GZ": + of0, cal0 = ( + of0 * PORTALSparameters["Pseudo_multipliers"][3], + cal0 * PORTALSparameters["Pseudo_multipliers"][3], + ) + elif var == "MtJm2": + of0, cal0 = ( + of0 * PORTALSparameters["Pseudo_multipliers"][4], + cal0 * PORTALSparameters["Pseudo_multipliers"][4], + ) + + of, cal = torch.cat((of, of0), dim=-1), torch.cat((cal, cal0), dim=-1) + + # ----------- + # Composition + # ----------- + + # 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 + res = -1 / source.shape[-1] * torch.norm(source, p=2, dim=-1) + + return of, cal, source, res + + +def calculate_residuals_distributions(powerstate, PORTALSparameters): + """ + - Works with tensors + - It should be independent on how many dimensions it has, except that the last dimension is the multi-ofs + """ + + # Prepare dictionary from powerstate (for use in Analysis) + + mapper = { + "Qe_tr_turb": "QeMWm2_tr_turb", + "Qi_tr_turb": "QiMWm2_tr_turb", + "Ge_tr_turb": "Ce_tr_turb", + "GZ_tr_turb": "CZ_tr_turb", + "Mt_tr_turb": "MtJm2_tr_turb", + "Qe_tr_neo": "QeMWm2_tr_neo", + "Qi_tr_neo": "QiMWm2_tr_neo", + "Ge_tr_neo": "Ce_tr_neo", + "GZ_tr_neo": "CZ_tr_neo", + "Mt_tr_neo": "MtJm2_tr_neo", + "Qe_tar": "QeMWm2", + "Qi_tar": "QiMWm2", + "Ge_tar": "Ce", + "GZ_tar": "CZ", + "Mt_tar": "MtJm2", + "PexchTurb": "PexchTurb" + } + + var_dict = {} + for ikey in mapper: + var_dict[ikey] = powerstate.plasma[mapper[ikey]][:, 1:] + if mapper[ikey] + "_stds" in powerstate.plasma: + var_dict[ikey + "_stds"] = powerstate.plasma[mapper[ikey] + "_stds"][:, 1:] + else: + var_dict[ikey + "_stds"] = None + + dfT = var_dict["Qe_tr_turb"] # as a reference for sizes + + # ------------------------------------------------------------------------- + # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added + # ------------------------------------------------------------------------- + + if PORTALSparameters["surrogateForTurbExch"]: + PexchTurb_integrated = computeTurbExchangeIndividual( + var_dict["PexchTurb"], powerstate + ) + PexchTurb_integrated_stds = computeTurbExchangeIndividual( + var_dict["PexchTurb_stds"], powerstate + ) + else: + PexchTurb_integrated = torch.zeros(dfT.shape).to(dfT) + PexchTurb_integrated_stds = torch.zeros(dfT.shape).to(dfT) + + # ------------------------------------------------------------------------ + # Go through each profile that needs to be predicted, calculate components + # ------------------------------------------------------------------------ + + of, cal = torch.Tensor().to(dfT), torch.Tensor().to(dfT) + ofE, calE = torch.Tensor().to(dfT), torch.Tensor().to(dfT) + for prof in powerstate.ProfilesPredicted: + if prof == "te": + var = "Qe" + elif prof == "ti": + var = "Qi" + elif prof == "ne": + var = "Ge" + elif prof == "nZ": + var = "GZ" + elif prof == "w0": + var = "MtJm2" + + """ + ----------------------------------------------------------------------------------- + Transport (_tr_turb+_tr_neo) + ----------------------------------------------------------------------------------- + """ + of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neo"] + of0E = (var_dict[f"{var}_tr_turb_stds"] ** 2 + var_dict[f"{var}_tr_neo_stds"] ** 2) ** 0.5 + + """ + ----------------------------------------------------------------------------------- + Target (Sum here the turbulent exchange power) + ----------------------------------------------------------------------------------- + """ + if var == "Qe": + cal0 = var_dict[f"{var}_tar"] + PexchTurb_integrated + cal0E = ( + var_dict[f"{var}_tar_stds"] ** 2 + PexchTurb_integrated_stds**2 + ) ** 0.5 + elif var == "Qi": + cal0 = var_dict[f"{var}_tar"] - PexchTurb_integrated + cal0E = ( + var_dict[f"{var}_tar_stds"] ** 2 + PexchTurb_integrated_stds**2 + ) ** 0.5 + else: + cal0 = var_dict[f"{var}_tar"] + cal0E = var_dict[f"{var}_tar_stds"] + + of, cal = torch.cat((of, of0), dim=-1), torch.cat((cal, cal0), dim=-1) + ofE, calE = torch.cat((ofE, of0E), dim=-1), torch.cat((calE, cal0E), dim=-1) + + return of, cal, ofE, calE diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 5078e5ae..81192888 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -6,8 +6,8 @@ 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.utils import PORTALSinteraction from mitim_modules.portals.utils import PORTALSplot +from mitim_modules.portals import PORTALStools from mitim_modules.powertorch import STATEtools from mitim_modules.powertorch.utils import POWERplot from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -228,7 +228,7 @@ def prep_metrics(self, ilast=None): # Residual definitions # ------------------------------------------------ - _, _, source, res = PORTALSinteraction.calculate_residuals( + _, _, source, res = PORTALStools.calculate_residuals( power, self.PORTALSparameters, ) @@ -272,7 +272,7 @@ def prep_metrics(self, ilast=None): y2, y1_std, y2_std, - ) = PORTALSinteraction.calculate_residuals_distributions( + ) = PORTALStools.calculate_residuals_distributions( power, self.PORTALSparameters, ) @@ -1075,3 +1075,5 @@ def plotMetrics(self, extra_lab="", **kwargs): axs[0].legend(prop={"size": 8}) axsGrads[0].legend(prop={"size": 8}) + + diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 01f7fa85..c17169ae 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -1,3 +1,4 @@ +import torch import copy import shutil import numpy as np @@ -151,7 +152,8 @@ def _postprocess_results(self, tgyro, label): impurityPosition = self.powerstate.impurityPosition_transport # Produce right quantities (TGYRO -> powerstate.plasma) - self.powerstate = tgyro.results[label].tgyro_to_powerstate( + self.powerstate = tgyro_to_powerstate( + tgyro.results[label], self.powerstate, useConvectiveFluxes=useConvectiveFluxes, includeFast=includeFast, @@ -872,3 +874,194 @@ def defineReferenceFluxes( return Qe_target, Qi_target, Ge_target_special, GZ_target_special, Mt_target + + +# ------------------------------------------------------------------------------------------------------------------------------------------------------ +# This is where the definitions for the summation variables happen for mitim and PORTALSplot +# ------------------------------------------------------------------------------------------------------------------------------------------------------ + +def tgyro_to_powerstate(TGYROresults, + powerstate, + useConvectiveFluxes=False, + forceZeroParticleFlux=False, + includeFast=False, + impurityPosition=1, + UseFineGridTargets=False, + OriginalFimp=1.0, + provideTurbulentExchange=False, + provideTargets=False + ): + """ + This function is used to extract the TGYRO results and store them in the powerstate object, from numpy arrays to torch tensors. + """ + + if "tgyro_stds" not in TGYROresults.__dict__: + TGYROresults.tgyro_stds = False + + if UseFineGridTargets: + TGYROresults.useFineGridTargets(impurityPosition=impurityPosition) + + nr = powerstate.plasma['rho'].shape[-1] + if powerstate.plasma['rho'].shape[-1] != TGYROresults.rho.shape[-1]: + print('\t- TGYRO was run with an extra point in the grid, treating it carefully now') + + # ********************************** + # *********** Electron Energy Fluxes + # ********************************** + + powerstate.plasma["QeMWm2_tr_turb"] = torch.Tensor(TGYROresults.Qe_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QeMWm2_tr_neo"] = torch.Tensor(TGYROresults.Qe_sim_neo[:, :nr]).to(powerstate.dfT) + + powerstate.plasma["QeMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Qe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QeMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Qe_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["QeMWm2"] = torch.Tensor(TGYROresults.Qe_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QeMWm2_stds"] = torch.Tensor(TGYROresults.Qe_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + # ********************************** + # *********** Ion Energy Fluxes + # ********************************** + + if includeFast: + + powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo[:, :nr]).to(powerstate.dfT) + + powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QiMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + else: + + powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr[:, :nr]).to(powerstate.dfT) + + powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QiMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["QiMWm2"] = torch.Tensor(TGYROresults.Qi_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_stds"] = torch.Tensor(TGYROresults.Qi_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + # ********************************** + # *********** Momentum Fluxes + # ********************************** + + powerstate.plasma["MtJm2_tr_turb"] = torch.Tensor(TGYROresults.Mt_sim_turb[:, :nr]).to(powerstate.dfT) # So far, let's include fast in momentum + powerstate.plasma["MtJm2_tr_neo"] = torch.Tensor(TGYROresults.Mt_sim_neo[:, :nr]).to(powerstate.dfT) + + powerstate.plasma["MtJm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Mt_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["MtJm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Mt_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["MtJm2"] = torch.Tensor(TGYROresults.Mt_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["MtJm2_stds"] = torch.Tensor(TGYROresults.Mt_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + # ********************************** + # *********** Particle Fluxes + # ********************************** + + # Store raw fluxes for better plotting later + powerstate.plasma["Ge1E20sm2_tr_turb"] = torch.Tensor(TGYROresults.Ge_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ge1E20sm2_tr_neo"] = torch.Tensor(TGYROresults.Ge_sim_neo[:, :nr]).to(powerstate.dfT) + + powerstate.plasma["Ge1E20sm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Ge_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ge1E20sm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Ge_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["Ge1E20sm2"] = torch.Tensor(TGYROresults.Ge_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ge1E20sm2_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + if not useConvectiveFluxes: + + powerstate.plasma["Ce_tr_turb"] = powerstate.plasma["Ge1E20sm2_tr_turb"] + powerstate.plasma["Ce_tr_neo"] = powerstate.plasma["Ge1E20sm2_tr_neo"] + + powerstate.plasma["Ce_tr_turb_stds"] = powerstate.plasma["Ge1E20sm2_tr_turb_stds"] + powerstate.plasma["Ce_tr_neo_stds"] = powerstate.plasma["Ge1E20sm2_tr_neo_stds"] + + if provideTargets: + powerstate.plasma["Ce"] = powerstate.plasma["Ge1E20sm2"] + powerstate.plasma["Ce_stds"] = powerstate.plasma["Ge1E20sm2_stds"] + + else: + + powerstate.plasma["Ce_tr_turb"] = torch.Tensor(TGYROresults.Ce_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ce_tr_neo"] = torch.Tensor(TGYROresults.Ce_sim_neo[:, :nr]).to(powerstate.dfT) + + powerstate.plasma["Ce_tr_turb_stds"] = torch.Tensor(TGYROresults.Ce_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ce_tr_neo_stds"] = torch.Tensor(TGYROresults.Ce_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["Ce"] = torch.Tensor(TGYROresults.Ce_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ce_stds"] = torch.Tensor(TGYROresults.Ce_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + # ********************************** + # *********** Impurity Fluxes + # ********************************** + + # Store raw fluxes for better plotting later + powerstate.plasma["CZ_raw_tr_turb"] = torch.Tensor(TGYROresults.Gi_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) + powerstate.plasma["CZ_raw_tr_neo"] = torch.Tensor(TGYROresults.Gi_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) + + powerstate.plasma["CZ_raw_tr_turb_stds"] = torch.Tensor(TGYROresults.Gi_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["CZ_raw_tr_neo_stds"] = torch.Tensor(TGYROresults.Gi_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["CZ_raw"] = torch.Tensor(TGYROresults.Gi_tar[impurityPosition, :, :nr]).to(powerstate.dfT) + powerstate.plasma["CZ_raw_stds"] = torch.Tensor(TGYROresults.Gi_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + if not useConvectiveFluxes: + + powerstate.plasma["CZ_tr_turb"] = powerstate.plasma["CZ_raw_tr_turb"] / OriginalFimp + powerstate.plasma["CZ_tr_neo"] = powerstate.plasma["CZ_raw_tr_neo"] / OriginalFimp + + powerstate.plasma["CZ_tr_turb_stds"] = powerstate.plasma["CZ_raw_tr_turb_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None + powerstate.plasma["CZ_tr_neo_stds"] = powerstate.plasma["CZ_raw_tr_neo_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["CZ"] = powerstate.plasma["CZ_raw"] / OriginalFimp + powerstate.plasma["CZ_stds"] = powerstate.plasma["CZ_raw_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None + + else: + + powerstate.plasma["CZ_tr_turb"] = torch.Tensor(TGYROresults.Ci_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp + powerstate.plasma["CZ_tr_neo"] = torch.Tensor(TGYROresults.Ci_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp + + powerstate.plasma["CZ_tr_turb_stds"] = torch.Tensor(TGYROresults.Ci_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None + powerstate.plasma["CZ_tr_neo_stds"] = torch.Tensor(TGYROresults.Ci_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["CZ"] = torch.Tensor(TGYROresults.Ci_tar[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp + powerstate.plasma["CZ_stds"] = torch.Tensor(TGYROresults.Ci_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None + + # ********************************** + # *********** Energy Exchange + # ********************************** + + if provideTurbulentExchange: + powerstate.plasma["PexchTurb"] = torch.Tensor(TGYROresults.EXe_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["PexchTurb_stds"] = torch.Tensor(TGYROresults.EXe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + else: + powerstate.plasma["PexchTurb"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + powerstate.plasma["PexchTurb_stds"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + + # ********************************** + # *********** Traget extra + # ********************************** + + if forceZeroParticleFlux and provideTargets: + powerstate.plasma["Ce"] = powerstate.plasma["Ce"] * 0.0 + powerstate.plasma["Ce_stds"] = powerstate.plasma["Ce_stds"] * 0.0 + + # ------------------------------------------------------------------------------------------------------------------------ + # Sum here turbulence and neoclassical, after modifications + # ------------------------------------------------------------------------------------------------------------------------ + + quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2', 'Ge1E20sm2', 'CZ_raw'] + for ikey in quantities: + powerstate.plasma[ikey+"_tr"] = powerstate.plasma[ikey+"_tr_turb"] + powerstate.plasma[ikey+"_tr_neo"] + + return powerstate + + diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 592bf70d..bc111811 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -11,7 +11,6 @@ from mitim_tools.gacode_tools.utils import GACODEdefaults from mitim_tools.transp_tools import CDFtools from mitim_tools.transp_tools.utils import TRANSPhelpers -from mitim_tools.gacode_tools.utils import PORTALSinteraction from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __version__ from IPython import embed @@ -3809,21 +3808,104 @@ def csv(self, file="input.gacode.xlsx"): IOtools.writeExcel_fromDict(dictExcel, file, fromRow=1) def parabolizePlasma(self): - PORTALSinteraction.parabolizePlasma(self) + _, T = PLASMAtools.parabolicProfile( + Tbar=self.derived["Te_vol"], + nu=self.derived["Te_peaking"], + rho=self.profiles["rho(-)"], + Tedge=self.profiles["te(keV)"][-1], + ) + _, Ti = PLASMAtools.parabolicProfile( + Tbar=self.derived["Ti_vol"], + nu=self.derived["Ti_peaking"], + rho=self.profiles["rho(-)"], + Tedge=self.profiles["ti(keV)"][-1, 0], + ) + _, n = PLASMAtools.parabolicProfile( + Tbar=self.derived["ne_vol20"] * 1e1, + nu=self.derived["ne_peaking"], + rho=self.profiles["rho(-)"], + Tedge=self.profiles["ne(10^19/m^3)"][-1], + ) + + self.profiles["te(keV)"] = T + + self.profiles["ti(keV)"][:, 0] = Ti + self.makeAllThermalIonsHaveSameTemp(refIon=0) + + factor_n = n / self.profiles["ne(10^19/m^3)"] + self.profiles["ne(10^19/m^3)"] = n + self.scaleAllThermalDensities(scaleFactor=factor_n) + + self.deriveQuantities() + def changeRFpower(self, PrfMW=25.0): - PORTALSinteraction.changeRFpower(self, PrfMW=PrfMW) + """ + keeps same partition + """ + print(f"- Changing the RF power from {self.derived['qRF_MWmiller'][-1]:.1f} MW to {PrfMW:.1f} MW",typeMsg="i",) + + if self.derived["qRF_MWmiller"][-1] == 0.0: + raise Exception("No RF power in the input.gacode, cannot modify the RF power") + + for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: + self.profiles[i] = self.profiles[i] * PrfMW / self.derived["qRF_MWmiller"][-1] def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0.2): - PORTALSinteraction.imposeBCtemps( - self, TkeV=TkeV, rho=rho, typeEdge=typeEdge, Tesep=Tesep, Tisep=Tisep - ) + + ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) + + self.profiles["te(keV)"] = self.profiles["te(keV)"] * TkeV / self.profiles["te(keV)"][ix] + + print(f"- Producing {typeEdge} boundary condition @ rho = {rho}, T = {TkeV} keV",typeMsg="i",) + + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "therm": + self.profiles["ti(keV)"][:, sp] = self.profiles["ti(keV)"][:, sp] * TkeV / self.profiles["ti(keV)"][ix, sp] + + if typeEdge == "linear": + self.profiles["te(keV)"][ix:] = np.linspace(TkeV, Tesep, len(self.profiles["rho(-)"][ix:])) + + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "therm": + self.profiles["ti(keV)"][ix:, sp] = np.linspace(TkeV, Tisep, len(self.profiles["rho(-)"][ix:])) + + elif typeEdge == "same": + pass + else: + raise Exception("no edge") + def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5): - PORTALSinteraction.imposeBCdens( - self, n20=n20, rho=rho, typeEdge=typeEdge, nedge20=nedge20 - ) + ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) + + print(f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") + factor = n20 / self.derived["ne_vol20"] + + for i in ["ne(10^19/m^3)", "ni(10^19/m^3)"]: + self.profiles[i] = self.profiles[i] * factor + + if typeEdge == "linear": + factor_x = ( + np.linspace( + self.profiles["ne(10^19/m^3)"][ix], + nedge20 * 1e1, + len(self.profiles["rho(-)"][ix:]), + ) + / self.profiles["ne(10^19/m^3)"][ix:] + ) + + self.profiles["ne(10^19/m^3)"][ix:] = self.profiles["ne(10^19/m^3)"][ix:] * factor_x + + for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): + self.profiles["ni(10^19/m^3)"][ix:, i] = self.profiles["ni(10^19/m^3)"][ix:, i] * factor_x + + elif typeEdge == "same": + pass + else: + raise Exception("no edge") + def addSawtoothEffectOnOhmic(self, PohTot, mixRadius=None, plotYN=False): """ This will implement a flat profile inside the mixRadius to reduce the ohmic power by certain amount diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index fb7cfcc8..79d3b4fe 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -12,15 +12,6 @@ from mitim_tools.gacode_tools.utils import GACODEinterpret, GACODEdefaults, GACODErun from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -import time - -try: - from mitim_tools.gacode_tools.utils import PORTALSinteraction -except: - print( - "- I could not import PORTALSinteraction, likely a consequence of botorch incompatbility", - typeMsg="w", - ) """ Same philosophy as the TGLFtools @@ -2100,9 +2091,7 @@ def derived(self): Note: This is only valid in the converged case??????????????? """ - if (self.profiles_final is not None) and ( - "derived" in self.profiles_final.__dict__ - ): + if (self.profiles_final is not None) and ("derived" in self.profiles_final.__dict__): prof = self.profiles_final elif (self.profiles is not None) and ("derived" in self.profiles.__dict__): prof = self.profiles @@ -2116,9 +2105,7 @@ def derived(self): self.Q_better = self.P_fusT_tgyro / self.P_inT - if (self.profiles_final is not None) and ( - "derived" in self.profiles_final.__dict__ - ): + if (self.profiles_final is not None) and ("derived" in self.profiles_final.__dict__): self.Q_best = self.profiles_final.derived["Q"] """ @@ -2170,10 +2157,7 @@ def useFineGridTargets(self, impurityPosition=1): ) # Profiles do not include ion fluxes for j in range(self.Gi_tar.shape[0]): - self.Gi_tar[j, i, :], self.Ci_tar[j, i, :] = ( - self.Ce_tar[i, :] * 0.0, - self.Ce_tar[i, :] * 0.0, - ) + self.Gi_tar[j, i, :], self.Ci_tar[j, i, :] = self.Ce_tar[i, :] * 0.0, self.Ce_tar[i, :] * 0.0 self.Mt_tar[i, :] = np.interp( rho_coarse, rho_fine, self.profiles_final.derived["mt_Jm2"] @@ -2185,9 +2169,6 @@ def useFineGridTargets(self, impurityPosition=1): self.Ge_tarMW = self.Ge_tar * self.dvoldr self.Ce_tarMW = self.Ce_tar * self.dvoldr - def tgyro_to_powerstate(self, *args, **kwargs): - return PORTALSinteraction.tgyro_to_powerstate(self, *args, **kwargs) - def plot(self, fn=None, label="", prelabel="", fn_color=None): if fn is None: from mitim_tools.misc_tools.GUItools import FigureNotebook diff --git a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py b/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py deleted file mode 100644 index 3ae6c89a..00000000 --- a/src/mitim_tools/gacode_tools/utils/PORTALSinteraction.py +++ /dev/null @@ -1,536 +0,0 @@ -import torch -import numpy as np -from mitim_tools.misc_tools import PLASMAtools -from mitim_modules.portals import PORTALStools -from mitim_tools.misc_tools.LOGtools import printMsg as print -from IPython import embed - -def parabolizePlasma(self): - _, T = PLASMAtools.parabolicProfile( - Tbar=self.derived["Te_vol"], - nu=self.derived["Te_peaking"], - rho=self.profiles["rho(-)"], - Tedge=self.profiles["te(keV)"][-1], - ) - _, Ti = PLASMAtools.parabolicProfile( - Tbar=self.derived["Ti_vol"], - nu=self.derived["Ti_peaking"], - rho=self.profiles["rho(-)"], - Tedge=self.profiles["ti(keV)"][-1, 0], - ) - _, n = PLASMAtools.parabolicProfile( - Tbar=self.derived["ne_vol20"] * 1e1, - nu=self.derived["ne_peaking"], - rho=self.profiles["rho(-)"], - Tedge=self.profiles["ne(10^19/m^3)"][-1], - ) - - self.profiles["te(keV)"] = T - - self.profiles["ti(keV)"][:, 0] = Ti - self.makeAllThermalIonsHaveSameTemp(refIon=0) - - factor_n = n / self.profiles["ne(10^19/m^3)"] - self.profiles["ne(10^19/m^3)"] = n - self.scaleAllThermalDensities(scaleFactor=factor_n) - - self.deriveQuantities() - - -def changeRFpower(self, PrfMW=25.0): - """ - keeps same partition - """ - print(f"- Changing the RF power from {self.derived['qRF_MWmiller'][-1]:.1f} MW to {PrfMW:.1f} MW",typeMsg="i",) - - if self.derived["qRF_MWmiller"][-1] == 0.0: - raise Exception("No RF power in the input.gacode, cannot modify the RF power") - - for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: - self.profiles[i] = self.profiles[i] * PrfMW / self.derived["qRF_MWmiller"][-1] - -def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0.2): - ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) - - self.profiles["te(keV)"] = self.profiles["te(keV)"] * TkeV / self.profiles["te(keV)"][ix] - - print(f"- Producing {typeEdge} boundary condition @ rho = {rho}, T = {TkeV} keV",typeMsg="i",) - - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm": - self.profiles["ti(keV)"][:, sp] = self.profiles["ti(keV)"][:, sp] * TkeV / self.profiles["ti(keV)"][ix, sp] - - if typeEdge == "linear": - self.profiles["te(keV)"][ix:] = np.linspace(TkeV, Tesep, len(self.profiles["rho(-)"][ix:])) - - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm": - self.profiles["ti(keV)"][ix:, sp] = np.linspace(TkeV, Tisep, len(self.profiles["rho(-)"][ix:])) - - elif typeEdge == "same": - pass - else: - raise Exception("no edge") - -def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5): - ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) - - print(f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") - - factor = n20 / self.derived["ne_vol20"] - - for i in ["ne(10^19/m^3)", "ni(10^19/m^3)"]: - self.profiles[i] = self.profiles[i] * factor - - if typeEdge == "linear": - factor_x = ( - np.linspace( - self.profiles["ne(10^19/m^3)"][ix], - nedge20 * 1e1, - len(self.profiles["rho(-)"][ix:]), - ) - / self.profiles["ne(10^19/m^3)"][ix:] - ) - - self.profiles["ne(10^19/m^3)"][ix:] = self.profiles["ne(10^19/m^3)"][ix:] * factor_x - - for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): - self.profiles["ni(10^19/m^3)"][ix:, i] = self.profiles["ni(10^19/m^3)"][ix:, i] * factor_x - - elif typeEdge == "same": - pass - else: - raise Exception("no edge") - - -# ------------------------------------------------------------------------------------------------------------------------------------------------------ -# This is where the definitions for the summation variables happen for mitim and PORTALSplot -# ------------------------------------------------------------------------------------------------------------------------------------------------------ - -def tgyro_to_powerstate(TGYROresults, - powerstate, - useConvectiveFluxes=False, - forceZeroParticleFlux=False, - includeFast=False, - impurityPosition=1, - UseFineGridTargets=False, - OriginalFimp=1.0, - provideTurbulentExchange=False, - provideTargets=False - ): - """ - This function is used to extract the TGYRO results and store them in the powerstate object, from numpy arrays to torch tensors. - """ - - if "tgyro_stds" not in TGYROresults.__dict__: - TGYROresults.tgyro_stds = False - - if UseFineGridTargets: - TGYROresults.useFineGridTargets(impurityPosition=impurityPosition) - - nr = powerstate.plasma['rho'].shape[-1] - if powerstate.plasma['rho'].shape[-1] != TGYROresults.rho.shape[-1]: - print('\t- TGYRO was run with an extra point in the grid, treating it carefully now') - - # ********************************** - # *********** Electron Energy Fluxes - # ********************************** - - powerstate.plasma["QeMWm2_tr_turb"] = torch.Tensor(TGYROresults.Qe_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QeMWm2_tr_neo"] = torch.Tensor(TGYROresults.Qe_sim_neo[:, :nr]).to(powerstate.dfT) - - powerstate.plasma["QeMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Qe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["QeMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Qe_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["QeMWm2"] = torch.Tensor(TGYROresults.Qe_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QeMWm2_stds"] = torch.Tensor(TGYROresults.Qe_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - # ********************************** - # *********** Ion Energy Fluxes - # ********************************** - - if includeFast: - - powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QiMWm2_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo[:, :nr]).to(powerstate.dfT) - - powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["QiMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - else: - - powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QiMWm2_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr[:, :nr]).to(powerstate.dfT) - - powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["QiMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["QiMWm2"] = torch.Tensor(TGYROresults.Qi_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QiMWm2_stds"] = torch.Tensor(TGYROresults.Qi_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - # ********************************** - # *********** Momentum Fluxes - # ********************************** - - powerstate.plasma["MtJm2_tr_turb"] = torch.Tensor(TGYROresults.Mt_sim_turb[:, :nr]).to(powerstate.dfT) # So far, let's include fast in momentum - powerstate.plasma["MtJm2_tr_neo"] = torch.Tensor(TGYROresults.Mt_sim_neo[:, :nr]).to(powerstate.dfT) - - powerstate.plasma["MtJm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Mt_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["MtJm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Mt_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["MtJm2"] = torch.Tensor(TGYROresults.Mt_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["MtJm2_stds"] = torch.Tensor(TGYROresults.Mt_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - # ********************************** - # *********** Particle Fluxes - # ********************************** - - # Store raw fluxes for better plotting later - powerstate.plasma["Ge1E20sm2_tr_turb"] = torch.Tensor(TGYROresults.Ge_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ge1E20sm2_tr_neo"] = torch.Tensor(TGYROresults.Ge_sim_neo[:, :nr]).to(powerstate.dfT) - - powerstate.plasma["Ge1E20sm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Ge_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Ge1E20sm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Ge_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["Ge1E20sm2"] = torch.Tensor(TGYROresults.Ge_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ge1E20sm2_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if not useConvectiveFluxes: - - powerstate.plasma["Ce_tr_turb"] = powerstate.plasma["Ge1E20sm2_tr_turb"] - powerstate.plasma["Ce_tr_neo"] = powerstate.plasma["Ge1E20sm2_tr_neo"] - - powerstate.plasma["Ce_tr_turb_stds"] = powerstate.plasma["Ge1E20sm2_tr_turb_stds"] - powerstate.plasma["Ce_tr_neo_stds"] = powerstate.plasma["Ge1E20sm2_tr_neo_stds"] - - if provideTargets: - powerstate.plasma["Ce"] = powerstate.plasma["Ge1E20sm2"] - powerstate.plasma["Ce_stds"] = powerstate.plasma["Ge1E20sm2_stds"] - - else: - - powerstate.plasma["Ce_tr_turb"] = torch.Tensor(TGYROresults.Ce_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_tr_neo"] = torch.Tensor(TGYROresults.Ce_sim_neo[:, :nr]).to(powerstate.dfT) - - powerstate.plasma["Ce_tr_turb_stds"] = torch.Tensor(TGYROresults.Ce_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Ce_tr_neo_stds"] = torch.Tensor(TGYROresults.Ce_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["Ce"] = torch.Tensor(TGYROresults.Ce_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_stds"] = torch.Tensor(TGYROresults.Ce_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - # ********************************** - # *********** Impurity Fluxes - # ********************************** - - # Store raw fluxes for better plotting later - powerstate.plasma["CZ_raw_tr_turb"] = torch.Tensor(TGYROresults.Gi_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) - powerstate.plasma["CZ_raw_tr_neo"] = torch.Tensor(TGYROresults.Gi_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) - - powerstate.plasma["CZ_raw_tr_turb_stds"] = torch.Tensor(TGYROresults.Gi_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_raw_tr_neo_stds"] = torch.Tensor(TGYROresults.Gi_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["CZ_raw"] = torch.Tensor(TGYROresults.Gi_tar[impurityPosition, :, :nr]).to(powerstate.dfT) - powerstate.plasma["CZ_raw_stds"] = torch.Tensor(TGYROresults.Gi_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if not useConvectiveFluxes: - - powerstate.plasma["CZ_tr_turb"] = powerstate.plasma["CZ_raw_tr_turb"] / OriginalFimp - powerstate.plasma["CZ_tr_neo"] = powerstate.plasma["CZ_raw_tr_neo"] / OriginalFimp - - powerstate.plasma["CZ_tr_turb_stds"] = powerstate.plasma["CZ_raw_tr_turb_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_tr_neo_stds"] = powerstate.plasma["CZ_raw_tr_neo_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["CZ"] = powerstate.plasma["CZ_raw"] / OriginalFimp - powerstate.plasma["CZ_stds"] = powerstate.plasma["CZ_raw_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None - - else: - - powerstate.plasma["CZ_tr_turb"] = torch.Tensor(TGYROresults.Ci_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - powerstate.plasma["CZ_tr_neo"] = torch.Tensor(TGYROresults.Ci_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - - powerstate.plasma["CZ_tr_turb_stds"] = torch.Tensor(TGYROresults.Ci_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_tr_neo_stds"] = torch.Tensor(TGYROresults.Ci_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["CZ"] = torch.Tensor(TGYROresults.Ci_tar[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - powerstate.plasma["CZ_stds"] = torch.Tensor(TGYROresults.Ci_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None - - # ********************************** - # *********** Energy Exchange - # ********************************** - - if provideTurbulentExchange: - powerstate.plasma["PexchTurb"] = torch.Tensor(TGYROresults.EXe_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["PexchTurb_stds"] = torch.Tensor(TGYROresults.EXe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - else: - powerstate.plasma["PexchTurb"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 - powerstate.plasma["PexchTurb_stds"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 - - # ********************************** - # *********** Traget extra - # ********************************** - - if forceZeroParticleFlux and provideTargets: - powerstate.plasma["Ce"] = powerstate.plasma["Ce"] * 0.0 - powerstate.plasma["Ce_stds"] = powerstate.plasma["Ce_stds"] * 0.0 - - # ------------------------------------------------------------------------------------------------------------------------ - # Sum here turbulence and neoclassical, after modifications - # ------------------------------------------------------------------------------------------------------------------------ - - quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2', 'Ge1E20sm2', 'CZ_raw'] - for ikey in quantities: - powerstate.plasma[ikey+"_tr"] = powerstate.plasma[ikey+"_tr_turb"] + powerstate.plasma[ikey+"_tr_neo"] - - return powerstate - - -def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): - """ - Notes - ----- - - Works with tensors - - It should be independent on how many dimensions it has, except that the last dimension is the multi-ofs - """ - - # Case where I have already constructed the dictionary (i.e. in scalarized objective) - if specific_vars is not None: - var_dict = specific_vars - # Prepare dictionary from powerstate (for use in Analysis) - else: - var_dict = {} - - mapper = { - "Qe_tr_turb": "QeMWm2_tr_turb", - "Qi_tr_turb": "QiMWm2_tr_turb", - "Ge_tr_turb": "Ce_tr_turb", - "GZ_tr_turb": "CZ_tr_turb", - "Mt_tr_turb": "MtJm2_tr_turb", - "Qe_tr_neo": "QeMWm2_tr_neo", - "Qi_tr_neo": "QiMWm2_tr_neo", - "Ge_tr_neo": "Ce_tr_neo", - "GZ_tr_neo": "CZ_tr_neo", - "Mt_tr_neo": "MtJm2_tr_neo", - "Qe_tar": "QeMWm2", - "Qi_tar": "QiMWm2", - "Ge_tar": "Ce", - "GZ_tar": "CZ", - "Mt_tar": "MtJm2", - "PexchTurb": "PexchTurb" - } - - for ikey in mapper: - var_dict[ikey] = powerstate.plasma[mapper[ikey]][..., 1:] - if mapper[ikey] + "_stds" in powerstate.plasma: - var_dict[ikey + "_stds"] = powerstate.plasma[mapper[ikey] + "_stds"][..., 1:] - else: - var_dict[ikey + "_stds"] = None - - dfT = list(var_dict.values())[0] # as a reference for sizes - - # ------------------------------------------------------------------------- - # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added - # ------------------------------------------------------------------------- - - if PORTALSparameters["surrogateForTurbExch"]: - PexchTurb_integrated = PORTALStools.computeTurbExchangeIndividual( - var_dict["PexchTurb"], powerstate - ) - else: - PexchTurb_integrated = torch.zeros(dfT.shape).to(dfT) - - # ------------------------------------------------------------------------ - # Go through each profile that needs to be predicted, calculate components - # ------------------------------------------------------------------------ - - of, cal, res = ( - torch.Tensor().to(dfT), - torch.Tensor().to(dfT), - torch.Tensor().to(dfT), - ) - for prof in powerstate.ProfilesPredicted: - if prof == "te": - var = "Qe" - elif prof == "ti": - var = "Qi" - elif prof == "ne": - var = "Ge" - elif prof == "nZ": - var = "GZ" - elif prof == "w0": - var = "Mt" - - """ - ----------------------------------------------------------------------------------- - Transport (_tr_turb+_tr_neo) - ----------------------------------------------------------------------------------- - """ - of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neo"] - - """ - ----------------------------------------------------------------------------------- - Target (Sum here the turbulent exchange power) - ----------------------------------------------------------------------------------- - """ - if var == "Qe": - cal0 = var_dict[f"{var}_tar"] + PexchTurb_integrated - elif var == "Qi": - cal0 = var_dict[f"{var}_tar"] - PexchTurb_integrated - else: - cal0 = var_dict[f"{var}_tar"] - - """ - ----------------------------------------------------------------------------------- - Ad-hoc modifications for different weighting - ----------------------------------------------------------------------------------- - """ - - if var == "Qe": - of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][0], - cal0 * PORTALSparameters["Pseudo_multipliers"][0], - ) - elif var == "Qi": - of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][1], - cal0 * PORTALSparameters["Pseudo_multipliers"][1], - ) - elif var == "Ge": - of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][2], - cal0 * PORTALSparameters["Pseudo_multipliers"][2], - ) - elif var == "GZ": - of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][3], - cal0 * PORTALSparameters["Pseudo_multipliers"][3], - ) - elif var == "MtJm2": - of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][4], - cal0 * PORTALSparameters["Pseudo_multipliers"][4], - ) - - of, cal = torch.cat((of, of0), dim=-1), torch.cat((cal, cal0), dim=-1) - - # ----------- - # Composition - # ----------- - - # 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 - res = -1 / source.shape[-1] * torch.norm(source, p=2, dim=-1) - - return of, cal, source, res - - -def calculate_residuals_distributions(powerstate, PORTALSparameters): - """ - - Works with tensors - - It should be independent on how many dimensions it has, except that the last dimension is the multi-ofs - """ - - # Prepare dictionary from powerstate (for use in Analysis) - - mapper = { - "Qe_tr_turb": "QeMWm2_tr_turb", - "Qi_tr_turb": "QiMWm2_tr_turb", - "Ge_tr_turb": "Ce_tr_turb", - "GZ_tr_turb": "CZ_tr_turb", - "Mt_tr_turb": "MtJm2_tr_turb", - "Qe_tr_neo": "QeMWm2_tr_neo", - "Qi_tr_neo": "QiMWm2_tr_neo", - "Ge_tr_neo": "Ce_tr_neo", - "GZ_tr_neo": "CZ_tr_neo", - "Mt_tr_neo": "MtJm2_tr_neo", - "Qe_tar": "QeMWm2", - "Qi_tar": "QiMWm2", - "Ge_tar": "Ce", - "GZ_tar": "CZ", - "Mt_tar": "MtJm2", - "PexchTurb": "PexchTurb" - } - - var_dict = {} - for ikey in mapper: - var_dict[ikey] = powerstate.plasma[mapper[ikey]][:, 1:] - if mapper[ikey] + "_stds" in powerstate.plasma: - var_dict[ikey + "_stds"] = powerstate.plasma[mapper[ikey] + "_stds"][:, 1:] - else: - var_dict[ikey + "_stds"] = None - - dfT = var_dict["Qe_tr_turb"] # as a reference for sizes - - # ------------------------------------------------------------------------- - # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added - # ------------------------------------------------------------------------- - - if PORTALSparameters["surrogateForTurbExch"]: - PexchTurb_integrated = PORTALStools.computeTurbExchangeIndividual( - var_dict["PexchTurb"], powerstate - ) - PexchTurb_integrated_stds = PORTALStools.computeTurbExchangeIndividual( - var_dict["PexchTurb_stds"], powerstate - ) - else: - PexchTurb_integrated = torch.zeros(dfT.shape).to(dfT) - PexchTurb_integrated_stds = torch.zeros(dfT.shape).to(dfT) - - # ------------------------------------------------------------------------ - # Go through each profile that needs to be predicted, calculate components - # ------------------------------------------------------------------------ - - of, cal = torch.Tensor().to(dfT), torch.Tensor().to(dfT) - ofE, calE = torch.Tensor().to(dfT), torch.Tensor().to(dfT) - for prof in powerstate.ProfilesPredicted: - if prof == "te": - var = "Qe" - elif prof == "ti": - var = "Qi" - elif prof == "ne": - var = "Ge" - elif prof == "nZ": - var = "GZ" - elif prof == "w0": - var = "MtJm2" - - """ - ----------------------------------------------------------------------------------- - Transport (_tr_turb+_tr_neo) - ----------------------------------------------------------------------------------- - """ - of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neo"] - of0E = (var_dict[f"{var}_tr_turb_stds"] ** 2 + var_dict[f"{var}_tr_neo_stds"] ** 2) ** 0.5 - - """ - ----------------------------------------------------------------------------------- - Target (Sum here the turbulent exchange power) - ----------------------------------------------------------------------------------- - """ - if var == "Qe": - cal0 = var_dict[f"{var}_tar"] + PexchTurb_integrated - cal0E = ( - var_dict[f"{var}_tar_stds"] ** 2 + PexchTurb_integrated_stds**2 - ) ** 0.5 - elif var == "Qi": - cal0 = var_dict[f"{var}_tar"] - PexchTurb_integrated - cal0E = ( - var_dict[f"{var}_tar_stds"] ** 2 + PexchTurb_integrated_stds**2 - ) ** 0.5 - else: - cal0 = var_dict[f"{var}_tar"] - cal0E = var_dict[f"{var}_tar_stds"] - - of, cal = torch.cat((of, of0), dim=-1), torch.cat((cal, cal0), dim=-1) - ofE, calE = torch.cat((ofE, of0E), dim=-1), torch.cat((calE, cal0E), dim=-1) - - return of, cal, ofE, calE From 4da41d0f05879fc304b6cb69d42020e388a744f2 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Apr 2025 09:02:10 -0700 Subject: [PATCH 030/385] Moved geometry tools to gacode --- src/mitim_tools/gacode_tools/PROFILEStools.py | 4 ++-- .../gacode_tools}/utils/GEOMETRYtools.py | 0 2 files changed, 2 insertions(+), 2 deletions(-) rename src/{mitim_modules/powertorch => mitim_tools/gacode_tools}/utils/GEOMETRYtools.py (100%) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index bc111811..76bbcd1c 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -5,10 +5,10 @@ import matplotlib.pyplot as plt from collections import OrderedDict from mitim_tools.misc_tools import GRAPHICStools, MATHtools, PLASMAtools, IOtools -from mitim_modules.powertorch.utils import GEOMETRYtools, CALCtools +from mitim_modules.powertorch.utils import CALCtools from mitim_tools.gs_tools import GEQtools from mitim_tools.gacode_tools import NEOtools -from mitim_tools.gacode_tools.utils import GACODEdefaults +from mitim_tools.gacode_tools.utils import GACODEdefaults, GEOMETRYtools from mitim_tools.transp_tools import CDFtools from mitim_tools.transp_tools.utils import TRANSPhelpers from mitim_tools.misc_tools.LOGtools import printMsg as print diff --git a/src/mitim_modules/powertorch/utils/GEOMETRYtools.py b/src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py similarity index 100% rename from src/mitim_modules/powertorch/utils/GEOMETRYtools.py rename to src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py From 450189d6a57f45dce8cb831debe4cd1c5838c5d3 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Apr 2025 09:38:35 -0700 Subject: [PATCH 031/385] Towards a tglf scan trick common to standalone TGLF (1) --- .../physics_models/transport_tgyro.py | 107 ++++++++---------- .../powertorch/utils/TRANSPORTtools.py | 3 +- 2 files changed, 52 insertions(+), 58 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index c17169ae..d6014854 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -197,16 +197,13 @@ def _profiles_to_store(self): print("\t- Could not move files", typeMsg="w") def tglf_scan_trick( - fluxesTGYRO, - tgyro, - label, + tglf, RadiisToRun, ProfilesPredicted, impurityPosition=1, includeFast=False, delta=0.02, cold_start=False, - check_coincidence_thr=1E-2, extra_name="", remove_folders_out = False, cores_per_tglf_instance = 4 # e.g. 4 core per radius, since this is going to launch ~ Nr=5 x (Nv=6 x Nd=2 + 1) = 65 TGLFs at once @@ -214,17 +211,7 @@ def tglf_scan_trick( print(f"\t- Running TGLF standalone scans ({delta = }) to determine relative errors") - # Grab fluxes from TGYRO - Qe_tgyro, Qi_tgyro, Ge_tgyro, GZ_tgyro, Mt_tgyro, Pexch_tgyro = fluxesTGYRO - - # ------------------------------------------------------------------------------------------------------------------------ - # TGLF scans - # ------------------------------------------------------------------------------------------------------------------------ - # Prepare scan - - tglf = tgyro.grab_tglf_objects(fromlabel=label, subfolder = 'tglf_explorations') - variables_to_scan = [] for i in ProfilesPredicted: if i == 'te': variables_to_scan.append('RLTS_1') @@ -295,28 +282,6 @@ def tglf_scan_trick( GZ[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Gi'] cont += jump - # ---------------------------------------------------- - # Do a check that TGLF scans are consistent with TGYRO - Qe_err = np.abs( (Qe[:,0] - Qe_tgyro) / Qe_tgyro ) if 'te' in ProfilesPredicted else np.zeros_like(Qe[:,0]) - Qi_err = np.abs( (Qi[:,0] - Qi_tgyro) / Qi_tgyro ) if 'ti' in ProfilesPredicted else np.zeros_like(Qi[:,0]) - Ge_err = np.abs( (Ge[:,0] - Ge_tgyro) / Ge_tgyro ) if 'ne' in ProfilesPredicted else np.zeros_like(Ge[:,0]) - GZ_err = np.abs( (GZ[:,0] - GZ_tgyro) / GZ_tgyro ) if 'nZ' in ProfilesPredicted else np.zeros_like(GZ[:,0]) - - F_err = np.concatenate((Qe_err, Qi_err, Ge_err, GZ_err)) - if F_err.max() > check_coincidence_thr: - print(f"\t- TGLF scans are not consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%",typeMsg="w") - if 'te' in ProfilesPredicted: - print('\t\t* Qe:',Qe_err) - if 'ti' in ProfilesPredicted: - print('\t\t* Qi:',Qi_err) - if 'ne' in ProfilesPredicted: - print('\t\t* Ge:',Ge_err) - if 'nZ' in ProfilesPredicted: - print('\t\t* GZ:',GZ_err) - else: - print(f"\t- TGLF scans are consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%") - # ---------------------------------------------------- - # Calculate the standard deviation of the scans, that's going to be the reported stds def calculate_mean_std(Q): @@ -338,14 +303,13 @@ def calculate_mean_std(Q): Ge_point, Ge_std = calculate_mean_std(Ge) GZ_point, GZ_std = calculate_mean_std(GZ) - #TODO: Implement Mt and Pexch - Mt_point, Pexch_point = Mt_tgyro, Pexch_tgyro - Mt_std, Pexch_std = abs(Mt_point) * 0.1, abs(Pexch_point) * 0.1 - #TODO: Careful with fast particles - return Qe_point, Qi_point, Ge_point, GZ_point, Mt_point, Pexch_point, Qe_std, Qi_std, Ge_std, GZ_std, Mt_std, Pexch_std + Flux_base = [Qe[:,0], Qi[:,0], Ge[:,0], GZ[:,0], None, None] #TODO (Mt, Pexch) + Flux_mean = [Qe_point, Qi_point, Ge_point, GZ_point, None, None] #TODO (Mt, Pexch) + Flux_std = [Qe_std, Qi_std, Ge_std, GZ_std, None, None] #TODO (Mt, Pexch) + return Flux_base, Flux_mean, Flux_std # ************************************************************************************************** # Functions # ************************************************************************************************** @@ -363,7 +327,8 @@ def curateTGYROfiles( use_tglf_scan_trick=None, cold_start=False, extra_name="", - cores_per_tglf_instance = 4 + cores_per_tglf_instance = 4, + check_coincidence_thr=1E-2, ): tgyro = tgyroObject.results[label] @@ -372,11 +337,7 @@ def curateTGYROfiles( relativeErrorNEO = percentError[1] / 100.0 relativeErrorTAR = percentError[2] / 100.0 - # ************************************************************************************************************************** - # TGLF - # ************************************************************************************************************************** - - # Grab fluxes + # Grab fluxes from TGYRO Qe = tgyro.Qe_sim_turb[0, 1:] Qi = tgyro.QiIons_sim_turb[0, 1:] if includeFast else tgyro.QiIons_sim_turb_thr[0, 1:] Ge = tgyro.Ge_sim_turb[0, 1:] @@ -387,17 +348,12 @@ def curateTGYROfiles( # Determine TGLF standard deviations if use_tglf_scan_trick is not None: - if provideTurbulentExchange: - print("> Turbulent exchange not implemented yet in TGLF scans", typeMsg="w") #TODO + # Grab TGLF object + tglfObject = tgyroObject.grab_tglf_objects(fromlabel=label, subfolder = 'tglf_explorations') - # -------------------------------------------------------------- - # If using the scan trick - # -------------------------------------------------------------- - - Qe, Qi, Ge, GZ, Mt, Pexch, QeE, QiE, GeE, GZE, MtE, PexchE = tglf_scan_trick( - [Qe, Qi, Ge, GZ, Mt, Pexch], - tgyroObject, - label, + # Run TGLF scan trick + Flux_base, Flux_mean, Flux_std = tglf_scan_trick( + tglfObject, RadiisToRun, ProfilesPredicted, impurityPosition=impurityPosition, @@ -408,6 +364,43 @@ def curateTGYROfiles( cores_per_tglf_instance=cores_per_tglf_instance ) + Qe, Qi, Ge, GZ, _, _ = Flux_mean + QeE, QiE, GeE, GZE, _, _ = Flux_std + + #TODO + MtE, PexchE = abs(Mt) * 0.1, abs(Pexch) * 0.1 + + # ---------------------------------------------------- + # Do a check that TGLF scans are consistent with TGYRO + + Qe_base, Qi_base, Ge_base, GZ_base, _, _ = Flux_base + + # Grab fluxes from TGYRO + Qe_tgyro = tgyro.Qe_sim_turb[0, 1:] + Qi_tgyro = tgyro.QiIons_sim_turb[0, 1:] if includeFast else tgyro.QiIons_sim_turb_thr[0, 1:] + Ge_tgyro = tgyro.Ge_sim_turb[0, 1:] + GZ_tgyro = tgyro.Gi_sim_turb[impurityPosition, 0, 1:] + + Qe_err = np.abs( (Qe_base - Qe_tgyro) / Qe_tgyro ) if 'te' in ProfilesPredicted else np.zeros_like(Qe_base) + Qi_err = np.abs( (Qi_base - Qi_tgyro) / Qi_tgyro ) if 'ti' in ProfilesPredicted else np.zeros_like(Qi_base) + Ge_err = np.abs( (Ge_base - Ge_tgyro) / Ge_tgyro ) if 'ne' in ProfilesPredicted else np.zeros_like(Ge_base) + GZ_err = np.abs( (GZ_base - GZ_tgyro) / GZ_tgyro ) if 'nZ' in ProfilesPredicted else np.zeros_like(GZ_base) + + F_err = np.concatenate((Qe_err, Qi_err, Ge_err, GZ_err)) + if F_err.max() > check_coincidence_thr: + print(f"\t- TGLF scans are not consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%",typeMsg="w") + if 'te' in ProfilesPredicted: + print('\t\t* Qe:',Qe_err) + if 'ti' in ProfilesPredicted: + print('\t\t* Qi:',Qi_err) + if 'ne' in ProfilesPredicted: + print('\t\t* Ge:',Ge_err) + if 'nZ' in ProfilesPredicted: + print('\t\t* GZ:',GZ_err) + else: + print(f"\t- TGLF scans are consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%") + # ---------------------------------------------------- + min_relative_error = 0.01 # To avoid problems with gpytorch, 1% error minimum QeE = QeE.clip(abs(Qe)*min_relative_error) diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index b1c8fc6c..c70d62e9 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -67,7 +67,8 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ } def produce_profiles(self): - # Only add self._produce_profiles() if it's needed (e.g. full TGLF), otherwise this is somewhat expensive (e.g. for flux matching) + # Only add self._produce_profiles() if it's needed (e.g. full TGLF), otherwise this is somewhat expensive + # (e.g. for flux matching of analytical models) pass def _produce_profiles(self,deriveQuantities=True): From 2cac8aaa88dad7fdd831980e9a107b9a59791c4c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Apr 2025 11:36:32 -0700 Subject: [PATCH 032/385] TGLF module now retrieves Momentum and Exchange --- src/mitim_tools/gacode_tools/PROFILEStools.py | 2 +- src/mitim_tools/gacode_tools/TGLFtools.py | 59 +++++++++++++------ src/mitim_tools/gacode_tools/TGYROtools.py | 4 +- .../gacode_tools/utils/NORMtools.py | 24 ++++---- 4 files changed, 55 insertions(+), 34 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 76bbcd1c..86a95e83 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -504,7 +504,7 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= 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"], self.derived["pi_gb"], self.derived["s_gb"], _ = PLASMAtools.gyrobohmUnits( self.profiles["te(keV)"], self.profiles["ne(10^19/m^3)"] * 1e-1, self.derived["mi_ref"], diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 4e63342f..76cc286a 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -2466,8 +2466,10 @@ def readScan( self.positionIon_scan = positionIon # ---- - x, Qe, Qi, Ge, Gi, ky, g, f, eta1, eta2, itg, tem, etg = [],[],[],[],[],[],[],[],[],[],[],[],[] - Qe_gb, Qi_gb, Ge_gb, Gi_gb = [], [], [], [] + x = [] + Qe, Qi, Ge, Gi, Mt, S = [],[],[],[],[],[] + Qe_gb, Qi_gb, Ge_gb, Gi_gb, Mt_gb, S_gb = [],[],[],[],[],[] + ky, g, f, eta1, eta2, itg, tem, etg = [],[],[],[],[],[],[],[] etalow_g, etalow_f, etalow_k = [], [], [] cont = 0 for ikey in self.results: @@ -2476,8 +2478,10 @@ def readScan( if isThisTheRightReadResults: self.scans[label]["results_tags"].append(ikey) - x0, Qe0, Qi0, Ge0, Gi0, ky0, g0, f0, eta10, eta20, itg0, tem0, etg0 = [],[],[],[],[],[],[],[],[],[],[],[],[] - Qe_gb0, Qi_gb0, Ge_gb0, Gi_gb0 = [], [], [], [] + x0 = [] + Qe0, Qi0, Ge0, Gi0, Mt0, S0 = [],[],[],[],[],[] + Qe_gb0, Qi_gb0, Ge_gb0, Gi_gb0, Mt_gb0, S_gb0 = [],[],[],[],[],[] + ky0, g0, f0, eta10, eta20, itg0, tem0, etg0 = [],[],[],[],[],[],[],[] etalow_g0, etalow_f0, etalow_k0 = [], [], [] for irho_cont in range(len(self.rhos)): irho = np.where(self.results[ikey]["x"] == self.rhos[irho_cont])[0][0] @@ -2488,6 +2492,8 @@ def readScan( Qi_gb0.append(self.results[ikey]["TGLFout"][irho].Qi) Ge_gb0.append(self.results[ikey]["TGLFout"][irho].Ge) Gi_gb0.append(self.results[ikey]["TGLFout"][irho].GiAll[self.positionIon_scan - 2]) + Mt_gb0.append(self.results[ikey]["TGLFout"][irho].Mt) + S_gb0.append(self.results[ikey]["TGLFout"][irho].Se) ky0.append(self.results[ikey]["TGLFout"][irho].ky) g0.append(self.results[ikey]["TGLFout"][irho].g) f0.append(self.results[ikey]["TGLFout"][irho].f) @@ -2505,6 +2511,8 @@ def readScan( Qi0.append(self.results[ikey]["TGLFout"][irho].Qi_unn) Ge0.append(self.results[ikey]["TGLFout"][irho].Ge_unn) Gi0.append(self.results[ikey]["TGLFout"][irho].GiAll_unn[self.positionIon_scan - 2]) + Mt0.append(self.results[ikey]["TGLFout"][irho].Mt_unn) + S0.append(self.results[ikey]["TGLFout"][irho].Se_unn) else: self.scans[label]["unnormalization_successful"] = False @@ -2517,6 +2525,8 @@ def readScan( Ge_gb.append(Ge_gb0) Gi_gb.append(Gi_gb0) Gi.append(Gi0) + Mt.append(Mt0) + S.append(S0) ky.append(ky0) g.append(g0) f.append(f0) @@ -2539,10 +2549,14 @@ def readScan( self.scans[label]["Qi_gb"] = np.atleast_2d(np.transpose(Qi_gb)) self.scans[label]["Ge_gb"] = np.atleast_2d(np.transpose(Ge_gb)) self.scans[label]["Gi_gb"] = np.atleast_2d(np.transpose(Gi_gb)) + self.scans[label]["Mt_gb"] = np.atleast_2d(np.transpose(Mt_gb)) + self.scans[label]["S_gb"] = np.atleast_2d(np.transpose(S_gb)) self.scans[label]["Qe"] = np.atleast_2d(np.transpose(Qe)) self.scans[label]["Qi"] = np.atleast_2d(np.transpose(Qi)) self.scans[label]["Ge"] = np.atleast_2d(np.transpose(Ge)) self.scans[label]["Gi"] = np.atleast_2d(np.transpose(Gi)) + self.scans[label]["Mt"] = np.atleast_2d(np.transpose(Mt)) + self.scans[label]["S"] = np.atleast_2d(np.transpose(S)) self.scans[label]["eta1"] = np.atleast_2d(np.transpose(eta1)) self.scans[label]["eta2"] = np.atleast_2d(np.transpose(eta2)) self.scans[label]["itg"] = np.atleast_2d(np.transpose(itg)) @@ -4491,13 +4505,9 @@ def __init__(self, FolderGACODE, suffix="",require_all_files=True): self.FolderGACODE, self.suffix = FolderGACODE, suffix if suffix == "": - print( - f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} without suffix" - ) + print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} without suffix") else: - print( - f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} with suffix {suffix}" - ) + print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} with suffix {suffix}") self.inputclass = TGLFinput(file=self.FolderGACODE / f"input.tglf{self.suffix}") self.roa = self.inputclass.geom["RMIN_LOC"] @@ -4552,17 +4562,29 @@ def read(self,require_all_files=True): self.Ge = data[0, 0] self.Qe = data[1, 0] + self.Me = data[2, 0] + self.Se = data[3, 0] self.GiAll = data[0, 1:] self.QiAll = data[1, 1:] + self.MiAll = data[2, 1:] + self.SiAll = data[3, 1:] - print( - f"\t\t- For Qi, summing contributions from ions {self.ions_included} (#0 is e-)", - typeMsg="i", - ) + print(f"\t\t- For Qi, summing contributions from ions {self.ions_included} (#0 is e-)",typeMsg="i",) self.Gi = data[0, self.ions_included].sum() self.Qi = data[1, self.ions_included].sum() + signMt = - self.inputclass.plasma['SIGN_IT'] # Following tgyro_flux.f90 + print(f"\t\t- Sign of Mt given by toroidal current direction (SIGN_IT={-signMt}): {signMt}",typeMsg="i",) + self.Me *= signMt + self.MiAll *= signMt + + print("\t\t- For Mt, summing all species contributions",typeMsg="i",) + self.Mt = self.Me + self.MiAll.sum() + + print("\t\t- For St, summing all ion species contributions",typeMsg="i",) + self.Si = self.SiAll.sum() + if require_all_files: # ------------------------------------------------------------------------ @@ -5009,14 +5031,14 @@ def read(self,require_all_files=True): lines = fi.readlines() self.inputFileTGLF = "".join(lines) - def unnormalize( - self, normalization, rho=None, convolution_fun_fluct=None, factorTot_to_Perp=1.0 - ): + def unnormalize(self, normalization, rho=None, convolution_fun_fluct=None, factorTot_to_Perp=1.0): if normalization is not None: rho_x = normalization["rho"] roa_x = normalization["roa"] q_gb = normalization["q_gb"] g_gb = normalization["g_gb"] + pi_gb = normalization["pi_gb"] + s_gb = normalization["s_gb"] rho_s = normalization["rho_s"] a = normalization["rmin"][-1] @@ -5036,6 +5058,9 @@ def unnormalize( self.Ge_unn = self.Ge * g_gb[ir] self.GiAll_unn = self.GiAll * g_gb[ir] + self.Mt_unn = self.Mt * pi_gb[ir] + self.Se_unn = self.Se * s_gb[ir] + self.AmplitudeSpectrum_Te_level = GACODErun.obtainFluctuationLevel( self.ky, self.AmplitudeSpectrum_Te, diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 79d3b4fe..5b83e895 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -1432,14 +1432,14 @@ def readFluxes(self): # Errors - Constructed outside of TGYRO call (e.g. powerstate) # *************************************************************** - if not (self.FolderTGYRO / f"out.tgyro.flux_e_stds").exists(): + if not (self.FolderTGYRO / "out.tgyro.flux_e_stds").exists(): self.tgyro_stds = False else: print("\t- Errors in TGYRO fluxes and targets found, adding to class") self.tgyro_stds = True - file = self.FolderTGYRO / f"out.tgyro.flux_e_stds" + file = self.FolderTGYRO / "out.tgyro.flux_e_stds" ( _, self.GeGB_sim_neo_stds, diff --git a/src/mitim_tools/gacode_tools/utils/NORMtools.py b/src/mitim_tools/gacode_tools/utils/NORMtools.py index 86ea444e..42bbae99 100644 --- a/src/mitim_tools/gacode_tools/utils/NORMtools.py +++ b/src/mitim_tools/gacode_tools/utils/NORMtools.py @@ -87,6 +87,8 @@ def normalizations_tgyro(tgyro, rho, roa): "rho": rho, "q_gb": np.interp(rho, x_tgyro, tgyro.Q_GB[iteration]), "g_gb": np.interp(rho, x_tgyro, tgyro.Gamma_GB[iteration]), + "pi_gb": np.interp(rho, x_tgyro, tgyro.Pi_GB[iteration]), + "s_gb": np.interp(rho, x_tgyro, tgyro.S_GB[iteration]), "c_s": np.interp(rho, x_tgyro, tgyro.c_s[iteration]), } @@ -94,6 +96,7 @@ def normalizations_tgyro(tgyro, rho, roa): def normalizations_profiles(profiles): + if profiles is not None: Set_norm = { "rho": profiles.profiles["rho(-)"], @@ -101,26 +104,18 @@ def normalizations_profiles(profiles): "rmin": np.abs(profiles.profiles["rmin(m)"]), "q_gb": np.abs(profiles.derived["q_gb"]), "g_gb": np.abs(profiles.derived["g_gb"]), - "exp_Qe": np.abs(profiles.derived["qe"]), - "exp_Qi": np.abs(profiles.derived["qi"]), - "exp_Ge": np.abs(profiles.derived["ge"]), + "pi_gb": np.abs(profiles.derived["pi_gb"]), + "s_gb": np.abs(profiles.derived["s_gb"]), "B_unit": np.abs(profiles.derived["B_unit"]), "rho_s": np.abs(profiles.derived["rho_s"]), "c_s": np.abs(profiles.derived["c_s"]), - "Te_keV": np.abs( - profiles.profiles[ - "te(keV)" if "te(keV)" in profiles.profiles else "Te(keV)" - ] - ), + "Te_keV": np.abs(profiles.profiles["te(keV)"]), "ne_20": np.abs(profiles.profiles["ne(10^19/m^3)"]) * 1e-1, "Ti_keV": np.abs(profiles.profiles["ti(keV)"][:, 0]), "ni_20": np.abs(profiles.derived["ni_thrAll"]) * 1e-1, - "exp_Qe": profiles.derived["qe_MWmiller"] - / profiles.derived["surfGACODE_miller"], # This is the same as qe_MWm2 - "exp_Qi": profiles.derived["qi_MWmiller"] - / profiles.derived["surfGACODE_miller"], - "exp_Ge": profiles.derived["ge_10E20miller"] - / profiles.derived["surfGACODE_miller"], + "exp_Qe": profiles.derived["qe_MWmiller"] / profiles.derived["surfGACODE_miller"], # This is the same as qe_MWm2 + "exp_Qi": profiles.derived["qi_MWmiller"] / profiles.derived["surfGACODE_miller"], + "exp_Ge": profiles.derived["ge_10E20miller"] / profiles.derived["surfGACODE_miller"], "mi_ref": profiles.derived["mi_ref"], } @@ -252,6 +247,7 @@ def plotNormalizations( colors=["b", "r", "g"], legYN=True, extralab="", + fn = None, ): if NormalizationSets is not None: if axs is None: From da8cbbecabc306c35fd5c5882203d3e398c7f482 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Apr 2025 11:36:54 -0700 Subject: [PATCH 033/385] Scan trick now correctly treats Momentum and Exchange --- .../physics_models/transport_tgyro.py | 41 ++++++++++++------- 1 file changed, 26 insertions(+), 15 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index d6014854..71ecaf4a 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -271,6 +271,8 @@ def tglf_scan_trick( Qi = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) Ge = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) GZ = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + Mt = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + S = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) cont = 0 for vari in variables_to_scan: @@ -280,6 +282,8 @@ def tglf_scan_trick( Qi[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qi'] Ge[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Ge'] GZ[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Gi'] + Mt[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Mt'] + S[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['S'] cont += jump # Calculate the standard deviation of the scans, that's going to be the reported stds @@ -302,12 +306,14 @@ def calculate_mean_std(Q): Qi_point, Qi_std = calculate_mean_std(Qi) Ge_point, Ge_std = calculate_mean_std(Ge) GZ_point, GZ_std = calculate_mean_std(GZ) + Mt_point, Mt_std = calculate_mean_std(Mt) + S_point, S_std = calculate_mean_std(S) #TODO: Careful with fast particles - Flux_base = [Qe[:,0], Qi[:,0], Ge[:,0], GZ[:,0], None, None] #TODO (Mt, Pexch) - Flux_mean = [Qe_point, Qi_point, Ge_point, GZ_point, None, None] #TODO (Mt, Pexch) - Flux_std = [Qe_std, Qi_std, Ge_std, GZ_std, None, None] #TODO (Mt, Pexch) + Flux_base = [Qe[:,0], Qi[:,0], Ge[:,0], GZ[:,0], Mt[:,0], S[:,0]] + Flux_mean = [Qe_point, Qi_point, Ge_point, GZ_point, Mt_point, S_point] + Flux_std = [Qe_std, Qi_std, Ge_std, GZ_std, Mt_std, S_std] return Flux_base, Flux_mean, Flux_std # ************************************************************************************************** @@ -364,39 +370,44 @@ def curateTGYROfiles( cores_per_tglf_instance=cores_per_tglf_instance ) - Qe, Qi, Ge, GZ, _, _ = Flux_mean - QeE, QiE, GeE, GZE, _, _ = Flux_std - - #TODO - MtE, PexchE = abs(Mt) * 0.1, abs(Pexch) * 0.1 + Qe, Qi, Ge, GZ, Mt, Pexch = Flux_mean + QeE, QiE, GeE, GZE, MtE, PexchE = Flux_std # ---------------------------------------------------- # Do a check that TGLF scans are consistent with TGYRO - Qe_base, Qi_base, Ge_base, GZ_base, _, _ = Flux_base + Qe_base, Qi_base, Ge_base, GZ_base, Mt_base, S_base = Flux_base # Grab fluxes from TGYRO Qe_tgyro = tgyro.Qe_sim_turb[0, 1:] Qi_tgyro = tgyro.QiIons_sim_turb[0, 1:] if includeFast else tgyro.QiIons_sim_turb_thr[0, 1:] Ge_tgyro = tgyro.Ge_sim_turb[0, 1:] GZ_tgyro = tgyro.Gi_sim_turb[impurityPosition, 0, 1:] + Mt_tgyro = tgyro.Mt_sim_turb[0, 1:] + Pexch_tgyro = tgyro.EXe_sim_turb[0, 1:] Qe_err = np.abs( (Qe_base - Qe_tgyro) / Qe_tgyro ) if 'te' in ProfilesPredicted else np.zeros_like(Qe_base) Qi_err = np.abs( (Qi_base - Qi_tgyro) / Qi_tgyro ) if 'ti' in ProfilesPredicted else np.zeros_like(Qi_base) Ge_err = np.abs( (Ge_base - Ge_tgyro) / Ge_tgyro ) if 'ne' in ProfilesPredicted else np.zeros_like(Ge_base) GZ_err = np.abs( (GZ_base - GZ_tgyro) / GZ_tgyro ) if 'nZ' in ProfilesPredicted else np.zeros_like(GZ_base) + Mt_err = np.abs( (Mt_base - Mt_tgyro) / Mt_tgyro ) if 'w0' in ProfilesPredicted else np.zeros_like(Mt_base) + Pexch_err = np.abs( (Pexch - Pexch_tgyro) / Pexch_tgyro ) if provideTurbulentExchange else np.zeros_like(Pexch) - F_err = np.concatenate((Qe_err, Qi_err, Ge_err, GZ_err)) + F_err = np.concatenate((Qe_err, Qi_err, Ge_err, GZ_err, Mt_err, Pexch_err)) if F_err.max() > check_coincidence_thr: - print(f"\t- TGLF scans are not consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%",typeMsg="w") - if 'te' in ProfilesPredicted: + print(f"\t- TGLF scans are not consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%, in quantity:",typeMsg="w") + if ('te' in ProfilesPredicted) and Qe_err.max() > check_coincidence_thr: print('\t\t* Qe:',Qe_err) - if 'ti' in ProfilesPredicted: + if ('ti' in ProfilesPredicted) and Qi_err.max() > check_coincidence_thr: print('\t\t* Qi:',Qi_err) - if 'ne' in ProfilesPredicted: + if ('ne' in ProfilesPredicted) and Ge_err.max() > check_coincidence_thr: print('\t\t* Ge:',Ge_err) - if 'nZ' in ProfilesPredicted: + if ('nZ' in ProfilesPredicted) and GZ_err.max() > check_coincidence_thr: print('\t\t* GZ:',GZ_err) + if ('w0' in ProfilesPredicted) and Mt_err.max() > check_coincidence_thr: + print('\t\t* Mt:',Mt_err) + if provideTurbulentExchange and Pexch_err.max() > check_coincidence_thr: + print('\t\t* Pexch:',Pexch_err) else: print(f"\t- TGLF scans are consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%") # ---------------------------------------------------- From 693ecb9c56e6f531769b0d97f73070c28fa08d06 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 28 Apr 2025 16:15:55 -0400 Subject: [PATCH 034/385] misc --- src/mitim_modules/powertorch/utils/TRANSPORTtools.py | 10 ++-------- src/mitim_tools/opt_tools/scripts/slurm.py | 2 +- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index c70d62e9..7ed58698 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -129,14 +129,8 @@ def _modify_profiles(self): def evaluate(self): ''' This needs to populate the following in self.powerstate.plasma - - Pe, Pe_tr, Pe_tr_turb, Pe_tr_neo -> MW/m^2 - - Pi, Pi_tr, Pi_tr_turb, Pi_tr_neo -> MW/m^2 - - Ce, Ce_tr, Ce_tr_turb, Ce_tr_neo -> MW/m^2 - * Ge1E20sm2, Ge1E20sm2_tr, Ge1E20sm2_tr_turb, Ge1E20sm2_tr_neo -> 10^20/s/m^2 - - CZ, CZ_tr, CZ_tr_turb, CZ_tr_neo -> MW/m^2 (but modified as needed, for example dividing by fZ0) - * CZ_raw, CZ_raw_tr, CZ_raw_tr_turb, CZ_raw_tr_neo -> 10^20/s/m^2 (NOT modified) - - Mt, Mt_tr, Mt_tr_turb, Mt_tr_neo -> J/m^2 - - PexchTurb -> MW/m^3 + - QeMWm2, QeMWm2_tr, QeMWm2_tr_turb, QeMWm2_tr_neo + Same for QiMWm2, Ce, CZ, MtJm2 and their respective standard deviations ''' diff --git a/src/mitim_tools/opt_tools/scripts/slurm.py b/src/mitim_tools/opt_tools/scripts/slurm.py index 4eacb7a7..392ff0bd 100644 --- a/src/mitim_tools/opt_tools/scripts/slurm.py +++ b/src/mitim_tools/opt_tools/scripts/slurm.py @@ -39,7 +39,7 @@ def run_slurm( command = [venv,script + (f" --seed {seed}" if seed is not None else "")] - nameJob = f"mitim_opt_{folder.name}{extra_name}" + nameJob = f"mitim_{folder.name}{extra_name}" _, fileSBATCH, _ = FARMINGtools.create_slurm_execution_files( command, From 46d9bc2c544348b9978d669adbd245bd296117bd Mon Sep 17 00:00:00 2001 From: audreysa Date: Tue, 29 Apr 2025 17:01:19 -0400 Subject: [PATCH 035/385] Fix print statement so single and double quotes are separate --- src/mitim_modules/maestro/utils/PORTALSbeat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 2b2c9f5b..bcc9bb34 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -56,7 +56,7 @@ def prepare(self, new_name = data_df['Ion'][iZ] - print(f"\t\t\t- Changing name of ion from {profiles.Species[i]["N"]} ({profiles.Species[i]["Z"]}) to {new_name} ({Z[iZ]})") + print(f'\t\t\t- Changing name of ion from {profiles.Species[i]["N"]} ({profiles.Species[i]["Z"]}) to {new_name} ({Z[iZ]})') profiles.profiles['name'][i] = profiles.Species[i]["N"] = new_name From 1bdad25c49ae3aa36a11769291a9ee6a5371b9e0 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 29 Apr 2025 17:05:57 -0400 Subject: [PATCH 036/385] Fixed case when no radiation compontents available --- src/mitim_tools/gacode_tools/PROFILEStools.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 86a95e83..6d74f122 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -87,6 +87,9 @@ def process(self, mi_ref=None, calculateDerived=True): "ptot(Pa)", # e.g. if I haven't written that info from ASTRA "zeta(-)", # e.g. if TGYRO is run with zeta=0, it won't write this column in .new "zmag(m)", + "qsync(MW/m^3)", + "qbrem(MW/m^3)", + "qline(MW/m^3)", self.varqpar, self.varqpar2, "shape_cos0(-)", From e18a76876c3d4d3b54ef9520c8104eb341139d32 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 29 Apr 2025 18:06:30 -0400 Subject: [PATCH 037/385] Minor bug of r/a legend in powertorch convergence --- src/mitim_modules/powertorch/utils/POWERplot.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index 9d17a272..436c3c6c 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -123,7 +123,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co position_in_batch = i * ( self.plasma['rho'].shape[-1] -1 ) + j - ax.plot(self.FluxMatch_Xopt[:,position_in_batch], "-o", color=colors[j], lw=1.0, label = f"r/a = {self.plasma['roa'][batch_num,j]:.2f}",markersize=0.5) + ax.plot(self.FluxMatch_Xopt[:,position_in_batch], "-o", color=colors[j], lw=1.0, label = f"r/a = {self.plasma['roa'][batch_num,j+1]:.2f}",markersize=0.5) if self.bounds_current is not None: for u in [0,1]: ax.axhline(y=self.bounds_current[u,position_in_batch], color=colors[j], linestyle='-.', lw=0.2) From 2998526f6e2fbe11b6846382071a29b4f2940eab Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 30 Apr 2025 09:35:04 -0400 Subject: [PATCH 038/385] Printing more info on gacode radiation components --- src/mitim_tools/gacode_tools/PROFILEStools.py | 94 ++++--------------- 1 file changed, 16 insertions(+), 78 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 6d74f122..2fe38315 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -1283,8 +1283,6 @@ def deriveContentByVolumes(self, rhos=[0.5], impurityPosition=3): def printInfo(self, label="", reDeriveIfNotFound=True): - if 'pfast_fraction' not in self.derived: self.derived['pfast_fraction'] = np.nan #TODO: remove this line - try: ImpurityText = "" for i in range(len(self.Species)): @@ -1298,90 +1296,30 @@ def printInfo(self, label="", reDeriveIfNotFound=True): print(f"\tkappa_sep = {self.profiles['kappa(-)'][-1]:.2f}, kappa_995 = {self.derived['kappa995']:.2f}, kappa_95 = {self.derived['kappa95']:.2f}, kappa_a = {self.derived['kappa_a']:.2f}") print(f"\tdelta_sep = {self.profiles['delta(-)'][-1]:.2f}, delta_995 = {self.derived['delta995']:.2f}, delta_95 = {self.derived['delta95']:.2f}") 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( - "\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_ne0.2 = {0:.2f} (nu_eff w/Zeff2 = {1:.2f})".format( - self.derived["ne_peaking0.2"], self.derived["nu_eff2"] - ) - ) + 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("\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_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}") print(f"\tp_vol = {self.derived['ptot_manual_vol']:.2f} MPa ({self.derived['pfast_fraction']*100.0:.1f}% fast)") - print( - f"\tBetaN = {self.derived['BetaN']:.3f} (BetaN w/B0 = {self.derived['BetaN_engineering']:.3f})" - ) - print( - "\tPrad = {0:.1f}MW ({1:.1f}% of total)".format( - self.derived["Prad"], - self.derived["Prad"] / self.derived["qHeat"] * 100.0, - ) - ) - print( - "\tPsol = {0:.1f}MW (fLH = {1:.2f})".format( - self.derived["Psol"], self.derived["LHratio"] - ) - ) - print( - "Operational point ( [,] = [{0:.2f},{1:.2f}] ) and species:".format( - self.derived["ne_vol20"], self.derived["Te_vol"] - ) - ) - print( - "\t = {0:.2f} keV (/ = {1:.2f}, Ti0/Te0 = {2:.2f})".format( - self.derived["Ti_vol"], - self.derived["tite_vol"], - self.derived["tite"][0], - ) - ) - print( - "\tfG = {0:.2f} ( = {1:.2f} * 10^20 m^-3)".format( - self.derived["fG"], self.derived["ne_vol20"] - ) - ) - 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']:.1e}]" - ) + print(f"\tBetaN = {self.derived['BetaN']:.3f} (BetaN w/B0 = {self.derived['BetaN_engineering']:.3f})") + print(f"\tPrad = {self.derived['Prad']:.1f}MW ({self.derived['Prad'] / self.derived['qHeat'] * 100.0:.1f}% of total) ({self.derived['Prad_brem']/self.derived['Prad'] * 100.0:.1f}% brem, {self.derived['Prad_line']/self.derived['Prad'] * 100.0:.1f}% line, {self.derived['Prad_sync']/self.derived['Prad'] * 100.0:.1f}% sync)") + print("\tPsol = {0:.1f}MW (fLH = {1:.2f})".format(self.derived["Psol"], self.derived["LHratio"])) + print("Operational point ( [,] = [{0:.2f},{1:.2f}] ) and species:".format(self.derived["ne_vol20"], self.derived["Te_vol"])) + print("\t = {0:.2f} keV (/ = {1:.2f}, Ti0/Te0 = {2:.2f})".format(self.derived["Ti_vol"],self.derived["tite_vol"],self.derived["tite"][0],)) + print("\tfG = {0:.2f} ( = {1:.2f} * 10^20 m^-3)".format(self.derived["fG"], self.derived["ne_vol20"])) + 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']:.1e}]") print(f"\tMach = {self.derived['MachNum_vol']:.2f} (vol avg)") 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( - f"\ttauE = { self.derived['tauE']:.3f} s, tauP = {self.derived['tauP']:.3f} s (tauP/tauE = {self.derived['tauPotauE']:.2f})" - ) + 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(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", - ) + 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) From b7e39d6e73a48042dd484b99d4166a83804e405d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 30 Apr 2025 09:35:43 -0400 Subject: [PATCH 039/385] Added capabiility to only print info from gacode, not plot --- src/mitim_tools/gacode_tools/scripts/read_gacode.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/mitim_tools/gacode_tools/scripts/read_gacode.py b/src/mitim_tools/gacode_tools/scripts/read_gacode.py index 9c994607..14b10374 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_gacode.py +++ b/src/mitim_tools/gacode_tools/scripts/read_gacode.py @@ -1,4 +1,5 @@ import argparse +from turtle import st from mitim_tools.gacode_tools import PROFILEStools """ @@ -11,13 +12,13 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("files", type=str, nargs="*") - parser.add_argument( - "--rho", type=float, required=False, default=0.89 - ) # Last rho for gradients plot + parser.add_argument("--rho", type=float, required=False, default=0.89) # Last rho for gradients plot + parser.add_argument("--print", required=False, default=False, action="store_true") # Last rho for gradients plot args = parser.parse_args() files = args.files rho = args.rho + print_only = args.print # Read profs = [] @@ -29,9 +30,11 @@ def main(): # Plot - fn = PROFILEStools.plotAll(profs, lastRhoGradients=rho) + if not print_only: - fn.show() + fn = PROFILEStools.plotAll(profs, lastRhoGradients=rho) + + fn.show() # Import IPython and embed an interactive session from IPython import embed From 82907cf2164faa0cd161fe506704f348350b024f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 30 Apr 2025 11:27:28 -0400 Subject: [PATCH 040/385] Store profiles_transport in powerstate, not the transport subclass, for better tracking --- .../powertorch/physics_models/transport_tgyro.py | 2 +- src/mitim_modules/powertorch/utils/TRANSPORTtools.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 71ecaf4a..ec332524 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -48,7 +48,7 @@ def _evaluate_tglf_neo(self): RadiisToRun = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] tgyro = TGYROtools.TGYRO(cdf=dummyCDF(self.folder, self.folder)) - tgyro.prep(self.folder, profilesclass_custom=self.profiles_transport) + tgyro.prep(self.folder, profilesclass_custom=self.powerstate.profiles_transport) if launchMODELviaSlurm: print("\t- Launching TGYRO evaluation as a batch job") diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 7ed58698..0242ce4f 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -87,7 +87,7 @@ def _produce_profiles(self,deriveQuantities=True): insert_highres_powers = deriveQuantities, # Insert powers so that Q, Pfus and all that it's consistent when read later ) - self.profiles_transport = copy.deepcopy(self.powerstate.profiles) + self.powerstate.profiles_transport = copy.deepcopy(self.powerstate.profiles) self._modify_profiles() @@ -104,7 +104,7 @@ def _modify_profiles(self): if profiles_postprocessing_fun is not None: print(f"\t- Modifying input.gacode to run transport calculations based on {profiles_postprocessing_fun}",typeMsg="i") - self.profiles_transport = profiles_postprocessing_fun(self.file_profs) + self.powerstate.profiles_transport = profiles_postprocessing_fun(self.file_profs) # Position of impurity ion may have changed p_old = PROFILEStools.PROFILES_GACODE(self.file_profs_unmod) From f38fffd3b42b0e54367579eed041ec7ba3831d9d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 30 Apr 2025 16:01:51 -0400 Subject: [PATCH 041/385] Added minimum absolute change in gradients during tglf scan trick --- src/mitim_modules/portals/PORTALStools.py | 3 +- .../physics_models/transport_tgyro.py | 8 ++++ .../powertorch/utils/TRANSPORTtools.py | 2 +- src/mitim_tools/gacode_tools/TGLFtools.py | 12 ++++++ .../gacode_tools/utils/GACODErun.py | 37 +++++++++++-------- 5 files changed, 43 insertions(+), 19 deletions(-) diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 1b46f72c..e97ed9f1 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -17,11 +17,10 @@ def surrogate_selection_portals(output, surrogate_options, CGYROrun=False): if output[3:6] == "tar": surrogate_options["TypeMean"] = 1 surrogate_options["TypeKernel"] = 2 # Constant kernel - # If it's not, stndard + # If it's not, standard case for fluxes else: surrogate_options["TypeMean"] = 2 # Linear in gradients, constant in rest surrogate_options["TypeKernel"] = 1 # RBF - # surrogate_options['ExtraNoise'] = True surrogate_options["additional_constraints"] = { 'lenghtscale_constraint': gpytorch.constraints.constraints.GreaterThan(0.01) # inputs normalized to [0,1], this is 1% lengthscale diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index ec332524..57ffb452 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -203,6 +203,7 @@ def tglf_scan_trick( impurityPosition=1, includeFast=False, delta=0.02, + minimum_abs_gradient=0.005, # This is 0.5% of aLx=1.0, to avoid extremely small scans when, for example, having aLn ~ 0.0 cold_start=False, extra_name="", remove_folders_out = False, @@ -230,6 +231,12 @@ def tglf_scan_trick( relative_scan = [1-delta, 1+delta] + # Enforce at least "minimum_abs_gradient" in gradient, to avoid zero gradient situations + minimum_delta_abs = {} + for ikey in variables_to_scan: + if 'RL' in ikey: + minimum_delta_abs[ikey] = minimum_abs_gradient + name = 'turb_drives' tglf.rhos = RadiisToRun # To avoid the case in which TGYRO was run with an extra rho point @@ -248,6 +255,7 @@ def tglf_scan_trick( subFolderTGLF = name, variablesDrives = variables_to_scan, varUpDown = relative_scan, + minimum_delta_abs = minimum_delta_abs, TGLFsettings = None, ApplyCorrections = False, add_baseline_to = 'first', diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 0242ce4f..02f5c395 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -120,7 +120,7 @@ def _modify_profiles(self): impurityPosition_new = self.powerstate.impurityPosition if impurityPosition_new != self.powerstate.impurityPosition: - print(f"\t- Impurity position has changed from {self.powerstate.impurityPosition} to {impurityPosition_new}",typeMsg="w") + print(f"\t- Impurity position has changed from {self.powerstate.impurityPosition} to {impurityPosition_new}",typeMsg="i") self.powerstate.impurityPosition_transport = p_new.Species.index(impurity_of_interest) # ---------------------------------------------------------------------------------------------------- diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 76cc286a..79712a3d 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -511,6 +511,7 @@ def run( TGLFsettings=None, extraOptions={}, multipliers={}, + minimum_delta_abs={}, runWaveForms=None, # e.g. runWaveForms = [0.3,1.0] forceClosestUnstableWF=True, # Look at the growth rate spectrum and run exactly the ky of the closest unstable ApplyCorrections=True, # Removing ions with too low density and that are fast species @@ -541,6 +542,7 @@ def run( TGLFsettings=TGLFsettings, extraOptions=extraOptions, multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, runWaveForms=runWaveForms, forceClosestUnstableWF=forceClosestUnstableWF, ApplyCorrections=ApplyCorrections, @@ -629,6 +631,7 @@ def _prepare_run_radii( TGLFsettings=None, extraOptions={}, multipliers={}, + minimum_delta_abs={}, ApplyCorrections=True, # Removing ions with too low density and that are fast species Quasineutral=False, # Ensures quasineutrality. By default is False because I may want to run the file directly launchSlurm=True, @@ -689,6 +692,7 @@ def _prepare_run_radii( TGLFsettings=TGLFsettings, extraOptions=extraOptions, multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, ApplyCorrections=ApplyCorrections, Quasineutral=Quasineutral, ) @@ -2326,6 +2330,7 @@ def runScan( self, subFolderTGLF, # 'scan1', multipliers={}, + minimum_delta_abs={}, variable="RLTS_1", varUpDown=[0.5, 1.0, 1.5], variables_scanTogether=[], @@ -2352,6 +2357,7 @@ def runScan( tglf_executor, tglf_executor_full, folders, varUpDown_new = self._prepare_scan( subFolderTGLF, multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, variable=variable, varUpDown=varUpDown_new, variables_scanTogether=variables_scanTogether, @@ -2380,6 +2386,7 @@ def _prepare_scan( self, subFolderTGLF, # 'scan1', multipliers={}, + minimum_delta_abs={}, variable="RLTS_1", varUpDown=[0.5, 1.0, 1.5], variables_scanTogether=[], @@ -2443,6 +2450,7 @@ def _prepare_scan( tglf_executor=tglf_executor, tglf_executor_full=tglf_executor_full, multipliers=multipliers_mod, + minimum_delta_abs=minimum_delta_abs, **kwargs_TGLFrun, ) @@ -3141,6 +3149,7 @@ def runScanTurbulenceDrives( variation=0.5, add_baseline_to = 'none', # 'all' or 'first' or 'none' variablesDrives=["RLTS_1", "RLTS_2", "RLNS_1", "XNUE", "TAUS_2"], + minimum_delta_abs={}, positionIon=2, **kwargs_TGLFrun, ): @@ -3172,6 +3181,7 @@ def runScanTurbulenceDrives( scan_name, variable=variable, varUpDown=varUpDown_dict[variable], + minimum_delta_abs=minimum_delta_abs, **kwargs_TGLFrun, ) @@ -3798,6 +3808,7 @@ def changeANDwrite_TGLF( TGLFsettings=None, extraOptions={}, multipliers={}, + minimum_delta_abs={}, ApplyCorrections=True, Quasineutral=False, ): @@ -3818,6 +3829,7 @@ def changeANDwrite_TGLF( Settings=TGLFsettings, extraOptions=extraOptions, multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, position_change=i, addControlFunction=GACODEdefaults.addTGLFcontrol, NS=NS, diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 74e7f0d0..8d1d7e16 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -99,6 +99,7 @@ def modifyInputs( Settings=None, extraOptions={}, multipliers={}, + minimum_delta_abs={}, position_change=0, addControlFunction=None, **kwargs_to_function, @@ -160,18 +161,12 @@ def modifyInputs( input_class.plasma[ikey] = var_new else: # If the variable in extraOptions wasn't in there, consider it a control param - print( - "\t\t- Variable to change did not exist previously, creating now", - typeMsg="i", - ) + print("\t\t- Variable to change did not exist previously, creating now",typeMsg="i") var_orig = None var_new = value_to_change_to input_class.controls[ikey] = var_new - print( - f"\t\t- Changing {ikey} from {var_orig} to {var_new}", - typeMsg="i", - ) + print(f"\t\t- Changing {ikey} from {var_orig} to {var_new}",typeMsg="i",) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Change with multipliers -> Input directly, not as multiplier @@ -184,31 +179,41 @@ def modifyInputs( specie = int(ikey.split("_")[-1]) varK = "_".join(ikey.split("_")[:-1]) var_orig = input_class.species[specie][varK] - var_new = var_orig * multipliers[ikey] + var_new = multiplier_tglf_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.species[specie][varK] = var_new else: if ikey in input_class.controls: var_orig = input_class.controls[ikey] - var_new = var_orig * multipliers[ikey] + var_new = multiplier_tglf_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.controls[ikey] = var_new + elif ikey in input_class.geom: var_orig = input_class.geom[ikey] - var_new = var_orig * multipliers[ikey] + var_new = multiplier_tglf_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.geom[ikey] = var_new + elif ikey in input_class.plasma: var_orig = input_class.plasma[ikey] - var_new = var_orig * multipliers[ikey] + var_new = multiplier_tglf_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.plasma[ikey] = var_new + else: - print( - "\t- Variable to scan did not exist in original file, add it as extraOptions first", - typeMsg="w", - ) + print("\t- Variable to scan did not exist in original file, add it as extraOptions first",typeMsg="w",) print(f"\t\t\t- Changing {ikey} from {var_orig} to {var_new} (x{multipliers[ikey]})") return input_class +def multiplier_tglf_input(var_orig, multiplier, minimum_delta_abs = None): + + delta = var_orig * (multiplier - 1.0) + + if minimum_delta_abs is not None: + if (multiplier != 1.0) and abs(delta) < minimum_delta_abs: + print(f"\t\t\t- delta = {delta} is smaller than minimum_delta_abs = {minimum_delta_abs}, enforcing",typeMsg="i") + delta = np.sign(delta) * minimum_delta_abs + + return var_orig + delta def findNamelist(LocationCDF, folderWork=None, nameRunid="10000", ForceFirst=True): # ----------------------------------------------------------- From 2f029dd141c14528fec9860cc8641b602544ac71 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 30 Apr 2025 18:53:17 -0400 Subject: [PATCH 042/385] misc --- .../opt_tools/utils/EVALUATORtools.py | 2 +- .../opt_tools/utils/SBOcorrections.py | 103 ++++-------------- 2 files changed, 24 insertions(+), 81 deletions(-) diff --git a/src/mitim_tools/opt_tools/utils/EVALUATORtools.py b/src/mitim_tools/opt_tools/utils/EVALUATORtools.py index 2a3bde11..399b4ab0 100644 --- a/src/mitim_tools/opt_tools/utils/EVALUATORtools.py +++ b/src/mitim_tools/opt_tools/utils/EVALUATORtools.py @@ -132,7 +132,7 @@ def mitimRun( if pd.Series(y).isna().any() or pd.Series(yE).isna().any(): print( f"--> Reading Tabular file failed or not evaluated yet for element {numEval}", - typeMsg="w", + typeMsg="i", ) cold_start = True else: diff --git a/src/mitim_tools/opt_tools/utils/SBOcorrections.py b/src/mitim_tools/opt_tools/utils/SBOcorrections.py index c97f4e03..0acd91b1 100644 --- a/src/mitim_tools/opt_tools/utils/SBOcorrections.py +++ b/src/mitim_tools/opt_tools/utils/SBOcorrections.py @@ -520,24 +520,12 @@ def updateMetrics(self, evaluatedPoints=1, IsThisAFreshIteration=True, position= Yvar = torch.from_numpy(self.train_Ystd).to(self.dfT) ** 2 if "steps" in self.__dict__: - self.BOmetrics["overall"]["Residual"] = ( - -self.steps[-1].evaluators["objective"](Y).detach().cpu().numpy() - ) - self.BOmetrics["overall"]["ResidualModeledLast"] = ( - -self.steps[-1].evaluators["residual_function"](X).detach().cpu().numpy() - ) - + self.BOmetrics["overall"]["Residual"] = -self.steps[-1].evaluators["objective"](Y).detach().cpu().numpy() + self.BOmetrics["overall"]["ResidualModeledLast"] = -self.steps[-1].evaluators["residual_function"](X).detach().cpu().numpy() else: - print( - "\t~ Cannot perform prediction at this iteration step, returning objective evaluation as acquisition", - typeMsg="w", - ) - self.BOmetrics["overall"]["Residual"] = ( - -self.scalarized_objective(Y)[2].detach().cpu().numpy() - ) - self.BOmetrics["overall"]["ResidualModeledLast"] = self.BOmetrics["overall"][ - "Residual" - ] + print("\t~ Cannot perform prediction at this iteration step, returning objective evaluation as acquisition",typeMsg="i",) + self.BOmetrics["overall"]["Residual"] = -self.scalarized_objective(Y)[2].detach().cpu().numpy() + self.BOmetrics["overall"]["ResidualModeledLast"] = self.BOmetrics["overall"]["Residual"] resi = self.BOmetrics["overall"]["Residual"] resiM = self.BOmetrics["overall"]["ResidualModeledLast"] @@ -547,50 +535,29 @@ def updateMetrics(self, evaluatedPoints=1, IsThisAFreshIteration=True, position= self.BOmetrics["overall"]["indBest"] = np.nanargmin(resi, axis=0) self.BOmetrics["overall"]["indBestModel"] = np.nanargmin(resiM, axis=0) - self.BOmetrics["overall"]["xBest"] = ( - X[self.BOmetrics["overall"]["indBest"], :].detach().cpu() - ) - self.BOmetrics["overall"]["yBest"] = ( - Y[self.BOmetrics["overall"]["indBest"], :].detach().cpu() - ) - self.BOmetrics["overall"]["yVarBest"] = ( - Yvar[self.BOmetrics["overall"]["indBest"], :].detach().cpu() - ) + self.BOmetrics["overall"]["xBest"] = X[self.BOmetrics["overall"]["indBest"], :].detach().cpu() + self.BOmetrics["overall"]["yBest"] = Y[self.BOmetrics["overall"]["indBest"], :].detach().cpu() + self.BOmetrics["overall"]["yVarBest"] = Yvar[self.BOmetrics["overall"]["indBest"], :].detach().cpu() # Best from last iteration zA = np.nanmin(resi[-evaluatedPoints:], axis=0) - self.BOmetrics["overall"]["indBestLast"] = np.nanargmin( - resi[-evaluatedPoints:], axis=0 - ) + self.BOmetrics["overall"]["indBestLast"] = np.nanargmin(resi[-evaluatedPoints:], axis=0) zM = np.nanmin(resiM[-evaluatedPoints:], axis=0) - self.BOmetrics["overall"]["indBestModelLast"] = np.nanargmin( - resiM[-evaluatedPoints:], axis=0 - ) + self.BOmetrics["overall"]["indBestModelLast"] = np.nanargmin(resiM[-evaluatedPoints:], axis=0) - self.BOmetrics["overall"]["xBestLast"] = X[ - self.BOmetrics["overall"]["indBestLast"], : - ] - self.BOmetrics["overall"]["yBestLast"] = Y[ - self.BOmetrics["overall"]["indBestLast"], : - ] + self.BOmetrics["overall"]["xBestLast"] = X[self.BOmetrics["overall"]["indBestLast"], :] + self.BOmetrics["overall"]["yBestLast"] = Y[self.BOmetrics["overall"]["indBestLast"], :] # Metric tracking of previous iterations - if ( - len(self.BOmetrics["overall"]["ResidualModeledLast"]) - == self.Originalinitial_training - ): + if len(self.BOmetrics["overall"]["ResidualModeledLast"]) == self.Originalinitial_training: ratio, metric = np.inf, 0.0 label = "\t(Initial batch only)" else: zA_prev = np.nanmin(resi[:-evaluatedPoints], axis=0) - self.BOmetrics["overall"]["indBestExceptLast"] = np.nanargmin( - resi[:-evaluatedPoints], axis=0 - ) + self.BOmetrics["overall"]["indBestExceptLast"] = np.nanargmin(resi[:-evaluatedPoints], axis=0) zM_prev = np.nanmin(resiM[:-evaluatedPoints], axis=0) - self.BOmetrics["overall"]["indBestModelExceptLast"] = np.nanargmin( - resiM[:-evaluatedPoints], axis=0 - ) + self.BOmetrics["overall"]["indBestModelExceptLast"] = np.nanargmin(resiM[:-evaluatedPoints], axis=0) ratio, metric, label = constructMetricsTR(zA, zM, zA_prev, zM_prev) @@ -609,9 +576,7 @@ def updateMetrics(self, evaluatedPoints=1, IsThisAFreshIteration=True, position= elif IsThisAFreshIteration: print(f"- BO fitness ratio = {ratio:.3f}, metric = {metric} ({label})") else: - print( - "- Note that this one was not fresh (e.g. post-correction), do not track metrics" - ) + print("- Note that this one was not fresh (e.g. post-correction), do not track metrics") print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") # ------------------------------------------------------------------------------------------------------------ @@ -630,13 +595,9 @@ def updateMetrics(self, evaluatedPoints=1, IsThisAFreshIteration=True, position= self.BOmetrics["xBest_track"][position] = self.BOmetrics["overall"]["xBest"] self.BOmetrics["yBest_track"][position] = self.BOmetrics["overall"]["yBest"] - self.BOmetrics["yVarBest_track"][position] = self.BOmetrics["overall"][ - "yVarBest" - ] + self.BOmetrics["yVarBest_track"][position] = self.BOmetrics["overall"]["yVarBest"] - self.BOmetrics["BOmetric_it"][position] = ( - X.shape[0] - 1 - ) # Evaluation position until now + self.BOmetrics["BOmetric_it"][position] = X.shape[0] - 1 # Evaluation position until now def constructMetricsTR(zA, zM, zA_prev, zM_prev): @@ -660,9 +621,7 @@ def constructMetricsTR(zA, zM, zA_prev, zM_prev): """ - label = "Evaluated previous = {2:.2e} --> Predicted new = {1:.2e} --> Evaluated new = {0:.2e}; ".format( - zA, zM, zA_prev, zM_prev - ) + label = f"Evaluated previous = {zA_prev:.2e} --> Predicted new = {zM:.2e} --> Evaluated new = {zA:.2e}; " # Relative improvements (Positive if it has gotten better = lower residual) zA = (zA_prev - zA) / np.abs(zA_prev) @@ -753,26 +712,10 @@ def plotTrustRegionInformation(self, fig=None): ratio = ratio[:lim] - metrics = np.array( - [self.BOmetrics["BOmetric"][i] for i in self.BOmetrics["BOmetric_it"]] - )[ - :lim - ] # self.BOmetrics['BOmetric']#[1:] - boundsX1 = np.array( - [self.BOmetrics["BoundsStorage"][i] for i in self.BOmetrics["BOmetric_it"]] - )[ - :lim - ] # self.BOmetrics['BoundsStorage'] - operat = np.array( - [self.BOmetrics["TRoperation"][i] for i in self.BOmetrics["BOmetric_it"]] - )[ - :lim - ] # self.BOmetrics['TRoperation'] #[1:] - evaluations = np.array( - [self.BOmetrics["BOmetric_it"][i] for i in self.BOmetrics["BOmetric_it"]] - )[ - :lim - ] # [1:] + metrics = np.array([self.BOmetrics["BOmetric"][i] for i in self.BOmetrics["BOmetric_it"]])[:lim] # self.BOmetrics['BOmetric']#[1:] + boundsX1 = np.array([self.BOmetrics["BoundsStorage"][i] for i in self.BOmetrics["BOmetric_it"]])[:lim] # self.BOmetrics['BoundsStorage'] + operat = np.array([self.BOmetrics["TRoperation"][i] for i in self.BOmetrics["BOmetric_it"]])[:lim] # self.BOmetrics['TRoperation'] #[1:] + evaluations = np.array([self.BOmetrics["BOmetric_it"][i] for i in self.BOmetrics["BOmetric_it"]])[:lim] # [1:] size, center = [], [] va = list(boundsX1[0].keys())[0] From 0327c44af272d829b45f8f85427dbfa38030d053 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 21 May 2025 09:15:38 -0400 Subject: [PATCH 043/385] misc --- src/mitim_tools/gs_tools/GEQtools.py | 4 ++-- src/mitim_tools/opt_tools/utils/EVALUATORtools.py | 15 +++------------ 2 files changed, 5 insertions(+), 14 deletions(-) diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 2e173604..896c334e 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -953,9 +953,9 @@ def check(self, warning_error = 0.01, plotYN = False): max_error = np.max([max_error, error]) if max_error > warning_error: - print(f"\t\t- Maximum error is {100*max_error:.2f}%", typeMsg='w') + print(f"\t\t- Maximum error in equilibrium quantities is {100*max_error:.2f}%", typeMsg='w') else: - print(f"\t\t- Maximum error is {100*max_error:.2f}%", typeMsg='i') + print(f"\t\t- Maximum error in equilibrium quantities is {100*max_error:.2f}%") # -------------------------------------------------------------- # Plotting diff --git a/src/mitim_tools/opt_tools/utils/EVALUATORtools.py b/src/mitim_tools/opt_tools/utils/EVALUATORtools.py index 399b4ab0..15781cd0 100644 --- a/src/mitim_tools/opt_tools/utils/EVALUATORtools.py +++ b/src/mitim_tools/opt_tools/utils/EVALUATORtools.py @@ -130,15 +130,10 @@ def mitimRun( y, yE, _ = optimization_data.grab_data_point(x) if pd.Series(y).isna().any() or pd.Series(yE).isna().any(): - print( - f"--> Reading Tabular file failed or not evaluated yet for element {numEval}", - typeMsg="i", - ) + print(f"--> Reading Tabular file failed or not evaluated yet for element {numEval}",typeMsg="i",) cold_start = True else: - print( - f"--> Reading Tabular file successful for element {numEval}", - ) + print(f"--> Reading Tabular file successful for element {numEval}",) if cold_start: # Create folder @@ -161,11 +156,7 @@ def mitimRun( try: y, yE = IOtools.readresults(resultsfile) except: - print( - "Could not read results file for {0}, printing error file to screen".format( - numEval - ) - ) + print(f"Could not read results file for {numEval}, printing error file to screen") # print(error) y, yE = np.array([np.nan]), np.array([np.nan]) From 5d5f0bc7763e5b6204a1a65f258df366065a7dd4 Mon Sep 17 00:00:00 2001 From: audreysa Date: Thu, 22 May 2025 11:07:30 -0400 Subject: [PATCH 044/385] Added runs_slurm_array function to slurm.py --- src/mitim_tools/opt_tools/scripts/slurm.py | 73 ++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/src/mitim_tools/opt_tools/scripts/slurm.py b/src/mitim_tools/opt_tools/scripts/slurm.py index 392ff0bd..d6503141 100644 --- a/src/mitim_tools/opt_tools/scripts/slurm.py +++ b/src/mitim_tools/opt_tools/scripts/slurm.py @@ -1,5 +1,6 @@ import os from mitim_tools.misc_tools import FARMINGtools, IOtools +from IPython import embed """ This script is used to launch a slurm job with a scpecific script like... python3 run_case.py 0 --R 6.0 @@ -65,3 +66,75 @@ def run_slurm( input_files=[fileSBATCH], job_name = nameJob, ) + + +def run_slurm_array( + scripts, + folder, + partition, + venv, + max_concurrent_jobs, + seeds=None, # If not None, assume that the script is able to receive --seeds # + hours=8, + n=32, + seed_specific=0, + machine="local", + exclude=None, + mem=None, +): + + folder = IOtools.expandPath(folder) + + if seeds is not None: + seeds_explore = [seed_specific] if seeds == 1 else list(range(seeds)) + else: + seeds_explore = [None] + + for seed in seeds_explore: + + extra_name = "" if (seed is None or seeds == 1) else f"_s{seed}" + + folder = IOtools.expandPath(folder) + folder = folder.with_name(folder.name + extra_name) + + print(f"* Launching slurm job of MITIM optimization with random seed = {seed}") + + folder.mkdir(parents=True, exist_ok=True) + + scan_size = len(scripts)cd + commands = [venv,scripts[0] + (f" --seed {seed}" if seed is not None else "")] + for script in scripts[1:]: + commands.append(script) + + + print('Commands' + str(commands)) + print('folder' + str(folder)) + + nameJob = f"mitim_{folder.name}{extra_name}" + + _, fileSBATCH, _ = FARMINGtools.create_slurm_execution_files( + command=commands, + folder_remote=folder, + folder_local=folder, + nameJob=nameJob, + slurm={"partition": partition, 'exclude': exclude}, + minutes=int(60 * hours), + ntasks=1, + cpuspertask=n, + memory_req_by_job=mem, + job_array=f'0-{scan_size}%{max_concurrent_jobs}', + + ) + + command_execution = f"sbatch {fileSBATCH}" + + if machine == "local": + os.system(command_execution) + else: + FARMINGtools.perform_quick_remote_execution( + folder, + machine, + command_execution, + input_files=[fileSBATCH], + job_name = nameJob, + ) From 84acd8b9989e5e9ce4a3912754bb7335a0e853d1 Mon Sep 17 00:00:00 2001 From: audreysa Date: Thu, 22 May 2025 14:22:48 -0400 Subject: [PATCH 045/385] Fixed run_slurm_array to actually launch as array, added tutorial for run_slurm_array --- src/mitim_tools/opt_tools/scripts/slurm.py | 19 +++++-------- .../run_slurm_array_tutorial/test_launcher.py | 27 +++++++++++++++++++ .../run_slurm_array_tutorial/test_script.py | 13 +++++++++ 3 files changed, 47 insertions(+), 12 deletions(-) create mode 100644 tutorials/run_slurm_array_tutorial/test_launcher.py create mode 100644 tutorials/run_slurm_array_tutorial/test_script.py diff --git a/src/mitim_tools/opt_tools/scripts/slurm.py b/src/mitim_tools/opt_tools/scripts/slurm.py index d6503141..c24dc425 100644 --- a/src/mitim_tools/opt_tools/scripts/slurm.py +++ b/src/mitim_tools/opt_tools/scripts/slurm.py @@ -69,11 +69,12 @@ def run_slurm( def run_slurm_array( - scripts, + script, + array_input, folder, partition, - venv, max_concurrent_jobs, + venv = '', seeds=None, # If not None, assume that the script is able to receive --seeds # hours=8, n=32, @@ -101,19 +102,13 @@ def run_slurm_array( folder.mkdir(parents=True, exist_ok=True) - scan_size = len(scripts)cd - commands = [venv,scripts[0] + (f" --seed {seed}" if seed is not None else "")] - for script in scripts[1:]: - commands.append(script) - - - print('Commands' + str(commands)) - print('folder' + str(folder)) + command = ['echo $SLURM_ARRAY_TASK_ID', venv, script + ' $SLURM_ARRAY_TASK_ID'+ (f" --seed {seed}" if seed is not None else "")] + string_of_array_input = ','.join([str(i) for i in array_input]) nameJob = f"mitim_{folder.name}{extra_name}" _, fileSBATCH, _ = FARMINGtools.create_slurm_execution_files( - command=commands, + command=command, folder_remote=folder, folder_local=folder, nameJob=nameJob, @@ -122,7 +117,7 @@ def run_slurm_array( ntasks=1, cpuspertask=n, memory_req_by_job=mem, - job_array=f'0-{scan_size}%{max_concurrent_jobs}', + job_array=f'{string_of_array_input}%{max_concurrent_jobs}', ) diff --git a/tutorials/run_slurm_array_tutorial/test_launcher.py b/tutorials/run_slurm_array_tutorial/test_launcher.py new file mode 100644 index 00000000..229afbf4 --- /dev/null +++ b/tutorials/run_slurm_array_tutorial/test_launcher.py @@ -0,0 +1,27 @@ +import json +import os +from mitim_tools.opt_tools.scripts.slurm import run_slurm_array +from mitim_tools.opt_tools.scripts.slurm import run_slurm +from mitim_tools import __mitimroot__ +from mitim_tools.misc_tools.CONFIGread import load_settings + +# You have to have a slurm partition specified for your local machine for this to work!! +partition = load_settings()['local']['slurm']['partition'] +print(f"Using partition: {partition}") + + +# Settings for slurm job +cpus = 2 +hours = 1 +memory = '100GB' +folder = __mitimroot__ / "tutorials" / "run_slurm_array_tutorial" / "scratch" +script = f'python {folder}/../test_script.py {folder} ' + +# Input the array runs +array_input = [62, 63, 81] + +# To use run_slurm_array +run_slurm_array(script, array_input, folder,partition, max_concurrent_jobs = 2, hours=hours,n=cpus,mem=memory) + +# For comparison run_slurm +# run_slurm(f'python test_script.py {folder} 84', '/home/audreysa/test_script/scratch_run_slurm', partition, environment, hours=hours, n=cpus, mem=memory) diff --git a/tutorials/run_slurm_array_tutorial/test_script.py b/tutorials/run_slurm_array_tutorial/test_script.py new file mode 100644 index 00000000..4f347cad --- /dev/null +++ b/tutorials/run_slurm_array_tutorial/test_script.py @@ -0,0 +1,13 @@ +import sys +import os + +folder = str(sys.argv[1]) +# print('Folder:', folder) + + +i = str(sys.argv[2]) +# print('i:', i) + +with open(f"{folder}/file_successfully_created_{i}.txt", "w") as f: + input_text = 'Successfully created file from job #:' + i + f.write(f"{input_text}\n") \ No newline at end of file From da8829e39efbd5547a28f310c8b41d2424206794 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 23 May 2025 09:19:28 -0400 Subject: [PATCH 046/385] If Bt and Ip are specified in namelist (not null), then they will be used instead of the geqdsk values --- .../maestro/utils/MAESTRObeat.py | 21 +++++++++++++++++++ src/mitim_tools/gs_tools/GEQtools.py | 1 - 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/utils/MAESTRObeat.py b/src/mitim_modules/maestro/utils/MAESTRObeat.py index b4ac6733..05024a7d 100644 --- a/src/mitim_modules/maestro/utils/MAESTRObeat.py +++ b/src/mitim_modules/maestro/utils/MAESTRObeat.py @@ -177,6 +177,27 @@ def __call__( print(f'\t- Converting geqdsk to profiles, using {coeffs_MXH = }') p = self.f.to_profiles(ne0_20 = netop_20, Zeff = Zeff, PichT = PichT_MW, coeffs_MXH = coeffs_MXH) + # Sometimes I may want to change Ip and Bt + if 'Ip_MA' in kwargs_profiles and kwargs_profiles['Ip_MA'] is not None: + Ip_in_geqdsk = p.profiles['current(MA)'][0] + if Ip_in_geqdsk != kwargs_profiles['Ip_MA']: + print(f'\t- Requested to ignore geqdsk current and use user-specified one, changing Ip from {Ip_in_geqdsk} to {kwargs_profiles["Ip_MA"]}', typeMsg = 'w') + p.profiles['current(MA)'][0] = kwargs_profiles['Ip_MA'] + print(f'\t\t* Scaling poloidal flux by same factor as Ip, {kwargs_profiles["Ip_MA"] / Ip_in_geqdsk:.2f}') + p.profiles['polflux(Wb/radian)'] *= kwargs_profiles['Ip_MA'] / Ip_in_geqdsk + print(f'\t\t* Scaling q-profile by same factor as Ip, {kwargs_profiles["Ip_MA"] / Ip_in_geqdsk:.2f}') + p.profiles['q(-)'] *= 1/(kwargs_profiles['Ip_MA'] / Ip_in_geqdsk) + + if 'B_T' in kwargs_profiles and kwargs_profiles['B_T'] is not None: + Bt_in_geqdsk = p.profiles['bcentr(T)'][0] + if Bt_in_geqdsk != kwargs_profiles['B_T']: + print(f'\t- Requested to ignore geqdsk B and use user-specified one, changing Bt from {Bt_in_geqdsk} to {kwargs_profiles["B_T"]}', typeMsg = 'w') + p.profiles['bcentr(T)'][0] = kwargs_profiles['B_T'] + print(f'\t\t* Scaling toroidal flux by same factor as Bt, {kwargs_profiles["B_T"] / Bt_in_geqdsk:.2f}') + p.profiles['torfluxa(Wb/radian)'] *= kwargs_profiles['B_T'] / Bt_in_geqdsk + print(f'\t\t* Scaling q-profile by same factor as Bt, {kwargs_profiles["B_T"] / Bt_in_geqdsk:.2f}') + p.profiles['q(-)'] *= kwargs_profiles['B_T'] / Bt_in_geqdsk + # Write it to initialization folder p.writeCurrentStatus(file=self.folder / 'input.geqdsk.gacode') diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 896c334e..3014a397 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -364,7 +364,6 @@ def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH profiles['ze'] = np.array([-1.0]) profiles['z'] = np.array([1.0, Z]) - profiles['torfluxa(Wb/radian)'] = np.array([torfluxa]) profiles['rcentr(m)'] = np.array([R0]) profiles['bcentr(T)'] = np.array([B0]) From b33b503867be5c5637c0fb01e56b006f0c23da46 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 23 May 2025 17:55:10 -0400 Subject: [PATCH 047/385] misc CGYRO --- src/mitim_tools/gacode_tools/CGYROtools.py | 15 ++++----------- .../gacode_tools/scripts/read_cgyro.py | 2 +- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index e1f0b7e8..7612be8e 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -199,11 +199,9 @@ def read(self, label="cgyro1", folder=None): folder = IOtools.expandPath(folder) if folder is not None else self.folderCGYRO try: - self.results[label] = cgyrodata_plot(folder) + self.results[label] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") except: - if ( - True - ): # print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): + if print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): os.chdir(folder) os.system("cgyro -t") self.results[label] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") @@ -211,14 +209,9 @@ def read(self, label="cgyro1", folder=None): # Extra postprocessing self.results[label].electron_flag = np.where(self.results[label].z == -1)[0][0] self.results[label].all_flags = np.arange(0, len(self.results[label].z), 1) - self.results[label].ions_flags = self.results[label].all_flags[ - self.results[label].all_flags != self.results[label].electron_flag - ] + self.results[label].ions_flags = self.results[label].all_flags[self.results[label].all_flags != self.results[label].electron_flag] - self.results[label].all_names = [ - f"{gacodefuncs.specmap(self.results[label].mass[i],self.results[label].z[i])}({self.results[label].z[i]},{self.results[label].mass[i]:.1f})" - for i in self.results[label].all_flags - ] + self.results[label].all_names = [f"{gacodefuncs.specmap(self.results[label].mass[i],self.results[label].z[i])}({self.results[label].z[i]},{self.results[label].mass[i]:.1f})" for i in self.results[label].all_flags] def plotLS(self, labels=["cgyro1"], fig=None): colors = GRAPHICStools.listColors() diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index facde86b..6f5dd9c7 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -4,7 +4,7 @@ from mitim_tools.gacode_tools import CGYROtools """ -e.g. plot_cgyro.py folder +e.g. read_cgyro.py folder """ def main(): From cd08304035599b88ac070d263cbf39b65fc5e0e3 Mon Sep 17 00:00:00 2001 From: audreysa Date: Fri, 23 May 2025 15:57:27 -0400 Subject: [PATCH 048/385] Add check for whether slurm_output exists in check_maestro.py --- src/mitim_modules/maestro/scripts/check_maestro.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/scripts/check_maestro.py b/src/mitim_modules/maestro/scripts/check_maestro.py index 089a4388..37d8c4eb 100644 --- a/src/mitim_modules/maestro/scripts/check_maestro.py +++ b/src/mitim_modules/maestro/scripts/check_maestro.py @@ -75,7 +75,7 @@ def check_cases(folders): minutes = (time_in_queue.seconds % 3600) // 60 job_status = f"{state.strip()} for {hours}h {minutes}m ({cores} cores on {partition})" - if (outputs_folder / 'beat_final').exists(): + if (outputs_folder / 'beat_final').exists() and slurm_output.exists(): mod_time = datetime.fromtimestamp((outputs_folder / 'beat_final').stat().st_mtime).strftime('%Y-%m-%d %H:%M:%S') with open(slurm_output, 'r') as f: lines = f.readlines() From cc615048d84a3da0723f2ae6ee2d9d42f7e01cb9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 28 May 2025 09:45:57 -0400 Subject: [PATCH 049/385] Improved CGYRO plotting --- src/mitim_tools/gacode_tools/CGYROtools.py | 5 +++++ src/mitim_tools/gacode_tools/scripts/read_cgyro.py | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 7612be8e..60c98d12 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -198,7 +198,10 @@ def read(self, label="cgyro1", folder=None): folder = IOtools.expandPath(folder) if folder is not None else self.folderCGYRO + original_dir = os.getcwd() + try: + print(f"\t- Reading CGYRO data from {folder.resolve()}") self.results[label] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") except: if print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): @@ -206,6 +209,8 @@ def read(self, label="cgyro1", folder=None): os.system("cgyro -t") self.results[label] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") + os.chdir(original_dir) + # Extra postprocessing self.results[label].electron_flag = np.where(self.results[label].z == -1)[0][0] self.results[label].all_flags = np.arange(0, len(self.results[label].z), 1) diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 6f5dd9c7..810e087e 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -20,7 +20,7 @@ def main(): labels = [] for i, folder in enumerate(folders): - labels.append(f"{IOtools.reducePathLevel(folder)[-1]}") + labels.append(f"case {i + 1}") c.read(label=labels[-1], folder=folder) c.plot(labels=labels) From 70e92070fb04db8397fe8ef3220ad637af913dc9 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Wed, 28 May 2025 11:28:03 -0400 Subject: [PATCH 050/385] 2pi factor in astra-gacode tool --- src/mitim_tools/astra_tools/ASTRAtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/astra_tools/ASTRAtools.py b/src/mitim_tools/astra_tools/ASTRAtools.py index b5aacac0..f377689e 100644 --- a/src/mitim_tools/astra_tools/ASTRAtools.py +++ b/src/mitim_tools/astra_tools/ASTRAtools.py @@ -263,7 +263,7 @@ def convert_ASTRA_to_gacode_from_transp_output(c, bcentr = np.array([c.BTOR[ai]]) ; params['bcentr(T)'] = bcentr current = np.array([c.IPL[ai]]) ; params['current(MA)'] = current rho = interp_to_nexp(c.rho[ai]/c.rho[ai,-1]); params['rho(-)'] = rho - polflux = interp_to_nexp(c.FP[ai]) ; params['polflux(Wb/radian)'] = polflux + polflux = interp_to_nexp(c.FP[ai]) ; params['polflux(Wb/radian)'] = polflux/2/np.pi polflux_norm = (polflux-polflux[0])/(polflux[-1]-polflux[0]) From de6d68ab98794ffc12b7b34f4db6e0b3a83ecdb3 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Thu, 29 May 2025 14:34:03 -0400 Subject: [PATCH 051/385] passthrough option for initializing ASTRA without eped --- src/mitim_tools/astra_tools/ASTRAtools.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/mitim_tools/astra_tools/ASTRAtools.py b/src/mitim_tools/astra_tools/ASTRAtools.py index f377689e..e59b0d44 100644 --- a/src/mitim_tools/astra_tools/ASTRAtools.py +++ b/src/mitim_tools/astra_tools/ASTRAtools.py @@ -544,6 +544,8 @@ def _betan_initial_conditions(x, geometry_object, rhotop, Ttop_keV, netop_19, ep T[BC_index:] = T_ped[BC_index:] print(f"Pedestal values: ne_ped = {ne_ped}, Te_ped = {Te_ped}") + else: + x=rho preamble_Temp = f""" 900052D3D 2 0 6 ;-SHOT #- F(X) DATA WRITEUF OMFIT ;-SHOT DATE- UFILES ASCII FILE SYSTEM From 2bbd2fcfcbe97e11a15c4131e4cbbd589b477dff Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 29 May 2025 17:16:00 -0400 Subject: [PATCH 052/385] Modifications to CGYRO data interpretation and plotting --- src/mitim_tools/gacode_tools/CGYROtools.py | 394 ++++++------------ .../gacode_tools/utils/CGYROutils.py | 367 ++++++++++++++++ 2 files changed, 496 insertions(+), 265 deletions(-) create mode 100644 src/mitim_tools/gacode_tools/utils/CGYROutils.py diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 60c98d12..e90dcb08 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -5,7 +5,7 @@ from pathlib import Path import numpy as np import matplotlib.pyplot as plt -from mitim_tools.gacode_tools.utils import GACODEdefaults, GACODErun +from mitim_tools.gacode_tools.utils import GACODEdefaults, GACODErun, CGYROutils from mitim_tools.misc_tools import IOtools, GRAPHICStools, FARMINGtools from mitim_tools.gacode_tools.utils import GACODEplotting from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -186,15 +186,13 @@ def get(self): self.cgyro_job.retrieve() self.cgyro_job.close() else: - print( - "- Not retrieving results because this was run command line (not slurm)" - ) + print("- Not retrieving results because this was run command line (not slurm)") # --------------------------------------------------------------------------------------------------------- # Reading and plotting # --------------------------------------------------------------------------------------------------------- - def read(self, label="cgyro1", folder=None): + def read(self, label="cgyro1", folder=None, tmin = 0.0): folder = IOtools.expandPath(folder) if folder is not None else self.folderCGYRO @@ -218,6 +216,132 @@ def read(self, label="cgyro1", folder=None): self.results[label].all_names = [f"{gacodefuncs.specmap(self.results[label].mass[i],self.results[label].z[i])}({self.results[label].z[i]},{self.results[label].mass[i]:.1f})" for i in self.results[label].all_flags] + # ************************ + # Calculations + # ************************ + + cgyro = self.results[label] + cgyro.getflux() + cgyro.getnorm("elec") + + self.results[label].t = self.results[label].tnorm + + ys = np.sum(cgyro.ky_flux, axis=(2, 3)) + + self.results[label].Qe = ys[-1, 1, :] / cgyro.qc + self.results[label].Qi_all = ys[:-1, 1, :] / cgyro.qc + self.results[label].Qi = self.results[label].Qi_all.sum(axis=0) + self.results[label].Ge = ys[-1, 0, :] + + self.results[label].Qe_mean, self.results[label].Qe_std = self.derive_statistics(self.results[label].t, self.results[label].Qe, x_min=tmin) + self.results[label].Qi_mean, self.results[label].Qi_std = self.derive_statistics(self.results[label].t, self.results[label].Qi, x_min=tmin) + self.results[label].Ge_mean, self.results[label].Ge_std = self.derive_statistics(self.results[label].t, self.results[label].Ge, x_min=tmin) + + roa,alne,self.results[label].aLTi,alte,self.results[label].Qi_mean,self.results[label].Qi_std,self.results[label].Qe_mean,self.results[label].Qe_std,self.results[label].Ge_mean, self.results[label].Ge_std,m_gimp,std_gimp,m_mo,std_mo,m_tur,std_tur,qgb,ggb,pgb,sgb,tstart,nt = CGYROutils.grab_cgyro_nth(str(folder.resolve()), tmin, False, False) + + + def derive_statistics(self,x,y,x_min=0.0): + + return y.mean(), y.std() + + def plot(self, labels=[""]): + from mitim_tools.misc_tools.GUItools import FigureNotebook + + self.fn = FigureNotebook("CGYRO Notebook", geometry="1600x1000") + + colors = GRAPHICStools.listColors() + + fig = self.fn.add_figure(label="Fluxes Time Traces") + axsFluxes_t = fig.subplot_mosaic( + """ + AC + BD + """ + ) + + for j in range(len(labels)): + self.plot_fluxes( + axs=axsFluxes_t, + label=labels[j], + c=colors[j], + plotLegend=j == len(labels) - 1, + ) + + def plot_fluxes(self, axs=None, label="", c="b", lw=1, plotLegend=True): + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + AB + CD + """ + ) + + ls = GRAPHICStools.listLS() + + # Electron energy flux + ax = axs["A"] + ax.plot( + self.results[label].t, + self.results[label].Qe, + ls=ls[0], + lw=lw, + c=c, + label=f"{label}, electron", + ) + ax.set_xlabel("$t$ ($a/c_s$)") + ax.set_ylabel("$Q_e$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron energy flux') + + # Ion energy fluxes + ax = axs["B"] + ax.plot( + self.results[label].t, + self.results[label].Qi, + ls=ls[0], + lw=lw, + c=c, + label=f"{label}", + ) + ax.set_ylabel("$Q_i$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion energy fluxes') + + ax = axs["C"] + for j, i in enumerate(self.results[label].ions_flags): + ax.plot( + self.results[label].t, + self.results[label].Qi_all[j], + ls=ls[j + 1], + lw=lw / 2, + c=c, + label=f"{label}, {self.results[label].all_names[i]}", + ) + ax.set_ylabel("$Q_i$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion energy fluxes (separate species)') + GRAPHICStools.addLegendApart(ax,ratio=0.95, withleg=True, size = 8) + + + # Electron particle flux + ax = axs["D"] + ax.plot( + self.results[label].t, + self.results[label].Ge, + ls=ls[0], + lw=lw, + c=c, + label=f"{label}, electron", + ) + ax.set_ylabel("$\\Gamma_e$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron particle flux') + + + def plotLS(self, labels=["cgyro1"], fig=None): colors = GRAPHICStools.listColors() @@ -384,266 +508,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") - 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 - ys = np.sum(cgyro.ky_flux, axis=(2, 3)) - if moment == "n": - y = ys[ispec, 0, :] - mtag = r"\Gamma" - elif moment == "e": - y = ys[ispec, 1, :] / cgyro.qc - mtag = r"Q" - elif moment == "v": - y = ys[ispec, 2, :] - mtag = r"\Pi" - elif moment == "s": - y = ys[ispec, 3, :] - mtag = r"S" - - name = gacodefuncs.specmap(cgyro.mass[ispec], cgyro.z[ispec]) - - if retrieveSpecieFlag: - flag = f"${mtag}_{{{name}}}$ (GB)" - else: - flag = f"${mtag}$ (GB)" - - return t, y, flag - - def plot_flux( - self, - ax=None, - label="", - moment="e", - ispecs=[0], - labelPlot="", - c="b", - lw=1, - tmax=None, - dense=True, - ls="-", - ): - if ax is None: - plt.ion() - fig, ax = plt.subplots() - - for i, ispec in enumerate(ispecs): - t, y0, yl = self.get_flux( - label=label, - moment=moment, - ispec=ispec, - retrieveSpecieFlag=len(ispecs) == 1, - ) - - if i == 0: - y = y0 - else: - y += y0 - - if tmax is not None: - it = np.argmin(np.abs(t - tmax)) - t = t[: it + 1] - y = y[: it + 1] - - ax.plot(t, y, ls=ls, lw=lw, c=c, label=labelPlot) - - ax.set_xlabel("$t$ ($a/c_s$)") - # ax.set_xlim(left=0) - ax.set_ylabel(yl) - - if dense: - GRAPHICStools.addDenseAxis(ax) - - def plot_fluxes(self, axs=None, label="", c="b", lw=1, plotLegend=True): - if axs is None: - plt.ion() - fig = plt.figure(figsize=(18, 9)) - - axs = fig.subplot_mosaic( - """ - ABC - DEF - """ - ) - - ls = GRAPHICStools.listLS() - - # Electron - ax = axs["A"] - self.plot_flux( - ax=ax, - label=label, - moment="e", - ispecs=[self.results[label].electron_flag], - labelPlot=label, - c=c, - lw=lw, - ) - ax.set_title("Electron energy flux") - - ax = axs["B"] - self.plot_flux( - ax=ax, - label=label, - moment="e", - ispecs=self.results[label].ions_flags, - labelPlot=f"{label}, sum", - c=c, - lw=lw, - ls=ls[0], - ) - for j, i in enumerate(self.results[label].ions_flags): - self.plot_flux( - ax=ax, - label=label, - moment="e", - ispecs=[i], - labelPlot=f"{label}, {self.results[label].all_names[i]}", - c=c, - lw=lw / 2, - ls=ls[j + 1], - ) - ax.set_title(f"Ion energy fluxes") - - # Ion - ax = axs["D"] - self.plot_flux( - ax=ax, - label=label, - moment="n", - ispecs=[self.results[label].electron_flag], - labelPlot=label, - c=c, - lw=lw, - ) - ax.set_title("Electron particle flux") - - ax = axs["E"] - for j, i in enumerate(self.results[label].ions_flags): - self.plot_flux( - ax=ax, - label=label, - moment="n", - ispecs=[i], - labelPlot=f"{label}, {self.results[label].all_names[i]}", - c=c, - lw=lw / 2, - ls=ls[j + 1], - ) - self.plot_flux( - ax=ax, - label=label, - moment="n", - ispecs=self.results[label].ions_flags, - labelPlot=f"{label}, sum", - c=c, - lw=lw, - ls=ls[0], - ) - ax.set_title("Ion particle fluxes") - - # Extra - ax = axs["C"] - for j, i in enumerate(self.results[label].all_flags): - self.plot_flux( - ax=ax, - label=label, - moment="v", - ispecs=[i], - labelPlot=f"{label}, {self.results[label].all_names[i]}", - c=c, - lw=lw / 2, - ls=ls[j + 1], - ) - self.plot_flux( - ax=ax, - label=label, - moment="v", - ispecs=self.results[label].all_flags, - labelPlot=f"{label}, sum", - c=c, - lw=lw, - ls=ls[0], - ) - ax.set_title("Momentum flux") - - ax = axs["F"] - try: - for j, i in enumerate(self.results[label].all_flags): - self.plot_flux( - ax=ax, - label=label, - moment="s", - ispecs=[self.results[label].electron_flag], - labelPlot=f"{label}, {self.results[label].all_names[i]}", - c=c, - lw=lw, - ) - worked = True - except: - print("Could not plot energy exchange", typeMsg="w") - worked = False - ax.set_title("Electron energy exchange") - - plt.subplots_adjust() - if plotLegend: - for n in ["B", "E", "C"]: - GRAPHICStools.addLegendApart( - axs[n], - ratio=0.7, - withleg=True, - extraPad=0, - size=7, - loc="upper left", - ) - for n in ["F"]: - GRAPHICStools.addLegendApart( - axs[n], - ratio=0.7, - withleg=worked, - extraPad=0, - size=7, - loc="upper left", - ) - for n in ["A", "D"]: - GRAPHICStools.addLegendApart( - axs[n], - ratio=0.7, - withleg=False, - extraPad=0, - size=7, - loc="upper left", - ) - - def plot(self, labels=[""]): - from mitim_tools.misc_tools.GUItools import FigureNotebook - - self.fn = FigureNotebook("CGYRO Notebook", geometry="1600x1000") - - colors = GRAPHICStools.listColors() - - fig = self.fn.add_figure(label="Fluxes Time Traces") - axsFluxes_t = fig.subplot_mosaic( - """ - ABC - DEF - """ - ) - - for j in range(len(labels)): - self.plot_fluxes( - axs=axsFluxes_t, - label=labels[j], - c=colors[j], - plotLegend=j == len(labels) - 1, - ) - class CGYROinput: def __init__(self, file=None): diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py new file mode 100644 index 00000000..ab28cc0a --- /dev/null +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -0,0 +1,367 @@ +''' +From NTH +''' + +import sys +import numpy as np +import matplotlib.pyplot as plt +from gacodefuncs import * +import statsmodels.api as sm +from cgyro.data import cgyrodata +import math +#from omfit_classes import omfit_gapy + +def grab_cgyro_nth(data_dir, tstart, plotflag, printflag, file = None): + + #Get all the relevant simulation quantities + data = cgyrodata(data_dir+'/') + print(dir(data)) + data.getflux() + rho_star=(data.vth_norm/((1.602e-19*data.b_unit)/(data.mass_norm*1e-27)))/data.a_meters + nt=data.n_time + t=data.t + dt=t[0] + ky=abs(data.ky) + n_n=data.n_n + n_spec=data.n_species + n_field=data.n_field + #print(n_field) + flux = data.ky_flux + #Shape should be (species,flux,field,ky,time) + #rint(dir(data)) + roa=data.rmin + #print data.__init__ + tmax=np.amax(t) + sflux = np.sum(flux,axis=3) + kflux=np.mean(flux,axis=4) + tkflux=np.sum(kflux,axis=2) + #Values of gradients and GB normalizations + alti=data.dlntdr[0] #technically ion 1 scale length + alte=data.dlntdr[n_spec-1] + alne=data.dlnndr[n_spec-1] + qgb=data.q_gb_norm + ggb=data.gamma_gb_norm + pgb=data.pi_gb_norm + sgb=qgb/data.a_meters + + #Total electron heat flux, sum over field + eflux_all=sflux[n_spec-1,1,:,:] + eflux=np.sum(eflux_all,axis=0) + + #ES electron heat flux, just electrostatic + efluxes=eflux_all[0,:] + + #EM electron heat flux, just A_|| + efluxem=eflux_all[1,:] + + #Total electron particle flux, sum over fields + epflux_all=sflux[n_spec-1,0,:,:] + epflux=np.sum(epflux_all,axis=0) + + #Total impurity particle flux, sum over fields (***NOTE THIS IS FOR SPECIES n_Spec -3******) + impflux_all=sflux[n_spec-3,0,:,:] + impflux=np.sum(impflux_all,axis=0) + + #Depending on if electrostatic of E&M output components + if n_field ==1: + epfluxp=epflux_all[0,:] + epfluxap=epflux_all[0,:]*0.0 #Zero out the A_par array + epfluxbp=epflux_all[0,:]*0.0 #Zero out the B_par array + + if n_field == 2: + epfluxp=epflux_all[0,:] + epfluxap=epflux_all[1,:] + epfluxbp=epflux_all[0,:]*0.0 #Zero out the B_par array + + if n_field ==3: + epfluxp=epflux_all[0,:] + epfluxap=epflux_all[1,:] + epfluxbp=epflux_all[2,:] + + + #Total (ES+EM) ion flux summed over ions + iflux_all=sflux[0:n_spec-1,1,:,:] + iflux=np.sum(iflux_all,axis=0) #Sum over ions + iflux=np.sum(iflux,axis=0) #Sum over fields + + #Sum the momentum flux (all species) + mflux_all=sflux[0:n_spec,2,:,:] + mflux=np.sum(mflux_all,axis=0) # Sum over species + mflux=np.sum(mflux,axis=0) # Sum over fields + + #Total electron turbulent exchange, sm over all fields + + #Put in an option if turbulent exchange is not enabled + if sflux.shape[1] == 4: + turflux_all=sflux[n_spec-1,3,:,:] + else: + turflux_all=sflux[n_spec-1,2,:,:]*0.0 #Fill in with dummy values and 0 + + turflux=np.sum(turflux_all,axis=0) + + + if np.amax(ky) > 1.0: + e_ind=np.where(ky > 1.0)[0] + eflux_elec_tmp1=flux[:,:,:,e_ind,:] + eflux_elec_tmp2=np.sum(eflux_elec_tmp1,axis=3) + eflux_elec_all=eflux_elec_tmp2[n_spec-1,1,:,:] + eflux_elec=np.sum(eflux_elec_all,axis=0) + + #Define max values for plot + imax=np.amax(iflux)*qgb + emax=np.amax(eflux)*qgb + smax=imax+emax + + #Determine if the time step has changed mid simulation + #Change the value of tstart internally if it has + tstart_save=tstart + tend=data.t[nt-1] + tstart=(np.abs(t - tstart)).argmin() + + #Take the mean values of the fluxes + m_qe= np.mean(eflux[int(tstart):int(nt)+1]) + m_qi= np.mean(iflux[int(tstart):int(nt)+1]) + m_ge= np.mean(epflux[int(tstart):int(nt)+1]) + m_qe_elec=np.mean(eflux_elec[int(tstart):int(nt)+1]) + m_ge_p=np.mean(epfluxp[int(tstart):int(nt)+1]) + m_ge_ap=np.mean(epfluxap[int(tstart):int(nt)+1]) + m_ge_bp=np.mean(epfluxbp[int(tstart):int(nt)+1]) + m_gimp=np.mean(impflux[int(tstart):int(nt)+1]) + m_mo=np.mean(mflux[int(tstart):int(nt)+1]) + m_tur=np.mean(turflux[int(tstart):int(nt)+1]) + print(nt) + + #Calculate the standard deviations of the fluxes based on autocorrelation + i_tmp=iflux[int(tstart):int(nt-1)] + e_tmp=eflux[int(tstart):int(nt-1)] + ep_tmp=epflux[int(tstart):int(nt-1)] + imp_tmp=impflux[int(tstart):int(nt-1)] + mo_tmp=mflux[int(tstart):int(nt-1)] + tur_tmp=turflux[int(tstart):int(nt-1)] + i_acf=sm.tsa.acf(i_tmp) + i_array=np.asarray(i_acf) + icor=(np.abs(i_array-0.36)).argmin() + e_acf=sm.tsa.acf(e_tmp) + e_array=np.asarray(e_acf) + ecor=(np.abs(e_array-0.36)).argmin() + ep_acf=sm.tsa.acf(ep_tmp) + ep_array=np.asarray(ep_acf) + epcor=(np.abs(ep_array-0.36)).argmin() + imp_acf=sm.tsa.acf(imp_tmp) + imp_array=np.asarray(imp_acf) + impcor=(np.abs(imp_array-0.36)).argmin() + mo_acf=sm.tsa.acf(mo_tmp) + mo_array=np.asarray(mo_acf) + mocor=(np.abs(mo_array-0.36)).argmin() + tur_acf=sm.tsa.acf(tur_tmp) + tur_array=np.asarray(tur_acf) + turcor=(np.abs(tur_array-0.36)).argmin() + + n_corr_i=(nt-tstart)/(3.0*icor) #Define "sample" as 3 x autocor time + n_corr_e=(nt-tstart)/(3.0*ecor) + n_corr_ep=(nt-tstart)/(3.0*epcor) + n_corr_imp=(nt-tstart)/(3.0*impcor) + n_corr_mo=(nt-tstart)/(3.0*mocor) + n_corr_tur=(nt-tstart)/(3.0*turcor) + print(n_corr_i) + print(n_corr_e) + print(n_corr_ep) + print(n_corr_imp) + print(n_corr_mo) + print(n_corr_tur) + std_qe=np.std(eflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_e) + std_qi=np.std(iflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_i) + std_ge=np.std(epflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_ep) + std_gimp=np.std(impflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_imp) + std_mo=np.std(mflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_mo) + std_tur=np.std(turflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_tur) + + m_qees= np.mean(efluxes[int(tstart):int(nt-1)]) + m_qeem= np.mean(efluxem[int(tstart):int(nt-1)]) + + #Make some string values so they can be truncated + s_alti=str(alti) + s_alte=str(alte) + s_alne=str(alne) + + #Print all the simulation information + print('') + print('========================================') + print('Time to start is:') + print(tstart) + print('Max simulation time is') + print(nt) + print('Simulation Gradients') + print('a/LTi = '+s_alti[:6]) + print('a/LTe = '+s_alte[:6]) + print('a/Lne = '+s_alne[:6]) + print('') + print('======================') + print('Heat Flux') + print('======================') + print('Q_gb is') + print(f"{qgb:.4f}") + print('Mean Qi (in GB)') + print(f"{m_qi:.4f}") + print('Qi Std deviation (in GB)') + print(f"{std_qi:.4f}") + print('Mean in MW/m^2') + print(f"{m_qi*qgb:.4f}") + print('----------------------') + print('Mean Qe (in GB)') + print(f"{m_qe:.4f}") + print('Qe Std deviation (in GB)') + print(f"{std_qe:.4f}") + print('Mean in MW/m^2') + print(f"{m_qe*qgb:.4f}") + print('Qi/Qe') + print(f"{m_qi/m_qe:.4f}") + print('High-k Qe (GB)') + print(f"{m_qe_elec:.4f}") + print('') + print('======================') + print('Particle Flux') + print('======================') + print('Gamma_gb is:') + print(f"{ggb:.4f}") + print('Mean Gamma_e (in GB)') + print(f"{m_ge:.4f}") + print('Gamma_e Std deviation (in GB)') + print(f"{std_ge:.4f}") + print('Mean in e19/m^2*s') + print(f"{m_ge*ggb:.4f}") + print('Mean Impurity Flux (GB)') + print(f"{m_gimp:.4e}") + print('Gamma_imp Std. deviation (in GB)') + print(f"{std_gimp:.4e}") + print('Electron Particle Flux Components(GB)') + print('Phi: 'f"{m_ge_p:.4f}") + print('A_||: 'f"{m_ge_ap:.4f}") + print('B_||: 'f"{m_ge_bp:.4f}") + print('') + print('======================') + print('Momentum Flux') + print('======================') + print('Pi_gb is:') + print(f"{pgb:.4f}") + print('Mean Pi (in GB)') + print(f"{m_mo:.4f}") + print('Momentum Flux Std deviation (in GB)') + print(f"{std_mo:.4f}") + print('Mean in J/m^2') + print(f"{m_mo*pgb:.4f}") + print('') + print('======================') + print('Turbulent Exchange') + print('======================') + print('S_gb is:') + print(f"{sgb:.4f}") + print('Mean Elec Turb Ex. (in GB)') + print(f"{m_tur:.4f}") + print('Turb Ex. Std deviation (in GB)') + print(f"{std_tur:.4f}") + + #Reset the starting time to correspond to the a/cs + tstart=tstart_save + + #Print the results to results file + if printflag ==1: + file1=open(file,"a") + file1.write('r/a = {0:4f} \t a/Lne = {1:4f} \t a/LTi = {2:4f} \t a/LTe = {3:4f} \t Qi = {4:4f} \t Qi_std = {5:4f} \t Qe = {6:4f} \t Qe_std = {7:4f} \t Ge = {8:4f} \t Ge_std = {9:4f} \t Gimp = {10:.4e} \t Gimp_std = {11:.4e} \t Pi = {12:4f} \t Pi_std = {13:4f} \t S = {14:4f} \t S_std = {15:4f} \t Q_gb = {16:4f} \t G_gb = {17:4f} \t Pi_gb = {18:4f} \t S_gb = {19:4f} \t tstart = {20:4f} \t tend = {21:4f} \n'.format(roa,alne,alti,alte,m_qi,std_qi,m_qe,std_qe,m_ge,std_ge,m_gimp,std_gimp,m_mo,std_mo,m_tur,std_tur,qgb,ggb,pgb,sgb,tstart,nt)) + file1.close() + + + + #Plot the results + if plotflag == 1: + + summax=2.0*(m_qe*qgb+m_qi*qgb) + maxf=np.amax([np.amax(iflux)*qgb,np.amax(eflux)*qgb]) + pmax=np.amin([2.0*(m_qe*qgb+m_qi*qgb),maxf]) + gmax=np.amax(epflux*ggb) + gmin=np.amin(epflux*ggb) + + # Setting the style for the plots + plt.rc('axes', labelsize=25) + plt.rc('xtick', labelsize=25) + plt.rc('ytick', labelsize=25) + + plt.rcParams['font.family'] = 'serif' + + # Create a figure with 3 subplots vertically stacked (3 rows, 1 column) + fig, ax = plt.subplots(2, 1, figsize=(10.0, 9.0)) + + # First plot (Q vs Time) + ax[0].plot(t, iflux * qgb, 'b-') + ax[0].plot([tstart, tend], [m_qi * qgb, m_qi * qgb], 'b-', linewidth=4.0) + ax[0].plot([tstart, tend], [(m_qi + std_qi) * qgb, (m_qi + std_qi) * qgb], 'b-', linewidth=2.0) + ax[0].plot([tstart, tend], [(m_qi - std_qi) * qgb, (m_qi - std_qi) * qgb], 'b-', linewidth=2.0) + ax[0].set_title('Q$_e$ and Q$_i$ vs Time',fontsize=18) + ax[0].set_xlabel('a/c$_s$',fontsize=15) + ax[0].set_ylabel('MW/m$^2$',fontsize=15) + ax[0].axis([0.0, tmax, 0.0, pmax]) + + # Second plot (Qe vs Time) + ax[0].plot(t, eflux * qgb, 'r-') + ax[0].plot([tstart, tend], [m_qe * qgb, m_qe * qgb], 'r-', linewidth=4.0) + ax[0].plot([tstart, tend], [(m_qe + std_qe) * qgb, (m_qe + std_qe) * qgb], 'r-', linewidth=2.0) + ax[0].plot([tstart, tend], [(m_qe - std_qe) * qgb, (m_qe - std_qe) * qgb], 'r-', linewidth=2.0) + + ax[0].fill_between([tstart,tend], [(m_qe - std_qe) * qgb, (m_qe - std_qe) * qgb], [(m_qe + std_qe) * qgb, (m_qe + std_qe) * qgb], color='red', alpha=0.2) + ax[0].fill_between([tstart,tend], [(m_qi - std_qi) * qgb, (m_qi - std_qi) * qgb], [(m_qi + std_qi) * qgb, (m_qi + std_qi) * qgb], color='blue', alpha=0.2) + + #ax[0].plot(t, eflux * qgb, 'g-') # You can customize this line further + ax[0].tick_params(axis='x', labelsize=12) # x-axis tick labels font size reduced + ax[0].tick_params(axis='y', labelsize=12) + factor=10 ** 3 + factor2=10 ** 2 + num=math.floor(m_qe*qgb*factor)/factor + num2=math.floor(std_qe*qgb*factor)/factor + ax[0].text(10,pmax*0.8,'$Q_e$='+str(num)+' +/- '+str(num2),fontsize=18,color='red') + + num=math.floor(m_qi*qgb*factor)/factor + num2=math.floor(std_qi*qgb*factor)/factor + num3=math.floor(m_qi/m_qe*factor2)/factor2 + ax[0].text(10,pmax*0.88,'$Q_i$='+str(num)+' +/- '+str(num2),fontsize=18,color='blue') + ax[0].text(10,pmax*0.72,'$Q_i/Q_e$='+str(num3),fontsize=18) + + # Third plot (epflux vs Time) + ax[1].plot(t, epflux * ggb, 'g-') # Plotting epflux on the third subplot + ax[1].set_title('$\Gamma_e$ vs Time',fontsize=18) + ax[1].set_xlabel('a/c$_s$',fontsize=15) + ax[1].set_ylabel('1e19/m$^2$s',fontsize=15) + ax[1].tick_params(axis='x', labelsize=12) # x-axis tick labels font size reduced + ax[1].plot([tstart, tend], [m_ge * ggb, m_ge * ggb], 'g-', linewidth=4.0) + ax[1].plot([tstart, tend], [(m_ge + std_ge) * ggb, (m_ge + std_ge) * ggb], 'g-', linewidth=2.0) + ax[1].plot([tstart, tend], [(m_ge - std_ge) * ggb, (m_ge - std_ge) * ggb], 'g-', linewidth=2.0) + ax[1].fill_between([tstart,tend], [(m_ge - std_ge) * ggb, (m_ge - std_ge) * ggb], [(m_ge + std_ge) * ggb, (m_ge + std_ge) * ggb], color='green', alpha=0.2) + + ax[1].plot(t, epflux * ggb, 'g-') + ax[1].plot(t,t*0.0,'k',linestyle='--') + ax[1].tick_params(axis='y', labelsize=12) + ax[1].axis([0.0, tmax, gmin, gmax]) + + factor=10 ** 3 + num=math.floor(m_ge*ggb*factor)/factor + num2=math.floor(std_ge*ggb*factor)/factor + ax[1].text(10,gmax*0.8,'$\Gamma_e$='+str(num)+' +/- '+str(num2),fontsize=18) + ax[1].text(10,gmin*0.95,str(int(tend))+' a/c$_s$',fontsize=18) + # Adjust spacing between subplots + plt.tight_layout() + + # Show the plot + plt.show() + + return roa,alne,alti,alte,m_qi,std_qi,m_qe,std_qe,m_ge,std_ge,m_gimp,std_gimp,m_mo,std_mo,m_tur,std_tur,qgb,ggb,pgb,sgb,tstart,nt + + +if __name__ == "__main__": + # Example usage + data_dir = sys.argv[1] # Directory containing the cgyro data + tstart = float(sys.argv[2]) # Start time for analysis + plotflag = int(sys.argv[3]) # Flag to indicate if plotting is required + printflag = int(sys.argv[4]) # Flag to indicate if printing results is required + file = sys.argv[5] # Print to this file + + grab_cgyro_nth(data_dir, tstart, plotflag, printflag, file = file) + From 880bcd92f7e30ddd332d6006afb936c0fb1d9528 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 29 May 2025 21:23:39 -0400 Subject: [PATCH 053/385] Fix units --- src/mitim_tools/gacode_tools/CGYROtools.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index e90dcb08..96e9792f 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -233,13 +233,14 @@ def read(self, label="cgyro1", folder=None, tmin = 0.0): self.results[label].Qi = self.results[label].Qi_all.sum(axis=0) self.results[label].Ge = ys[-1, 0, :] - self.results[label].Qe_mean, self.results[label].Qe_std = self.derive_statistics(self.results[label].t, self.results[label].Qe, x_min=tmin) - self.results[label].Qi_mean, self.results[label].Qi_std = self.derive_statistics(self.results[label].t, self.results[label].Qi, x_min=tmin) - self.results[label].Ge_mean, self.results[label].Ge_std = self.derive_statistics(self.results[label].t, self.results[label].Ge, x_min=tmin) - roa,alne,self.results[label].aLTi,alte,self.results[label].Qi_mean,self.results[label].Qi_std,self.results[label].Qe_mean,self.results[label].Qe_std,self.results[label].Ge_mean, self.results[label].Ge_std,m_gimp,std_gimp,m_mo,std_mo,m_tur,std_tur,qgb,ggb,pgb,sgb,tstart,nt = CGYROutils.grab_cgyro_nth(str(folder.resolve()), tmin, False, False) + + self.results[label].Qi_mean *= qgb + self.results[label].Qi_std *= qgb + self.results[label].Qe_mean *= qgb + self.results[label].Qe_std *= qgb - + def derive_statistics(self,x,y,x_min=0.0): return y.mean(), y.std() From c5268ca402bcadcfb97d63031dec1d8605a3cb7a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 30 May 2025 10:20:05 -0400 Subject: [PATCH 054/385] bug fix forgot from name transition --- src/mitim_modules/powertorch/physics_models/transport_cgyro.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 031ed0a0..c8580d1a 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -122,7 +122,7 @@ def _print_info(self): for var, varn in zip( ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ce (MW/m^2)", "CZ (MW/m^2)", "MtJm2 (J/m^2) "], - ["QeMWm2", "Pi", "Ce", "CZ", "MtJm2"], + ["QeMWm2", "QiMWm2", "Ce", "CZ", "MtJm2"], ): txt += f"\n{var} = " for j in range(self.powerstate.plasma["rho"].shape[1] - 1): From 58280243f1630d5e252f603df133602884a309e8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 30 May 2025 13:51:17 -0400 Subject: [PATCH 055/385] Changed BCdens function --- src/mitim_tools/gacode_tools/PROFILEStools.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 2fe38315..6967e4e4 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -3792,6 +3792,8 @@ def changeRFpower(self, PrfMW=25.0): for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: self.profiles[i] = self.profiles[i] * PrfMW / self.derived["qRF_MWmiller"][-1] + self.deriveQuantities() + def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0.2): ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) @@ -3817,25 +3819,25 @@ def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0 raise Exception("no edge") - def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5): + def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5, isn20_edge=True): ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) - print(f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") - - factor = n20 / self.derived["ne_vol20"] + # Determine the factor to scale the density (either average or at rho) + if not isn20_edge: + print(f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") + factor = n20 / self.derived["ne_vol20"] + else: + print(f"- Changing the density at rho={rho} from {self.profiles['ne(10^19/m^3)'][ix]*1E-1:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") + factor = n20 / (self.profiles["ne(10^19/m^3)"][ix]*1E-1) + # ------------------------------------------------------------------ + # Scale the density profiles for i in ["ne(10^19/m^3)", "ni(10^19/m^3)"]: self.profiles[i] = self.profiles[i] * factor + # Apply the edge condition if typeEdge == "linear": - factor_x = ( - np.linspace( - self.profiles["ne(10^19/m^3)"][ix], - nedge20 * 1e1, - len(self.profiles["rho(-)"][ix:]), - ) - / self.profiles["ne(10^19/m^3)"][ix:] - ) + factor_x = np.linspace(self.profiles["ne(10^19/m^3)"][ix],nedge20 * 1e1,len(self.profiles["rho(-)"][ix:]),)/ self.profiles["ne(10^19/m^3)"][ix:] self.profiles["ne(10^19/m^3)"][ix:] = self.profiles["ne(10^19/m^3)"][ix:] * factor_x From fd186e9964903758b09c4fefb054edc70dfb746e Mon Sep 17 00:00:00 2001 From: Audrey Saltzman <37987951+AudreySaltzman@users.noreply.github.com> Date: Fri, 16 May 2025 16:20:46 -0400 Subject: [PATCH 056/385] Add f90nml as dependency --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index bae000fd..04b1e190 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,8 @@ dependencies = [ "scikit-image", # Stricly not for MITIM, but good to have for pygacode "psutil", "onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models - "tensorflow" + "tensorflow", + "f90nml", ] [project.optional-dependencies] From 973ee815cb30ea3c6ab818d14d6c97617a768885 Mon Sep 17 00:00:00 2001 From: Audrey Saltzman <37987951+AudreySaltzman@users.noreply.github.com> Date: Fri, 16 May 2025 17:49:59 -0400 Subject: [PATCH 057/385] EPEDtools --- src/mitim_tools/eped_tools/EPEDtools.py | 202 ++++++++++++++++++++++++ 1 file changed, 202 insertions(+) create mode 100644 src/mitim_tools/eped_tools/EPEDtools.py diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py new file mode 100644 index 00000000..e276077b --- /dev/null +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -0,0 +1,202 @@ +import os +import re +import subprocess +import f90nml +from pathlib import Path +import numpy as np +import pandas as pd +import xarray as xr + + +def convert_to_dimensional(df): + #ee = 1.60217663e-19 + mu0 = 1.25663706127e-6 + df['a'] = df['r'] / df['epsilon'] + df['ip'] = 1.0e-6 * (2.0 * np.pi * np.square(df['a']) * df['kappa'] * df['bt']) / (df['qstar'] * df['r'] * mu0) + df['neped'] = 10.0 * df['fgped'] * df['ip'] / (np.pi * np.square(df['a'])) + df['nesep'] = 0.25 * df['neped'] + #df['teped'] = 2500 * df['bt'] * df['ip'] * df['betan'] / (3 * df['a'] * 1.5 * df['neped']) + #df['teped'] = df['teped'].clip(upper=8000) + df['teped'] = df['r'] * 0.0 - 1.0 + return df + + +def convert_to_dimensionless(df): + mu0 = 1.25663706127e-6 + df['epsilon'] = df['r'] / df['a'] + df['fgped'] = df['neped'] * np.pi * np.square(df['a']) / (10.0 * df['ip']) + df['qstar'] = (2.0 * np.pi * np.square(df['a']) * df['kappa'] * df['bt']) / (1.0e6 * df['ip'] * df['r'] * mu0) + df['nesep'] = 0.25 * df['neped'] + df['teped'] = df['r'] * 0.0 - 1.0 + return df + + +def setup_eped(output_path, inputs_list, template_path): + + output_path = Path(output_path).resolve() # Ensure absolute path + output_path.mkdir(parents=True, exist_ok=True) + + subprocess.run(['cp', str(template_path.resolve() / 'exec_eped.sh'), str(output_path)]) + subprocess.run(['cp', str(template_path.resolve() / 'submit_eped_array_psfc.batch'), str(output_path)]) + subprocess.run(['cp', str(template_path.resolve() / 'postprocessing.py'), str(output_path)]) + rpaths = [] + + for run_num, inputs in enumerate(inputs_list): + run_id = f'run{run_num + 1:03d}' + rpath = output_path / run_id # Construct the absolute path for the run directory + subprocess.run(['cp', '-r', str(template_path.resolve() / 'eped_run_template'), str(rpath)]) + + # Edit input file + input_file = rpath / 'eped.input' + contents = f90nml.read(str(input_file)) + for param, value in inputs.items(): + contents['eped_input'][param] = value + contents.write(str(input_file), force=True) + + rpaths.append(rpath) + + #logger.info(f'{len(inputs_list)} Runs created at {output_path}') + + return rpaths + + +def setup_array_batch(launch_path, rpaths, maxqueue=5): + + # Convert to Path object and ensure absolute path + launch_path = Path(launch_path).resolve() + + s = '' + for path in rpaths: + if s: + s += '\n' + s += f'"./exec_eped.sh {path.resolve()}"' + + # Use proper Path object for file operations + batch_file = launch_path / 'submit_eped_array_psfc.batch' + with batch_file.open('r') as f: + content = f.read() + new_content = re.sub('', str(len(rpaths) - 1), content) + new_content = re.sub('', str(maxqueue), new_content) + new_content = re.sub('', str(launch_path), new_content) # Convert to string for substitution + new_content = re.sub('#', s, new_content) + with batch_file.open('w') as f: + f.write(new_content) + + #logger.info('Batch array created') + + return batch_file + + +def postprocess_eped(data, diamagnetic_stab_rule, stability_threshold): + + coords = {k: data[k].values for k in ['dim_height', 'dim_widths', 'dim_nmodes', 'dim_rho', 'dim_three', 'dim_one']} + data = data.assign_coords(coords) + + x = data['eq_betanped'].data + index = np.where(x < 0)[0] + if diamagnetic_stab_rule == 'G': + y = data['gamma'].data.copy() + elif diamagnetic_stab_rule in ['GH', 'HG']: + y = data['gamma_PB'].data.copy() + y *= data['gamma'].data.copy() + elif diamagnetic_stab_rule == 'H': + y = data['gamma_PB'].data.copy() + else: + y = data['gamma'].data.copy() + y[index, :] = np.nan + + data['stability'] = (('dim_height', 'dim_nmodes'), y) + y0 = np.nanmax(y, 1) + y0 = np.where(y0 == None, 0, y0) + indices = np.where(y0 > stability_threshold)[0] + if len(indices): + step = indices[0] + else: + step = -1 + + dims = ('dim_one') + data['stability_rule'] = (dims, [diamagnetic_stab_rule]) + data['stability_threshold'] = (dims, np.array([stability_threshold])) + if step > 0: + data['stability_index'] = (dims, np.array([step])) + data['pped'] = (dims, np.array([data['eq_pped'].data[step] * 1.0e3])) + data['ptop'] = (dims, np.array([data['eq_ptop'].data[step] * 1.0e3])) + data['tped'] = (dims, np.array([data['eq_tped'].data[step]])) + data['ttop'] = (dims, np.array([data['eq_ttop'].data[step]])) + data['wpped'] = (dims, np.array([data['eq_wped_psi'].data[step]])) + data['wptop'] = (dims, np.array([data['eq_wped_psi'].data[step] * 1.5])) + data['wrped'] = (dims, np.array([data['eq_wped_rho'].data[step]])) + if np.any(data['tesep'].data < 0): + data['tesep'] = (dims, np.array([75.0])) + data['nesep'] = 0.25 * data['neped'] + + return data + +def read_eped_file(ipaths): + invars = ['ip', 'bt', 'r' , 'a', 'kappa', 'delta', 'neped', 'betan', 'zeffped', 'nesep', 'tesep'] + data_arrays = [] + for ipath in ipaths: + dummy_coords = { + 'dim_height': np.empty((0, ), dtype=int), + 'dim_nmodes': np.empty((0, ), dtype=int), + 'dim_widths': np.empty((0, ), dtype=int), + 'dim_rho': np.empty((0, ), dtype=int), + 'dim_three': np.empty((0, ), dtype=int), + 'dim_one': np.arange(1), + } + set_inputs = f90nml.read(str(ipath.parent.parent.parent / 'eped.input')) + dummy_vars = {k: (['dim_one'], [v]) for k, v in set_inputs['eped_input'].items() if k in invars} + data = xr.Dataset(coords=dummy_coords, data_vars=dummy_vars) + if ipath.is_file(): + data = xr.open_dataset(f'{ipath.resolve()}', engine='netcdf4') + data = postprocess_eped(data, 'G', 0.03) + data_arrays.append(data.expand_dims({'filename': [ipath.parent.parent.parent.name]})) + + dataset = xr.merge(data_arrays, join='outer', fill_value=np.nan).sortby('filename') + return dataset + +def launch_eped_slurm(base_input_path, scan_params, nscan, output_path, template_path, run_tag, wait=False): + input_params = f90nml.read(str(base_input_path)).todict().get('eped_input', {}) + input_params.update(scan_params) + data = {} + for var, val in input_params.items(): + if isinstance(val, (tuple, list, np.ndarray)) and len(val) > 1: + data[var] = np.linspace(val[0], val[1], nscan) + else: + data[var] = np.zeros((nscan, )) + val + #if scan_var == 'qstar': # Use for ip scan + # data['fgped'] = (0.5 / 3.5) * data['qstar'] + inp = pd.DataFrame(data=data, index=pd.RangeIndex(nscan)) + #inp = convert_to_dimensional(inp) + inputs = [{ivar: inp[ivar].iloc[i] for ivar in ivars} for i in range(len(inp))] + run_paths = setup_eped(output_path, inputs, template_path) + spath = setup_array_batch(output_path, run_paths) + inp.to_hdf(output_path / f'{output_path.name}.h5', key='/data') + command = ['sbatch'] + if wait: + command.append('--wait') + command.append(f'{spath.resolve()}') + subprocess.run(command) + + return run_paths + +def main(): + + rootdir = Path(os.environ.get('PIXI_PROJECT_ROOT', './')) + run_tag = 'scan_tesep' + ivars = ['ip', 'bt', 'r', 'a', 'kappa', 'delta', 'neped', 'betan', 'zeffped', 'nesep', 'tesep', 'teped'] + base_input_path = Path('./') / 'eped.input' + scan_params = { + # 'tesep': [50.0, 300.0], + } + nscan = 1 + output_path = Path('./') / f'eped_{run_tag}' + template_path = rootdir / 'ips-eped-master' / 'template' / 'engaging' + wait = False + + launch_eped_slurm(base_input_path, scan_params, nscan, output_path, template_path, run_tag, wait=wait) + + + +if __name__ == '__main__': + main() From 6b0d0440de5a03f62626620ca7974fee79f912c0 Mon Sep 17 00:00:00 2001 From: audreysa Date: Fri, 16 May 2025 18:00:33 -0400 Subject: [PATCH 058/385] Fix Launch EPED t=to have ivars --- src/mitim_tools/eped_tools/EPEDtools.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index e276077b..9136875f 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -156,6 +156,7 @@ def read_eped_file(ipaths): return dataset def launch_eped_slurm(base_input_path, scan_params, nscan, output_path, template_path, run_tag, wait=False): + ivars = ['ip', 'bt', 'r', 'a', 'kappa', 'delta', 'neped', 'betan', 'zeffped', 'nesep', 'tesep', 'teped'] input_params = f90nml.read(str(base_input_path)).todict().get('eped_input', {}) input_params.update(scan_params) data = {} @@ -183,8 +184,7 @@ def launch_eped_slurm(base_input_path, scan_params, nscan, output_path, template def main(): rootdir = Path(os.environ.get('PIXI_PROJECT_ROOT', './')) - run_tag = 'scan_tesep' - ivars = ['ip', 'bt', 'r', 'a', 'kappa', 'delta', 'neped', 'betan', 'zeffped', 'nesep', 'tesep', 'teped'] + run_tag = 'mitim_eped_test' base_input_path = Path('./') / 'eped.input' scan_params = { # 'tesep': [50.0, 300.0], @@ -194,8 +194,8 @@ def main(): template_path = rootdir / 'ips-eped-master' / 'template' / 'engaging' wait = False - launch_eped_slurm(base_input_path, scan_params, nscan, output_path, template_path, run_tag, wait=wait) + launch_eped_slurm(base_input_path, scan_params, nscan, output_path, template_path, run_tag, wait=wait) if __name__ == '__main__': From 79d85700f4dbb1cbc762ddd35720cc98e2372cd6 Mon Sep 17 00:00:00 2001 From: audreysa Date: Fri, 16 May 2025 18:27:10 -0400 Subject: [PATCH 059/385] Changed launch to accept namelist object --- src/mitim_tools/eped_tools/EPEDtools.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index 9136875f..77aa11e6 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -155,9 +155,8 @@ def read_eped_file(ipaths): dataset = xr.merge(data_arrays, join='outer', fill_value=np.nan).sortby('filename') return dataset -def launch_eped_slurm(base_input_path, scan_params, nscan, output_path, template_path, run_tag, wait=False): +def launch_eped_slurm(input_params, scan_params, nscan, output_path, template_path, run_tag, wait=False): ivars = ['ip', 'bt', 'r', 'a', 'kappa', 'delta', 'neped', 'betan', 'zeffped', 'nesep', 'tesep', 'teped'] - input_params = f90nml.read(str(base_input_path)).todict().get('eped_input', {}) input_params.update(scan_params) data = {} for var, val in input_params.items(): @@ -194,8 +193,9 @@ def main(): template_path = rootdir / 'ips-eped-master' / 'template' / 'engaging' wait = False + input_params = f90nml.read(str(base_input_path)).todict().get('eped_input', {}) - launch_eped_slurm(base_input_path, scan_params, nscan, output_path, template_path, run_tag, wait=wait) + launch_eped_slurm(input_params, scan_params, nscan, output_path, template_path, run_tag, wait=wait) if __name__ == '__main__': From 114bce3153710b1e3ce871a4262a2aab996c649e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 23 May 2025 17:13:43 -0400 Subject: [PATCH 060/385] Progress on EPED class --- src/mitim_tools/eped_tools/EPEDtools.py | 112 ++++++++++++++++++++++++ tests/EPED_workflow.py | 28 ++++++ 2 files changed, 140 insertions(+) create mode 100644 tests/EPED_workflow.py diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index 77aa11e6..dcfbbed6 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -3,10 +3,122 @@ import subprocess import f90nml from pathlib import Path +from mitim_tools.misc_tools import FARMINGtools import numpy as np import pandas as pd import xarray as xr +from IPython import embed + +class EPED: + def __init__( + self, + folder + ): + + self.folder = Path(folder) + self.folder.mkdir(parents=True, exist_ok=True) + + self.results = {} + + def run( + self, + subfolder = 'run1', + input_params = None, + nproc = 64, + ): + + # Set up folder + self.folder_run = self.folder / subfolder + self.folder_run.mkdir(parents=True, exist_ok=True) + + # ------------------------------------ + # Write input file to EPED + # ------------------------------------ + + eped_input_file = self.folder_run / 'eped.input' + + shot = 0 + timeid = 0 + + # Update with fixed parameters + input_params.update( + {'num_scan': 1, + 'shot': shot, + 'timeid': timeid, + 'runid': 0, + 'm': 2, + 'z': 1, + 'mi': 20, + 'zi': 10, + 'tewid': 0.03, + 'ptotwid': 0.03, + 'teped': -1, + 'ptotped': -1, + } + ) + + eped_input = {'eped_input': input_params} + nml = f90nml.Namelist(eped_input) + + # Write the input file + f90nml.write(nml, eped_input_file, force=True) + + # Initialize Job + self.eped_job = FARMINGtools.mitim_job(self.folder_run) + + self.eped_job.define_machine( + "eped", + "mitim_eped", + launchSlurm=False, + ) + + # ------------------------------------- + # Executable commands + # ------------------------------------- + + EPEDcommand = f'cp $EPED_SOURCE_PATH/template/engaging/eped_run_template/* {self.eped_job.folderExecution}/. && export NPROC_EPED={nproc} && ips.py --config=eped.config --platform=psfc_cluster.conf' + + # ------------------------------------- + # Execute + # ------------------------------------- + + output_file = f'e{shot:06d}.{timeid:05d}' + + self.eped_job.prep( + EPEDcommand, + input_files=[eped_input_file], + output_files=[f'eped/SUMMARY/{output_file}'], + ) + + self.eped_job.run(removeScratchFolders=False) + + + # Rename output file + os.system(f'mv {self.folder_run / output_file} {self.folder_run / "output.nc"}') + + def read( + self, + subfolder = 'run1', + name = 'run1', + ): + + output_file = self.folder / subfolder / 'output.nc' + + data = xr.open_dataset(f'{output_file.resolve()}', engine='netcdf4') + data = postprocess_eped(data, 'G', 0.03) + + self.results[name] = data + + def plot( + self, + name = 'run1', + ): + + data = self.results[name] + +# ************************************************************************************************************ +# ************************************************************************************************************ def convert_to_dimensional(df): #ee = 1.60217663e-19 diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py new file mode 100644 index 00000000..d907d6d2 --- /dev/null +++ b/tests/EPED_workflow.py @@ -0,0 +1,28 @@ +from re import sub +from mitim_tools.eped_tools import EPEDtools +from mitim_tools import __mitimroot__ + +folder = __mitimroot__ / "tests" / "scratch" / "eped_test" + +eped = EPEDtools.EPED(folder=folder) + +eped.run( + subfolder = 'run1', + input_params = { + 'ip': 0.5, + 'bt': 1.0, + 'r': 1.0, + 'a': 0.5, + 'kappa': 1.5, + 'delta': 0.5, + 'neped': 1.0, + 'betan': 0.5, + 'zeffped': 1.0, + 'nesep': 0.25, + 'tesep': 1.0, + }, + nproc = 64 +) + +eped.read(subfolder='run1') + From 537979b4335fb75440a227dab22c8715ea5e7711 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 27 May 2025 13:59:47 -0400 Subject: [PATCH 061/385] MVP EPED working --- src/mitim_tools/eped_tools/EPEDtools.py | 61 ++++++++++++++++--- .../misc_tools/scripts/compare_namelist.py | 2 +- tests/EPED_workflow.py | 5 +- 3 files changed, 58 insertions(+), 10 deletions(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index dcfbbed6..b9e0df40 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -1,12 +1,14 @@ import os import re import subprocess +import matplotlib.pyplot as plt import f90nml from pathlib import Path -from mitim_tools.misc_tools import FARMINGtools +from mitim_tools.misc_tools import FARMINGtools, GRAPHICStools import numpy as np import pandas as pd import xarray as xr +from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed class EPED: @@ -21,13 +23,31 @@ def __init__( self.results = {} + + def cold_start_checker( + self, + subfolder = 'run1' + ): + + # Check if the run folder exists + self.folder_run = self.folder / subfolder + + output_file = self.folder_run / 'output.nc' + + return not output_file.exists() or not output_file.is_file() + def run( self, subfolder = 'run1', input_params = None, nproc = 64, + cold_start = False, ): + if (not cold_start) and not self.cold_start_checker(subfolder): + print(f'\t> Run {subfolder} already exists and cold_start is set to False. Skipping run.', typeMsg='i') + return + # Set up folder self.folder_run = self.folder / subfolder self.folder_run.mkdir(parents=True, exist_ok=True) @@ -70,9 +90,12 @@ def run( self.eped_job.define_machine( "eped", "mitim_eped", - launchSlurm=False, + slurm_settings={ + 'name': 'mitim_eped', + 'minutes': 20, + 'ntasks': nproc, + } ) - # ------------------------------------- # Executable commands # ------------------------------------- @@ -93,9 +116,8 @@ def run( self.eped_job.run(removeScratchFolders=False) - # Rename output file - os.system(f'mv {self.folder_run / output_file} {self.folder_run / "output.nc"}') + os.system(f'mv {self.folder_run / 'eped' / 'SUMMARY' /output_file} {self.folder_run / "output.nc"}') def read( self, @@ -112,10 +134,35 @@ def read( def plot( self, - name = 'run1', + labels = ['run1'], ): - data = self.results[name] + plt.ion(); fig, axs = plt.subplots(2, 1, figsize=(10, 6)) + + colors = GRAPHICStools.listColors() + + for i,name in enumerate(labels): + + data = self.results[name] + + neped = float(data['neped']) + ptop = float(data['ptop']) + wtop = float(data['wptop']) + + axs[0].plot(neped,ptop,'-s', c = colors[i], ms = 12) + axs[1].plot(neped,wtop,'-s', c = colors[i], ms = 12) + + ax = axs[0] + ax.set_xlabel('neped ($10^{19}m^{-3}$)') + ax.set_ylabel('ptop (kPa)') + ax.set_ylim(bottom=0) + GRAPHICStools.addDenseAxis(ax) + + ax = axs[1] + ax.set_xlabel('neped ($10^{19}m^{-3}$)') + ax.set_ylabel('wptop (psi_pol)') + ax.set_ylim(bottom=0) + GRAPHICStools.addDenseAxis(ax) # ************************************************************************************************************ # ************************************************************************************************************ diff --git a/src/mitim_tools/misc_tools/scripts/compare_namelist.py b/src/mitim_tools/misc_tools/scripts/compare_namelist.py index 0897e832..9ad08b77 100644 --- a/src/mitim_tools/misc_tools/scripts/compare_namelist.py +++ b/src/mitim_tools/misc_tools/scripts/compare_namelist.py @@ -147,7 +147,7 @@ def printTable(diff, warning_percent=1e-1): perc = np.nan print( f"{key:>15}{str(diff[key][0]):>25}{str(diff[key][1]):>25} (~{perc:.0e}%)", - typeMsg="w" if perc > warning_percent else "", + typeMsg="i" if perc > warning_percent else "", ) else: print(f"{key:>15}{str(diff[key][0]):>25}{'':>25}") diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py index d907d6d2..1d0af189 100644 --- a/tests/EPED_workflow.py +++ b/tests/EPED_workflow.py @@ -1,4 +1,3 @@ -from re import sub from mitim_tools.eped_tools import EPEDtools from mitim_tools import __mitimroot__ @@ -21,8 +20,10 @@ 'nesep': 0.25, 'tesep': 1.0, }, - nproc = 64 + nproc = 64, + cold_start = True, ) eped.read(subfolder='run1') +eped.plot(labels=['run1']) From 79f6ee46d3d42cc8f3ac86d6f69676f7c088ee0a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 27 May 2025 14:46:27 -0400 Subject: [PATCH 062/385] Fixes to input overwritten and cold start --- src/mitim_tools/eped_tools/EPEDtools.py | 38 +++++++++++++++++-------- tests/EPED_workflow.py | 21 +++++++------- 2 files changed, 37 insertions(+), 22 deletions(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index b9e0df40..0ff3d33e 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -4,7 +4,7 @@ import matplotlib.pyplot as plt import f90nml from pathlib import Path -from mitim_tools.misc_tools import FARMINGtools, GRAPHICStools +from mitim_tools.misc_tools import FARMINGtools, GRAPHICStools, IOtools import numpy as np import pandas as pd import xarray as xr @@ -24,7 +24,7 @@ def __init__( self.results = {} - def cold_start_checker( + def case_exists( self, subfolder = 'run1' ): @@ -34,29 +34,41 @@ def cold_start_checker( output_file = self.folder_run / 'output.nc' - return not output_file.exists() or not output_file.is_file() + return output_file.exists() and output_file.is_file() def run( self, subfolder = 'run1', input_params = None, nproc = 64, + minutes_slurm = 60, cold_start = False, ): - if (not cold_start) and not self.cold_start_checker(subfolder): - print(f'\t> Run {subfolder} already exists and cold_start is set to False. Skipping run.', typeMsg='i') - return + self.folder_run = self.folder / subfolder + if self.case_exists(subfolder): + if cold_start: + res = print(f'\t> Run {subfolder} already exists but cold_start is set to True. Running from scratch.', typeMsg='q') + if res: + IOtools.shutil_rmtree(self.folder_run) + else: + return + else: + print(f'\t> Run {subfolder} already exists and cold_start is set to False. Skipping run.', typeMsg='i') + return + # Set up folder - self.folder_run = self.folder / subfolder + self.folder_run.mkdir(parents=True, exist_ok=True) # ------------------------------------ # Write input file to EPED # ------------------------------------ - eped_input_file = self.folder_run / 'eped.input' + eped_input_file_suffix = 'eped.input.1' # Do not call it directly 'eped.input' as it may be overwritten by the job script template copying commands + + eped_input_file = self.folder_run / eped_input_file_suffix shot = 0 timeid = 0 @@ -92,7 +104,7 @@ def run( "mitim_eped", slurm_settings={ 'name': 'mitim_eped', - 'minutes': 20, + 'minutes': minutes_slurm, 'ntasks': nproc, } ) @@ -100,7 +112,9 @@ def run( # Executable commands # ------------------------------------- - EPEDcommand = f'cp $EPED_SOURCE_PATH/template/engaging/eped_run_template/* {self.eped_job.folderExecution}/. && export NPROC_EPED={nproc} && ips.py --config=eped.config --platform=psfc_cluster.conf' + required_files_folder = '$EPED_SOURCE_PATH/template/engaging/eped_run_template' + + EPEDcommand = f'cp {required_files_folder}/* {self.eped_job.folderExecution}/. && mv {self.eped_job.folderExecution}/{eped_input_file_suffix} {self.eped_job.folderExecution}/eped.input && export NPROC_EPED={nproc} && ips.py --config=eped.config --platform=psfc_cluster.conf' # ------------------------------------- # Execute @@ -122,7 +136,7 @@ def run( def read( self, subfolder = 'run1', - name = 'run1', + label = None, ): output_file = self.folder / subfolder / 'output.nc' @@ -130,7 +144,7 @@ def read( data = xr.open_dataset(f'{output_file.resolve()}', engine='netcdf4') data = postprocess_eped(data, 'G', 0.03) - self.results[name] = data + self.results[label if label is not None else subfolder] = data def plot( self, diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py index 1d0af189..bc1680e3 100644 --- a/tests/EPED_workflow.py +++ b/tests/EPED_workflow.py @@ -8,17 +8,17 @@ eped.run( subfolder = 'run1', input_params = { - 'ip': 0.5, - 'bt': 1.0, - 'r': 1.0, - 'a': 0.5, - 'kappa': 1.5, + 'ip': 12.0, + 'bt': 12.16, + 'r': 1.85, + 'a': 0.57, + 'kappa': 1.9, 'delta': 0.5, - 'neped': 1.0, - 'betan': 0.5, - 'zeffped': 1.0, - 'nesep': 0.25, - 'tesep': 1.0, + 'neped': 30.0, + 'betan': 1.0, + 'zeffped': 1.5, + 'nesep': 10.0, + 'tesep': 100.0, }, nproc = 64, cold_start = True, @@ -26,4 +26,5 @@ eped.read(subfolder='run1') + eped.plot(labels=['run1']) From 2a38e6a23636419e04b1d30f9149fd97c0892d1d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 28 May 2025 17:43:10 -0400 Subject: [PATCH 063/385] return True for printMsg that do not provide questions --- src/mitim_tools/misc_tools/LOGtools.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/mitim_tools/misc_tools/LOGtools.py b/src/mitim_tools/misc_tools/LOGtools.py index e3c582ae..4646cfbb 100644 --- a/src/mitim_tools/misc_tools/LOGtools.py +++ b/src/mitim_tools/misc_tools/LOGtools.py @@ -83,6 +83,8 @@ def printMsg(*args, typeMsg=""): if typeMsg == "q": return query_yes_no("\t\t>> Do you want to continue?", extra=extra) + return True # Default return value if no specific typeMsg is provided + if not sys.platform.startswith('win'): import termios From c5fe791e1a680232f1b2c1bf2916f5350c29cd4f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 28 May 2025 17:45:00 -0400 Subject: [PATCH 064/385] source modules before shellPreCommands --- src/mitim_tools/misc_tools/FARMINGtools.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index a21c9fb9..ce60badb 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -141,6 +141,10 @@ def prep( # Pass to class self.command = command + + if not isinstance(self.command, list): + self.command = [self.command] + self.input_files = input_files if isinstance(input_files, list) else [] self.input_folders = input_folders if isinstance(input_folders, list) else [] self.output_files = output_files if isinstance(output_files, list) else [] @@ -164,8 +168,11 @@ def run( removeScratchFolders = False # Always start by going to the folder (inside sbatch file) - command_str_mod = [f"cd {self.folderExecution}", f"{self.command}"] + command_str_mod = [f"cd {self.folderExecution}"] + for command in self.command: + command_str_mod += [command] + # ****** Prepare SLURM job ***************************** comm, fileSBATCH, fileSHELL = create_slurm_execution_files( command_str_mod, @@ -196,12 +203,8 @@ def run( self.output_files = curateOutFiles(self.output_files) # Relative paths - self.input_files = [ - path.relative_to(self.folder_local) for path in self.input_files - ] - self.input_folders = [ - path.relative_to(self.folder_local) for path in self.input_folders - ] + self.input_files = [path.relative_to(self.folder_local) for path in self.input_files] + self.input_folders = [path.relative_to(self.folder_local) for path in self.input_folders] # Process self.full_process( @@ -1135,10 +1138,13 @@ def create_slurm_execution_files( """ commandSHELL = ["#!/usr/bin/env bash"] - commandSHELL.extend(copy.deepcopy(shellPreCommands)) + commandSHELL.append("") if modules_remote is not None: commandSHELL.append(modules_remote) + + commandSHELL.extend(copy.deepcopy(shellPreCommands)) + commandSHELL.append(f"{launch} {fileSBATCH_remote}") commandSHELL.append("") for i in range(len(shellPostCommands)): From b0b24220860ca85709c2d3dd97ce688a2ea0142c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 29 May 2025 17:39:42 -0400 Subject: [PATCH 065/385] Job array and scans working --- src/mitim_tools/eped_tools/EPEDtools.py | 226 +++++++++++++++--------- src/mitim_tools/misc_tools/LOGtools.py | 6 +- tests/EPED_workflow.py | 13 +- 3 files changed, 156 insertions(+), 89 deletions(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index 0ff3d33e..4f8547ff 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -1,5 +1,6 @@ import os import re +import copy import subprocess import matplotlib.pyplot as plt import f90nml @@ -23,53 +24,137 @@ def __init__( self.results = {} - - def case_exists( - self, - subfolder = 'run1' - ): - - # Check if the run folder exists - self.folder_run = self.folder / subfolder - - output_file = self.folder_run / 'output.nc' - - return output_file.exists() and output_file.is_file() - def run( self, subfolder = 'run1', - input_params = None, - nproc = 64, - minutes_slurm = 60, + input_params = None, # {'ip': 12.0, 'bt': 12.16, 'r': 1.85, 'a': 0.57, 'kappa': 1.9, 'delta': 0.5, 'neped': 30.0, 'betan': 1.0, 'zeffped': 1.5, 'nesep': 10.0, 'tesep': 100.0}, + scan_param = None, # {'variable': 'neped', 'values': [10.0, 20.0, 30.0]} + keep_nsep_ratio = None, # Ratio of neped to nesep + nproc_per_run = 64, + minutes_slurm = 30, cold_start = False, ): + # ------------------------------------ + # Prepare job + # ------------------------------------ + + # Prepare folder structure self.folder_run = self.folder / subfolder - if self.case_exists(subfolder): - if cold_start: - res = print(f'\t> Run {subfolder} already exists but cold_start is set to True. Running from scratch.', typeMsg='q') - if res: - IOtools.shutil_rmtree(self.folder_run) - else: - return - else: - print(f'\t> Run {subfolder} already exists and cold_start is set to False. Skipping run.', typeMsg='i') - return - - # Set up folder - - self.folder_run.mkdir(parents=True, exist_ok=True) + # Prepare scan parameters + scan_param_variable = scan_param['variable'] if scan_param is not None else None + scan_param_values = scan_param['values'] if scan_param is not None else [None] + + # Prepare job array setup + job_array = '' + for i in range(len(scan_param_values)): + job_array += f'{i+1}' if i == 0 else f',{i+1}' + + # Initialize Job + self.eped_job = FARMINGtools.mitim_job(self.folder_run) + + self.eped_job.define_machine( + "eped", + "mitim_eped", + slurm_settings={ + 'name': 'mitim_eped', + 'minutes': minutes_slurm, + 'ntasks': nproc_per_run, + 'job_array': job_array + } + ) # ------------------------------------ - # Write input file to EPED + # Prepare each individual case # ------------------------------------ - eped_input_file_suffix = 'eped.input.1' # Do not call it directly 'eped.input' as it may be overwritten by the job script template copying commands + folder_cases, output_files, shellPreCommands = [], [], [] + for i,value in enumerate(scan_param_values): + + # Folder structure + subfolder = f'run{i+1}' + folder_case = self.folder_run / subfolder + + # Prepare input parameters + if scan_param_variable is not None: + input_params_new = input_params.copy() if input_params is not None else {} + input_params_new[scan_param_variable] = value + else: + input_params_new = input_params + + if keep_nsep_ratio is not None: + print(f'\t> Setting nesep to {keep_nsep_ratio} * neped') + input_params_new['nesep'] = keep_nsep_ratio * input_params_new['neped'] + + # ******************************* + # Check if the case should be run + run_case = True + force_res = False + if (self.folder_run / f'output_{subfolder}.nc').exists(): + if cold_start: + res = print(f'\t> Run {subfolder} already exists but cold_start is set to True. Running from scratch.', typeMsg='i' if force_res else 'q') + if res: + IOtools.shutil_rmtree(folder_case) + (self.folder_run / f'output_{subfolder}.nc').unlink(missing_ok=True) + force_res = True + else: + run_case = False + else: + print(f'\t> Run {subfolder} already exists and cold_start is set to False. Skipping run.', typeMsg='i') + run_case = False + + if not run_case: + continue + # ******************************* + + # Set up folder + folder_case.mkdir(parents=True, exist_ok=True) + + # Preparation of the run folder by copying the template files + eped_input_file = 'eped.input.1' + required_files_folder = '$EPED_SOURCE_PATH/template/engaging/eped_run_template' + shellPreCommands.append(f'cp {required_files_folder}/* {self.eped_job.folderExecution}/{subfolder}/. && mv {self.eped_job.folderExecution}/{subfolder}/{eped_input_file} {self.eped_job.folderExecution}/{subfolder}/eped.input') + + # Write input file to EPED, determining the expected output file + output_file = self._prep_input_file(folder_case,input_params=input_params_new,eped_input_file=eped_input_file) + + output_files.append(output_file.as_posix()) + folder_cases.append(folder_case) + + # ------------------------------------- + # Execute + # ------------------------------------- + + # Command to execute by each job in the array + EPEDcommand = f'cd {self.eped_job.folderExecution}/run"$SLURM_ARRAY_TASK_ID" && export NPROC_EPED={nproc_per_run} && ips.py --config=eped.config --platform=psfc_cluster.conf' - eped_input_file = self.folder_run / eped_input_file_suffix + # Prepare the job script + self.eped_job.prep(EPEDcommand,input_folders=folder_cases,output_files=copy.deepcopy(output_files),shellPreCommands=shellPreCommands) + # Run the job + self.eped_job.run() #removeScratchFolders=False) + + # ------------------------------------- + # Postprocessing + # ------------------------------------- + + # Remove potential output files from previous runs + output_files_old = sorted(list(self.folder_run.glob("*.nc"))) + for output_file in output_files_old: + output_file.unlink() + + # Rename output files + for i in range(len(output_files)): + os.system(f'mv {self.folder_run / output_files[i]} {self.folder_run / f"output_run{i+1}.nc"}') + + def _prep_input_file( + self, + folder_case, + input_params = None, + eped_input_file = 'eped.input.1', # Do not call it directly 'eped.input' as it may be overwritten by the job script template copying commands + ): + shot = 0 timeid = 0 @@ -94,44 +179,12 @@ def run( nml = f90nml.Namelist(eped_input) # Write the input file - f90nml.write(nml, eped_input_file, force=True) - - # Initialize Job - self.eped_job = FARMINGtools.mitim_job(self.folder_run) - - self.eped_job.define_machine( - "eped", - "mitim_eped", - slurm_settings={ - 'name': 'mitim_eped', - 'minutes': minutes_slurm, - 'ntasks': nproc, - } - ) - # ------------------------------------- - # Executable commands - # ------------------------------------- + f90nml.write(nml, folder_case / eped_input_file, force=True) - required_files_folder = '$EPED_SOURCE_PATH/template/engaging/eped_run_template' - - EPEDcommand = f'cp {required_files_folder}/* {self.eped_job.folderExecution}/. && mv {self.eped_job.folderExecution}/{eped_input_file_suffix} {self.eped_job.folderExecution}/eped.input && export NPROC_EPED={nproc} && ips.py --config=eped.config --platform=psfc_cluster.conf' - - # ------------------------------------- - # Execute - # ------------------------------------- - - output_file = f'e{shot:06d}.{timeid:05d}' - - self.eped_job.prep( - EPEDcommand, - input_files=[eped_input_file], - output_files=[f'eped/SUMMARY/{output_file}'], - ) + # What's the expected output file? + output_file = folder_case.relative_to(self.folder_run) / 'eped' / 'SUMMARY' / f'e{shot:06d}.{timeid:05d}' - self.eped_job.run(removeScratchFolders=False) - - # Rename output file - os.system(f'mv {self.folder_run / 'eped' / 'SUMMARY' /output_file} {self.folder_run / "output.nc"}') + return output_file def read( self, @@ -139,12 +192,19 @@ def read( label = None, ): - output_file = self.folder / subfolder / 'output.nc' + self.results[label if label is not None else subfolder] = {} + + output_files = sorted(list((self.folder / subfolder).glob("*.nc"))) + + for output_file in output_files: + + + data = xr.open_dataset(f'{output_file.resolve()}', engine='netcdf4') + data = postprocess_eped(data, 'G', 0.03) - data = xr.open_dataset(f'{output_file.resolve()}', engine='netcdf4') - data = postprocess_eped(data, 'G', 0.03) + sublabel = output_file.name.split('_')[-1].split('.')[0] - self.results[label if label is not None else subfolder] = data + self.results[label if label is not None else subfolder][sublabel] = data def plot( self, @@ -159,12 +219,18 @@ def plot( data = self.results[name] - neped = float(data['neped']) - ptop = float(data['ptop']) - wtop = float(data['wptop']) - - axs[0].plot(neped,ptop,'-s', c = colors[i], ms = 12) - axs[1].plot(neped,wtop,'-s', c = colors[i], ms = 12) + neped, ptop, wtop = [], [], [] + for sublabel in data: + neped.append(float(data[sublabel]['neped'])) + if 'ptop' in data[sublabel].data_vars: + ptop.append(float(data[sublabel]['ptop'])) + wtop.append(float(data[sublabel]['wptop'])) + else: + ptop.append(0.0) + wtop.append(0.0) + + axs[0].plot(neped,ptop,'-s', c = colors[i], ms = 10) + axs[1].plot(neped,wtop,'-s', c = colors[i], ms = 10) ax = axs[0] ax.set_xlabel('neped ($10^{19}m^{-3}$)') @@ -178,6 +244,8 @@ def plot( ax.set_ylim(bottom=0) GRAPHICStools.addDenseAxis(ax) + plt.tight_layout() + # ************************************************************************************************************ # ************************************************************************************************************ diff --git a/src/mitim_tools/misc_tools/LOGtools.py b/src/mitim_tools/misc_tools/LOGtools.py index 4646cfbb..48b7ba02 100644 --- a/src/mitim_tools/misc_tools/LOGtools.py +++ b/src/mitim_tools/misc_tools/LOGtools.py @@ -27,7 +27,7 @@ def printMsg(*args, typeMsg=""): verbose = read_verbose_level() if verbose == 0: - return False + return True else: # ----------------------------------------------------------------------------- @@ -57,8 +57,6 @@ def printMsg(*args, typeMsg=""): # Print if typeMsg in ["w"]: print(*total) - # Question result - return False elif verbose == 2: # Print @@ -83,7 +81,7 @@ def printMsg(*args, typeMsg=""): if typeMsg == "q": return query_yes_no("\t\t>> Do you want to continue?", extra=extra) - return True # Default return value if no specific typeMsg is provided + return True # Default return value if no specific typeMsg is provided if not sys.platform.startswith('win'): diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py index bc1680e3..be03ff77 100644 --- a/tests/EPED_workflow.py +++ b/tests/EPED_workflow.py @@ -6,9 +6,9 @@ eped = EPEDtools.EPED(folder=folder) eped.run( - subfolder = 'run1', + subfolder = 'case1', input_params = { - 'ip': 12.0, + 'ip': 8.7, 'bt': 12.16, 'r': 1.85, 'a': 0.57, @@ -20,11 +20,12 @@ 'nesep': 10.0, 'tesep': 100.0, }, - nproc = 64, + scan_param = {'variable': 'neped', 'values': [15.0, 30.0, 45.0, 60.0, 75.0]}, + keep_nsep_ratio = 0.4, + nproc_per_run = 64, cold_start = True, ) -eped.read(subfolder='run1') +eped.read(subfolder='case1') - -eped.plot(labels=['run1']) +eped.plot(labels=['case1']) From 7ce4b0e07605259737ea8d4f928ba6c400014516 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 30 May 2025 13:56:47 -0400 Subject: [PATCH 066/385] If no cases to run, do not farm it --- src/mitim_tools/eped_tools/EPEDtools.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index 4f8547ff..b10463bb 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -122,6 +122,10 @@ def run( output_files.append(output_file.as_posix()) folder_cases.append(folder_case) + # If no cases to run, exit + if len(folder_cases) == 0: + return + # ------------------------------------- # Execute # ------------------------------------- From 1a2cb16209c154ee1399663a18aef1838ca0da75 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 30 May 2025 15:04:36 -0400 Subject: [PATCH 067/385] Added printing of results when reading --- src/mitim_tools/eped_tools/EPEDtools.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index b10463bb..61cd0ed5 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -24,6 +24,8 @@ def __init__( self.results = {} + self.inputs_potential = ['ip', 'bt', 'r', 'a', 'kappa', 'delta', 'neped', 'betan', 'zeffped', 'nesep', 'tesep'] + def run( self, subfolder = 'run1', @@ -193,6 +195,7 @@ def _prep_input_file( def read( self, subfolder = 'run1', + print_results = True, label = None, ): @@ -210,6 +213,26 @@ def read( self.results[label if label is not None else subfolder][sublabel] = data + if print_results: + self.print(label if label is not None else subfolder, sublabel) + + def print(self,label,sublabel): + + print(f'\n\t> EPED results {sublabel}:') + data = self.results[label][sublabel] + + print('\t\t> Inputs:') + for input_param in self.inputs_potential: + print(f'\t\t\t{input_param}: {data[input_param].values[0]}') + + + print('\t\t> Outputs:') + if 'ptop' in data.data_vars: + print(f'\t\t\tptop: {data["ptop"].values[0]:.2f} kPa') + print(f'\t\t\twptop: {data["wptop"].values[0]:.3f} psi_pol') + else: + print('\t\t\tptop: Not available',typeMsg='w') + def plot( self, labels = ['run1'], From a864cffce787ee8b2ea5d23ce5419bf47b71e7c4 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 31 May 2025 21:19:07 -0400 Subject: [PATCH 068/385] Plotting linear of CGYRO from alias --- src/mitim_tools/gacode_tools/scripts/read_cgyro.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 810e087e..22fde9cd 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -11,9 +11,11 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("folders", type=str, nargs="*") + parser.add_argument("--linear", action="store_true", help="linear run") args = parser.parse_args() folders = args.folders + linear = args.linear # Read c = CGYROtools.CGYRO() @@ -23,7 +25,11 @@ def main(): labels.append(f"case {i + 1}") c.read(label=labels[-1], folder=folder) - c.plot(labels=labels) + if linear: + # Plot linear spectrum + c.plotLS(labels=labels) + else: + c.plot(labels=labels) c.fn.show() embed() From d4278046ddb9d693b32e3de4c9f88626983e03d4 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 31 May 2025 21:19:20 -0400 Subject: [PATCH 069/385] Changed default cgyro controls --- templates/input.cgyro.controls | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/templates/input.cgyro.controls b/templates/input.cgyro.controls index 220481f8..20286741 100644 --- a/templates/input.cgyro.controls +++ b/templates/input.cgyro.controls @@ -99,6 +99,9 @@ PRINT_STEP=100 #Max simulation time (units a/c_s) MAX_TIME=100 #900 +#Frequency Tolerance +FREQ_TOL=0.001 + #Number of data outputs before saving a cold_start (DELTA_T*PRINT_STEP is one data output) RESTART_STEP=10 @@ -113,9 +116,9 @@ COLLISION_MODEL=4 #Rotation Scaling of Exp. Values ROTATION_MODEL=2 SHEAR_METHOD=2 -GAMMA_E_SCALE=0.0 -GAMMA_P_SCALE=0.0 -MACH_SCALE=0.0 +GAMMA_E_SCALE=1.0 +GAMMA_P_SCALE=1.0 +MACH_SCALE=1.0 #Scaling of Electron beta and Lambda Debye BETAE_UNIT_SCALE=1.0 From e8f4538dcc90031bbf67d0cdb5460b45492df51c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 31 May 2025 21:19:57 -0400 Subject: [PATCH 070/385] Farming can now remove remote folder even if I do not check later --- src/mitim_tools/misc_tools/FARMINGtools.py | 29 ++++++++++++---------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index ce60badb..ff82275c 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -160,12 +160,17 @@ def run( waitYN=True, timeoutSecs=1e6, removeScratchFolders=True, + removeScratchFolders_goingIn=None, check_if_files_received=True, attempts_execution=1, ): + removeScratchFolders_goingOut = removeScratchFolders + if removeScratchFolders_goingIn is None: + removeScratchFolders_goingIn = removeScratchFolders + if not waitYN: - removeScratchFolders = False + removeScratchFolders_goingOut = False # Always start by going to the folder (inside sbatch file) command_str_mod = [f"cd {self.folderExecution}"] @@ -209,7 +214,8 @@ def run( # Process self.full_process( comm, - removeScratchFolders=removeScratchFolders, + removeScratchFolders_goingIn=removeScratchFolders_goingIn, + removeScratchFolders_goingOut=removeScratchFolders_goingOut, timeoutSecs=timeoutSecs, check_if_files_received=waitYN and check_if_files_received, check_files_in_folder=self.check_files_in_folder, @@ -237,7 +243,8 @@ def full_process( self, comm, timeoutSecs=1e6, - removeScratchFolders=True, + removeScratchFolders_goingIn=True, + removeScratchFolders_goingOut=True, check_if_files_received=True, check_files_in_folder={}, attempts_execution = 1, @@ -256,7 +263,7 @@ def full_process( self.connect(log_file=self.folder_local / "paramiko.log") # ~~~~~~ Prepare scratch folder - if removeScratchFolders: + if removeScratchFolders_goingIn: self.remove_scratch_folder() self.create_scratch_folder() @@ -289,7 +296,7 @@ def full_process( # ~~~~~~ Remove scratch folder if received: - if wait_for_all_commands and removeScratchFolders: + if wait_for_all_commands and removeScratchFolders_goingOut: self.remove_scratch_folder() else: @@ -365,7 +372,7 @@ def define_server(self, disabled_algorithms=None): self.target_host, username=self.target_user, disabled_algorithms=disabled_algorithms, - key_filename=self.key_filename, + key_filename=str(self.key_filename) if self.key_filename is not None else None, port=self.port, sock=self.sock, allow_agent=True, @@ -377,7 +384,7 @@ def define_server(self, disabled_algorithms=None): self.target_host, username=self.target_user, disabled_algorithms=disabled_algorithms, - key_filename=self.key_filename, + key_filename=str(self.key_filename) if self.key_filename is not None else None, port=self.port, sock=self.sock, allow_agent=True, @@ -566,14 +573,10 @@ def execute_local(self, command_str, printYN=False, timeoutSecs=None, **kwargs): return output, error def retrieve(self, check_if_files_received=True, check_files_in_folder={}): - print( - f'\t* Retrieving files{" from remote server" if self.ssh is not None else ""}:' - ) + print(f'\t* Retrieving files{" from remote server" if self.ssh is not None else ""}:') # Create a tarball of the output files & folders on the remote machine - print( - "\t\t- Removing local output files & folders that potentially exist from previous runs" - ) + print("\t\t- Removing local output files & folders that potentially exist from previous runs") for file in self.output_files: (self.folder_local / file).unlink(missing_ok=True) for folder in self.output_folders: From 710f215586086859ce72180d8aa6a0baf9310139 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 31 May 2025 21:20:19 -0400 Subject: [PATCH 071/385] Controls checker for CGYRO --- src/mitim_tools/gacode_tools/utils/GACODEdefaults.py | 6 +++--- src/mitim_tools/gacode_tools/utils/GACODErun.py | 5 +++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index adae8bb2..aa890fcd 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -414,9 +414,9 @@ def convolution_CECE(d_perp_dict, dRdx=1.0): return fun, factorTot_to_Perp -def review_controls(TGLFoptions): +def review_controls(TGLFoptions, control = "input.tglf.controls"): - TGLFoptions_check = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.tglf.controls", caseInsensitive=False) + TGLFoptions_check = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / control, caseInsensitive=False) # Add plasma too potential_flags = ['NS', 'SIGN_BT', 'SIGN_IT', 'VEXB', 'VEXB_SHEAR', 'BETAE', 'XNUE', 'ZEFF', 'DEBYE'] @@ -431,4 +431,4 @@ def review_controls(TGLFoptions): isGeometry = option.split('_')[-1] in ['LOC'] if (not isSpecie) and (not isGeometry) and (option not in TGLFoptions_check): - print(f"\t- TGLF option {option} not in input.tglf.controls, prone to errors", typeMsg="q") + print(f"\t- Option {option} not in {control}, prone to errors", typeMsg="q") diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 8d1d7e16..a0e6eaef 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -102,12 +102,13 @@ def modifyInputs( minimum_delta_abs={}, position_change=0, addControlFunction=None, + control_file = 'input.tglf.controls', **kwargs_to_function, ): # Check that those are valid flags - GACODEdefaults.review_controls(extraOptions) - GACODEdefaults.review_controls(multipliers) + GACODEdefaults.review_controls(extraOptions, control = control_file) + GACODEdefaults.review_controls(multipliers, control = control_file) # ------------------------------------------- if Settings is not None: From bf8f49fc1951a23c087f2e6c2da61ea2f5d5ef2c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 31 May 2025 23:37:42 -0400 Subject: [PATCH 072/385] CGYRO scans working with job arrays --- src/mitim_tools/gacode_tools/CGYROtools.py | 262 ++++++++++++++++++--- 1 file changed, 223 insertions(+), 39 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 96e9792f..16166756 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -1,4 +1,5 @@ import os +from re import sub import shutil import datetime import time @@ -74,19 +75,16 @@ def prep(self, folder, inputgacode_file): self.inputgacode_file = self.folder / "input.gacode" shutil.copy2(IOtools.expandPath(inputgacode_file), self.inputgacode_file) - def run( + def _prerun( self, subFolderCGYRO, roa=0.55, CGYROsettings=None, extraOptions={}, multipliers={}, - test_run=False, - n=16, - nomp=1, ): - self.folderCGYRO = self.folder / f"{subFolderCGYRO}_{roa:.6f}" + self.folderCGYRO = self.folder / Path(subFolderCGYRO) self.folderCGYRO.mkdir(parents=True, exist_ok=True) @@ -109,52 +107,204 @@ def run( multipliers=multipliers, addControlFunction=GACODEdefaults.addCGYROcontrol, rmin=roa, + control_file = 'input.cgyro.controls' ) inputCGYRO.writeCurrentStatus() + return input_cgyro_file, inputgacode_file_this + + def run_test( + self, + subFolderCGYRO, + roa=0.55, + CGYROsettings=None, + extraOptions={}, + multipliers={}, + **kwargs + ): + + if 'scan_param' in kwargs: + print("\t- Cannot run CGYRO tests with scan_param, running just the base",typeMsg="i") + + input_cgyro_file, inputgacode_file_this = self._prerun( + subFolderCGYRO, + roa=roa, + CGYROsettings=CGYROsettings, + extraOptions=extraOptions, + multipliers=multipliers, + ) + self.cgyro_job = FARMINGtools.mitim_job(self.folderCGYRO) - name = f'mitim_cgyro_{subFolderCGYRO}_{roa:.6f}{"_test" if test_run else ""}' + name = f'mitim_cgyro_{subFolderCGYRO}_{roa:.6f}_test' + + self.cgyro_job.define_machine( + "cgyro", + name, + slurm_settings={ + "name": name, + "minutes": 5, + "cpuspertask": 1, + "ntasks": 1, + }, + ) + + CGYROcommand = "cgyro -t ." + + self.cgyro_job.prep( + CGYROcommand, + input_files=[input_cgyro_file, inputgacode_file_this], + output_files=self.output_files_test, + ) + + self.cgyro_job.run() + + def run(self,subFolderCGYRO,test_run=False,**kwargs): if test_run: + self.run_test(subFolderCGYRO,**kwargs) + else: + self.run_full(subFolderCGYRO,**kwargs) + + def run_full( + self, + subFolderCGYRO, + roa=0.55, + CGYROsettings=None, + extraOptions={}, + multipliers={}, + scan_param = None, # {'variable': 'KY', 'values': [0.2,0.3,0.4]} + minutes = 5, + n = 16, + nomp = 1, + submit_via_qsub=True, + clean_folder_going_in=True, # Make sure the scratch folder is removed before running (unless I want a restart!) + ): + + input_cgyro_file, inputgacode_file_this = self._prerun( + subFolderCGYRO, + roa=roa, + CGYROsettings=CGYROsettings, + extraOptions=extraOptions, + multipliers=multipliers, + ) + + self.cgyro_job = FARMINGtools.mitim_job(self.folderCGYRO) + + name = f'mitim_cgyro_{subFolderCGYRO}_{roa:.6f}' + + if scan_param is not None and submit_via_qsub: + raise Exception(" Cannot use scan_param with submit_via_qsub=True, because it requires a different job for each value of the scan parameter.") + + if submit_via_qsub: self.cgyro_job.define_machine( "cgyro", name, - slurm_settings={ - "name": name, - "minutes": 5, - "cpuspertask": 1, - "ntasks": 1, - }, + launchSlurm=False, ) - CGYROcommand = "cgyro -t ." + CGYROcommand = f'gacode_qsub -e . -n {n} -nomp {nomp} -queue {self.cgyro_job.machineSettings['slurm']['partition']} -w 0:{minutes}:00 -s' + + if "account" in self.cgyro_job.machineSettings["slurm"] and self.cgyro_job.machineSettings["slurm"]["account"] is not None: + CGYROcommand += f" -repo {self.cgyro_job.machineSettings['slurm']['account']}" + + self.slurm_output = "batch.out" else: + if scan_param is None: + job_array = None + folder = 'scan0' + scan_param = {'variable': None, 'values': [0]} # Dummy scan parameter to avoid issues with the code below + else: + # Array + job_array = '' + for i,value in enumerate(scan_param['values']): + if job_array != '': + job_array += ',' + job_array += str(i) + + folder = 'scan"$SLURM_ARRAY_TASK_ID"' + + # Machine self.cgyro_job.define_machine( "cgyro", name, - launchSlurm=False, + slurm_settings={ + "name": name, + "minutes": minutes, + "ntasks": n, + "job_array": job_array, + }, ) - if self.cgyro_job.launchSlurm: - CGYROcommand = f'gacode_qsub -e . -n {n} -nomp {nomp} -repo {self.cgyro_job.machineSettings["slurm"]["account"]} -queue {self.cgyro_job.machineSettings["slurm"]["partition"]} -w 0:10:00 -s' - else: + if not self.cgyro_job.launchSlurm: + raise Exception(" Cannot run CGYRO scans without slurm") + + # Command to run cgyro + CGYROcommand = f'cgyro -e {folder} -n {n} -nomp {nomp} -p {self.cgyro_job.folderExecution}' + + # Scans + input_folders = [] + output_folders = [] + for i,value in enumerate(scan_param['values']): + subfolder = f"scan{i}" + folder_run = self.folderCGYRO / subfolder + folder_run.mkdir(parents=True, exist_ok=True) + + # Copy the input.cgyro in the subfolder + input_cgyro_file_this = folder_run / "input.cgyro" + shutil.copy2(input_cgyro_file, input_cgyro_file_this) + + # Modify the input.cgyro file with the scan parameter + extraOptions_this = extraOptions.copy() + if scan_param['variable'] is not None: + extraOptions_this[scan_param['variable']] = value + inputCGYRO = CGYROinput(file=input_cgyro_file_this) + input_cgyro_file_this = GACODErun.modifyInputs( + inputCGYRO, + Settings=CGYROsettings, + extraOptions=extraOptions_this, + multipliers=multipliers, + addControlFunction=GACODEdefaults.addCGYROcontrol, + rmin=roa, + control_file = 'input.cgyro.controls' + ) + + input_cgyro_file_this.writeCurrentStatus() - CGYROcommand = f"cgyro -e . -n {n} -nomp {nomp}" + # Copy the input.gacode file in the subfolder + inputgacode_file_this = folder_run / "input.gacode" + shutil.copy2(self.inputgacode_file, inputgacode_file_this) + # Prepare the input and output folders + input_folders.append(folder_run) + output_folders.append(subfolder) + + self.slurm_output = "slurm_output.dat" + + # First submit the job with gacode_qsub, which will submit the cgyro job via slurm, with name self.cgyro_job.prep( CGYROcommand, - input_files=[input_cgyro_file, inputgacode_file_this], - output_files=self.output_files if not test_run else self.output_files_test, - ) + input_folders = input_folders, + output_folders=output_folders, + ) self.cgyro_job.run( - waitYN=not self.cgyro_job.launchSlurm - ) # ,removeScratchFolders=False) + waitYN=False, + check_if_files_received=False, + removeScratchFolders=False, + removeScratchFolders_goingIn=clean_folder_going_in, + ) + + # Prepare how to search for the job without waiting for it + name_default_submission_qsub = Path(self.cgyro_job.folderExecution).name + + self.cgyro_job.launchSlurm = True + self.cgyro_job.slurm_settings['name'] = name_default_submission_qsub + def check(self, every_n_minutes=5): @@ -162,7 +312,7 @@ def check(self, every_n_minutes=5): print("- Checker job status") while True: - self.cgyro_job.check() + self.cgyro_job.check(file_output = self.slurm_output) print( f'\t- Current status (as of {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}): {self.cgyro_job.status} ({self.cgyro_job.infoSLURM["STATE"]})' ) @@ -176,11 +326,13 @@ def check(self, every_n_minutes=5): print("\t- Job considered finished") - def get(self): + def fetch(self): """ For a job that has been submitted but not waited for, once it is done, get the results """ + print("\n\n\t- Fetching results") + if self.cgyro_job.launchSlurm: self.cgyro_job.connect() self.cgyro_job.retrieve() @@ -188,6 +340,19 @@ def get(self): else: print("- Not retrieving results because this was run command line (not slurm)") + def delete(self): + + print("\n\n\t- Deleting job") + + self.cgyro_job.launchSlurm = False + + self.cgyro_job.prep( + f"scancel -n {self.cgyro_job.slurm_settings['name']}", + label_log_files="_finish", + ) + + self.cgyro_job.run() + # --------------------------------------------------------------------------------------------------------- # Reading and plotting # --------------------------------------------------------------------------------------------------------- @@ -196,18 +361,37 @@ def read(self, label="cgyro1", folder=None, tmin = 0.0): folder = IOtools.expandPath(folder) if folder is not None else self.folderCGYRO - original_dir = os.getcwd() + folders = sorted(list((folder).glob("scan*"))) + + if len(folders) == 0: + folders = [folder] + else: + print(f"\t- Found {len(folders)} scan folders in {folder.resolve()}:") + for f in folders: + print(f"\t\t- {f.name}") + + for folder in folders: + + original_dir = os.getcwd() + + try: + print(f"\t- Reading CGYRO data from {folder.resolve()}") + self.results[f'{label}_{folder.name}'] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") + except: + if print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): + os.chdir(folder) + os.system("cgyro -t") + self.results[f'{label}_{folder.name}'] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") + + os.chdir(original_dir) - try: - print(f"\t- Reading CGYRO data from {folder.resolve()}") - self.results[label] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") - except: - if print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): - os.chdir(folder) - os.system("cgyro -t") - self.results[label] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") + try: + self._postprocess(f'{label}_{folder.name}', folder, tmin=tmin) + except Exception as e: + print(f"\t- Error during postprocessing: {e}") + print("\t- Skipping postprocessing, results may be incomplete or linear run") - os.chdir(original_dir) + def _postprocess(self, label, folder, tmin=0.0): # Extra postprocessing self.results[label].electron_flag = np.where(self.results[label].z == -1)[0][0] @@ -351,12 +535,12 @@ def plotLS(self, labels=["cgyro1"], fig=None): from mitim_tools.misc_tools.GUItools import FigureNotebook - self.fnLS = FigureNotebook( + self.fn = FigureNotebook( "Linear CGYRO Notebook", geometry="1600x1000", ) - fig1 = self.fnLS.add_figure(label="Linear Stability") - fig2 = self.fnLS.add_figure(label="Ballooning") + fig1 = self.fn.add_figure(label="Linear Stability") + fig2 = self.fn.add_figure(label="Ballooning") grid = plt.GridSpec(2, 2, hspace=0.3, wspace=0.3) ax00 = fig1.add_subplot(grid[0, 0]) @@ -430,7 +614,7 @@ def plotLS(self, labels=["cgyro1"], fig=None): ax.set_xlim(left=0) ax = ax01 - ax.set_xlim([5e-2, 50.0]) + #ax.set_xlim([5e-2, 50.0]) grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) ax00 = fig2.add_subplot(grid[0, 0]) From 937a241c655869e0a50e45a2e1a594096355c499 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 2 Jun 2025 08:27:13 -0400 Subject: [PATCH 073/385] better standalone CGYRO --- src/mitim_tools/gacode_tools/CGYROtools.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 16166756..62fda53a 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -365,28 +365,35 @@ def read(self, label="cgyro1", folder=None, tmin = 0.0): if len(folders) == 0: folders = [folder] + attach_name = False else: print(f"\t- Found {len(folders)} scan folders in {folder.resolve()}:") for f in folders: print(f"\t\t- {f.name}") + attach_name = True for folder in folders: original_dir = os.getcwd() + if attach_name: + label_new = f"{label}_{folder.name}" + else: + label_new = label + try: print(f"\t- Reading CGYRO data from {folder.resolve()}") - self.results[f'{label}_{folder.name}'] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") + self.results[label_new] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") except: if print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): os.chdir(folder) os.system("cgyro -t") - self.results[f'{label}_{folder.name}'] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") + self.results[label_new] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") os.chdir(original_dir) try: - self._postprocess(f'{label}_{folder.name}', folder, tmin=tmin) + self._postprocess(label_new, folder, tmin=tmin) except Exception as e: print(f"\t- Error during postprocessing: {e}") print("\t- Skipping postprocessing, results may be incomplete or linear run") From 688ee1f85b8d573471debf7ce62612df61f08ee2 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 5 Jun 2025 10:00:55 +0200 Subject: [PATCH 074/385] misc --- src/mitim_tools/gacode_tools/CGYROtools.py | 22 +++++++++++++-- tests/CGYRO_workflow.py | 33 ++++++++-------------- 2 files changed, 31 insertions(+), 24 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 62fda53a..37e82741 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -178,7 +178,7 @@ def run_full( minutes = 5, n = 16, nomp = 1, - submit_via_qsub=True, + submit_via_qsub=True, #TODO fix this, works only at NERSC? no scans? clean_folder_going_in=True, # Make sure the scratch folder is removed before running (unless I want a restart!) ): @@ -205,13 +205,31 @@ def run_full( launchSlurm=False, ) - CGYROcommand = f'gacode_qsub -e . -n {n} -nomp {nomp} -queue {self.cgyro_job.machineSettings['slurm']['partition']} -w 0:{minutes}:00 -s' + subfolder = "scan0" + + CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} -queue {self.cgyro_job.machineSettings['slurm']['partition']} -w 0:{minutes}:00 -s' if "account" in self.cgyro_job.machineSettings["slurm"] and self.cgyro_job.machineSettings["slurm"]["account"] is not None: CGYROcommand += f" -repo {self.cgyro_job.machineSettings['slurm']['account']}" self.slurm_output = "batch.out" + # --- + folder_run = self.folderCGYRO / subfolder + folder_run.mkdir(parents=True, exist_ok=True) + + # Copy the input.cgyro in the subfolder + input_cgyro_file_this = folder_run / "input.cgyro" + shutil.copy2(input_cgyro_file, input_cgyro_file_this) + + # Copy the input.gacode file in the subfolder + inputgacode_file_this = folder_run / "input.gacode" + shutil.copy2(self.inputgacode_file, inputgacode_file_this) + + # Prepare the input and output folders + input_folders = [folder_run] + output_folders = [subfolder] + else: if scan_param is None: diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index 557f085f..8ec4be0e 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -5,7 +5,7 @@ cold_start = True gacode_file = __mitimroot__ / "tests" / "data" / "input.gacode" -folder = __mitimroot__ / "tests" / "scratch" / "cgyro_test" +folder = __mitimroot__ / "tests" / "scratch" / "cgyro_test2" if cold_start and folder.exists(): os.system(f"rm -r {folder}") @@ -21,28 +21,17 @@ roa = 0.55, CGYROsettings=0, extraOptions={ - 'KY':0.3 - }) -cgyro.read(label="cgyro1") + 'KY':0.3, + 'MAX_TIME': 1E1, # Short, I just want to test the run + }, + submit_via_qsub=True # NERSC: True #TODO change this + ) -cgyro.run( - 'linear', - roa = 0.55, - CGYROsettings=0, - extraOptions={ - 'KY':0.5 - }) -cgyro.read(label="cgyro2") - -cgyro.run( - 'linear', - roa = 0.55, - CGYROsettings=0, - extraOptions={ - 'KY':0.7 - }) -cgyro.read(label="cgyro3") +cgyro.check(every_n_minutes=1) +cgyro.fetch() +cgyro.delete() +cgyro.read(label="cgyro1") cgyro.plotLS() -cgyro.fnLS.show() +cgyro.fn.show() From c9b8dcfdc489fa8cd8e660e4834aa43712047381 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 11 Jun 2025 15:07:51 +0200 Subject: [PATCH 075/385] Furthering CGYRO package --- src/mitim_tools/gacode_tools/CGYROtools.py | 24 +++++++++++++++------- templates/input.cgyro.controls | 5 +++++ 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 37e82741..aff78516 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -175,11 +175,13 @@ def run_full( extraOptions={}, multipliers={}, scan_param = None, # {'variable': 'KY', 'values': [0.2,0.3,0.4]} + enforce_equality = None, # e.g. {'DLNTDR_SCALE_2': 'DLNTDR_SCALE_1', 'DLNTDR_SCALE_3': 'DLNTDR_SCALE_1'} minutes = 5, n = 16, nomp = 1, submit_via_qsub=True, #TODO fix this, works only at NERSC? no scans? clean_folder_going_in=True, # Make sure the scratch folder is removed before running (unless I want a restart!) + submit_run=True, # False if I just want to check and fetch the job that was already submitted (e.g. via qsub or slurm) ): input_cgyro_file, inputgacode_file_this = self._prerun( @@ -280,6 +282,13 @@ def run_full( extraOptions_this = extraOptions.copy() if scan_param['variable'] is not None: extraOptions_this[scan_param['variable']] = value + + + # If there is an enforce_equality, apply it + if enforce_equality is not None: + for key in enforce_equality: + extraOptions_this[key] = extraOptions_this[enforce_equality[key]] + inputCGYRO = CGYROinput(file=input_cgyro_file_this) input_cgyro_file_this = GACODErun.modifyInputs( inputCGYRO, @@ -310,12 +319,13 @@ def run_full( output_folders=output_folders, ) - self.cgyro_job.run( - waitYN=False, - check_if_files_received=False, - removeScratchFolders=False, - removeScratchFolders_goingIn=clean_folder_going_in, - ) + if submit_run: + self.cgyro_job.run( + waitYN=False, + check_if_files_received=False, + removeScratchFolders=False, + removeScratchFolders_goingIn=clean_folder_going_in, + ) # Prepare how to search for the job without waiting for it name_default_submission_qsub = Path(self.cgyro_job.folderExecution).name @@ -342,7 +352,7 @@ def check(self, every_n_minutes=5): else: print("- Not checking status because this was run command line (not slurm)") - print("\t- Job considered finished") + print("\n\t* Job considered finished",typeMsg="i") def fetch(self): """ diff --git a/templates/input.cgyro.controls b/templates/input.cgyro.controls index 20286741..e8e84341 100644 --- a/templates/input.cgyro.controls +++ b/templates/input.cgyro.controls @@ -135,3 +135,8 @@ N_SPECIES=2 EXCH_FLAG=1 +DLNTDR_SCALE_1 = 1.0 +DLNTDR_SCALE_2 = 1.0 +DLNTDR_SCALE_3 = 1.0 +DLNTDR_SCALE_4 = 1.0 +DLNTDR_SCALE_5 = 1.0 \ No newline at end of file From 1719adf25a0f42c041f616761a8922ddb5039bee Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 11 Jun 2025 15:07:57 +0200 Subject: [PATCH 076/385] misc --- src/mitim_modules/powertorch/utils/CALCtools.py | 2 +- src/mitim_tools/gacode_tools/utils/GACODErun.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/powertorch/utils/CALCtools.py b/src/mitim_modules/powertorch/utils/CALCtools.py index 93b0b58c..980869cc 100644 --- a/src/mitim_modules/powertorch/utils/CALCtools.py +++ b/src/mitim_modules/powertorch/utils/CALCtools.py @@ -85,7 +85,7 @@ def integrateQuadPoly(r, s, p=None): Computes int(s*dr), so if s is s*dV/dr, then int(s*dV), which is the full integral From tgyro_volume_int.f90 - r - minor raidus + r - minor radius s - s*volp (Modified to avoid if statements and for loops) diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index a0e6eaef..b40baf4f 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -152,11 +152,11 @@ def modifyInputs( var_orig = input_class.controls[ikey] var_new = value_to_change_to input_class.controls[ikey] = var_new - elif ikey in input_class.geom: + elif 'geom' in input_class.__dict__ and ikey in input_class.geom: var_orig = input_class.geom[ikey] var_new = value_to_change_to input_class.geom[ikey] = var_new - elif ikey in input_class.plasma: + elif 'plasma' in input_class.__dict__ and ikey in input_class.plasma: var_orig = input_class.plasma[ikey] var_new = value_to_change_to input_class.plasma[ikey] = var_new From b51509532f7c713d0f6f177b73b0e0c7072669f5 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Fri, 6 Jun 2025 10:23:11 -0400 Subject: [PATCH 077/385] don't include partition info if not specified in config --- src/mitim_tools/gacode_tools/CGYROtools.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index aff78516..37ce636d 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -208,8 +208,9 @@ def run_full( ) subfolder = "scan0" + queue = "-queue " + self.cgyro_job.machineSettings['slurm']['partition'] if "partition" in self.cgyro_job.machineSettings['slurm'] else "" - CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} -queue {self.cgyro_job.machineSettings['slurm']['partition']} -w 0:{minutes}:00 -s' + CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} {queue} -w 0:{minutes}:00 -s' if "account" in self.cgyro_job.machineSettings["slurm"] and self.cgyro_job.machineSettings["slurm"]["account"] is not None: CGYROcommand += f" -repo {self.cgyro_job.machineSettings['slurm']['account']}" From e2451cc5b43c62f3bcf73c3983afed36ee177a9e Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Fri, 6 Jun 2025 10:44:16 -0400 Subject: [PATCH 078/385] fix None partition in farmingtools --- src/mitim_tools/misc_tools/FARMINGtools.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index ff82275c..c9ee8ad0 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -1067,7 +1067,8 @@ def create_slurm_execution_files( commandSBATCH.append("#SBATCH --mail-user=" + email) # ******* Partition / Billing - commandSBATCH.append(f"#SBATCH --partition {partition}") + if partition is not None: + commandSBATCH.append(f"#SBATCH --partition {partition}") if account is not None: commandSBATCH.append(f"#SBATCH --account {account}") From aac9c6516909d85e1a6643c7944a495b4309c985 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 14:39:13 +0200 Subject: [PATCH 079/385] Bug fix pathlib plasmastate --- .../transp_tools/utils/PLASMASTATEtools.py | 32 ++++++------------- 1 file changed, 10 insertions(+), 22 deletions(-) diff --git a/src/mitim_tools/transp_tools/utils/PLASMASTATEtools.py b/src/mitim_tools/transp_tools/utils/PLASMASTATEtools.py index af29a8ca..2c9da34a 100644 --- a/src/mitim_tools/transp_tools/utils/PLASMASTATEtools.py +++ b/src/mitim_tools/transp_tools/utils/PLASMASTATEtools.py @@ -1,15 +1,9 @@ import copy import numpy as np +import matplotlib.pyplot as plt from collections import OrderedDict - -try: - import xarray as xr -except: - pass -try: - from IPython import embed -except: - pass +import xarray as xr +from IPython import embed """ This set of routines is used to create a Plasmastate class by reading a standard @@ -54,18 +48,14 @@ def modify_default( print(f"\t- Modifying {self.CDFfile} Plasmastate file...") - self.lumpChargeStates(self.CDFfile_new + "_1") + self.lumpChargeStates( self.CDFfile_new.with_name(self.CDFfile_new.name + '_1')) + try: - self.removeExtraFusionIons( - self.CDFfile_new + "_2", speciesNames=RemoveFusionIons - ) + self.removeExtraFusionIons( self.CDFfile_new.with_name(self.CDFfile_new.name + '_2'), speciesNames=RemoveFusionIons) except: - print( - " --> I could not remove extra fusion ions. Probably because TRANSP was run without fusion reactions", - ) - self.removeExtraTHERMALIons( - self.CDFfile_new + "_3", speciesNames=RemoveTHERMALIons - ) + print(" --> I could not remove extra fusion ions. Probably because TRANSP was run without fusion reactions",) + + self.removeExtraTHERMALIons( self.CDFfile_new.with_name(self.CDFfile_new.name + '_3'), speciesNames=RemoveTHERMALIons) self.addShotNumber(self.CDFfile_new, shotNumber) @@ -129,9 +119,7 @@ def lumpChargeStates(self, fileNew): self.plasma = xr.Dataset(NewData) else: - print( - f"Charge States for impurity {impName} could not be found...", - ) + print(f"Charge States for impurity {impName} could not be found...",) # ------------------------------------------------------------------ # Writting New PLASMASTATE From 67438359faa25cbac2bab6737ae1ecc6de5b1924 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 14:39:51 +0200 Subject: [PATCH 080/385] Creation of parent state class with profiles inheriting from it --- .../maestro/utils/MAESTROplot.py | 3 +- .../portals/utils/PORTALSplot.py | 3 +- .../powertorch/utils/POWERplot.py | 4 +- .../powertorch/utils/TRANSFORMtools.py | 4 +- src/mitim_tools/gacode_tools/NEOtools.py | 2 +- src/mitim_tools/gacode_tools/PROFILEStools.py | 4358 +---------------- src/mitim_tools/gacode_tools/TGLFtools.py | 6 +- src/mitim_tools/gacode_tools/TGYROtools.py | 26 +- .../gacode_tools/scripts/read_gacode.py | 4 +- .../gacode_tools/utils/GACODErun.py | 2 +- src/mitim_tools/gs_tools/GEQtools.py | 4 +- .../plasmastate_tools/MITIMstate.py | 4259 ++++++++++++++++ src/mitim_tools/transp_tools/CDFtools.py | 3 +- .../transp_tools/utils/TRANSPhelpers.py | 3 +- 14 files changed, 4346 insertions(+), 4335 deletions(-) create mode 100644 src/mitim_tools/plasmastate_tools/MITIMstate.py diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index 321eb904..708e396e 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -2,6 +2,7 @@ from collections import OrderedDict from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.misc_tools import LOGtools, GRAPHICStools +from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.gs_tools import GEQtools from pathlib import Path from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -99,7 +100,7 @@ def plot_results(self, fn): figs = PROFILEStools.add_figures(fn,fnlab_pre = 'MAESTRO - ') log_file = self.folder_logs/'plot_maestro.log' if (not self.terminal_outputs) else None with LOGtools.conditional_log_to_file(log_file=log_file): - PROFILEStools.plotAll(ps[-maxPlot:], extralabs=ps_lab[-maxPlot:], figs=figs) + MITIMstate.plotAll(ps[-maxPlot:], extralabs=ps_lab[-maxPlot:], figs=figs) for p,pl in zip(ps,ps_lab): p.printInfo(label = pl) diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index fcae089e..e82ea97b 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -7,6 +7,7 @@ from mitim_modules.portals import PORTALStools from mitim_modules.powertorch import STATEtools from mitim_modules.powertorch.utils import POWERplot +from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -1829,7 +1830,7 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): figs = PROFILEStools.add_figures(fn,fnlab_pre = "PROFILES - ") if indecesPlot[0] < len(self.powerstates): - _ = PROFILEStools.plotAll( + _ = MITIMstate.plotAll( [ self.powerstates[indecesPlot[1]].profiles, self.powerstates[indecesPlot[0]].profiles, diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index 436c3c6c..6d87f477 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -1,4 +1,4 @@ -from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.misc_tools import GRAPHICStools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -15,7 +15,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co profiles_new = self.from_powerstate(insert_highres_powers=True) # Plot the inserted profiles together with the original ones - _ = PROFILEStools.plotAll([self.profiles, profiles_new], figs=figs) + _ = MITIMstate.plotAll([self.profiles, profiles_new], figs=figs) # ----------------------------------------------------------------------------------------------------------- # ---- Plot plasma state diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 68c283a8..2b85ac23 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -4,7 +4,7 @@ import numpy as np import pandas as pd from mitim_tools.misc_tools import LOGtools, IOtools -from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.plasmastate_tools import MITIMstate from mitim_modules.powertorch.physics_models import targets_analytic, parameterizers from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __mitimroot__ @@ -548,7 +548,7 @@ def debug_transformation(p, p_new, s): print(f'Profile mean error: {np.mean(err_prof):.2f}%', typeMsg='i' if np.mean(err_prof) < 1e-0 else 'w') print(f'Gradient mean error (ignoring 0.0): {np.mean(err_grad):.2f}%', typeMsg='i' if np.mean(err_grad) < 1e-0 else 'w') - fn = PROFILEStools.plotAll([p,p_new],extralabs=['Original','New'],lastRhoGradients=rho[-1].item()+0.01) + fn = MITIMstate.plotAll([p,p_new],extralabs=['Original','New'],lastRhoGradients=rho[-1].item()+0.01) axs = fn.figure_handles[3].figure.axes diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index d09d4ed6..8a407662 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -1,5 +1,4 @@ from mitim_tools.misc_tools import IOtools -from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.gacode_tools.utils import GACODErun from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -67,6 +66,7 @@ def run_vgen(self, subfolder="vgen1", vgenOptions={}, cold_start=False): # ---- Postprocess + from mitim_tools.gacode_tools import PROFILEStools self.inputgacode_vgen = PROFILEStools.PROFILES_GACODE( file_new, calculateDerived=True, mi_ref=self.inputgacode.mi_ref ) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 6967e4e4..d11d8f59 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -1,27 +1,24 @@ import copy -import torch -import csv import numpy as np -import matplotlib.pyplot as plt from collections import OrderedDict -from mitim_tools.misc_tools import GRAPHICStools, MATHtools, PLASMAtools, IOtools -from mitim_modules.powertorch.utils import CALCtools -from mitim_tools.gs_tools import GEQtools -from mitim_tools.gacode_tools import NEOtools -from mitim_tools.gacode_tools.utils import GACODEdefaults, GEOMETRYtools -from mitim_tools.transp_tools import CDFtools -from mitim_tools.transp_tools.utils import TRANSPhelpers +from mitim_tools.plasmastate_tools.MITIMstate import mitim_state from mitim_tools.misc_tools.LOGtools import printMsg as print -from mitim_tools import __version__ from IPython import embed -# ------------------------------------------------------------------------------------- -# input.gacode -# ------------------------------------------------------------------------------------- +class PROFILES_GACODE(mitim_state): + ''' + Class to read and manipulate GACODE profiles files (input.gacode). + It inherits from the main MITIMstate class, which provides basic + functionality for plasma state management. + The class reads the GACODE profiles file, extracts relevant data, + and writes them in the way that MITIMstate class expects. + ''' -class PROFILES_GACODE: def __init__(self, file, calculateDerived=True, mi_ref=None): + + super().__init__(type_file='input.gacode') + """ Depending on resolution, derived can be expensive, so I mmay not do it every time """ @@ -46,191 +43,20 @@ def __init__(self, file, calculateDerived=True, mi_ref=None): self.lines = f.readlines() # Read file and store raw data - self.readHeader() - self.readProfiles() - - # Process - self.process(mi_ref=mi_ref, calculateDerived=calculateDerived) - - def process(self, mi_ref=None, calculateDerived=True): - """ - Perform MITIM derivations (can be expensive, only if requested) - Note: One can force what mi_ref to use (in a.m.u.). This is because, by default, MITIM - will use the mass of the first thermal ion to produce quantities such as Q_GB, rho_s, etc. - However, in some ocasions (like when running TGLF), the normalization that must be used - for those quantities is a fixed one (e.g. Deuterium) - """ - - # Calculate necessary quantities - - if "qpar_beam(MW/m^3)" in self.profiles: - self.varqpar, self.varqpar2 = "qpar_beam(MW/m^3)", "qpar_wall(MW/m^3)" - else: - self.varqpar, self.varqpar2 = "qpar_beam(1/m^3/s)", "qpar_wall(1/m^3/s)" - - if "qmom(Nm)" in self.profiles: - self.varqmom = "qmom(Nm)" # Old, wrong one. But Candy fixed it as of 02/24/2023 - else: - self.varqmom = "qmom(N/m^2)" # CORRECT ONE - - # ------------------------------------------------------------------------------------------------------------------- - # Insert zeros in those cases whose column are not there - # ------------------------------------------------------------------------------------------------------------------- - - some_times_are_not_here = [ - "qei(MW/m^3)", - "qohme(MW/m^3)", - "johm(MA/m^2)", - "jbs(MA/m^2)", - "jbstor(MA/m^2)", - "w0(rad/s)", - "ptot(Pa)", # e.g. if I haven't written that info from ASTRA - "zeta(-)", # e.g. if TGYRO is run with zeta=0, it won't write this column in .new - "zmag(m)", - "qsync(MW/m^3)", - "qbrem(MW/m^3)", - "qline(MW/m^3)", - self.varqpar, - self.varqpar2, - "shape_cos0(-)", - self.varqmom, - ] - - num_moments = 6 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros - for i in range(num_moments): - some_times_are_not_here.append(f"shape_cos{i + 1}(-)") - if i > 1: - some_times_are_not_here.append(f"shape_sin{i + 1}(-)") - - for ikey in some_times_are_not_here: - if ikey not in self.profiles.keys(): - self.profiles[ikey] = copy.deepcopy(self.profiles["rmin(m)"]) * 0.0 - - self.deriveQuantities(mi_ref=mi_ref, calculateDerived=calculateDerived) - - def deriveQuantities(self, mi_ref=None, calculateDerived=True, n_theta_geo=1001, rederiveGeometry=True): - - # ------------------------------------------------------------------------------------------------------------------- - self.readSpecies() - self.produce_shape_lists() - self.mi_first = self.Species[0]["A"] - self.DTplasma() - self.sumFast() - # ------------------------------------- - - if "derived" not in self.__dict__: - self.derived = {} - - if mi_ref is not None: - self.derived["mi_ref"] = mi_ref - print(f"\t* Reference mass ({self.derived['mi_ref']:.2f}) to use was forced by class initialization",typeMsg="w") - else: - self.derived["mi_ref"] = self.mi_first - print(f"\t* Reference mass ({self.derived['mi_ref']}) from first ion",typeMsg="i") - - # Useful to have gradients in the basic ---------------------------------------------------------- - self.derived["aLTe"] = aLT(self.profiles["rmin(m)"], self.profiles["te(keV)"]) - self.derived["aLne"] = aLT( - self.profiles["rmin(m)"], self.profiles["ne(10^19/m^3)"] - ) - - self.derived["aLTi"] = self.profiles["ti(keV)"] * 0.0 - self.derived["aLni"] = [] - for i in range(self.profiles["ti(keV)"].shape[1]): - self.derived["aLTi"][:, i] = aLT( - self.profiles["rmin(m)"], self.profiles["ti(keV)"][:, i] - ) - self.derived["aLni"].append( - aLT(self.profiles["rmin(m)"], self.profiles["ni(10^19/m^3)"][:, i]) - ) - self.derived["aLni"] = np.transpose(np.array(self.derived["aLni"])) - # ------------------------------------------------------------------------------------------------ - - if calculateDerived: - self.deriveQuantities_full(rederiveGeometry=rederiveGeometry) - - # ------------------------------------------------------------------------------------- - # Method to write a scratch file - # ------------------------------------------------------------------------------------- - - @classmethod - def scratch(cls, profiles, label_header='', **kwargs_process): - instance = cls(None) - - # Header - instance.header = f''' -# Created from scratch with MITIM version {__version__} -# {label_header} -# -''' - # Add data to profiles - instance.profiles = profiles - - instance.process(**kwargs_process) + self._read_header() + self._read_profiles() - return instance + self._ensure_existence() - # ------------------------------------------------------------------------------------- + self.deriveQuantities(mi_ref=mi_ref, calculateDerived=calculateDerived) - def calculate_Er( - self, - folder, - rhos=None, - vgenOptions={}, - name="vgen1", - includeAll=False, - write_new_file=None, - cold_start=False, - ): - 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, cold_start=cold_start) - - 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)", - ] - - 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] - - self.deriveQuantities() - - if write_new_file is not None: - self.writeCurrentStatus(file=write_new_file) - - # ***************** - - def readHeader(self): + def _read_header(self): for i in range(len(self.lines)): if "# nexp" in self.lines[i]: istartProfs = i self.header = self.lines[:istartProfs] - def readProfiles(self): + def _read_profiles(self): singleLine, title, var = None, None, None # for ruff complaints # --- @@ -289,4123 +115,49 @@ def readProfiles(self): self.profiles["w0(rad/s)"] = self.profiles["omega0(rad/s)"] del self.profiles["omega0(rad/s)"] + def _ensure_existence(self): + # Calculate necessary quantities - def produce_shape_lists(self): - self.shape_cos = [ - self.profiles["shape_cos0(-)"], # tilt - self.profiles["shape_cos1(-)"], - self.profiles["shape_cos2(-)"], - self.profiles["shape_cos3(-)"], - self.profiles["shape_cos4(-)"], - self.profiles["shape_cos5(-)"], - self.profiles["shape_cos6(-)"], - ] - self.shape_sin = [ - None, - None, # s1 is arcsin(triangularity) - None, # s2 is minus squareness - self.profiles["shape_sin3(-)"], - self.profiles["shape_sin4(-)"], - self.profiles["shape_sin5(-)"], - self.profiles["shape_sin6(-)"], - ] - - def readSpecies(self, maxSpecies=100): - maxSpecies = int(self.profiles["nion"][0]) - - Species = [] - for j in range(maxSpecies): - # To determine later if this specie has zero density - niT = self.profiles["ni(10^19/m^3)"][0, j] - - sp = { - "N": self.profiles["name"][j], - "Z": float(self.profiles["z"][j]), - "A": float(self.profiles["mass"][j]), - "S": self.profiles["type"][j].split("[")[-1].split("]")[0], - "n0": niT, - } - - Species.append(sp) - - self.Species = Species - - def sumFast(self): - self.nFast = self.profiles["ne(10^19/m^3)"] * 0.0 - self.nZFast = self.profiles["ne(10^19/m^3)"] * 0.0 - self.nThermal = self.profiles["ne(10^19/m^3)"] * 0.0 - self.nZThermal = self.profiles["ne(10^19/m^3)"] * 0.0 - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "fast": - self.nFast += self.profiles["ni(10^19/m^3)"][:, sp] - self.nZFast += ( - self.profiles["ni(10^19/m^3)"][:, sp] * self.profiles["z"][sp] - ) - else: - self.nThermal += self.profiles["ni(10^19/m^3)"][:, sp] - self.nZThermal += ( - self.profiles["ni(10^19/m^3)"][:, sp] * self.profiles["z"][sp] - ) - - def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry=True): - """ - deriving geometry is expensive, so if I'm just updating profiles it may not be needed - """ - - self.varqmom = "qmom(N/m^2)" - if self.varqmom not in self.profiles: - self.profiles[self.varqmom] = self.profiles["rho(-)"] * 0.0 - - if "derived" not in self.__dict__: - self.derived = {} - - # --------------------------------------------------------------------------------------------------------------------- - # --------- MAIN (useful for STATEtools) - # --------------------------------------------------------------------------------------------------------------------- - - self.derived["a"] = self.profiles["rmin(m)"][-1] - # self.derived['epsX'] = self.profiles['rmaj(m)'] / self.profiles['rmin(m)'] - # self.derived['eps'] = self.derived['epsX'][-1] - self.derived["eps"] = self.profiles["rmin(m)"][-1] / self.profiles["rmaj(m)"][-1] - - self.derived["roa"] = self.profiles["rmin(m)"] / self.derived["a"] - self.derived["Rmajoa"] = self.profiles["rmaj(m)"] / self.derived["a"] - self.derived["Zmagoa"] = self.profiles["zmag(m)"] / self.derived["a"] - - self.derived["torflux"] = ( - float(self.profiles["torfluxa(Wb/radian)"][0]) - * 2 - * np.pi - * self.profiles["rho(-)"] ** 2 - ) # Wb - self.derived["B_unit"] = PLASMAtools.Bunit( - self.derived["torflux"], self.profiles["rmin(m)"] - ) - - self.derived["psi_pol_n"] = ( - self.profiles["polflux(Wb/radian)"] - self.profiles["polflux(Wb/radian)"][0] - ) / ( - self.profiles["polflux(Wb/radian)"][-1] - - self.profiles["polflux(Wb/radian)"][0] - ) - self.derived["rho_pol"] = self.derived["psi_pol_n"] ** 0.5 - - self.derived["q95"] = np.interp( - 0.95, self.derived["psi_pol_n"], self.profiles["q(-)"] - ) - - self.derived["q0"] = self.profiles["q(-)"][0] - - if self.profiles["q(-)"].min() > 1.0: - self.derived["rho_saw"] = np.nan + if "qpar_beam(MW/m^3)" in self.profiles: + self.varqpar, self.varqpar2 = "qpar_beam(MW/m^3)", "qpar_wall(MW/m^3)" else: - self.derived["rho_saw"] = np.interp( - 1.0, self.profiles["q(-)"], self.profiles["rho(-)"] - ) - - # --------- Geometry (only if it doesn't exist or if I ask to recalculate) - - if rederiveGeometry or ("volp_miller" not in self.derived): - - self.produce_shape_lists() - - ( - self.derived["volp_miller"], - self.derived["surf_miller"], - self.derived["gradr_miller"], - self.derived["bp2_miller"], - self.derived["bt2_miller"], - self.derived["geo_bt"], - ) = GEOMETRYtools.calculateGeometricFactors( - self, - n_theta=n_theta_geo, - ) - - # Calculate flux surfaces - cn = np.array(self.shape_cos).T - sn = copy.deepcopy(self.shape_sin) - sn[0] = self.profiles["rmaj(m)"]*0.0 - sn[1] = np.arcsin(self.profiles["delta(-)"]) - sn[2] = -self.profiles["zeta(-)"] - sn = np.array(sn).T - flux_surfaces = GEQtools.mitim_flux_surfaces() - flux_surfaces.reconstruct_from_mxh_moments( - self.profiles["rmaj(m)"], - self.profiles["rmin(m)"], - self.profiles["kappa(-)"], - self.profiles["zmag(m)"], - cn, - sn) - self.derived["R_surface"],self.derived["Z_surface"] = flux_surfaces.R, flux_surfaces.Z - # ----------------------------------------------- - - #cross-sectional area of each flux surface - self.derived["surfXS"] = GEOMETRYtools.xsec_area_RZ( - self.derived["R_surface"], - self.derived["Z_surface"] - ) - - self.derived["R_LF"] = self.derived["R_surface"].max( - axis=1 - ) # self.profiles['rmaj(m)'][0]+self.profiles['rmin(m)'] - - # For Synchrotron - self.derived["B_ref"] = np.abs( - self.derived["B_unit"] * self.derived["geo_bt"] - ) - - # -------------------------------------------------------------------------- - # Reference mass - # -------------------------------------------------------------------------- - - # Forcing mass from this specific deriveQuantities call - if mi_ref is not None: - self.derived["mi_ref"] = mi_ref - print(f'\t- Using mi_ref={self.derived["mi_ref"]} provided in this particular deriveQuantities method, subtituting initialization one',typeMsg='i') - - # --------------------------------------------------------------------------------------------------------------------- - # --------- Important for scaling laws - # --------------------------------------------------------------------------------------------------------------------- - - self.derived["kappa95"] = np.interp( - 0.95, self.derived["psi_pol_n"], self.profiles["kappa(-)"] - ) - - self.derived["kappa995"] = np.interp( - 0.995, self.derived["psi_pol_n"], self.profiles["kappa(-)"] - ) - - self.derived["kappa_a"] = self.derived["surfXS"][-1] / np.pi / self.derived["a"] ** 2 - - self.derived["delta95"] = np.interp( - 0.95, self.derived["psi_pol_n"], self.profiles["delta(-)"] - ) - - self.derived["delta995"] = np.interp( - 0.995, self.derived["psi_pol_n"], self.profiles["delta(-)"] - ) - - self.derived["Rgeo"] = float(self.profiles["rcentr(m)"][-1]) - self.derived["B0"] = np.abs(float(self.profiles["bcentr(T)"][-1])) - - # --------------------------------------------------------------------------------------------------------------------- - - """ - surf_miller is truly surface area, but because of the GACODE definitions of flux, - Surf = V' <|grad r|> - Surf_GACODE = V' - """ - - self.derived["surfGACODE_miller"] = (self.derived["surf_miller"] / self.derived["gradr_miller"]) - - self.derived["surfGACODE_miller"][np.isnan(self.derived["surfGACODE_miller"])] = 0 - - self.derived["c_s"] = PLASMAtools.c_s( - self.profiles["te(keV)"], self.derived["mi_ref"] - ) - self.derived["rho_s"] = PLASMAtools.rho_s( - self.profiles["te(keV)"], self.derived["mi_ref"], self.derived["B_unit"] - ) - - self.derived["q_gb"], self.derived["g_gb"], self.derived["pi_gb"], self.derived["s_gb"], _ = PLASMAtools.gyrobohmUnits( - self.profiles["te(keV)"], - self.profiles["ne(10^19/m^3)"] * 1e-1, - self.derived["mi_ref"], - np.abs(self.derived["B_unit"]), - self.profiles["rmin(m)"][-1], - ) - - """ - In prgen_map_plasmastate: - qspow_e = expro_qohme+expro_qbeame+expro_qrfe+expro_qfuse-expro_qei & - -expro_qsync-expro_qbrem-expro_qline - qspow_i = expro_qbeami+expro_qrfi+expro_qfusi+expro_qei - """ - - qe_terms = { - "qohme(MW/m^3)": 1, - "qbeame(MW/m^3)": 1, - "qrfe(MW/m^3)": 1, - "qfuse(MW/m^3)": 1, - "qei(MW/m^3)": -1, - "qsync(MW/m^3)": -1, - "qbrem(MW/m^3)": -1, - "qline(MW/m^3)": -1, - "qione(MW/m^3)": 1, - } - - self.derived["qe"] = np.zeros(len(self.profiles["rho(-)"])) - for i in qe_terms: - if i in self.profiles: - self.derived["qe"] += qe_terms[i] * self.profiles[i] - - qrad = { - "qsync(MW/m^3)": 1, - "qbrem(MW/m^3)": 1, - "qline(MW/m^3)": 1, - } - - self.derived["qrad"] = np.zeros(len(self.profiles["rho(-)"])) - for i in qrad: - if i in self.profiles: - self.derived["qrad"] += qrad[i] * self.profiles[i] - - qi_terms = { - "qbeami(MW/m^3)": 1, - "qrfi(MW/m^3)": 1, - "qfusi(MW/m^3)": 1, - "qei(MW/m^3)": 1, - "qioni(MW/m^3)": 1, - } - - self.derived["qi"] = np.zeros(len(self.profiles["rho(-)"])) - for i in qi_terms: - if i in self.profiles: - self.derived["qi"] += qi_terms[i] * self.profiles[i] - - # Depends on GACODE version - ge_terms = {self.varqpar: 1, self.varqpar2: 1} - - self.derived["ge"] = np.zeros(len(self.profiles["rho(-)"])) - for i in ge_terms: - if i in self.profiles: - self.derived["ge"] += ge_terms[i] * self.profiles[i] - - """ - Careful, that's in MW/m^3. I need to find the volumes. Using here the Miller - calculation. Should be consistent with TGYRO - - profiles_gen puts any missing power into the CX: qioni, qione - """ - - r = self.profiles["rmin(m)"] - volp = self.derived["volp_miller"] - - self.derived["qe_MWmiller"] = CALCtools.integrateFS(self.derived["qe"], r, volp) - self.derived["qi_MWmiller"] = CALCtools.integrateFS(self.derived["qi"], r, volp) - self.derived["ge_10E20miller"] = CALCtools.integrateFS( - 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) - - self.derived["QiQe"] = self.derived["qi_MWm2"] / np.where(self.derived["qe_MWm2"] == 0, 1e-10, self.derived["qe_MWm2"]) # to avoid division by zero - - # "Convective" flux - self.derived["ce_MWmiller"] = PLASMAtools.convective_flux( - self.profiles["te(keV)"], self.derived["ge_10E20miller"] - ) - self.derived["ce_MWm2"] = PLASMAtools.convective_flux( - self.profiles["te(keV)"], self.derived["ge_10E20m2"] - ) - - # qmom - self.derived["mt_Jmiller"] = CALCtools.integrateFS( - self.profiles[self.varqmom], r, volp - ) - self.derived["mt_Jm2"] = self.derived["mt_Jmiller"] / (volp) - - # Extras for plotting in TGYRO for comparison - P = np.zeros(len(self.profiles["rmin(m)"])) - if "qsync(MW/m^3)" in self.profiles: - P += self.profiles["qsync(MW/m^3)"] - if "qbrem(MW/m^3)" in self.profiles: - P += self.profiles["qbrem(MW/m^3)"] - if "qline(MW/m^3)" in self.profiles: - P += self.profiles["qline(MW/m^3)"] - self.derived["qe_rad_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - P = self.profiles["qei(MW/m^3)"] - self.derived["qe_exc_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - """ - --------------------------------------------------------------------------------------------------------------------- - Note that the real auxiliary power is RF+BEAMS+OHMIC, - The QIONE is added by TGYRO, but sometimes it includes radiation and direct RF to electrons - --------------------------------------------------------------------------------------------------------------------- - """ - - # ** Electrons - - P = np.zeros(len(self.profiles["rho(-)"])) - for i in ["qrfe(MW/m^3)", "qohme(MW/m^3)", "qbeame(MW/m^3)"]: - if i in self.profiles: - P += self.profiles[i] - - self.derived["qe_auxONLY"] = copy.deepcopy(P) - self.derived["qe_auxONLY_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - for i in ["qione(MW/m^3)"]: - if i in self.profiles: - P += self.profiles[i] - - self.derived["qe_aux"] = copy.deepcopy(P) - self.derived["qe_aux_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - # ** Ions - - P = np.zeros(len(self.profiles["rho(-)"])) - for i in ["qrfi(MW/m^3)", "qbeami(MW/m^3)"]: - if i in self.profiles: - P += self.profiles[i] - - self.derived["qi_auxONLY"] = copy.deepcopy(P) - self.derived["qi_auxONLY_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - for i in ["qioni(MW/m^3)"]: - if i in self.profiles: - P += self.profiles[i] - - self.derived["qi_aux"] = copy.deepcopy(P) - self.derived["qi_aux_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - # ** General - - P = np.zeros(len(self.profiles["rho(-)"])) - for i in ["qohme(MW/m^3)"]: - if i in self.profiles: - P += self.profiles[i] - self.derived["qOhm_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - P = np.zeros(len(self.profiles["rho(-)"])) - for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: - if i in self.profiles: - P += self.profiles[i] - self.derived["qRF_MWmiller"] = CALCtools.integrateFS(P, r, volp) - if "qrfe(MW/m^3)" in self.profiles: - self.derived["qRFe_MWmiller"] = CALCtools.integrateFS( - self.profiles["qrfe(MW/m^3)"], r, volp - ) - if "qrfi(MW/m^3)" in self.profiles: - self.derived["qRFi_MWmiller"] = CALCtools.integrateFS( - self.profiles["qrfi(MW/m^3)"], r, volp - ) - - P = np.zeros(len(self.profiles["rho(-)"])) - for i in ["qbeame(MW/m^3)", "qbeami(MW/m^3)"]: - if i in self.profiles: - P += self.profiles[i] - self.derived["qBEAM_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.varqpar, self.varqpar2 = "qpar_beam(1/m^3/s)", "qpar_wall(1/m^3/s)" - self.derived["qrad_MWmiller"] = CALCtools.integrateFS(self.derived["qrad"], r, volp) - if "qsync(MW/m^3)" in self.profiles: - self.derived["qrad_sync_MWmiller"] = CALCtools.integrateFS(self.profiles["qsync(MW/m^3)"], r, volp) - else: - self.derived["qrad_sync_MWmiller"] = self.derived["qrad_MWmiller"]*0.0 - if "qbrem(MW/m^3)" in self.profiles: - self.derived["qrad_brem_MWmiller"] = CALCtools.integrateFS(self.profiles["qbrem(MW/m^3)"], r, volp) - else: - self.derived["qrad_brem_MWmiller"] = self.derived["qrad_MWmiller"]*0.0 - if "qline(MW/m^3)" in self.profiles: - self.derived["qrad_line_MWmiller"] = CALCtools.integrateFS(self.profiles["qline(MW/m^3)"], r, volp) + if "qmom(Nm)" in self.profiles: + self.varqmom = "qmom(Nm)" # Old, wrong one. But Candy fixed it as of 02/24/2023 else: - self.derived["qrad_line_MWmiller"] = self.derived["qrad_MWmiller"]*0.0 - - P = np.zeros(len(self.profiles["rho(-)"])) - for i in ["qfuse(MW/m^3)", "qfusi(MW/m^3)"]: - if i in self.profiles: - P += self.profiles[i] - self.derived["qFus_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - P = np.zeros(len(self.profiles["rho(-)"])) - for i in ["qioni(MW/m^3)", "qione(MW/m^3)"]: - if i in self.profiles: - P += self.profiles[i] - self.derived["qz_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - self.derived["q_MWmiller"] = ( - self.derived["qe_MWmiller"] + self.derived["qi_MWmiller"] - ) - - # --------------------------------------------------------------------------------------------------------------------- - # --------------------------------------------------------------------------------------------------------------------- - - P = np.zeros(len(self.profiles["rho(-)"])) - if "qfuse(MW/m^3)" in self.profiles: - P = self.profiles["qfuse(MW/m^3)"] - self.derived["qe_fus_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - P = np.zeros(len(self.profiles["rho(-)"])) - if "qfusi(MW/m^3)" in self.profiles: - P = self.profiles["qfusi(MW/m^3)"] - self.derived["qi_fus_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - P = np.zeros(len(self.profiles["rho(-)"])) - if "qfusi(MW/m^3)" in self.profiles: - self.derived["q_fus"] = ( - self.profiles["qfuse(MW/m^3)"] + self.profiles["qfusi(MW/m^3)"] - ) * 5 - P = self.derived["q_fus"] - self.derived["q_fus"] = P - self.derived["q_fus_MWmiller"] = CALCtools.integrateFS(P, r, volp) - - """ - Derivatives - """ - self.derived["aLTe"] = aLT(self.profiles["rmin(m)"], self.profiles["te(keV)"]) - self.derived["aLTi"] = self.profiles["ti(keV)"] * 0.0 - for i in range(self.profiles["ti(keV)"].shape[1]): - self.derived["aLTi"][:, i] = aLT( - self.profiles["rmin(m)"], self.profiles["ti(keV)"][:, i] - ) - self.derived["aLne"] = aLT( - self.profiles["rmin(m)"], self.profiles["ne(10^19/m^3)"] - ) - self.derived["aLni"] = [] - for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): - self.derived["aLni"].append( - aLT(self.profiles["rmin(m)"], self.profiles["ni(10^19/m^3)"][:, i]) - ) - self.derived["aLni"] = np.transpose(np.array(self.derived["aLni"])) - - if "w0(rad/s)" not in self.profiles: - self.profiles["w0(rad/s)"] = self.profiles["rho(-)"] * 0.0 - self.derived["aLw0"] = aLT(self.profiles["rmin(m)"], self.profiles["w0(rad/s)"]) - self.derived["dw0dr"] = -grad( - self.profiles["rmin(m)"], self.profiles["w0(rad/s)"] - ) - - self.derived["dqdr"] = grad(self.profiles["rmin(m)"], self.profiles["q(-)"]) - - """ - Other, performance - """ - qFus = self.derived["qe_fus_MWmiller"] + self.derived["qi_fus_MWmiller"] - self.derived["Pfus"] = qFus[-1] * 5 - - # Note that in cases with NPRAD=0 in TRANPS, this includes radiation! no way to deal wit this... - qIn = self.derived["qe_aux_MWmiller"] + self.derived["qi_aux_MWmiller"] - self.derived["qIn"] = qIn[-1] - self.derived["Q"] = self.derived["Pfus"] / self.derived["qIn"] - self.derived["qHeat"] = qIn[-1] + qFus[-1] - - self.derived["qTr"] = ( - self.derived["qe_aux_MWmiller"] - + self.derived["qi_aux_MWmiller"] - + (self.derived["qe_fus_MWmiller"] + self.derived["qi_fus_MWmiller"]) - - self.derived["qrad_MWmiller"] - ) - - self.derived["Prad"] = self.derived["qrad_MWmiller"][-1] - self.derived["Prad_sync"] = self.derived["qrad_sync_MWmiller"][-1] - self.derived["Prad_brem"] = self.derived["qrad_brem_MWmiller"][-1] - self.derived["Prad_line"] = self.derived["qrad_line_MWmiller"][-1] - self.derived["Psol"] = self.derived["qHeat"] - self.derived["Prad"] - - self.derived["ni_thr"] = [] - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm": - self.derived["ni_thr"].append(self.profiles["ni(10^19/m^3)"][:, sp]) - self.derived["ni_thr"] = np.transpose(self.derived["ni_thr"]) - self.derived["ni_thrAll"] = self.derived["ni_thr"].sum(axis=1) - - self.derived["ni_All"] = self.profiles["ni(10^19/m^3)"].sum(axis=1) - - - ( - self.derived["ptot_manual"], - self.derived["pe"], - self.derived["pi"], - ) = PLASMAtools.calculatePressure( - np.expand_dims(self.profiles["te(keV)"], 0), - np.expand_dims(np.transpose(self.profiles["ti(keV)"]), 0), - np.expand_dims(self.profiles["ne(10^19/m^3)"] * 0.1, 0), - np.expand_dims(np.transpose(self.profiles["ni(10^19/m^3)"] * 0.1), 0), - ) - self.derived["ptot_manual"], self.derived["pe"], self.derived["pi"] = ( - self.derived["ptot_manual"][0], - self.derived["pe"][0], - self.derived["pi"][0], - ) - - ( - self.derived["pthr_manual"], - _, - self.derived["pi_thr"], - ) = PLASMAtools.calculatePressure( - np.expand_dims(self.profiles["te(keV)"], 0), - np.expand_dims(np.transpose(self.profiles["ti(keV)"]), 0), - np.expand_dims(self.profiles["ne(10^19/m^3)"] * 0.1, 0), - np.expand_dims(np.transpose(self.derived["ni_thr"] * 0.1), 0), - ) - self.derived["pthr_manual"], self.derived["pi_thr"] = ( - self.derived["pthr_manual"][0], - self.derived["pi_thr"][0], - ) - - # ------- - # Content - # ------- - - ( - self.derived["We"], - self.derived["Wi_thr"], - self.derived["Ne"], - self.derived["Ni_thr"], - ) = PLASMAtools.calculateContent( - np.expand_dims(r, 0), - np.expand_dims(self.profiles["te(keV)"], 0), - np.expand_dims(np.transpose(self.profiles["ti(keV)"]), 0), - np.expand_dims(self.profiles["ne(10^19/m^3)"] * 0.1, 0), - np.expand_dims(np.transpose(self.derived["ni_thr"] * 0.1), 0), - np.expand_dims(volp, 0), - ) - - ( - self.derived["We"], - self.derived["Wi_thr"], - self.derived["Ne"], - self.derived["Ni_thr"], - ) = ( - self.derived["We"][0], - self.derived["Wi_thr"][0], - self.derived["Ne"][0], - self.derived["Ni_thr"][0], - ) - - self.derived["Nthr"] = self.derived["Ne"] + self.derived["Ni_thr"] - self.derived["Wthr"] = self.derived["We"] + self.derived["Wi_thr"] # Thermal - - self.derived["tauE"] = self.derived["Wthr"] / self.derived["qHeat"] # Seconds - - self.derived["tauP"] = np.where(self.derived["geIn"] != 0, self.derived["Ne"] / self.derived["geIn"], np.inf) # 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)"] - ).transpose().repeat(self.profiles["ni(10^19/m^3)"].shape[1], axis=1) - - # Vol-avg density - self.derived["volume"] = CALCtools.integrateFS(np.ones(r.shape[0]), r, volp)[ - -1 - ] # m^3 - self.derived["ne_vol20"] = ( - CALCtools.integrateFS(self.profiles["ne(10^19/m^3)"] * 0.1, r, volp)[-1] - / self.derived["volume"] - ) # 1E20/m^3 - - self.derived["ni_vol20"] = np.zeros(self.profiles["ni(10^19/m^3)"].shape[1]) - self.derived["fi_vol"] = np.zeros(self.profiles["ni(10^19/m^3)"].shape[1]) - for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): - self.derived["ni_vol20"][i] = ( - CALCtools.integrateFS( - self.profiles["ni(10^19/m^3)"][:, i] * 0.1, r, volp - )[-1] - / self.derived["volume"] - ) # 1E20/m^3 - self.derived["fi_vol"][i] = ( - self.derived["ni_vol20"][i] / self.derived["ne_vol20"] - ) - - self.derived["fi_onlyions_vol"] = self.derived["ni_vol20"] / np.sum( - self.derived["ni_vol20"] - ) - - self.derived["ne_peaking"] = ( - self.profiles["ne(10^19/m^3)"][0] * 0.1 / self.derived["ne_vol20"] - ) - - xcoord = self.derived[ - "rho_pol" - ] # to find the peaking at rho_pol (with square root) as in Angioni PRL 2003 - self.derived["ne_peaking0.2"] = ( - self.profiles["ne(10^19/m^3)"][np.argmin(np.abs(xcoord - 0.2))] - * 0.1 - / self.derived["ne_vol20"] - ) - - self.derived["Te_vol"] = ( - CALCtools.integrateFS(self.profiles["te(keV)"], r, volp)[-1] - / self.derived["volume"] - ) # keV - self.derived["Te_peaking"] = ( - self.profiles["te(keV)"][0] / self.derived["Te_vol"] - ) - self.derived["Ti_vol"] = ( - CALCtools.integrateFS(self.profiles["ti(keV)"][:, 0], r, volp)[-1] - / self.derived["volume"] - ) # keV - self.derived["Ti_peaking"] = ( - self.profiles["ti(keV)"][0, 0] / self.derived["Ti_vol"] - ) - - self.derived["ptot_manual_vol"] = ( - CALCtools.integrateFS(self.derived["ptot_manual"], r, volp)[-1] - / self.derived["volume"] - ) # MPa - self.derived["pthr_manual_vol"] = ( - CALCtools.integrateFS(self.derived["pthr_manual"], r, volp)[-1] - / self.derived["volume"] - ) # MPa - - self.derived['pfast_manual'] = self.derived['ptot_manual'] - self.derived['pthr_manual'] - self.derived["pfast_manual_vol"] = ( - CALCtools.integrateFS(self.derived["pfast_manual"], r, volp)[-1] - / self.derived["volume"] - ) # MPa - - self.derived['pfast_fraction'] = self.derived['pfast_manual_vol'] / self.derived['ptot_manual_vol'] - - #approximate pedestal top density - self.derived['ptop(Pa)'] = np.interp(0.90, self.profiles['rho(-)'], self.profiles['ptot(Pa)']) - - # Quasineutrality - self.derived["QN_Error"] = np.abs( - 1 - np.sum(self.derived["fi_vol"] * self.profiles["z"]) - ) - self.derived["Zeff"] = ( - np.sum(self.profiles["ni(10^19/m^3)"] * self.profiles["z"] ** 2, axis=1) - / self.profiles["ne(10^19/m^3)"] - ) - self.derived["Zeff_vol"] = ( - CALCtools.integrateFS(self.derived["Zeff"], r, volp)[-1] - / self.derived["volume"] - ) - - self.derived["nu_eff"] = PLASMAtools.coll_Angioni07( - self.derived["ne_vol20"] * 1e1, - self.derived["Te_vol"], - self.derived["Rgeo"], - Zeff=self.derived["Zeff_vol"], - ) - - self.derived["nu_eff2"] = PLASMAtools.coll_Angioni07( - self.derived["ne_vol20"] * 1e1, - self.derived["Te_vol"], - self.derived["Rgeo"], - Zeff=2.0, - ) - - # Avg mass - self.calculateMass() - - params_set_scaling = ( - np.abs(float(self.profiles["current(MA)"][-1])), - self.derived["Rgeo"], - self.derived["kappa_a"], - self.derived["ne_vol20"], - self.derived["a"] / self.derived["Rgeo"], - self.derived["B0"], - self.derived["mbg_main"], - self.derived["qHeat"], - ) - - self.derived["tau98y2"], self.derived["H98"] = PLASMAtools.tau98y2( - *params_set_scaling, tauE=self.derived["tauE"] - ) - self.derived["tau89p"], self.derived["H89"] = PLASMAtools.tau89p( - *params_set_scaling, tauE=self.derived["tauE"] - ) - self.derived["tau97L"], self.derived["H97L"] = PLASMAtools.tau97L( - *params_set_scaling, tauE=self.derived["tauE"] - ) - - """ - Mach number - """ - - Vtor_LF_Mach1 = PLASMAtools.constructVtorFromMach( - 1.0, self.profiles["ti(keV)"][:, 0], self.derived["mbg"] - ) # m/s - w0_Mach1 = Vtor_LF_Mach1 / (self.derived["R_LF"]) # rad/s - self.derived["MachNum"] = self.profiles["w0(rad/s)"] / w0_Mach1 - self.derived["MachNum_vol"] = ( - CALCtools.integrateFS(self.derived["MachNum"], r, volp)[-1] - / self.derived["volume"] - ) - - # Retain the old beta definition for comparison with 0D modeling - Beta_old = (self.derived["ptot_manual_vol"]* 1e6 / (self.derived["B0"] ** 2 / (2 * 4 * np.pi * 1e-7))) - self.derived["BetaN_engineering"] = (Beta_old / - (np.abs(float(self.profiles["current(MA)"][-1])) / - (self.derived["a"] * self.derived["B0"]) - )* 100.0 - ) # expressed in percent - - ''' - --------------------------------------------------------------------------------------------------- - Using B_unit, derive and for betap and betat calculations. - Equivalent to GACODE expro_bp2, expro_bt2 - --------------------------------------------------------------------------------------------------- - ''' - - self.derived["bp2_exp"] = self.derived["bp2_miller"] * self.derived["B_unit"] ** 2 - self.derived["bt2_exp"] = self.derived["bt2_miller"] * self.derived["B_unit"] ** 2 - - # Calculate the volume averages of bt2 and bp2 - - P = self.derived["bp2_exp"] - self.derived["bp2_vol_avg"] = CALCtools.integrateFS(P, r, volp)[-1] / self.derived["volume"] - P = self.derived["bt2_exp"] - self.derived["bt2_vol_avg"] = CALCtools.integrateFS(P, r, volp)[-1] / self.derived["volume"] - - # calculate beta_poloidal and beta_toroidal using volume averaged values - # mu0 = 4pi x 10^-7, also need to convert MPa to Pa - - self.derived["Beta_p"] = (2 * 4 * np.pi * 1e-7)*self.derived["ptot_manual_vol"]* 1e6/self.derived["bp2_vol_avg"] - self.derived["Beta_t"] = (2 * 4 * np.pi * 1e-7)*self.derived["ptot_manual_vol"]* 1e6/self.derived["bt2_vol_avg"] - - self.derived["Beta"] = 1/(1/self.derived["Beta_p"]+1/self.derived["Beta_t"]) - - TroyonFactor = np.abs(float(self.profiles["current(MA)"][-1])) / (self.derived["a"] * self.derived["B0"]) - - self.derived["BetaN"] = self.derived["Beta"] / TroyonFactor * 100.0 - - # --- - - nG = PLASMAtools.Greenwald_density( - np.abs(float(self.profiles["current(MA)"][-1])), - float(self.profiles["rmin(m)"][-1]), - ) - self.derived["fG"] = self.derived["ne_vol20"] / nG - self.derived["fG_x"] = self.profiles["ne(10^19/m^3)"]* 0.1 / nG - - self.derived["tite"] = self.profiles["ti(keV)"][:, 0] / self.profiles["te(keV)"] - self.derived["tite_vol"] = self.derived["Ti_vol"] / self.derived["Te_vol"] - - self.derived["LH_nmin"] = PLASMAtools.LHthreshold_nmin( - np.abs(float(self.profiles["current(MA)"][-1])), - self.derived["B0"], - self.derived["a"], - self.derived["Rgeo"], - ) - - self.derived["LH_Martin2"] = ( - PLASMAtools.LHthreshold_Martin2( - self.derived["ne_vol20"], - self.derived["B0"], - self.derived["a"], - self.derived["Rgeo"], - nmin=self.derived["LH_nmin"], - ) - * (2 / self.derived["mbg_main"]) ** 1.11 - ) + self.varqmom = "qmom(N/m^2)" # CORRECT ONE - self.derived["LHratio"] = self.derived["Psol"] / self.derived["LH_Martin2"] + # ------------------------------------------------------------------------------------------------------------------- + # Insert zeros in those cases whose column are not there + # ------------------------------------------------------------------------------------------------------------------- - self.readSpecies() + some_times_are_not_here = [ + "qei(MW/m^3)", + "qohme(MW/m^3)", + "johm(MA/m^2)", + "jbs(MA/m^2)", + "jbstor(MA/m^2)", + "w0(rad/s)", + "ptot(Pa)", # e.g. if I haven't written that info from ASTRA + "zeta(-)", # e.g. if TGYRO is run with zeta=0, it won't write this column in .new + "zmag(m)", + "qsync(MW/m^3)", + "qbrem(MW/m^3)", + "qline(MW/m^3)", + self.varqpar, + self.varqpar2, + "shape_cos0(-)", + self.varqmom, + ] - # ------------------------------------------------------- - # q-star - # ------------------------------------------------------- + num_moments = 6 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros + for i in range(num_moments): + some_times_are_not_here.append(f"shape_cos{i + 1}(-)") + if i > 1: + some_times_are_not_here.append(f"shape_sin{i + 1}(-)") - self.derived["qstar"] = PLASMAtools.evaluate_qstar( - self.profiles['current(MA)'][0], - self.profiles['rcentr(m)'], - self.derived['kappa95'], - self.profiles['bcentr(T)'], - self.derived['eps'], - self.derived['delta95'], - ITERcorrection=False, - includeShaping=True, - )[0] - self.derived["qstar_ITER"] = PLASMAtools.evaluate_qstar( - self.profiles['current(MA)'][0], - self.profiles['rcentr(m)'], - self.derived['kappa95'], - self.profiles['bcentr(T)'], - self.derived['eps'], - self.derived['delta95'], - ITERcorrection=True, - includeShaping=True, - )[0] - - # ------------------------------------------------------- - # Separatrix estimations - # ------------------------------------------------------- - - # ~~~~ Estimate lambda_q - pressure_atm = self.derived["ptot_manual_vol"] * 1e6 / 101325.0 - Lambda_q = PLASMAtools.calculateHeatFluxWidth_Brunner(pressure_atm) - - # ~~~~ Estimate upstream temperature - Bt = self.profiles["bcentr(T)"][0] - Bp = self.derived["eps"] * Bt / self.derived["q95"] #TODO: VERY ROUGH APPROXIMATION!!!! - - self.derived['Te_lcfs_estimate'] = PLASMAtools.calculateUpstreamTemperature( - Lambda_q, - self.derived["q95"], - self.derived["ne_vol20"], - self.derived["Psol"], - self.profiles["rcentr(m)"][0], - Bp, - Bt - )[0] - - # ~~~~ Estimate upstream density - self.derived['ne_lcfs_estimate'] = self.derived["ne_vol20"] * 0.6 - - # ------------------------------------------------------- - # TGLF-relevant quantities - # ------------------------------------------------------- - - self.tglf_plasma() - - def tglf_plasma(self): - - def deriv_gacode(y): - return grad(self.profiles["rmin(m)"],y).cpu().numpy() - - self.derived["tite_all"] = self.profiles["ti(keV)"] / self.profiles["te(keV)"][:, np.newaxis] - - self.derived['betae'] = PLASMAtools.betae( - self.profiles['te(keV)'], - self.profiles['ne(10^19/m^3)']*0.1, - self.derived["B_unit"]) - - self.derived['xnue'] = PLASMAtools.xnue( - torch.from_numpy(self.profiles['te(keV)']).to(torch.double), - torch.from_numpy(self.profiles['ne(10^19/m^3)']*0.1).to(torch.double), - self.derived["a"], - mref_u=self.derived["mi_ref"]).cpu().numpy() - - self.derived['debye'] = PLASMAtools.debye( - self.profiles['te(keV)'], - self.profiles['ne(10^19/m^3)']*0.1, - self.derived["mi_ref"], - self.derived["B_unit"]) - - self.derived['pprime'] = 1E-7 * self.profiles["q(-)"]*self.derived['a']**2/self.profiles["rmin(m)"]/self.derived["B_unit"]**2*deriv_gacode(self.profiles["ptot(Pa)"]) - self.derived['pprime'][0] = 0.0 - - self.derived['drmin/dr'] = deriv_gacode(self.profiles["rmin(m)"]) - self.derived['dRmaj/dr'] = deriv_gacode(self.profiles["rmaj(m)"]) - self.derived['dZmaj/dr'] = deriv_gacode(self.profiles["zmag(m)"]) - - self.derived['s_kappa'] = self.profiles["rmin(m)"] / self.profiles["kappa(-)"] * deriv_gacode(self.profiles["kappa(-)"]) - self.derived['s_delta'] = self.profiles["rmin(m)"] * deriv_gacode(self.profiles["delta(-)"]) - self.derived['s_zeta'] = self.profiles["rmin(m)"] * deriv_gacode(self.profiles["zeta(-)"]) - - s = self.profiles["rmin(m)"] / self.profiles["q(-)"]*deriv_gacode(self.profiles["q(-)"]) - self.derived['s_q'] = np.concatenate([np.array([0.0]),(self.profiles["q(-)"][1:] / self.derived['roa'][1:])**2 * s[1:]]) # infinite in first location - - ''' - Rotations - -------------------------------------------------------- - From TGYRO/TGLF definitions - w0p = expro_w0p(:)/100.0 - f_rot(:) = w0p(:)/w0_norm - gamma_p0 = -r_maj(i_r)*f_rot(i_r)*w0_norm - gamma_eb0 = gamma_p0*r(i_r)/(q_abs*r_maj(i_r)) - ''' - - w0p = deriv_gacode(self.profiles["w0(rad/s)"]) - gamma_p0 = -self.profiles["rmaj(m)"]*w0p - gamma_eb0 = -deriv_gacode(self.profiles["w0(rad/s)"]) * self.profiles["rmin(m)"]/self.profiles["q(-)"] - - self.derived['vexb_shear'] = gamma_eb0 * self.derived["a"]/self.derived['c_s'] - self.derived['vpar_shear'] = gamma_p0 * self.derived["a"]/self.derived['c_s'] - self.derived['vpar'] = self.profiles["rmaj(m)"]*self.profiles["w0(rad/s)"]/self.derived['c_s'] - - def calculateMass(self): - self.derived["mbg"] = 0.0 - self.derived["fmain"] = 0.0 - for i in range(self.derived["ni_vol20"].shape[0]): - self.derived["mbg"] += ( - float(self.profiles["mass"][i]) * self.derived["fi_onlyions_vol"][i] - ) - - if self.DTplasmaBool: - self.derived["mbg_main"] = ( - self.profiles["mass"][self.Dion] - * self.derived["fi_onlyions_vol"][self.Dion] - + self.profiles["mass"][self.Tion] - * self.derived["fi_onlyions_vol"][self.Tion] - ) / ( - self.derived["fi_onlyions_vol"][self.Dion] - + self.derived["fi_onlyions_vol"][self.Tion] - ) - self.derived["fmain"] = ( - self.derived["fi_vol"][self.Dion] + self.derived["fi_vol"][self.Tion] - ) - else: - self.derived["mbg_main"] = self.profiles["mass"][self.Mion] - self.derived["fmain"] = self.derived["fi_vol"][self.Mion] - - def deriveContentByVolumes(self, rhos=[0.5], impurityPosition=3): - """ - Calculates total particles and energy for ions and electrons, at a given volume - It fails near axis because of the polynomial integral, requiring a number of poitns - """ - - min_number_points = 3 - - We_x = np.zeros(self.profiles["te(keV)"].shape[0]) - Wi_x = np.zeros(self.profiles["te(keV)"].shape[0]) - Ne_x = np.zeros(self.profiles["te(keV)"].shape[0]) - Ni_x = np.zeros(self.profiles["te(keV)"].shape[0]) - for j in range(self.profiles["te(keV)"].shape[0] - min_number_points): - i = j + min_number_points - We_x[i], Wi_x[i], Ne_x[i], _ = PLASMAtools.calculateContent( - np.expand_dims(self.profiles["rmin(m)"][:i], 0), - np.expand_dims(self.profiles["te(keV)"][:i], 0), - np.expand_dims(np.transpose(self.profiles["ti(keV)"][:i]), 0), - np.expand_dims(self.profiles["ne(10^19/m^3)"][:i] * 0.1, 0), - np.expand_dims( - np.transpose(self.profiles["ni(10^19/m^3)"][:i] * 0.1), 0 - ), - np.expand_dims(self.derived["volp_miller"][:i], 0), - ) - - _, _, Ni_x[i], _ = PLASMAtools.calculateContent( - np.expand_dims(self.profiles["rmin(m)"][:i], 0), - np.expand_dims(self.profiles["te(keV)"][:i], 0), - np.expand_dims(np.transpose(self.profiles["ti(keV)"][:i]), 0), - np.expand_dims( - self.profiles["ni(10^19/m^3)"][:i, impurityPosition] * 0.1, 0 - ), - np.expand_dims( - np.transpose(self.profiles["ni(10^19/m^3)"][:i] * 0.1), 0 - ), - np.expand_dims(self.derived["volp_miller"][:i], 0), - ) - - We, Wi, Ne, Ni = ( - np.zeros(len(rhos)), - np.zeros(len(rhos)), - np.zeros(len(rhos)), - np.zeros(len(rhos)), - ) - for i in range(len(rhos)): - We[i] = np.interp(rhos[i], self.profiles["rho(-)"], We_x) - Wi[i] = np.interp(rhos[i], self.profiles["rho(-)"], Wi_x) - Ne[i] = np.interp(rhos[i], self.profiles["rho(-)"], Ne_x) - Ni[i] = np.interp(rhos[i], self.profiles["rho(-)"], Ni_x) - - return We, Wi, Ne, Ni - - 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("Engineering Parameters:") - print(f"\tBt = {self.profiles['bcentr(T)'][0]:.2f}T, Ip = {self.profiles['current(MA)'][0]:.2f}MA (q95 = {self.derived['q95']:.2f}, q* = {self.derived['qstar']:.2f}, q*ITER = {self.derived['qstar_ITER']:.2f}), Pin = {self.derived['qIn']:.2f}MW") - print(f"\tR = {self.profiles['rcentr(m)'][0]:.2f}m, a = {self.derived['a']:.2f}m (eps = {self.derived['eps']:.3f})") - print(f"\tkappa_sep = {self.profiles['kappa(-)'][-1]:.2f}, kappa_995 = {self.derived['kappa995']:.2f}, kappa_95 = {self.derived['kappa95']:.2f}, kappa_a = {self.derived['kappa_a']:.2f}") - print(f"\tdelta_sep = {self.profiles['delta(-)'][-1]:.2f}, delta_995 = {self.derived['delta995']:.2f}, delta_95 = {self.derived['delta95']:.2f}") - 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("\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_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}") - print(f"\tp_vol = {self.derived['ptot_manual_vol']:.2f} MPa ({self.derived['pfast_fraction']*100.0:.1f}% fast)") - print(f"\tBetaN = {self.derived['BetaN']:.3f} (BetaN w/B0 = {self.derived['BetaN_engineering']:.3f})") - print(f"\tPrad = {self.derived['Prad']:.1f}MW ({self.derived['Prad'] / self.derived['qHeat'] * 100.0:.1f}% of total) ({self.derived['Prad_brem']/self.derived['Prad'] * 100.0:.1f}% brem, {self.derived['Prad_line']/self.derived['Prad'] * 100.0:.1f}% line, {self.derived['Prad_sync']/self.derived['Prad'] * 100.0:.1f}% sync)") - print("\tPsol = {0:.1f}MW (fLH = {1:.2f})".format(self.derived["Psol"], self.derived["LHratio"])) - print("Operational point ( [,] = [{0:.2f},{1:.2f}] ) and species:".format(self.derived["ne_vol20"], self.derived["Te_vol"])) - print("\t = {0:.2f} keV (/ = {1:.2f}, Ti0/Te0 = {2:.2f})".format(self.derived["Ti_vol"],self.derived["tite_vol"],self.derived["tite"][0],)) - print("\tfG = {0:.2f} ( = {1:.2f} * 10^20 m^-3)".format(self.derived["fG"], self.derived["ne_vol20"])) - 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']:.1e}]") - print(f"\tMach = {self.derived['MachNum_vol']:.2f} (vol avg)") - 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(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 export_to_table(self, table=None, name=None): - - if table is None: - table = DataTable() - - data = [name] - for var in table.variables: - if table.variables[var][1] is not None: - if table.variables[var][1].split("_")[0] == "rho": - ix = np.argmin( - np.abs( - self.profiles["rho(-)"] - - float(table.variables[var][1].split("_")[1]) - ) - ) - elif table.variables[var][1].split("_")[0] == "psi": - ix = np.argmin( - np.abs( - self.derived["psi_pol_n"] - - float(table.variables[var][1].split("_")[1]) - ) - ) - elif table.variables[var][1].split("_")[0] == "pos": - ix = int(table.variables[var][1].split("_")[1]) - vari = self.__dict__[table.variables[var][2]][table.variables[var][0]][ - ix - ] - else: - vari = self.__dict__[table.variables[var][2]][table.variables[var][0]] - - data.append(f"{vari*table.variables[var][4]:{table.variables[var][3]}}") - - table.data.append(data) - print(f"\t* Exported {name} to table") - - return table - - def makeAllThermalIonsHaveSameTemp(self, refIon=0): - SpecRef = self.Species[refIon]["N"] - tiRef = self.profiles["ti(keV)"][:, refIon] - - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm" and sp != refIon: - print( - f"\t\t\t- Temperature forcing {self.Species[sp]['N']} --> {SpecRef}" - ) - self.profiles["ti(keV)"][:, sp] = tiRef - - def scaleAllThermalDensities(self, scaleFactor=1.0): - scaleFactor_ions = scaleFactor - - 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_ions):.3f}" - ) - ni_orig = self.profiles["ni(10^19/m^3)"][:, sp] - self.profiles["ni(10^19/m^3)"][:, sp] = scaleFactor_ions * ni_orig - - def toNumpyArrays(self): - self.profiles.update({key: tensor.cpu().detach().cpu().numpy() for key, tensor in self.profiles.items() if isinstance(tensor, torch.Tensor)}) - self.derived.update({key: tensor.cpu().detach().cpu().numpy() for key, tensor in self.derived.items() if isinstance(tensor, torch.Tensor)}) - - def writeCurrentStatus(self, file=None, limitedNames=False): - print("\t- Writting input.gacode file") - - if file is None: - file = self.file - - with open(file, "w") as f: - for line in self.header: - f.write(line) - - for i in self.profiles: - if "(" not in i: - f.write(f"# {i}\n") - else: - f.write(f"# {i.split('(')[0]} | {i.split('(')[-1].split(')')[0]}\n") - - if i in self.titles_single: - if i == "name" and limitedNames: - newlist = [self.profiles[i][0]] - for k in self.profiles[i][1:]: - if k not in [ - "D", - "H", - "T", - "He4", - "he4", - "C", - "O", - "Ar", - "W", - ]: - newlist.append("C") - else: - newlist.append(k) - print( - f"\n\n!! Correcting ion names from {self.profiles[i]} to {newlist} to avoid TGYRO radiation error (to solve in future?)\n\n", - typeMsg="w", - ) - listWrite = newlist - else: - listWrite = self.profiles[i] - - if IOtools.isfloat(listWrite[0]): - listWrite = [f"{i:.7e}".rjust(14) for i in listWrite] - f.write(f"{''.join(listWrite)}\n") - else: - f.write(f"{' '.join(listWrite)}\n") - - else: - if len(self.profiles[i].shape) == 1: - for j, val in enumerate(self.profiles[i]): - pos = f"{j + 1}".rjust(3) - valt = f"{round(val,99):.7e}".rjust(15) - f.write(f"{pos}{valt}\n") - else: - for j, val in enumerate(self.profiles[i]): - pos = f"{j + 1}".rjust(3) - txt = "".join([f"{k:.7e}".rjust(15) for k in val]) - f.write(f"{pos}{txt}\n") - - print(f"\t\t~ File {IOtools.clipstr(file)} written") - - # Update file - self.file = file - - def writeMiminalKinetic(self, file): - setProfs = [ - "rho(-)", - "polflux(Wb/radian)", - "q(-)", - "te(keV)", - "ti(keV)", - "ne(10^19/m^3)", - ] - - with open(file, "w") as f: - for i in setProfs: - if "(" not in i: - f.write(f"# {i}\n") - else: - f.write(f"# {i.split('(')[0]} | {i.split('(')[-1].split(')')[0]}\n") - - if len(self.profiles[i].shape) > 1: - p = self.profiles[i][:, 0] - else: - p = self.profiles[i] - - for j, val in enumerate(p): - pos = f"{j + 1}".rjust(3) - valt = f"{val:.7e}".rjust(15) - f.write(f"{pos}{valt}\n") - - def changeResolution(self, n=100, rho_new=None, interpolation_function=MATHtools.extrapolateCubicSpline): - rho = copy.deepcopy(self.profiles["rho(-)"]) - - if rho_new is None: - n = int(n) - rho_new = np.linspace(rho[0], rho[-1], n) - else: - rho_new = np.unique(np.sort(rho_new)) - n = len(rho_new) - - self.profiles["nexp"] = [str(n)] - - pro = self.profiles - for i in pro: - if i not in self.titles_single: - if len(pro[i].shape) == 1: - pro[i] = interpolation_function(rho_new, rho, pro[i]) - else: - prof = [] - for j in range(pro[i].shape[1]): - pp = interpolation_function(rho_new, rho, pro[i][:, j]) - prof.append(pp) - prof = np.array(prof) - - pro[i] = np.transpose(prof) - - self.produce_shape_lists() - - self.deriveQuantities() - - print( - f"\t\t- Resolution of profiles changed to {n} points with function {interpolation_function}" - ) - - def DTplasma(self): - self.Dion, self.Tion = None, None - try: - self.Dion = np.where(self.profiles["name"] == "D")[0][0] - except: - pass - try: - self.Tion = np.where(self.profiles["name"] == "T")[0][0] - except: - pass - - if self.Dion is not None and self.Tion is not None: - self.DTplasmaBool = True - else: - self.DTplasmaBool = False - if self.Dion is not None: - self.Mion = self.Dion # Main - elif self.Tion is not None: - self.Mion = self.Tion # Main - else: - self.Mion = ( - 0 # If no D or T, assume that the main ion is the first and only - ) - - self.ion_list_main = [] - if self.DTplasmaBool: - self.ion_list_main = [self.Dion+1, self.Tion+1] - else: - self.ion_list_main = [self.Mion+1] - - self.ion_list_impurities = [i+1 for i in range(len(self.Species)) if i+1 not in self.ion_list_main] - - def remove(self, ions_list): - # First order them - ions_list.sort() - print( - "\t\t- Removing ions in positions (of ions order, no zero): ", - ions_list, - typeMsg="i", - ) - - ions_list = [i - 1 for i in ions_list] - - fail = False - - var_changes = ["name", "type", "mass", "z"] - for i in var_changes: - try: - self.profiles[i] = np.delete(self.profiles[i], ions_list) - except: - print( - f"\t\t\t* Ions {[k+1 for k in ions_list]} could not be removed", - typeMsg="w", - ) - fail = True - break - - if not fail: - var_changes = ["ni(10^19/m^3)", "ti(keV)"] - for i in var_changes: - self.profiles[i] = np.delete(self.profiles[i], ions_list, axis=1) - - if not fail: - # Ensure we extract the scalar value from the array - self.profiles["nion"] = np.array( - [str(int(self.profiles["nion"][0]) - len(ions_list))] - ) - - self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) - - print("\t\t\t- Set of ions in updated profiles: ", self.profiles["name"]) - - def lumpSpecies( - self, ions_list=[2, 3], allthermal=False, forcename=None, force_integer=False, force_mass=None - ): - """ - if (D,Z1,Z2), lumping Z1 and Z2 requires ions_list = [2,3] - - if force_integer, the Zeff won't be kept exactly - """ - - # All thermal except first - if allthermal: - ions_list = [] - for i in range(len(self.Species) - 1): - if self.Species[i + 1]["S"] == "therm": - ions_list.append(i + 2) - lab = "therm" - else: - lab = "therm" - - print( - "\t\t- Lumping ions in positions (of ions order, no zero): ", - ions_list, - typeMsg="i", - ) - - if forcename is None: - forcename = "LUMPED" - - # Contributions to dilution and to Zeff - fZ1 = np.zeros(self.derived["fi"].shape[0]) - fZ2 = np.zeros(self.derived["fi"].shape[0]) - for i in ions_list: - fZ1 += self.Species[i - 1]["Z"] * self.derived["fi"][:, i - 1] - fZ2 += self.Species[i - 1]["Z"] ** 2 * self.derived["fi"][:, i - 1] - - Zr = fZ2 / fZ1 - Zr_vol = ( - CALCtools.integrateFS( - Zr, self.profiles["rmin(m)"], self.derived["volp_miller"] - )[-1] - / self.derived["volume"] - ) - - print(f'\t\t\t* Original plasma had Zeff_vol={self.derived["Zeff_vol"]:.2f}, QN error={self.derived["QN_Error"]:.4f}') - - # New specie parameters - if force_integer: - Z = round(Zr_vol) - print(f"\t\t\t* Lumped Z forced to be an integer ({Zr_vol}->{Z}), so plasma may not be quasineutral or fulfill original Zeff",typeMsg="w",) - else: - Z = Zr_vol - - A = Z * 2 if force_mass is None else force_mass - nZ = fZ1 / Z * self.profiles["ne(10^19/m^3)"] - - print(f"\t\t\t* New lumped impurity has Z={Z:.2f}, A={A:.2f} (calculated as 2*Z)") - - # Insert cases - self.profiles["nion"] = np.array([f"{int(self.profiles['nion'][0])+1}"]) - self.profiles["name"] = np.append(self.profiles["name"], forcename) - self.profiles["mass"] = np.append(self.profiles["mass"], A) - self.profiles["z"] = np.append(self.profiles["z"], Z) - self.profiles["type"] = np.append(self.profiles["type"], f"[{lab}]") - self.profiles["ni(10^19/m^3)"] = np.append( - self.profiles["ni(10^19/m^3)"], np.transpose(np.atleast_2d(nZ)), axis=1 - ) - self.profiles["ti(keV)"] = np.append( - self.profiles["ti(keV)"], - np.transpose(np.atleast_2d(self.profiles["ti(keV)"][:, 0])), - axis=1, - ) - - self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) - - # Remove species - self.remove(ions_list) - - # Contributions to dilution and to Zeff - print( - f'\t\t\t* New plasma has Zeff_vol={self.derived["Zeff_vol"]:.2f}, QN error={self.derived["QN_Error"]:.4f}' - ) - - def lumpImpurities(self): - - self.lumpSpecies(ions_list=self.ion_list_impurities) - - def lumpDT(self): - - if self.DTplasmaBool: - self.lumpSpecies(ions_list=self.ion_list_main, forcename="DT", force_mass=2.5) - else: - print('\t\t- No DT plasma, so no lumping of main ions') - - self.moveSpecie(pos=len(self.Species), pos_new=1) - - def changeZeff(self, Zeff, ion_pos=2, quasineutral_ions=None, enforceSameGradients=False): - """ - if (D,Z1,Z2), pos 1 -> change Z1 - """ - - if quasineutral_ions is None: - if self.DTplasmaBool: - quasineutral_ions = [self.Dion, self.Tion] - else: - quasineutral_ions = [self.Mion] - - print(f'\t\t- Changing Zeff (from {self.derived["Zeff_vol"]:.3f} to {Zeff=:.3f}) by changing content of ion in position {ion_pos} {self.Species[ion_pos]["N"],self.Species[ion_pos]["Z"]}, quasineutralized by ions {quasineutral_ions}',typeMsg="i",) - - # Plasma needs to be in quasineutrality to start with - self.enforceQuasineutrality() - - # ------------------------------------------------------ - # Contributions to equations - # ------------------------------------------------------ - Zq = np.zeros(self.derived["fi"].shape[0]) - Zq2 = np.zeros(self.derived["fi"].shape[0]) - fZj = np.zeros(self.derived["fi"].shape[0]) - fZj2 = np.zeros(self.derived["fi"].shape[0]) - for i in range(len(self.Species)): - if i in quasineutral_ions: - Zq += self.Species[i]["Z"] - Zq2 += self.Species[i]["Z"] ** 2 - elif i != ion_pos: - fZj += self.Species[i]["Z"] * self.derived["fi"][:, i] - fZj2 += self.Species[i]["Z"] ** 2 * self.derived["fi"][:, i] - else: - Zk = self.Species[i]["Z"] - - # ------------------------------------------------------ - # Find free parameters (fk and fq) - # ------------------------------------------------------ - - fk = ( Zeff - (1-fZj)*Zq2/Zq - fZj2 ) / ( Zk**2 - Zk*Zq2/Zq) - fq = ( 1 - fZj - fk*Zk ) / Zq - - if (fq<0).any(): - raise ValueError(f"Zeff cannot be reduced by changing ion #{ion_pos} because it would require negative densities for quasineutral ions") - - # ------------------------------------------------------ - # Insert - # ------------------------------------------------------ - - fi_orig = self.derived["fi"][:, ion_pos] - - self.profiles["ni(10^19/m^3)"][:, ion_pos] = fk * self.profiles["ne(10^19/m^3)"] - for i in quasineutral_ions: - self.profiles["ni(10^19/m^3)"][:, i] = fq * self.profiles["ne(10^19/m^3)"] - - self.readSpecies() - - self.deriveQuantities(rederiveGeometry=False) - - if enforceSameGradients: - self.scaleAllThermalDensities() - self.deriveQuantities(rederiveGeometry=False) - - print(f'\t\t\t* Dilution changed from {fi_orig.mean():.2e} (vol avg) to { self.derived["fi"][:, ion_pos].mean():.2e} to achieve Zeff={self.derived["Zeff_vol"]:.3f} (fDT={self.derived["fmain"]:.3f}) [quasineutrality error = {self.derived["QN_Error"]:.1e}]') - - def moveSpecie(self, pos=2, pos_new=1): - """ - if (D,Z1,Z2), pos 1 pos_new 2-> (Z1,D,Z2) - """ - - if pos_new > pos: - pos, pos_new = pos_new, pos - - position_to_moveFROM_in_profiles = pos - 1 - position_to_moveTO_in_profiles = pos_new - 1 - - print(f'\t\t- Moving ion in position (of ions order, no zero) {pos} ({self.profiles["name"][position_to_moveFROM_in_profiles]}) to {pos_new}',typeMsg="i",) - - self.profiles["nion"] = np.array([f"{int(self.profiles['nion'][0])+1}"]) - - for ikey in ["name", "mass", "z", "type", "ni(10^19/m^3)", "ti(keV)"]: - if len(self.profiles[ikey].shape) > 1: - axis = 1 - newly = self.profiles[ikey][:, position_to_moveFROM_in_profiles] - else: - axis = 0 - newly = self.profiles[ikey][position_to_moveFROM_in_profiles] - self.profiles[ikey] = np.insert( - self.profiles[ikey], position_to_moveTO_in_profiles, newly, axis=axis - ) - - self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) - - if position_to_moveTO_in_profiles > position_to_moveFROM_in_profiles: - self.remove([position_to_moveFROM_in_profiles + 1]) - else: - self.remove([position_to_moveFROM_in_profiles + 2]) - - def addSpecie(self, Z=5.0, mass=10.0, fi_vol=0.1, forcename=None): - print( - f"\t\t- Creating new specie with Z={Z}, mass={mass}, fi_vol={fi_vol}", - typeMsg="i", - ) - - if forcename is None: - forcename = "LUMPED" - - lab = "therm" - nZ = fi_vol * self.profiles["ne(10^19/m^3)"] - - self.profiles["nion"] = np.array([f"{int(self.profiles['nion'][0])+1}"]) - self.profiles["name"] = np.append(self.profiles["name"], forcename) - self.profiles["mass"] = np.append(self.profiles["mass"], mass) - self.profiles["z"] = np.append(self.profiles["z"], Z) - self.profiles["type"] = np.append(self.profiles["type"], f"[{lab}]") - self.profiles["ni(10^19/m^3)"] = np.append( - self.profiles["ni(10^19/m^3)"], np.transpose(np.atleast_2d(nZ)), axis=1 - ) - self.profiles["ti(keV)"] = np.append( - self.profiles["ti(keV)"], - np.transpose(np.atleast_2d(self.profiles["ti(keV)"][:, 0])), - axis=1, - ) - if "vtor(m/s)" in self.profiles: - self.profiles["vtor(m/s)"] = np.append( - self.profiles["vtor(m/s)"], - np.transpose(np.atleast_2d(self.profiles["vtor(m/s)"][:, 0])), - axis=1, - ) - - self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) - - def correct(self, options={}, write=False, new_file=None): - """ - if name= T D LUMPED, and I want to eliminate D, removeIons = [2] - """ - - recompute_ptot = options.get("recompute_ptot", True) # Only done by default - removeIons = options.get("removeIons", []) - removeFast = options.get("removeFast", False) - quasineutrality = options.get("quasineutrality", False) - sameDensityGradients = options.get("sameDensityGradients", False) - groupQIONE = options.get("groupQIONE", False) - ensurePostiveGamma = options.get("ensurePostiveGamma", False) - ensureMachNumber = options.get("ensureMachNumber", None) - FastIsThermal = options.get("FastIsThermal", False) - - print("\t- Custom correction of input.gacode file has been requested") - - # ---------------------------------------------------------------------- - # Correct - # ---------------------------------------------------------------------- - - # Remove desired ions - if len(removeIons) > 0: - self.remove(removeIons) - - # Remove fast - if removeFast: - ions_fast = [] - for sp in range(len(self.Species)): - if self.Species[sp]["S"] != "therm": - ions_fast.append(sp + 1) - if len(ions_fast) > 0: - print( - f"\t\t- Detected fast ions in positions {ions_fast}, removing them..." - ) - self.remove(ions_fast) - # Fast as thermal - elif FastIsThermal: - self.make_fast_ions_thermal() - - # Correct LUMPED - for i in range(len(self.profiles["name"])): - if self.profiles["name"][i] in ["LUMPED", "None"]: - name = ionName( - int(self.profiles["z"][i]), int(self.profiles["mass"][i]) - ) - if name is not None: - print( - f'\t\t- Ion in position #{i+1} was named LUMPED with Z={self.profiles["z"][i]}, now it is renamed to {name}', - typeMsg="i", - ) - self.profiles["name"][i] = name - else: - print( - f'\t\t- Ion in position #{i+1} was named LUMPED with Z={self.profiles["z"][i]}, but I could not find what element it is, so doing nothing', - typeMsg="w", - ) - - # Correct qione - if groupQIONE and (np.abs(self.profiles["qione(MW/m^3)"].sum()) > 1e-14): - print('\t\t- Inserting "qione" into "qrfe"', typeMsg="i") - self.profiles["qrfe(MW/m^3)"] += self.profiles["qione(MW/m^3)"] - self.profiles["qione(MW/m^3)"] = self.profiles["qione(MW/m^3)"] * 0.0 - - # Make all thermal ions have the same gradient as the electron density, by keeping volume average constant - if sameDensityGradients: - self.enforce_same_density_gradients() - - # Enforce quasineutrality - if quasineutrality: - self.enforceQuasineutrality() - - print(f"\t\t\t* Quasineutrality error = {self.derived['QN_Error']:.1e}") - - # Recompute ptot - if recompute_ptot: - self.deriveQuantities(rederiveGeometry=False) - self.selfconsistentPTOT() - - # If I don't trust the negative particle flux in the core that comes from TRANSP... - if ensurePostiveGamma: - print("\t\t- Making particle flux always positive", typeMsg="i") - self.profiles[self.varqpar] = self.profiles[self.varqpar].clip(0) - self.profiles[self.varqpar2] = self.profiles[self.varqpar2].clip(0) - - # Mach - if ensureMachNumber is not None: - self.introduceRotationProfile(Mach_LF=ensureMachNumber) - - # ---------------------------------------------------------------------- - # Re-derive - # ---------------------------------------------------------------------- - - self.deriveQuantities(rederiveGeometry=False) - - # ---------------------------------------------------------------------- - # Write - # ---------------------------------------------------------------------- - if write: - self.writeCurrentStatus(file=new_file) - self.printInfo() - - def enforce_same_density_gradients(self): - txt = "" - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm": - self.profiles["ni(10^19/m^3)"][:, sp] = self.derived["fi_vol"][sp] * self.profiles["ne(10^19/m^3)"] - txt += f"{self.Species[sp]['N']} " - print(f"\t\t- Making all thermal ions ({txt}) have the same a/Ln as electrons (making them an exact flat fraction)",typeMsg="i",) - self.deriveQuantities(rederiveGeometry=False) - - def make_fast_ions_thermal(self): - modified_num = 0 - for i in range(len(self.Species)): - if self.Species[i]["S"] != "therm": - print( - f'\t\t- Specie {i} ({self.profiles["name"][i]}) was fast, but now it is considered thermal' - ) - self.Species[i]["S"] = "therm" - self.profiles["type"][i] = "[therm]" - self.profiles["ti(keV)"][:, i] = self.profiles["ti(keV)"][:, 0] - modified_num += 1 - if modified_num > 0: - print("\t- Making fast species as if they were thermal (to keep dilution effect and Qi-sum of fluxes)",typeMsg="w") - - def selfconsistentPTOT(self): - print(f"\t\t* Recomputing ptot and inserting it as ptot(Pa), changed from p0 = {self.profiles['ptot(Pa)'][0] * 1e-3:.1f} to {self.derived['ptot_manual'][0]*1e+3:.1f} kPa",typeMsg="i") - self.profiles["ptot(Pa)"] = self.derived["ptot_manual"] * 1e6 - - def enforceQuasineutrality(self): - print(f"\t\t- Enforcing quasineutrality (error = {self.derived['QN_Error']:.1e})",typeMsg="i",) - - # What's the lack of quasineutrality? - ni = self.profiles["ne(10^19/m^3)"] * 0.0 - for sp in range(len(self.Species)): - ni += self.profiles["ni(10^19/m^3)"][:, sp] * self.profiles["z"][sp] - ne_missing = self.profiles["ne(10^19/m^3)"] - ni - - # What ion to modify? - if self.DTplasmaBool: - print("\t\t\t* Enforcing quasineutrality by modifying D and T equally") - prev_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Dion]) - self.profiles["ni(10^19/m^3)"][:, self.Dion] += ne_missing / 2 - self.profiles["ni(10^19/m^3)"][:, self.Tion] += ne_missing / 2 - new_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Dion]) - else: - print( - f"\t\t\t* Enforcing quasineutrality by modifying main ion (position #{self.Mion})" - ) - prev_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Mion]) - self.profiles["ni(10^19/m^3)"][:, self.Mion] += ne_missing - new_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Mion]) - - print( - f"\t\t\t\t- Changed on-axis density from n0 = {prev_on_axis:.2f} to {new_on_axis:.2f} ({100*(new_on_axis-prev_on_axis)/prev_on_axis:.1f}%)" - ) - - self.deriveQuantities(rederiveGeometry=False) - - def introduceRotationProfile(self, Mach_LF=1.0, new_file=None): - print(f"\t- Enforcing Mach Number in LF of {Mach_LF}") - self.deriveQuantities() - Vtor_LF = PLASMAtools.constructVtorFromMach( - Mach_LF, self.profiles["ti(keV)"][:, 0], self.derived["mbg"] - ) # m/s - - self.profiles["w0(rad/s)"] = Vtor_LF / (self.derived["R_LF"]) # rad/s - - self.deriveQuantities() - - if new_file is not None: - self.writeCurrentStatus(file=new_file) - - def plot( - self, - axs1=None, - axs2=None, - axs3=None, - axs4=None, - axsFlows=None, - axs6=None, - axsImps=None, - color="b", - legYN=True, - extralab="", - fn=None, - fnlab="", - lsFlows="-", - legFlows=True, - showtexts=True, - lastRhoGradients=0.89, - ): - if axs1 is None: - if fn is None: - from mitim_tools.misc_tools.GUItools import FigureNotebook - - self.fn = FigureNotebook("PROFILES Notebook", geometry="1600x1000") - - fig, fig2, fig3, fig4, fig5, fig6, fig7 = add_figures(self.fn, fnlab=fnlab) - - grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) - axs1 = [ - fig.add_subplot(grid[0, 0]), - fig.add_subplot(grid[1, 0]), - fig.add_subplot(grid[2, 0]), - fig.add_subplot(grid[0, 1]), - fig.add_subplot(grid[1, 1]), - fig.add_subplot(grid[2, 1]), - fig.add_subplot(grid[0, 2]), - fig.add_subplot(grid[1, 2]), - fig.add_subplot(grid[2, 2]), - ] - - - grid = plt.GridSpec(3, 2, hspace=0.3, wspace=0.3) - axs2 = [ - fig2.add_subplot(grid[0, 0]), - fig2.add_subplot(grid[0, 1]), - fig2.add_subplot(grid[1, 0]), - fig2.add_subplot(grid[1, 1]), - fig2.add_subplot(grid[2, 0]), - fig2.add_subplot(grid[2, 1]), - ] - - - grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.5) - ax00c = fig3.add_subplot(grid[0, 0]) - axs3 = [ - ax00c, - fig3.add_subplot(grid[1, 0], sharex=ax00c), - fig3.add_subplot(grid[2, 0], sharex=ax00c), - fig3.add_subplot(grid[0, 1], sharex=ax00c), - fig3.add_subplot(grid[1, 1], sharex=ax00c), - fig3.add_subplot(grid[2, 1], sharex=ax00c), - fig3.add_subplot(grid[0, 2], sharex=ax00c), - fig3.add_subplot(grid[1, 2], sharex=ax00c), - fig3.add_subplot(grid[2, 2], sharex=ax00c), - fig3.add_subplot(grid[0, 3], sharex=ax00c), - fig3.add_subplot(grid[1, 3], sharex=ax00c), - fig3.add_subplot(grid[2, 3], sharex=ax00c), - ] - - - grid = plt.GridSpec(2, 3, 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]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - - axsFlows = [ - fig5.add_subplot(grid[0, 0]), - fig5.add_subplot(grid[1, 0]), - fig5.add_subplot(grid[0, 1]), - fig5.add_subplot(grid[0, 2]), - fig5.add_subplot(grid[1, 1]), - fig5.add_subplot(grid[1, 2]), - ] - - - grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) - axs6 = [ - fig6.add_subplot(grid[0, 0]), - fig6.add_subplot(grid[:, 1]), - fig6.add_subplot(grid[0, 2]), - fig6.add_subplot(grid[1, 0]), - fig6.add_subplot(grid[1, 2]), - fig6.add_subplot(grid[0, 3]), - fig6.add_subplot(grid[1, 3]), - ] - - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsImps = [ - fig7.add_subplot(grid[0, 0]), - fig7.add_subplot(grid[0, 1]), - fig7.add_subplot(grid[0, 2]), - fig7.add_subplot(grid[1, 0]), - fig7.add_subplot(grid[1, 1]), - fig7.add_subplot(grid[1, 2]), - ] - - [ax00, ax10, ax20, ax01, ax11, ax21, ax02, ax12, ax22] = axs1 - [ax00b, ax01b, ax10b, ax11b, ax20b, ax21b] = axs2 - [ - ax00c, - ax10c, - ax20c, - ax01c, - ax11c, - ax21c, - ax02c, - ax12c, - ax22c, - ax03c, - ax13c, - ax23c, - ] = axs3 - - lw = 1 - fs = 6 - rho = self.profiles["rho(-)"] - - lines = ["-", "--", "-.", ":", "-", "--", "-."] - - self.plot_temps(ax=ax00, leg=legYN, col=color, lw=lw, fs=fs, extralab=extralab) - self.plot_dens(ax=ax01, leg=legYN, col=color, lw=lw, fs=fs, extralab=extralab) - - ax = ax10 - cont = 0 - for i in range(len(self.Species)): - if self.Species[i]["S"] == "therm": - var = self.profiles["ti(keV)"][:, i] - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - varL = "Thermal $T_i$ (keV)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax20 - cont = 0 - for i in range(len(self.Species)): - if self.Species[i]["S"] == "fast": - var = self.profiles["ti(keV)"][:, i] - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - varL = "Fast $T_i$ (keV)" - ax.plot( - rho, - self.profiles["ti(keV)"][:, 0], - lw=0.5, - ls="-", - alpha=0.5, - c=color, - label=extralab + "$T_{i,1}$", - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax11 - cont = 0 - for i in range(len(self.Species)): - if self.Species[i]["S"] == "therm": - var = self.profiles["ni(10^19/m^3)"][:, i] * 1e-1 - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - varL = "Thermal $n_i$ ($10^{20}/m^3$)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax21 - cont = 0 - for i in range(len(self.Species)): - if self.Species[i]["S"] == "fast": - var = self.profiles["ni(10^19/m^3)"][:, i] * 1e-1 * 1e5 - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - varL = "Fast $n_i$ ($10^{15}/m^3$)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN and cont>0: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax02 - var = self.profiles["w0(rad/s)"] - ax.plot(rho, var, lw=lw, ls="-", c=color) - varL = "$\\omega_{0}$ (rad/s)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax12 - var = self.profiles["ptot(Pa)"] * 1e-6 - ax.plot(rho, var, lw=lw, ls="-", c=color, label=extralab + "ptot") - if "ptot_manual" in self.derived: - ax.plot( - rho, - self.derived["ptot_manual"], - lw=lw, - ls="--", - c=color, - label=extralab + "check", - ) - # ax.plot(rho,np.abs(var-self.derived['ptot_manual']),lw=lw,ls='-.',c=color,label=extralab+'diff') - - ax.plot( - rho, - self.derived["pthr_manual"], - lw=lw, - ls="-.", - c=color, - label=extralab + "check, thrm", - ) - - - varL = "$p$ (MPa)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - # ax.set_ylim(bottom=0) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax00b - varL = "$MW/m^3$" - cont = 0 - var = -self.profiles["qei(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "i->e", c=color) - cont += 1 - if "qrfe(MW/m^3)" in self.profiles: - var = self.profiles["qrfe(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "rf", c=color) - cont += 1 - if "qfuse(MW/m^3)" in self.profiles: - var = self.profiles["qfuse(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "fus", c=color) - cont += 1 - if "qbeame(MW/m^3)" in self.profiles: - var = self.profiles["qbeame(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "beam", c=color) - cont += 1 - if "qione(MW/m^3)" in self.profiles: - var = self.profiles["qione(MW/m^3)"] - ax.plot( - rho, var, lw=lw / 2, ls=lines[cont], label=extralab + "extra", c=color - ) - cont += 1 - if "qohme(MW/m^3)" in self.profiles: - var = self.profiles["qohme(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "ohmic", c=color) - cont += 1 - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - ax.set_title("Electron Power Density") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax01b - if "varqmom" not in self.__dict__: - self.varqmom = "qmom(N/m^2)" - self.profiles[self.varqmom] = self.profiles["rho(-)"] * 0.0 - - ax.plot(rho, self.profiles[self.varqmom], lw=lw, ls="-", c=color) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$N/m^2$, $J/m^3$") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - ax.set_title("Momentum Source Density") - - ax = ax21b - ax.plot( - rho, self.derived["qe_MWm2"], lw=lw, ls="-", label=extralab + "qe", c=color - ) - ax.plot( - rho, self.derived["qi_MWm2"], lw=lw, ls="--", label=extralab + "qi", c=color - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("Heat Flux ($MW/m^2$)") - if legYN: - ax.legend(loc="lower left", fontsize=fs) - ax.set_title("Flux per unit area (gacode: P/V')") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax21b.twinx() - ax.plot( - rho, - self.derived["ge_10E20m2"], - lw=lw, - ls="-.", - label=extralab + "$\\Gamma_e$", - c=color, - ) - ax.set_ylabel("Particle Flux ($10^{20}/m^2/s$)") - if legYN: - ax.legend(loc="lower right", fontsize=fs) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax20b - varL = "$Q_{rad}$ ($MW/m^3$)" - if "qbrem(MW/m^3)" in self.profiles: - var = self.profiles["qbrem(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls="-", label=extralab + "brem", c=color) - if "qline(MW/m^3)" in self.profiles: - var = self.profiles["qline(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls="--", label=extralab + "line", c=color) - if "qsync(MW/m^3)" in self.profiles: - var = self.profiles["qsync(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=":", label=extralab + "sync", c=color) - - var = self.derived["qrad"] - ax.plot(rho, var, lw=lw * 1.5, ls="-", label=extralab + "Total", c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - ax.set_title("Radiation Contributions") - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax10b - varL = "$MW/m^3$" - cont = 0 - var = self.profiles["qei(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "e->i", c=color) - cont += 1 - if "qrfi(MW/m^3)" in self.profiles: - var = self.profiles["qrfi(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "rf", c=color) - cont += 1 - if "qfusi(MW/m^3)" in self.profiles: - var = self.profiles["qfusi(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "fus", c=color) - cont += 1 - if "qbeami(MW/m^3)" in self.profiles: - var = self.profiles["qbeami(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "beam", c=color) - cont += 1 - if "qioni(MW/m^3)" in self.profiles: - var = self.profiles["qioni(MW/m^3)"] - ax.plot( - rho, var, lw=lw / 2, ls=lines[cont], label=extralab + "extra", c=color - ) - cont += 1 - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - ax.set_title("Ion Power Density") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - """ - Note that in prgen_map_plasmastate, that variable: - expro_qpar_beam(i) = plst_sn_trans(i-1)/dvol - - Note that in prgen_read_plasmastate, that variable: - ! Particle source - err = nf90_inq_varid(ncid,trim('sn_trans'),varid) - err = nf90_get_var(ncid,varid,plst_sn_trans(1:nx-1)) - plst_sn_trans(nx) = 0.0 - - Note that in the plasmastate file, the variable "sn_trans": - - long_name: particle transport (loss) - units: #/sec - component: PLASMA - section: STATE_PROFILES - specification: R|units=#/sec|step*dV sn_trans(~nrho,0:nspec_th) - - So, this means that expro_qpar_beam is in units of #/sec/m^3, meaning that - it is a particle flux DENSITY. It therefore requires volume integral and - divide by surface to produce a flux. - - The units of this qpar_beam column is NOT MW/m^3. In the gacode source codes - they also say that those units are wrong. - - """ - - ax = ax11b - cont = 0 - var = self.profiles[self.varqpar] * 1e-20 - ax.plot(rho, var, lw=lw, ls=lines[0], c=color, label=extralab + "beam") - var = self.profiles[self.varqpar2] * 1e-20 - ax.plot(rho, var, lw=lw, ls=lines[1], c=color, label=extralab + "wall") - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.axhline(y=0, lw=0.5, ls="--", c="k") - ax.set_ylabel("$10^{20}m^{-3}s^{-1}$") - ax.set_title("Particle Source Density") - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax00c - varL = "cos Shape Params" - yl = 0 - cont = 0 - - for i, s in enumerate(self.shape_cos): - if s is not None: - valmax = np.abs(s).max() - if valmax > 1e-10: - lab = f"c{i}" - ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) - cont += 1 - - yl = np.max([yl, valmax]) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - - - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - if legYN: - ax.legend(loc="best", fontsize=fs) - - ax = ax01c - varL = "sin Shape Params" - cont = 0 - for i, s in enumerate(self.shape_sin): - if s is not None: - valmax = np.abs(s).max() - if valmax > 1e-10: - lab = f"s{i}" - ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) - cont += 1 - - yl = np.max([yl, valmax]) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax02c - var = self.profiles["q(-)"] - ax.plot(rho, var, lw=lw, ls="-", c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0) - ax.set_ylabel("q") - - ax.axhline(y=1, ls="--", c="k", lw=1) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0.0) - - - - ax = ax12c - var = self.profiles["polflux(Wb/radian)"] - ax.plot(rho, var, lw=lw, ls="-", c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("Poloidal $\\psi$ ($Wb/rad$)") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax10c - - var = self.profiles["rho(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0) - ax.set_ylabel("$\\rho$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax11c - - var = self.profiles["rmin(m)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylim(bottom=0) - ax.set_ylabel("$r_{min}$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax20c - - var = self.profiles["rmaj(m)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$R_{maj}$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax21c - - var = self.profiles["zmag(m)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - yl = np.max([0.1, np.max(np.abs(var))]) - ax.set_ylim([-yl, yl]) - ax.set_ylabel("$Z_{maj}$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax22c - - var = self.profiles["kappa(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$\\kappa$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=1) - - ax = ax03c - - var = self.profiles["delta(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$\\delta$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax13c - - var = self.profiles["zeta(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("zeta") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax23c - - var = self.profiles["johm(MA/m^2)"] - ax.plot(rho, var, "-", lw=lw, c=color, label=extralab + "$J_{OH}$") - var = self.profiles["jbs(MA/m^2)"] - ax.plot(rho, var, "--", lw=lw, c=color, label=extralab + "$J_{BS,par}$") - var = self.profiles["jbstor(MA/m^2)"] - ax.plot(rho, var, "-.", lw=lw, c=color, label=extralab + "$J_{BS,tor}$") - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylim(bottom=0) - ax.set_ylabel("J ($MA/m^2$)") - if legYN: - ax.legend(loc="best", prop={"size": 7}) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - # Derived - self.plotGradients( - axs4, color=color, lw=lw, lastRho=lastRhoGradients, label=extralab - ) - - # Others - ax = axs6[0] - ax.plot(self.profiles["rho(-)"], self.derived["dw0dr"] * 1e-5, c=color, lw=lw) - ax.set_ylabel("$-d\\omega_0/dr$ (krad/s/cm)") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - ax.axhline(y=0, lw=1.0, c="k", ls="--") - - ax = axs6[2] - ax.plot(self.profiles["rho(-)"], self.derived["q_fus"], c=color, lw=lw) - ax.set_ylabel("$q_{fus}$ ($MW/m^3$)") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axs6[3] - ax.plot(self.profiles["rho(-)"], self.derived["q_fus_MWmiller"], c=color, lw=lw) - ax.set_ylabel("$P_{fus}$ ($MW$)") - ax.set_xlim([0, 1]) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axs6[4] - ax.plot(self.profiles["rho(-)"], self.derived["tite"], c=color, lw=lw) - ax.set_ylabel("$T_i/T_e$") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - ax.axhline(y=1, ls="--", lw=1.0, c="k") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = axs6[5] - if "MachNum" in self.derived: - ax.plot(self.profiles["rho(-)"], self.derived["MachNum"], c=color, lw=lw) - ax.set_ylabel("Mach Number") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - ax.axhline(y=0, ls="--", c="k", lw=0.5) - ax.axhline(y=1, ls="--", c="k", lw=0.5) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = axs6[6] - safe_division = np.divide( - self.derived["qi_MWm2"], - self.derived["qe_MWm2"], - where=self.derived["qe_MWm2"] != 0, - out=np.full_like(self.derived["qi_MWm2"], np.nan), - ) - ax.plot( - self.profiles["rho(-)"], - safe_division, - c=color, - lw=lw, - label=extralab + "$Q_i/Q_e$", - ) - safe_division = np.divide( - self.derived["qi_aux_MWmiller"], - self.derived["qe_aux_MWmiller"], - where=self.derived["qe_aux_MWmiller"] != 0, - out=np.full_like(self.derived["qi_aux_MWmiller"], np.nan), - ) - ax.plot( - self.profiles["rho(-)"], - safe_division, - c=color, - lw=lw, - ls="--", - label=extralab + "$P_i/P_e$", - ) - ax.set_ylabel("Power ratios") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - ax.axhline(y=1.0, ls="--", c="k", lw=1.0) - GRAPHICStools.addDenseAxis(ax) - # GRAPHICStools.autoscale_y(ax,bottomy=0) - ax.set_ylim(bottom=0) - ax.legend(loc="best", fontsize=fs) - - # Final - if axsFlows is not None: - self.plotBalance( - axs=axsFlows, ls=lsFlows, leg=legFlows, showtexts=showtexts - ) - - # Geometry - ax = axs6[1] - self.plotGeometry(ax=ax, color=color) - - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - GRAPHICStools.addDenseAxis(ax) - - # Impurities - ax = axsImps[0] - for i in range(len(self.Species)): - var = ( - self.profiles["ni(10^19/m^3)"][:, i] - / self.profiles["ni(10^19/m^3)"][0, i] - ) - ax.plot( - rho, - var, - lw=lw, - ls=lines[i], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - varL = "$n_i/n_{i,0}$" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axsImps[1] - for i in range(len(self.Species)): - var = self.derived["fi"][:, i] - ax.plot( - rho, - var, - lw=lw, - ls=lines[i], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - varL = "$f_i$" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - ax.set_ylim([0, 1]) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axsImps[2] - - lastRho = 0.9 - - ix = np.argmin(np.abs(self.profiles["rho(-)"] - lastRho)) + 1 - ax.plot( - rho[:ix], self.derived["aLne"][:ix], lw=lw * 3, ls="-", c=color, label="e" - ) - for i in range(len(self.Species)): - var = self.derived["aLni"][:, i] - ax.plot( - rho[:ix], - var[:ix], - lw=lw, - ls=lines[i], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - varL = "$a/L_{ni}$" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axsImps[5] - - ax = axsImps[3] - ax.plot(self.profiles["rho(-)"], self.derived["Zeff"], c=color, lw=lw) - ax.set_ylabel("$Z_{eff}$") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axsImps[4] - cont = 0 - if "vtor(m/s)" in self.profiles: - for i in range(len(self.Species)): - try: # REMOVE FOR FUTURE - var = self.profiles["vtor(m/s)"][:, i] * 1e-3 - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - except: - break - varL = "$V_{tor}$ (km/s)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if "vtor(m/s)" in self.profiles and legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - def plotGradients( - self, - axs4, - color="b", - lw=1.0, - label="", - ls="-o", - lastRho=0.89, - ms=2, - alpha=1.0, - useRoa=False, - RhoLocationsPlot=None, - plotImpurity=None, - plotRotation=False, - autoscale=True, - ): - - if RhoLocationsPlot is None: RhoLocationsPlot=[] - - if axs4 is None: - plt.ion() - fig, axs = plt.subplots( - ncols=3 + int(plotImpurity is not None) + int(plotRotation), - nrows=2, - figsize=(12, 5), - ) - - axs4 = [] - for i in range(axs.shape[-1]): - axs4.append(axs[0, i]) - axs4.append(axs[1, i]) - - ix = np.argmin(np.abs(self.profiles["rho(-)"] - lastRho)) + 1 - - xcoord = self.profiles["rho(-)"] if (not useRoa) else self.derived["roa"] - labelx = "$\\rho$" if (not useRoa) else "$r/a$" - - ax = axs4[0] - ax.plot( - xcoord, - self.profiles["te(keV)"], - ls, - c=color, - lw=lw, - label=label, - markersize=ms, - alpha=alpha, - ) - ax = axs4[2] - ax.plot( - xcoord, - self.profiles["ti(keV)"][:, 0], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - ax = axs4[4] - ax.plot( - xcoord, - self.profiles["ne(10^19/m^3)"] * 1e-1, - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - - if "derived" in self.__dict__: - ax = axs4[1] - ax.plot( - xcoord[:ix], - self.derived["aLTe"][:ix], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - ax = axs4[3] - ax.plot( - xcoord[:ix], - self.derived["aLTi"][:ix, 0], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - ax = axs4[5] - ax.plot( - xcoord[:ix], - self.derived["aLne"][:ix], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - - for ax in axs4: - ax.set_xlim([0, 1]) - - ax = axs4[0] - ax.set_ylabel("$T_e$ (keV)") - ax.set_xlabel(labelx) - if autoscale: - 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) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax = axs4[4] - ax.set_ylabel("$n_e$ ($10^{20}m^{-3}$)") - ax.set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axs4[1] - ax.set_ylabel("$a/L_{Te}$") - ax.set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax = axs4[3] - ax.set_ylabel("$a/L_{Ti}$") - ax.set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax = axs4[5] - ax.set_ylabel("$a/L_{ne}$") - ax.axhline(y=0, ls="--", lw=0.5, c="k") - ax.set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - - cont = 0 - if plotImpurity is not None: - axs4[6 + cont].plot( - xcoord, - self.profiles["ni(10^19/m^3)"][:, plotImpurity] * 1e-1, - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - axs4[6 + cont].set_ylabel("$n_Z$ ($10^{20}m^{-3}$)") - axs4[6].set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - if "derived" in self.__dict__: - axs4[7 + cont].plot( - xcoord[:ix], - self.derived["aLni"][:ix, plotImpurity], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - axs4[7 + cont].set_ylabel("$a/L_{nZ}$") - axs4[7 + cont].axhline(y=0, ls="--", lw=0.5, c="k") - axs4[7 + cont].set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - cont += 2 - - if plotRotation: - axs4[6 + cont].plot( - xcoord, - self.profiles["w0(rad/s)"] * 1e-3, - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - axs4[6 + cont].set_ylabel("$w_0$ (krad/s)") - axs4[6 + cont].set_xlabel(labelx) - if "derived" in self.__dict__: - axs4[7 + cont].plot( - xcoord[:ix], - self.derived["dw0dr"][:ix] * 1e-5, - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - axs4[7 + cont].set_ylabel("-$d\\omega_0/dr$ (krad/s/cm)") - axs4[7 + cont].axhline(y=0, ls="--", lw=0.5, c="k") - axs4[7 + cont].set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - cont += 2 - - for x0 in RhoLocationsPlot: - ix = np.argmin(np.abs(self.profiles["rho(-)"] - x0)) - for ax in axs4: - ax.axvline(x=xcoord[ix], ls="--", lw=0.5, c=color) - - for i in range(len(axs4)): - ax = axs4[i] - GRAPHICStools.addDenseAxis(ax) - - def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): - if axs is None: - fig1 = plt.figure() - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - - axs = [ - fig1.add_subplot(grid[0, 0]), - fig1.add_subplot(grid[1, 0]), - fig1.add_subplot(grid[0, 1]), - fig1.add_subplot(grid[0, 2]), - fig1.add_subplot(grid[1, 1]), - fig1.add_subplot(grid[1, 2]), - ] - - # Profiles - - ax = axs[0] - axT = axs[1] - roa = self.profiles["rmin(m)"] / self.profiles["rmin(m)"][-1] - Te = self.profiles["te(keV)"] - ne = self.profiles["ne(10^19/m^3)"] * 1e-1 - ni = self.profiles["ni(10^19/m^3)"] * 1e-1 - niT = np.sum(ni, axis=1) - Ti = self.profiles["ti(keV)"][:, 0] - ax.plot(roa, Te, lw=2, c="r", label="$T_e$" if leg else "", ls=ls) - ax.plot(roa, Ti, lw=2, c="b", label="$T_i$" if leg else "", ls=ls) - axT.plot(roa, ne, lw=2, c="m", label="$n_e$" if leg else "", ls=ls) - axT.plot(roa, niT, lw=2, c="c", label="$\\sum n_i$" if leg else "", ls=ls) - if limits is not None: - [roa_first, roa_last] = limits - ax.plot(roa_last, np.interp(roa_last, roa, Te), "s", c="r", markersize=3) - ax.plot(roa_first, np.interp(roa_first, roa, Te), "s", c="r", markersize=3) - ax.plot(roa_last, np.interp(roa_last, roa, Ti), "s", c="b", markersize=3) - ax.plot(roa_first, np.interp(roa_first, roa, Ti), "s", c="b", markersize=3) - axT.plot(roa_last, np.interp(roa_last, roa, ne), "s", c="m", markersize=3) - axT.plot(roa_first, np.interp(roa_first, roa, ne), "s", c="m", markersize=3) - - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - axT.set_xlabel("r/a") - axT.set_xlim([0, 1]) - ax.set_ylabel("$T$ (keV)") - ax.set_ylim(bottom=0) - axT.set_ylabel("$n$ ($10^{20}m^{-3}$)") - axT.set_ylim(bottom=0) - # axT.set_ylim([0,np.max(ne)*1.5]) - ax.legend() - axT.legend() - ax.set_title("Final Temperature profiles") - axT.set_title("Final Density profiles") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - GRAPHICStools.addDenseAxis(axT) - GRAPHICStools.autoscale_y(axT, bottomy=0) - - if showtexts: - if self.derived["Q"] > 0.005: - ax.text( - 0.05, - 0.05, - f"Pfus = {self.derived['Pfus']:.1f}MW, Q = {self.derived['Q']:.2f}", - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - transform=ax.transAxes, - ) - - axT.text( - 0.05, - 0.4, - "ne_20 = {0:.1f} (fG = {1:.2f}), Zeff = {2:.1f}".format( - self.derived["ne_vol20"], - self.derived["fG"], - self.derived["Zeff_vol"], - ), - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - transform=axT.transAxes, - ) - - # F - ax = axs[2] - P = ( - self.derived["qe_fus_MWmiller"] - + self.derived["qe_aux_MWmiller"] - + -self.derived["qe_rad_MWmiller"] - + -self.derived["qe_exc_MWmiller"] - ) - - ax.plot( - roa, - -self.derived["qe_MWmiller"], - c="g", - lw=2, - label="$P_{e}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qe_fus_MWmiller"], - c="r", - lw=2, - label="$P_{fus,e}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qe_aux_MWmiller"], - c="b", - lw=2, - label="$P_{aux,e}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - -self.derived["qe_exc_MWmiller"], - c="m", - lw=2, - label="$P_{exc,e}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - -self.derived["qe_rad_MWmiller"], - c="c", - lw=2, - label="$P_{rad,e}$" if leg else "", - ls=ls, - ) - ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) - - # Pe = self.profiles['te(keV)']*1E3*e_J*self.profiles['ne(10^19/m^3)']*1E-1*1E20 *1E-6 - # ax.plot(roa,Pe,ls='-',lw=3,alpha=0.1,c='k',label='$W_e$ (MJ/m^3)') - - ax.plot( - roa, - -self.derived["ce_MWmiller"], - c="k", - lw=1, - label="($P_{conv,e}$)" if leg else "", - ) - - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - ax.set_ylabel("$P$ (MW)") - # ax.set_ylim(bottom=0) - ax.set_title("Electron Thermal Flows") - ax.axhline(y=0.0, lw=0.5, ls="--", c="k") - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" - ) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axs[3] - P = ( - self.derived["qi_fus_MWmiller"] - + self.derived["qi_aux_MWmiller"] - + self.derived["qe_exc_MWmiller"] - ) - - ax.plot( - roa, - -self.derived["qi_MWmiller"], - c="g", - lw=2, - label="$P_{i}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qi_fus_MWmiller"], - c="r", - lw=2, - label="$P_{fus,i}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qi_aux_MWmiller"], - c="b", - lw=2, - label="$P_{aux,i}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qe_exc_MWmiller"], - c="m", - lw=2, - label="$P_{exc,i}$" if leg else "", - ls=ls, - ) - ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) - - # Pi = self.profiles['ti(keV)'][:,0]*1E3*e_J*self.profiles['ni(10^19/m^3)'][:,0]*1E-1*1E20 *1E-6 - # ax.plot(roa,Pi,ls='-',lw=3,alpha=0.1,c='k',label='$W_$ (MJ/m^3)') - - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - ax.set_ylabel("$P$ (MW)") - # ax.set_ylim(bottom=0) - ax.set_title("Ion Thermal Flows") - ax.axhline(y=0.0, lw=0.5, ls="--", c="k") - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" - ) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - # F - ax = axs[4] - - ax.plot( - roa, - self.derived["ge_10E20miller"], - c="g", - lw=2, - label="$\\Gamma_{e}$" if leg else "", - ls=ls, - ) - # ax.plot(roa,self.profiles['ne(10^19/m^3)']*1E-1,lw=3,alpha=0.1,c='k',label='$n_e$ ($10^{20}/m^3$)' if leg else '',ls=ls) - - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - ax.set_ylabel("$\\Gamma$ ($10^{20}/s$)") - ax.set_title("Particle Flows") - ax.axhline(y=0.0, lw=0.5, ls="--", c="k") - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" - ) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - # TOTAL - ax = axs[5] - P = ( - self.derived["qOhm_MWmiller"] - + self.derived["qRF_MWmiller"] - + self.derived["qFus_MWmiller"] - + -self.derived["qe_rad_MWmiller"] - + self.derived["qz_MWmiller"] - + self.derived["qBEAM_MWmiller"] - ) - - ax.plot( - roa, - -self.derived["q_MWmiller"], - c="g", - lw=2, - label="$P$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qOhm_MWmiller"], - c="k", - lw=2, - label="$P_{Oh}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qRF_MWmiller"], - c="b", - lw=2, - label="$P_{RF}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qBEAM_MWmiller"], - c="pink", - lw=2, - label="$P_{NBI}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qFus_MWmiller"], - c="r", - lw=2, - label="$P_{fus}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - -self.derived["qe_rad_MWmiller"], - c="c", - lw=2, - label="$P_{rad}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qz_MWmiller"], - c="orange", - lw=1, - label="$P_{ionz.}$" if leg else "", - ls=ls, - ) - - # P = Pe+Pi - # ax.plot(roa,P,ls='-',lw=3,alpha=0.1,c='k',label='$W$ (MJ)') - - ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - ax.set_ylabel("$P$ (MW)") - # ax.set_ylim(bottom=0) - ax.set_title("Total Thermal Flows") - - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" - ) - - ax.axhline(y=0.0, lw=0.5, ls="--", c="k") - # GRAPHICStools.drawLineWithTxt(ax,0.0,label='',orientation='vertical',color='k',lw=1,ls='--',alpha=1.0,fontsize=10,fromtop=0.85,fontweight='normal', - # verticalalignment='bottom',horizontalalignment='left',separation=0) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - def plot_temps(self, ax=None, leg=False, col="b", lw=2, extralab="", fs=10): - if ax is None: - fig, ax = plt.subplots() - - rho = self.profiles["rho(-)"] - - var = self.profiles["te(keV)"] - varL = "$T_e$ , $T_i$ (keV)" - if leg: - lab = extralab + "e" - else: - lab = "" - ax.plot(rho, var, lw=lw, ls="-", label=lab, c=col) - var = self.profiles["ti(keV)"][:, 0] - if leg: - lab = extralab + "i" - else: - lab = "" - ax.plot(rho, var, lw=lw, ls="--", label=lab, c=col) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - - def plot_dens(self, ax=None, leg=False, col="b", lw=2, extralab="", fs=10): - if ax is None: - fig, ax = plt.subplots() - - rho = self.profiles["rho(-)"] - - var = self.profiles["ne(10^19/m^3)"] * 1e-1 - varL = "$n_e$ ($10^{20}/m^3$)" - if leg: - lab = extralab + "e" - else: - lab = "" - ax.plot(rho, var, lw=lw, ls="-", label=lab, c=col) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - - def plotGeometry(self, ax=None, surfaces_rho=np.linspace(0, 1, 11), color="b", label = '', lw=1.0, lw1=2.0): - if ("R_surface" in self.derived) and (self.derived["R_surface"] is not None): - if ax is None: - plt.ion() - fig, ax = plt.subplots() - provided = False - else: - provided = True - - for rho in surfaces_rho: - ir = np.argmin(np.abs(self.profiles["rho(-)"] - rho)) - - ax.plot( - self.derived["R_surface"][ir, :], - self.derived["Z_surface"][ir, :], - "-", - lw=lw if rho<1.0 else lw1, - c=color, - ) - - ax.axhline(y=0, ls="--", lw=0.2, c="k") - ax.plot( - [self.profiles["rmaj(m)"][0]], - [self.profiles["zmag(m)"][0]], - "o", - markersize=2, - c=color, - label = label - ) - - if not provided: - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - ax.set_title("Surfaces @ rho=" + str(surfaces_rho), fontsize=8) - ax.set_aspect("equal") - else: - print("\t- Cannot plot flux surface geometry", typeMsg="w") - - def plotPeaking( - self, ax, c="b", marker="*", label="", debugPlot=False, printVals=False - ): - nu_effCGYRO = self.derived["nu_eff"] * 2 / self.derived["Zeff_vol"] - ne_peaking = self.derived["ne_peaking0.2"] - ax.scatter([nu_effCGYRO], [ne_peaking], s=400, c=c, marker=marker, label=label) - - if printVals: - print(f"\t- nu_eff = {nu_effCGYRO}, ne_peaking = {ne_peaking}") - - # Extra - r = self.profiles["rmin(m)"] - volp = self.derived["volp_miller"] - ix = np.argmin(np.abs(self.profiles["rho(-)"] - 0.9)) - - if debugPlot: - fig, axq = plt.subplots() - - ne = self.profiles["ne(10^19/m^3)"] - axq.plot(self.profiles["rho(-)"], ne, color="m") - ne_vol = ( - CALCtools.integrateFS(ne * 0.1, r, volp)[-1] / self.derived["volume"] - ) - axq.axhline(y=ne_vol * 10, color="m") - - ne = copy.deepcopy(self.profiles["ne(10^19/m^3)"]) - ne[ix:] = (0,) * len(ne[ix:]) - ne_vol = CALCtools.integrateFS(ne * 0.1, r, volp)[-1] / self.derived["volume"] - ne_peaking0 = ( - ne[np.argmin(np.abs(self.derived["rho_pol"] - 0.2))] * 0.1 / ne_vol - ) - - if debugPlot: - axq.plot(self.profiles["rho(-)"], ne, color="r") - axq.axhline(y=ne_vol * 10, color="r") - - ne = copy.deepcopy(self.profiles["ne(10^19/m^3)"]) - ne[ix:] = (ne[ix],) * len(ne[ix:]) - ne_vol = CALCtools.integrateFS(ne * 0.1, r, volp)[-1] / self.derived["volume"] - ne_peaking1 = ( - ne[np.argmin(np.abs(self.derived["rho_pol"] - 0.2))] * 0.1 / ne_vol - ) - - ne_peaking0 = ne_peaking - - ax.errorbar( - [nu_effCGYRO], - [ne_peaking], - yerr=[[ne_peaking - ne_peaking1], [ne_peaking0 - ne_peaking]], - marker=marker, - c=c, - markersize=16, - capsize=2.0, - fmt="s", - elinewidth=1.0, - capthick=1.0, - ) - - if debugPlot: - axq.plot(self.profiles["rho(-)"], ne, color="b") - axq.axhline(y=ne_vol * 10, color="b") - plt.show() - - # print(f'{ne_peaking0}-{ne_peaking}-{ne_peaking1}') - - return nu_effCGYRO, ne_peaking - - def plotRelevant(self, axs = None, color = 'b', label ='', lw = 1, ms = 1): - - if axs is None: - fig = plt.figure() - axs = fig.subplot_mosaic( - """ - ABCDH - AEFGI - """ - ) - axs = [axs['A'], axs['B'], axs['C'], axs['D'], axs['E'], axs['F'], axs['G'], axs['H'], axs['I']] - - # ---------------------------------- - # Equilibria - # ---------------------------------- - - ax = axs[0] - rho = np.linspace(0, 1, 21) - - self.plotGeometry(ax=ax, surfaces_rho=rho, label=label, color=color, lw=lw, lw1=lw*3) - - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - ax.set_aspect("equal") - ax.legend(prop={'size':8}) - GRAPHICStools.addDenseAxis(ax) - ax.set_title("Equilibria") - - # ---------------------------------- - # Kinetic Profiles - # ---------------------------------- - - # T profiles - ax = axs[1] - - ax.plot(self.profiles['rho(-)'], self.profiles['te(keV)'], '-o', markersize=ms, lw = lw, label=label+', e', color=color) - ax.plot(self.profiles['rho(-)'], self.profiles['ti(keV)'][:,0], '--*', markersize=ms, lw = lw, label=label+', i', color=color) - - ax.set_xlabel("$\\rho_N$") - ax.set_ylabel("$T$ (keV)") - #ax.set_ylim(bottom = 0) - ax.set_xlim(0,1) - ax.legend(prop={'size':8}) - GRAPHICStools.addDenseAxis(ax) - ax.set_title("Temperatures") - - # ne profiles - ax = axs[2] - - ax.plot(self.profiles['rho(-)'], self.profiles['ne(10^19/m^3)']*1E-1, '-o', markersize=ms, lw = lw, label=label, color=color) - - ax.set_xlabel("$\\rho_N$") - ax.set_ylabel("$n_e$ ($10^{20}m^{-3}$)") - #ax.set_ylim(bottom = 0) - ax.set_xlim(0,1) - ax.legend(prop={'size':8}) - GRAPHICStools.addDenseAxis(ax) - ax.set_title("Electron Density") - - # ---------------------------------- - # Pressure - # ---------------------------------- - - ax = axs[3] - - ax.plot(self.profiles['rho(-)'], self.derived['ptot_manual'], '-o', markersize=ms, lw = lw, label=label, color=color) - - ax.set_xlabel("$\\rho_N$") - ax.set_ylabel("$p_{kin}$ (MPa)") - #ax.set_ylim(bottom = 0) - ax.set_xlim(0,1) - ax.legend(prop={'size':8}) - GRAPHICStools.addDenseAxis(ax) - ax.set_title("Total Pressure") - - # ---------------------------------- - # Current - # ---------------------------------- - - # q-profile - ax = axs[4] - - ax.plot(self.profiles['rho(-)'], self.profiles['q(-)'], '-o', markersize=ms, lw = lw, label=label, color=color) - - ax.set_xlabel("$\\rho_N$") - ax.set_ylabel("$q$") - #ax.set_ylim(bottom = 0) - ax.set_xlim(0,1) - ax.legend(prop={'size':8}) - GRAPHICStools.addDenseAxis(ax) - ax.set_title("Safety Factor") - - # ---------------------------------- - # Powers - # ---------------------------------- - - # RF - ax = axs[5] - - ax.plot(self.profiles['rho(-)'], self.profiles['qrfe(MW/m^3)'], '-o', markersize=ms, lw = lw, label=label+', e', color=color) - ax.plot(self.profiles['rho(-)'], self.profiles['qrfi(MW/m^3)'], '--*', markersize=ms, lw = lw, label=label+', i', color=color) - - ax.set_xlabel("$\\rho_N$") - ax.set_ylabel("$P_{ich}$ (MW/m$^3$)") - #ax.set_ylim(bottom = 0) - ax.set_xlim(0,1) - ax.legend(prop={'size':8}) - GRAPHICStools.addDenseAxis(ax) - ax.set_title("ICH Power Deposition") - - # Ohmic - ax = axs[6] - - ax.plot(self.profiles['rho(-)'], self.profiles['qohme(MW/m^3)'], '-o', markersize=ms, lw = lw, label=label, color=color) - - ax.set_xlabel("$\\rho_N$") - ax.set_ylabel("$P_{oh}$ (MW/m$^3$)") - #ax.set_ylim(bottom = 0) - ax.set_xlim(0,1) - ax.legend(prop={'size':8}) - GRAPHICStools.addDenseAxis(ax) - ax.set_title("Ohmic Power Deposition") - - # ---------------------------------- - # Heat fluxes - # ---------------------------------- - - ax = axs[7] - - ax.plot(self.profiles['rho(-)'], self.derived['qe_MWm2'], '-o', markersize=ms, lw = lw, label=label+', e', color=color) - ax.plot(self.profiles['rho(-)'], self.derived['qi_MWm2'], '--*', markersize=ms, lw = lw, label=label+', i', color=color) - - ax.set_xlabel("$\\rho_N$") - ax.set_ylabel("$Q$ ($MW/m^2$)") - #ax.set_ylim(bottom = 0) - ax.set_xlim(0,1) - ax.legend(prop={'size':8}) - GRAPHICStools.addDenseAxis(ax) - ax.set_title("Energy Fluxes") - - # ---------------------------------- - # Dynamic targets - # ---------------------------------- - - ax = axs[8] - - ax.plot(self.profiles['rho(-)'], self.derived['qrad'], '-o', markersize=ms, lw = lw, label=label+', rad', color=color) - ax.plot(self.profiles['rho(-)'], self.profiles['qei(MW/m^3)'], '--*', markersize=ms, lw = lw, label=label+', exc', color=color) - if 'qfuse(MW/m^3)' in self.profiles: - ax.plot(self.profiles['rho(-)'], self.profiles['qfuse(MW/m^3)']+self.profiles['qfusi(MW/m^3)'], '-.s', markersize=ms, lw = lw, label=label+', fus', color=color) - - ax.set_xlabel("$\\rho_N$") - ax.set_ylabel("$Q$ ($MW/m^2$)") - #ax.set_ylim(bottom = 0) - ax.set_xlim(0,1) - ax.legend(prop={'size':8}) - GRAPHICStools.addDenseAxis(ax) - ax.set_title("Dynamic Targets") - - - def csv(self, file="input.gacode.xlsx"): - dictExcel = IOtools.OrderedDict() - - for ikey in self.profiles: - print(ikey) - if len(self.profiles[ikey].shape) == 1: - dictExcel[ikey] = self.profiles[ikey] - else: - dictExcel[ikey] = self.profiles[ikey][:, 0] - - IOtools.writeExcel_fromDict(dictExcel, file, fromRow=1) - - def parabolizePlasma(self): - _, T = PLASMAtools.parabolicProfile( - Tbar=self.derived["Te_vol"], - nu=self.derived["Te_peaking"], - rho=self.profiles["rho(-)"], - Tedge=self.profiles["te(keV)"][-1], - ) - _, Ti = PLASMAtools.parabolicProfile( - Tbar=self.derived["Ti_vol"], - nu=self.derived["Ti_peaking"], - rho=self.profiles["rho(-)"], - Tedge=self.profiles["ti(keV)"][-1, 0], - ) - _, n = PLASMAtools.parabolicProfile( - Tbar=self.derived["ne_vol20"] * 1e1, - nu=self.derived["ne_peaking"], - rho=self.profiles["rho(-)"], - Tedge=self.profiles["ne(10^19/m^3)"][-1], - ) - - self.profiles["te(keV)"] = T - - self.profiles["ti(keV)"][:, 0] = Ti - self.makeAllThermalIonsHaveSameTemp(refIon=0) - - factor_n = n / self.profiles["ne(10^19/m^3)"] - self.profiles["ne(10^19/m^3)"] = n - self.scaleAllThermalDensities(scaleFactor=factor_n) - - self.deriveQuantities() - - - def changeRFpower(self, PrfMW=25.0): - """ - keeps same partition - """ - print(f"- Changing the RF power from {self.derived['qRF_MWmiller'][-1]:.1f} MW to {PrfMW:.1f} MW",typeMsg="i",) - - if self.derived["qRF_MWmiller"][-1] == 0.0: - raise Exception("No RF power in the input.gacode, cannot modify the RF power") - - for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: - self.profiles[i] = self.profiles[i] * PrfMW / self.derived["qRF_MWmiller"][-1] - - self.deriveQuantities() - - def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0.2): - - ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) - - self.profiles["te(keV)"] = self.profiles["te(keV)"] * TkeV / self.profiles["te(keV)"][ix] - - print(f"- Producing {typeEdge} boundary condition @ rho = {rho}, T = {TkeV} keV",typeMsg="i",) - - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm": - self.profiles["ti(keV)"][:, sp] = self.profiles["ti(keV)"][:, sp] * TkeV / self.profiles["ti(keV)"][ix, sp] - - if typeEdge == "linear": - self.profiles["te(keV)"][ix:] = np.linspace(TkeV, Tesep, len(self.profiles["rho(-)"][ix:])) - - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm": - self.profiles["ti(keV)"][ix:, sp] = np.linspace(TkeV, Tisep, len(self.profiles["rho(-)"][ix:])) - - elif typeEdge == "same": - pass - else: - raise Exception("no edge") - - - def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5, isn20_edge=True): - ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) - - # Determine the factor to scale the density (either average or at rho) - if not isn20_edge: - print(f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") - factor = n20 / self.derived["ne_vol20"] - else: - print(f"- Changing the density at rho={rho} from {self.profiles['ne(10^19/m^3)'][ix]*1E-1:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") - factor = n20 / (self.profiles["ne(10^19/m^3)"][ix]*1E-1) - # ------------------------------------------------------------------ - - # Scale the density profiles - for i in ["ne(10^19/m^3)", "ni(10^19/m^3)"]: - self.profiles[i] = self.profiles[i] * factor - - # Apply the edge condition - if typeEdge == "linear": - factor_x = np.linspace(self.profiles["ne(10^19/m^3)"][ix],nedge20 * 1e1,len(self.profiles["rho(-)"][ix:]),)/ self.profiles["ne(10^19/m^3)"][ix:] - - self.profiles["ne(10^19/m^3)"][ix:] = self.profiles["ne(10^19/m^3)"][ix:] * factor_x - - for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): - self.profiles["ni(10^19/m^3)"][ix:, i] = self.profiles["ni(10^19/m^3)"][ix:, i] * factor_x - - elif typeEdge == "same": - pass - else: - raise Exception("no edge") - - def addSawtoothEffectOnOhmic(self, PohTot, mixRadius=None, plotYN=False): - """ - This will implement a flat profile inside the mixRadius to reduce the ohmic power by certain amount - """ - - if mixRadius is None: - mixRadius = self.profiles["rho(-)"][np.where(self.profiles["q(-)"] > 1)][0] - - print(f"\t- Original Ohmic power: {self.derived['qOhm_MWmiller'][-1]:.2f}MW") - Ohmic_old = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) - - dvol = self.derived["volp_miller"] * np.append( - [0], np.diff(self.profiles["rmin(m)"]) - ) - - print( - f"\t- Will implement sawtooth ohmic power correction inside rho={mixRadius}" - ) - Psaw = CDFtools.profilePower( - self.profiles["rho(-)"], - dvol, - PohTot - self.derived["qOhm_MWmiller"][-1], - mixRadius, - ) - self.profiles["qohme(MW/m^3)"] += Psaw - self.deriveQuantities() - - print(f"\t- New Ohmic power: {self.derived['qOhm_MWmiller'][-1]:.2f}MW") - Ohmic_new = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) - - if plotYN: - fig, ax = plt.subplots() - ax.plot(self.profiles["rho(-)"], Ohmic_old, "r", lw=2) - ax.plot(self.profiles["rho(-)"], Ohmic_new, "g", lw=2) - plt.show() - - # --------------------------------------------------------------------------------------------------------------------------------------- - # Code conversions - # --------------------------------------------------------------------------------------------------------------------------------------- - - def to_tglf(self, rhos=[0.5], TGLFsettings=1): - - # <> Function to interpolate a curve <> - from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function - - inputsTGLF = {} - for rho in rhos: - - # --------------------------------------------------------------------------------------------------------------------------------------- - # Define interpolator at this rho - # --------------------------------------------------------------------------------------------------------------------------------------- - - def interpolator(y): - return interpolation_function(rho, self.profiles['rho(-)'],y).item() - - TGLFinput, TGLFoptions, label = GACODEdefaults.addTGLFcontrol(TGLFsettings) - - # --------------------------------------------------------------------------------------------------------------------------------------- - # Controls come from options - # --------------------------------------------------------------------------------------------------------------------------------------- - - controls = TGLFoptions - - # --------------------------------------------------------------------------------------------------------------------------------------- - # Species come from profiles - # --------------------------------------------------------------------------------------------------------------------------------------- - - #mass_ref = self.derived["mi_ref"] - # input.gacode uses the deuterium mass as reference already (https://github.com/gafusion/gacode/issues/398), so this should be 2.0 - mass_ref = 2.0 - - mass_e = 0.000272445 * mass_ref - - species = { - 1: { - 'ZS': -1.0, - 'MASS': mass_e/mass_ref, - 'RLNS': interpolator(self.derived['aLne']), - 'RLTS': interpolator(self.derived['aLTe']), - 'TAUS': 1.0, - 'AS': 1.0, - 'VPAR': interpolator(self.derived['vpar']), - 'VPAR_SHEAR': interpolator(self.derived['vpar_shear']), - 'VNS_SHEAR': 0.0, - 'VTS_SHEAR': 0.0}, - } - - for i in range(len(self.Species)): - species[i+2] = { - 'ZS': self.Species[i]['Z'], - 'MASS': self.Species[i]['A']/mass_ref, - 'RLNS': interpolator(self.derived['aLni'][:,i]), - 'RLTS': interpolator(self.derived['aLTi'][:,0] if self.Species[i]['S'] == 'therm' else self.derived["aLTi"][:,i]), - 'TAUS': interpolator(self.derived["tite_all"][:,i]), - 'AS': interpolator(self.derived['fi'][:,i]), - 'VPAR': interpolator(self.derived['vpar']), - 'VPAR_SHEAR': interpolator(self.derived['vpar_shear']), - 'VNS_SHEAR': 0.0, - 'VTS_SHEAR': 0.0 - } - - # --------------------------------------------------------------------------------------------------------------------------------------- - # Plasma comes from profiles - # --------------------------------------------------------------------------------------------------------------------------------------- - - plasma = { - 'NS': len(species)+1, - 'SIGN_BT': -1.0, - 'SIGN_IT': -1.0, - 'VEXB': 0.0, - 'VEXB_SHEAR': interpolator(self.derived['vexb_shear']), - 'BETAE': interpolator(self.derived['betae']), - 'XNUE': interpolator(self.derived['xnue']), - 'ZEFF': interpolator(self.derived['Zeff']), - 'DEBYE': interpolator(self.derived['debye']), - } - - # --------------------------------------------------------------------------------------------------------------------------------------- - # Geometry comes from profiles - # --------------------------------------------------------------------------------------------------------------------------------------- - - parameters = { - 'RMIN_LOC': self.derived['roa'], - 'RMAJ_LOC': self.derived['Rmajoa'], - 'ZMAJ_LOC': self.derived["Zmagoa"], - 'DRMINDX_LOC': self.derived['drmin/dr'], - 'DRMAJDX_LOC': self.derived['dRmaj/dr'], - 'DZMAJDX_LOC': self.derived['dZmaj/dr'], - 'Q_LOC': self.profiles["q(-)"], - 'KAPPA_LOC': self.profiles["kappa(-)"], - 'S_KAPPA_LOC': self.derived['s_kappa'], - 'DELTA_LOC': self.profiles["delta(-)"], - 'S_DELTA_LOC': self.derived['s_delta'], - 'ZETA_LOC': self.profiles["zeta(-)"], - 'S_ZETA_LOC': self.derived['s_zeta'], - 'P_PRIME_LOC': self.derived['pprime'], - 'Q_PRIME_LOC': self.derived['s_q'], - } - - geom = {} - for k in parameters: - par = torch.nan_to_num(torch.from_numpy(parameters[k]) if type(parameters[k]) is np.ndarray else parameters[k], nan=0.0, posinf=1E10, neginf=-1E10) - geom[k] = interpolator(par) - - geom['BETA_LOC'] = 0.0 - geom['KX0_LOC'] = 0.0 - - # --------------------------------------------------------------------------------------------------------------------------------------- - # Merging - # --------------------------------------------------------------------------------------------------------------------------------------- - - input_dict = {**controls, **plasma, **geom} - - for i in range(len(species)): - for k in species[i+1]: - input_dict[f'{k}_{i+1}'] = species[i+1][k] - - inputsTGLF[rho] = input_dict - - return inputsTGLF - - def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', times = [0.0,1.0], Vsurf = 0.0): - - print("\t- Converting to TRANSP") - folder = IOtools.expandPath(folder) - folder.mkdir(parents=True, exist_ok=True) - - transp = TRANSPhelpers.transp_run(folder, shot, runid) - for time in times: - transp.populate_time.from_profiles(time,self, Vsurf = Vsurf) - - transp.write_ufiles() - - return transp - - def to_eped(self, ped_rho = 0.95): - - neped_19 = np.interp(ped_rho, self.profiles['rho(-)'], self.profiles['ne(10^19/m^3)']) - - eped_evaluation = { - 'Ip': np.abs(self.profiles['current(MA)'][0]), - 'Bt': np.abs(self.profiles['bcentr(T)'][0]), - 'R': np.abs(self.profiles['rcentr(m)'][0]), - 'a': np.abs(self.derived['a']), - 'kappa995': np.abs(self.derived['kappa995']), - 'delta995': np.abs(self.derived['delta995']), - 'neped': np.abs(neped_19), - 'betan': np.abs(self.derived['BetaN_engineering']), - 'zeff': np.abs(self.derived['Zeff_vol']), - 'tesep': np.abs(self.profiles['te(keV)'][-1])*1E3, - 'nesep_ratio': np.abs(self.profiles['ne(10^19/m^3)'][-1] / neped_19), - } - - return eped_evaluation - -class DataTable: - def __init__(self, variables=None): - - if variables is not None: - self.variables = variables - else: - - # Default for confinement mode access studies (JWH 03/2024) - self.variables = { - "Rgeo": ["rcentr(m)", "pos_0", "profiles", ".2f", 1, "m"], - "ageo": ["a", None, "derived", ".2f", 1, "m"], - "volume": ["volume", None, "derived", ".2f", 1, "m"], - "kappa @psi=0.95": ["kappa(-)", "psi_0.95", "profiles", ".2f", 1, None], - "delta @psi=0.95": ["delta(-)", "psi_0.95", "profiles", ".2f", 1, None], - "Bt": ["bcentr(T)", "pos_0", "profiles", ".1f", 1, "T"], - "Ip": ["current(MA)", "pos_0", "profiles", ".1f", 1, "MA"], - "Pin": ["qIn", None, "derived", ".1f", 1, "MW"], - "Te @rho=0.9": ["te(keV)", "rho_0.90", "profiles", ".2f", 1, "keV"], - "Ti/Te @rho=0.9": ["tite", "rho_0.90", "derived", ".2f", 1, None], - "ne @rho=0.9": [ - "ne(10^19/m^3)", - "rho_0.90", - "profiles", - ".2f", - 0.1, - "E20m-3", - ], - "ptot @rho=0.9": [ - "ptot_manual", - "rho_0.90", - "derived", - ".1f", - 1e3, - "kPa", - ], - "Zeff": ["Zeff_vol", None, "derived", ".1f", 1, None], - "fDT": ["fmain", None, "derived", ".2f", 1, None], - "H89p": ["H89", None, "derived", ".2f", 1, None], - "H98y2": ["H98", None, "derived", ".2f", 1, None], - "ne (vol avg)": ["ne_vol20", None, "derived", ".2f", 1, "E20m-3"], - "Ptop": ["ptop", None, "derived", ".1f", 1, "Pa"], - "fG": ["fG", None, "derived", ".2f", 1, None], - "Pfus": ["Pfus", None, "derived", ".1f", 1, "MW"], - "Prad": ["Prad", None, "derived", ".1f", 1, "MW"], - "Q": ["Q", None, "derived", ".2f", 1, None], - "Pnet @rho=0.9": ["qTr", "rho_0.90", "derived", ".1f", 1, "MW"], - "Qi/Qe @rho=0.9": ["QiQe", "rho_0.90", "derived", ".2f", 1, None], - } - - self.data = [] - - def export_to_csv(self, filename, title=None): - - title_data = [""] - for key in self.variables: - if self.variables[key][5] is None: - title_data.append(f"{key}") - else: - title_data.append(f"{key} ({self.variables[key][5]})") - - # Open a file with the given filename in write mode - with open(filename, mode="w", newline="") as file: - writer = csv.writer(file) - - # Write the title row first if it is provided - if title: - writer.writerow([title] + [""] * (len(self.data[0]) - 1)) - - writer.writerow(title_data) - - # Write each row in self.data to the CSV file - for row in self.data: - writer.writerow(row) - -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: - from mitim_tools.misc_tools.GUItools import FigureNotebook - - fn = FigureNotebook("Profiles", geometry="1800x900") - figProf_1, figProf_2, figProf_3, figProf_4, figFlows, figProf_6, fig7 = add_figures(fn) - - grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) - axsProf_1 = [ - figProf_1.add_subplot(grid[0, 0]), - figProf_1.add_subplot(grid[1, 0]), - figProf_1.add_subplot(grid[2, 0]), - figProf_1.add_subplot(grid[0, 1]), - figProf_1.add_subplot(grid[1, 1]), - figProf_1.add_subplot(grid[2, 1]), - figProf_1.add_subplot(grid[0, 2]), - figProf_1.add_subplot(grid[1, 2]), - figProf_1.add_subplot(grid[2, 2]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsProf_2 = [ - figProf_2.add_subplot(grid[0, 0]), - figProf_2.add_subplot(grid[0, 1]), - figProf_2.add_subplot(grid[1, 0]), - figProf_2.add_subplot(grid[1, 1]), - figProf_2.add_subplot(grid[0, 2]), - figProf_2.add_subplot(grid[1, 2]), - ] - grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.3) - ax00c = figProf_3.add_subplot(grid[0, 0]) - axsProf_3 = [ - ax00c, - figProf_3.add_subplot(grid[1, 0], sharex=ax00c), - figProf_3.add_subplot(grid[2, 0]), - figProf_3.add_subplot(grid[0, 1]), - figProf_3.add_subplot(grid[1, 1]), - figProf_3.add_subplot(grid[2, 1]), - figProf_3.add_subplot(grid[0, 2]), - figProf_3.add_subplot(grid[1, 2]), - figProf_3.add_subplot(grid[2, 2]), - figProf_3.add_subplot(grid[0, 3]), - figProf_3.add_subplot(grid[1, 3]), - figProf_3.add_subplot(grid[2, 3]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsProf_4 = [ - figProf_4.add_subplot(grid[0, 0]), - figProf_4.add_subplot(grid[1, 0]), - figProf_4.add_subplot(grid[0, 1]), - figProf_4.add_subplot(grid[1, 1]), - figProf_4.add_subplot(grid[0, 2]), - figProf_4.add_subplot(grid[1, 2]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsFlows = [ - figFlows.add_subplot(grid[0, 0]), - figFlows.add_subplot(grid[1, 0]), - figFlows.add_subplot(grid[0, 1]), - figFlows.add_subplot(grid[0, 2]), - figFlows.add_subplot(grid[1, 1]), - figFlows.add_subplot(grid[1, 2]), - ] - - grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) - axsProf_6 = [ - figProf_6.add_subplot(grid[0, 0]), - figProf_6.add_subplot(grid[:, 1]), - figProf_6.add_subplot(grid[0, 2]), - figProf_6.add_subplot(grid[1, 0]), - figProf_6.add_subplot(grid[1, 2]), - figProf_6.add_subplot(grid[0, 3]), - figProf_6.add_subplot(grid[1, 3]), - ] - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsImps = [ - fig7.add_subplot(grid[0, 0]), - fig7.add_subplot(grid[0, 1]), - fig7.add_subplot(grid[0, 2]), - fig7.add_subplot(grid[1, 0]), - fig7.add_subplot(grid[1, 1]), - fig7.add_subplot(grid[1, 2]), - ] - - ls = GRAPHICStools.listLS() - colors = GRAPHICStools.listColors() - for i, profiles in enumerate(profiles_list): - if extralabs is None: - extralab = f"#{i}, " - else: - extralab = f"{extralabs[i]}, " - profiles.plot( - axs1=axsProf_1, - axs2=axsProf_2, - axs3=axsProf_3, - axs4=axsProf_4, - axsFlows=axsFlows, - axs6=axsProf_6, - axsImps=axsImps, - color=colors[i], - legYN=True, - extralab=extralab, - lsFlows=ls[i], - legFlows=i == 0, - showtexts=False, - lastRhoGradients=lastRhoGradients, - ) - - return fn - - -def readTGYRO_profile_extra(file, varLabel="B_unit (T)"): - with open(file) as f: - aux = f.readlines() - - lenn = int(aux[36].split()[-1]) - - i = 38 - allVec = [] - while i < len(aux): - vec = np.array([float(j) for j in aux[i : i + lenn]]) - i += lenn - allVec.append(vec) - allVec = np.array(allVec) - - dictL = OrderedDict() - for line in aux[2:35]: - lab = line.split("(:)")[-1].split("\n")[0] - try: - dictL[lab] = int(line.split()[1]) - except: - dictL[lab] = [int(j) for j in line.split()[1].split("-")] - - for i in dictL: - if i.strip(" ") == varLabel: - val = allVec[dictL[i] - 1] - break - - return val - - -def aLT(r, p): - return ( - r[-1] - * CALCtools.produceGradient( - torch.from_numpy(r).to(torch.double), torch.from_numpy(p).to(torch.double) - ) - .cpu() - .cpu().numpy() - ) - - -def grad(r, p): - return MATHtools.deriv(torch.from_numpy(r), torch.from_numpy(p), array=False) - - -def ionName(Z, A): - # Based on Z - if Z == 2: - return "He" - elif Z == 9: - return "F" - elif Z == 6: - return "C" - elif Z == 11: - return "Na" - elif Z == 30: - return "Zn" - elif Z == 31: - return "Ga" - - # # Based on Mass (this is the correct way, since the radiation needs to be calculated with the full element) - # if A in [3,4]: return 'He' - # elif A == 18: return 'F' - # elif A == 12: return 'C' - # elif A == 22: return 'Na' - # elif A == 60: return 'Zn' - # elif A == 69: return 'Ga' - - -def gradientsMerger(p0, p_true, roa=0.46, blending=0.1): - p = copy.deepcopy(p0) - - aLTe_true = np.interp( - p.derived["roa"], p_true.derived["roa"], p_true.derived["aLTe"] - ) - aLTi_true = np.interp( - p.derived["roa"], p_true.derived["roa"], p_true.derived["aLTi"][:, 0] - ) - aLne_true = np.interp( - p.derived["roa"], p_true.derived["roa"], p_true.derived["aLne"] - ) - - ix1 = np.argmin(np.abs(p.derived["roa"] - roa + blending)) - ix2 = np.argmin(np.abs(p.derived["roa"] - roa)) - - aLT0 = aLTe_true[: ix1 + 1] - aLT2 = p.derived["aLTe"][ix2:] - aLT1 = np.interp( - p.derived["roa"][ix1 : ix2 + 1], - [p.derived["roa"][ix1], p.derived["roa"][ix2]], - [aLT0[-1], aLT2[0]], - )[1:-1] - - aLTe = np.append(np.append(aLT0, aLT1), aLT2) - Te = ( - CALCtools.integrateGradient( - torch.from_numpy(p.derived["roa"]).unsqueeze(0), - torch.Tensor(aLTe).unsqueeze(0), - p.profiles["te(keV)"][-1], - ) - .cpu() - .cpu().numpy()[0] - ) - - aLT0 = aLTi_true[: ix1 + 1] - aLT2 = p.derived["aLTi"][ix2:, 0] - aLT1 = np.interp( - p.derived["roa"][ix1 : ix2 + 1], - [p.derived["roa"][ix1], p.derived["roa"][ix2]], - [aLT0[-1], aLT2[0]], - )[1:-1] - - aLTi = np.append(np.append(aLT0, aLT1), aLT2) - Ti = ( - CALCtools.integrateGradient( - torch.from_numpy(p.derived["roa"]).unsqueeze(0), - torch.Tensor(aLTi).unsqueeze(0), - p.profiles["ti(keV)"][-1, 0], - ) - .cpu() - .cpu().numpy()[0] - ) - - aLT0 = aLne_true[: ix1 + 1] - aLT2 = p.derived["aLne"][ix2:] - aLT1 = np.interp( - p.derived["roa"][ix1 : ix2 + 1], - [p.derived["roa"][ix1], p.derived["roa"][ix2]], - [aLT0[-1], aLT2[0]], - )[1:-1] - - aLne = np.append(np.append(aLT0, aLT1), aLT2) - ne = ( - CALCtools.integrateGradient( - torch.from_numpy(p.derived["roa"]).unsqueeze(0), - torch.Tensor(aLne).unsqueeze(0), - p.profiles["ne(10^19/m^3)"][-1], - ) - .cpu() - .cpu().numpy()[0] - ) - - p.profiles["te(keV)"] = Te - p.profiles["ti(keV)"][:, 0] = Ti - p.profiles["ne(10^19/m^3)"] = ne - - p.deriveQuantities() - - return p - -def add_figures(fn, fnlab='', fnlab_pre='', tab_color=None): - - figProf_1 = fn.add_figure(label= fnlab_pre + "Profiles" + fnlab, tab_color=tab_color) - figProf_2 = fn.add_figure(label= fnlab_pre + "Powers" + fnlab, tab_color=tab_color) - figProf_3 = fn.add_figure(label= fnlab_pre + "Geometry" + fnlab, tab_color=tab_color) - figProf_4 = fn.add_figure(label= fnlab_pre + "Gradients" + fnlab, tab_color=tab_color) - figFlows = fn.add_figure(label= fnlab_pre + "Flows" + fnlab, tab_color=tab_color) - figProf_6 = fn.add_figure(label= fnlab_pre + "Other" + fnlab, tab_color=tab_color) - fig7 = fn.add_figure(label= fnlab_pre + "Impurities" + fnlab, tab_color=tab_color) - figs = [figProf_1, figProf_2, figProf_3, figProf_4, figFlows, figProf_6, fig7] - - return figs - -def impurity_location(profiles, impurity_of_interest): - - position_of_impurity = None - for i in range(len(profiles.Species)): - if profiles.Species[i]["N"] == impurity_of_interest: - if position_of_impurity is not None: - raise ValueError(f"[MITIM] Species {impurity_of_interest} found at positions {position_of_impurity} and {i}") - position_of_impurity = i - if position_of_impurity is None: - raise ValueError(f"[MITIM] Species {impurity_of_interest} not found in profiles") + for ikey in some_times_are_not_here: + if ikey not in self.profiles.keys(): + self.profiles[ikey] = copy.deepcopy(self.profiles["rmin(m)"]) * 0.0 - return position_of_impurity \ No newline at end of file diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 79712a3d..6a1dc3b0 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -4,7 +4,8 @@ import numpy as np import xarray as xr import matplotlib.pyplot as plt -from mitim_tools.gacode_tools import TGYROtools, PROFILEStools + +from mitim_tools.gacode_tools import TGYROtools from mitim_tools.misc_tools import ( IOtools, GRAPHICStools, @@ -232,6 +233,7 @@ def prep( # PROFILES class. + from mitim_tools.gacode_tools import PROFILEStools profiles = ( PROFILEStools.PROFILES_GACODE(inputgacode) if inputgacode is not None @@ -371,6 +373,7 @@ def prep_direct_tglf( # PROFILES class. + from mitim_tools.gacode_tools import PROFILEStools self.profiles = ( PROFILEStools.PROFILES_GACODE(inputgacode) if inputgacode is not None @@ -475,6 +478,7 @@ def prep_from_tglf( self.FolderGACODE = IOtools.expandPath(FolderGACODE) # Main folder where things are + from mitim_tools.gacode_tools import PROFILEStools self.NormalizationSets, _ = NORMtools.normalizations( PROFILEStools.PROFILES_GACODE(input_gacode) if input_gacode is not None diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 5b83e895..e14bba87 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -8,7 +8,7 @@ GRAPHICStools, PLASMAtools, ) -from mitim_tools.gacode_tools import TGLFtools, PROFILEStools +from mitim_tools.gacode_tools import TGLFtools from mitim_tools.gacode_tools.utils import GACODEinterpret, GACODEdefaults, GACODErun from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -140,6 +140,7 @@ def prep( ) self.file_input_profiles = self.FolderGACODE / "input.gacode" + from mitim_tools.gacode_tools import PROFILEStools self.profiles = PROFILEStools.PROFILES_GACODE(self.file_input_profiles) if correctPROFILES: @@ -459,14 +460,10 @@ def run( ------------------------------------------------------------------------------------------------------------------------ """ if modify_inputgacodenew: - print( - "\t- It was requested that input.gacode.new is modified according to what TypeTarget was", - typeMsg="i", - ) + print("\t- It was requested that input.gacode.new is modified according to what TypeTarget was",typeMsg="i",) - inputgacode_new = PROFILEStools.PROFILES_GACODE( - self.FolderTGYRO_tmp / "input.gacode.new" - ) + from mitim_tools.gacode_tools import PROFILEStools + inputgacode_new = PROFILEStools.PROFILES_GACODE(self.FolderTGYRO_tmp / "input.gacode.new") if TGYRO_physics_options["TypeTarget"] < 3: for ikey in [ @@ -540,6 +537,7 @@ def read(self, label="tgyro1", folder=None, file_input_profiles=None): else: prof = self.profiles else: + from mitim_tools.gacode_tools import PROFILEStools prof = PROFILEStools.PROFILES_GACODE(file_input_profiles) self.results[label] = TGYROoutput(folder, profiles=prof) @@ -1203,10 +1201,9 @@ class TGYROoutput: def __init__(self, FolderTGYRO, profiles=None): self.FolderTGYRO = FolderTGYRO - if (profiles is None) and (FolderTGYRO / f"input.gacode").exists(): - profiles = PROFILEStools.PROFILES_GACODE( - FolderTGYRO / f"input.gacode", calculateDerived=False - ) + if (profiles is None) and (FolderTGYRO / "input.gacode").exists(): + from mitim_tools.gacode_tools import PROFILEStools + profiles = PROFILEStools.PROFILES_GACODE(FolderTGYRO / f"input.gacode", calculateDerived=False) self.profiles = profiles @@ -1220,10 +1217,7 @@ def __init__(self, FolderTGYRO, profiles=None): calculateDerived = True try: - self.profiles_final = PROFILEStools.PROFILES_GACODE( - self.FolderTGYRO / f"input.gacode.new", - calculateDerived=calculateDerived, - ) + self.profiles_final = PROFILEStools.PROFILES_GACODE(self.FolderTGYRO / "input.gacode.new",calculateDerived=calculateDerived,) except: self.profiles_final = None diff --git a/src/mitim_tools/gacode_tools/scripts/read_gacode.py b/src/mitim_tools/gacode_tools/scripts/read_gacode.py index 14b10374..778fe07d 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_gacode.py +++ b/src/mitim_tools/gacode_tools/scripts/read_gacode.py @@ -1,5 +1,5 @@ import argparse -from turtle import st +from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.gacode_tools import PROFILEStools """ @@ -32,7 +32,7 @@ def main(): if not print_only: - fn = PROFILEStools.plotAll(profs, lastRhoGradients=rho) + fn = MITIMstate.plotAll(profs, lastRhoGradients=rho) fn.show() diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index b40baf4f..223d247a 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -480,7 +480,7 @@ def runPROFILES_GEN( txt += f" -g {nameFiles}.geq\n" else: txt += "\n" - with open(FolderTGLF + "profiles_gen.sh", "w") as f: + with open(FolderTGLF / "profiles_gen.sh", "w") as f: f.write(txt) # ****************** diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 3014a397..6dd8f32c 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -5,7 +5,7 @@ import numpy as np import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, PLASMAtools -from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.gs_tools.utils import GEQplotting from shapely.geometry import LineString from scipy.integrate import quad @@ -418,7 +418,7 @@ def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH _, profiles["qrfe(MW/m^3)"] = PLASMAtools.parabolicProfile(Tbar=1.0,nu=5.0,rho=rhotor,Tedge=0.0) - p = PROFILEStools.PROFILES_GACODE.scratch(profiles) + p = MITIMstate.mitim_state.scratch(profiles) p.profiles["qrfe(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] * PichT/p.derived['qRF_MWmiller'][-1] /2 p.profiles["qrfi(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py new file mode 100644 index 00000000..ec50e760 --- /dev/null +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -0,0 +1,4259 @@ +import copy +import torch +import csv +import numpy as np +import matplotlib.pyplot as plt +from collections import OrderedDict +from mitim_tools.misc_tools import GRAPHICStools, MATHtools, PLASMAtools, IOtools +from mitim_modules.powertorch.utils import CALCtools +from mitim_tools.gs_tools import GEQtools +from mitim_tools.gacode_tools import NEOtools +from mitim_tools.gacode_tools.utils import GACODEdefaults, GEOMETRYtools +from mitim_tools.transp_tools import CDFtools +from mitim_tools.transp_tools.utils import TRANSPhelpers +from mitim_tools.misc_tools.LOGtools import printMsg as print +from mitim_tools import __version__ +from IPython import embed + +class mitim_state: + ''' + Class to manipulate the plasma state in MITIM. + ''' + + def __init__(self, type_file = 'input.gacode'): + + self.type = type_file + + def deriveQuantities(self, mi_ref=None, calculateDerived=True, n_theta_geo=1001, rederiveGeometry=True): + + # ------------------------------------------------------------------------------------------------------------------- + self.readSpecies() + self.produce_shape_lists() + self.mi_first = self.Species[0]["A"] + self.DTplasma() + self.sumFast() + # ------------------------------------- + + if "derived" not in self.__dict__: + self.derived = {} + + if mi_ref is not None: + self.derived["mi_ref"] = mi_ref + print(f"\t* Reference mass ({self.derived['mi_ref']:.2f}) to use was forced by class initialization",typeMsg="w") + else: + self.derived["mi_ref"] = self.mi_first + print(f"\t* Reference mass ({self.derived['mi_ref']}) from first ion",typeMsg="i") + + # Useful to have gradients in the basic ---------------------------------------------------------- + self.derived["aLTe"] = aLT(self.profiles["rmin(m)"], self.profiles["te(keV)"]) + self.derived["aLne"] = aLT( + self.profiles["rmin(m)"], self.profiles["ne(10^19/m^3)"] + ) + + self.derived["aLTi"] = self.profiles["ti(keV)"] * 0.0 + self.derived["aLni"] = [] + for i in range(self.profiles["ti(keV)"].shape[1]): + self.derived["aLTi"][:, i] = aLT( + self.profiles["rmin(m)"], self.profiles["ti(keV)"][:, i] + ) + self.derived["aLni"].append( + aLT(self.profiles["rmin(m)"], self.profiles["ni(10^19/m^3)"][:, i]) + ) + self.derived["aLni"] = np.transpose(np.array(self.derived["aLni"])) + # ------------------------------------------------------------------------------------------------ + + if calculateDerived: + self.deriveQuantities_full(rederiveGeometry=rederiveGeometry) + + # ------------------------------------------------------------------------------------- + # Method to write a scratch file + # ------------------------------------------------------------------------------------- + + @classmethod + def scratch(cls, profiles, label_header='', **kwargs_process): + instance = cls(None) + + # Header + instance.header = f''' +# Created from scratch with MITIM version {__version__} +# {label_header} +# +''' + # Add data to profiles + instance.profiles = profiles + + instance.process(**kwargs_process) + + return instance + + # ------------------------------------------------------------------------------------- + + def calculate_Er( + self, + folder, + rhos=None, + vgenOptions={}, + name="vgen1", + includeAll=False, + write_new_file=None, + cold_start=False, + ): + 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, cold_start=cold_start) + + 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)", + ] + + 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] + + self.deriveQuantities() + + if write_new_file is not None: + self.writeCurrentStatus(file=write_new_file) + + def produce_shape_lists(self): + self.shape_cos = [ + self.profiles["shape_cos0(-)"], # tilt + self.profiles["shape_cos1(-)"], + self.profiles["shape_cos2(-)"], + self.profiles["shape_cos3(-)"], + self.profiles["shape_cos4(-)"], + self.profiles["shape_cos5(-)"], + self.profiles["shape_cos6(-)"], + ] + self.shape_sin = [ + None, + None, # s1 is arcsin(triangularity) + None, # s2 is minus squareness + self.profiles["shape_sin3(-)"], + self.profiles["shape_sin4(-)"], + self.profiles["shape_sin5(-)"], + self.profiles["shape_sin6(-)"], + ] + + def readSpecies(self, maxSpecies=100): + maxSpecies = int(self.profiles["nion"][0]) + + Species = [] + for j in range(maxSpecies): + # To determine later if this specie has zero density + niT = self.profiles["ni(10^19/m^3)"][0, j] + + sp = { + "N": self.profiles["name"][j], + "Z": float(self.profiles["z"][j]), + "A": float(self.profiles["mass"][j]), + "S": self.profiles["type"][j].split("[")[-1].split("]")[0], + "n0": niT, + } + + Species.append(sp) + + self.Species = Species + + def sumFast(self): + self.nFast = self.profiles["ne(10^19/m^3)"] * 0.0 + self.nZFast = self.profiles["ne(10^19/m^3)"] * 0.0 + self.nThermal = self.profiles["ne(10^19/m^3)"] * 0.0 + self.nZThermal = self.profiles["ne(10^19/m^3)"] * 0.0 + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "fast": + self.nFast += self.profiles["ni(10^19/m^3)"][:, sp] + self.nZFast += ( + self.profiles["ni(10^19/m^3)"][:, sp] * self.profiles["z"][sp] + ) + else: + self.nThermal += self.profiles["ni(10^19/m^3)"][:, sp] + self.nZThermal += ( + self.profiles["ni(10^19/m^3)"][:, sp] * self.profiles["z"][sp] + ) + + def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry=True): + """ + deriving geometry is expensive, so if I'm just updating profiles it may not be needed + """ + + self.varqmom = "qmom(N/m^2)" + if self.varqmom not in self.profiles: + self.profiles[self.varqmom] = self.profiles["rho(-)"] * 0.0 + + if "derived" not in self.__dict__: + self.derived = {} + + # --------------------------------------------------------------------------------------------------------------------- + # --------- MAIN (useful for STATEtools) + # --------------------------------------------------------------------------------------------------------------------- + + self.derived["a"] = self.profiles["rmin(m)"][-1] + # self.derived['epsX'] = self.profiles['rmaj(m)'] / self.profiles['rmin(m)'] + # self.derived['eps'] = self.derived['epsX'][-1] + self.derived["eps"] = self.profiles["rmin(m)"][-1] / self.profiles["rmaj(m)"][-1] + + self.derived["roa"] = self.profiles["rmin(m)"] / self.derived["a"] + self.derived["Rmajoa"] = self.profiles["rmaj(m)"] / self.derived["a"] + self.derived["Zmagoa"] = self.profiles["zmag(m)"] / self.derived["a"] + + self.derived["torflux"] = ( + float(self.profiles["torfluxa(Wb/radian)"][0]) + * 2 + * np.pi + * self.profiles["rho(-)"] ** 2 + ) # Wb + self.derived["B_unit"] = PLASMAtools.Bunit( + self.derived["torflux"], self.profiles["rmin(m)"] + ) + + self.derived["psi_pol_n"] = ( + self.profiles["polflux(Wb/radian)"] - self.profiles["polflux(Wb/radian)"][0] + ) / ( + self.profiles["polflux(Wb/radian)"][-1] + - self.profiles["polflux(Wb/radian)"][0] + ) + self.derived["rho_pol"] = self.derived["psi_pol_n"] ** 0.5 + + self.derived["q95"] = np.interp( + 0.95, self.derived["psi_pol_n"], self.profiles["q(-)"] + ) + + self.derived["q0"] = self.profiles["q(-)"][0] + + if self.profiles["q(-)"].min() > 1.0: + self.derived["rho_saw"] = np.nan + else: + self.derived["rho_saw"] = np.interp( + 1.0, self.profiles["q(-)"], self.profiles["rho(-)"] + ) + + # --------- Geometry (only if it doesn't exist or if I ask to recalculate) + + if rederiveGeometry or ("volp_miller" not in self.derived): + + self.produce_shape_lists() + + ( + self.derived["volp_miller"], + self.derived["surf_miller"], + self.derived["gradr_miller"], + self.derived["bp2_miller"], + self.derived["bt2_miller"], + self.derived["geo_bt"], + ) = GEOMETRYtools.calculateGeometricFactors( + self, + n_theta=n_theta_geo, + ) + + # Calculate flux surfaces + cn = np.array(self.shape_cos).T + sn = copy.deepcopy(self.shape_sin) + sn[0] = self.profiles["rmaj(m)"]*0.0 + sn[1] = np.arcsin(self.profiles["delta(-)"]) + sn[2] = -self.profiles["zeta(-)"] + sn = np.array(sn).T + flux_surfaces = GEQtools.mitim_flux_surfaces() + flux_surfaces.reconstruct_from_mxh_moments( + self.profiles["rmaj(m)"], + self.profiles["rmin(m)"], + self.profiles["kappa(-)"], + self.profiles["zmag(m)"], + cn, + sn) + self.derived["R_surface"],self.derived["Z_surface"] = flux_surfaces.R, flux_surfaces.Z + # ----------------------------------------------- + + #cross-sectional area of each flux surface + self.derived["surfXS"] = GEOMETRYtools.xsec_area_RZ( + self.derived["R_surface"], + self.derived["Z_surface"] + ) + + self.derived["R_LF"] = self.derived["R_surface"].max( + axis=1 + ) # self.profiles['rmaj(m)'][0]+self.profiles['rmin(m)'] + + # For Synchrotron + self.derived["B_ref"] = np.abs( + self.derived["B_unit"] * self.derived["geo_bt"] + ) + + # -------------------------------------------------------------------------- + # Reference mass + # -------------------------------------------------------------------------- + + # Forcing mass from this specific deriveQuantities call + if mi_ref is not None: + self.derived["mi_ref"] = mi_ref + print(f'\t- Using mi_ref={self.derived["mi_ref"]} provided in this particular deriveQuantities method, subtituting initialization one',typeMsg='i') + + # --------------------------------------------------------------------------------------------------------------------- + # --------- Important for scaling laws + # --------------------------------------------------------------------------------------------------------------------- + + self.derived["kappa95"] = np.interp( + 0.95, self.derived["psi_pol_n"], self.profiles["kappa(-)"] + ) + + self.derived["kappa995"] = np.interp( + 0.995, self.derived["psi_pol_n"], self.profiles["kappa(-)"] + ) + + self.derived["kappa_a"] = self.derived["surfXS"][-1] / np.pi / self.derived["a"] ** 2 + + self.derived["delta95"] = np.interp( + 0.95, self.derived["psi_pol_n"], self.profiles["delta(-)"] + ) + + self.derived["delta995"] = np.interp( + 0.995, self.derived["psi_pol_n"], self.profiles["delta(-)"] + ) + + self.derived["Rgeo"] = float(self.profiles["rcentr(m)"][-1]) + self.derived["B0"] = np.abs(float(self.profiles["bcentr(T)"][-1])) + + # --------------------------------------------------------------------------------------------------------------------- + + """ + surf_miller is truly surface area, but because of the GACODE definitions of flux, + Surf = V' <|grad r|> + Surf_GACODE = V' + """ + + self.derived["surfGACODE_miller"] = (self.derived["surf_miller"] / self.derived["gradr_miller"]) + + self.derived["surfGACODE_miller"][np.isnan(self.derived["surfGACODE_miller"])] = 0 + + self.derived["c_s"] = PLASMAtools.c_s( + self.profiles["te(keV)"], self.derived["mi_ref"] + ) + self.derived["rho_s"] = PLASMAtools.rho_s( + self.profiles["te(keV)"], self.derived["mi_ref"], self.derived["B_unit"] + ) + + self.derived["q_gb"], self.derived["g_gb"], self.derived["pi_gb"], self.derived["s_gb"], _ = PLASMAtools.gyrobohmUnits( + self.profiles["te(keV)"], + self.profiles["ne(10^19/m^3)"] * 1e-1, + self.derived["mi_ref"], + np.abs(self.derived["B_unit"]), + self.profiles["rmin(m)"][-1], + ) + + """ + In prgen_map_plasmastate: + qspow_e = expro_qohme+expro_qbeame+expro_qrfe+expro_qfuse-expro_qei & + -expro_qsync-expro_qbrem-expro_qline + qspow_i = expro_qbeami+expro_qrfi+expro_qfusi+expro_qei + """ + + qe_terms = { + "qohme(MW/m^3)": 1, + "qbeame(MW/m^3)": 1, + "qrfe(MW/m^3)": 1, + "qfuse(MW/m^3)": 1, + "qei(MW/m^3)": -1, + "qsync(MW/m^3)": -1, + "qbrem(MW/m^3)": -1, + "qline(MW/m^3)": -1, + "qione(MW/m^3)": 1, + } + + self.derived["qe"] = np.zeros(len(self.profiles["rho(-)"])) + for i in qe_terms: + if i in self.profiles: + self.derived["qe"] += qe_terms[i] * self.profiles[i] + + qrad = { + "qsync(MW/m^3)": 1, + "qbrem(MW/m^3)": 1, + "qline(MW/m^3)": 1, + } + + self.derived["qrad"] = np.zeros(len(self.profiles["rho(-)"])) + for i in qrad: + if i in self.profiles: + self.derived["qrad"] += qrad[i] * self.profiles[i] + + qi_terms = { + "qbeami(MW/m^3)": 1, + "qrfi(MW/m^3)": 1, + "qfusi(MW/m^3)": 1, + "qei(MW/m^3)": 1, + "qioni(MW/m^3)": 1, + } + + self.derived["qi"] = np.zeros(len(self.profiles["rho(-)"])) + for i in qi_terms: + if i in self.profiles: + self.derived["qi"] += qi_terms[i] * self.profiles[i] + + # Depends on GACODE version + ge_terms = {self.varqpar: 1, self.varqpar2: 1} + + self.derived["ge"] = np.zeros(len(self.profiles["rho(-)"])) + for i in ge_terms: + if i in self.profiles: + self.derived["ge"] += ge_terms[i] * self.profiles[i] + + """ + Careful, that's in MW/m^3. I need to find the volumes. Using here the Miller + calculation. Should be consistent with TGYRO + + profiles_gen puts any missing power into the CX: qioni, qione + """ + + r = self.profiles["rmin(m)"] + volp = self.derived["volp_miller"] + + self.derived["qe_MWmiller"] = CALCtools.integrateFS(self.derived["qe"], r, volp) + self.derived["qi_MWmiller"] = CALCtools.integrateFS(self.derived["qi"], r, volp) + self.derived["ge_10E20miller"] = CALCtools.integrateFS( + 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) + + self.derived["QiQe"] = self.derived["qi_MWm2"] / np.where(self.derived["qe_MWm2"] == 0, 1e-10, self.derived["qe_MWm2"]) # to avoid division by zero + + # "Convective" flux + self.derived["ce_MWmiller"] = PLASMAtools.convective_flux( + self.profiles["te(keV)"], self.derived["ge_10E20miller"] + ) + self.derived["ce_MWm2"] = PLASMAtools.convective_flux( + self.profiles["te(keV)"], self.derived["ge_10E20m2"] + ) + + # qmom + self.derived["mt_Jmiller"] = CALCtools.integrateFS( + self.profiles[self.varqmom], r, volp + ) + self.derived["mt_Jm2"] = self.derived["mt_Jmiller"] / (volp) + + # Extras for plotting in TGYRO for comparison + P = np.zeros(len(self.profiles["rmin(m)"])) + if "qsync(MW/m^3)" in self.profiles: + P += self.profiles["qsync(MW/m^3)"] + if "qbrem(MW/m^3)" in self.profiles: + P += self.profiles["qbrem(MW/m^3)"] + if "qline(MW/m^3)" in self.profiles: + P += self.profiles["qline(MW/m^3)"] + self.derived["qe_rad_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + P = self.profiles["qei(MW/m^3)"] + self.derived["qe_exc_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + """ + --------------------------------------------------------------------------------------------------------------------- + Note that the real auxiliary power is RF+BEAMS+OHMIC, + The QIONE is added by TGYRO, but sometimes it includes radiation and direct RF to electrons + --------------------------------------------------------------------------------------------------------------------- + """ + + # ** Electrons + + P = np.zeros(len(self.profiles["rho(-)"])) + for i in ["qrfe(MW/m^3)", "qohme(MW/m^3)", "qbeame(MW/m^3)"]: + if i in self.profiles: + P += self.profiles[i] + + self.derived["qe_auxONLY"] = copy.deepcopy(P) + self.derived["qe_auxONLY_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + for i in ["qione(MW/m^3)"]: + if i in self.profiles: + P += self.profiles[i] + + self.derived["qe_aux"] = copy.deepcopy(P) + self.derived["qe_aux_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + # ** Ions + + P = np.zeros(len(self.profiles["rho(-)"])) + for i in ["qrfi(MW/m^3)", "qbeami(MW/m^3)"]: + if i in self.profiles: + P += self.profiles[i] + + self.derived["qi_auxONLY"] = copy.deepcopy(P) + self.derived["qi_auxONLY_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + for i in ["qioni(MW/m^3)"]: + if i in self.profiles: + P += self.profiles[i] + + self.derived["qi_aux"] = copy.deepcopy(P) + self.derived["qi_aux_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + # ** General + + P = np.zeros(len(self.profiles["rho(-)"])) + for i in ["qohme(MW/m^3)"]: + if i in self.profiles: + P += self.profiles[i] + self.derived["qOhm_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + P = np.zeros(len(self.profiles["rho(-)"])) + for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: + if i in self.profiles: + P += self.profiles[i] + self.derived["qRF_MWmiller"] = CALCtools.integrateFS(P, r, volp) + if "qrfe(MW/m^3)" in self.profiles: + self.derived["qRFe_MWmiller"] = CALCtools.integrateFS( + self.profiles["qrfe(MW/m^3)"], r, volp + ) + if "qrfi(MW/m^3)" in self.profiles: + self.derived["qRFi_MWmiller"] = CALCtools.integrateFS( + self.profiles["qrfi(MW/m^3)"], r, volp + ) + + P = np.zeros(len(self.profiles["rho(-)"])) + for i in ["qbeame(MW/m^3)", "qbeami(MW/m^3)"]: + if i in self.profiles: + P += self.profiles[i] + self.derived["qBEAM_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + self.derived["qrad_MWmiller"] = CALCtools.integrateFS(self.derived["qrad"], r, volp) + if "qsync(MW/m^3)" in self.profiles: + self.derived["qrad_sync_MWmiller"] = CALCtools.integrateFS(self.profiles["qsync(MW/m^3)"], r, volp) + else: + self.derived["qrad_sync_MWmiller"] = self.derived["qrad_MWmiller"]*0.0 + if "qbrem(MW/m^3)" in self.profiles: + self.derived["qrad_brem_MWmiller"] = CALCtools.integrateFS(self.profiles["qbrem(MW/m^3)"], r, volp) + else: + self.derived["qrad_brem_MWmiller"] = self.derived["qrad_MWmiller"]*0.0 + if "qline(MW/m^3)" in self.profiles: + self.derived["qrad_line_MWmiller"] = CALCtools.integrateFS(self.profiles["qline(MW/m^3)"], r, volp) + else: + self.derived["qrad_line_MWmiller"] = self.derived["qrad_MWmiller"]*0.0 + + P = np.zeros(len(self.profiles["rho(-)"])) + for i in ["qfuse(MW/m^3)", "qfusi(MW/m^3)"]: + if i in self.profiles: + P += self.profiles[i] + self.derived["qFus_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + P = np.zeros(len(self.profiles["rho(-)"])) + for i in ["qioni(MW/m^3)", "qione(MW/m^3)"]: + if i in self.profiles: + P += self.profiles[i] + self.derived["qz_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + self.derived["q_MWmiller"] = ( + self.derived["qe_MWmiller"] + self.derived["qi_MWmiller"] + ) + + # --------------------------------------------------------------------------------------------------------------------- + # --------------------------------------------------------------------------------------------------------------------- + + P = np.zeros(len(self.profiles["rho(-)"])) + if "qfuse(MW/m^3)" in self.profiles: + P = self.profiles["qfuse(MW/m^3)"] + self.derived["qe_fus_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + P = np.zeros(len(self.profiles["rho(-)"])) + if "qfusi(MW/m^3)" in self.profiles: + P = self.profiles["qfusi(MW/m^3)"] + self.derived["qi_fus_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + P = np.zeros(len(self.profiles["rho(-)"])) + if "qfusi(MW/m^3)" in self.profiles: + self.derived["q_fus"] = ( + self.profiles["qfuse(MW/m^3)"] + self.profiles["qfusi(MW/m^3)"] + ) * 5 + P = self.derived["q_fus"] + self.derived["q_fus"] = P + self.derived["q_fus_MWmiller"] = CALCtools.integrateFS(P, r, volp) + + """ + Derivatives + """ + self.derived["aLTe"] = aLT(self.profiles["rmin(m)"], self.profiles["te(keV)"]) + self.derived["aLTi"] = self.profiles["ti(keV)"] * 0.0 + for i in range(self.profiles["ti(keV)"].shape[1]): + self.derived["aLTi"][:, i] = aLT( + self.profiles["rmin(m)"], self.profiles["ti(keV)"][:, i] + ) + self.derived["aLne"] = aLT( + self.profiles["rmin(m)"], self.profiles["ne(10^19/m^3)"] + ) + self.derived["aLni"] = [] + for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): + self.derived["aLni"].append( + aLT(self.profiles["rmin(m)"], self.profiles["ni(10^19/m^3)"][:, i]) + ) + self.derived["aLni"] = np.transpose(np.array(self.derived["aLni"])) + + if "w0(rad/s)" not in self.profiles: + self.profiles["w0(rad/s)"] = self.profiles["rho(-)"] * 0.0 + self.derived["aLw0"] = aLT(self.profiles["rmin(m)"], self.profiles["w0(rad/s)"]) + self.derived["dw0dr"] = -grad( + self.profiles["rmin(m)"], self.profiles["w0(rad/s)"] + ) + + self.derived["dqdr"] = grad(self.profiles["rmin(m)"], self.profiles["q(-)"]) + + """ + Other, performance + """ + qFus = self.derived["qe_fus_MWmiller"] + self.derived["qi_fus_MWmiller"] + self.derived["Pfus"] = qFus[-1] * 5 + + # Note that in cases with NPRAD=0 in TRANPS, this includes radiation! no way to deal wit this... + qIn = self.derived["qe_aux_MWmiller"] + self.derived["qi_aux_MWmiller"] + self.derived["qIn"] = qIn[-1] + self.derived["Q"] = self.derived["Pfus"] / self.derived["qIn"] + self.derived["qHeat"] = qIn[-1] + qFus[-1] + + self.derived["qTr"] = ( + self.derived["qe_aux_MWmiller"] + + self.derived["qi_aux_MWmiller"] + + (self.derived["qe_fus_MWmiller"] + self.derived["qi_fus_MWmiller"]) + - self.derived["qrad_MWmiller"] + ) + + self.derived["Prad"] = self.derived["qrad_MWmiller"][-1] + self.derived["Prad_sync"] = self.derived["qrad_sync_MWmiller"][-1] + self.derived["Prad_brem"] = self.derived["qrad_brem_MWmiller"][-1] + self.derived["Prad_line"] = self.derived["qrad_line_MWmiller"][-1] + self.derived["Psol"] = self.derived["qHeat"] - self.derived["Prad"] + + self.derived["ni_thr"] = [] + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "therm": + self.derived["ni_thr"].append(self.profiles["ni(10^19/m^3)"][:, sp]) + self.derived["ni_thr"] = np.transpose(self.derived["ni_thr"]) + self.derived["ni_thrAll"] = self.derived["ni_thr"].sum(axis=1) + + self.derived["ni_All"] = self.profiles["ni(10^19/m^3)"].sum(axis=1) + + + ( + self.derived["ptot_manual"], + self.derived["pe"], + self.derived["pi"], + ) = PLASMAtools.calculatePressure( + np.expand_dims(self.profiles["te(keV)"], 0), + np.expand_dims(np.transpose(self.profiles["ti(keV)"]), 0), + np.expand_dims(self.profiles["ne(10^19/m^3)"] * 0.1, 0), + np.expand_dims(np.transpose(self.profiles["ni(10^19/m^3)"] * 0.1), 0), + ) + self.derived["ptot_manual"], self.derived["pe"], self.derived["pi"] = ( + self.derived["ptot_manual"][0], + self.derived["pe"][0], + self.derived["pi"][0], + ) + + ( + self.derived["pthr_manual"], + _, + self.derived["pi_thr"], + ) = PLASMAtools.calculatePressure( + np.expand_dims(self.profiles["te(keV)"], 0), + np.expand_dims(np.transpose(self.profiles["ti(keV)"]), 0), + np.expand_dims(self.profiles["ne(10^19/m^3)"] * 0.1, 0), + np.expand_dims(np.transpose(self.derived["ni_thr"] * 0.1), 0), + ) + self.derived["pthr_manual"], self.derived["pi_thr"] = ( + self.derived["pthr_manual"][0], + self.derived["pi_thr"][0], + ) + + # ------- + # Content + # ------- + + ( + self.derived["We"], + self.derived["Wi_thr"], + self.derived["Ne"], + self.derived["Ni_thr"], + ) = PLASMAtools.calculateContent( + np.expand_dims(r, 0), + np.expand_dims(self.profiles["te(keV)"], 0), + np.expand_dims(np.transpose(self.profiles["ti(keV)"]), 0), + np.expand_dims(self.profiles["ne(10^19/m^3)"] * 0.1, 0), + np.expand_dims(np.transpose(self.derived["ni_thr"] * 0.1), 0), + np.expand_dims(volp, 0), + ) + + ( + self.derived["We"], + self.derived["Wi_thr"], + self.derived["Ne"], + self.derived["Ni_thr"], + ) = ( + self.derived["We"][0], + self.derived["Wi_thr"][0], + self.derived["Ne"][0], + self.derived["Ni_thr"][0], + ) + + self.derived["Nthr"] = self.derived["Ne"] + self.derived["Ni_thr"] + self.derived["Wthr"] = self.derived["We"] + self.derived["Wi_thr"] # Thermal + + self.derived["tauE"] = self.derived["Wthr"] / self.derived["qHeat"] # Seconds + + self.derived["tauP"] = np.where(self.derived["geIn"] != 0, self.derived["Ne"] / self.derived["geIn"], np.inf) # 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)"] + ).transpose().repeat(self.profiles["ni(10^19/m^3)"].shape[1], axis=1) + + # Vol-avg density + self.derived["volume"] = CALCtools.integrateFS(np.ones(r.shape[0]), r, volp)[ + -1 + ] # m^3 + self.derived["ne_vol20"] = ( + CALCtools.integrateFS(self.profiles["ne(10^19/m^3)"] * 0.1, r, volp)[-1] + / self.derived["volume"] + ) # 1E20/m^3 + + self.derived["ni_vol20"] = np.zeros(self.profiles["ni(10^19/m^3)"].shape[1]) + self.derived["fi_vol"] = np.zeros(self.profiles["ni(10^19/m^3)"].shape[1]) + for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): + self.derived["ni_vol20"][i] = ( + CALCtools.integrateFS( + self.profiles["ni(10^19/m^3)"][:, i] * 0.1, r, volp + )[-1] + / self.derived["volume"] + ) # 1E20/m^3 + self.derived["fi_vol"][i] = ( + self.derived["ni_vol20"][i] / self.derived["ne_vol20"] + ) + + self.derived["fi_onlyions_vol"] = self.derived["ni_vol20"] / np.sum( + self.derived["ni_vol20"] + ) + + self.derived["ne_peaking"] = ( + self.profiles["ne(10^19/m^3)"][0] * 0.1 / self.derived["ne_vol20"] + ) + + xcoord = self.derived[ + "rho_pol" + ] # to find the peaking at rho_pol (with square root) as in Angioni PRL 2003 + self.derived["ne_peaking0.2"] = ( + self.profiles["ne(10^19/m^3)"][np.argmin(np.abs(xcoord - 0.2))] + * 0.1 + / self.derived["ne_vol20"] + ) + + self.derived["Te_vol"] = ( + CALCtools.integrateFS(self.profiles["te(keV)"], r, volp)[-1] + / self.derived["volume"] + ) # keV + self.derived["Te_peaking"] = ( + self.profiles["te(keV)"][0] / self.derived["Te_vol"] + ) + self.derived["Ti_vol"] = ( + CALCtools.integrateFS(self.profiles["ti(keV)"][:, 0], r, volp)[-1] + / self.derived["volume"] + ) # keV + self.derived["Ti_peaking"] = ( + self.profiles["ti(keV)"][0, 0] / self.derived["Ti_vol"] + ) + + self.derived["ptot_manual_vol"] = ( + CALCtools.integrateFS(self.derived["ptot_manual"], r, volp)[-1] + / self.derived["volume"] + ) # MPa + self.derived["pthr_manual_vol"] = ( + CALCtools.integrateFS(self.derived["pthr_manual"], r, volp)[-1] + / self.derived["volume"] + ) # MPa + + self.derived['pfast_manual'] = self.derived['ptot_manual'] - self.derived['pthr_manual'] + self.derived["pfast_manual_vol"] = ( + CALCtools.integrateFS(self.derived["pfast_manual"], r, volp)[-1] + / self.derived["volume"] + ) # MPa + + self.derived['pfast_fraction'] = self.derived['pfast_manual_vol'] / self.derived['ptot_manual_vol'] + + #approximate pedestal top density + self.derived['ptop(Pa)'] = np.interp(0.90, self.profiles['rho(-)'], self.profiles['ptot(Pa)']) + + # Quasineutrality + self.derived["QN_Error"] = np.abs( + 1 - np.sum(self.derived["fi_vol"] * self.profiles["z"]) + ) + self.derived["Zeff"] = ( + np.sum(self.profiles["ni(10^19/m^3)"] * self.profiles["z"] ** 2, axis=1) + / self.profiles["ne(10^19/m^3)"] + ) + self.derived["Zeff_vol"] = ( + CALCtools.integrateFS(self.derived["Zeff"], r, volp)[-1] + / self.derived["volume"] + ) + + self.derived["nu_eff"] = PLASMAtools.coll_Angioni07( + self.derived["ne_vol20"] * 1e1, + self.derived["Te_vol"], + self.derived["Rgeo"], + Zeff=self.derived["Zeff_vol"], + ) + + self.derived["nu_eff2"] = PLASMAtools.coll_Angioni07( + self.derived["ne_vol20"] * 1e1, + self.derived["Te_vol"], + self.derived["Rgeo"], + Zeff=2.0, + ) + + # Avg mass + self.calculateMass() + + params_set_scaling = ( + np.abs(float(self.profiles["current(MA)"][-1])), + self.derived["Rgeo"], + self.derived["kappa_a"], + self.derived["ne_vol20"], + self.derived["a"] / self.derived["Rgeo"], + self.derived["B0"], + self.derived["mbg_main"], + self.derived["qHeat"], + ) + + self.derived["tau98y2"], self.derived["H98"] = PLASMAtools.tau98y2( + *params_set_scaling, tauE=self.derived["tauE"] + ) + self.derived["tau89p"], self.derived["H89"] = PLASMAtools.tau89p( + *params_set_scaling, tauE=self.derived["tauE"] + ) + self.derived["tau97L"], self.derived["H97L"] = PLASMAtools.tau97L( + *params_set_scaling, tauE=self.derived["tauE"] + ) + + """ + Mach number + """ + + Vtor_LF_Mach1 = PLASMAtools.constructVtorFromMach( + 1.0, self.profiles["ti(keV)"][:, 0], self.derived["mbg"] + ) # m/s + w0_Mach1 = Vtor_LF_Mach1 / (self.derived["R_LF"]) # rad/s + self.derived["MachNum"] = self.profiles["w0(rad/s)"] / w0_Mach1 + self.derived["MachNum_vol"] = ( + CALCtools.integrateFS(self.derived["MachNum"], r, volp)[-1] + / self.derived["volume"] + ) + + # Retain the old beta definition for comparison with 0D modeling + Beta_old = (self.derived["ptot_manual_vol"]* 1e6 / (self.derived["B0"] ** 2 / (2 * 4 * np.pi * 1e-7))) + self.derived["BetaN_engineering"] = (Beta_old / + (np.abs(float(self.profiles["current(MA)"][-1])) / + (self.derived["a"] * self.derived["B0"]) + )* 100.0 + ) # expressed in percent + + ''' + --------------------------------------------------------------------------------------------------- + Using B_unit, derive and for betap and betat calculations. + Equivalent to GACODE expro_bp2, expro_bt2 + --------------------------------------------------------------------------------------------------- + ''' + + self.derived["bp2_exp"] = self.derived["bp2_miller"] * self.derived["B_unit"] ** 2 + self.derived["bt2_exp"] = self.derived["bt2_miller"] * self.derived["B_unit"] ** 2 + + # Calculate the volume averages of bt2 and bp2 + + P = self.derived["bp2_exp"] + self.derived["bp2_vol_avg"] = CALCtools.integrateFS(P, r, volp)[-1] / self.derived["volume"] + P = self.derived["bt2_exp"] + self.derived["bt2_vol_avg"] = CALCtools.integrateFS(P, r, volp)[-1] / self.derived["volume"] + + # calculate beta_poloidal and beta_toroidal using volume averaged values + # mu0 = 4pi x 10^-7, also need to convert MPa to Pa + + self.derived["Beta_p"] = (2 * 4 * np.pi * 1e-7)*self.derived["ptot_manual_vol"]* 1e6/self.derived["bp2_vol_avg"] + self.derived["Beta_t"] = (2 * 4 * np.pi * 1e-7)*self.derived["ptot_manual_vol"]* 1e6/self.derived["bt2_vol_avg"] + + self.derived["Beta"] = 1/(1/self.derived["Beta_p"]+1/self.derived["Beta_t"]) + + TroyonFactor = np.abs(float(self.profiles["current(MA)"][-1])) / (self.derived["a"] * self.derived["B0"]) + + self.derived["BetaN"] = self.derived["Beta"] / TroyonFactor * 100.0 + + # --- + + nG = PLASMAtools.Greenwald_density( + np.abs(float(self.profiles["current(MA)"][-1])), + float(self.profiles["rmin(m)"][-1]), + ) + self.derived["fG"] = self.derived["ne_vol20"] / nG + self.derived["fG_x"] = self.profiles["ne(10^19/m^3)"]* 0.1 / nG + + self.derived["tite"] = self.profiles["ti(keV)"][:, 0] / self.profiles["te(keV)"] + self.derived["tite_vol"] = self.derived["Ti_vol"] / self.derived["Te_vol"] + + self.derived["LH_nmin"] = PLASMAtools.LHthreshold_nmin( + np.abs(float(self.profiles["current(MA)"][-1])), + self.derived["B0"], + self.derived["a"], + self.derived["Rgeo"], + ) + + self.derived["LH_Martin2"] = ( + PLASMAtools.LHthreshold_Martin2( + self.derived["ne_vol20"], + self.derived["B0"], + self.derived["a"], + self.derived["Rgeo"], + nmin=self.derived["LH_nmin"], + ) + * (2 / self.derived["mbg_main"]) ** 1.11 + ) + + self.derived["LHratio"] = self.derived["Psol"] / self.derived["LH_Martin2"] + + self.readSpecies() + + # ------------------------------------------------------- + # q-star + # ------------------------------------------------------- + + self.derived["qstar"] = PLASMAtools.evaluate_qstar( + self.profiles['current(MA)'][0], + self.profiles['rcentr(m)'], + self.derived['kappa95'], + self.profiles['bcentr(T)'], + self.derived['eps'], + self.derived['delta95'], + ITERcorrection=False, + includeShaping=True, + )[0] + self.derived["qstar_ITER"] = PLASMAtools.evaluate_qstar( + self.profiles['current(MA)'][0], + self.profiles['rcentr(m)'], + self.derived['kappa95'], + self.profiles['bcentr(T)'], + self.derived['eps'], + self.derived['delta95'], + ITERcorrection=True, + includeShaping=True, + )[0] + + # ------------------------------------------------------- + # Separatrix estimations + # ------------------------------------------------------- + + # ~~~~ Estimate lambda_q + pressure_atm = self.derived["ptot_manual_vol"] * 1e6 / 101325.0 + Lambda_q = PLASMAtools.calculateHeatFluxWidth_Brunner(pressure_atm) + + # ~~~~ Estimate upstream temperature + Bt = self.profiles["bcentr(T)"][0] + Bp = self.derived["eps"] * Bt / self.derived["q95"] #TODO: VERY ROUGH APPROXIMATION!!!! + + self.derived['Te_lcfs_estimate'] = PLASMAtools.calculateUpstreamTemperature( + Lambda_q, + self.derived["q95"], + self.derived["ne_vol20"], + self.derived["Psol"], + self.profiles["rcentr(m)"][0], + Bp, + Bt + )[0] + + # ~~~~ Estimate upstream density + self.derived['ne_lcfs_estimate'] = self.derived["ne_vol20"] * 0.6 + + # ------------------------------------------------------- + # TGLF-relevant quantities + # ------------------------------------------------------- + + self.tglf_plasma() + + def tglf_plasma(self): + + def deriv_gacode(y): + return grad(self.profiles["rmin(m)"],y).cpu().numpy() + + self.derived["tite_all"] = self.profiles["ti(keV)"] / self.profiles["te(keV)"][:, np.newaxis] + + self.derived['betae'] = PLASMAtools.betae( + self.profiles['te(keV)'], + self.profiles['ne(10^19/m^3)']*0.1, + self.derived["B_unit"]) + + self.derived['xnue'] = PLASMAtools.xnue( + torch.from_numpy(self.profiles['te(keV)']).to(torch.double), + torch.from_numpy(self.profiles['ne(10^19/m^3)']*0.1).to(torch.double), + self.derived["a"], + mref_u=self.derived["mi_ref"]).cpu().numpy() + + self.derived['debye'] = PLASMAtools.debye( + self.profiles['te(keV)'], + self.profiles['ne(10^19/m^3)']*0.1, + self.derived["mi_ref"], + self.derived["B_unit"]) + + self.derived['pprime'] = 1E-7 * self.profiles["q(-)"]*self.derived['a']**2/self.profiles["rmin(m)"]/self.derived["B_unit"]**2*deriv_gacode(self.profiles["ptot(Pa)"]) + self.derived['pprime'][0] = 0.0 + + self.derived['drmin/dr'] = deriv_gacode(self.profiles["rmin(m)"]) + self.derived['dRmaj/dr'] = deriv_gacode(self.profiles["rmaj(m)"]) + self.derived['dZmaj/dr'] = deriv_gacode(self.profiles["zmag(m)"]) + + self.derived['s_kappa'] = self.profiles["rmin(m)"] / self.profiles["kappa(-)"] * deriv_gacode(self.profiles["kappa(-)"]) + self.derived['s_delta'] = self.profiles["rmin(m)"] * deriv_gacode(self.profiles["delta(-)"]) + self.derived['s_zeta'] = self.profiles["rmin(m)"] * deriv_gacode(self.profiles["zeta(-)"]) + + s = self.profiles["rmin(m)"] / self.profiles["q(-)"]*deriv_gacode(self.profiles["q(-)"]) + self.derived['s_q'] = np.concatenate([np.array([0.0]),(self.profiles["q(-)"][1:] / self.derived['roa'][1:])**2 * s[1:]]) # infinite in first location + + ''' + Rotations + -------------------------------------------------------- + From TGYRO/TGLF definitions + w0p = expro_w0p(:)/100.0 + f_rot(:) = w0p(:)/w0_norm + gamma_p0 = -r_maj(i_r)*f_rot(i_r)*w0_norm + gamma_eb0 = gamma_p0*r(i_r)/(q_abs*r_maj(i_r)) + ''' + + w0p = deriv_gacode(self.profiles["w0(rad/s)"]) + gamma_p0 = -self.profiles["rmaj(m)"]*w0p + gamma_eb0 = -deriv_gacode(self.profiles["w0(rad/s)"]) * self.profiles["rmin(m)"]/self.profiles["q(-)"] + + self.derived['vexb_shear'] = gamma_eb0 * self.derived["a"]/self.derived['c_s'] + self.derived['vpar_shear'] = gamma_p0 * self.derived["a"]/self.derived['c_s'] + self.derived['vpar'] = self.profiles["rmaj(m)"]*self.profiles["w0(rad/s)"]/self.derived['c_s'] + + def calculateMass(self): + self.derived["mbg"] = 0.0 + self.derived["fmain"] = 0.0 + for i in range(self.derived["ni_vol20"].shape[0]): + self.derived["mbg"] += ( + float(self.profiles["mass"][i]) * self.derived["fi_onlyions_vol"][i] + ) + + if self.DTplasmaBool: + self.derived["mbg_main"] = ( + self.profiles["mass"][self.Dion] + * self.derived["fi_onlyions_vol"][self.Dion] + + self.profiles["mass"][self.Tion] + * self.derived["fi_onlyions_vol"][self.Tion] + ) / ( + self.derived["fi_onlyions_vol"][self.Dion] + + self.derived["fi_onlyions_vol"][self.Tion] + ) + self.derived["fmain"] = ( + self.derived["fi_vol"][self.Dion] + self.derived["fi_vol"][self.Tion] + ) + else: + self.derived["mbg_main"] = self.profiles["mass"][self.Mion] + self.derived["fmain"] = self.derived["fi_vol"][self.Mion] + + def deriveContentByVolumes(self, rhos=[0.5], impurityPosition=3): + """ + Calculates total particles and energy for ions and electrons, at a given volume + It fails near axis because of the polynomial integral, requiring a number of poitns + """ + + min_number_points = 3 + + We_x = np.zeros(self.profiles["te(keV)"].shape[0]) + Wi_x = np.zeros(self.profiles["te(keV)"].shape[0]) + Ne_x = np.zeros(self.profiles["te(keV)"].shape[0]) + Ni_x = np.zeros(self.profiles["te(keV)"].shape[0]) + for j in range(self.profiles["te(keV)"].shape[0] - min_number_points): + i = j + min_number_points + We_x[i], Wi_x[i], Ne_x[i], _ = PLASMAtools.calculateContent( + np.expand_dims(self.profiles["rmin(m)"][:i], 0), + np.expand_dims(self.profiles["te(keV)"][:i], 0), + np.expand_dims(np.transpose(self.profiles["ti(keV)"][:i]), 0), + np.expand_dims(self.profiles["ne(10^19/m^3)"][:i] * 0.1, 0), + np.expand_dims( + np.transpose(self.profiles["ni(10^19/m^3)"][:i] * 0.1), 0 + ), + np.expand_dims(self.derived["volp_miller"][:i], 0), + ) + + _, _, Ni_x[i], _ = PLASMAtools.calculateContent( + np.expand_dims(self.profiles["rmin(m)"][:i], 0), + np.expand_dims(self.profiles["te(keV)"][:i], 0), + np.expand_dims(np.transpose(self.profiles["ti(keV)"][:i]), 0), + np.expand_dims( + self.profiles["ni(10^19/m^3)"][:i, impurityPosition] * 0.1, 0 + ), + np.expand_dims( + np.transpose(self.profiles["ni(10^19/m^3)"][:i] * 0.1), 0 + ), + np.expand_dims(self.derived["volp_miller"][:i], 0), + ) + + We, Wi, Ne, Ni = ( + np.zeros(len(rhos)), + np.zeros(len(rhos)), + np.zeros(len(rhos)), + np.zeros(len(rhos)), + ) + for i in range(len(rhos)): + We[i] = np.interp(rhos[i], self.profiles["rho(-)"], We_x) + Wi[i] = np.interp(rhos[i], self.profiles["rho(-)"], Wi_x) + Ne[i] = np.interp(rhos[i], self.profiles["rho(-)"], Ne_x) + Ni[i] = np.interp(rhos[i], self.profiles["rho(-)"], Ni_x) + + return We, Wi, Ne, Ni + + 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("Engineering Parameters:") + print(f"\tBt = {self.profiles['bcentr(T)'][0]:.2f}T, Ip = {self.profiles['current(MA)'][0]:.2f}MA (q95 = {self.derived['q95']:.2f}, q* = {self.derived['qstar']:.2f}, q*ITER = {self.derived['qstar_ITER']:.2f}), Pin = {self.derived['qIn']:.2f}MW") + print(f"\tR = {self.profiles['rcentr(m)'][0]:.2f}m, a = {self.derived['a']:.2f}m (eps = {self.derived['eps']:.3f})") + print(f"\tkappa_sep = {self.profiles['kappa(-)'][-1]:.2f}, kappa_995 = {self.derived['kappa995']:.2f}, kappa_95 = {self.derived['kappa95']:.2f}, kappa_a = {self.derived['kappa_a']:.2f}") + print(f"\tdelta_sep = {self.profiles['delta(-)'][-1]:.2f}, delta_995 = {self.derived['delta995']:.2f}, delta_95 = {self.derived['delta95']:.2f}") + 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("\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_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}") + print(f"\tp_vol = {self.derived['ptot_manual_vol']:.2f} MPa ({self.derived['pfast_fraction']*100.0:.1f}% fast)") + print(f"\tBetaN = {self.derived['BetaN']:.3f} (BetaN w/B0 = {self.derived['BetaN_engineering']:.3f})") + print(f"\tPrad = {self.derived['Prad']:.1f}MW ({self.derived['Prad'] / self.derived['qHeat'] * 100.0:.1f}% of total) ({self.derived['Prad_brem']/self.derived['Prad'] * 100.0:.1f}% brem, {self.derived['Prad_line']/self.derived['Prad'] * 100.0:.1f}% line, {self.derived['Prad_sync']/self.derived['Prad'] * 100.0:.1f}% sync)") + print("\tPsol = {0:.1f}MW (fLH = {1:.2f})".format(self.derived["Psol"], self.derived["LHratio"])) + print("Operational point ( [,] = [{0:.2f},{1:.2f}] ) and species:".format(self.derived["ne_vol20"], self.derived["Te_vol"])) + print("\t = {0:.2f} keV (/ = {1:.2f}, Ti0/Te0 = {2:.2f})".format(self.derived["Ti_vol"],self.derived["tite_vol"],self.derived["tite"][0],)) + print("\tfG = {0:.2f} ( = {1:.2f} * 10^20 m^-3)".format(self.derived["fG"], self.derived["ne_vol20"])) + 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']:.1e}]") + print(f"\tMach = {self.derived['MachNum_vol']:.2f} (vol avg)") + 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(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 export_to_table(self, table=None, name=None): + + if table is None: + table = DataTable() + + data = [name] + for var in table.variables: + if table.variables[var][1] is not None: + if table.variables[var][1].split("_")[0] == "rho": + ix = np.argmin( + np.abs( + self.profiles["rho(-)"] + - float(table.variables[var][1].split("_")[1]) + ) + ) + elif table.variables[var][1].split("_")[0] == "psi": + ix = np.argmin( + np.abs( + self.derived["psi_pol_n"] + - float(table.variables[var][1].split("_")[1]) + ) + ) + elif table.variables[var][1].split("_")[0] == "pos": + ix = int(table.variables[var][1].split("_")[1]) + vari = self.__dict__[table.variables[var][2]][table.variables[var][0]][ + ix + ] + else: + vari = self.__dict__[table.variables[var][2]][table.variables[var][0]] + + data.append(f"{vari*table.variables[var][4]:{table.variables[var][3]}}") + + table.data.append(data) + print(f"\t* Exported {name} to table") + + return table + + def makeAllThermalIonsHaveSameTemp(self, refIon=0): + SpecRef = self.Species[refIon]["N"] + tiRef = self.profiles["ti(keV)"][:, refIon] + + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "therm" and sp != refIon: + print( + f"\t\t\t- Temperature forcing {self.Species[sp]['N']} --> {SpecRef}" + ) + self.profiles["ti(keV)"][:, sp] = tiRef + + def scaleAllThermalDensities(self, scaleFactor=1.0): + scaleFactor_ions = scaleFactor + + 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_ions):.3f}" + ) + ni_orig = self.profiles["ni(10^19/m^3)"][:, sp] + self.profiles["ni(10^19/m^3)"][:, sp] = scaleFactor_ions * ni_orig + + def toNumpyArrays(self): + self.profiles.update({key: tensor.cpu().detach().cpu().numpy() for key, tensor in self.profiles.items() if isinstance(tensor, torch.Tensor)}) + self.derived.update({key: tensor.cpu().detach().cpu().numpy() for key, tensor in self.derived.items() if isinstance(tensor, torch.Tensor)}) + + def writeCurrentStatus(self, file=None, limitedNames=False): + print("\t- Writting input.gacode file") + + if file is None: + file = self.file + + with open(file, "w") as f: + for line in self.header: + f.write(line) + + for i in self.profiles: + if "(" not in i: + f.write(f"# {i}\n") + else: + f.write(f"# {i.split('(')[0]} | {i.split('(')[-1].split(')')[0]}\n") + + if i in self.titles_single: + if i == "name" and limitedNames: + newlist = [self.profiles[i][0]] + for k in self.profiles[i][1:]: + if k not in [ + "D", + "H", + "T", + "He4", + "he4", + "C", + "O", + "Ar", + "W", + ]: + newlist.append("C") + else: + newlist.append(k) + print( + f"\n\n!! Correcting ion names from {self.profiles[i]} to {newlist} to avoid TGYRO radiation error (to solve in future?)\n\n", + typeMsg="w", + ) + listWrite = newlist + else: + listWrite = self.profiles[i] + + if IOtools.isfloat(listWrite[0]): + listWrite = [f"{i:.7e}".rjust(14) for i in listWrite] + f.write(f"{''.join(listWrite)}\n") + else: + f.write(f"{' '.join(listWrite)}\n") + + else: + if len(self.profiles[i].shape) == 1: + for j, val in enumerate(self.profiles[i]): + pos = f"{j + 1}".rjust(3) + valt = f"{round(val,99):.7e}".rjust(15) + f.write(f"{pos}{valt}\n") + else: + for j, val in enumerate(self.profiles[i]): + pos = f"{j + 1}".rjust(3) + txt = "".join([f"{k:.7e}".rjust(15) for k in val]) + f.write(f"{pos}{txt}\n") + + print(f"\t\t~ File {IOtools.clipstr(file)} written") + + # Update file + self.file = file + + def writeMiminalKinetic(self, file): + setProfs = [ + "rho(-)", + "polflux(Wb/radian)", + "q(-)", + "te(keV)", + "ti(keV)", + "ne(10^19/m^3)", + ] + + with open(file, "w") as f: + for i in setProfs: + if "(" not in i: + f.write(f"# {i}\n") + else: + f.write(f"# {i.split('(')[0]} | {i.split('(')[-1].split(')')[0]}\n") + + if len(self.profiles[i].shape) > 1: + p = self.profiles[i][:, 0] + else: + p = self.profiles[i] + + for j, val in enumerate(p): + pos = f"{j + 1}".rjust(3) + valt = f"{val:.7e}".rjust(15) + f.write(f"{pos}{valt}\n") + + def changeResolution(self, n=100, rho_new=None, interpolation_function=MATHtools.extrapolateCubicSpline): + rho = copy.deepcopy(self.profiles["rho(-)"]) + + if rho_new is None: + n = int(n) + rho_new = np.linspace(rho[0], rho[-1], n) + else: + rho_new = np.unique(np.sort(rho_new)) + n = len(rho_new) + + self.profiles["nexp"] = [str(n)] + + pro = self.profiles + for i in pro: + if i not in self.titles_single: + if len(pro[i].shape) == 1: + pro[i] = interpolation_function(rho_new, rho, pro[i]) + else: + prof = [] + for j in range(pro[i].shape[1]): + pp = interpolation_function(rho_new, rho, pro[i][:, j]) + prof.append(pp) + prof = np.array(prof) + + pro[i] = np.transpose(prof) + + self.produce_shape_lists() + + self.deriveQuantities() + + print( + f"\t\t- Resolution of profiles changed to {n} points with function {interpolation_function}" + ) + + def DTplasma(self): + self.Dion, self.Tion = None, None + try: + self.Dion = np.where(self.profiles["name"] == "D")[0][0] + except: + pass + try: + self.Tion = np.where(self.profiles["name"] == "T")[0][0] + except: + pass + + if self.Dion is not None and self.Tion is not None: + self.DTplasmaBool = True + else: + self.DTplasmaBool = False + if self.Dion is not None: + self.Mion = self.Dion # Main + elif self.Tion is not None: + self.Mion = self.Tion # Main + else: + self.Mion = ( + 0 # If no D or T, assume that the main ion is the first and only + ) + + self.ion_list_main = [] + if self.DTplasmaBool: + self.ion_list_main = [self.Dion+1, self.Tion+1] + else: + self.ion_list_main = [self.Mion+1] + + self.ion_list_impurities = [i+1 for i in range(len(self.Species)) if i+1 not in self.ion_list_main] + + def remove(self, ions_list): + # First order them + ions_list.sort() + print( + "\t\t- Removing ions in positions (of ions order, no zero): ", + ions_list, + typeMsg="i", + ) + + ions_list = [i - 1 for i in ions_list] + + fail = False + + var_changes = ["name", "type", "mass", "z"] + for i in var_changes: + try: + self.profiles[i] = np.delete(self.profiles[i], ions_list) + except: + print( + f"\t\t\t* Ions {[k+1 for k in ions_list]} could not be removed", + typeMsg="w", + ) + fail = True + break + + if not fail: + var_changes = ["ni(10^19/m^3)", "ti(keV)"] + for i in var_changes: + self.profiles[i] = np.delete(self.profiles[i], ions_list, axis=1) + + if not fail: + # Ensure we extract the scalar value from the array + self.profiles["nion"] = np.array( + [str(int(self.profiles["nion"][0]) - len(ions_list))] + ) + + self.readSpecies() + self.deriveQuantities(rederiveGeometry=False) + + print("\t\t\t- Set of ions in updated profiles: ", self.profiles["name"]) + + def lumpSpecies( + self, ions_list=[2, 3], allthermal=False, forcename=None, force_integer=False, force_mass=None + ): + """ + if (D,Z1,Z2), lumping Z1 and Z2 requires ions_list = [2,3] + + if force_integer, the Zeff won't be kept exactly + """ + + # All thermal except first + if allthermal: + ions_list = [] + for i in range(len(self.Species) - 1): + if self.Species[i + 1]["S"] == "therm": + ions_list.append(i + 2) + lab = "therm" + else: + lab = "therm" + + print( + "\t\t- Lumping ions in positions (of ions order, no zero): ", + ions_list, + typeMsg="i", + ) + + if forcename is None: + forcename = "LUMPED" + + # Contributions to dilution and to Zeff + fZ1 = np.zeros(self.derived["fi"].shape[0]) + fZ2 = np.zeros(self.derived["fi"].shape[0]) + for i in ions_list: + fZ1 += self.Species[i - 1]["Z"] * self.derived["fi"][:, i - 1] + fZ2 += self.Species[i - 1]["Z"] ** 2 * self.derived["fi"][:, i - 1] + + Zr = fZ2 / fZ1 + Zr_vol = ( + CALCtools.integrateFS( + Zr, self.profiles["rmin(m)"], self.derived["volp_miller"] + )[-1] + / self.derived["volume"] + ) + + print(f'\t\t\t* Original plasma had Zeff_vol={self.derived["Zeff_vol"]:.2f}, QN error={self.derived["QN_Error"]:.4f}') + + # New specie parameters + if force_integer: + Z = round(Zr_vol) + print(f"\t\t\t* Lumped Z forced to be an integer ({Zr_vol}->{Z}), so plasma may not be quasineutral or fulfill original Zeff",typeMsg="w",) + else: + Z = Zr_vol + + A = Z * 2 if force_mass is None else force_mass + nZ = fZ1 / Z * self.profiles["ne(10^19/m^3)"] + + print(f"\t\t\t* New lumped impurity has Z={Z:.2f}, A={A:.2f} (calculated as 2*Z)") + + # Insert cases + self.profiles["nion"] = np.array([f"{int(self.profiles['nion'][0])+1}"]) + self.profiles["name"] = np.append(self.profiles["name"], forcename) + self.profiles["mass"] = np.append(self.profiles["mass"], A) + self.profiles["z"] = np.append(self.profiles["z"], Z) + self.profiles["type"] = np.append(self.profiles["type"], f"[{lab}]") + self.profiles["ni(10^19/m^3)"] = np.append( + self.profiles["ni(10^19/m^3)"], np.transpose(np.atleast_2d(nZ)), axis=1 + ) + self.profiles["ti(keV)"] = np.append( + self.profiles["ti(keV)"], + np.transpose(np.atleast_2d(self.profiles["ti(keV)"][:, 0])), + axis=1, + ) + + self.readSpecies() + self.deriveQuantities(rederiveGeometry=False) + + # Remove species + self.remove(ions_list) + + # Contributions to dilution and to Zeff + print( + f'\t\t\t* New plasma has Zeff_vol={self.derived["Zeff_vol"]:.2f}, QN error={self.derived["QN_Error"]:.4f}' + ) + + def lumpImpurities(self): + + self.lumpSpecies(ions_list=self.ion_list_impurities) + + def lumpDT(self): + + if self.DTplasmaBool: + self.lumpSpecies(ions_list=self.ion_list_main, forcename="DT", force_mass=2.5) + else: + print('\t\t- No DT plasma, so no lumping of main ions') + + self.moveSpecie(pos=len(self.Species), pos_new=1) + + def changeZeff(self, Zeff, ion_pos=2, quasineutral_ions=None, enforceSameGradients=False): + """ + if (D,Z1,Z2), pos 1 -> change Z1 + """ + + if quasineutral_ions is None: + if self.DTplasmaBool: + quasineutral_ions = [self.Dion, self.Tion] + else: + quasineutral_ions = [self.Mion] + + print(f'\t\t- Changing Zeff (from {self.derived["Zeff_vol"]:.3f} to {Zeff=:.3f}) by changing content of ion in position {ion_pos} {self.Species[ion_pos]["N"],self.Species[ion_pos]["Z"]}, quasineutralized by ions {quasineutral_ions}',typeMsg="i",) + + # Plasma needs to be in quasineutrality to start with + self.enforceQuasineutrality() + + # ------------------------------------------------------ + # Contributions to equations + # ------------------------------------------------------ + Zq = np.zeros(self.derived["fi"].shape[0]) + Zq2 = np.zeros(self.derived["fi"].shape[0]) + fZj = np.zeros(self.derived["fi"].shape[0]) + fZj2 = np.zeros(self.derived["fi"].shape[0]) + for i in range(len(self.Species)): + if i in quasineutral_ions: + Zq += self.Species[i]["Z"] + Zq2 += self.Species[i]["Z"] ** 2 + elif i != ion_pos: + fZj += self.Species[i]["Z"] * self.derived["fi"][:, i] + fZj2 += self.Species[i]["Z"] ** 2 * self.derived["fi"][:, i] + else: + Zk = self.Species[i]["Z"] + + # ------------------------------------------------------ + # Find free parameters (fk and fq) + # ------------------------------------------------------ + + fk = ( Zeff - (1-fZj)*Zq2/Zq - fZj2 ) / ( Zk**2 - Zk*Zq2/Zq) + fq = ( 1 - fZj - fk*Zk ) / Zq + + if (fq<0).any(): + raise ValueError(f"Zeff cannot be reduced by changing ion #{ion_pos} because it would require negative densities for quasineutral ions") + + # ------------------------------------------------------ + # Insert + # ------------------------------------------------------ + + fi_orig = self.derived["fi"][:, ion_pos] + + self.profiles["ni(10^19/m^3)"][:, ion_pos] = fk * self.profiles["ne(10^19/m^3)"] + for i in quasineutral_ions: + self.profiles["ni(10^19/m^3)"][:, i] = fq * self.profiles["ne(10^19/m^3)"] + + self.readSpecies() + + self.deriveQuantities(rederiveGeometry=False) + + if enforceSameGradients: + self.scaleAllThermalDensities() + self.deriveQuantities(rederiveGeometry=False) + + print(f'\t\t\t* Dilution changed from {fi_orig.mean():.2e} (vol avg) to { self.derived["fi"][:, ion_pos].mean():.2e} to achieve Zeff={self.derived["Zeff_vol"]:.3f} (fDT={self.derived["fmain"]:.3f}) [quasineutrality error = {self.derived["QN_Error"]:.1e}]') + + def moveSpecie(self, pos=2, pos_new=1): + """ + if (D,Z1,Z2), pos 1 pos_new 2-> (Z1,D,Z2) + """ + + if pos_new > pos: + pos, pos_new = pos_new, pos + + position_to_moveFROM_in_profiles = pos - 1 + position_to_moveTO_in_profiles = pos_new - 1 + + print(f'\t\t- Moving ion in position (of ions order, no zero) {pos} ({self.profiles["name"][position_to_moveFROM_in_profiles]}) to {pos_new}',typeMsg="i",) + + self.profiles["nion"] = np.array([f"{int(self.profiles['nion'][0])+1}"]) + + for ikey in ["name", "mass", "z", "type", "ni(10^19/m^3)", "ti(keV)"]: + if len(self.profiles[ikey].shape) > 1: + axis = 1 + newly = self.profiles[ikey][:, position_to_moveFROM_in_profiles] + else: + axis = 0 + newly = self.profiles[ikey][position_to_moveFROM_in_profiles] + self.profiles[ikey] = np.insert( + self.profiles[ikey], position_to_moveTO_in_profiles, newly, axis=axis + ) + + self.readSpecies() + self.deriveQuantities(rederiveGeometry=False) + + if position_to_moveTO_in_profiles > position_to_moveFROM_in_profiles: + self.remove([position_to_moveFROM_in_profiles + 1]) + else: + self.remove([position_to_moveFROM_in_profiles + 2]) + + def addSpecie(self, Z=5.0, mass=10.0, fi_vol=0.1, forcename=None): + print( + f"\t\t- Creating new specie with Z={Z}, mass={mass}, fi_vol={fi_vol}", + typeMsg="i", + ) + + if forcename is None: + forcename = "LUMPED" + + lab = "therm" + nZ = fi_vol * self.profiles["ne(10^19/m^3)"] + + self.profiles["nion"] = np.array([f"{int(self.profiles['nion'][0])+1}"]) + self.profiles["name"] = np.append(self.profiles["name"], forcename) + self.profiles["mass"] = np.append(self.profiles["mass"], mass) + self.profiles["z"] = np.append(self.profiles["z"], Z) + self.profiles["type"] = np.append(self.profiles["type"], f"[{lab}]") + self.profiles["ni(10^19/m^3)"] = np.append( + self.profiles["ni(10^19/m^3)"], np.transpose(np.atleast_2d(nZ)), axis=1 + ) + self.profiles["ti(keV)"] = np.append( + self.profiles["ti(keV)"], + np.transpose(np.atleast_2d(self.profiles["ti(keV)"][:, 0])), + axis=1, + ) + if "vtor(m/s)" in self.profiles: + self.profiles["vtor(m/s)"] = np.append( + self.profiles["vtor(m/s)"], + np.transpose(np.atleast_2d(self.profiles["vtor(m/s)"][:, 0])), + axis=1, + ) + + self.readSpecies() + self.deriveQuantities(rederiveGeometry=False) + + def correct(self, options={}, write=False, new_file=None): + """ + if name= T D LUMPED, and I want to eliminate D, removeIons = [2] + """ + + recompute_ptot = options.get("recompute_ptot", True) # Only done by default + removeIons = options.get("removeIons", []) + removeFast = options.get("removeFast", False) + quasineutrality = options.get("quasineutrality", False) + sameDensityGradients = options.get("sameDensityGradients", False) + groupQIONE = options.get("groupQIONE", False) + ensurePostiveGamma = options.get("ensurePostiveGamma", False) + ensureMachNumber = options.get("ensureMachNumber", None) + FastIsThermal = options.get("FastIsThermal", False) + + print("\t- Custom correction of input.gacode file has been requested") + + # ---------------------------------------------------------------------- + # Correct + # ---------------------------------------------------------------------- + + # Remove desired ions + if len(removeIons) > 0: + self.remove(removeIons) + + # Remove fast + if removeFast: + ions_fast = [] + for sp in range(len(self.Species)): + if self.Species[sp]["S"] != "therm": + ions_fast.append(sp + 1) + if len(ions_fast) > 0: + print( + f"\t\t- Detected fast ions in positions {ions_fast}, removing them..." + ) + self.remove(ions_fast) + # Fast as thermal + elif FastIsThermal: + self.make_fast_ions_thermal() + + # Correct LUMPED + for i in range(len(self.profiles["name"])): + if self.profiles["name"][i] in ["LUMPED", "None"]: + name = ionName( + int(self.profiles["z"][i]), int(self.profiles["mass"][i]) + ) + if name is not None: + print( + f'\t\t- Ion in position #{i+1} was named LUMPED with Z={self.profiles["z"][i]}, now it is renamed to {name}', + typeMsg="i", + ) + self.profiles["name"][i] = name + else: + print( + f'\t\t- Ion in position #{i+1} was named LUMPED with Z={self.profiles["z"][i]}, but I could not find what element it is, so doing nothing', + typeMsg="w", + ) + + # Correct qione + if groupQIONE and (np.abs(self.profiles["qione(MW/m^3)"].sum()) > 1e-14): + print('\t\t- Inserting "qione" into "qrfe"', typeMsg="i") + self.profiles["qrfe(MW/m^3)"] += self.profiles["qione(MW/m^3)"] + self.profiles["qione(MW/m^3)"] = self.profiles["qione(MW/m^3)"] * 0.0 + + # Make all thermal ions have the same gradient as the electron density, by keeping volume average constant + if sameDensityGradients: + self.enforce_same_density_gradients() + + # Enforce quasineutrality + if quasineutrality: + self.enforceQuasineutrality() + + print(f"\t\t\t* Quasineutrality error = {self.derived['QN_Error']:.1e}") + + # Recompute ptot + if recompute_ptot: + self.deriveQuantities(rederiveGeometry=False) + self.selfconsistentPTOT() + + # If I don't trust the negative particle flux in the core that comes from TRANSP... + if ensurePostiveGamma: + print("\t\t- Making particle flux always positive", typeMsg="i") + self.profiles[self.varqpar] = self.profiles[self.varqpar].clip(0) + self.profiles[self.varqpar2] = self.profiles[self.varqpar2].clip(0) + + # Mach + if ensureMachNumber is not None: + self.introduceRotationProfile(Mach_LF=ensureMachNumber) + + # ---------------------------------------------------------------------- + # Re-derive + # ---------------------------------------------------------------------- + + self.deriveQuantities(rederiveGeometry=False) + + # ---------------------------------------------------------------------- + # Write + # ---------------------------------------------------------------------- + if write: + self.writeCurrentStatus(file=new_file) + self.printInfo() + + def enforce_same_density_gradients(self): + txt = "" + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "therm": + self.profiles["ni(10^19/m^3)"][:, sp] = self.derived["fi_vol"][sp] * self.profiles["ne(10^19/m^3)"] + txt += f"{self.Species[sp]['N']} " + print(f"\t\t- Making all thermal ions ({txt}) have the same a/Ln as electrons (making them an exact flat fraction)",typeMsg="i",) + self.deriveQuantities(rederiveGeometry=False) + + def make_fast_ions_thermal(self): + modified_num = 0 + for i in range(len(self.Species)): + if self.Species[i]["S"] != "therm": + print( + f'\t\t- Specie {i} ({self.profiles["name"][i]}) was fast, but now it is considered thermal' + ) + self.Species[i]["S"] = "therm" + self.profiles["type"][i] = "[therm]" + self.profiles["ti(keV)"][:, i] = self.profiles["ti(keV)"][:, 0] + modified_num += 1 + if modified_num > 0: + print("\t- Making fast species as if they were thermal (to keep dilution effect and Qi-sum of fluxes)",typeMsg="w") + + def selfconsistentPTOT(self): + print(f"\t\t* Recomputing ptot and inserting it as ptot(Pa), changed from p0 = {self.profiles['ptot(Pa)'][0] * 1e-3:.1f} to {self.derived['ptot_manual'][0]*1e+3:.1f} kPa",typeMsg="i") + self.profiles["ptot(Pa)"] = self.derived["ptot_manual"] * 1e6 + + def enforceQuasineutrality(self): + print(f"\t\t- Enforcing quasineutrality (error = {self.derived['QN_Error']:.1e})",typeMsg="i",) + + # What's the lack of quasineutrality? + ni = self.profiles["ne(10^19/m^3)"] * 0.0 + for sp in range(len(self.Species)): + ni += self.profiles["ni(10^19/m^3)"][:, sp] * self.profiles["z"][sp] + ne_missing = self.profiles["ne(10^19/m^3)"] - ni + + # What ion to modify? + if self.DTplasmaBool: + print("\t\t\t* Enforcing quasineutrality by modifying D and T equally") + prev_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Dion]) + self.profiles["ni(10^19/m^3)"][:, self.Dion] += ne_missing / 2 + self.profiles["ni(10^19/m^3)"][:, self.Tion] += ne_missing / 2 + new_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Dion]) + else: + print( + f"\t\t\t* Enforcing quasineutrality by modifying main ion (position #{self.Mion})" + ) + prev_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Mion]) + self.profiles["ni(10^19/m^3)"][:, self.Mion] += ne_missing + new_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Mion]) + + print( + f"\t\t\t\t- Changed on-axis density from n0 = {prev_on_axis:.2f} to {new_on_axis:.2f} ({100*(new_on_axis-prev_on_axis)/prev_on_axis:.1f}%)" + ) + + self.deriveQuantities(rederiveGeometry=False) + + def introduceRotationProfile(self, Mach_LF=1.0, new_file=None): + print(f"\t- Enforcing Mach Number in LF of {Mach_LF}") + self.deriveQuantities() + Vtor_LF = PLASMAtools.constructVtorFromMach( + Mach_LF, self.profiles["ti(keV)"][:, 0], self.derived["mbg"] + ) # m/s + + self.profiles["w0(rad/s)"] = Vtor_LF / (self.derived["R_LF"]) # rad/s + + self.deriveQuantities() + + if new_file is not None: + self.writeCurrentStatus(file=new_file) + + def plot( + self, + axs1=None, + axs2=None, + axs3=None, + axs4=None, + axsFlows=None, + axs6=None, + axsImps=None, + color="b", + legYN=True, + extralab="", + fn=None, + fnlab="", + lsFlows="-", + legFlows=True, + showtexts=True, + lastRhoGradients=0.89, + ): + if axs1 is None: + if fn is None: + from mitim_tools.misc_tools.GUItools import FigureNotebook + + self.fn = FigureNotebook("PROFILES Notebook", geometry="1600x1000") + + fig, fig2, fig3, fig4, fig5, fig6, fig7 = add_figures(self.fn, fnlab=fnlab) + + grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) + axs1 = [ + fig.add_subplot(grid[0, 0]), + fig.add_subplot(grid[1, 0]), + fig.add_subplot(grid[2, 0]), + fig.add_subplot(grid[0, 1]), + fig.add_subplot(grid[1, 1]), + fig.add_subplot(grid[2, 1]), + fig.add_subplot(grid[0, 2]), + fig.add_subplot(grid[1, 2]), + fig.add_subplot(grid[2, 2]), + ] + + + grid = plt.GridSpec(3, 2, hspace=0.3, wspace=0.3) + axs2 = [ + fig2.add_subplot(grid[0, 0]), + fig2.add_subplot(grid[0, 1]), + fig2.add_subplot(grid[1, 0]), + fig2.add_subplot(grid[1, 1]), + fig2.add_subplot(grid[2, 0]), + fig2.add_subplot(grid[2, 1]), + ] + + + grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.5) + ax00c = fig3.add_subplot(grid[0, 0]) + axs3 = [ + ax00c, + fig3.add_subplot(grid[1, 0], sharex=ax00c), + fig3.add_subplot(grid[2, 0], sharex=ax00c), + fig3.add_subplot(grid[0, 1], sharex=ax00c), + fig3.add_subplot(grid[1, 1], sharex=ax00c), + fig3.add_subplot(grid[2, 1], sharex=ax00c), + fig3.add_subplot(grid[0, 2], sharex=ax00c), + fig3.add_subplot(grid[1, 2], sharex=ax00c), + fig3.add_subplot(grid[2, 2], sharex=ax00c), + fig3.add_subplot(grid[0, 3], sharex=ax00c), + fig3.add_subplot(grid[1, 3], sharex=ax00c), + fig3.add_subplot(grid[2, 3], sharex=ax00c), + ] + + + grid = plt.GridSpec(2, 3, 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]), + ] + + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + + axsFlows = [ + fig5.add_subplot(grid[0, 0]), + fig5.add_subplot(grid[1, 0]), + fig5.add_subplot(grid[0, 1]), + fig5.add_subplot(grid[0, 2]), + fig5.add_subplot(grid[1, 1]), + fig5.add_subplot(grid[1, 2]), + ] + + + grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) + axs6 = [ + fig6.add_subplot(grid[0, 0]), + fig6.add_subplot(grid[:, 1]), + fig6.add_subplot(grid[0, 2]), + fig6.add_subplot(grid[1, 0]), + fig6.add_subplot(grid[1, 2]), + fig6.add_subplot(grid[0, 3]), + fig6.add_subplot(grid[1, 3]), + ] + + + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + axsImps = [ + fig7.add_subplot(grid[0, 0]), + fig7.add_subplot(grid[0, 1]), + fig7.add_subplot(grid[0, 2]), + fig7.add_subplot(grid[1, 0]), + fig7.add_subplot(grid[1, 1]), + fig7.add_subplot(grid[1, 2]), + ] + + [ax00, ax10, ax20, ax01, ax11, ax21, ax02, ax12, ax22] = axs1 + [ax00b, ax01b, ax10b, ax11b, ax20b, ax21b] = axs2 + [ + ax00c, + ax10c, + ax20c, + ax01c, + ax11c, + ax21c, + ax02c, + ax12c, + ax22c, + ax03c, + ax13c, + ax23c, + ] = axs3 + + lw = 1 + fs = 6 + rho = self.profiles["rho(-)"] + + lines = ["-", "--", "-.", ":", "-", "--", "-."] + + self.plot_temps(ax=ax00, leg=legYN, col=color, lw=lw, fs=fs, extralab=extralab) + self.plot_dens(ax=ax01, leg=legYN, col=color, lw=lw, fs=fs, extralab=extralab) + + ax = ax10 + cont = 0 + for i in range(len(self.Species)): + if self.Species[i]["S"] == "therm": + var = self.profiles["ti(keV)"][:, i] + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + varL = "Thermal $T_i$ (keV)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax20 + cont = 0 + for i in range(len(self.Species)): + if self.Species[i]["S"] == "fast": + var = self.profiles["ti(keV)"][:, i] + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + varL = "Fast $T_i$ (keV)" + ax.plot( + rho, + self.profiles["ti(keV)"][:, 0], + lw=0.5, + ls="-", + alpha=0.5, + c=color, + label=extralab + "$T_{i,1}$", + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax11 + cont = 0 + for i in range(len(self.Species)): + if self.Species[i]["S"] == "therm": + var = self.profiles["ni(10^19/m^3)"][:, i] * 1e-1 + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + varL = "Thermal $n_i$ ($10^{20}/m^3$)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax21 + cont = 0 + for i in range(len(self.Species)): + if self.Species[i]["S"] == "fast": + var = self.profiles["ni(10^19/m^3)"][:, i] * 1e-1 * 1e5 + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + varL = "Fast $n_i$ ($10^{15}/m^3$)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN and cont>0: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax02 + var = self.profiles["w0(rad/s)"] + ax.plot(rho, var, lw=lw, ls="-", c=color) + varL = "$\\omega_{0}$ (rad/s)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax12 + var = self.profiles["ptot(Pa)"] * 1e-6 + ax.plot(rho, var, lw=lw, ls="-", c=color, label=extralab + "ptot") + if "ptot_manual" in self.derived: + ax.plot( + rho, + self.derived["ptot_manual"], + lw=lw, + ls="--", + c=color, + label=extralab + "check", + ) + # ax.plot(rho,np.abs(var-self.derived['ptot_manual']),lw=lw,ls='-.',c=color,label=extralab+'diff') + + ax.plot( + rho, + self.derived["pthr_manual"], + lw=lw, + ls="-.", + c=color, + label=extralab + "check, thrm", + ) + + + varL = "$p$ (MPa)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + # ax.set_ylim(bottom=0) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax00b + varL = "$MW/m^3$" + cont = 0 + var = -self.profiles["qei(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "i->e", c=color) + cont += 1 + if "qrfe(MW/m^3)" in self.profiles: + var = self.profiles["qrfe(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "rf", c=color) + cont += 1 + if "qfuse(MW/m^3)" in self.profiles: + var = self.profiles["qfuse(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "fus", c=color) + cont += 1 + if "qbeame(MW/m^3)" in self.profiles: + var = self.profiles["qbeame(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "beam", c=color) + cont += 1 + if "qione(MW/m^3)" in self.profiles: + var = self.profiles["qione(MW/m^3)"] + ax.plot( + rho, var, lw=lw / 2, ls=lines[cont], label=extralab + "extra", c=color + ) + cont += 1 + if "qohme(MW/m^3)" in self.profiles: + var = self.profiles["qohme(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "ohmic", c=color) + cont += 1 + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + ax.set_title("Electron Power Density") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax01b + if "varqmom" not in self.__dict__: + self.varqmom = "qmom(N/m^2)" + self.profiles[self.varqmom] = self.profiles["rho(-)"] * 0.0 + + ax.plot(rho, self.profiles[self.varqmom], lw=lw, ls="-", c=color) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$N/m^2$, $J/m^3$") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + ax.set_title("Momentum Source Density") + + ax = ax21b + ax.plot( + rho, self.derived["qe_MWm2"], lw=lw, ls="-", label=extralab + "qe", c=color + ) + ax.plot( + rho, self.derived["qi_MWm2"], lw=lw, ls="--", label=extralab + "qi", c=color + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("Heat Flux ($MW/m^2$)") + if legYN: + ax.legend(loc="lower left", fontsize=fs) + ax.set_title("Flux per unit area (gacode: P/V')") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax21b.twinx() + ax.plot( + rho, + self.derived["ge_10E20m2"], + lw=lw, + ls="-.", + label=extralab + "$\\Gamma_e$", + c=color, + ) + ax.set_ylabel("Particle Flux ($10^{20}/m^2/s$)") + if legYN: + ax.legend(loc="lower right", fontsize=fs) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax20b + varL = "$Q_{rad}$ ($MW/m^3$)" + if "qbrem(MW/m^3)" in self.profiles: + var = self.profiles["qbrem(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls="-", label=extralab + "brem", c=color) + if "qline(MW/m^3)" in self.profiles: + var = self.profiles["qline(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls="--", label=extralab + "line", c=color) + if "qsync(MW/m^3)" in self.profiles: + var = self.profiles["qsync(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=":", label=extralab + "sync", c=color) + + var = self.derived["qrad"] + ax.plot(rho, var, lw=lw * 1.5, ls="-", label=extralab + "Total", c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + ax.set_title("Radiation Contributions") + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax10b + varL = "$MW/m^3$" + cont = 0 + var = self.profiles["qei(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "e->i", c=color) + cont += 1 + if "qrfi(MW/m^3)" in self.profiles: + var = self.profiles["qrfi(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "rf", c=color) + cont += 1 + if "qfusi(MW/m^3)" in self.profiles: + var = self.profiles["qfusi(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "fus", c=color) + cont += 1 + if "qbeami(MW/m^3)" in self.profiles: + var = self.profiles["qbeami(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "beam", c=color) + cont += 1 + if "qioni(MW/m^3)" in self.profiles: + var = self.profiles["qioni(MW/m^3)"] + ax.plot( + rho, var, lw=lw / 2, ls=lines[cont], label=extralab + "extra", c=color + ) + cont += 1 + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + ax.set_title("Ion Power Density") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + """ + Note that in prgen_map_plasmastate, that variable: + expro_qpar_beam(i) = plst_sn_trans(i-1)/dvol + + Note that in prgen_read_plasmastate, that variable: + ! Particle source + err = nf90_inq_varid(ncid,trim('sn_trans'),varid) + err = nf90_get_var(ncid,varid,plst_sn_trans(1:nx-1)) + plst_sn_trans(nx) = 0.0 + + Note that in the plasmastate file, the variable "sn_trans": + + long_name: particle transport (loss) + units: #/sec + component: PLASMA + section: STATE_PROFILES + specification: R|units=#/sec|step*dV sn_trans(~nrho,0:nspec_th) + + So, this means that expro_qpar_beam is in units of #/sec/m^3, meaning that + it is a particle flux DENSITY. It therefore requires volume integral and + divide by surface to produce a flux. + + The units of this qpar_beam column is NOT MW/m^3. In the gacode source codes + they also say that those units are wrong. + + """ + + ax = ax11b + cont = 0 + var = self.profiles[self.varqpar] * 1e-20 + ax.plot(rho, var, lw=lw, ls=lines[0], c=color, label=extralab + "beam") + var = self.profiles[self.varqpar2] * 1e-20 + ax.plot(rho, var, lw=lw, ls=lines[1], c=color, label=extralab + "wall") + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.axhline(y=0, lw=0.5, ls="--", c="k") + ax.set_ylabel("$10^{20}m^{-3}s^{-1}$") + ax.set_title("Particle Source Density") + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax00c + varL = "cos Shape Params" + yl = 0 + cont = 0 + + for i, s in enumerate(self.shape_cos): + if s is not None: + valmax = np.abs(s).max() + if valmax > 1e-10: + lab = f"c{i}" + ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) + cont += 1 + + yl = np.max([yl, valmax]) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + + + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + if legYN: + ax.legend(loc="best", fontsize=fs) + + ax = ax01c + varL = "sin Shape Params" + cont = 0 + for i, s in enumerate(self.shape_sin): + if s is not None: + valmax = np.abs(s).max() + if valmax > 1e-10: + lab = f"s{i}" + ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) + cont += 1 + + yl = np.max([yl, valmax]) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax02c + var = self.profiles["q(-)"] + ax.plot(rho, var, lw=lw, ls="-", c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0) + ax.set_ylabel("q") + + ax.axhline(y=1, ls="--", c="k", lw=1) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0.0) + + + + ax = ax12c + var = self.profiles["polflux(Wb/radian)"] + ax.plot(rho, var, lw=lw, ls="-", c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("Poloidal $\\psi$ ($Wb/rad$)") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax10c + + var = self.profiles["rho(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0) + ax.set_ylabel("$\\rho$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax11c + + var = self.profiles["rmin(m)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylim(bottom=0) + ax.set_ylabel("$r_{min}$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax20c + + var = self.profiles["rmaj(m)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$R_{maj}$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax21c + + var = self.profiles["zmag(m)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + yl = np.max([0.1, np.max(np.abs(var))]) + ax.set_ylim([-yl, yl]) + ax.set_ylabel("$Z_{maj}$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax22c + + var = self.profiles["kappa(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$\\kappa$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=1) + + ax = ax03c + + var = self.profiles["delta(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$\\delta$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax13c + + var = self.profiles["zeta(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("zeta") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax23c + + var = self.profiles["johm(MA/m^2)"] + ax.plot(rho, var, "-", lw=lw, c=color, label=extralab + "$J_{OH}$") + var = self.profiles["jbs(MA/m^2)"] + ax.plot(rho, var, "--", lw=lw, c=color, label=extralab + "$J_{BS,par}$") + var = self.profiles["jbstor(MA/m^2)"] + ax.plot(rho, var, "-.", lw=lw, c=color, label=extralab + "$J_{BS,tor}$") + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylim(bottom=0) + ax.set_ylabel("J ($MA/m^2$)") + if legYN: + ax.legend(loc="best", prop={"size": 7}) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + # Derived + self.plotGradients( + axs4, color=color, lw=lw, lastRho=lastRhoGradients, label=extralab + ) + + # Others + ax = axs6[0] + ax.plot(self.profiles["rho(-)"], self.derived["dw0dr"] * 1e-5, c=color, lw=lw) + ax.set_ylabel("$-d\\omega_0/dr$ (krad/s/cm)") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + ax.axhline(y=0, lw=1.0, c="k", ls="--") + + ax = axs6[2] + ax.plot(self.profiles["rho(-)"], self.derived["q_fus"], c=color, lw=lw) + ax.set_ylabel("$q_{fus}$ ($MW/m^3$)") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axs6[3] + ax.plot(self.profiles["rho(-)"], self.derived["q_fus_MWmiller"], c=color, lw=lw) + ax.set_ylabel("$P_{fus}$ ($MW$)") + ax.set_xlim([0, 1]) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axs6[4] + ax.plot(self.profiles["rho(-)"], self.derived["tite"], c=color, lw=lw) + ax.set_ylabel("$T_i/T_e$") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + ax.axhline(y=1, ls="--", lw=1.0, c="k") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = axs6[5] + if "MachNum" in self.derived: + ax.plot(self.profiles["rho(-)"], self.derived["MachNum"], c=color, lw=lw) + ax.set_ylabel("Mach Number") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + ax.axhline(y=0, ls="--", c="k", lw=0.5) + ax.axhline(y=1, ls="--", c="k", lw=0.5) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = axs6[6] + safe_division = np.divide( + self.derived["qi_MWm2"], + self.derived["qe_MWm2"], + where=self.derived["qe_MWm2"] != 0, + out=np.full_like(self.derived["qi_MWm2"], np.nan), + ) + ax.plot( + self.profiles["rho(-)"], + safe_division, + c=color, + lw=lw, + label=extralab + "$Q_i/Q_e$", + ) + safe_division = np.divide( + self.derived["qi_aux_MWmiller"], + self.derived["qe_aux_MWmiller"], + where=self.derived["qe_aux_MWmiller"] != 0, + out=np.full_like(self.derived["qi_aux_MWmiller"], np.nan), + ) + ax.plot( + self.profiles["rho(-)"], + safe_division, + c=color, + lw=lw, + ls="--", + label=extralab + "$P_i/P_e$", + ) + ax.set_ylabel("Power ratios") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + ax.axhline(y=1.0, ls="--", c="k", lw=1.0) + GRAPHICStools.addDenseAxis(ax) + # GRAPHICStools.autoscale_y(ax,bottomy=0) + ax.set_ylim(bottom=0) + ax.legend(loc="best", fontsize=fs) + + # Final + if axsFlows is not None: + self.plotBalance( + axs=axsFlows, ls=lsFlows, leg=legFlows, showtexts=showtexts + ) + + # Geometry + ax = axs6[1] + self.plotGeometry(ax=ax, color=color) + + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + GRAPHICStools.addDenseAxis(ax) + + # Impurities + ax = axsImps[0] + for i in range(len(self.Species)): + var = ( + self.profiles["ni(10^19/m^3)"][:, i] + / self.profiles["ni(10^19/m^3)"][0, i] + ) + ax.plot( + rho, + var, + lw=lw, + ls=lines[i], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + varL = "$n_i/n_{i,0}$" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axsImps[1] + for i in range(len(self.Species)): + var = self.derived["fi"][:, i] + ax.plot( + rho, + var, + lw=lw, + ls=lines[i], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + varL = "$f_i$" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + ax.set_ylim([0, 1]) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axsImps[2] + + lastRho = 0.9 + + ix = np.argmin(np.abs(self.profiles["rho(-)"] - lastRho)) + 1 + ax.plot( + rho[:ix], self.derived["aLne"][:ix], lw=lw * 3, ls="-", c=color, label="e" + ) + for i in range(len(self.Species)): + var = self.derived["aLni"][:, i] + ax.plot( + rho[:ix], + var[:ix], + lw=lw, + ls=lines[i], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + varL = "$a/L_{ni}$" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axsImps[5] + + ax = axsImps[3] + ax.plot(self.profiles["rho(-)"], self.derived["Zeff"], c=color, lw=lw) + ax.set_ylabel("$Z_{eff}$") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axsImps[4] + cont = 0 + if "vtor(m/s)" in self.profiles: + for i in range(len(self.Species)): + try: # REMOVE FOR FUTURE + var = self.profiles["vtor(m/s)"][:, i] * 1e-3 + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + except: + break + varL = "$V_{tor}$ (km/s)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if "vtor(m/s)" in self.profiles and legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + def plotGradients( + self, + axs4, + color="b", + lw=1.0, + label="", + ls="-o", + lastRho=0.89, + ms=2, + alpha=1.0, + useRoa=False, + RhoLocationsPlot=None, + plotImpurity=None, + plotRotation=False, + autoscale=True, + ): + + if RhoLocationsPlot is None: RhoLocationsPlot=[] + + if axs4 is None: + plt.ion() + fig, axs = plt.subplots( + ncols=3 + int(plotImpurity is not None) + int(plotRotation), + nrows=2, + figsize=(12, 5), + ) + + axs4 = [] + for i in range(axs.shape[-1]): + axs4.append(axs[0, i]) + axs4.append(axs[1, i]) + + ix = np.argmin(np.abs(self.profiles["rho(-)"] - lastRho)) + 1 + + xcoord = self.profiles["rho(-)"] if (not useRoa) else self.derived["roa"] + labelx = "$\\rho$" if (not useRoa) else "$r/a$" + + ax = axs4[0] + ax.plot( + xcoord, + self.profiles["te(keV)"], + ls, + c=color, + lw=lw, + label=label, + markersize=ms, + alpha=alpha, + ) + ax = axs4[2] + ax.plot( + xcoord, + self.profiles["ti(keV)"][:, 0], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + ax = axs4[4] + ax.plot( + xcoord, + self.profiles["ne(10^19/m^3)"] * 1e-1, + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + + if "derived" in self.__dict__: + ax = axs4[1] + ax.plot( + xcoord[:ix], + self.derived["aLTe"][:ix], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + ax = axs4[3] + ax.plot( + xcoord[:ix], + self.derived["aLTi"][:ix, 0], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + ax = axs4[5] + ax.plot( + xcoord[:ix], + self.derived["aLne"][:ix], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + + for ax in axs4: + ax.set_xlim([0, 1]) + + ax = axs4[0] + ax.set_ylabel("$T_e$ (keV)") + ax.set_xlabel(labelx) + if autoscale: + 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) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = axs4[4] + ax.set_ylabel("$n_e$ ($10^{20}m^{-3}$)") + ax.set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axs4[1] + ax.set_ylabel("$a/L_{Te}$") + ax.set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = axs4[3] + ax.set_ylabel("$a/L_{Ti}$") + ax.set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = axs4[5] + ax.set_ylabel("$a/L_{ne}$") + ax.axhline(y=0, ls="--", lw=0.5, c="k") + ax.set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + + cont = 0 + if plotImpurity is not None: + axs4[6 + cont].plot( + xcoord, + self.profiles["ni(10^19/m^3)"][:, plotImpurity] * 1e-1, + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + axs4[6 + cont].set_ylabel("$n_Z$ ($10^{20}m^{-3}$)") + axs4[6].set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + if "derived" in self.__dict__: + axs4[7 + cont].plot( + xcoord[:ix], + self.derived["aLni"][:ix, plotImpurity], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + axs4[7 + cont].set_ylabel("$a/L_{nZ}$") + axs4[7 + cont].axhline(y=0, ls="--", lw=0.5, c="k") + axs4[7 + cont].set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + cont += 2 + + if plotRotation: + axs4[6 + cont].plot( + xcoord, + self.profiles["w0(rad/s)"] * 1e-3, + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + axs4[6 + cont].set_ylabel("$w_0$ (krad/s)") + axs4[6 + cont].set_xlabel(labelx) + if "derived" in self.__dict__: + axs4[7 + cont].plot( + xcoord[:ix], + self.derived["dw0dr"][:ix] * 1e-5, + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + axs4[7 + cont].set_ylabel("-$d\\omega_0/dr$ (krad/s/cm)") + axs4[7 + cont].axhline(y=0, ls="--", lw=0.5, c="k") + axs4[7 + cont].set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + cont += 2 + + for x0 in RhoLocationsPlot: + ix = np.argmin(np.abs(self.profiles["rho(-)"] - x0)) + for ax in axs4: + ax.axvline(x=xcoord[ix], ls="--", lw=0.5, c=color) + + for i in range(len(axs4)): + ax = axs4[i] + GRAPHICStools.addDenseAxis(ax) + + def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): + if axs is None: + fig1 = plt.figure() + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + + axs = [ + fig1.add_subplot(grid[0, 0]), + fig1.add_subplot(grid[1, 0]), + fig1.add_subplot(grid[0, 1]), + fig1.add_subplot(grid[0, 2]), + fig1.add_subplot(grid[1, 1]), + fig1.add_subplot(grid[1, 2]), + ] + + # Profiles + + ax = axs[0] + axT = axs[1] + roa = self.profiles["rmin(m)"] / self.profiles["rmin(m)"][-1] + Te = self.profiles["te(keV)"] + ne = self.profiles["ne(10^19/m^3)"] * 1e-1 + ni = self.profiles["ni(10^19/m^3)"] * 1e-1 + niT = np.sum(ni, axis=1) + Ti = self.profiles["ti(keV)"][:, 0] + ax.plot(roa, Te, lw=2, c="r", label="$T_e$" if leg else "", ls=ls) + ax.plot(roa, Ti, lw=2, c="b", label="$T_i$" if leg else "", ls=ls) + axT.plot(roa, ne, lw=2, c="m", label="$n_e$" if leg else "", ls=ls) + axT.plot(roa, niT, lw=2, c="c", label="$\\sum n_i$" if leg else "", ls=ls) + if limits is not None: + [roa_first, roa_last] = limits + ax.plot(roa_last, np.interp(roa_last, roa, Te), "s", c="r", markersize=3) + ax.plot(roa_first, np.interp(roa_first, roa, Te), "s", c="r", markersize=3) + ax.plot(roa_last, np.interp(roa_last, roa, Ti), "s", c="b", markersize=3) + ax.plot(roa_first, np.interp(roa_first, roa, Ti), "s", c="b", markersize=3) + axT.plot(roa_last, np.interp(roa_last, roa, ne), "s", c="m", markersize=3) + axT.plot(roa_first, np.interp(roa_first, roa, ne), "s", c="m", markersize=3) + + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + axT.set_xlabel("r/a") + axT.set_xlim([0, 1]) + ax.set_ylabel("$T$ (keV)") + ax.set_ylim(bottom=0) + axT.set_ylabel("$n$ ($10^{20}m^{-3}$)") + axT.set_ylim(bottom=0) + # axT.set_ylim([0,np.max(ne)*1.5]) + ax.legend() + axT.legend() + ax.set_title("Final Temperature profiles") + axT.set_title("Final Density profiles") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + GRAPHICStools.addDenseAxis(axT) + GRAPHICStools.autoscale_y(axT, bottomy=0) + + if showtexts: + if self.derived["Q"] > 0.005: + ax.text( + 0.05, + 0.05, + f"Pfus = {self.derived['Pfus']:.1f}MW, Q = {self.derived['Q']:.2f}", + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + transform=ax.transAxes, + ) + + axT.text( + 0.05, + 0.4, + "ne_20 = {0:.1f} (fG = {1:.2f}), Zeff = {2:.1f}".format( + self.derived["ne_vol20"], + self.derived["fG"], + self.derived["Zeff_vol"], + ), + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + transform=axT.transAxes, + ) + + # F + ax = axs[2] + P = ( + self.derived["qe_fus_MWmiller"] + + self.derived["qe_aux_MWmiller"] + + -self.derived["qe_rad_MWmiller"] + + -self.derived["qe_exc_MWmiller"] + ) + + ax.plot( + roa, + -self.derived["qe_MWmiller"], + c="g", + lw=2, + label="$P_{e}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qe_fus_MWmiller"], + c="r", + lw=2, + label="$P_{fus,e}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qe_aux_MWmiller"], + c="b", + lw=2, + label="$P_{aux,e}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + -self.derived["qe_exc_MWmiller"], + c="m", + lw=2, + label="$P_{exc,e}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + -self.derived["qe_rad_MWmiller"], + c="c", + lw=2, + label="$P_{rad,e}$" if leg else "", + ls=ls, + ) + ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) + + # Pe = self.profiles['te(keV)']*1E3*e_J*self.profiles['ne(10^19/m^3)']*1E-1*1E20 *1E-6 + # ax.plot(roa,Pe,ls='-',lw=3,alpha=0.1,c='k',label='$W_e$ (MJ/m^3)') + + ax.plot( + roa, + -self.derived["ce_MWmiller"], + c="k", + lw=1, + label="($P_{conv,e}$)" if leg else "", + ) + + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + ax.set_ylabel("$P$ (MW)") + # ax.set_ylim(bottom=0) + ax.set_title("Electron Thermal Flows") + ax.axhline(y=0.0, lw=0.5, ls="--", c="k") + GRAPHICStools.addLegendApart( + ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" + ) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axs[3] + P = ( + self.derived["qi_fus_MWmiller"] + + self.derived["qi_aux_MWmiller"] + + self.derived["qe_exc_MWmiller"] + ) + + ax.plot( + roa, + -self.derived["qi_MWmiller"], + c="g", + lw=2, + label="$P_{i}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qi_fus_MWmiller"], + c="r", + lw=2, + label="$P_{fus,i}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qi_aux_MWmiller"], + c="b", + lw=2, + label="$P_{aux,i}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qe_exc_MWmiller"], + c="m", + lw=2, + label="$P_{exc,i}$" if leg else "", + ls=ls, + ) + ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) + + # Pi = self.profiles['ti(keV)'][:,0]*1E3*e_J*self.profiles['ni(10^19/m^3)'][:,0]*1E-1*1E20 *1E-6 + # ax.plot(roa,Pi,ls='-',lw=3,alpha=0.1,c='k',label='$W_$ (MJ/m^3)') + + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + ax.set_ylabel("$P$ (MW)") + # ax.set_ylim(bottom=0) + ax.set_title("Ion Thermal Flows") + ax.axhline(y=0.0, lw=0.5, ls="--", c="k") + GRAPHICStools.addLegendApart( + ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" + ) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + # F + ax = axs[4] + + ax.plot( + roa, + self.derived["ge_10E20miller"], + c="g", + lw=2, + label="$\\Gamma_{e}$" if leg else "", + ls=ls, + ) + # ax.plot(roa,self.profiles['ne(10^19/m^3)']*1E-1,lw=3,alpha=0.1,c='k',label='$n_e$ ($10^{20}/m^3$)' if leg else '',ls=ls) + + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + ax.set_ylabel("$\\Gamma$ ($10^{20}/s$)") + ax.set_title("Particle Flows") + ax.axhline(y=0.0, lw=0.5, ls="--", c="k") + GRAPHICStools.addLegendApart( + ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" + ) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + # TOTAL + ax = axs[5] + P = ( + self.derived["qOhm_MWmiller"] + + self.derived["qRF_MWmiller"] + + self.derived["qFus_MWmiller"] + + -self.derived["qe_rad_MWmiller"] + + self.derived["qz_MWmiller"] + + self.derived["qBEAM_MWmiller"] + ) + + ax.plot( + roa, + -self.derived["q_MWmiller"], + c="g", + lw=2, + label="$P$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qOhm_MWmiller"], + c="k", + lw=2, + label="$P_{Oh}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qRF_MWmiller"], + c="b", + lw=2, + label="$P_{RF}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qBEAM_MWmiller"], + c="pink", + lw=2, + label="$P_{NBI}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qFus_MWmiller"], + c="r", + lw=2, + label="$P_{fus}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + -self.derived["qe_rad_MWmiller"], + c="c", + lw=2, + label="$P_{rad}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qz_MWmiller"], + c="orange", + lw=1, + label="$P_{ionz.}$" if leg else "", + ls=ls, + ) + + # P = Pe+Pi + # ax.plot(roa,P,ls='-',lw=3,alpha=0.1,c='k',label='$W$ (MJ)') + + ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + ax.set_ylabel("$P$ (MW)") + # ax.set_ylim(bottom=0) + ax.set_title("Total Thermal Flows") + + GRAPHICStools.addLegendApart( + ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" + ) + + ax.axhline(y=0.0, lw=0.5, ls="--", c="k") + # GRAPHICStools.drawLineWithTxt(ax,0.0,label='',orientation='vertical',color='k',lw=1,ls='--',alpha=1.0,fontsize=10,fromtop=0.85,fontweight='normal', + # verticalalignment='bottom',horizontalalignment='left',separation=0) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + def plot_temps(self, ax=None, leg=False, col="b", lw=2, extralab="", fs=10): + if ax is None: + fig, ax = plt.subplots() + + rho = self.profiles["rho(-)"] + + var = self.profiles["te(keV)"] + varL = "$T_e$ , $T_i$ (keV)" + if leg: + lab = extralab + "e" + else: + lab = "" + ax.plot(rho, var, lw=lw, ls="-", label=lab, c=col) + var = self.profiles["ti(keV)"][:, 0] + if leg: + lab = extralab + "i" + else: + lab = "" + ax.plot(rho, var, lw=lw, ls="--", label=lab, c=col) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + + def plot_dens(self, ax=None, leg=False, col="b", lw=2, extralab="", fs=10): + if ax is None: + fig, ax = plt.subplots() + + rho = self.profiles["rho(-)"] + + var = self.profiles["ne(10^19/m^3)"] * 1e-1 + varL = "$n_e$ ($10^{20}/m^3$)" + if leg: + lab = extralab + "e" + else: + lab = "" + ax.plot(rho, var, lw=lw, ls="-", label=lab, c=col) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + + def plotGeometry(self, ax=None, surfaces_rho=np.linspace(0, 1, 11), color="b", label = '', lw=1.0, lw1=2.0): + if ("R_surface" in self.derived) and (self.derived["R_surface"] is not None): + if ax is None: + plt.ion() + fig, ax = plt.subplots() + provided = False + else: + provided = True + + for rho in surfaces_rho: + ir = np.argmin(np.abs(self.profiles["rho(-)"] - rho)) + + ax.plot( + self.derived["R_surface"][ir, :], + self.derived["Z_surface"][ir, :], + "-", + lw=lw if rho<1.0 else lw1, + c=color, + ) + + ax.axhline(y=0, ls="--", lw=0.2, c="k") + ax.plot( + [self.profiles["rmaj(m)"][0]], + [self.profiles["zmag(m)"][0]], + "o", + markersize=2, + c=color, + label = label + ) + + if not provided: + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_title("Surfaces @ rho=" + str(surfaces_rho), fontsize=8) + ax.set_aspect("equal") + else: + print("\t- Cannot plot flux surface geometry", typeMsg="w") + + def plotPeaking( + self, ax, c="b", marker="*", label="", debugPlot=False, printVals=False + ): + nu_effCGYRO = self.derived["nu_eff"] * 2 / self.derived["Zeff_vol"] + ne_peaking = self.derived["ne_peaking0.2"] + ax.scatter([nu_effCGYRO], [ne_peaking], s=400, c=c, marker=marker, label=label) + + if printVals: + print(f"\t- nu_eff = {nu_effCGYRO}, ne_peaking = {ne_peaking}") + + # Extra + r = self.profiles["rmin(m)"] + volp = self.derived["volp_miller"] + ix = np.argmin(np.abs(self.profiles["rho(-)"] - 0.9)) + + if debugPlot: + fig, axq = plt.subplots() + + ne = self.profiles["ne(10^19/m^3)"] + axq.plot(self.profiles["rho(-)"], ne, color="m") + ne_vol = ( + CALCtools.integrateFS(ne * 0.1, r, volp)[-1] / self.derived["volume"] + ) + axq.axhline(y=ne_vol * 10, color="m") + + ne = copy.deepcopy(self.profiles["ne(10^19/m^3)"]) + ne[ix:] = (0,) * len(ne[ix:]) + ne_vol = CALCtools.integrateFS(ne * 0.1, r, volp)[-1] / self.derived["volume"] + ne_peaking0 = ( + ne[np.argmin(np.abs(self.derived["rho_pol"] - 0.2))] * 0.1 / ne_vol + ) + + if debugPlot: + axq.plot(self.profiles["rho(-)"], ne, color="r") + axq.axhline(y=ne_vol * 10, color="r") + + ne = copy.deepcopy(self.profiles["ne(10^19/m^3)"]) + ne[ix:] = (ne[ix],) * len(ne[ix:]) + ne_vol = CALCtools.integrateFS(ne * 0.1, r, volp)[-1] / self.derived["volume"] + ne_peaking1 = ( + ne[np.argmin(np.abs(self.derived["rho_pol"] - 0.2))] * 0.1 / ne_vol + ) + + ne_peaking0 = ne_peaking + + ax.errorbar( + [nu_effCGYRO], + [ne_peaking], + yerr=[[ne_peaking - ne_peaking1], [ne_peaking0 - ne_peaking]], + marker=marker, + c=c, + markersize=16, + capsize=2.0, + fmt="s", + elinewidth=1.0, + capthick=1.0, + ) + + if debugPlot: + axq.plot(self.profiles["rho(-)"], ne, color="b") + axq.axhline(y=ne_vol * 10, color="b") + plt.show() + + # print(f'{ne_peaking0}-{ne_peaking}-{ne_peaking1}') + + return nu_effCGYRO, ne_peaking + + def plotRelevant(self, axs = None, color = 'b', label ='', lw = 1, ms = 1): + + if axs is None: + fig = plt.figure() + axs = fig.subplot_mosaic( + """ + ABCDH + AEFGI + """ + ) + axs = [axs['A'], axs['B'], axs['C'], axs['D'], axs['E'], axs['F'], axs['G'], axs['H'], axs['I']] + + # ---------------------------------- + # Equilibria + # ---------------------------------- + + ax = axs[0] + rho = np.linspace(0, 1, 21) + + self.plotGeometry(ax=ax, surfaces_rho=rho, label=label, color=color, lw=lw, lw1=lw*3) + + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_aspect("equal") + ax.legend(prop={'size':8}) + GRAPHICStools.addDenseAxis(ax) + ax.set_title("Equilibria") + + # ---------------------------------- + # Kinetic Profiles + # ---------------------------------- + + # T profiles + ax = axs[1] + + ax.plot(self.profiles['rho(-)'], self.profiles['te(keV)'], '-o', markersize=ms, lw = lw, label=label+', e', color=color) + ax.plot(self.profiles['rho(-)'], self.profiles['ti(keV)'][:,0], '--*', markersize=ms, lw = lw, label=label+', i', color=color) + + ax.set_xlabel("$\\rho_N$") + ax.set_ylabel("$T$ (keV)") + #ax.set_ylim(bottom = 0) + ax.set_xlim(0,1) + ax.legend(prop={'size':8}) + GRAPHICStools.addDenseAxis(ax) + ax.set_title("Temperatures") + + # ne profiles + ax = axs[2] + + ax.plot(self.profiles['rho(-)'], self.profiles['ne(10^19/m^3)']*1E-1, '-o', markersize=ms, lw = lw, label=label, color=color) + + ax.set_xlabel("$\\rho_N$") + ax.set_ylabel("$n_e$ ($10^{20}m^{-3}$)") + #ax.set_ylim(bottom = 0) + ax.set_xlim(0,1) + ax.legend(prop={'size':8}) + GRAPHICStools.addDenseAxis(ax) + ax.set_title("Electron Density") + + # ---------------------------------- + # Pressure + # ---------------------------------- + + ax = axs[3] + + ax.plot(self.profiles['rho(-)'], self.derived['ptot_manual'], '-o', markersize=ms, lw = lw, label=label, color=color) + + ax.set_xlabel("$\\rho_N$") + ax.set_ylabel("$p_{kin}$ (MPa)") + #ax.set_ylim(bottom = 0) + ax.set_xlim(0,1) + ax.legend(prop={'size':8}) + GRAPHICStools.addDenseAxis(ax) + ax.set_title("Total Pressure") + + # ---------------------------------- + # Current + # ---------------------------------- + + # q-profile + ax = axs[4] + + ax.plot(self.profiles['rho(-)'], self.profiles['q(-)'], '-o', markersize=ms, lw = lw, label=label, color=color) + + ax.set_xlabel("$\\rho_N$") + ax.set_ylabel("$q$") + #ax.set_ylim(bottom = 0) + ax.set_xlim(0,1) + ax.legend(prop={'size':8}) + GRAPHICStools.addDenseAxis(ax) + ax.set_title("Safety Factor") + + # ---------------------------------- + # Powers + # ---------------------------------- + + # RF + ax = axs[5] + + ax.plot(self.profiles['rho(-)'], self.profiles['qrfe(MW/m^3)'], '-o', markersize=ms, lw = lw, label=label+', e', color=color) + ax.plot(self.profiles['rho(-)'], self.profiles['qrfi(MW/m^3)'], '--*', markersize=ms, lw = lw, label=label+', i', color=color) + + ax.set_xlabel("$\\rho_N$") + ax.set_ylabel("$P_{ich}$ (MW/m$^3$)") + #ax.set_ylim(bottom = 0) + ax.set_xlim(0,1) + ax.legend(prop={'size':8}) + GRAPHICStools.addDenseAxis(ax) + ax.set_title("ICH Power Deposition") + + # Ohmic + ax = axs[6] + + ax.plot(self.profiles['rho(-)'], self.profiles['qohme(MW/m^3)'], '-o', markersize=ms, lw = lw, label=label, color=color) + + ax.set_xlabel("$\\rho_N$") + ax.set_ylabel("$P_{oh}$ (MW/m$^3$)") + #ax.set_ylim(bottom = 0) + ax.set_xlim(0,1) + ax.legend(prop={'size':8}) + GRAPHICStools.addDenseAxis(ax) + ax.set_title("Ohmic Power Deposition") + + # ---------------------------------- + # Heat fluxes + # ---------------------------------- + + ax = axs[7] + + ax.plot(self.profiles['rho(-)'], self.derived['qe_MWm2'], '-o', markersize=ms, lw = lw, label=label+', e', color=color) + ax.plot(self.profiles['rho(-)'], self.derived['qi_MWm2'], '--*', markersize=ms, lw = lw, label=label+', i', color=color) + + ax.set_xlabel("$\\rho_N$") + ax.set_ylabel("$Q$ ($MW/m^2$)") + #ax.set_ylim(bottom = 0) + ax.set_xlim(0,1) + ax.legend(prop={'size':8}) + GRAPHICStools.addDenseAxis(ax) + ax.set_title("Energy Fluxes") + + # ---------------------------------- + # Dynamic targets + # ---------------------------------- + + ax = axs[8] + + ax.plot(self.profiles['rho(-)'], self.derived['qrad'], '-o', markersize=ms, lw = lw, label=label+', rad', color=color) + ax.plot(self.profiles['rho(-)'], self.profiles['qei(MW/m^3)'], '--*', markersize=ms, lw = lw, label=label+', exc', color=color) + if 'qfuse(MW/m^3)' in self.profiles: + ax.plot(self.profiles['rho(-)'], self.profiles['qfuse(MW/m^3)']+self.profiles['qfusi(MW/m^3)'], '-.s', markersize=ms, lw = lw, label=label+', fus', color=color) + + ax.set_xlabel("$\\rho_N$") + ax.set_ylabel("$Q$ ($MW/m^2$)") + #ax.set_ylim(bottom = 0) + ax.set_xlim(0,1) + ax.legend(prop={'size':8}) + GRAPHICStools.addDenseAxis(ax) + ax.set_title("Dynamic Targets") + + + def csv(self, file="input.gacode.xlsx"): + dictExcel = IOtools.OrderedDict() + + for ikey in self.profiles: + print(ikey) + if len(self.profiles[ikey].shape) == 1: + dictExcel[ikey] = self.profiles[ikey] + else: + dictExcel[ikey] = self.profiles[ikey][:, 0] + + IOtools.writeExcel_fromDict(dictExcel, file, fromRow=1) + + def parabolizePlasma(self): + _, T = PLASMAtools.parabolicProfile( + Tbar=self.derived["Te_vol"], + nu=self.derived["Te_peaking"], + rho=self.profiles["rho(-)"], + Tedge=self.profiles["te(keV)"][-1], + ) + _, Ti = PLASMAtools.parabolicProfile( + Tbar=self.derived["Ti_vol"], + nu=self.derived["Ti_peaking"], + rho=self.profiles["rho(-)"], + Tedge=self.profiles["ti(keV)"][-1, 0], + ) + _, n = PLASMAtools.parabolicProfile( + Tbar=self.derived["ne_vol20"] * 1e1, + nu=self.derived["ne_peaking"], + rho=self.profiles["rho(-)"], + Tedge=self.profiles["ne(10^19/m^3)"][-1], + ) + + self.profiles["te(keV)"] = T + + self.profiles["ti(keV)"][:, 0] = Ti + self.makeAllThermalIonsHaveSameTemp(refIon=0) + + factor_n = n / self.profiles["ne(10^19/m^3)"] + self.profiles["ne(10^19/m^3)"] = n + self.scaleAllThermalDensities(scaleFactor=factor_n) + + self.deriveQuantities() + + + def changeRFpower(self, PrfMW=25.0): + """ + keeps same partition + """ + print(f"- Changing the RF power from {self.derived['qRF_MWmiller'][-1]:.1f} MW to {PrfMW:.1f} MW",typeMsg="i",) + + if self.derived["qRF_MWmiller"][-1] == 0.0: + raise Exception("No RF power in the input.gacode, cannot modify the RF power") + + for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: + self.profiles[i] = self.profiles[i] * PrfMW / self.derived["qRF_MWmiller"][-1] + + self.deriveQuantities() + + def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0.2): + + ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) + + self.profiles["te(keV)"] = self.profiles["te(keV)"] * TkeV / self.profiles["te(keV)"][ix] + + print(f"- Producing {typeEdge} boundary condition @ rho = {rho}, T = {TkeV} keV",typeMsg="i",) + + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "therm": + self.profiles["ti(keV)"][:, sp] = self.profiles["ti(keV)"][:, sp] * TkeV / self.profiles["ti(keV)"][ix, sp] + + if typeEdge == "linear": + self.profiles["te(keV)"][ix:] = np.linspace(TkeV, Tesep, len(self.profiles["rho(-)"][ix:])) + + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "therm": + self.profiles["ti(keV)"][ix:, sp] = np.linspace(TkeV, Tisep, len(self.profiles["rho(-)"][ix:])) + + elif typeEdge == "same": + pass + else: + raise Exception("no edge") + + + def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5, isn20_edge=True): + ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) + + # Determine the factor to scale the density (either average or at rho) + if not isn20_edge: + print(f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") + factor = n20 / self.derived["ne_vol20"] + else: + print(f"- Changing the density at rho={rho} from {self.profiles['ne(10^19/m^3)'][ix]*1E-1:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") + factor = n20 / (self.profiles["ne(10^19/m^3)"][ix]*1E-1) + # ------------------------------------------------------------------ + + # Scale the density profiles + for i in ["ne(10^19/m^3)", "ni(10^19/m^3)"]: + self.profiles[i] = self.profiles[i] * factor + + # Apply the edge condition + if typeEdge == "linear": + factor_x = np.linspace(self.profiles["ne(10^19/m^3)"][ix],nedge20 * 1e1,len(self.profiles["rho(-)"][ix:]),)/ self.profiles["ne(10^19/m^3)"][ix:] + + self.profiles["ne(10^19/m^3)"][ix:] = self.profiles["ne(10^19/m^3)"][ix:] * factor_x + + for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): + self.profiles["ni(10^19/m^3)"][ix:, i] = self.profiles["ni(10^19/m^3)"][ix:, i] * factor_x + + elif typeEdge == "same": + pass + else: + raise Exception("no edge") + + def addSawtoothEffectOnOhmic(self, PohTot, mixRadius=None, plotYN=False): + """ + This will implement a flat profile inside the mixRadius to reduce the ohmic power by certain amount + """ + + if mixRadius is None: + mixRadius = self.profiles["rho(-)"][np.where(self.profiles["q(-)"] > 1)][0] + + print(f"\t- Original Ohmic power: {self.derived['qOhm_MWmiller'][-1]:.2f}MW") + Ohmic_old = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) + + dvol = self.derived["volp_miller"] * np.append( + [0], np.diff(self.profiles["rmin(m)"]) + ) + + print( + f"\t- Will implement sawtooth ohmic power correction inside rho={mixRadius}" + ) + Psaw = CDFtools.profilePower( + self.profiles["rho(-)"], + dvol, + PohTot - self.derived["qOhm_MWmiller"][-1], + mixRadius, + ) + self.profiles["qohme(MW/m^3)"] += Psaw + self.deriveQuantities() + + print(f"\t- New Ohmic power: {self.derived['qOhm_MWmiller'][-1]:.2f}MW") + Ohmic_new = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) + + if plotYN: + fig, ax = plt.subplots() + ax.plot(self.profiles["rho(-)"], Ohmic_old, "r", lw=2) + ax.plot(self.profiles["rho(-)"], Ohmic_new, "g", lw=2) + plt.show() + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Code conversions + # --------------------------------------------------------------------------------------------------------------------------------------- + + def to_tglf(self, rhos=[0.5], TGLFsettings=1): + + # <> Function to interpolate a curve <> + from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function + + inputsTGLF = {} + for rho in rhos: + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Define interpolator at this rho + # --------------------------------------------------------------------------------------------------------------------------------------- + + def interpolator(y): + return interpolation_function(rho, self.profiles['rho(-)'],y).item() + + TGLFinput, TGLFoptions, label = GACODEdefaults.addTGLFcontrol(TGLFsettings) + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Controls come from options + # --------------------------------------------------------------------------------------------------------------------------------------- + + controls = TGLFoptions + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Species come from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + #mass_ref = self.derived["mi_ref"] + # input.gacode uses the deuterium mass as reference already (https://github.com/gafusion/gacode/issues/398), so this should be 2.0 + mass_ref = 2.0 + + mass_e = 0.000272445 * mass_ref + + species = { + 1: { + 'ZS': -1.0, + 'MASS': mass_e/mass_ref, + 'RLNS': interpolator(self.derived['aLne']), + 'RLTS': interpolator(self.derived['aLTe']), + 'TAUS': 1.0, + 'AS': 1.0, + 'VPAR': interpolator(self.derived['vpar']), + 'VPAR_SHEAR': interpolator(self.derived['vpar_shear']), + 'VNS_SHEAR': 0.0, + 'VTS_SHEAR': 0.0}, + } + + for i in range(len(self.Species)): + species[i+2] = { + 'ZS': self.Species[i]['Z'], + 'MASS': self.Species[i]['A']/mass_ref, + 'RLNS': interpolator(self.derived['aLni'][:,i]), + 'RLTS': interpolator(self.derived['aLTi'][:,0] if self.Species[i]['S'] == 'therm' else self.derived["aLTi"][:,i]), + 'TAUS': interpolator(self.derived["tite_all"][:,i]), + 'AS': interpolator(self.derived['fi'][:,i]), + 'VPAR': interpolator(self.derived['vpar']), + 'VPAR_SHEAR': interpolator(self.derived['vpar_shear']), + 'VNS_SHEAR': 0.0, + 'VTS_SHEAR': 0.0 + } + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Plasma comes from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + plasma = { + 'NS': len(species)+1, + 'SIGN_BT': -1.0, + 'SIGN_IT': -1.0, + 'VEXB': 0.0, + 'VEXB_SHEAR': interpolator(self.derived['vexb_shear']), + 'BETAE': interpolator(self.derived['betae']), + 'XNUE': interpolator(self.derived['xnue']), + 'ZEFF': interpolator(self.derived['Zeff']), + 'DEBYE': interpolator(self.derived['debye']), + } + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Geometry comes from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + parameters = { + 'RMIN_LOC': self.derived['roa'], + 'RMAJ_LOC': self.derived['Rmajoa'], + 'ZMAJ_LOC': self.derived["Zmagoa"], + 'DRMINDX_LOC': self.derived['drmin/dr'], + 'DRMAJDX_LOC': self.derived['dRmaj/dr'], + 'DZMAJDX_LOC': self.derived['dZmaj/dr'], + 'Q_LOC': self.profiles["q(-)"], + 'KAPPA_LOC': self.profiles["kappa(-)"], + 'S_KAPPA_LOC': self.derived['s_kappa'], + 'DELTA_LOC': self.profiles["delta(-)"], + 'S_DELTA_LOC': self.derived['s_delta'], + 'ZETA_LOC': self.profiles["zeta(-)"], + 'S_ZETA_LOC': self.derived['s_zeta'], + 'P_PRIME_LOC': self.derived['pprime'], + 'Q_PRIME_LOC': self.derived['s_q'], + } + + geom = {} + for k in parameters: + par = torch.nan_to_num(torch.from_numpy(parameters[k]) if type(parameters[k]) is np.ndarray else parameters[k], nan=0.0, posinf=1E10, neginf=-1E10) + geom[k] = interpolator(par) + + geom['BETA_LOC'] = 0.0 + geom['KX0_LOC'] = 0.0 + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Merging + # --------------------------------------------------------------------------------------------------------------------------------------- + + input_dict = {**controls, **plasma, **geom} + + for i in range(len(species)): + for k in species[i+1]: + input_dict[f'{k}_{i+1}'] = species[i+1][k] + + inputsTGLF[rho] = input_dict + + return inputsTGLF + + def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', times = [0.0,1.0], Vsurf = 0.0): + + print("\t- Converting to TRANSP") + folder = IOtools.expandPath(folder) + folder.mkdir(parents=True, exist_ok=True) + + transp = TRANSPhelpers.transp_run(folder, shot, runid) + for time in times: + transp.populate_time.from_profiles(time,self, Vsurf = Vsurf) + + transp.write_ufiles() + + return transp + + def to_eped(self, ped_rho = 0.95): + + neped_19 = np.interp(ped_rho, self.profiles['rho(-)'], self.profiles['ne(10^19/m^3)']) + + eped_evaluation = { + 'Ip': np.abs(self.profiles['current(MA)'][0]), + 'Bt': np.abs(self.profiles['bcentr(T)'][0]), + 'R': np.abs(self.profiles['rcentr(m)'][0]), + 'a': np.abs(self.derived['a']), + 'kappa995': np.abs(self.derived['kappa995']), + 'delta995': np.abs(self.derived['delta995']), + 'neped': np.abs(neped_19), + 'betan': np.abs(self.derived['BetaN_engineering']), + 'zeff': np.abs(self.derived['Zeff_vol']), + 'tesep': np.abs(self.profiles['te(keV)'][-1])*1E3, + 'nesep_ratio': np.abs(self.profiles['ne(10^19/m^3)'][-1] / neped_19), + } + + return eped_evaluation + +class DataTable: + def __init__(self, variables=None): + + if variables is not None: + self.variables = variables + else: + + # Default for confinement mode access studies (JWH 03/2024) + self.variables = { + "Rgeo": ["rcentr(m)", "pos_0", "profiles", ".2f", 1, "m"], + "ageo": ["a", None, "derived", ".2f", 1, "m"], + "volume": ["volume", None, "derived", ".2f", 1, "m"], + "kappa @psi=0.95": ["kappa(-)", "psi_0.95", "profiles", ".2f", 1, None], + "delta @psi=0.95": ["delta(-)", "psi_0.95", "profiles", ".2f", 1, None], + "Bt": ["bcentr(T)", "pos_0", "profiles", ".1f", 1, "T"], + "Ip": ["current(MA)", "pos_0", "profiles", ".1f", 1, "MA"], + "Pin": ["qIn", None, "derived", ".1f", 1, "MW"], + "Te @rho=0.9": ["te(keV)", "rho_0.90", "profiles", ".2f", 1, "keV"], + "Ti/Te @rho=0.9": ["tite", "rho_0.90", "derived", ".2f", 1, None], + "ne @rho=0.9": [ + "ne(10^19/m^3)", + "rho_0.90", + "profiles", + ".2f", + 0.1, + "E20m-3", + ], + "ptot @rho=0.9": [ + "ptot_manual", + "rho_0.90", + "derived", + ".1f", + 1e3, + "kPa", + ], + "Zeff": ["Zeff_vol", None, "derived", ".1f", 1, None], + "fDT": ["fmain", None, "derived", ".2f", 1, None], + "H89p": ["H89", None, "derived", ".2f", 1, None], + "H98y2": ["H98", None, "derived", ".2f", 1, None], + "ne (vol avg)": ["ne_vol20", None, "derived", ".2f", 1, "E20m-3"], + "Ptop": ["ptop", None, "derived", ".1f", 1, "Pa"], + "fG": ["fG", None, "derived", ".2f", 1, None], + "Pfus": ["Pfus", None, "derived", ".1f", 1, "MW"], + "Prad": ["Prad", None, "derived", ".1f", 1, "MW"], + "Q": ["Q", None, "derived", ".2f", 1, None], + "Pnet @rho=0.9": ["qTr", "rho_0.90", "derived", ".1f", 1, "MW"], + "Qi/Qe @rho=0.9": ["QiQe", "rho_0.90", "derived", ".2f", 1, None], + } + + self.data = [] + + def export_to_csv(self, filename, title=None): + + title_data = [""] + for key in self.variables: + if self.variables[key][5] is None: + title_data.append(f"{key}") + else: + title_data.append(f"{key} ({self.variables[key][5]})") + + # Open a file with the given filename in write mode + with open(filename, mode="w", newline="") as file: + writer = csv.writer(file) + + # Write the title row first if it is provided + if title: + writer.writerow([title] + [""] * (len(self.data[0]) - 1)) + + writer.writerow(title_data) + + # Write each row in self.data to the CSV file + for row in self.data: + writer.writerow(row) + +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: + from mitim_tools.misc_tools.GUItools import FigureNotebook + + fn = FigureNotebook("Profiles", geometry="1800x900") + figProf_1, figProf_2, figProf_3, figProf_4, figFlows, figProf_6, fig7 = add_figures(fn) + + grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) + axsProf_1 = [ + figProf_1.add_subplot(grid[0, 0]), + figProf_1.add_subplot(grid[1, 0]), + figProf_1.add_subplot(grid[2, 0]), + figProf_1.add_subplot(grid[0, 1]), + figProf_1.add_subplot(grid[1, 1]), + figProf_1.add_subplot(grid[2, 1]), + figProf_1.add_subplot(grid[0, 2]), + figProf_1.add_subplot(grid[1, 2]), + figProf_1.add_subplot(grid[2, 2]), + ] + + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + axsProf_2 = [ + figProf_2.add_subplot(grid[0, 0]), + figProf_2.add_subplot(grid[0, 1]), + figProf_2.add_subplot(grid[1, 0]), + figProf_2.add_subplot(grid[1, 1]), + figProf_2.add_subplot(grid[0, 2]), + figProf_2.add_subplot(grid[1, 2]), + ] + grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.3) + ax00c = figProf_3.add_subplot(grid[0, 0]) + axsProf_3 = [ + ax00c, + figProf_3.add_subplot(grid[1, 0], sharex=ax00c), + figProf_3.add_subplot(grid[2, 0]), + figProf_3.add_subplot(grid[0, 1]), + figProf_3.add_subplot(grid[1, 1]), + figProf_3.add_subplot(grid[2, 1]), + figProf_3.add_subplot(grid[0, 2]), + figProf_3.add_subplot(grid[1, 2]), + figProf_3.add_subplot(grid[2, 2]), + figProf_3.add_subplot(grid[0, 3]), + figProf_3.add_subplot(grid[1, 3]), + figProf_3.add_subplot(grid[2, 3]), + ] + + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + axsProf_4 = [ + figProf_4.add_subplot(grid[0, 0]), + figProf_4.add_subplot(grid[1, 0]), + figProf_4.add_subplot(grid[0, 1]), + figProf_4.add_subplot(grid[1, 1]), + figProf_4.add_subplot(grid[0, 2]), + figProf_4.add_subplot(grid[1, 2]), + ] + + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + axsFlows = [ + figFlows.add_subplot(grid[0, 0]), + figFlows.add_subplot(grid[1, 0]), + figFlows.add_subplot(grid[0, 1]), + figFlows.add_subplot(grid[0, 2]), + figFlows.add_subplot(grid[1, 1]), + figFlows.add_subplot(grid[1, 2]), + ] + + grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) + axsProf_6 = [ + figProf_6.add_subplot(grid[0, 0]), + figProf_6.add_subplot(grid[:, 1]), + figProf_6.add_subplot(grid[0, 2]), + figProf_6.add_subplot(grid[1, 0]), + figProf_6.add_subplot(grid[1, 2]), + figProf_6.add_subplot(grid[0, 3]), + figProf_6.add_subplot(grid[1, 3]), + ] + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + axsImps = [ + fig7.add_subplot(grid[0, 0]), + fig7.add_subplot(grid[0, 1]), + fig7.add_subplot(grid[0, 2]), + fig7.add_subplot(grid[1, 0]), + fig7.add_subplot(grid[1, 1]), + fig7.add_subplot(grid[1, 2]), + ] + + ls = GRAPHICStools.listLS() + colors = GRAPHICStools.listColors() + for i, profiles in enumerate(profiles_list): + if extralabs is None: + extralab = f"#{i}, " + else: + extralab = f"{extralabs[i]}, " + profiles.plot( + axs1=axsProf_1, + axs2=axsProf_2, + axs3=axsProf_3, + axs4=axsProf_4, + axsFlows=axsFlows, + axs6=axsProf_6, + axsImps=axsImps, + color=colors[i], + legYN=True, + extralab=extralab, + lsFlows=ls[i], + legFlows=i == 0, + showtexts=False, + lastRhoGradients=lastRhoGradients, + ) + + return fn + + +def readTGYRO_profile_extra(file, varLabel="B_unit (T)"): + with open(file) as f: + aux = f.readlines() + + lenn = int(aux[36].split()[-1]) + + i = 38 + allVec = [] + while i < len(aux): + vec = np.array([float(j) for j in aux[i : i + lenn]]) + i += lenn + allVec.append(vec) + allVec = np.array(allVec) + + dictL = OrderedDict() + for line in aux[2:35]: + lab = line.split("(:)")[-1].split("\n")[0] + try: + dictL[lab] = int(line.split()[1]) + except: + dictL[lab] = [int(j) for j in line.split()[1].split("-")] + + for i in dictL: + if i.strip(" ") == varLabel: + val = allVec[dictL[i] - 1] + break + + return val + + +def aLT(r, p): + return ( + r[-1] + * CALCtools.produceGradient( + torch.from_numpy(r).to(torch.double), torch.from_numpy(p).to(torch.double) + ) + .cpu() + .cpu().numpy() + ) + + +def grad(r, p): + return MATHtools.deriv(torch.from_numpy(r), torch.from_numpy(p), array=False) + + +def ionName(Z, A): + # Based on Z + if Z == 2: + return "He" + elif Z == 9: + return "F" + elif Z == 6: + return "C" + elif Z == 11: + return "Na" + elif Z == 30: + return "Zn" + elif Z == 31: + return "Ga" + + # # Based on Mass (this is the correct way, since the radiation needs to be calculated with the full element) + # if A in [3,4]: return 'He' + # elif A == 18: return 'F' + # elif A == 12: return 'C' + # elif A == 22: return 'Na' + # elif A == 60: return 'Zn' + # elif A == 69: return 'Ga' + + +def gradientsMerger(p0, p_true, roa=0.46, blending=0.1): + p = copy.deepcopy(p0) + + aLTe_true = np.interp( + p.derived["roa"], p_true.derived["roa"], p_true.derived["aLTe"] + ) + aLTi_true = np.interp( + p.derived["roa"], p_true.derived["roa"], p_true.derived["aLTi"][:, 0] + ) + aLne_true = np.interp( + p.derived["roa"], p_true.derived["roa"], p_true.derived["aLne"] + ) + + ix1 = np.argmin(np.abs(p.derived["roa"] - roa + blending)) + ix2 = np.argmin(np.abs(p.derived["roa"] - roa)) + + aLT0 = aLTe_true[: ix1 + 1] + aLT2 = p.derived["aLTe"][ix2:] + aLT1 = np.interp( + p.derived["roa"][ix1 : ix2 + 1], + [p.derived["roa"][ix1], p.derived["roa"][ix2]], + [aLT0[-1], aLT2[0]], + )[1:-1] + + aLTe = np.append(np.append(aLT0, aLT1), aLT2) + Te = ( + CALCtools.integrateGradient( + torch.from_numpy(p.derived["roa"]).unsqueeze(0), + torch.Tensor(aLTe).unsqueeze(0), + p.profiles["te(keV)"][-1], + ) + .cpu() + .cpu().numpy()[0] + ) + + aLT0 = aLTi_true[: ix1 + 1] + aLT2 = p.derived["aLTi"][ix2:, 0] + aLT1 = np.interp( + p.derived["roa"][ix1 : ix2 + 1], + [p.derived["roa"][ix1], p.derived["roa"][ix2]], + [aLT0[-1], aLT2[0]], + )[1:-1] + + aLTi = np.append(np.append(aLT0, aLT1), aLT2) + Ti = ( + CALCtools.integrateGradient( + torch.from_numpy(p.derived["roa"]).unsqueeze(0), + torch.Tensor(aLTi).unsqueeze(0), + p.profiles["ti(keV)"][-1, 0], + ) + .cpu() + .cpu().numpy()[0] + ) + + aLT0 = aLne_true[: ix1 + 1] + aLT2 = p.derived["aLne"][ix2:] + aLT1 = np.interp( + p.derived["roa"][ix1 : ix2 + 1], + [p.derived["roa"][ix1], p.derived["roa"][ix2]], + [aLT0[-1], aLT2[0]], + )[1:-1] + + aLne = np.append(np.append(aLT0, aLT1), aLT2) + ne = ( + CALCtools.integrateGradient( + torch.from_numpy(p.derived["roa"]).unsqueeze(0), + torch.Tensor(aLne).unsqueeze(0), + p.profiles["ne(10^19/m^3)"][-1], + ) + .cpu() + .cpu().numpy()[0] + ) + + p.profiles["te(keV)"] = Te + p.profiles["ti(keV)"][:, 0] = Ti + p.profiles["ne(10^19/m^3)"] = ne + + p.deriveQuantities() + + return p + +def add_figures(fn, fnlab='', fnlab_pre='', tab_color=None): + + figProf_1 = fn.add_figure(label= fnlab_pre + "Profiles" + fnlab, tab_color=tab_color) + figProf_2 = fn.add_figure(label= fnlab_pre + "Powers" + fnlab, tab_color=tab_color) + figProf_3 = fn.add_figure(label= fnlab_pre + "Geometry" + fnlab, tab_color=tab_color) + figProf_4 = fn.add_figure(label= fnlab_pre + "Gradients" + fnlab, tab_color=tab_color) + figFlows = fn.add_figure(label= fnlab_pre + "Flows" + fnlab, tab_color=tab_color) + figProf_6 = fn.add_figure(label= fnlab_pre + "Other" + fnlab, tab_color=tab_color) + fig7 = fn.add_figure(label= fnlab_pre + "Impurities" + fnlab, tab_color=tab_color) + figs = [figProf_1, figProf_2, figProf_3, figProf_4, figFlows, figProf_6, fig7] + + return figs + +def impurity_location(profiles, impurity_of_interest): + + position_of_impurity = None + for i in range(len(profiles.Species)): + if profiles.Species[i]["N"] == impurity_of_interest: + if position_of_impurity is not None: + raise ValueError(f"[MITIM] Species {impurity_of_interest} found at positions {position_of_impurity} and {i}") + position_of_impurity = i + if position_of_impurity is None: + raise ValueError(f"[MITIM] Species {impurity_of_interest} not found in profiles") + + return position_of_impurity \ No newline at end of file diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index fcb1745d..a5500f66 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -15,7 +15,7 @@ GRAPHICStools, ) from mitim_tools.transp_tools import UFILEStools -from mitim_tools.gacode_tools import TGLFtools, TGYROtools, PROFILEStools +from mitim_tools.gacode_tools import TGLFtools, TGYROtools from mitim_tools.gacode_tools.utils import GACODEplotting, GACODErun, TRANSPgacode from mitim_tools.transp_tools.utils import ( FBMtools, @@ -15425,6 +15425,7 @@ def grid_interpolation_method_to_zero(x,y): for key in ['ne(10^19/m^3)', 'ni(10^19/m^3)', 'te(keV)', 'ti(keV)', 'rmin(m)']: profiles[key] = profiles[key].clip(min=minimum) + from mitim_tools.gacode_tools import PROFILEStools p = PROFILEStools.PROFILES_GACODE.scratch(profiles) return p diff --git a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py index 0b78795e..5c2a6058 100644 --- a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py +++ b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py @@ -1,10 +1,8 @@ -import os import shutil import numpy as np import matplotlib.pyplot as plt from mitim_tools.transp_tools import TRANSPtools, CDFtools, UFILEStools, NMLtools from mitim_tools.gs_tools import GEQtools -from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.misc_tools import IOtools, MATHtools, PLASMAtools, GRAPHICStools, FARMINGtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -762,6 +760,7 @@ def from_profiles(self, time, profiles_file, Vsurf = 0.0): self.time = time if isinstance(profiles_file, str): + from mitim_tools.gacode_tools import PROFILEStools self.p = PROFILEStools.PROFILES_GACODE(profiles_file) else: self.p = profiles_file From f9fb2e314133e4b1b955970421aabcf6885f8236 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 16:11:22 +0200 Subject: [PATCH 081/385] Moved geometric functions outside of PROFILES and into MITIMstate --- src/mitim_modules/maestro/MAESTROmain.py | 2 +- src/mitim_modules/maestro/utils/EPEDbeat.py | 2 +- .../maestro/utils/MAESTRObeat.py | 2 +- .../maestro/utils/PORTALSbeat.py | 2 +- src/mitim_modules/maestro/utils/TRANSPbeat.py | 12 +- .../portals/utils/PORTALSanalysis.py | 6 +- src/mitim_modules/powertorch/STATEtools.py | 2 +- .../powertorch/scripts/calculateTargets.py | 2 +- .../powertorch/scripts/compareWithTGYRO.py | 2 +- .../powertorch/utils/TRANSFORMtools.py | 22 +- .../powertorch/utils/TRANSPORTtools.py | 6 +- src/mitim_tools/astra_tools/ASTRAtools.py | 4 +- src/mitim_tools/gacode_tools/PROFILEStools.py | 99 ++++- src/mitim_tools/gacode_tools/TGLFtools.py | 6 +- src/mitim_tools/gacode_tools/TGYROtools.py | 48 +-- .../gacode_tools/utils/GEOMETRYtools.py | 12 +- .../gacode_tools/utils/NORMtools.py | 6 +- src/mitim_tools/gs_tools/GEQtools.py | 4 +- .../plasmastate_tools/MITIMstate.py | 397 +++++++----------- src/mitim_tools/popcon_tools/POPCONtools.py | 10 +- src/mitim_tools/popcon_tools/RAPIDStools.py | 4 +- .../transp_tools/utils/TRANSPhelpers.py | 2 +- 22 files changed, 307 insertions(+), 345 deletions(-) diff --git a/src/mitim_modules/maestro/MAESTROmain.py b/src/mitim_modules/maestro/MAESTROmain.py index 73fcd981..8d1f8a72 100644 --- a/src/mitim_modules/maestro/MAESTROmain.py +++ b/src/mitim_modules/maestro/MAESTROmain.py @@ -209,7 +209,7 @@ def prepare(self, *args, **kwargs): self.beat.initialize() # ----------------------------- - self.beat.profiles_current.deriveQuantities() + self.beat.profiles_current.derive_quantities() self.beat.prepare(*args, **kwargs) else: diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index 0b9303e8..4f6c1d8f 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -599,6 +599,6 @@ def eped_profiler(profiles, xp_old, rhotop, Tetop_keV, Titop_keV, netop_20, mini # Re-derive # --------------------------------- - profiles_output.deriveQuantities(rederiveGeometry=False) + profiles_output.derive_quantities(rederiveGeometry=False) return profiles_output \ No newline at end of file diff --git a/src/mitim_modules/maestro/utils/MAESTRObeat.py b/src/mitim_modules/maestro/utils/MAESTRObeat.py index 05024a7d..f32e2830 100644 --- a/src/mitim_modules/maestro/utils/MAESTRObeat.py +++ b/src/mitim_modules/maestro/utils/MAESTRObeat.py @@ -311,7 +311,7 @@ def __call__(self): self.initialize_instance.profiles_current.profiles['ni(10^19/m^3)'] = self.initialize_instance.profiles_current.profiles['ni(10^19/m^3)'] * (self.initialize_instance.profiles_current.profiles['ne(10^19/m^3)']/old_density)[:,np.newaxis] # Update derived - self.initialize_instance.profiles_current.deriveQuantities() + self.initialize_instance.profiles_current.derive_quantities() def _inform_save(self, **kwargs): pass diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index bcc9bb34..03cc8d8c 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -230,7 +230,7 @@ def merge_parameters(self): # -------------------------------------------------------------------------------------------- # Write to final input.gacode - self.profiles_output.deriveQuantities() + self.profiles_output.derive_quantities() self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') def grab_output(self, full = False): diff --git a/src/mitim_modules/maestro/utils/TRANSPbeat.py b/src/mitim_modules/maestro/utils/TRANSPbeat.py index f0d6d395..2635fc92 100644 --- a/src/mitim_modules/maestro/utils/TRANSPbeat.py +++ b/src/mitim_modules/maestro/utils/TRANSPbeat.py @@ -104,7 +104,7 @@ def prepare( self._additional_operations_add_initialization() # ICRF on - PichT_MW = self.profiles_current.derived['qRF_MWmiller'][-1] + PichT_MW = self.profiles_current.derived['qRF_MW'][-1] if freq_ICH is None: @@ -181,11 +181,11 @@ def _add_heating_profiles(self, force_auxiliary_heating_at_output = {'Pe': None, force_auxiliary_heating_at_output['Pe'] has the shaping function (takes rho) and the integrated value ''' - for key, pkey, ikey in zip(['Pe','Pi'], ['qrfe(MW/m^3)', 'qrfi(MW/m^3)'], ['qRFe_MWmiller', 'qRFi_MWmiller']): + for key, pkey, ikey in zip(['Pe','Pi'], ['qrfe(MW/m^3)', 'qrfi(MW/m^3)'], ['qRFe_MW', 'qRFi_MW']): if force_auxiliary_heating_at_output[key] is not None: self.profiles_output.profiles[pkey] = force_auxiliary_heating_at_output[key][0](self.profiles_output.profiles['rho(-)']) - self.profiles_output.deriveQuantities() + self.profiles_output.derive_quantities() self.profiles_output.profiles[pkey] = self.profiles_output.profiles[pkey] * force_auxiliary_heating_at_output[key][1]/self.profiles_output.derived[ikey][-1] def merge_parameters(self): @@ -230,13 +230,13 @@ def merge_parameters(self): self.profiles_output.profiles[key] = p_frozen.profiles[key] # Power scale - self.profiles_output.profiles['qrfe(MW/m^3)'] *= p_frozen.derived['qRF_MWmiller'][-1] / self.profiles_output.derived['qRF_MWmiller'][-1] - self.profiles_output.profiles['qrfi(MW/m^3)'] *= p_frozen.derived['qRF_MWmiller'][-1] / self.profiles_output.derived['qRF_MWmiller'][-1] + self.profiles_output.profiles['qrfe(MW/m^3)'] *= p_frozen.derived['qRF_MW'][-1] / self.profiles_output.derived['qRF_MW'][-1] + self.profiles_output.profiles['qrfi(MW/m^3)'] *= p_frozen.derived['qRF_MW'][-1] / self.profiles_output.derived['qRF_MW'][-1] # -------------------------------------------------------------------------------------------- # Write to final input.gacode - self.profiles_output.deriveQuantities() + self.profiles_output.derive_quantities() self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') def grab_output(self): diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 81192888..288b1d57 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -193,7 +193,7 @@ def prep_metrics(self, ilast=None): self.profiles_next_new.printInfo(label="NEXT") else: self.profiles_next_new = self.profiles_next - self.profiles_next_new.deriveQuantities() + self.profiles_next_new.derive_quantities() else: print("\t\t- Could not read next profile to evaluate (from folder)") @@ -217,7 +217,7 @@ def prep_metrics(self, ilast=None): print(f"\t\t- Processing evaluation {i}/{len(self.powerstates)-1}") if 'Q' not in power.profiles.derived: - power.profiles.deriveQuantities() + power.profiles.derive_quantities() self.evaluations.append(i) self.FusionGain.append(power.profiles.derived["Q"]) @@ -966,7 +966,7 @@ def __init__(self, folder): except FileNotFoundError: break - p.profiles.deriveQuantities() + p.profiles.derive_quantities() self.powerstates.append(p) self.fn = None diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 703d0500..446dcdcd 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -141,7 +141,7 @@ def _ensure_ne_before_nz(lst): # Use a copy because I'm deriving, it may be expensive and I don't want to carry that out outside of this class self.profiles = copy.deepcopy(profiles_object) if "derived" not in self.profiles.__dict__: - self.profiles.deriveQuantities() + self.profiles.derive_quantities() else: raise ValueError("[MITIM] The input profile object is not recognized, please use PROFILES_GACODE") diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 26dcf732..046f8cc3 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -110,7 +110,7 @@ def calculator( p.volume_integrate(p.plasma["qrad"]) * p.plasma["volp"] )[..., -1] - p.profiles.deriveQuantities() + p.profiles.derive_quantities() p.from_powerstate( write_input_gacode=folder / "input.gacode.new.powerstate", diff --git a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py index fe103393..564e784c 100644 --- a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py +++ b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py @@ -22,7 +22,7 @@ # TGYRO t = TGYROtools.TGYROoutput(folderTGYRO) -t.profiles.deriveQuantities() +t.profiles.derive_quantities() t.useFineGridTargets() diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 2b85ac23..aeea2304 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -69,7 +69,7 @@ def gacode_to_powerstate(self, rho_vec=None): ["rho", "rho(-)", None, True, False], ["roa", "roa", None, True, True], ["Rmajoa", "Rmajoa", None, True, True], - ["volp", "volp_miller", None, True, True], + ["volp", "volp_geo", None, True, True], ["rmin", "rmin(m)", None, True, False], ["te", "te(keV)", None, True, False], ["ti", "ti(keV)", 0, True, False], @@ -114,21 +114,21 @@ def gacode_to_powerstate(self, rho_vec=None): # ********************************************************************************************* quantitites = {} - quantitites["QeMWm2_fixedtargets"] = input_gacode.derived["qe_aux_MWmiller"] - quantitites["QiMWm2_fixedtargets"] = input_gacode.derived["qi_aux_MWmiller"] - quantitites["Ge_fixedtargets"] = input_gacode.derived["ge_10E20miller"] - quantitites["GZ_fixedtargets"] = input_gacode.derived["ge_10E20miller"] * 0.0 + quantitites["QeMWm2_fixedtargets"] = input_gacode.derived["qe_aux_MW"] + quantitites["QiMWm2_fixedtargets"] = input_gacode.derived["qi_aux_MW"] + quantitites["Ge_fixedtargets"] = input_gacode.derived["ge_10E20"] + quantitites["GZ_fixedtargets"] = input_gacode.derived["ge_10E20"] * 0.0 quantitites["MtJm2_fixedtargets"] = input_gacode.derived["mt_Jmiller"] if self.TargetOptions["ModelOptions"]["TypeTarget"] < 3: # Fusion and radiation fixed if 1,2 - quantitites["QeMWm2_fixedtargets"] += input_gacode.derived["qe_fus_MWmiller"] - input_gacode.derived["qrad_MWmiller"] - quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qi_fus_MWmiller"] + quantitites["QeMWm2_fixedtargets"] += input_gacode.derived["qe_fus_MW"] - input_gacode.derived["qrad_MW"] + quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qi_fus_MW"] if self.TargetOptions["ModelOptions"]["TypeTarget"] < 2: # Exchange fixed if 1 - quantitites["QeMWm2_fixedtargets"] -= input_gacode.derived["qe_exc_MWmiller"] - quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qe_exc_MWmiller"] + quantitites["QeMWm2_fixedtargets"] -= input_gacode.derived["qe_exc_MW"] + quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qe_exc_MW"] for key in quantitites: @@ -335,7 +335,7 @@ def powerstate_to_gacode( # ------------------------------------------------------------------------------------------ if rederive or recompute_ptot: - profiles.deriveQuantities(rederiveGeometry=False) + profiles.derive_quantities(rederiveGeometry=False) if recompute_ptot: profiles.selfconsistentPTOT() @@ -347,7 +347,7 @@ def powerstate_to_gacode( def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): - profiles.deriveQuantities(rederiveGeometry=False) + profiles.derive_quantities(rederiveGeometry=False) print("\t- Insering powers") diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 02f5c395..a9461805 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -71,7 +71,7 @@ def produce_profiles(self): # (e.g. for flux matching of analytical models) pass - def _produce_profiles(self,deriveQuantities=True): + def _produce_profiles(self,derive_quantities=True): self.applyCorrections = self.powerstate.TransportOptions["ModelOptions"].get("MODELparameters", {}).get("applyCorrections", {}) @@ -83,8 +83,8 @@ def _produce_profiles(self,deriveQuantities=True): self.powerstate.profiles = powerstate_detached.from_powerstate( write_input_gacode=self.file_profs, postprocess_input_gacode=self.applyCorrections, - rederive_profiles = deriveQuantities, # Derive quantities so that it's ready for analysis and plotting later - insert_highres_powers = deriveQuantities, # Insert powers so that Q, Pfus and all that it's consistent when read later + rederive_profiles = derive_quantities, # Derive quantities so that it's ready for analysis and plotting later + insert_highres_powers = derive_quantities, # Insert powers so that Q, Pfus and all that it's consistent when read later ) self.powerstate.profiles_transport = copy.deepcopy(self.powerstate.profiles) diff --git a/src/mitim_tools/astra_tools/ASTRAtools.py b/src/mitim_tools/astra_tools/ASTRAtools.py index e59b0d44..3795d92f 100644 --- a/src/mitim_tools/astra_tools/ASTRAtools.py +++ b/src/mitim_tools/astra_tools/ASTRAtools.py @@ -342,7 +342,7 @@ def convert_ASTRA_to_gacode_from_transp_output(c, p.profiles['rho(-)'][0] = 0.0 # rederive quantities - p.deriveQuantities() + p.derive_quantities() # Print output to check Q, Pfus, etc. p.printInfo() @@ -501,7 +501,7 @@ def _betan_initial_conditions(x, geometry_object, rhotop, Ttop_keV, netop_19, ep profiles.makeAllThermalIonsHaveSameTemp() profiles.profiles['ni(10^19/m^3)'][:,0] = profiles.profiles['ne(10^19/m^3)'] profiles.enforceQuasineutrality() - profiles.deriveQuantities() + profiles.derive_quantities() print("residual:", ((profiles.derived['BetaN_engineering']-betan_desired) / betan_desired)**2) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index d11d8f59..bf8fef43 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -2,6 +2,8 @@ import numpy as np from collections import OrderedDict from mitim_tools.plasmastate_tools.MITIMstate import mitim_state +from mitim_tools.gs_tools import GEQtools +from mitim_tools.gacode_tools.utils import GEOMETRYtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -15,6 +17,10 @@ class PROFILES_GACODE(mitim_state): and writes them in the way that MITIMstate class expects. ''' + # ------------------------------------------------------------------ + # Reading and interpreting + # ------------------------------------------------------------------ + def __init__(self, file, calculateDerived=True, mi_ref=None): super().__init__(type_file='input.gacode') @@ -23,21 +29,12 @@ def __init__(self, file, calculateDerived=True, mi_ref=None): Depending on resolution, derived can be expensive, so I mmay not do it every time """ - self.titles_singleNum = ["nexp", "nion", "shot", "name", "type", "time"] - self.titles_singleArr = [ - "masse", - "mass", - "ze", - "z", - "torfluxa(Wb/radian)", - "rcentr(m)", - "bcentr(T)", - "current(MA)", - ] - self.titles_single = self.titles_singleNum + self.titles_singleArr - self.file = file + self.titles_singleNum = ["nexp", "nion", "shot", "name", "type", "time"] + self.titles_singleArr = ["masse","mass","ze","z","torfluxa(Wb/radian)","rcentr(m)","bcentr(T)","current(MA)"] + self.titles_single = self.titles_singleNum + self.titles_singleArr + if self.file is not None: with open(self.file, "r") as f: self.lines = f.readlines() @@ -45,10 +42,10 @@ def __init__(self, file, calculateDerived=True, mi_ref=None): # Read file and store raw data self._read_header() self._read_profiles() - self._ensure_existence() - self.deriveQuantities(mi_ref=mi_ref, calculateDerived=calculateDerived) + # Derive + self.derive_quantities(mi_ref=mi_ref, calculateDerived=calculateDerived) def _read_header(self): for i in range(len(self.lines)): @@ -161,3 +158,75 @@ def _ensure_existence(self): if ikey not in self.profiles.keys(): self.profiles[ikey] = copy.deepcopy(self.profiles["rmin(m)"]) * 0.0 + # ------------------------------------------------------------------ + # Derivation (different from MITIMstate) + # ------------------------------------------------------------------ + + def derive_quantities(self, **kwargs): + + self._produce_shape_lists() + + super().derive_quantities(**kwargs) + + def _produce_shape_lists(self): + self.shape_cos = [ + self.profiles["shape_cos0(-)"], # tilt + self.profiles["shape_cos1(-)"], + self.profiles["shape_cos2(-)"], + self.profiles["shape_cos3(-)"], + self.profiles["shape_cos4(-)"], + self.profiles["shape_cos5(-)"], + self.profiles["shape_cos6(-)"], + ] + self.shape_sin = [ + None, + None, # s1 is arcsin(triangularity) + None, # s2 is minus squareness + self.profiles["shape_sin3(-)"], + self.profiles["shape_sin4(-)"], + self.profiles["shape_sin5(-)"], + self.profiles["shape_sin6(-)"], + ] + + # ------------------------------------------------------------------ + # Geometry + # ------------------------------------------------------------------ + + def derive_geometry(self, n_theta_geo=1001): + + self._produce_shape_lists() + + ( + self.derived["volp_geo"], + self.derived["surf_geo"], + self.derived["gradr_geo"], + self.derived["bp2_geo"], + self.derived["bt2_geo"], + self.derived["bt_geo"], + ) = GEOMETRYtools.calculateGeometricFactors(self,n_theta=n_theta_geo) + + # Calculate flux surfaces + cn = np.array(self.shape_cos).T + sn = copy.deepcopy(self.shape_sin) + sn[0] = self.profiles["rmaj(m)"]*0.0 + sn[1] = np.arcsin(self.profiles["delta(-)"]) + sn[2] = -self.profiles["zeta(-)"] + sn = np.array(sn).T + flux_surfaces = GEQtools.mitim_flux_surfaces() + flux_surfaces.reconstruct_from_mxh_moments( + self.profiles["rmaj(m)"], + self.profiles["rmin(m)"], + self.profiles["kappa(-)"], + self.profiles["zmag(m)"], + cn, + sn) + self.derived["R_surface"],self.derived["Z_surface"] = flux_surfaces.R, flux_surfaces.Z + # ----------------------------------------------- + + #cross-sectional area of each flux surface + self.derived["surfXS"] = GEOMETRYtools.xsec_area_RZ(self.derived["R_surface"],self.derived["Z_surface"]) + + self.derived["R_LF"] = self.derived["R_surface"].max(axis=1) # self.profiles['rmaj(m)'][0]+self.profiles['rmin(m)'] + + # For Synchrotron + self.derived["B_ref"] = np.abs(self.derived["B_unit"] * self.derived["bt_geo"]) diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 6a1dc3b0..6f7cb8c2 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -338,7 +338,7 @@ def prep( print("> Setting up normalizations") print("\t- Using mass of deuterium to unnormalize TGLF (not necesarily the first ion)",typeMsg="i") - self.tgyro.profiles.deriveQuantities(mi_ref=mi_D) + self.tgyro.profiles.derive_quantities(mi_ref=mi_D) self.NormalizationSets, cdf = NORMtools.normalizations( self.tgyro.profiles, @@ -398,7 +398,7 @@ def prep_direct_tglf( self.profiles = self.tgyro.profiles - self.profiles.deriveQuantities(mi_ref=mi_D) + self.profiles.derive_quantities(mi_ref=mi_D) self.profiles.correct(options={'recompute_ptot':recalculatePTOT,'removeFast':onlyThermal_TGYRO}) @@ -451,7 +451,7 @@ def prep_direct_tglf( "\t- Using mass of deuterium to normalize things (not necesarily the first ion)", typeMsg="w", ) - self.profiles.deriveQuantities(mi_ref=mi_D) + self.profiles.derive_quantities(mi_ref=mi_D) self.NormalizationSets, cdf = NORMtools.normalizations( self.profiles, diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index e14bba87..b70068de 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -3641,7 +3641,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles.derived["qe_fus_MWmiller"] + P = self.profiles.derived["qe_fus_MW"] ax.plot(roa, P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: @@ -3649,7 +3649,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles_final.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles_final.derived["qe_fus_MWmiller"] + P = self.profiles_final.derived["qe_fus_MW"] ax.plot(roa, P, "--", c="k", label="profiles_new (miller)", lw=1.0) GRAPHICStools.addDenseAxis(ax) @@ -3682,7 +3682,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles.derived["qe_aux_MWmiller"] + P = self.profiles.derived["qe_aux_MW"] ax.plot(roa, P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: @@ -3690,7 +3690,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles_final.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles_final.derived["qe_aux_MWmiller"] + P = self.profiles_final.derived["qe_aux_MW"] ax.plot(roa, P, "--", c="k", label="profiles_new (miller)", lw=1.0) GRAPHICStools.addDenseAxis(ax) @@ -3763,7 +3763,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = -1 * sign * self.profiles.derived["qe_rad_MWmiller"] + P = -1 * sign * self.profiles.derived["qe_rad_MW"] ax.plot(roa, P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: @@ -3771,7 +3771,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles_final.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles_final.derived["qe_rad_MWmiller"] + P = self.profiles_final.derived["qe_rad_MW"] ax.plot(roa, P, "--", c="k", label="profiles_new (miller)", lw=1.0) ax.legend(prop={"size": 6}) @@ -3835,14 +3835,14 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = -self.profiles.derived["qe_exc_MWmiller"] + P = -self.profiles.derived["qe_exc_MW"] ax.plot(roa, P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: roa = ( self.profiles_final.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = -self.profiles_final.derived["qe_exc_MWmiller"] + P = -self.profiles_final.derived["qe_exc_MW"] ax.plot(roa, P, "--", c="k", label="profiles_new (miller)", lw=1.0) ax.legend(prop={"size": 6}) @@ -3880,14 +3880,14 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles.derived["qe_MWmiller"] + P = self.profiles.derived["qe_MW"] ax.plot(roa, P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: roa = ( self.profiles_final.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles_final.derived["qe_MWmiller"] + P = self.profiles_final.derived["qe_MW"] ax.plot(roa, P, "--", c="k", label="profiles_new (miller)", lw=1.0) ax.legend(prop={"size": 6}) @@ -3921,14 +3921,14 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles.derived["qi_fus_MWmiller"] + P = self.profiles.derived["qi_fus_MW"] ax.plot(roa, P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: roa = ( self.profiles_final.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles_final.derived["qi_fus_MWmiller"] + P = self.profiles_final.derived["qi_fus_MW"] ax.plot(roa, P, "--", c="k", label="profiles_new (miller)", lw=1.0) GRAPHICStools.addDenseAxis(ax) @@ -3970,20 +3970,20 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles.derived["qi_aux_MWmiller"] + P = self.profiles.derived["qi_aux_MW"] ax.plot(roa, P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: roa = ( self.profiles_final.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles_final.derived["qi_aux_MWmiller"] + P = self.profiles_final.derived["qi_aux_MW"] ax.plot(roa, P, "--", c="k", label="profiles_new (miller)", lw=1.0) ax.plot( roa, - self.profiles_final.derived["qi_aux_MWmiller"] - + self.profiles_final.derived["qe_aux_MWmiller"], + self.profiles_final.derived["qi_aux_MW"] + + self.profiles_final.derived["qe_aux_MW"], "-.", c="y", lw=0.5, @@ -4009,8 +4009,8 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): / self.profiles.profiles["rmin(m)"][-1] ) P = ( - self.profiles.derived["qe_fus_MWmiller"] - + self.profiles.derived["qi_fus_MWmiller"] + self.profiles.derived["qe_fus_MW"] + + self.profiles.derived["qi_fus_MW"] ) ax.plot(roa, 5 * P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: @@ -4019,8 +4019,8 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): / self.profiles_final.profiles["rmin(m)"][-1] ) P = ( - self.profiles_final.derived["qe_fus_MWmiller"] - + self.profiles_final.derived["qi_fus_MWmiller"] + self.profiles_final.derived["qe_fus_MW"] + + self.profiles_final.derived["qi_fus_MW"] ) ax.plot(roa, 5 * P, "--", c="k", label="profiles_new (miller)", lw=1.0) @@ -4060,7 +4060,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles.derived["qe_exc_MWmiller"] + P = self.profiles.derived["qe_exc_MW"] ax.plot(roa, P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: @@ -4068,7 +4068,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles_final.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles_final.derived["qe_exc_MWmiller"] + P = self.profiles_final.derived["qe_exc_MW"] ax.plot(roa, P, "--", c="k", label="profiles_new (miller)", lw=1.0) GRAPHICStools.addDenseAxis(ax) @@ -4094,7 +4094,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles.derived["qi_MWmiller"] + P = self.profiles.derived["qi_MW"] ax.plot(roa, P, "--", c="green", label="profiles (miller)", lw=1.0) if self.profiles_final is not None: @@ -4102,7 +4102,7 @@ def plot(self, fn=None, label="", prelabel="", fn_color=None): self.profiles_final.profiles["rmin(m)"] / self.profiles.profiles["rmin(m)"][-1] ) - P = self.profiles_final.derived["qi_MWmiller"] + P = self.profiles_final.derived["qi_MW"] ax.plot(roa, P, "--", c="k", label="profiles_new (miller)", lw=1.0) GRAPHICStools.addDenseAxis(ax) diff --git a/src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py b/src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py index 1b09ef46..dbccba56 100644 --- a/src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py +++ b/src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py @@ -64,8 +64,8 @@ def calculateGeometricFactors(profiles, n_theta=1001): geo_fluxsurfave_grad_r, geo_fluxsurfave_bp2, geo_fluxsurfave_bt2, - geo_bt0, - ) = volp_surf_Miller_vectorized( + bt_geo0, + ) = volp_surf_geo_vectorized( R, r, delta, @@ -92,9 +92,9 @@ def calculateGeometricFactors(profiles, n_theta=1001): volp = geo_volume_prime * profiles.profiles["rmin(m)"][-1] ** 2 surf = geo_surf * profiles.profiles["rmin(m)"][-1] ** 2 - return volp, surf, geo_fluxsurfave_grad_r, geo_fluxsurfave_bp2, geo_fluxsurfave_bt2, geo_bt0 + return volp, surf, geo_fluxsurfave_grad_r, geo_fluxsurfave_bp2, geo_fluxsurfave_bt2, bt_geo0 -def volp_surf_Miller_vectorized( +def volp_surf_geo_vectorized( geo_rmaj_in, geo_rmin_in, geo_delta_in, @@ -356,7 +356,7 @@ def volp_surf_Miller_vectorized( z = (x0 - x1) / dx if i2 == n_theta: i2 -= 1 - geo_bt0 = geov_bt[i1] + (geov_bt[i2] - geov_bt[i1]) * z + bt_geo0 = geov_bt[i1] + (geov_bt[i2] - geov_bt[i1]) * z denom = 0 for i in range(n_theta - 1): @@ -383,7 +383,7 @@ def volp_surf_Miller_vectorized( + geov_bp ** 2 * geov_g_theta[i] / geov_b[i] / denom ) - return geo_volume_prime, geo_surf, geo_fluxsurfave_grad_r, geo_fluxsurfave__bp2, geo_fluxsurfave_bt2, geo_bt0 + return geo_volume_prime, geo_surf, geo_fluxsurfave_grad_r, geo_fluxsurfave__bp2, geo_fluxsurfave_bt2, bt_geo0 def xsec_area_RZ( R, diff --git a/src/mitim_tools/gacode_tools/utils/NORMtools.py b/src/mitim_tools/gacode_tools/utils/NORMtools.py index 42bbae99..b17c44c2 100644 --- a/src/mitim_tools/gacode_tools/utils/NORMtools.py +++ b/src/mitim_tools/gacode_tools/utils/NORMtools.py @@ -113,9 +113,9 @@ def normalizations_profiles(profiles): "ne_20": np.abs(profiles.profiles["ne(10^19/m^3)"]) * 1e-1, "Ti_keV": np.abs(profiles.profiles["ti(keV)"][:, 0]), "ni_20": np.abs(profiles.derived["ni_thrAll"]) * 1e-1, - "exp_Qe": profiles.derived["qe_MWmiller"] / profiles.derived["surfGACODE_miller"], # This is the same as qe_MWm2 - "exp_Qi": profiles.derived["qi_MWmiller"] / profiles.derived["surfGACODE_miller"], - "exp_Ge": profiles.derived["ge_10E20miller"] / profiles.derived["surfGACODE_miller"], + "exp_Qe": profiles.derived["qe_MW"] / profiles.derived["surfGACODE_geo"], # This is the same as qe_MWm2 + "exp_Qi": profiles.derived["qi_MW"] / profiles.derived["surfGACODE_geo"], + "exp_Ge": profiles.derived["ge_10E20"] / profiles.derived["surfGACODE_geo"], "mi_ref": profiles.derived["mi_ref"], } diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 6dd8f32c..e27774f7 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -420,14 +420,14 @@ def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH p = MITIMstate.mitim_state.scratch(profiles) - p.profiles["qrfe(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] * PichT/p.derived['qRF_MWmiller'][-1] /2 + p.profiles["qrfe(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] * PichT/p.derived['qRF_MW'][-1] /2 p.profiles["qrfi(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] # ------------------------------------------------------------------------------------------------------- # Ready to go # ------------------------------------------------------------------------------------------------------- - p.deriveQuantities() + p.derive_quantities() # ------------------------------------------------------------------------------------------------------- # Plotting diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index ec50e760..ab344587 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -6,9 +6,8 @@ from collections import OrderedDict from mitim_tools.misc_tools import GRAPHICStools, MATHtools, PLASMAtools, IOtools from mitim_modules.powertorch.utils import CALCtools -from mitim_tools.gs_tools import GEQtools from mitim_tools.gacode_tools import NEOtools -from mitim_tools.gacode_tools.utils import GACODEdefaults, GEOMETRYtools +from mitim_tools.gacode_tools.utils import GACODEdefaults from mitim_tools.transp_tools import CDFtools from mitim_tools.transp_tools.utils import TRANSPhelpers from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -24,11 +23,10 @@ def __init__(self, type_file = 'input.gacode'): self.type = type_file - def deriveQuantities(self, mi_ref=None, calculateDerived=True, n_theta_geo=1001, rederiveGeometry=True): + def derive_quantities(self, mi_ref=None, calculateDerived=True, rederiveGeometry=True): # ------------------------------------------------------------------------------------------------------------------- self.readSpecies() - self.produce_shape_lists() self.mi_first = self.Species[0]["A"] self.DTplasma() self.sumFast() @@ -46,24 +44,21 @@ def deriveQuantities(self, mi_ref=None, calculateDerived=True, n_theta_geo=1001, # Useful to have gradients in the basic ---------------------------------------------------------- self.derived["aLTe"] = aLT(self.profiles["rmin(m)"], self.profiles["te(keV)"]) - self.derived["aLne"] = aLT( - self.profiles["rmin(m)"], self.profiles["ne(10^19/m^3)"] - ) + self.derived["aLne"] = aLT(self.profiles["rmin(m)"], self.profiles["ne(10^19/m^3)"]) self.derived["aLTi"] = self.profiles["ti(keV)"] * 0.0 self.derived["aLni"] = [] for i in range(self.profiles["ti(keV)"].shape[1]): - self.derived["aLTi"][:, i] = aLT( - self.profiles["rmin(m)"], self.profiles["ti(keV)"][:, i] - ) - self.derived["aLni"].append( - aLT(self.profiles["rmin(m)"], self.profiles["ni(10^19/m^3)"][:, i]) - ) + self.derived["aLTi"][:, i] = aLT(self.profiles["rmin(m)"], self.profiles["ti(keV)"][:, i]) + self.derived["aLni"].append(aLT(self.profiles["rmin(m)"], self.profiles["ni(10^19/m^3)"][:, i])) self.derived["aLni"] = np.transpose(np.array(self.derived["aLni"])) # ------------------------------------------------------------------------------------------------ if calculateDerived: - self.deriveQuantities_full(rederiveGeometry=rederiveGeometry) + self.derive_quantities_full(rederiveGeometry=rederiveGeometry) + + def derive_geometry(self, **kwargs): + raise Exception('[MITIM] This method is not implemented in the base class. Please use a derived class that implements it.') # ------------------------------------------------------------------------------------- # Method to write a scratch file @@ -128,35 +123,14 @@ def calculate_Er( 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 ""}' - ) + 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] - self.deriveQuantities() + self.derive_quantities() if write_new_file is not None: self.writeCurrentStatus(file=write_new_file) - - def produce_shape_lists(self): - self.shape_cos = [ - self.profiles["shape_cos0(-)"], # tilt - self.profiles["shape_cos1(-)"], - self.profiles["shape_cos2(-)"], - self.profiles["shape_cos3(-)"], - self.profiles["shape_cos4(-)"], - self.profiles["shape_cos5(-)"], - self.profiles["shape_cos6(-)"], - ] - self.shape_sin = [ - None, - None, # s1 is arcsin(triangularity) - None, # s2 is minus squareness - self.profiles["shape_sin3(-)"], - self.profiles["shape_sin4(-)"], - self.profiles["shape_sin5(-)"], - self.profiles["shape_sin6(-)"], - ] + def readSpecies(self, maxSpecies=100): maxSpecies = int(self.profiles["nion"][0]) @@ -191,11 +165,9 @@ def sumFast(self): ) else: self.nThermal += self.profiles["ni(10^19/m^3)"][:, sp] - self.nZThermal += ( - self.profiles["ni(10^19/m^3)"][:, sp] * self.profiles["z"][sp] - ) + self.nZThermal += self.profiles["ni(10^19/m^3)"][:, sp] * self.profiles["z"][sp] - def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry=True): + def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): """ deriving geometry is expensive, so if I'm just updating profiles it may not be needed """ @@ -220,15 +192,8 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= self.derived["Rmajoa"] = self.profiles["rmaj(m)"] / self.derived["a"] self.derived["Zmagoa"] = self.profiles["zmag(m)"] / self.derived["a"] - self.derived["torflux"] = ( - float(self.profiles["torfluxa(Wb/radian)"][0]) - * 2 - * np.pi - * self.profiles["rho(-)"] ** 2 - ) # Wb - self.derived["B_unit"] = PLASMAtools.Bunit( - self.derived["torflux"], self.profiles["rmin(m)"] - ) + self.derived["torflux"] = float(self.profiles["torfluxa(Wb/radian)"][0])* 2* np.pi* self.profiles["rho(-)"] ** 2 # Wb + self.derived["B_unit"] = PLASMAtools.Bunit(self.derived["torflux"], self.profiles["rmin(m)"]) self.derived["psi_pol_n"] = ( self.profiles["polflux(Wb/radian)"] - self.profiles["polflux(Wb/radian)"][0] @@ -238,100 +203,42 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= ) self.derived["rho_pol"] = self.derived["psi_pol_n"] ** 0.5 - self.derived["q95"] = np.interp( - 0.95, self.derived["psi_pol_n"], self.profiles["q(-)"] - ) + self.derived["q95"] = np.interp(0.95, self.derived["psi_pol_n"], self.profiles["q(-)"]) self.derived["q0"] = self.profiles["q(-)"][0] if self.profiles["q(-)"].min() > 1.0: self.derived["rho_saw"] = np.nan else: - self.derived["rho_saw"] = np.interp( - 1.0, self.profiles["q(-)"], self.profiles["rho(-)"] - ) + self.derived["rho_saw"] = np.interp(1.0, self.profiles["q(-)"], self.profiles["rho(-)"]) # --------- Geometry (only if it doesn't exist or if I ask to recalculate) - if rederiveGeometry or ("volp_miller" not in self.derived): - - self.produce_shape_lists() - - ( - self.derived["volp_miller"], - self.derived["surf_miller"], - self.derived["gradr_miller"], - self.derived["bp2_miller"], - self.derived["bt2_miller"], - self.derived["geo_bt"], - ) = GEOMETRYtools.calculateGeometricFactors( - self, - n_theta=n_theta_geo, - ) - - # Calculate flux surfaces - cn = np.array(self.shape_cos).T - sn = copy.deepcopy(self.shape_sin) - sn[0] = self.profiles["rmaj(m)"]*0.0 - sn[1] = np.arcsin(self.profiles["delta(-)"]) - sn[2] = -self.profiles["zeta(-)"] - sn = np.array(sn).T - flux_surfaces = GEQtools.mitim_flux_surfaces() - flux_surfaces.reconstruct_from_mxh_moments( - self.profiles["rmaj(m)"], - self.profiles["rmin(m)"], - self.profiles["kappa(-)"], - self.profiles["zmag(m)"], - cn, - sn) - self.derived["R_surface"],self.derived["Z_surface"] = flux_surfaces.R, flux_surfaces.Z - # ----------------------------------------------- - - #cross-sectional area of each flux surface - self.derived["surfXS"] = GEOMETRYtools.xsec_area_RZ( - self.derived["R_surface"], - self.derived["Z_surface"] - ) - - self.derived["R_LF"] = self.derived["R_surface"].max( - axis=1 - ) # self.profiles['rmaj(m)'][0]+self.profiles['rmin(m)'] - - # For Synchrotron - self.derived["B_ref"] = np.abs( - self.derived["B_unit"] * self.derived["geo_bt"] - ) + if rederiveGeometry or ("volp_geo" not in self.derived): + self.derive_geometry() # -------------------------------------------------------------------------- # Reference mass # -------------------------------------------------------------------------- - # Forcing mass from this specific deriveQuantities call + # Forcing mass from this specific derive_quantities call if mi_ref is not None: self.derived["mi_ref"] = mi_ref - print(f'\t- Using mi_ref={self.derived["mi_ref"]} provided in this particular deriveQuantities method, subtituting initialization one',typeMsg='i') + print(f'\t- Using mi_ref={self.derived["mi_ref"]} provided in this particular derive_quantities method, subtituting initialization one',typeMsg='i') # --------------------------------------------------------------------------------------------------------------------- # --------- Important for scaling laws # --------------------------------------------------------------------------------------------------------------------- - self.derived["kappa95"] = np.interp( - 0.95, self.derived["psi_pol_n"], self.profiles["kappa(-)"] - ) + self.derived["kappa95"] = np.interp(0.95, self.derived["psi_pol_n"], self.profiles["kappa(-)"]) - self.derived["kappa995"] = np.interp( - 0.995, self.derived["psi_pol_n"], self.profiles["kappa(-)"] - ) + self.derived["kappa995"] = np.interp(0.995, self.derived["psi_pol_n"], self.profiles["kappa(-)"]) self.derived["kappa_a"] = self.derived["surfXS"][-1] / np.pi / self.derived["a"] ** 2 - self.derived["delta95"] = np.interp( - 0.95, self.derived["psi_pol_n"], self.profiles["delta(-)"] - ) + self.derived["delta95"] = np.interp(0.95, self.derived["psi_pol_n"], self.profiles["delta(-)"]) - self.derived["delta995"] = np.interp( - 0.995, self.derived["psi_pol_n"], self.profiles["delta(-)"] - ) + self.derived["delta995"] = np.interp(0.995, self.derived["psi_pol_n"], self.profiles["delta(-)"]) self.derived["Rgeo"] = float(self.profiles["rcentr(m)"][-1]) self.derived["B0"] = np.abs(float(self.profiles["bcentr(T)"][-1])) @@ -339,21 +246,17 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= # --------------------------------------------------------------------------------------------------------------------- """ - surf_miller is truly surface area, but because of the GACODE definitions of flux, + surf_geo is truly surface area, but because of the GACODE definitions of flux, Surf = V' <|grad r|> Surf_GACODE = V' """ - self.derived["surfGACODE_miller"] = (self.derived["surf_miller"] / self.derived["gradr_miller"]) + self.derived["surfGACODE_geo"] = (self.derived["surf_geo"] / self.derived["gradr_geo"]) - self.derived["surfGACODE_miller"][np.isnan(self.derived["surfGACODE_miller"])] = 0 + self.derived["surfGACODE_geo"][np.isnan(self.derived["surfGACODE_geo"])] = 0 - self.derived["c_s"] = PLASMAtools.c_s( - self.profiles["te(keV)"], self.derived["mi_ref"] - ) - self.derived["rho_s"] = PLASMAtools.rho_s( - self.profiles["te(keV)"], self.derived["mi_ref"], self.derived["B_unit"] - ) + self.derived["c_s"] = PLASMAtools.c_s(self.profiles["te(keV)"], self.derived["mi_ref"]) + self.derived["rho_s"] = PLASMAtools.rho_s(self.profiles["te(keV)"], self.derived["mi_ref"], self.derived["B_unit"]) self.derived["q_gb"], self.derived["g_gb"], self.derived["pi_gb"], self.derived["s_gb"], _ = PLASMAtools.gyrobohmUnits( self.profiles["te(keV)"], @@ -427,25 +330,23 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= """ r = self.profiles["rmin(m)"] - volp = self.derived["volp_miller"] + volp = self.derived["volp_geo"] - self.derived["qe_MWmiller"] = CALCtools.integrateFS(self.derived["qe"], r, volp) - self.derived["qi_MWmiller"] = CALCtools.integrateFS(self.derived["qi"], r, volp) - self.derived["ge_10E20miller"] = CALCtools.integrateFS( - self.derived["ge"] * 1e-20, r, volp - ) # Because the units were #/sec/m^3 + self.derived["qe_MW"] = CALCtools.integrateFS(self.derived["qe"], r, volp) + self.derived["qi_MW"] = CALCtools.integrateFS(self.derived["qi"], r, volp) + self.derived["ge_10E20"] = CALCtools.integrateFS(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["geIn"] = self.derived["ge_10E20"][-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) + self.derived["qe_MWm2"] = self.derived["qe_MW"] / (volp) + self.derived["qi_MWm2"] = self.derived["qi_MW"] / (volp) + self.derived["ge_10E20m2"] = self.derived["ge_10E20"] / (volp) self.derived["QiQe"] = self.derived["qi_MWm2"] / np.where(self.derived["qe_MWm2"] == 0, 1e-10, self.derived["qe_MWm2"]) # to avoid division by zero # "Convective" flux - self.derived["ce_MWmiller"] = PLASMAtools.convective_flux( - self.profiles["te(keV)"], self.derived["ge_10E20miller"] + self.derived["ce_MW"] = PLASMAtools.convective_flux( + self.profiles["te(keV)"], self.derived["ge_10E20"] ) self.derived["ce_MWm2"] = PLASMAtools.convective_flux( self.profiles["te(keV)"], self.derived["ge_10E20m2"] @@ -465,10 +366,10 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= P += self.profiles["qbrem(MW/m^3)"] if "qline(MW/m^3)" in self.profiles: P += self.profiles["qline(MW/m^3)"] - self.derived["qe_rad_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_rad_MW"] = CALCtools.integrateFS(P, r, volp) P = self.profiles["qei(MW/m^3)"] - self.derived["qe_exc_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_exc_MW"] = CALCtools.integrateFS(P, r, volp) """ --------------------------------------------------------------------------------------------------------------------- @@ -485,14 +386,14 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= P += self.profiles[i] self.derived["qe_auxONLY"] = copy.deepcopy(P) - self.derived["qe_auxONLY_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_auxONLY_MW"] = CALCtools.integrateFS(P, r, volp) for i in ["qione(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] self.derived["qe_aux"] = copy.deepcopy(P) - self.derived["qe_aux_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_aux_MW"] = CALCtools.integrateFS(P, r, volp) # ** Ions @@ -502,14 +403,14 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= P += self.profiles[i] self.derived["qi_auxONLY"] = copy.deepcopy(P) - self.derived["qi_auxONLY_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qi_auxONLY_MW"] = CALCtools.integrateFS(P, r, volp) for i in ["qioni(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] self.derived["qi_aux"] = copy.deepcopy(P) - self.derived["qi_aux_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qi_aux_MW"] = CALCtools.integrateFS(P, r, volp) # ** General @@ -517,19 +418,19 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= for i in ["qohme(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qOhm_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qOhm_MW"] = CALCtools.integrateFS(P, r, volp) P = np.zeros(len(self.profiles["rho(-)"])) for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qRF_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qRF_MW"] = CALCtools.integrateFS(P, r, volp) if "qrfe(MW/m^3)" in self.profiles: - self.derived["qRFe_MWmiller"] = CALCtools.integrateFS( + self.derived["qRFe_MW"] = CALCtools.integrateFS( self.profiles["qrfe(MW/m^3)"], r, volp ) if "qrfi(MW/m^3)" in self.profiles: - self.derived["qRFi_MWmiller"] = CALCtools.integrateFS( + self.derived["qRFi_MW"] = CALCtools.integrateFS( self.profiles["qrfi(MW/m^3)"], r, volp ) @@ -537,36 +438,36 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= for i in ["qbeame(MW/m^3)", "qbeami(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qBEAM_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qBEAM_MW"] = CALCtools.integrateFS(P, r, volp) - self.derived["qrad_MWmiller"] = CALCtools.integrateFS(self.derived["qrad"], r, volp) + self.derived["qrad_MW"] = CALCtools.integrateFS(self.derived["qrad"], r, volp) if "qsync(MW/m^3)" in self.profiles: - self.derived["qrad_sync_MWmiller"] = CALCtools.integrateFS(self.profiles["qsync(MW/m^3)"], r, volp) + self.derived["qrad_sync_MW"] = CALCtools.integrateFS(self.profiles["qsync(MW/m^3)"], r, volp) else: - self.derived["qrad_sync_MWmiller"] = self.derived["qrad_MWmiller"]*0.0 + self.derived["qrad_sync_MW"] = self.derived["qrad_MW"]*0.0 if "qbrem(MW/m^3)" in self.profiles: - self.derived["qrad_brem_MWmiller"] = CALCtools.integrateFS(self.profiles["qbrem(MW/m^3)"], r, volp) + self.derived["qrad_brem_MW"] = CALCtools.integrateFS(self.profiles["qbrem(MW/m^3)"], r, volp) else: - self.derived["qrad_brem_MWmiller"] = self.derived["qrad_MWmiller"]*0.0 + self.derived["qrad_brem_MW"] = self.derived["qrad_MW"]*0.0 if "qline(MW/m^3)" in self.profiles: - self.derived["qrad_line_MWmiller"] = CALCtools.integrateFS(self.profiles["qline(MW/m^3)"], r, volp) + self.derived["qrad_line_MW"] = CALCtools.integrateFS(self.profiles["qline(MW/m^3)"], r, volp) else: - self.derived["qrad_line_MWmiller"] = self.derived["qrad_MWmiller"]*0.0 + self.derived["qrad_line_MW"] = self.derived["qrad_MW"]*0.0 P = np.zeros(len(self.profiles["rho(-)"])) for i in ["qfuse(MW/m^3)", "qfusi(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qFus_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qFus_MW"] = CALCtools.integrateFS(P, r, volp) P = np.zeros(len(self.profiles["rho(-)"])) for i in ["qioni(MW/m^3)", "qione(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qz_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qz_MW"] = CALCtools.integrateFS(P, r, volp) - self.derived["q_MWmiller"] = ( - self.derived["qe_MWmiller"] + self.derived["qi_MWmiller"] + self.derived["q_MW"] = ( + self.derived["qe_MW"] + self.derived["qi_MW"] ) # --------------------------------------------------------------------------------------------------------------------- @@ -575,12 +476,12 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= P = np.zeros(len(self.profiles["rho(-)"])) if "qfuse(MW/m^3)" in self.profiles: P = self.profiles["qfuse(MW/m^3)"] - self.derived["qe_fus_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_fus_MW"] = CALCtools.integrateFS(P, r, volp) P = np.zeros(len(self.profiles["rho(-)"])) if "qfusi(MW/m^3)" in self.profiles: P = self.profiles["qfusi(MW/m^3)"] - self.derived["qi_fus_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["qi_fus_MW"] = CALCtools.integrateFS(P, r, volp) P = np.zeros(len(self.profiles["rho(-)"])) if "qfusi(MW/m^3)" in self.profiles: @@ -589,7 +490,7 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= ) * 5 P = self.derived["q_fus"] self.derived["q_fus"] = P - self.derived["q_fus_MWmiller"] = CALCtools.integrateFS(P, r, volp) + self.derived["q_fus_MW"] = CALCtools.integrateFS(P, r, volp) """ Derivatives @@ -622,26 +523,26 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= """ Other, performance """ - qFus = self.derived["qe_fus_MWmiller"] + self.derived["qi_fus_MWmiller"] + qFus = self.derived["qe_fus_MW"] + self.derived["qi_fus_MW"] self.derived["Pfus"] = qFus[-1] * 5 # Note that in cases with NPRAD=0 in TRANPS, this includes radiation! no way to deal wit this... - qIn = self.derived["qe_aux_MWmiller"] + self.derived["qi_aux_MWmiller"] + qIn = self.derived["qe_aux_MW"] + self.derived["qi_aux_MW"] self.derived["qIn"] = qIn[-1] self.derived["Q"] = self.derived["Pfus"] / self.derived["qIn"] self.derived["qHeat"] = qIn[-1] + qFus[-1] self.derived["qTr"] = ( - self.derived["qe_aux_MWmiller"] - + self.derived["qi_aux_MWmiller"] - + (self.derived["qe_fus_MWmiller"] + self.derived["qi_fus_MWmiller"]) - - self.derived["qrad_MWmiller"] + self.derived["qe_aux_MW"] + + self.derived["qi_aux_MW"] + + (self.derived["qe_fus_MW"] + self.derived["qi_fus_MW"]) + - self.derived["qrad_MW"] ) - self.derived["Prad"] = self.derived["qrad_MWmiller"][-1] - self.derived["Prad_sync"] = self.derived["qrad_sync_MWmiller"][-1] - self.derived["Prad_brem"] = self.derived["qrad_brem_MWmiller"][-1] - self.derived["Prad_line"] = self.derived["qrad_line_MWmiller"][-1] + self.derived["Prad"] = self.derived["qrad_MW"][-1] + self.derived["Prad_sync"] = self.derived["qrad_sync_MW"][-1] + self.derived["Prad_brem"] = self.derived["qrad_brem_MW"][-1] + self.derived["Prad_line"] = self.derived["qrad_line_MW"][-1] self.derived["Psol"] = self.derived["qHeat"] - self.derived["Prad"] self.derived["ni_thr"] = [] @@ -884,8 +785,8 @@ def deriveQuantities_full(self, mi_ref=None, n_theta_geo=1001, rederiveGeometry= --------------------------------------------------------------------------------------------------- ''' - self.derived["bp2_exp"] = self.derived["bp2_miller"] * self.derived["B_unit"] ** 2 - self.derived["bt2_exp"] = self.derived["bt2_miller"] * self.derived["B_unit"] ** 2 + self.derived["bp2_exp"] = self.derived["bp2_geo"] * self.derived["B_unit"] ** 2 + self.derived["bt2_exp"] = self.derived["bt2_geo"] * self.derived["B_unit"] ** 2 # Calculate the volume averages of bt2 and bp2 @@ -1099,7 +1000,7 @@ def deriveContentByVolumes(self, rhos=[0.5], impurityPosition=3): np.expand_dims( np.transpose(self.profiles["ni(10^19/m^3)"][:i] * 0.1), 0 ), - np.expand_dims(self.derived["volp_miller"][:i], 0), + np.expand_dims(self.derived["volp_geo"][:i], 0), ) _, _, Ni_x[i], _ = PLASMAtools.calculateContent( @@ -1112,7 +1013,7 @@ def deriveContentByVolumes(self, rhos=[0.5], impurityPosition=3): np.expand_dims( np.transpose(self.profiles["ni(10^19/m^3)"][:i] * 0.1), 0 ), - np.expand_dims(self.derived["volp_miller"][:i], 0), + np.expand_dims(self.derived["volp_geo"][:i], 0), ) We, Wi, Ne, Ni = ( @@ -1169,7 +1070,7 @@ def printInfo(self, label="", reDeriveIfNotFound=True): 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.derive_quantities() self.printInfo(label=label, reDeriveIfNotFound=False) def export_to_table(self, table=None, name=None): @@ -1353,13 +1254,9 @@ def changeResolution(self, n=100, rho_new=None, interpolation_function=MATHtools pro[i] = np.transpose(prof) - self.produce_shape_lists() + self.derive_quantities() - self.deriveQuantities() - - print( - f"\t\t- Resolution of profiles changed to {n} points with function {interpolation_function}" - ) + print(f"\t\t- Resolution of profiles changed to {n} points with function {interpolation_function}") def DTplasma(self): self.Dion, self.Tion = None, None @@ -1430,7 +1327,7 @@ def remove(self, ions_list): ) self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) print("\t\t\t- Set of ions in updated profiles: ", self.profiles["name"]) @@ -1453,11 +1350,7 @@ def lumpSpecies( else: lab = "therm" - print( - "\t\t- Lumping ions in positions (of ions order, no zero): ", - ions_list, - typeMsg="i", - ) + print("\t\t- Lumping ions in positions (of ions order, no zero): ",ions_list,typeMsg="i",) if forcename is None: forcename = "LUMPED" @@ -1472,7 +1365,7 @@ def lumpSpecies( Zr = fZ2 / fZ1 Zr_vol = ( CALCtools.integrateFS( - Zr, self.profiles["rmin(m)"], self.derived["volp_miller"] + Zr, self.profiles["rmin(m)"], self.derived["volp_geo"] )[-1] / self.derived["volume"] ) @@ -1507,7 +1400,7 @@ def lumpSpecies( ) self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) # Remove species self.remove(ions_list) @@ -1585,11 +1478,11 @@ def changeZeff(self, Zeff, ion_pos=2, quasineutral_ions=None, enforceSameGradien self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) if enforceSameGradients: self.scaleAllThermalDensities() - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) print(f'\t\t\t* Dilution changed from {fi_orig.mean():.2e} (vol avg) to { self.derived["fi"][:, ion_pos].mean():.2e} to achieve Zeff={self.derived["Zeff_vol"]:.3f} (fDT={self.derived["fmain"]:.3f}) [quasineutrality error = {self.derived["QN_Error"]:.1e}]') @@ -1620,7 +1513,7 @@ def moveSpecie(self, pos=2, pos_new=1): ) self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) if position_to_moveTO_in_profiles > position_to_moveFROM_in_profiles: self.remove([position_to_moveFROM_in_profiles + 1]) @@ -1660,7 +1553,7 @@ def addSpecie(self, Z=5.0, mass=10.0, fi_vol=0.1, forcename=None): ) self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) def correct(self, options={}, write=False, new_file=None): """ @@ -1738,7 +1631,7 @@ def correct(self, options={}, write=False, new_file=None): # Recompute ptot if recompute_ptot: - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) self.selfconsistentPTOT() # If I don't trust the negative particle flux in the core that comes from TRANSP... @@ -1755,7 +1648,7 @@ def correct(self, options={}, write=False, new_file=None): # Re-derive # ---------------------------------------------------------------------- - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) # ---------------------------------------------------------------------- # Write @@ -1771,7 +1664,7 @@ def enforce_same_density_gradients(self): self.profiles["ni(10^19/m^3)"][:, sp] = self.derived["fi_vol"][sp] * self.profiles["ne(10^19/m^3)"] txt += f"{self.Species[sp]['N']} " print(f"\t\t- Making all thermal ions ({txt}) have the same a/Ln as electrons (making them an exact flat fraction)",typeMsg="i",) - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) def make_fast_ions_thermal(self): modified_num = 0 @@ -1819,18 +1712,18 @@ def enforceQuasineutrality(self): f"\t\t\t\t- Changed on-axis density from n0 = {prev_on_axis:.2f} to {new_on_axis:.2f} ({100*(new_on_axis-prev_on_axis)/prev_on_axis:.1f}%)" ) - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) def introduceRotationProfile(self, Mach_LF=1.0, new_file=None): print(f"\t- Enforcing Mach Number in LF of {Mach_LF}") - self.deriveQuantities() + self.derive_quantities() Vtor_LF = PLASMAtools.constructVtorFromMach( Mach_LF, self.profiles["ti(keV)"][:, 0], self.derived["mbg"] ) # m/s self.profiles["w0(rad/s)"] = Vtor_LF / (self.derived["R_LF"]) # rad/s - self.deriveQuantities() + self.derive_quantities() if new_file is not None: self.writeCurrentStatus(file=new_file) @@ -2532,7 +2425,7 @@ def plot( GRAPHICStools.autoscale_y(ax, bottomy=0) ax = axs6[3] - ax.plot(self.profiles["rho(-)"], self.derived["q_fus_MWmiller"], c=color, lw=lw) + ax.plot(self.profiles["rho(-)"], self.derived["q_fus_MW"], c=color, lw=lw) ax.set_ylabel("$P_{fus}$ ($MW$)") ax.set_xlim([0, 1]) @@ -2576,10 +2469,10 @@ def plot( label=extralab + "$Q_i/Q_e$", ) safe_division = np.divide( - self.derived["qi_aux_MWmiller"], - self.derived["qe_aux_MWmiller"], - where=self.derived["qe_aux_MWmiller"] != 0, - out=np.full_like(self.derived["qi_aux_MWmiller"], np.nan), + self.derived["qi_aux_MW"], + self.derived["qe_aux_MW"], + where=self.derived["qe_aux_MW"] != 0, + out=np.full_like(self.derived["qi_aux_MW"], np.nan), ) ax.plot( self.profiles["rho(-)"], @@ -3028,15 +2921,15 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): # F ax = axs[2] P = ( - self.derived["qe_fus_MWmiller"] - + self.derived["qe_aux_MWmiller"] - + -self.derived["qe_rad_MWmiller"] - + -self.derived["qe_exc_MWmiller"] + self.derived["qe_fus_MW"] + + self.derived["qe_aux_MW"] + + -self.derived["qe_rad_MW"] + + -self.derived["qe_exc_MW"] ) ax.plot( roa, - -self.derived["qe_MWmiller"], + -self.derived["qe_MW"], c="g", lw=2, label="$P_{e}$" if leg else "", @@ -3044,7 +2937,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qe_fus_MWmiller"], + self.derived["qe_fus_MW"], c="r", lw=2, label="$P_{fus,e}$" if leg else "", @@ -3052,7 +2945,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qe_aux_MWmiller"], + self.derived["qe_aux_MW"], c="b", lw=2, label="$P_{aux,e}$" if leg else "", @@ -3060,7 +2953,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - -self.derived["qe_exc_MWmiller"], + -self.derived["qe_exc_MW"], c="m", lw=2, label="$P_{exc,e}$" if leg else "", @@ -3068,7 +2961,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - -self.derived["qe_rad_MWmiller"], + -self.derived["qe_rad_MW"], c="c", lw=2, label="$P_{rad,e}$" if leg else "", @@ -3081,7 +2974,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ax.plot( roa, - -self.derived["ce_MWmiller"], + -self.derived["ce_MW"], c="k", lw=1, label="($P_{conv,e}$)" if leg else "", @@ -3102,14 +2995,14 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ax = axs[3] P = ( - self.derived["qi_fus_MWmiller"] - + self.derived["qi_aux_MWmiller"] - + self.derived["qe_exc_MWmiller"] + self.derived["qi_fus_MW"] + + self.derived["qi_aux_MW"] + + self.derived["qe_exc_MW"] ) ax.plot( roa, - -self.derived["qi_MWmiller"], + -self.derived["qi_MW"], c="g", lw=2, label="$P_{i}$" if leg else "", @@ -3117,7 +3010,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qi_fus_MWmiller"], + self.derived["qi_fus_MW"], c="r", lw=2, label="$P_{fus,i}$" if leg else "", @@ -3125,7 +3018,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qi_aux_MWmiller"], + self.derived["qi_aux_MW"], c="b", lw=2, label="$P_{aux,i}$" if leg else "", @@ -3133,7 +3026,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qe_exc_MWmiller"], + self.derived["qe_exc_MW"], c="m", lw=2, label="$P_{exc,i}$" if leg else "", @@ -3162,7 +3055,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ax.plot( roa, - self.derived["ge_10E20miller"], + self.derived["ge_10E20"], c="g", lw=2, label="$\\Gamma_{e}$" if leg else "", @@ -3185,17 +3078,17 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): # TOTAL ax = axs[5] P = ( - self.derived["qOhm_MWmiller"] - + self.derived["qRF_MWmiller"] - + self.derived["qFus_MWmiller"] - + -self.derived["qe_rad_MWmiller"] - + self.derived["qz_MWmiller"] - + self.derived["qBEAM_MWmiller"] + self.derived["qOhm_MW"] + + self.derived["qRF_MW"] + + self.derived["qFus_MW"] + + -self.derived["qe_rad_MW"] + + self.derived["qz_MW"] + + self.derived["qBEAM_MW"] ) ax.plot( roa, - -self.derived["q_MWmiller"], + -self.derived["q_MW"], c="g", lw=2, label="$P$" if leg else "", @@ -3203,7 +3096,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qOhm_MWmiller"], + self.derived["qOhm_MW"], c="k", lw=2, label="$P_{Oh}$" if leg else "", @@ -3211,7 +3104,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qRF_MWmiller"], + self.derived["qRF_MW"], c="b", lw=2, label="$P_{RF}$" if leg else "", @@ -3219,7 +3112,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qBEAM_MWmiller"], + self.derived["qBEAM_MW"], c="pink", lw=2, label="$P_{NBI}$" if leg else "", @@ -3227,7 +3120,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qFus_MWmiller"], + self.derived["qFus_MW"], c="r", lw=2, label="$P_{fus}$" if leg else "", @@ -3235,7 +3128,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - -self.derived["qe_rad_MWmiller"], + -self.derived["qe_rad_MW"], c="c", lw=2, label="$P_{rad}$" if leg else "", @@ -3243,7 +3136,7 @@ def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ) ax.plot( roa, - self.derived["qz_MWmiller"], + self.derived["qz_MW"], c="orange", lw=1, label="$P_{ionz.}$" if leg else "", @@ -3369,7 +3262,7 @@ def plotPeaking( # Extra r = self.profiles["rmin(m)"] - volp = self.derived["volp_miller"] + volp = self.derived["volp_geo"] ix = np.argmin(np.abs(self.profiles["rho(-)"] - 0.9)) if debugPlot: @@ -3625,22 +3518,22 @@ def parabolizePlasma(self): self.profiles["ne(10^19/m^3)"] = n self.scaleAllThermalDensities(scaleFactor=factor_n) - self.deriveQuantities() + self.derive_quantities() def changeRFpower(self, PrfMW=25.0): """ keeps same partition """ - print(f"- Changing the RF power from {self.derived['qRF_MWmiller'][-1]:.1f} MW to {PrfMW:.1f} MW",typeMsg="i",) + print(f"- Changing the RF power from {self.derived['qRF_MW'][-1]:.1f} MW to {PrfMW:.1f} MW",typeMsg="i",) - if self.derived["qRF_MWmiller"][-1] == 0.0: + if self.derived["qRF_MW"][-1] == 0.0: raise Exception("No RF power in the input.gacode, cannot modify the RF power") for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: - self.profiles[i] = self.profiles[i] * PrfMW / self.derived["qRF_MWmiller"][-1] + self.profiles[i] = self.profiles[i] * PrfMW / self.derived["qRF_MW"][-1] - self.deriveQuantities() + self.derive_quantities() def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0.2): @@ -3705,10 +3598,10 @@ def addSawtoothEffectOnOhmic(self, PohTot, mixRadius=None, plotYN=False): if mixRadius is None: mixRadius = self.profiles["rho(-)"][np.where(self.profiles["q(-)"] > 1)][0] - print(f"\t- Original Ohmic power: {self.derived['qOhm_MWmiller'][-1]:.2f}MW") + print(f"\t- Original Ohmic power: {self.derived['qOhm_MW'][-1]:.2f}MW") Ohmic_old = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) - dvol = self.derived["volp_miller"] * np.append( + dvol = self.derived["volp_geo"] * np.append( [0], np.diff(self.profiles["rmin(m)"]) ) @@ -3718,13 +3611,13 @@ def addSawtoothEffectOnOhmic(self, PohTot, mixRadius=None, plotYN=False): Psaw = CDFtools.profilePower( self.profiles["rho(-)"], dvol, - PohTot - self.derived["qOhm_MWmiller"][-1], + PohTot - self.derived["qOhm_MW"][-1], mixRadius, ) self.profiles["qohme(MW/m^3)"] += Psaw - self.deriveQuantities() + self.derive_quantities() - print(f"\t- New Ohmic power: {self.derived['qOhm_MWmiller'][-1]:.2f}MW") + print(f"\t- New Ohmic power: {self.derived['qOhm_MW'][-1]:.2f}MW") Ohmic_new = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) if plotYN: @@ -4228,7 +4121,7 @@ def gradientsMerger(p0, p_true, roa=0.46, blending=0.1): p.profiles["ti(keV)"][:, 0] = Ti p.profiles["ne(10^19/m^3)"] = ne - p.deriveQuantities() + p.derive_quantities() return p diff --git a/src/mitim_tools/popcon_tools/POPCONtools.py b/src/mitim_tools/popcon_tools/POPCONtools.py index f5dde84c..71c3426e 100644 --- a/src/mitim_tools/popcon_tools/POPCONtools.py +++ b/src/mitim_tools/popcon_tools/POPCONtools.py @@ -202,8 +202,8 @@ def match_pfus(self, Prad_residual = (point['P_radiation'].data.magnitude - profiles_gacode.derived['Prad']) / profiles_gacode.derived['Prad'] - Pin_derived = (profiles_gacode.derived['qi_aux_MWmiller'][-1] - +profiles_gacode.derived['qe_aux_MWmiller'][-1] + Pin_derived = (profiles_gacode.derived['qi_aux_MW'][-1] + +profiles_gacode.derived['qe_aux_MW'][-1] ) Pin_residual = (point['P_auxillary_launched'].data.magnitude - Pin_derived) / Pin_derived @@ -428,13 +428,13 @@ def print_data(self, print(f"P_sol: ", f"POPCON: {(point['P_LH_thresh'].data.magnitude *point['ratio_of_P_SOL_to_P_LH'].data.magnitude):.2f}", f"GACODE: {profiles_gacode.derived['Psol']:.2f}","(MW)", f"(~%{point['ratio_of_P_SOL_to_P_LH'].data.magnitude*1e2:.2f} of LH threshold)") print(f"P_aux: ", f"POPCON: {point['P_auxillary_launched'].data.magnitude:.2f}", - f"GACODE: {(profiles_gacode.derived['qi_aux_MWmiller'][-1]+profiles_gacode.derived['qe_aux_MWmiller'][-1]):.2f}", + f"GACODE: {(profiles_gacode.derived['qi_aux_MW'][-1]+profiles_gacode.derived['qe_aux_MW'][-1]):.2f}", "(MW)") print(f"P_rad: ", f"POPCON: {point['P_radiation'].data.magnitude:.2f}",f"GACODE: {profiles_gacode.derived['Prad']:.2f}","(MW)") print(f"P_ext: ", f"POPCON: {point['P_external'].data.magnitude:.2f}","(MW)") - print(f"P_ohm: ", f"POPCON: {point['P_ohmic'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['qOhm_MWmiller'][-1]:.2f}","(MW)") + print(f"P_ohm: ", f"POPCON: {point['P_ohmic'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['qOhm_MW'][-1]:.2f}","(MW)") print(f"P_in: ", f"POPCON: {point['P_in'].data.magnitude:.2f}", - f"GACODE: {(profiles_gacode.derived['qOhm_MWmiller'][-1]+profiles_gacode.derived['qi_aux_MWmiller'][-1]+profiles_gacode.derived['qe_aux_MWmiller'][-1]+profiles_gacode.derived['Pfus']*0.2):.2f}", + f"GACODE: {(profiles_gacode.derived['qOhm_MW'][-1]+profiles_gacode.derived['qi_aux_MW'][-1]+profiles_gacode.derived['qe_aux_MW'][-1]+profiles_gacode.derived['Pfus']*0.2):.2f}", "(MW)") print(f"q95: ", f"POPCON: {point['q_star'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['q95']:.2f}") print(f"Wtot: ", f"POPCON: {point['plasma_stored_energy'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['Wthr']:.2f}", "(MJ)") diff --git a/src/mitim_tools/popcon_tools/RAPIDStools.py b/src/mitim_tools/popcon_tools/RAPIDStools.py index 77f0a6e1..d435555d 100644 --- a/src/mitim_tools/popcon_tools/RAPIDStools.py +++ b/src/mitim_tools/popcon_tools/RAPIDStools.py @@ -94,7 +94,7 @@ def rapids_evaluator(nn, aLT, aLn, TiTe, p_base, R=None, a=None, Bt=None, Ip=Non p.profiles['ne(10^19/m^3)'] = np.interp(p.derived['roa'], roa, ne) p.profiles['ni(10^19/m^3)'] = p_base.profiles['ni(10^19/m^3)'] * np.transpose(np.atleast_2d((p.profiles['ne(10^19/m^3)']/p_base.profiles['ne(10^19/m^3)']))) - p.deriveQuantities() + p.derive_quantities() # Change Zeff p.changeZeff(Zeff, ion_pos=3) @@ -116,7 +116,7 @@ def pedestal(p): p = eped_profiler(p, rhotop_assume, rhotop, Tetop_keV, Titop_keV, netop_20) # Derive quantities - p.deriveQuantities(rederiveGeometry=False) + p.derive_quantities(rederiveGeometry=False) BetaN_used = p.derived['BetaN_engineering'] * BetaN_multiplier diff --git a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py index 5c2a6058..6a331280 100644 --- a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py +++ b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py @@ -808,7 +808,7 @@ def _produce_quantity_profiles(self, var = 'Te', Vsurf = None): z = self.p.derived['Zeff'] elif var == 'PichT': x = [1] - z = self.p.derived['qRF_MWmiller'][-1]*1E6 + z = self.p.derived['qRF_MW'][-1]*1E6 return x,z From e650b4aad3b0bafc42eacf2a868c727a85f03094 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 16:16:27 +0200 Subject: [PATCH 082/385] Started VMEC tools --- src/mitim_tools/gacode_tools/PROFILEStools.py | 24 ++++++++--------- .../plasmastate_tools/utils/VMECtools.py | 26 +++++++++++++++++++ 2 files changed, 37 insertions(+), 13 deletions(-) create mode 100644 src/mitim_tools/plasmastate_tools/utils/VMECtools.py diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index bf8fef43..ffdf273d 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -25,12 +25,16 @@ def __init__(self, file, calculateDerived=True, mi_ref=None): super().__init__(type_file='input.gacode') - """ - Depending on resolution, derived can be expensive, so I mmay not do it every time - """ - self.file = file + self._read_inputgacocde() + + if self.file is not None: + # Derive (Depending on resolution, derived can be expensive, so I mmay not do it every time) + self.derive_quantities(mi_ref=mi_ref, calculateDerived=calculateDerived) + + def _read_inputgacocde(self): + self.titles_singleNum = ["nexp", "nion", "shot", "name", "type", "time"] self.titles_singleArr = ["masse","mass","ze","z","torfluxa(Wb/radian)","rcentr(m)","bcentr(T)","current(MA)"] self.titles_single = self.titles_singleNum + self.titles_singleArr @@ -44,9 +48,6 @@ def __init__(self, file, calculateDerived=True, mi_ref=None): self._read_profiles() self._ensure_existence() - # Derive - self.derive_quantities(mi_ref=mi_ref, calculateDerived=calculateDerived) - def _read_header(self): for i in range(len(self.lines)): if "# nexp" in self.lines[i]: @@ -54,7 +55,7 @@ def _read_header(self): self.header = self.lines[:istartProfs] def _read_profiles(self): - singleLine, title, var = None, None, None # for ruff complaints + singleLine, title, var = None, None, None # --- found = False @@ -92,10 +93,7 @@ def _read_profiles(self): """ Sometimes there's a bug in TGYRO, where the powers may be too low (E-191) that cannot be properly written """ - varT = [ - float(j) if (j[-4].upper() == "E" or "." in j) else 0.0 - for j in var0[1:] - ] + varT = [float(j) if (j[-4].upper() == "E" or "." in j) else 0.0for j in var0[1:]] var.append(varT) @@ -148,7 +146,7 @@ def _ensure_existence(self): self.varqmom, ] - num_moments = 6 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros + num_moments = 7 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros for i in range(num_moments): some_times_are_not_here.append(f"shape_cos{i + 1}(-)") if i > 1: diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py new file mode 100644 index 00000000..8480922d --- /dev/null +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -0,0 +1,26 @@ + +from mitim_tools.plasmastate_tools.MITIMstate import mitim_state + +class vmec_state(mitim_state): + ''' + Class to read and manipulate VMEC files + ''' + + # ------------------------------------------------------------------ + # Reading and interpreting + # ------------------------------------------------------------------ + + def __init__(self, file, calculateDerived=True, mi_ref=None): + + super().__init__(type_file='vmec') + + self.file = file + + self._read_vmec() + + if self.file is not None: + # Derive (Depending on resolution, derived can be expensive, so I mmay not do it every time) + self.derive_quantities(mi_ref=mi_ref, calculateDerived=calculateDerived) + + def _read_vmec(self): + pass From 6762771e8ba6624e3914a044e73461996ab63e8d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 16:17:19 +0200 Subject: [PATCH 083/385] named changed from PROFILES_GACODE to gacode_state --- docs/capabilities/misc_capabilities.rst | 2 +- docs/capabilities/tgyro_capabilities.rst | 4 +- src/mitim_modules/maestro/MAESTROmain.py | 4 +- .../maestro/scripts/run_maestro.py | 2 +- src/mitim_modules/maestro/utils/EPEDbeat.py | 6 +- .../maestro/utils/MAESTRObeat.py | 2 +- .../maestro/utils/MAESTROplot.py | 4 +- .../maestro/utils/PORTALSbeat.py | 2 +- src/mitim_modules/maestro/utils/TRANSPbeat.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 4 +- .../portals/utils/PORTALSanalysis.py | 6 +- .../portals/utils/PORTALSinit.py | 2 +- src/mitim_modules/powertorch/STATEtools.py | 6 +- .../physics_models/transport_cgyro.py | 4 +- .../powertorch/scripts/calculateTargets.py | 2 +- .../scripts/compareRadialResolution.py | 2 +- .../powertorch/utils/TRANSFORMtools.py | 2 +- .../powertorch/utils/TRANSPORTtools.py | 4 +- src/mitim_tools/astra_tools/ASTRAtools.py | 4 +- src/mitim_tools/gacode_tools/NEOtools.py | 2 +- src/mitim_tools/gacode_tools/PROFILEStools.py | 2 +- src/mitim_tools/gacode_tools/TGLFtools.py | 6 +- src/mitim_tools/gacode_tools/TGYROtools.py | 10 +- .../gacode_tools/scripts/compare_MXH3.py | 2 +- .../gacode_tools/scripts/read_gacode.py | 2 +- .../gacode_tools/scripts/read_tgyro.py | 2 +- src/mitim_tools/popcon_tools/POPCONtools.py | 116 +++++++++--------- src/mitim_tools/transp_tools/CDFtools.py | 2 +- .../transp_tools/utils/TRANSPhelpers.py | 2 +- tests/PORTALS_workflow.py | 2 +- tests/POWERTORCH_workflow.py | 2 +- tests/TGYRO_workflow.py | 2 +- 32 files changed, 108 insertions(+), 108 deletions(-) diff --git a/docs/capabilities/misc_capabilities.rst b/docs/capabilities/misc_capabilities.rst index 49d3ccdc..41e23b1e 100644 --- a/docs/capabilities/misc_capabilities.rst +++ b/docs/capabilities/misc_capabilities.rst @@ -47,7 +47,7 @@ To open and plot an ``input.gacode`` file: .. code-block:: python from mitim_tools.gacode_tools import PROFILEStools - p = PROFILEStools.PROFILES_GACODE(file) + p = PROFILEStools.gacode_state(file) p.plot() It will plot results in a notebook-like plot with different tabs: diff --git a/docs/capabilities/tgyro_capabilities.rst b/docs/capabilities/tgyro_capabilities.rst index a099b69b..e1a35c1b 100644 --- a/docs/capabilities/tgyro_capabilities.rst +++ b/docs/capabilities/tgyro_capabilities.rst @@ -35,7 +35,7 @@ Create a PROFILES class from the input.gacode file: .. code-block:: python - profiles = PROFILEStools.PROFILES_GACODE(gacode_file) + profiles = PROFILEStools.gacode_state(gacode_file) .. tip:: @@ -120,7 +120,7 @@ Create a profiles class with the `input.gacode` file that TGYRO used to run and gacode_file = Path('MITIM-fusion/tests/data/input.gacode') folder = Path('MITIM-fusion/tests/scratch/tgyro_tut/run1') - profiles = PROFILEStools.PROFILES_GACODE(gacode_file) + profiles = PROFILEStools.gacode_state(gacode_file) tgyro_out = TGYROtools.TGYROoutput(folder,profiles=profiles) Plot results: diff --git a/src/mitim_modules/maestro/MAESTROmain.py b/src/mitim_modules/maestro/MAESTROmain.py index 8d1f8a72..9b3923b2 100644 --- a/src/mitim_modules/maestro/MAESTROmain.py +++ b/src/mitim_modules/maestro/MAESTROmain.py @@ -193,7 +193,7 @@ def initialize(self, *args, **kwargs): if self.profiles_with_engineering_parameters is None: # First initialization, freeze engineering parameters - self._freeze_parameters(profiles = PROFILEStools.PROFILES_GACODE(self.beat.initialize.folder / 'input.gacode')) + self._freeze_parameters(profiles = PROFILEStools.gacode_state(self.beat.initialize.folder / 'input.gacode')) @mitim_timer('\t\t* Preparation', name_timer=None) def prepare(self, *args, **kwargs): @@ -255,7 +255,7 @@ def run(self, **kwargs): def _freeze_parameters(self, profiles = None): if profiles is None: - profiles = PROFILEStools.PROFILES_GACODE(self.beat.folder_output / 'input.gacode') + profiles = PROFILEStools.gacode_state(self.beat.folder_output / 'input.gacode') print('\t\t- Freezing engineering parameters from MAESTRO') self.profiles_with_engineering_parameters = copy.deepcopy(profiles) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index 56279d35..b1a14749 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -139,7 +139,7 @@ def parse_maestro_nml(file_path): # add postprocessing function def profiles_postprocessing_fun(file_profs): - p = PROFILEStools.PROFILES_GACODE(file_profs) + p = PROFILEStools.gacode_state(file_profs) if lumpImpurities: p.lumpImpurities() if enforce_same_density_gradients: diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index 4f6c1d8f..2ac10ae7 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -291,7 +291,7 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F def finalize(self, **kwargs): - self.profiles_output = PROFILEStools.PROFILES_GACODE(self.folder / 'input.gacode.eped') + self.profiles_output = PROFILEStools.gacode_state(self.folder / 'input.gacode.eped') self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') @@ -307,7 +307,7 @@ def grab_output(self): loaded_results = np.load(self.folder_output / 'eped_results.npy', allow_pickle=True).item() - profiles = PROFILEStools.PROFILES_GACODE(self.folder_output / 'input.gacode') if isitfinished else None + profiles = PROFILEStools.gacode_state(self.folder_output / 'input.gacode') if isitfinished else None else: @@ -332,7 +332,7 @@ def plot(self, fn = None, counter = 0, full_plot = True): loaded_results, profiles = self.grab_output() - profiles_current = PROFILEStools.PROFILES_GACODE(self.folder / 'input.gacode') + profiles_current = PROFILEStools.gacode_state(self.folder / 'input.gacode') profiles_current.plotRelevant(axs = axs, color = 'b', label = 'orig') diff --git a/src/mitim_modules/maestro/utils/MAESTRObeat.py b/src/mitim_modules/maestro/utils/MAESTRObeat.py index f32e2830..a4fa0e9c 100644 --- a/src/mitim_modules/maestro/utils/MAESTRObeat.py +++ b/src/mitim_modules/maestro/utils/MAESTRObeat.py @@ -92,7 +92,7 @@ def __init__(self, beat_instance, label = 'profiles'): def __call__(self, profiles_file = None, Vsurf = None, **kwargs_beat): # Load profiles - self.profiles_current = PROFILEStools.PROFILES_GACODE(profiles_file) + self.profiles_current = PROFILEStools.gacode_state(profiles_file) # -------------------------------------------------------------------------------------------- # Operations diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index 708e396e..f97f2849 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -63,7 +63,7 @@ def plot_results(self, fn): # ******************************************************************************************************** # Collect initialization - ini = {'geqdsk': None, 'profiles': PROFILEStools.PROFILES_GACODE(f'{self.beats[1].initialize.folder}/input.gacode')} + ini = {'geqdsk': None, 'profiles': PROFILEStools.gacode_state(f'{self.beats[1].initialize.folder}/input.gacode')} if (self.beats[1].initialize.folder / 'input.geqdsk').exists(): ini['geqdsk'] = GEQtools.MITIMgeqdsk(self.beats[1].initialize.folder / 'input.geqdsk') @@ -90,7 +90,7 @@ def plot_results(self, fn): # ******************************************************************************************************** ps, ps_lab = [], [] for label in objs: - if isinstance(objs[label], PROFILEStools.PROFILES_GACODE): + if isinstance(objs[label], PROFILEStools.gacode_state): ps.append(objs[label]) ps_lab.append(label) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 03cc8d8c..0f80b716 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -241,7 +241,7 @@ def grab_output(self, full = False): opt_fun = STRATEGYtools.opt_evaluator(folder) if full else PORTALSanalysis.PORTALSanalyzer.from_folder(folder) - profiles = PROFILEStools.PROFILES_GACODE(self.folder_output / 'input.gacode') if isitfinished else None + profiles = PROFILEStools.gacode_state(self.folder_output / 'input.gacode') if isitfinished else None return opt_fun, profiles diff --git a/src/mitim_modules/maestro/utils/TRANSPbeat.py b/src/mitim_modules/maestro/utils/TRANSPbeat.py index 2635fc92..739e0ddc 100644 --- a/src/mitim_modules/maestro/utils/TRANSPbeat.py +++ b/src/mitim_modules/maestro/utils/TRANSPbeat.py @@ -245,7 +245,7 @@ def grab_output(self): if isitfinished: c = CDFtools.transp_output(self.folder_output) - profiles = PROFILEStools.PROFILES_GACODE(self.folder_output / 'input.gacode') + profiles = PROFILEStools.gacode_state(self.folder_output / 'input.gacode') else: # Trying to see if there's an intermediate CDF in folder print('\t\t- Searching for intermediate CDF in folder') diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index a75d9351..ea56f4c5 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -396,10 +396,10 @@ def run(self, paramsfile, resultsfile): if self.optimization_extra is not None: dictStore = IOtools.unpickle_mitim(self.optimization_extra) #TODO: This will fail in future versions of torch dictStore[int(numPORTALS)] = {"powerstate": powerstate} - dictStore["profiles_modified"] = PROFILEStools.PROFILES_GACODE( + dictStore["profiles_modified"] = PROFILEStools.gacode_state( self.folder / "Initialization" / "input.gacode_modified" ) - dictStore["profiles_original"] = PROFILEStools.PROFILES_GACODE( + dictStore["profiles_original"] = PROFILEStools.gacode_state( self.folder / "Initialization" / "input.gacode_original" ) with open(self.optimization_extra, "wb") as handle: diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 288b1d57..eab4fb4b 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -183,11 +183,11 @@ def prep_metrics(self, ilast=None): file = self.opt_fun.folder / "Execution" / f"Evaluation.{x_train_num}" / "model_complete" / "input.gacode_unmodified" if file.exists(): print("\t\t- Reading next profile to evaluate (from folder)") - self.profiles_next = PROFILEStools.PROFILES_GACODE(file, calculateDerived=False) + self.profiles_next = PROFILEStools.gacode_state(file, calculateDerived=False) file = self.opt_fun.folder / "Execution" / f"Evaluation.{x_train_num}" / "model_complete" / "input.gacode.new" if file.exists(): - self.profiles_next_new = PROFILEStools.PROFILES_GACODE( + self.profiles_next_new = PROFILEStools.gacode_state( file, calculateDerived=False ) self.profiles_next_new.printInfo(label="NEXT") @@ -951,7 +951,7 @@ def __init__(self, folder): self.profiles = [] for i in range(100): try: - prof = PROFILEStools.PROFILES_GACODE( + prof = PROFILEStools.gacode_state( self.folder / "Outputs" / "portals_profiles" / f"input.gacode.{i}" ) except FileNotFoundError: diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 9d8ea742..a10dd250 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -69,7 +69,7 @@ def initializeProblem( # ---- Initialize file to modify and increase resolution initialization_file = FolderInitialization / "input.gacode" - profiles = PROFILEStools.PROFILES_GACODE(initialization_file) + profiles = PROFILEStools.gacode_state(initialization_file) # About radial locations if portals_fun.MODELparameters["RoaLocations"] is not None: diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 446dcdcd..dcca4b3d 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -30,7 +30,7 @@ def __init__( ): ''' Inputs: - - profiles_object: Object for PROFILES_GACODE or others + - profiles_object: Object for gacode_state or others - EvolutionOptions: - rhoPredicted: radial grid (MUST NOT CONTAIN ZERO, it will be added internally) - ProfilesPredicted: list of profiles to predict @@ -134,7 +134,7 @@ def _ensure_ne_before_nz(lst): # Object type (e.g. input.gacode) # ------------------------------------------------------------------------------------- - if isinstance(profiles_object, PROFILEStools.PROFILES_GACODE): + if isinstance(profiles_object, PROFILEStools.gacode_state): self.to_powerstate = TRANSFORMtools.gacode_to_powerstate self.from_powerstate = MethodType(TRANSFORMtools.to_gacode, self) @@ -144,7 +144,7 @@ def _ensure_ne_before_nz(lst): self.profiles.derive_quantities() else: - raise ValueError("[MITIM] The input profile object is not recognized, please use PROFILES_GACODE") + raise ValueError("[MITIM] The input profile object is not recognized, please use gacode_state") # ------------------------------------------------------------------------------------- # Fine targets (need to do it here so that it's only once per definition of powerstate) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index c8580d1a..f7067722 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -172,7 +172,7 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod useConvectiveFluxes = PORTALSparameters["useConvectiveFluxes"] try: - impurityPosition = PROFILEStools.impurity_location(PROFILEStools.PROFILES_GACODE(unmodified_profiles), PORTALSparameters["ImpurityOfInterest"]) + impurityPosition = PROFILEStools.impurity_location(PROFILEStools.gacode_state(unmodified_profiles), PORTALSparameters["ImpurityOfInterest"]) except ValueError: if 'nZ' in ProfilesPredicted: raise ValueError(f"Impurity {PORTALSparameters['ImpurityOfInterest']} not found in the profiles and needed for CGYRO evaluation") @@ -246,7 +246,7 @@ def cgyroing( print(f"\t- Modifying {IOtools.clipstr(FolderEvaluation)} with position {k} in CGYRO results file {IOtools.clipstr(file)}") # Get TGYRO - tgyro = TGYROtools.TGYROoutput(FolderEvaluation,profiles=PROFILEStools.PROFILES_GACODE(unmodified_profiles)) + tgyro = TGYROtools.TGYROoutput(FolderEvaluation,profiles=PROFILEStools.gacode_state(unmodified_profiles)) # Quick checker of correct file wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro) diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 046f8cc3..3cf769de 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -21,7 +21,7 @@ def calculator( fineTargetsResolution = None, ): profiles = ( - input_gacode if profProvided else PROFILEStools.PROFILES_GACODE(input_gacode) + input_gacode if profProvided else PROFILEStools.gacode_state(input_gacode) ) # Calculate using TGYRO diff --git a/src/mitim_modules/powertorch/scripts/compareRadialResolution.py b/src/mitim_modules/powertorch/scripts/compareRadialResolution.py index bbdeea21..d5a6646f 100644 --- a/src/mitim_modules/powertorch/scripts/compareRadialResolution.py +++ b/src/mitim_modules/powertorch/scripts/compareRadialResolution.py @@ -28,7 +28,7 @@ inputgacode = IOtools.expandPath(args.file) rho = np.array([float(i) for i in args.rhos]) -profiles = PROFILEStools.PROFILES_GACODE(inputgacode) +profiles = PROFILEStools.gacode_state(inputgacode) markersize_coarse = 6 markersize_fine = 3 diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index aeea2304..60398143 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -248,7 +248,7 @@ def powerstate_to_gacode( ): """ Notes: - - This function assumes that "profiles" is the PROFILES_GACODE that everything started with. + - This function assumes that "profiles" is the gacode_state that everything started with. - We assume that what changes is only the kinetic profiles allowed to vary. - This only works for a single profile, in position_in_powerstate_batch - rederive is expensive, so I'm not re-deriving the geometry which is the most expensive diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index a9461805..dd94010b 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -107,8 +107,8 @@ def _modify_profiles(self): self.powerstate.profiles_transport = profiles_postprocessing_fun(self.file_profs) # Position of impurity ion may have changed - p_old = PROFILEStools.PROFILES_GACODE(self.file_profs_unmod) - p_new = PROFILEStools.PROFILES_GACODE(self.file_profs) + p_old = PROFILEStools.gacode_state(self.file_profs_unmod) + p_new = PROFILEStools.gacode_state(self.file_profs) impurity_of_interest = p_old.Species[self.powerstate.impurityPosition] diff --git a/src/mitim_tools/astra_tools/ASTRAtools.py b/src/mitim_tools/astra_tools/ASTRAtools.py index 3795d92f..ccfbfba1 100644 --- a/src/mitim_tools/astra_tools/ASTRAtools.py +++ b/src/mitim_tools/astra_tools/ASTRAtools.py @@ -158,7 +158,7 @@ def convert_ASTRA_to_gacode_fromCDF(astra_cdf, """ template_path = __mitimroot__ / "tests" / "data"/ "input.gacode" - p = PROFILEStools.PROFILES_GACODE(template_path) + p = PROFILEStools.gacode_state(template_path) params = p.profiles # Extract CDF file @@ -204,7 +204,7 @@ def convert_ASTRA_to_gacode_from_transp_output(c, """ template_path = __mitimroot__ / "tests" / "data"/ "input.gacode" - p = PROFILEStools.PROFILES_GACODE(template_path) + p = PROFILEStools.gacode_state(template_path) params = p.profiles #c.calcProfiles() diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 8a407662..c5f1d15a 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -67,7 +67,7 @@ def run_vgen(self, subfolder="vgen1", vgenOptions={}, cold_start=False): # ---- Postprocess from mitim_tools.gacode_tools import PROFILEStools - self.inputgacode_vgen = PROFILEStools.PROFILES_GACODE( + self.inputgacode_vgen = PROFILEStools.gacode_state( file_new, calculateDerived=True, mi_ref=self.inputgacode.mi_ref ) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index ffdf273d..455ccf3f 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -7,7 +7,7 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -class PROFILES_GACODE(mitim_state): +class gacode_state(mitim_state): ''' Class to read and manipulate GACODE profiles files (input.gacode). It inherits from the main MITIMstate class, which provides basic diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 6f7cb8c2..4604b1e2 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -235,7 +235,7 @@ def prep( from mitim_tools.gacode_tools import PROFILEStools profiles = ( - PROFILEStools.PROFILES_GACODE(inputgacode) + PROFILEStools.gacode_state(inputgacode) if inputgacode is not None else None ) @@ -375,7 +375,7 @@ def prep_direct_tglf( from mitim_tools.gacode_tools import PROFILEStools self.profiles = ( - PROFILEStools.PROFILES_GACODE(inputgacode) + PROFILEStools.gacode_state(inputgacode) if inputgacode is not None else None ) @@ -480,7 +480,7 @@ def prep_from_tglf( # Main folder where things are from mitim_tools.gacode_tools import PROFILEStools self.NormalizationSets, _ = NORMtools.normalizations( - PROFILEStools.PROFILES_GACODE(input_gacode) + PROFILEStools.gacode_state(input_gacode) if input_gacode is not None else None ) diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index b70068de..d2a9a152 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -141,7 +141,7 @@ def prep( self.file_input_profiles = self.FolderGACODE / "input.gacode" from mitim_tools.gacode_tools import PROFILEStools - self.profiles = PROFILEStools.PROFILES_GACODE(self.file_input_profiles) + self.profiles = PROFILEStools.gacode_state(self.file_input_profiles) if correctPROFILES: self.profiles.correct(write=True) @@ -463,7 +463,7 @@ def run( print("\t- It was requested that input.gacode.new is modified according to what TypeTarget was",typeMsg="i",) from mitim_tools.gacode_tools import PROFILEStools - inputgacode_new = PROFILEStools.PROFILES_GACODE(self.FolderTGYRO_tmp / "input.gacode.new") + inputgacode_new = PROFILEStools.gacode_state(self.FolderTGYRO_tmp / "input.gacode.new") if TGYRO_physics_options["TypeTarget"] < 3: for ikey in [ @@ -538,7 +538,7 @@ def read(self, label="tgyro1", folder=None, file_input_profiles=None): prof = self.profiles else: from mitim_tools.gacode_tools import PROFILEStools - prof = PROFILEStools.PROFILES_GACODE(file_input_profiles) + prof = PROFILEStools.gacode_state(file_input_profiles) self.results[label] = TGYROoutput(folder, profiles=prof) @@ -1203,7 +1203,7 @@ def __init__(self, FolderTGYRO, profiles=None): if (profiles is None) and (FolderTGYRO / "input.gacode").exists(): from mitim_tools.gacode_tools import PROFILEStools - profiles = PROFILEStools.PROFILES_GACODE(FolderTGYRO / f"input.gacode", calculateDerived=False) + profiles = PROFILEStools.gacode_state(FolderTGYRO / f"input.gacode", calculateDerived=False) self.profiles = profiles @@ -1217,7 +1217,7 @@ def __init__(self, FolderTGYRO, profiles=None): calculateDerived = True try: - self.profiles_final = PROFILEStools.PROFILES_GACODE(self.FolderTGYRO / "input.gacode.new",calculateDerived=calculateDerived,) + self.profiles_final = PROFILEStools.gacode_state(self.FolderTGYRO / "input.gacode.new",calculateDerived=calculateDerived,) except: self.profiles_final = None diff --git a/src/mitim_tools/gacode_tools/scripts/compare_MXH3.py b/src/mitim_tools/gacode_tools/scripts/compare_MXH3.py index 227d45fa..26a01adb 100644 --- a/src/mitim_tools/gacode_tools/scripts/compare_MXH3.py +++ b/src/mitim_tools/gacode_tools/scripts/compare_MXH3.py @@ -12,7 +12,7 @@ """ file_input_gacode = sys.argv[1] -p = PROFILEStools.PROFILES_GACODE(file_input_gacode) +p = PROFILEStools.gacode_state(file_input_gacode) file_geq = sys.argv[2] g = GEQtools.MITIMgeqdsk(file_geq) diff --git a/src/mitim_tools/gacode_tools/scripts/read_gacode.py b/src/mitim_tools/gacode_tools/scripts/read_gacode.py index 778fe07d..c668a9ea 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_gacode.py +++ b/src/mitim_tools/gacode_tools/scripts/read_gacode.py @@ -23,7 +23,7 @@ def main(): # Read profs = [] for file in files: - p = PROFILEStools.PROFILES_GACODE(file) + p = PROFILEStools.gacode_state(file) profs.append(p) p.printInfo() diff --git a/src/mitim_tools/gacode_tools/scripts/read_tgyro.py b/src/mitim_tools/gacode_tools/scripts/read_tgyro.py index 679b050f..25856a60 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_tgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_tgyro.py @@ -21,7 +21,7 @@ def main(): tgyros = [] for folder in folders: prof_file = folder / "input.gacode" - prof = PROFILEStools.PROFILES_GACODE(prof_file) + prof = PROFILEStools.gacode_state(prof_file) p = TGYROtools.TGYROoutput(folder, profiles=prof) tgyros.append(p) diff --git a/src/mitim_tools/popcon_tools/POPCONtools.py b/src/mitim_tools/popcon_tools/POPCONtools.py index 71c3426e..44f7d95c 100644 --- a/src/mitim_tools/popcon_tools/POPCONtools.py +++ b/src/mitim_tools/popcon_tools/POPCONtools.py @@ -21,42 +21,42 @@ def evaluate(self): self.results = self.algorithm.update_dataset(self.dataset) def update_from_gacode(self, - profiles_gacode: PROFILEStools.PROFILES_GACODE, + gacode_state: PROFILEStools.gacode_state, confinement_type="H98" ): - self.dataset['major_radius'].data = profiles_gacode.profiles['rcentr(m)'][-1] * ureg.meter + self.dataset['major_radius'].data = gacode_state.profiles['rcentr(m)'][-1] * ureg.meter - self.dataset["magnetic_field_on_axis"].data = np.abs(profiles_gacode.derived["B0"]) * ureg.tesla + self.dataset["magnetic_field_on_axis"].data = np.abs(gacode_state.derived["B0"]) * ureg.tesla - rmin = profiles_gacode.profiles["rmin(m)"][-1] - rmaj = profiles_gacode.profiles["rmaj(m)"][-1] + rmin = gacode_state.profiles["rmin(m)"][-1] + rmaj = gacode_state.profiles["rmaj(m)"][-1] self.dataset["inverse_aspect_ratio"].data = (rmin / rmaj) * ureg.dimensionless - kappa_a = 1.5 #profiles_gacode.derived["kappa_a"] - kappa_sep = profiles_gacode.profiles["kappa(-)"][-1] + kappa_a = 1.5 #gacode_state.derived["kappa_a"] + kappa_sep = gacode_state.profiles["kappa(-)"][-1] self.dataset["areal_elongation"].data = kappa_a * ureg.dimensionless self.dataset["elongation_ratio_sep_to_areal"].data = (kappa_sep / kappa_a) * ureg.dimensionless - delta_95 = np.interp(0.95, profiles_gacode.derived["psi_pol_n"], profiles_gacode.profiles["delta(-)"]) - delta_sep = profiles_gacode.profiles["delta(-)"][-1] + delta_95 = np.interp(0.95, gacode_state.derived["psi_pol_n"], gacode_state.profiles["delta(-)"]) + delta_sep = gacode_state.profiles["delta(-)"][-1] self.dataset["triangularity_psi95"].data = delta_95 * ureg.dimensionless self.dataset["triangularity_ratio_sep_to_psi95"].data = (delta_sep / delta_95) * ureg.dimensionless - ip = np.abs(profiles_gacode.profiles["current(MA)"][-1]) * 1e6 + ip = np.abs(gacode_state.profiles["current(MA)"][-1]) * 1e6 self.dataset["plasma_current"].data = ip * ureg.ampere - ne_vol_19 = profiles_gacode.derived["ne_vol20"] * 10 + ne_vol_19 = gacode_state.derived["ne_vol20"] * 10 self.dataset = self.dataset.assign_coords(dim_average_electron_density=np.array([ne_vol_19])) self.dataset["average_electron_density"].data = np.array([ne_vol_19]) * ureg._1e19_per_cubic_metre - #self.dataset["nesep_over_nebar"].data = profiles_gacode.profiles["ne(10^19/m^3)"][-1] / ne_vol_19 + #self.dataset["nesep_over_nebar"].data = gacode_state.profiles["ne(10^19/m^3)"][-1] / ne_vol_19 - te_vol_keV = profiles_gacode.derived["Te_vol"] + te_vol_keV = gacode_state.derived["Te_vol"] self.dataset = self.dataset.assign_coords(dim_average_electron_temp=np.array([te_vol_keV])) self.dataset["average_electron_temp"].data = np.array([te_vol_keV]) * ureg.kiloelectron_volt - impurity_zs = profiles_gacode.profiles["z"][np.where(profiles_gacode.profiles["z"] > 1)] - imputity_fs = profiles_gacode.derived["fi_vol"][np.where(profiles_gacode.profiles["z"] > 1)] + impurity_zs = gacode_state.profiles["z"][np.where(gacode_state.profiles["z"] > 1)] + imputity_fs = gacode_state.derived["fi_vol"][np.where(gacode_state.profiles["z"] > 1)] impurities = [] concentrations = [] named_options_array = [1,2,3,4,6,7,8,10,18,36,54,74] # atomicspecies built into cfspopcon @@ -79,33 +79,33 @@ def update_from_gacode(self, #from .formulas.impurities.impurity_array_helpers import make_impurity_concentration_array self.dataset['impurities'] = cfspopcon.formulas.impurities.impurity_array_helpers.make_impurity_concentration_array(impurities, concentrations) - arg_min_rho = np.argmin(np.abs(profiles_gacode.profiles["rho(-)"] - 0.4)) - arg_max_rho = np.argmin(np.abs(profiles_gacode.profiles["rho(-)"] - 0.8)) + arg_min_rho = np.argmin(np.abs(gacode_state.profiles["rho(-)"] - 0.4)) + arg_max_rho = np.argmin(np.abs(gacode_state.profiles["rho(-)"] - 0.8)) # calculate the predicted density peaking using the Angioni 2007 scaling - beta_percent = (profiles_gacode.derived["BetaN"] - *profiles_gacode.profiles["current(MA)"][-1] - / profiles_gacode.profiles["rmin(m)"][-1] - / profiles_gacode.profiles['bcentr(T)'][-1]) + beta_percent = (gacode_state.derived["BetaN"] + *gacode_state.profiles["current(MA)"][-1] + / gacode_state.profiles["rmin(m)"][-1] + / gacode_state.profiles['bcentr(T)'][-1]) - nu_n_scaling = cfspopcon.formulas.plasma_profiles.calc_density_peaking(profiles_gacode.derived["nu_eff"], + nu_n_scaling = cfspopcon.formulas.plasma_profiles.calc_density_peaking(gacode_state.derived["nu_eff"], beta_percent*1e-2, nu_noffset=0.0) - aLTe = profiles_gacode.derived["aLTe"][arg_min_rho:arg_max_rho].mean() + aLTe = gacode_state.derived["aLTe"][arg_min_rho:arg_max_rho].mean() self.dataset["normalized_inverse_temp_scale_length"].data = aLTe * ureg.dimensionless - nu_ne_offset = (nu_n_scaling - profiles_gacode.derived["ne_peaking"]) + nu_ne_offset = (nu_n_scaling - gacode_state.derived["ne_peaking"]) self.dataset["electron_density_peaking_offset"].data = nu_ne_offset * ureg.dimensionless self.dataset["ion_density_peaking_offset"].data = nu_ne_offset * ureg.dimensionless - nu_te = profiles_gacode.derived["Te_peaking"] + nu_te = gacode_state.derived["Te_peaking"] self.dataset["temperature_peaking"].data = nu_te * ureg.dimensionless - confinement_scalar = profiles_gacode.derived[confinement_type] + confinement_scalar = gacode_state.derived[confinement_type] self.dataset["confinement_time_scalar"].data = confinement_scalar * ureg.dimensionless - ti_over_te = profiles_gacode.derived["tite_vol"] + ti_over_te = gacode_state.derived["tite_vol"] self.dataset["ion_to_electron_temp_ratio"].data = ti_over_te * ureg.dimensionless def update_transport(self, @@ -124,7 +124,7 @@ def update_transport(self, self.dataset["ion_to_electron_temp_ratio"].data = ti_over_te * ureg.dimensionless def match_to_gacode(self, - profiles_gacode: PROFILEStools.PROFILES_GACODE, + gacode_state: PROFILEStools.gacode_state, confinement_type="H98", print_progress=False, plot_convergence=True, @@ -139,7 +139,7 @@ def match_to_gacode(self, print("Starting optimization...") print("... Initializing POPCON with GACODE parameters") - self.update_from_gacode(profiles_gacode=profiles_gacode, + self.update_from_gacode(gacode_state=gacode_state, confinement_type=confinement_type ) @@ -163,7 +163,7 @@ def match_to_gacode(self, res = minimize(self.match_pfus, x0, - args=(profiles_gacode, print_progress), + args=(gacode_state, print_progress), method='Nelder-Mead', bounds=bounds, options={'disp': True}, @@ -179,7 +179,7 @@ def match_to_gacode(self, def match_pfus(self, x, - profiles_gacode, + gacode_state, print_progress=False ): @@ -194,29 +194,29 @@ def match_pfus(self, point = self.results.isel(dim_average_electron_temp=0, dim_average_electron_density=0) - Pfus_residual = (point['P_fusion'].data.magnitude - profiles_gacode.derived['Pfus']) / profiles_gacode.derived['Pfus'] + Pfus_residual = (point['P_fusion'].data.magnitude - gacode_state.derived['Pfus']) / gacode_state.derived['Pfus'] #Psol_residual = ((point['P_LH_thresh'].data.magnitude #* point['ratio_of_P_SOL_to_P_LH'].data.magnitude) - #- profiles_gacode.derived['Psol']) / profiles_gacode.derived['Psol'] + #- gacode_state.derived['Psol']) / gacode_state.derived['Psol'] - Prad_residual = (point['P_radiation'].data.magnitude - profiles_gacode.derived['Prad']) / profiles_gacode.derived['Prad'] + Prad_residual = (point['P_radiation'].data.magnitude - gacode_state.derived['Prad']) / gacode_state.derived['Prad'] - Pin_derived = (profiles_gacode.derived['qi_aux_MW'][-1] - +profiles_gacode.derived['qe_aux_MW'][-1] + Pin_derived = (gacode_state.derived['qi_aux_MW'][-1] + +gacode_state.derived['qe_aux_MW'][-1] ) Pin_residual = (point['P_auxillary_launched'].data.magnitude - Pin_derived) / Pin_derived - te_profiles = np.interp(point["dim_rho"], point["dim_rho"].size*profiles_gacode.profiles["rho(-)"] ,profiles_gacode.profiles["te(keV)"]) + te_profiles = np.interp(point["dim_rho"], point["dim_rho"].size*gacode_state.profiles["rho(-)"] ,gacode_state.profiles["te(keV)"]) te_popcon = point["electron_temp_profile"].data.magnitude te_L2 = np.sum(((te_profiles-te_popcon)/te_profiles)**2) - ti_profiles = np.interp(point["dim_rho"], point["dim_rho"].size*profiles_gacode.profiles["rho(-)"] ,profiles_gacode.profiles["ti(keV)"][:, 0]) + ti_profiles = np.interp(point["dim_rho"], point["dim_rho"].size*gacode_state.profiles["rho(-)"] ,gacode_state.profiles["ti(keV)"][:, 0]) ti_popcon = point["ion_temp_profile"].data.magnitude ti_L2 = np.sum(((ti_profiles-ti_popcon)/ti_profiles)**2) - ne19_profiles = np.interp(point["dim_rho"], point["dim_rho"].size*profiles_gacode.profiles["rho(-)"] ,profiles_gacode.profiles["ne(10^19/m^3)"]) + ne19_profiles = np.interp(point["dim_rho"], point["dim_rho"].size*gacode_state.profiles["rho(-)"] ,gacode_state.profiles["ne(10^19/m^3)"]) ne19_popcon = point["electron_density_profile"].data.magnitude ne19_L2 = np.sum(((ne19_profiles-ne19_popcon)/ne19_profiles)**2) @@ -239,7 +239,7 @@ def match_pfus(self, def match_radiation(self, x, - profiles_gacode + gacode_state ): radiated_power_scalar = x[0] @@ -252,7 +252,7 @@ def match_radiation(self, point = self.results.isel(dim_average_electron_temp=0, dim_average_electron_density=0) - Prad_residual = (point['P_radiation'].data.magnitude - profiles_gacode.derived['Prad']) / profiles_gacode.derived['Prad'] + Prad_residual = (point['P_radiation'].data.magnitude - gacode_state.derived['Prad']) / gacode_state.derived['Prad'] return Prad_residual**2 @@ -292,7 +292,7 @@ def plot_convergence(self): plt.tight_layout() plt.show() - def plot_profile_comparison(self, profiles_gacode: PROFILEStools.PROFILES_GACODE): + def plot_profile_comparison(self, gacode_state: PROFILEStools.gacode_state): point = self.results.isel(dim_average_electron_temp=0, dim_average_electron_density=0) @@ -302,18 +302,18 @@ def plot_profile_comparison(self, profiles_gacode: PROFILEStools.PROFILES_GACODE x = np.linspace(0,1, point["dim_rho"].size) # normalized toroidal flux ax.plot(x, - np.interp(point["dim_rho"], point["dim_rho"].size*profiles_gacode.profiles["rho(-)"] ,profiles_gacode.profiles["te(keV)"]), + np.interp(point["dim_rho"], point["dim_rho"].size*gacode_state.profiles["rho(-)"] ,gacode_state.profiles["te(keV)"]), label="Te (GACODE)",color='tab:red') ax.plot(x,point["electron_temp_profile"], label="Te (POPCON)",ls="--",color='tab:red') ax.plot(x, - np.interp(point["dim_rho"], point["dim_rho"].size*profiles_gacode.profiles["rho(-)"] ,profiles_gacode.profiles["ti(keV)"][:, 0]), + np.interp(point["dim_rho"], point["dim_rho"].size*gacode_state.profiles["rho(-)"] ,gacode_state.profiles["ti(keV)"][:, 0]), label="Ti (GACODE)",color='tab:purple') ax.plot(x,point["ion_temp_profile"], label="Ti (POPCON)",ls="--",color='tab:purple') ax2.plot(x, - np.interp(point["dim_rho"], point["dim_rho"].size*profiles_gacode.profiles["rho(-)"] ,profiles_gacode.profiles["ne(10^19/m^3)"]), + np.interp(point["dim_rho"], point["dim_rho"].size*gacode_state.profiles["rho(-)"] ,gacode_state.profiles["ne(10^19/m^3)"]), label="ne (GACODE)",color='tab:blue') ax2.plot(x,point["electron_density_profile"], label="ne (POPCON)",ls="--",color='tab:blue') @@ -405,7 +405,7 @@ def plot(self, def print_data(self, compare_to_gacode=False, - profiles_gacode=None + gacode_state=None ): try: @@ -418,26 +418,26 @@ def print_data(self, print(f"Operational point: ={point['average_electron_density'].data.magnitude}, ={point['average_electron_temp'].data.magnitude}") if compare_to_gacode: - if profiles_gacode is None: + if gacode_state is None: raise ValueError("No GACODE profiles passed to compare to.") - print(f"Pfus: ", f"POPCON: {point['P_fusion'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['Pfus']:.2f}", "(MW)") - print(f"Q: ", f"POPCON: {point['Q'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['Q']:.2f}") - print(f"TauE: ", f"POPCON: {point['energy_confinement_time'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['tauE']:.2f}", "(s)") - print(f"Beta_N:", f"POPCON: {point['normalized_beta'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['BetaN']:.2f}") + print(f"Pfus: ", f"POPCON: {point['P_fusion'].data.magnitude:.2f}", f"GACODE: {gacode_state.derived['Pfus']:.2f}", "(MW)") + print(f"Q: ", f"POPCON: {point['Q'].data.magnitude:.2f}", f"GACODE: {gacode_state.derived['Q']:.2f}") + print(f"TauE: ", f"POPCON: {point['energy_confinement_time'].data.magnitude:.2f}", f"GACODE: {gacode_state.derived['tauE']:.2f}", "(s)") + print(f"Beta_N:", f"POPCON: {point['normalized_beta'].data.magnitude:.2f}", f"GACODE: {gacode_state.derived['BetaN']:.2f}") print(f"P_sol: ", f"POPCON: {(point['P_LH_thresh'].data.magnitude *point['ratio_of_P_SOL_to_P_LH'].data.magnitude):.2f}", - f"GACODE: {profiles_gacode.derived['Psol']:.2f}","(MW)", f"(~%{point['ratio_of_P_SOL_to_P_LH'].data.magnitude*1e2:.2f} of LH threshold)") + f"GACODE: {gacode_state.derived['Psol']:.2f}","(MW)", f"(~%{point['ratio_of_P_SOL_to_P_LH'].data.magnitude*1e2:.2f} of LH threshold)") print(f"P_aux: ", f"POPCON: {point['P_auxillary_launched'].data.magnitude:.2f}", - f"GACODE: {(profiles_gacode.derived['qi_aux_MW'][-1]+profiles_gacode.derived['qe_aux_MW'][-1]):.2f}", + f"GACODE: {(gacode_state.derived['qi_aux_MW'][-1]+gacode_state.derived['qe_aux_MW'][-1]):.2f}", "(MW)") - print(f"P_rad: ", f"POPCON: {point['P_radiation'].data.magnitude:.2f}",f"GACODE: {profiles_gacode.derived['Prad']:.2f}","(MW)") + print(f"P_rad: ", f"POPCON: {point['P_radiation'].data.magnitude:.2f}",f"GACODE: {gacode_state.derived['Prad']:.2f}","(MW)") print(f"P_ext: ", f"POPCON: {point['P_external'].data.magnitude:.2f}","(MW)") - print(f"P_ohm: ", f"POPCON: {point['P_ohmic'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['qOhm_MW'][-1]:.2f}","(MW)") + print(f"P_ohm: ", f"POPCON: {point['P_ohmic'].data.magnitude:.2f}", f"GACODE: {gacode_state.derived['qOhm_MW'][-1]:.2f}","(MW)") print(f"P_in: ", f"POPCON: {point['P_in'].data.magnitude:.2f}", - f"GACODE: {(profiles_gacode.derived['qOhm_MW'][-1]+profiles_gacode.derived['qi_aux_MW'][-1]+profiles_gacode.derived['qe_aux_MW'][-1]+profiles_gacode.derived['Pfus']*0.2):.2f}", + f"GACODE: {(gacode_state.derived['qOhm_MW'][-1]+gacode_state.derived['qi_aux_MW'][-1]+gacode_state.derived['qe_aux_MW'][-1]+gacode_state.derived['Pfus']*0.2):.2f}", "(MW)") - print(f"q95: ", f"POPCON: {point['q_star'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['q95']:.2f}") - print(f"Wtot: ", f"POPCON: {point['plasma_stored_energy'].data.magnitude:.2f}", f"GACODE: {profiles_gacode.derived['Wthr']:.2f}", "(MJ)") + print(f"q95: ", f"POPCON: {point['q_star'].data.magnitude:.2f}", f"GACODE: {gacode_state.derived['q95']:.2f}") + print(f"Wtot: ", f"POPCON: {point['plasma_stored_energy'].data.magnitude:.2f}", f"GACODE: {gacode_state.derived['Wthr']:.2f}", "(MJ)") print(" ") print("Transport Parameters:") diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index a5500f66..0502824b 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -15426,7 +15426,7 @@ def grid_interpolation_method_to_zero(x,y): profiles[key] = profiles[key].clip(min=minimum) from mitim_tools.gacode_tools import PROFILEStools - p = PROFILEStools.PROFILES_GACODE.scratch(profiles) + p = PROFILEStools.gacode_state.scratch(profiles) return p diff --git a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py index 6a331280..afd675af 100644 --- a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py +++ b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py @@ -761,7 +761,7 @@ def from_profiles(self, time, profiles_file, Vsurf = 0.0): if isinstance(profiles_file, str): from mitim_tools.gacode_tools import PROFILEStools - self.p = PROFILEStools.PROFILES_GACODE(profiles_file) + self.p = PROFILEStools.gacode_state(profiles_file) else: self.p = profiles_file diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index 421d3251..e6b3a6fe 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -53,7 +53,7 @@ # For fun and to show capabilities, let's do a flux match of the current surrogates and plot in the same notebook PORTALSoptimization.flux_match_surrogate( - mitim_bo.steps[-1],PROFILEStools.PROFILES_GACODE(inputgacode), + mitim_bo.steps[-1],PROFILEStools.gacode_state(inputgacode), fn = portals_fun.fn, plot_results = True, keep_within_bounds = False diff --git a/tests/POWERTORCH_workflow.py b/tests/POWERTORCH_workflow.py index 6175176f..d33b34e2 100644 --- a/tests/POWERTORCH_workflow.py +++ b/tests/POWERTORCH_workflow.py @@ -7,7 +7,7 @@ from mitim_tools import __mitimroot__ # Inputs -inputgacode = PROFILEStools.PROFILES_GACODE(__mitimroot__ / "tests" / "data" / "input.gacode") +inputgacode = PROFILEStools.gacode_state(__mitimroot__ / "tests" / "data" / "input.gacode") rho = torch.from_numpy(np.linspace(0.1,0.9,9)).to(dtype=torch.double) s = STATEtools.powerstate(inputgacode, diff --git a/tests/TGYRO_workflow.py b/tests/TGYRO_workflow.py index 2f6f2cd7..b2ec7557 100644 --- a/tests/TGYRO_workflow.py +++ b/tests/TGYRO_workflow.py @@ -18,7 +18,7 @@ if cold_start and folder.exists(): os.system(f"rm -r {folder.resolve()}") -profiles = PROFILEStools.PROFILES_GACODE(gacode_file) +profiles = PROFILEStools.gacode_state(gacode_file) tgyro = TGYROtools.TGYRO() tgyro.prep(folder, profilesclass_custom=profiles, cold_start=True, forceIfcold_start=True) From a2b873fad04bac89930fd2fcb07c500d765927d0 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 16:23:13 +0200 Subject: [PATCH 084/385] Removed xarray requirement from [omfit] since that's not needed with PIXI, but carefuly with simple PIP --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 04b1e190..5ffd4e92 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,7 +55,7 @@ omfit = [ "omfit_classes>3.2024.19.2", # Otherwise, it will need an old version of matplotlib, matplotlib<3.6 "scipy<1.14.0", # As of 08/08/2024, because of https://github.com/gafusion/OMFIT-source/issues/7104 "numpy<2.0.0", # For the xarray requirement below to work - "xarray==2022.6.0", # As of 08/08/2024, because https://github.com/gafusion/OMFIT-source/issues/7104. Froze this version because of the PLASMAstate xr reader (importlib_metadata issues) + # "xarray==2022.6.0", # As of 08/08/2024, because https://github.com/gafusion/OMFIT-source/issues/7104. Froze this version because of the PLASMAstate xr reader (importlib_metadata issues) "omas", "fortranformat", "openpyxl", From 077d6ddeea2ef432ef65e301c9fc93325e0f4c66 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 16:24:14 +0200 Subject: [PATCH 085/385] renamed NTCC plasmastate --- src/mitim_tools/gacode_tools/utils/GACODErun.py | 4 ++-- .../transp_tools/utils/{PLASMASTATEtools.py => NTCCtools.py} | 0 2 files changed, 2 insertions(+), 2 deletions(-) rename src/mitim_tools/transp_tools/utils/{PLASMASTATEtools.py => NTCCtools.py} (100%) diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 223d247a..34a6e908 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -4,7 +4,7 @@ import matplotlib.pyplot as plt from scipy.interpolate import interp1d from mitim_tools.gacode_tools.utils import GACODEdefaults -from mitim_tools.transp_tools.utils import PLASMASTATEtools +from mitim_tools.transp_tools.utils import NTCCtools from mitim_tools.misc_tools import FARMINGtools, IOtools, MATHtools, GRAPHICStools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -463,7 +463,7 @@ def runPROFILES_GEN( if UseMITIMmodification: print("\t\t- Running modifyPlasmaState") shutil.copy2(FolderTGLF / f"{nameFiles}.cdf", FolderTGLF / f"{nameFiles}.cdf_old") - pls = PLASMASTATEtools.Plasmastate(FolderTGLF / f"{nameFiles}.cdf_old") + pls = NTCCtools.Plasmastate(FolderTGLF / f"{nameFiles}.cdf_old") pls.modify_default(FolderTGLF / f"{nameFiles}.cdf") inputFiles = [ diff --git a/src/mitim_tools/transp_tools/utils/PLASMASTATEtools.py b/src/mitim_tools/transp_tools/utils/NTCCtools.py similarity index 100% rename from src/mitim_tools/transp_tools/utils/PLASMASTATEtools.py rename to src/mitim_tools/transp_tools/utils/NTCCtools.py From c9d5744b14b7497d46aa0d2020223ed89b2ee871 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 16:37:39 +0200 Subject: [PATCH 086/385] PORTALS working fine --- src/mitim_modules/portals/PORTALStools.py | 2 +- .../portals/utils/PORTALSanalysis.py | 1 + .../portals/utils/PORTALSinit.py | 3 +- .../portals/utils/PORTALSplot.py | 3 +- src/mitim_tools/gacode_tools/PROFILEStools.py | 404 +++++++++++++++++- .../gacode_tools/utils/GEOMETRYtools.py | 401 ----------------- 6 files changed, 406 insertions(+), 408 deletions(-) delete mode 100644 src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index e97ed9f1..9668a31a 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -615,7 +615,7 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): elif prof == "nZ": var = "GZ" elif prof == "w0": - var = "MtJm2" + var = "Mt" """ ----------------------------------------------------------------------------------- diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index eab4fb4b..12b1e2a1 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -297,6 +297,7 @@ def prep_metrics(self, ilast=None): ] ) except: + embed() print("\t- Could not calculate Ricci metric", typeMsg="w") calculateRicci = None self.qR_Ricci, self.chiR_Ricci, self.points_Ricci = None, None, None diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index a10dd250..b436fbb8 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -6,6 +6,7 @@ from collections import OrderedDict from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.plasmastate_tools import MITIMstate from mitim_modules.powertorch import STATEtools from mitim_modules.portals import PORTALStools from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -90,7 +91,7 @@ def initializeProblem( profiles.correct(options=INITparameters) if portals_fun.PORTALSparameters["ImpurityOfInterest"] is not None: - position_of_impurity = PROFILEStools.impurity_location(profiles, portals_fun.PORTALSparameters["ImpurityOfInterest"]) + position_of_impurity = MITIMstate.impurity_location(profiles, portals_fun.PORTALSparameters["ImpurityOfInterest"]) else: position_of_impurity = 1 diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index e82ea97b..728c8639 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -3,7 +3,6 @@ import numpy as np import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools -from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.portals import PORTALStools from mitim_modules.powertorch import STATEtools from mitim_modules.powertorch.utils import POWERplot @@ -1827,7 +1826,7 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): # Plot PROFILES # ------------------------------------------------------- - figs = PROFILEStools.add_figures(fn,fnlab_pre = "PROFILES - ") + figs = MITIMstate.add_figures(fn,fnlab_pre = "PROFILES - ") if indecesPlot[0] < len(self.powerstates): _ = MITIMstate.plotAll( diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 455ccf3f..ab9a64d4 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -3,7 +3,7 @@ from collections import OrderedDict from mitim_tools.plasmastate_tools.MITIMstate import mitim_state from mitim_tools.gs_tools import GEQtools -from mitim_tools.gacode_tools.utils import GEOMETRYtools +from mitim_tools.misc_tools import MATHtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -201,7 +201,7 @@ def derive_geometry(self, n_theta_geo=1001): self.derived["bp2_geo"], self.derived["bt2_geo"], self.derived["bt_geo"], - ) = GEOMETRYtools.calculateGeometricFactors(self,n_theta=n_theta_geo) + ) = calculateGeometricFactors(self,n_theta=n_theta_geo) # Calculate flux surfaces cn = np.array(self.shape_cos).T @@ -222,9 +222,407 @@ def derive_geometry(self, n_theta_geo=1001): # ----------------------------------------------- #cross-sectional area of each flux surface - self.derived["surfXS"] = GEOMETRYtools.xsec_area_RZ(self.derived["R_surface"],self.derived["Z_surface"]) + self.derived["surfXS"] = xsec_area_RZ(self.derived["R_surface"],self.derived["Z_surface"]) self.derived["R_LF"] = self.derived["R_surface"].max(axis=1) # self.profiles['rmaj(m)'][0]+self.profiles['rmin(m)'] # For Synchrotron self.derived["B_ref"] = np.abs(self.derived["B_unit"] * self.derived["bt_geo"]) + + + + +def calculateGeometricFactors(profiles, n_theta=1001): + + # ---------------------------------------- + # Raw parameters from the file + # in expro_util.f90, it performs those divisions to pass to geo library + # ---------------------------------------- + + r = profiles.profiles["rmin(m)"] / profiles.profiles["rmin(m)"][-1] + R = profiles.profiles["rmaj(m)"] / profiles.profiles["rmin(m)"][-1] + kappa = profiles.profiles["kappa(-)"] + delta = profiles.profiles["delta(-)"] + zeta = profiles.profiles["zeta(-)"] + zmag = profiles.profiles["zmag(m)"] / profiles.profiles["rmin(m)"][-1] + q = profiles.profiles["q(-)"] + + shape_coeffs = profiles.shape_cos + profiles.shape_sin + + # ---------------------------------------- + # Derivatives as defined in expro_util.f90 + # ---------------------------------------- + + s_delta = r * MATHtools.deriv(r, delta) + s_kappa = r / kappa * MATHtools.deriv(r, kappa) + s_zeta = r * MATHtools.deriv(r, zeta) + dzmag = MATHtools.deriv(r, zmag) + dRmag = MATHtools.deriv(r, R) + + s_shape_coeffs = [] + for i in range(len(shape_coeffs)): + if shape_coeffs[i] is not None: + s_shape_coeffs.append(r * MATHtools.deriv(r, shape_coeffs[i])) + else: + s_shape_coeffs.append(None) + + # ---------------------------------------- + # Calculate the differencial volume at each radii + # from f2py/geo/geo.f90 in gacode source we have geo_volume_prime. + # ---------------------------------------- + + # Prepare cos_sins + cos_sin = [] + cos_sin_s = [] + for j in range(len(R)): + cos_sin0 = [] + cos_sin_s0 = [] + for k in range(len(shape_coeffs)): + if shape_coeffs[k] is not None: + cos_sin0.append(shape_coeffs[k][j]) + cos_sin_s0.append(s_shape_coeffs[k][j]) + else: + cos_sin0.append(None) + cos_sin_s0.append(None) + cos_sin.append(cos_sin0) + cos_sin_s.append(cos_sin_s0) + + ( + geo_volume_prime, + geo_surf, + geo_fluxsurfave_grad_r, + geo_fluxsurfave_bp2, + geo_fluxsurfave_bt2, + bt_geo0, + ) = volp_surf_geo_vectorized( + R, + r, + delta, + kappa, + cos_sin, + cos_sin_s, + zeta, + zmag, + s_delta, + s_kappa, + s_zeta, + dzmag, + dRmag, + q, + n_theta=n_theta, + ) + + """ + from expro_util.f90 we have: + expro_volp(i) = geo_volume_prime*r_min**2, where r_min = expro_rmin(expro_n_exp) + expro_surf(i) = geo_surf*r_min**2 + """ + + volp = geo_volume_prime * profiles.profiles["rmin(m)"][-1] ** 2 + surf = geo_surf * profiles.profiles["rmin(m)"][-1] ** 2 + + return volp, surf, geo_fluxsurfave_grad_r, geo_fluxsurfave_bp2, geo_fluxsurfave_bt2, bt_geo0 + +def volp_surf_geo_vectorized( + geo_rmaj_in, + geo_rmin_in, + geo_delta_in, + geo_kappa_in, + cos_sin, + cos_sin_s, + geo_zeta_in, + geo_zmag_in, + geo_s_delta_in, + geo_s_kappa_in, + geo_s_zeta_in, + geo_dzmag_in, + geo_drmaj_in, + geo_q_in, + n_theta=1001): + """ + Completety from f2py/geo/geo.f90 + """ + + geo_rmin_in = geo_rmin_in.clip( + 1e-10 + ) # To avoid problems at 0 (Implemented by PRF, not sure how TGYRO deals with this) + + geo_q_in = geo_q_in.clip(1e-2) # To avoid problems at 0 with some geqdsk files that are corrupted... + + + [ + geo_shape_cos0_in, + geo_shape_cos1_in, + geo_shape_cos2_in, + geo_shape_cos3_in, + geo_shape_cos4_in, + geo_shape_cos5_in, + geo_shape_cos6_in, + _, + _, + _, + geo_shape_sin3_in, + geo_shape_sin4_in, + geo_shape_sin5_in, + geo_shape_sin6_in, + ] = np.array(cos_sin).astype(float).T + + [ + geo_shape_s_cos0_in, + geo_shape_s_cos1_in, + geo_shape_s_cos2_in, + geo_shape_s_cos3_in, + geo_shape_s_cos4_in, + geo_shape_s_cos5_in, + geo_shape_s_cos6_in, + _, + _, + _, + geo_shape_s_sin3_in, + geo_shape_s_sin4_in, + geo_shape_s_sin5_in, + geo_shape_s_sin6_in, + ] = np.array(cos_sin_s).astype(float).T + + geo_signb_in = 1.0 + + geov_theta = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_bigr = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_bigr_r = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_bigr_t = np.zeros((n_theta,geo_rmin_in.shape[0])) + bigz = np.zeros((n_theta,geo_rmin_in.shape[0])) + bigz_r = np.zeros((n_theta,geo_rmin_in.shape[0])) + bigz_t = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_jac_r = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_grad_r = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_l_t = np.zeros((n_theta,geo_rmin_in.shape[0])) + r_c = np.zeros((n_theta,geo_rmin_in.shape[0])) + bigz_l = np.zeros((n_theta,geo_rmin_in.shape[0])) + bigr_l = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_l_r = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_nsin = np.zeros((n_theta,geo_rmin_in.shape[0])) + + pi_2 = 8.0 * np.arctan(1.0) + d_theta = pi_2 / (n_theta - 1) + + for i in range(n_theta): + #!----------------------------------------- + #! Generalized Miller-type parameterization + #!----------------------------------------- + + theta = -0.5 * pi_2 + (i - 1) * d_theta + + geov_theta[i] = theta + + x = np.arcsin(geo_delta_in) + + #! A + #! dA/dtheta + #! d^2A/dtheta^2 + a = ( + theta + + geo_shape_cos0_in + + geo_shape_cos1_in * np.cos(theta) + + geo_shape_cos2_in * np.cos(2 * theta) + + geo_shape_cos3_in * np.cos(3 * theta) + + geo_shape_cos4_in * np.cos(4 * theta) + + geo_shape_cos5_in * np.cos(5 * theta) + + geo_shape_cos6_in * np.cos(6 * theta) + + geo_shape_sin3_in * np.sin(3 * theta) + + x * np.sin(theta) + - geo_zeta_in * np.sin(2 * theta) + + geo_shape_sin3_in * np.sin(3 * theta) + + geo_shape_sin4_in * np.sin(4 * theta) + + geo_shape_sin5_in * np.sin(5 * theta) + + geo_shape_sin6_in * np.sin(6 * theta) + ) + a_t = ( + 1.0 + - geo_shape_cos1_in * np.sin(theta) + - 2 * geo_shape_cos2_in * np.sin(2 * theta) + - 3 * geo_shape_cos3_in * np.sin(3 * theta) + - 4 * geo_shape_cos4_in * np.sin(4 * theta) + - 5 * geo_shape_cos5_in * np.sin(5 * theta) + - 6 * geo_shape_cos6_in * np.sin(6 * theta) + + x * np.cos(theta) + - 2 * geo_zeta_in * np.cos(2 * theta) + + 3 * geo_shape_sin3_in * np.cos(3 * theta) + + 4 * geo_shape_sin4_in * np.cos(4 * theta) + + 5 * geo_shape_sin5_in * np.cos(5 * theta) + + 6 * geo_shape_sin6_in * np.cos(6 * theta) + ) + a_tt = ( + -geo_shape_cos1_in * np.cos(theta) + - 4 * geo_shape_cos2_in * np.cos(2 * theta) + - 9 * geo_shape_cos3_in * np.cos(3 * theta) + - 16 * geo_shape_cos4_in * np.cos(4 * theta) + - 25 * geo_shape_cos5_in * np.cos(5 * theta) + - 36 * geo_shape_cos6_in * np.cos(6 * theta) + - x * np.sin(theta) + + 4 * geo_zeta_in * np.sin(2 * theta) + - 9 * geo_shape_sin3_in * np.sin(3 * theta) + - 16 * geo_shape_sin4_in * np.sin(4 * theta) + - 25 * geo_shape_sin5_in * np.sin(5 * theta) + - 36 * geo_shape_sin6_in * np.sin(6 * theta) + ) + + #! R(theta) + #! dR/dr + #! dR/dtheta + #! d^2R/dtheta^2 + geov_bigr[i] = geo_rmaj_in + geo_rmin_in * np.cos(a) + geov_bigr_r[i] = ( + geo_drmaj_in + + np.cos(a) + - np.sin(a) + * ( + geo_shape_s_cos0_in + + geo_shape_s_cos1_in * np.cos(theta) + + geo_shape_s_cos2_in * np.cos(2 * theta) + + geo_shape_s_cos3_in * np.cos(3 * theta) + + geo_shape_s_cos4_in * np.cos(4 * theta) + + geo_shape_s_cos5_in * np.cos(5 * theta) + + geo_shape_s_cos6_in * np.cos(6 * theta) + + geo_s_delta_in / np.cos(x) * np.sin(theta) + - geo_s_zeta_in * np.sin(2 * theta) + + geo_shape_s_sin3_in * np.sin(3 * theta) + + geo_shape_s_sin4_in * np.sin(4 * theta) + + geo_shape_s_sin5_in * np.sin(5 * theta) + + geo_shape_s_sin6_in * np.sin(6 * theta) + ) + ) + geov_bigr_t[i] = -geo_rmin_in * a_t * np.sin(a) + bigr_tt = -geo_rmin_in * a_t**2 * np.cos(a) - geo_rmin_in * a_tt * np.sin(a) + + #!----------------------------------------------------------- + + #! A + #! dA/dtheta + #! d^2A/dtheta^2 + a = theta + a_t = 1.0 + a_tt = 0.0 + + #! Z(theta) + #! dZ/dr + #! dZ/dtheta + #! d^2Z/dtheta^2 + bigz[i] = geo_zmag_in + geo_kappa_in * geo_rmin_in * np.sin(a) + bigz_r[i] = geo_dzmag_in + geo_kappa_in * (1.0 + geo_s_kappa_in) * np.sin(a) + bigz_t[i] = geo_kappa_in * geo_rmin_in * np.cos(a) * a_t + bigz_tt = ( + -geo_kappa_in * geo_rmin_in * np.sin(a) * a_t**2 + + geo_kappa_in * geo_rmin_in * np.cos(a) * a_tt + ) + + g_tt = geov_bigr_t[i] ** 2 + bigz_t[i] ** 2 + + geov_jac_r[i] = geov_bigr[i] * ( + geov_bigr_r[i] * bigz_t[i] - geov_bigr_t[i] * bigz_r[i] + ) + + geov_grad_r[i] = geov_bigr[i] * np.sqrt(g_tt) / geov_jac_r[i] + + geov_l_t[i] = np.sqrt(g_tt) + + r_c[i] = geov_l_t[i] ** 3 / (geov_bigr_t[i] * bigz_tt - bigz_t[i] * bigr_tt) + + bigz_l[i] = bigz_t[i] / geov_l_t[i] + + bigr_l[i] = geov_bigr_t[i] / geov_l_t[i] + + geov_l_r[i] = bigz_l[i] * bigz_r[i] + bigr_l[i] * geov_bigr_r[i] + + geov_nsin[i] = ( + geov_bigr_r[i] * geov_bigr_t[i] + bigz_r[i] * bigz_t[i] + ) / geov_l_t[i] + + c = 0.0 + for i in range(n_theta): + c = c + geov_l_t[i] / (geov_bigr[i] * geov_grad_r[i]) + + f = geo_rmin_in / (c * d_theta / pi_2) + + c = 0.0 + for i in range(n_theta - 1): + c = c + geov_l_t[i] * geov_bigr[i] / geov_grad_r[i] + + geo_volume_prime = pi_2 * c * d_theta + + # Line 716 in geo.f90 + geo_surf = 0.0 + for i in range(n_theta - 1): + geo_surf = geo_surf + geov_l_t[i] * geov_bigr[i] + geo_surf = pi_2 * geo_surf * d_theta + + # ----- + c = 0.0 + for i in range(n_theta - 1): + c = c + geov_l_t[i] / (geov_bigr[i] * geov_grad_r[i]) + f = geo_rmin_in / (c * d_theta / pi_2) + + geov_b = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_g_theta = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_bt = np.zeros((n_theta,geo_rmin_in.shape[0])) + for i in range(n_theta): + geov_bt[i] = f / geov_bigr[i] + geov_bp = (geo_rmin_in / geo_q_in) * geov_grad_r[i] / geov_bigr[i] + + geov_b[i] = geo_signb_in * (geov_bt[i] ** 2 + geov_bp**2) ** 0.5 + geov_g_theta[i] = ( + geov_bigr[i] + * geov_b[i] + * geov_l_t[i] + / (geo_rmin_in * geo_rmaj_in * geov_grad_r[i]) + ) + + theta_0 = 0 + dx = geov_theta[1,0] - geov_theta[0,0] + x0 = theta_0 - geov_theta[0,0] + i1 = int(x0 / dx) + 1 + i2 = i1 + 1 + x1 = (i1 - 1) * dx + z = (x0 - x1) / dx + if i2 == n_theta: + i2 -= 1 + bt_geo0 = geov_bt[i1] + (geov_bt[i2] - geov_bt[i1]) * z + + denom = 0 + for i in range(n_theta - 1): + denom = denom + geov_g_theta[i] / geov_b[i] + + geo_fluxsurfave_grad_r = 0 + for i in range(n_theta - 1): + geo_fluxsurfave_grad_r = ( + geo_fluxsurfave_grad_r + + geov_grad_r[i] * geov_g_theta[i] / geov_b[i] / denom + ) + + geo_fluxsurfave__bp2 = 0 + for i in range(n_theta - 1): + geo_fluxsurfave__bp2 = ( + geo_fluxsurfave__bp2 + + geov_bt[i] ** 2 * geov_g_theta[i] / geov_b[i] / denom + ) + + geo_fluxsurfave_bt2 = 0 + for i in range(n_theta - 1): + geo_fluxsurfave_bt2 = ( + geo_fluxsurfave_bt2 + + geov_bp ** 2 * geov_g_theta[i] / geov_b[i] / denom + ) + + return geo_volume_prime, geo_surf, geo_fluxsurfave_grad_r, geo_fluxsurfave__bp2, geo_fluxsurfave_bt2, bt_geo0 + +def xsec_area_RZ(R,Z): + # calculates the cross-sectional area of the plasma for each flux surface + xsec_area = [] + for i in range(R.shape[0]): + R0 = np.max(R[i,:]) - np.min(R[i,:]) + Z0 = np.max(Z[i,:]) - np.min(Z[i,:]) + xsec_area.append(np.trapz(R[i], Z[i])) + + xsec_area = np.array(xsec_area) + + return xsec_area + diff --git a/src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py b/src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py deleted file mode 100644 index dbccba56..00000000 --- a/src/mitim_tools/gacode_tools/utils/GEOMETRYtools.py +++ /dev/null @@ -1,401 +0,0 @@ -import numpy as np -from mitim_tools.misc_tools import MATHtools -from IPython import embed - - -def calculateGeometricFactors(profiles, n_theta=1001): - - # ---------------------------------------- - # Raw parameters from the file - # in expro_util.f90, it performs those divisions to pass to geo library - # ---------------------------------------- - - r = profiles.profiles["rmin(m)"] / profiles.profiles["rmin(m)"][-1] - R = profiles.profiles["rmaj(m)"] / profiles.profiles["rmin(m)"][-1] - kappa = profiles.profiles["kappa(-)"] - delta = profiles.profiles["delta(-)"] - zeta = profiles.profiles["zeta(-)"] - zmag = profiles.profiles["zmag(m)"] / profiles.profiles["rmin(m)"][-1] - q = profiles.profiles["q(-)"] - - shape_coeffs = profiles.shape_cos + profiles.shape_sin - - # ---------------------------------------- - # Derivatives as defined in expro_util.f90 - # ---------------------------------------- - - s_delta = r * MATHtools.deriv(r, delta) - s_kappa = r / kappa * MATHtools.deriv(r, kappa) - s_zeta = r * MATHtools.deriv(r, zeta) - dzmag = MATHtools.deriv(r, zmag) - dRmag = MATHtools.deriv(r, R) - - s_shape_coeffs = [] - for i in range(len(shape_coeffs)): - if shape_coeffs[i] is not None: - s_shape_coeffs.append(r * MATHtools.deriv(r, shape_coeffs[i])) - else: - s_shape_coeffs.append(None) - - # ---------------------------------------- - # Calculate the differencial volume at each radii - # from f2py/geo/geo.f90 in gacode source we have geo_volume_prime. - # ---------------------------------------- - - # Prepare cos_sins - cos_sin = [] - cos_sin_s = [] - for j in range(len(R)): - cos_sin0 = [] - cos_sin_s0 = [] - for k in range(len(shape_coeffs)): - if shape_coeffs[k] is not None: - cos_sin0.append(shape_coeffs[k][j]) - cos_sin_s0.append(s_shape_coeffs[k][j]) - else: - cos_sin0.append(None) - cos_sin_s0.append(None) - cos_sin.append(cos_sin0) - cos_sin_s.append(cos_sin_s0) - - ( - geo_volume_prime, - geo_surf, - geo_fluxsurfave_grad_r, - geo_fluxsurfave_bp2, - geo_fluxsurfave_bt2, - bt_geo0, - ) = volp_surf_geo_vectorized( - R, - r, - delta, - kappa, - cos_sin, - cos_sin_s, - zeta, - zmag, - s_delta, - s_kappa, - s_zeta, - dzmag, - dRmag, - q, - n_theta=n_theta, - ) - - """ - from expro_util.f90 we have: - expro_volp(i) = geo_volume_prime*r_min**2, where r_min = expro_rmin(expro_n_exp) - expro_surf(i) = geo_surf*r_min**2 - """ - - volp = geo_volume_prime * profiles.profiles["rmin(m)"][-1] ** 2 - surf = geo_surf * profiles.profiles["rmin(m)"][-1] ** 2 - - return volp, surf, geo_fluxsurfave_grad_r, geo_fluxsurfave_bp2, geo_fluxsurfave_bt2, bt_geo0 - -def volp_surf_geo_vectorized( - geo_rmaj_in, - geo_rmin_in, - geo_delta_in, - geo_kappa_in, - cos_sin, - cos_sin_s, - geo_zeta_in, - geo_zmag_in, - geo_s_delta_in, - geo_s_kappa_in, - geo_s_zeta_in, - geo_dzmag_in, - geo_drmaj_in, - geo_q_in, - n_theta=1001): - """ - Completety from f2py/geo/geo.f90 - """ - - geo_rmin_in = geo_rmin_in.clip( - 1e-10 - ) # To avoid problems at 0 (Implemented by PRF, not sure how TGYRO deals with this) - - geo_q_in = geo_q_in.clip(1e-2) # To avoid problems at 0 with some geqdsk files that are corrupted... - - - [ - geo_shape_cos0_in, - geo_shape_cos1_in, - geo_shape_cos2_in, - geo_shape_cos3_in, - geo_shape_cos4_in, - geo_shape_cos5_in, - geo_shape_cos6_in, - _, - _, - _, - geo_shape_sin3_in, - geo_shape_sin4_in, - geo_shape_sin5_in, - geo_shape_sin6_in, - ] = np.array(cos_sin).astype(float).T - - [ - geo_shape_s_cos0_in, - geo_shape_s_cos1_in, - geo_shape_s_cos2_in, - geo_shape_s_cos3_in, - geo_shape_s_cos4_in, - geo_shape_s_cos5_in, - geo_shape_s_cos6_in, - _, - _, - _, - geo_shape_s_sin3_in, - geo_shape_s_sin4_in, - geo_shape_s_sin5_in, - geo_shape_s_sin6_in, - ] = np.array(cos_sin_s).astype(float).T - - geo_signb_in = 1.0 - - geov_theta = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_bigr = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_bigr_r = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_bigr_t = np.zeros((n_theta,geo_rmin_in.shape[0])) - bigz = np.zeros((n_theta,geo_rmin_in.shape[0])) - bigz_r = np.zeros((n_theta,geo_rmin_in.shape[0])) - bigz_t = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_jac_r = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_grad_r = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_l_t = np.zeros((n_theta,geo_rmin_in.shape[0])) - r_c = np.zeros((n_theta,geo_rmin_in.shape[0])) - bigz_l = np.zeros((n_theta,geo_rmin_in.shape[0])) - bigr_l = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_l_r = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_nsin = np.zeros((n_theta,geo_rmin_in.shape[0])) - - pi_2 = 8.0 * np.arctan(1.0) - d_theta = pi_2 / (n_theta - 1) - - for i in range(n_theta): - #!----------------------------------------- - #! Generalized Miller-type parameterization - #!----------------------------------------- - - theta = -0.5 * pi_2 + (i - 1) * d_theta - - geov_theta[i] = theta - - x = np.arcsin(geo_delta_in) - - #! A - #! dA/dtheta - #! d^2A/dtheta^2 - a = ( - theta - + geo_shape_cos0_in - + geo_shape_cos1_in * np.cos(theta) - + geo_shape_cos2_in * np.cos(2 * theta) - + geo_shape_cos3_in * np.cos(3 * theta) - + geo_shape_cos4_in * np.cos(4 * theta) - + geo_shape_cos5_in * np.cos(5 * theta) - + geo_shape_cos6_in * np.cos(6 * theta) - + geo_shape_sin3_in * np.sin(3 * theta) - + x * np.sin(theta) - - geo_zeta_in * np.sin(2 * theta) - + geo_shape_sin3_in * np.sin(3 * theta) - + geo_shape_sin4_in * np.sin(4 * theta) - + geo_shape_sin5_in * np.sin(5 * theta) - + geo_shape_sin6_in * np.sin(6 * theta) - ) - a_t = ( - 1.0 - - geo_shape_cos1_in * np.sin(theta) - - 2 * geo_shape_cos2_in * np.sin(2 * theta) - - 3 * geo_shape_cos3_in * np.sin(3 * theta) - - 4 * geo_shape_cos4_in * np.sin(4 * theta) - - 5 * geo_shape_cos5_in * np.sin(5 * theta) - - 6 * geo_shape_cos6_in * np.sin(6 * theta) - + x * np.cos(theta) - - 2 * geo_zeta_in * np.cos(2 * theta) - + 3 * geo_shape_sin3_in * np.cos(3 * theta) - + 4 * geo_shape_sin4_in * np.cos(4 * theta) - + 5 * geo_shape_sin5_in * np.cos(5 * theta) - + 6 * geo_shape_sin6_in * np.cos(6 * theta) - ) - a_tt = ( - -geo_shape_cos1_in * np.cos(theta) - - 4 * geo_shape_cos2_in * np.cos(2 * theta) - - 9 * geo_shape_cos3_in * np.cos(3 * theta) - - 16 * geo_shape_cos4_in * np.cos(4 * theta) - - 25 * geo_shape_cos5_in * np.cos(5 * theta) - - 36 * geo_shape_cos6_in * np.cos(6 * theta) - - x * np.sin(theta) - + 4 * geo_zeta_in * np.sin(2 * theta) - - 9 * geo_shape_sin3_in * np.sin(3 * theta) - - 16 * geo_shape_sin4_in * np.sin(4 * theta) - - 25 * geo_shape_sin5_in * np.sin(5 * theta) - - 36 * geo_shape_sin6_in * np.sin(6 * theta) - ) - - #! R(theta) - #! dR/dr - #! dR/dtheta - #! d^2R/dtheta^2 - geov_bigr[i] = geo_rmaj_in + geo_rmin_in * np.cos(a) - geov_bigr_r[i] = ( - geo_drmaj_in - + np.cos(a) - - np.sin(a) - * ( - geo_shape_s_cos0_in - + geo_shape_s_cos1_in * np.cos(theta) - + geo_shape_s_cos2_in * np.cos(2 * theta) - + geo_shape_s_cos3_in * np.cos(3 * theta) - + geo_shape_s_cos4_in * np.cos(4 * theta) - + geo_shape_s_cos5_in * np.cos(5 * theta) - + geo_shape_s_cos6_in * np.cos(6 * theta) - + geo_s_delta_in / np.cos(x) * np.sin(theta) - - geo_s_zeta_in * np.sin(2 * theta) - + geo_shape_s_sin3_in * np.sin(3 * theta) - + geo_shape_s_sin4_in * np.sin(4 * theta) - + geo_shape_s_sin5_in * np.sin(5 * theta) - + geo_shape_s_sin6_in * np.sin(6 * theta) - ) - ) - geov_bigr_t[i] = -geo_rmin_in * a_t * np.sin(a) - bigr_tt = -geo_rmin_in * a_t**2 * np.cos(a) - geo_rmin_in * a_tt * np.sin(a) - - #!----------------------------------------------------------- - - #! A - #! dA/dtheta - #! d^2A/dtheta^2 - a = theta - a_t = 1.0 - a_tt = 0.0 - - #! Z(theta) - #! dZ/dr - #! dZ/dtheta - #! d^2Z/dtheta^2 - bigz[i] = geo_zmag_in + geo_kappa_in * geo_rmin_in * np.sin(a) - bigz_r[i] = geo_dzmag_in + geo_kappa_in * (1.0 + geo_s_kappa_in) * np.sin(a) - bigz_t[i] = geo_kappa_in * geo_rmin_in * np.cos(a) * a_t - bigz_tt = ( - -geo_kappa_in * geo_rmin_in * np.sin(a) * a_t**2 - + geo_kappa_in * geo_rmin_in * np.cos(a) * a_tt - ) - - g_tt = geov_bigr_t[i] ** 2 + bigz_t[i] ** 2 - - geov_jac_r[i] = geov_bigr[i] * ( - geov_bigr_r[i] * bigz_t[i] - geov_bigr_t[i] * bigz_r[i] - ) - - geov_grad_r[i] = geov_bigr[i] * np.sqrt(g_tt) / geov_jac_r[i] - - geov_l_t[i] = np.sqrt(g_tt) - - r_c[i] = geov_l_t[i] ** 3 / (geov_bigr_t[i] * bigz_tt - bigz_t[i] * bigr_tt) - - bigz_l[i] = bigz_t[i] / geov_l_t[i] - - bigr_l[i] = geov_bigr_t[i] / geov_l_t[i] - - geov_l_r[i] = bigz_l[i] * bigz_r[i] + bigr_l[i] * geov_bigr_r[i] - - geov_nsin[i] = ( - geov_bigr_r[i] * geov_bigr_t[i] + bigz_r[i] * bigz_t[i] - ) / geov_l_t[i] - - c = 0.0 - for i in range(n_theta): - c = c + geov_l_t[i] / (geov_bigr[i] * geov_grad_r[i]) - - f = geo_rmin_in / (c * d_theta / pi_2) - - c = 0.0 - for i in range(n_theta - 1): - c = c + geov_l_t[i] * geov_bigr[i] / geov_grad_r[i] - - geo_volume_prime = pi_2 * c * d_theta - - # Line 716 in geo.f90 - geo_surf = 0.0 - for i in range(n_theta - 1): - geo_surf = geo_surf + geov_l_t[i] * geov_bigr[i] - geo_surf = pi_2 * geo_surf * d_theta - - # ----- - c = 0.0 - for i in range(n_theta - 1): - c = c + geov_l_t[i] / (geov_bigr[i] * geov_grad_r[i]) - f = geo_rmin_in / (c * d_theta / pi_2) - - geov_b = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_g_theta = np.zeros((n_theta,geo_rmin_in.shape[0])) - geov_bt = np.zeros((n_theta,geo_rmin_in.shape[0])) - for i in range(n_theta): - geov_bt[i] = f / geov_bigr[i] - geov_bp = (geo_rmin_in / geo_q_in) * geov_grad_r[i] / geov_bigr[i] - - geov_b[i] = geo_signb_in * (geov_bt[i] ** 2 + geov_bp**2) ** 0.5 - geov_g_theta[i] = ( - geov_bigr[i] - * geov_b[i] - * geov_l_t[i] - / (geo_rmin_in * geo_rmaj_in * geov_grad_r[i]) - ) - - theta_0 = 0 - dx = geov_theta[1,0] - geov_theta[0,0] - x0 = theta_0 - geov_theta[0,0] - i1 = int(x0 / dx) + 1 - i2 = i1 + 1 - x1 = (i1 - 1) * dx - z = (x0 - x1) / dx - if i2 == n_theta: - i2 -= 1 - bt_geo0 = geov_bt[i1] + (geov_bt[i2] - geov_bt[i1]) * z - - denom = 0 - for i in range(n_theta - 1): - denom = denom + geov_g_theta[i] / geov_b[i] - - geo_fluxsurfave_grad_r = 0 - for i in range(n_theta - 1): - geo_fluxsurfave_grad_r = ( - geo_fluxsurfave_grad_r - + geov_grad_r[i] * geov_g_theta[i] / geov_b[i] / denom - ) - - geo_fluxsurfave__bp2 = 0 - for i in range(n_theta - 1): - geo_fluxsurfave__bp2 = ( - geo_fluxsurfave__bp2 - + geov_bt[i] ** 2 * geov_g_theta[i] / geov_b[i] / denom - ) - - geo_fluxsurfave_bt2 = 0 - for i in range(n_theta - 1): - geo_fluxsurfave_bt2 = ( - geo_fluxsurfave_bt2 - + geov_bp ** 2 * geov_g_theta[i] / geov_b[i] / denom - ) - - return geo_volume_prime, geo_surf, geo_fluxsurfave_grad_r, geo_fluxsurfave__bp2, geo_fluxsurfave_bt2, bt_geo0 - -def xsec_area_RZ( - R, - Z -): - # calculates the cross-sectional area of the plasma for each flux surface - xsec_area = [] - for i in range(R.shape[0]): - R0 = np.max(R[i,:]) - np.min(R[i,:]) - Z0 = np.max(Z[i,:]) - np.min(Z[i,:]) - xsec_area.append(np.trapz(R[i], Z[i])) - - xsec_area = np.array(xsec_area) - - return xsec_area \ No newline at end of file From c060a4eb610cde820019680b2cf57eae7e7e14ea Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 20:27:59 +0200 Subject: [PATCH 087/385] Cleaned up volume integration --- src/mitim_modules/powertorch/STATEtools.py | 12 +- .../powertorch/utils/CALCtools.py | 94 +++++---------- .../powertorch/utils/TARGETStools.py | 2 +- src/mitim_tools/misc_tools/MATHtools.py | 51 ++++++++ src/mitim_tools/misc_tools/PLASMAtools.py | 12 +- .../plasmastate_tools/MITIMstate.py | 111 +++++++----------- .../popcon_tools/utils/FUNCTIONALScalc.py | 7 +- 7 files changed, 143 insertions(+), 146 deletions(-) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index dcca4b3d..30df15e7 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -752,13 +752,13 @@ def volume_integrate(self, var, force_dim=None): """ if force_dim is None: - return CALCtools.integrateQuadPoly( - self.plasma["rmin"], var * self.plasma["volp"] - ) / self.plasma["volp"] + return CALCtools.volume_integration( + var, self.plasma["rmin"], self.plasma["volp"] + ) / self.plasma["volp"] else: - return CALCtools.integrateQuadPoly( - self.plasma["rmin"][0,:].repeat(force_dim,1), var * self.plasma["volp"][0,:].repeat(force_dim,1), - ) / self.plasma["volp"][0,:].repeat(force_dim,1) + return CALCtools.volume_integration( + var, self.plasma["rmin"][0,:].repeat(force_dim,1), self.plasma["volp"][0,:].repeat(force_dim,1) + ) / self.plasma["volp"][0,:].repeat(force_dim,1) def add_axes_powerstate_plot(figMain, num_kp=3): diff --git a/src/mitim_modules/powertorch/utils/CALCtools.py b/src/mitim_modules/powertorch/utils/CALCtools.py index 980869cc..77188182 100644 --- a/src/mitim_modules/powertorch/utils/CALCtools.py +++ b/src/mitim_modules/powertorch/utils/CALCtools.py @@ -1,5 +1,6 @@ import torch import numpy as np +import contextlib from mitim_tools.misc_tools import MATHtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -76,72 +77,45 @@ def produceGradient_lin(r, p): return z +def _to_2d(x, xp): + """Ensure shape (batch, N) for either NumPy or Torch.""" + if xp is np: + return np.atleast_2d(x) + else: # torch + return x.unsqueeze(0) if x.ndim == 1 else x - -def integrateQuadPoly(r, s, p=None): +def volume_integration(p, r, volp): """ - (batch,dim) - - Computes int(s*dr), so if s is s*dV/dr, then int(s*dV), which is the full integral - - From tgyro_volume_int.f90 - r - minor radius - s - s*volp - - (Modified to avoid if statements and for loops) - + Compute the volume integral ∫ p · dV with dV/dr = volp. + + Parameters + ---------- + p, r, volp : 1-D or 2-D NumPy arrays (shape: (N,) or (M, N)) + • If they are 1-D, each represents a single radial profile. + • If they are 2-D -> (batch,dim_radius) + + Returns + ------- + out : ndarray + • 1-D array if the inputs were 1-D + • 2-D array (same leading dimension as the inputs) otherwise """ - if p is None: - p = torch.zeros((r.shape[0], r.shape[1])).to(r) - - # First point - - x1, x2, x3 = r[..., 0], r[..., 1], r[..., 2] - f1, f2, f3 = s[..., 0], s[..., 1], s[..., 2] + # Decide backend from *p* only + xp = torch if isinstance(p, torch.Tensor) else np - p[..., 1] = (x2 - x1) * ( - (3 * x3 - x2 - 2 * x1) * f1 / 6 / (x3 - x1) - + (3 * x3 - 2 * x2 - x1) * f2 / 6 / (x3 - x2) - - (x2 - x1) ** 2 * f3 / 6 / (x3 - x1) / (x3 - x2) - ) + # Remember whether the caller passed 1-D profiles + one_dim = (p.ndim == 1) and (r.ndim == 1) and (volp.ndim == 1) - # Next points - x1, x2, x3 = r[..., :-2], r[..., 1:-1], r[..., 2:] - f1, f2, f3 = s[..., :-2], s[..., 1:-1], s[..., 2:] + # Promote to 2-D for vectorised processing + r_2d = _to_2d(r, xp) + pdVdr_2d = _to_2d(p * volp, xp) - p[..., 2:] = ( - (x3 - x2) - / (x3 - x1) - / 6 - * ( - (2 * x3 + x2 - 3 * x1) * f3 - + (x3 + 2 * x2 - 3 * x1) * f2 * (x3 - x1) / (x2 - x1) - - (x3 - x2) ** 2 * f1 / (x2 - x1) - ) - ) - - try: - p = torch.cumsum(p, 1) - except: - p = np.cumsum(p, 1) + # Integrate row-wise (using your original routine) + result_2d = MATHtools.integrateQuadPoly(r_2d, pdVdr_2d) - return p - - -def integrateFS(P, r, volp): - """ - Based on the idea that volp = dV/dr, whatever r is - - Ptot = int_V P*dV = int_r P*V'*dr - - """ - - I = integrateQuadPoly( - np.atleast_2d(r), np.atleast_2d(P * volp), p=np.zeros((1, P.shape[0])) - )[0, :] - - return I + # Collapse back to original rank if necessary + return result_2d[0] if one_dim else result_2d """ @@ -150,10 +124,6 @@ def integrateFS(P, r, volp): ---------------------------------------------------------------------------------------------------------------- """ -import torch -import contextlib - - class Interp1d(torch.autograd.Function): def __call__(self, x, y, xnew, out=None): return self.forward(x, y, xnew, out) diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 1ae63c80..ae420886 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -62,7 +62,7 @@ def flux_integrate(self): """ ************************************************************************************************** Calculate integral of all targets, and then sum aux. - Reason why I do it this convoluted way is to make it faster in mitim, not to run integrateQuadPoly all the time. + Reason why I do it this convoluted way is to make it faster in mitim, not to run the volume integral all the time. Run once for all the batch and also for electrons and ions (in MW/m^2) ************************************************************************************************** diff --git a/src/mitim_tools/misc_tools/MATHtools.py b/src/mitim_tools/misc_tools/MATHtools.py index cfe023ea..e192301c 100644 --- a/src/mitim_tools/misc_tools/MATHtools.py +++ b/src/mitim_tools/misc_tools/MATHtools.py @@ -273,6 +273,57 @@ def integrate_definite(x, y, rangex=None): return 0 +def integrateQuadPoly(r, s): + """ + (batch,dim) + + Computes int(s*dr), so if s is s*dV/dr, then int(s*dV), which is the full integral + + From tgyro_volume_int.f90 + r - minor radius + s - s*volp + + (Modified to avoid if statements and for loops) + + """ + + if isinstance(s, torch.Tensor): + p = torch.zeros((r.shape[0], r.shape[1])).to(r) + else: + p = np.zeros((r.shape[0], r.shape[1])) + + # First point + + x1, x2, x3 = r[..., 0], r[..., 1], r[..., 2] + f1, f2, f3 = s[..., 0], s[..., 1], s[..., 2] + + p[..., 1] = (x2 - x1) * ( + (3 * x3 - x2 - 2 * x1) * f1 / 6 / (x3 - x1) + + (3 * x3 - 2 * x2 - x1) * f2 / 6 / (x3 - x2) + - (x2 - x1) ** 2 * f3 / 6 / (x3 - x1) / (x3 - x2) + ) + + # Next points + x1, x2, x3 = r[..., :-2], r[..., 1:-1], r[..., 2:] + f1, f2, f3 = s[..., :-2], s[..., 1:-1], s[..., 2:] + + p[..., 2:] = ( + (x3 - x2) + / (x3 - x1) + / 6 + * ( + (2 * x3 + x2 - 3 * x1) * f3 + + (x3 + 2 * x2 - 3 * x1) * f2 * (x3 - x1) / (x2 - x1) + - (x3 - x2) ** 2 * f1 / (x2 - x1) + ) + ) + + if isinstance(p, torch.Tensor): + return torch.cumsum(p, 1) + else: + return np.cumsum(p, 1) + + def extrapolate(x, xp, yp, order=3): s = InterpolatedUnivariateSpline(xp, yp, k=order) diff --git a/src/mitim_tools/misc_tools/PLASMAtools.py b/src/mitim_tools/misc_tools/PLASMAtools.py index e044737a..6ef18d70 100644 --- a/src/mitim_tools/misc_tools/PLASMAtools.py +++ b/src/mitim_tools/misc_tools/PLASMAtools.py @@ -306,9 +306,9 @@ def calculatePressure(Te, Ti, ne, ni): def calculateVolumeAverage(rmin, var, dVdr): W, vals = [], [] for it in range(rmin.shape[0]): - W.append(CALCtools.integrateFS(var[it, :], rmin[it, :], dVdr[it, :])[-1]) + W.append(CALCtools.volume_integration(var[it, :], rmin[it, :], dVdr[it, :])[-1]) vals.append( - CALCtools.integrateFS(np.ones(rmin.shape[1]), rmin[it, :], dVdr[it, :])[-1] + CALCtools.volume_integration(np.ones(rmin.shape[1]), rmin[it, :], dVdr[it, :])[-1] ) W = np.array(W) / np.array(vals) @@ -344,7 +344,7 @@ def calculateContent(rmin, Te, Ti, ne, ni, dVdr): for it in range(rmin.shape[0]): # Number of electrons Ne.append( - CALCtools.integrateFS(ne[it, :], rmin[it, :], dVdr[it, :])[-1] + CALCtools.volume_integration(ne[it, :], rmin[it, :], dVdr[it, :])[-1] ) # Number of particles total # Number of ions @@ -352,14 +352,14 @@ def calculateContent(rmin, Te, Ti, ne, ni, dVdr): for i in range(ni.shape[1]): ni0 += ni[it, i, :] Ni.append( - CALCtools.integrateFS(ni0, rmin[it, :], dVdr[it, :])[-1] + CALCtools.volume_integration(ni0, rmin[it, :], dVdr[it, :])[-1] ) # Number of particles total # Electron stored energy Wx = 3 / 2 * pe[it, :] - We.append(CALCtools.integrateFS(Wx, rmin[it, :], dVdr[it, :])[-1]) # MJ + We.append(CALCtools.volume_integration(Wx, rmin[it, :], dVdr[it, :])[-1]) # MJ Wx = 3 / 2 * pi[it, :] - Wi.append(CALCtools.integrateFS(Wx, rmin[it, :], dVdr[it, :])[-1]) # MJ + Wi.append(CALCtools.volume_integration(Wx, rmin[it, :], dVdr[it, :])[-1]) # MJ We = np.array(We) Wi = np.array(Wi) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index ab344587..76e06ef6 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -332,9 +332,9 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): r = self.profiles["rmin(m)"] volp = self.derived["volp_geo"] - self.derived["qe_MW"] = CALCtools.integrateFS(self.derived["qe"], r, volp) - self.derived["qi_MW"] = CALCtools.integrateFS(self.derived["qi"], r, volp) - self.derived["ge_10E20"] = CALCtools.integrateFS(self.derived["ge"] * 1e-20, r, volp) # Because the units were #/sec/m^3 + self.derived["qe_MW"] = CALCtools.volume_integration(self.derived["qe"], r, volp) + self.derived["qi_MW"] = CALCtools.volume_integration(self.derived["qi"], r, volp) + self.derived["ge_10E20"] = CALCtools.volume_integration(self.derived["ge"] * 1e-20, r, volp) # Because the units were #/sec/m^3 self.derived["geIn"] = self.derived["ge_10E20"][-1] # 1E20 particles/sec @@ -353,7 +353,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): ) # qmom - self.derived["mt_Jmiller"] = CALCtools.integrateFS( + self.derived["mt_Jmiller"] = CALCtools.volume_integration( self.profiles[self.varqmom], r, volp ) self.derived["mt_Jm2"] = self.derived["mt_Jmiller"] / (volp) @@ -366,10 +366,10 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): P += self.profiles["qbrem(MW/m^3)"] if "qline(MW/m^3)" in self.profiles: P += self.profiles["qline(MW/m^3)"] - self.derived["qe_rad_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_rad_MW"] = CALCtools.volume_integration(P, r, volp) P = self.profiles["qei(MW/m^3)"] - self.derived["qe_exc_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_exc_MW"] = CALCtools.volume_integration(P, r, volp) """ --------------------------------------------------------------------------------------------------------------------- @@ -386,14 +386,14 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): P += self.profiles[i] self.derived["qe_auxONLY"] = copy.deepcopy(P) - self.derived["qe_auxONLY_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_auxONLY_MW"] = CALCtools.volume_integration(P, r, volp) for i in ["qione(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] self.derived["qe_aux"] = copy.deepcopy(P) - self.derived["qe_aux_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_aux_MW"] = CALCtools.volume_integration(P, r, volp) # ** Ions @@ -403,14 +403,14 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): P += self.profiles[i] self.derived["qi_auxONLY"] = copy.deepcopy(P) - self.derived["qi_auxONLY_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qi_auxONLY_MW"] = CALCtools.volume_integration(P, r, volp) for i in ["qioni(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] self.derived["qi_aux"] = copy.deepcopy(P) - self.derived["qi_aux_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qi_aux_MW"] = CALCtools.volume_integration(P, r, volp) # ** General @@ -418,19 +418,19 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): for i in ["qohme(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qOhm_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qOhm_MW"] = CALCtools.volume_integration(P, r, volp) P = np.zeros(len(self.profiles["rho(-)"])) for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qRF_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qRF_MW"] = CALCtools.volume_integration(P, r, volp) if "qrfe(MW/m^3)" in self.profiles: - self.derived["qRFe_MW"] = CALCtools.integrateFS( + self.derived["qRFe_MW"] = CALCtools.volume_integration( self.profiles["qrfe(MW/m^3)"], r, volp ) if "qrfi(MW/m^3)" in self.profiles: - self.derived["qRFi_MW"] = CALCtools.integrateFS( + self.derived["qRFi_MW"] = CALCtools.volume_integration( self.profiles["qrfi(MW/m^3)"], r, volp ) @@ -438,19 +438,19 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): for i in ["qbeame(MW/m^3)", "qbeami(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qBEAM_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qBEAM_MW"] = CALCtools.volume_integration(P, r, volp) - self.derived["qrad_MW"] = CALCtools.integrateFS(self.derived["qrad"], r, volp) + self.derived["qrad_MW"] = CALCtools.volume_integration(self.derived["qrad"], r, volp) if "qsync(MW/m^3)" in self.profiles: - self.derived["qrad_sync_MW"] = CALCtools.integrateFS(self.profiles["qsync(MW/m^3)"], r, volp) + self.derived["qrad_sync_MW"] = CALCtools.volume_integration(self.profiles["qsync(MW/m^3)"], r, volp) else: self.derived["qrad_sync_MW"] = self.derived["qrad_MW"]*0.0 if "qbrem(MW/m^3)" in self.profiles: - self.derived["qrad_brem_MW"] = CALCtools.integrateFS(self.profiles["qbrem(MW/m^3)"], r, volp) + self.derived["qrad_brem_MW"] = CALCtools.volume_integration(self.profiles["qbrem(MW/m^3)"], r, volp) else: self.derived["qrad_brem_MW"] = self.derived["qrad_MW"]*0.0 if "qline(MW/m^3)" in self.profiles: - self.derived["qrad_line_MW"] = CALCtools.integrateFS(self.profiles["qline(MW/m^3)"], r, volp) + self.derived["qrad_line_MW"] = CALCtools.volume_integration(self.profiles["qline(MW/m^3)"], r, volp) else: self.derived["qrad_line_MW"] = self.derived["qrad_MW"]*0.0 @@ -458,13 +458,13 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): for i in ["qfuse(MW/m^3)", "qfusi(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qFus_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qFus_MW"] = CALCtools.volume_integration(P, r, volp) P = np.zeros(len(self.profiles["rho(-)"])) for i in ["qioni(MW/m^3)", "qione(MW/m^3)"]: if i in self.profiles: P += self.profiles[i] - self.derived["qz_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qz_MW"] = CALCtools.volume_integration(P, r, volp) self.derived["q_MW"] = ( self.derived["qe_MW"] + self.derived["qi_MW"] @@ -476,12 +476,12 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): P = np.zeros(len(self.profiles["rho(-)"])) if "qfuse(MW/m^3)" in self.profiles: P = self.profiles["qfuse(MW/m^3)"] - self.derived["qe_fus_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qe_fus_MW"] = CALCtools.volume_integration(P, r, volp) P = np.zeros(len(self.profiles["rho(-)"])) if "qfusi(MW/m^3)" in self.profiles: P = self.profiles["qfusi(MW/m^3)"] - self.derived["qi_fus_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["qi_fus_MW"] = CALCtools.volume_integration(P, r, volp) P = np.zeros(len(self.profiles["rho(-)"])) if "qfusi(MW/m^3)" in self.profiles: @@ -490,7 +490,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): ) * 5 P = self.derived["q_fus"] self.derived["q_fus"] = P - self.derived["q_fus_MW"] = CALCtools.integrateFS(P, r, volp) + self.derived["q_fus_MW"] = CALCtools.volume_integration(P, r, volp) """ Derivatives @@ -632,11 +632,11 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): ).transpose().repeat(self.profiles["ni(10^19/m^3)"].shape[1], axis=1) # Vol-avg density - self.derived["volume"] = CALCtools.integrateFS(np.ones(r.shape[0]), r, volp)[ + self.derived["volume"] = CALCtools.volume_integration(np.ones(r.shape[0]), r, volp)[ -1 ] # m^3 self.derived["ne_vol20"] = ( - CALCtools.integrateFS(self.profiles["ne(10^19/m^3)"] * 0.1, r, volp)[-1] + CALCtools.volume_integration(self.profiles["ne(10^19/m^3)"] * 0.1, r, volp)[-1] / self.derived["volume"] ) # 1E20/m^3 @@ -644,7 +644,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["fi_vol"] = np.zeros(self.profiles["ni(10^19/m^3)"].shape[1]) for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): self.derived["ni_vol20"][i] = ( - CALCtools.integrateFS( + CALCtools.volume_integration( self.profiles["ni(10^19/m^3)"][:, i] * 0.1, r, volp )[-1] / self.derived["volume"] @@ -671,14 +671,14 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): ) self.derived["Te_vol"] = ( - CALCtools.integrateFS(self.profiles["te(keV)"], r, volp)[-1] + CALCtools.volume_integration(self.profiles["te(keV)"], r, volp)[-1] / self.derived["volume"] ) # keV self.derived["Te_peaking"] = ( self.profiles["te(keV)"][0] / self.derived["Te_vol"] ) self.derived["Ti_vol"] = ( - CALCtools.integrateFS(self.profiles["ti(keV)"][:, 0], r, volp)[-1] + CALCtools.volume_integration(self.profiles["ti(keV)"][:, 0], r, volp)[-1] / self.derived["volume"] ) # keV self.derived["Ti_peaking"] = ( @@ -686,17 +686,17 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): ) self.derived["ptot_manual_vol"] = ( - CALCtools.integrateFS(self.derived["ptot_manual"], r, volp)[-1] + CALCtools.volume_integration(self.derived["ptot_manual"], r, volp)[-1] / self.derived["volume"] ) # MPa self.derived["pthr_manual_vol"] = ( - CALCtools.integrateFS(self.derived["pthr_manual"], r, volp)[-1] + CALCtools.volume_integration(self.derived["pthr_manual"], r, volp)[-1] / self.derived["volume"] ) # MPa self.derived['pfast_manual'] = self.derived['ptot_manual'] - self.derived['pthr_manual'] self.derived["pfast_manual_vol"] = ( - CALCtools.integrateFS(self.derived["pfast_manual"], r, volp)[-1] + CALCtools.volume_integration(self.derived["pfast_manual"], r, volp)[-1] / self.derived["volume"] ) # MPa @@ -714,7 +714,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): / self.profiles["ne(10^19/m^3)"] ) self.derived["Zeff_vol"] = ( - CALCtools.integrateFS(self.derived["Zeff"], r, volp)[-1] + CALCtools.volume_integration(self.derived["Zeff"], r, volp)[-1] / self.derived["volume"] ) @@ -766,7 +766,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): w0_Mach1 = Vtor_LF_Mach1 / (self.derived["R_LF"]) # rad/s self.derived["MachNum"] = self.profiles["w0(rad/s)"] / w0_Mach1 self.derived["MachNum_vol"] = ( - CALCtools.integrateFS(self.derived["MachNum"], r, volp)[-1] + CALCtools.volume_integration(self.derived["MachNum"], r, volp)[-1] / self.derived["volume"] ) @@ -791,9 +791,9 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): # Calculate the volume averages of bt2 and bp2 P = self.derived["bp2_exp"] - self.derived["bp2_vol_avg"] = CALCtools.integrateFS(P, r, volp)[-1] / self.derived["volume"] + self.derived["bp2_vol_avg"] = CALCtools.volume_integration(P, r, volp)[-1] / self.derived["volume"] P = self.derived["bt2_exp"] - self.derived["bt2_vol_avg"] = CALCtools.integrateFS(P, r, volp)[-1] / self.derived["volume"] + self.derived["bt2_vol_avg"] = CALCtools.volume_integration(P, r, volp)[-1] / self.derived["volume"] # calculate beta_poloidal and beta_toroidal using volume averaged values # mu0 = 4pi x 10^-7, also need to convert MPa to Pa @@ -1278,26 +1278,15 @@ def DTplasma(self): elif self.Tion is not None: self.Mion = self.Tion # Main else: - self.Mion = ( - 0 # If no D or T, assume that the main ion is the first and only - ) + self.Mion = 0 # If no D or T, assume that the main ion is the first and only - self.ion_list_main = [] - if self.DTplasmaBool: - self.ion_list_main = [self.Dion+1, self.Tion+1] - else: - self.ion_list_main = [self.Mion+1] - + self.ion_list_main = [self.Dion+1, self.Tion+1] if self.DTplasmaBool else [self.Mion+1] self.ion_list_impurities = [i+1 for i in range(len(self.Species)) if i+1 not in self.ion_list_main] def remove(self, ions_list): # First order them ions_list.sort() - print( - "\t\t- Removing ions in positions (of ions order, no zero): ", - ions_list, - typeMsg="i", - ) + print("\t\t- Removing ions in positions (of ions order, no zero): ",ions_list,typeMsg="i",) ions_list = [i - 1 for i in ions_list] @@ -1308,10 +1297,7 @@ def remove(self, ions_list): try: self.profiles[i] = np.delete(self.profiles[i], ions_list) except: - print( - f"\t\t\t* Ions {[k+1 for k in ions_list]} could not be removed", - typeMsg="w", - ) + print(f"\t\t\t* Ions {[k+1 for k in ions_list]} could not be removed",typeMsg="w") fail = True break @@ -1322,9 +1308,7 @@ def remove(self, ions_list): if not fail: # Ensure we extract the scalar value from the array - self.profiles["nion"] = np.array( - [str(int(self.profiles["nion"][0]) - len(ions_list))] - ) + self.profiles["nion"] = np.array([str(int(self.profiles["nion"][0]) - len(ions_list))]) self.readSpecies() self.derive_quantities(rederiveGeometry=False) @@ -1363,12 +1347,7 @@ def lumpSpecies( fZ2 += self.Species[i - 1]["Z"] ** 2 * self.derived["fi"][:, i - 1] Zr = fZ2 / fZ1 - Zr_vol = ( - CALCtools.integrateFS( - Zr, self.profiles["rmin(m)"], self.derived["volp_geo"] - )[-1] - / self.derived["volume"] - ) + Zr_vol = CALCtools.volume_integration(Zr, self.profiles["rmin(m)"], self.derived["volp_geo"])[-1] / self.derived["volume"] print(f'\t\t\t* Original plasma had Zeff_vol={self.derived["Zeff_vol"]:.2f}, QN error={self.derived["QN_Error"]:.4f}') @@ -3271,13 +3250,13 @@ def plotPeaking( ne = self.profiles["ne(10^19/m^3)"] axq.plot(self.profiles["rho(-)"], ne, color="m") ne_vol = ( - CALCtools.integrateFS(ne * 0.1, r, volp)[-1] / self.derived["volume"] + CALCtools.volume_integration(ne * 0.1, r, volp)[-1] / self.derived["volume"] ) axq.axhline(y=ne_vol * 10, color="m") ne = copy.deepcopy(self.profiles["ne(10^19/m^3)"]) ne[ix:] = (0,) * len(ne[ix:]) - ne_vol = CALCtools.integrateFS(ne * 0.1, r, volp)[-1] / self.derived["volume"] + ne_vol = CALCtools.volume_integration(ne * 0.1, r, volp)[-1] / self.derived["volume"] ne_peaking0 = ( ne[np.argmin(np.abs(self.derived["rho_pol"] - 0.2))] * 0.1 / ne_vol ) @@ -3288,7 +3267,7 @@ def plotPeaking( ne = copy.deepcopy(self.profiles["ne(10^19/m^3)"]) ne[ix:] = (ne[ix],) * len(ne[ix:]) - ne_vol = CALCtools.integrateFS(ne * 0.1, r, volp)[-1] / self.derived["volume"] + ne_vol = CALCtools.volume_integration(ne * 0.1, r, volp)[-1] / self.derived["volume"] ne_peaking1 = ( ne[np.argmin(np.abs(self.derived["rho_pol"] - 0.2))] * 0.1 / ne_vol ) diff --git a/src/mitim_tools/popcon_tools/utils/FUNCTIONALScalc.py b/src/mitim_tools/popcon_tools/utils/FUNCTIONALScalc.py index db8ec4f8..b2a8e6e6 100644 --- a/src/mitim_tools/popcon_tools/utils/FUNCTIONALScalc.py +++ b/src/mitim_tools/popcon_tools/utils/FUNCTIONALScalc.py @@ -35,13 +35,10 @@ def doubleLinear_aLT(x, g1, g2, t, T1): def calculate_simplified_volavg(x, T): x = np.atleast_2d(x) dVdr = 2 * x - vol = CALCtools.integrateQuadPoly(torch.from_numpy(x), torch.ones(x.shape) * dVdr) + vol = CALCtools.volume_integration(torch.ones(x.shape), torch.from_numpy(x), dVdr) Tvol = ( - CALCtools.integrateQuadPoly(torch.from_numpy(x), torch.from_numpy(T) * dVdr)[ - :, -1 - ] - / vol[:, -1] + CALCtools.volume_integration(torch.from_numpy(T), torch.from_numpy(x), dVdr)[:, -1] / vol[:, -1] ).cpu().numpy() return Tvol From 9cb42a40d1a2a6ff24297f4fab10baff16eece08 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 20:40:28 +0200 Subject: [PATCH 088/385] Cleaned up integration and derivation functions --- src/mitim_modules/maestro/utils/EPEDbeat.py | 8 +-- .../maestro/utils/MAESTROplot.py | 2 +- src/mitim_modules/powertorch/STATEtools.py | 3 +- .../physics_models/parameterizers.py | 8 +-- .../powertorch/utils/CALCtools.py | 49 ++++++++++--------- src/mitim_tools/gacode_tools/TGLFtools.py | 2 +- .../plasmastate_tools/MITIMstate.py | 8 +-- .../popcon_tools/FunctionalForms.py | 4 +- .../popcon_tools/scripts/test_functionals.py | 12 ++--- 9 files changed, 51 insertions(+), 45 deletions(-) diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index 2ac10ae7..63267924 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -488,12 +488,12 @@ def scale_profile_by_stretching( x, y, xp, yp, xp_old, plotYN=False, label='', k print('\t\t\t* Keeping old aLT profile in the core-predicted region, using r/a for it') # Calculate gradient in entire region - aLy = CALCtools.produceGradient( torch.from_numpy(roa), torch.from_numpy(y) ) + aLy = CALCtools.derivation_into_Lx( torch.from_numpy(roa), torch.from_numpy(y) ) # I'm only interested in core region, plus one ghost point with the same gradient aLy = torch.cat( (aLy[:ibc+1], aLy[ibc].unsqueeze(0)) ) - y_mod = CALCtools.integrateGradient( torch.from_numpy(roa[:ibc+2]).unsqueeze(0), aLy.unsqueeze(0), torch.from_numpy(np.array(ynew[ibc+1])).unsqueeze(0) ).squeeze().numpy() + y_mod = CALCtools.integration_Lx( torch.from_numpy(roa[:ibc+2]).unsqueeze(0), aLy.unsqueeze(0), torch.from_numpy(np.array(ynew[ibc+1])).unsqueeze(0) ).squeeze().numpy() ynew[:ibc+2] = y_mod @@ -511,11 +511,11 @@ def scale_profile_by_stretching( x, y, xp, yp, xp_old, plotYN=False, label='', k ax.legend() ax = axs[1] - aLy = CALCtools.produceGradient( torch.from_numpy(roa), torch.from_numpy(y) ) + aLy = CALCtools.derivation_into_Lx( torch.from_numpy(roa), torch.from_numpy(y) ) ax.plot(x,aLy,'-o',color='b', label='old') ax.axvline(x=xp_old,color='b',ls='--') - aLy = CALCtools.produceGradient( torch.from_numpy(roa), torch.from_numpy(ynew) ) + aLy = CALCtools.derivation_into_Lx( torch.from_numpy(roa), torch.from_numpy(ynew) ) ax.plot(x,aLy,'-o',color='r', label='new') ax.axvline(x=xp,color='r',ls='--') diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index f97f2849..41be692c 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -97,7 +97,7 @@ def plot_results(self, fn): maxPlot = 5 if len(ps) > 0: # Plot profiles - figs = PROFILEStools.add_figures(fn,fnlab_pre = 'MAESTRO - ') + figs = MITIMstate.add_figures(fn,fnlab_pre = 'MAESTRO - ') log_file = self.folder_logs/'plot_maestro.log' if (not self.terminal_outputs) else None with LOGtools.conditional_log_to_file(log_file=log_file): MITIMstate.plotAll(ps[-maxPlot:], extralabs=ps_lab[-maxPlot:], figs=figs) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 30df15e7..e3e9440f 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -7,6 +7,7 @@ import dill as pickle from mitim_tools.misc_tools import PLASMAtools, IOtools from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.plasmastate_tools import MITIMstate from mitim_modules.powertorch.utils import TRANSFORMtools, POWERplot from mitim_tools.opt_tools.optimizers import optim from mitim_modules.powertorch.utils import TARGETStools, CALCtools, TRANSPORTtools @@ -449,7 +450,7 @@ def plot(self, axs=None, axsRes=None, axsMetrics=None, figs=None, fn=None,c="r", axsRes.append(figOpt.add_subplot(grid[j, i+1])) # Profiles - figs = PROFILEStools.add_figures(fn, tab_color='b') + figs = MITIMstate.add_figures(fn, tab_color='b') axs, axsMetrics = add_axes_powerstate_plot(figMain, num_kp = len(self.ProfilesPredicted)) diff --git a/src/mitim_modules/powertorch/physics_models/parameterizers.py b/src/mitim_modules/powertorch/physics_models/parameterizers.py index 36bfe27b..0f99e5d8 100644 --- a/src/mitim_modules/powertorch/physics_models/parameterizers.py +++ b/src/mitim_modules/powertorch/physics_models/parameterizers.py @@ -27,14 +27,14 @@ def piecewise_linear( if parameterize_in_aLx: # 1/Lx = -1/X*dX/dr integrator_function, derivator_function = ( - CALCtools.integrateGradient, - CALCtools.produceGradient, + CALCtools.integration_Lx, + CALCtools.derivation_into_Lx, ) else: # -dX/dr integrator_function, derivator_function = ( - CALCtools.integrateGradient_lin, - CALCtools.produceGradient_lin, + CALCtools.integration_dxdr, + CALCtools.derivation_into_dxdr, ) y_coord = torch.from_numpy(y_coord_raw).to(x_coarse_tensor) * multiplier_quantity diff --git a/src/mitim_modules/powertorch/utils/CALCtools.py b/src/mitim_modules/powertorch/utils/CALCtools.py index 77188182..7a910d8c 100644 --- a/src/mitim_modules/powertorch/utils/CALCtools.py +++ b/src/mitim_modules/powertorch/utils/CALCtools.py @@ -6,16 +6,22 @@ from IPython import embed # ******************************************************************************************************************** -# Normalized logaritmic gradient calculations +# Gradient calculations # ******************************************************************************************************************** -def integrateGradient(x, z, z0_bound): +def integration_Lx(x, z, f_bound): """ - inputs as - (batch,dim) - From tgyro_profile_functions.f90 - x is r - z is 1/LT = =-1/T*dT/dr + Integrates the gradient scale length into the profile + (adapted from tgyro_profile_functions.f90) + Inputs as + (batch,dim) + + x is r + z is 1/LT = =-1/T*dT/dr + f_bound is T @ at boundary condition (last point of the given profile) + + Notes: + - If x is r/a, then z is a/LT """ @@ -24,31 +30,25 @@ def integrateGradient(x, z, z0_bound): f1 = b / torch.cumprod(b, 1) * torch.prod(b, 1, keepdims=True) # Add the extra point of bounday condition - f = torch.cat((f1, torch.ones(z.shape[0], 1).to(f1)), dim=1) * z0_bound + f = torch.cat((f1, torch.ones(z.shape[0], 1).to(f1)), dim=1) * f_bound return f - -def produceGradient(r, p): +def derivation_into_Lx(r, p): """ Produces -1/p * dp/dr - or if r is roa: a/Lp + (adapted from expro_util.f90, bound_deriv) + + Notes: + - if r is r/a: a/Lp """ - # This is the same as it happens in expro_util.f90, bound_deriv z = MATHtools.deriv(r, -torch.log(p), array=False) - # # COMMENTED because this should happen at the coarse grid - # z = tgyro_math_zfind(r,p,z=z) - - return z # .nan_to_num(0.0) # Added this so that, when evaluating things like rotation shear, it doesn't blow - -# ******************************************************************************************************************** -# Linear gradient calculations -# ******************************************************************************************************************** + return z -def integrateGradient_lin(x, z, z0_bound): +def integration_dxdr(x, z, z0_bound): """ (batch,dim) From tgyro_profile_functions.f90 @@ -67,7 +67,7 @@ def integrateGradient_lin(x, z, z0_bound): return f -def produceGradient_lin(r, p): +def derivation_into_dxdr(r, p): """ Produces -dp/dr """ @@ -77,6 +77,11 @@ def produceGradient_lin(r, p): return z +# ******************************************************************************************************************** +# Volume calculations +# ******************************************************************************************************************** + + def _to_2d(x, xp): """Ensure shape (batch, N) for either NumPy or Torch.""" if xp is np: diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 4604b1e2..4df28493 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -3682,7 +3682,7 @@ def plotAnalysis(self, labels=["analysis1"], analysisType="chi_e", figs=None): from mitim_modules.powertorch.utils import CALCtools BC = 1.0 - T = CALCtools.integrateGradient( + T = CALCtools.integration_Lx( torch.from_numpy(rho_mod).unsqueeze(0), torch.from_numpy(aLn).unsqueeze(0), BC, diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 76e06ef6..7fcef888 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -3987,7 +3987,7 @@ def readTGYRO_profile_extra(file, varLabel="B_unit (T)"): def aLT(r, p): return ( r[-1] - * CALCtools.produceGradient( + * CALCtools.derivation_into_Lx( torch.from_numpy(r).to(torch.double), torch.from_numpy(p).to(torch.double) ) .cpu() @@ -4049,7 +4049,7 @@ def gradientsMerger(p0, p_true, roa=0.46, blending=0.1): aLTe = np.append(np.append(aLT0, aLT1), aLT2) Te = ( - CALCtools.integrateGradient( + CALCtools.integration_Lx( torch.from_numpy(p.derived["roa"]).unsqueeze(0), torch.Tensor(aLTe).unsqueeze(0), p.profiles["te(keV)"][-1], @@ -4068,7 +4068,7 @@ def gradientsMerger(p0, p_true, roa=0.46, blending=0.1): aLTi = np.append(np.append(aLT0, aLT1), aLT2) Ti = ( - CALCtools.integrateGradient( + CALCtools.integration_Lx( torch.from_numpy(p.derived["roa"]).unsqueeze(0), torch.Tensor(aLTi).unsqueeze(0), p.profiles["ti(keV)"][-1, 0], @@ -4087,7 +4087,7 @@ def gradientsMerger(p0, p_true, roa=0.46, blending=0.1): aLne = np.append(np.append(aLT0, aLT1), aLT2) ne = ( - CALCtools.integrateGradient( + CALCtools.integration_Lx( torch.from_numpy(p.derived["roa"]).unsqueeze(0), torch.Tensor(aLne).unsqueeze(0), p.profiles["ne(10^19/m^3)"][-1], diff --git a/src/mitim_tools/popcon_tools/FunctionalForms.py b/src/mitim_tools/popcon_tools/FunctionalForms.py index f2d9cf7f..c34d67b9 100644 --- a/src/mitim_tools/popcon_tools/FunctionalForms.py +++ b/src/mitim_tools/popcon_tools/FunctionalForms.py @@ -156,7 +156,7 @@ def MITIMfunctional_aLyTanh( aLy_profile[~linear_region] = aLy # Create core profile - Ycore = CALCtools.integrateGradient(torch.from_numpy(xcore).unsqueeze(0), + Ycore = CALCtools.integration_Lx(torch.from_numpy(xcore).unsqueeze(0), torch.from_numpy(aLy_profile).unsqueeze(0), y_top ).cpu().numpy()[0] @@ -182,7 +182,7 @@ def MITIMfunctional_aLyTanh( GRAPHICStools.addDenseAxis(ax) ax = axs[1] - aLy_reconstructed = CALCtools.produceGradient(torch.from_numpy(x), torch.from_numpy(y)).cpu().numpy() + aLy_reconstructed = CALCtools.derivation_into_Lx(torch.from_numpy(x), torch.from_numpy(y)).cpu().numpy() ax.plot(x, aLy_reconstructed, '-o', c='b', markersize=3, lw=1.0) ax.plot(xcore, aLy_profile, '-*', c='r', markersize=1, lw=1.0) diff --git a/src/mitim_tools/popcon_tools/scripts/test_functionals.py b/src/mitim_tools/popcon_tools/scripts/test_functionals.py index 7e8ceb21..7d6a1993 100644 --- a/src/mitim_tools/popcon_tools/scripts/test_functionals.py +++ b/src/mitim_tools/popcon_tools/scripts/test_functionals.py @@ -22,13 +22,13 @@ x, n = parabolic(Tbar=n_avol, nu=nu_n) axs[0, 0].plot(x, T, "-", c="b", label="Parabolic") -aLT = CALCtools.produceGradient(torch.from_numpy(x), torch.from_numpy(T)).cpu().numpy() +aLT = CALCtools.derivation_into_Lx(torch.from_numpy(x), torch.from_numpy(T)).cpu().numpy() axs[1, 0].plot(x, aLT, "-", c="b") print(f" Parabolic = {FUNCTIONALScalc.calculate_simplified_volavg(x,T)[0]:.3f}keV") # axs[0,1].plot(x,n,'-',c='b') -# aLn = CALCtools.produceGradient(torch.from_numpy(x),torch.from_numpy(n)).cpu().numpy() +# aLn = CALCtools.derivation_into_Lx(torch.from_numpy(x),torch.from_numpy(n)).cpu().numpy() # axs[1,1].plot(x,aLn,'-',c='b') # print(f' Parabolic = {FUNCTIONALScalc.calculate_simplified_volavg(x,n)[0]:.3f}') @@ -36,13 +36,13 @@ # x, T, n = PRFfunctionals_Hmode( T_avol, n_avol, nu_T, nu_n, aLT = 2.0 ) # axs[0,0].plot(x,T,'-',c='r',label='PRFfunctionals (H-mode)') -# aLT = CALCtools.produceGradient(torch.from_numpy(x),torch.from_numpy(T)).cpu().numpy() +# aLT = CALCtools.derivation_into_Lx(torch.from_numpy(x),torch.from_numpy(T)).cpu().numpy() # axs[1,0].plot(x,aLT,'-',c='r') # print(f' PRF H = {FUNCTIONALScalc.calculate_simplified_volavg(x,T)[0]:.3f}keV') # axs[0,1].plot(x,n,'-',c='r') -# aLn = CALCtools.produceGradient(torch.from_numpy(x),torch.from_numpy(n)).cpu().numpy() +# aLn = CALCtools.derivation_into_Lx(torch.from_numpy(x),torch.from_numpy(n)).cpu().numpy() # axs[1,1].plot(x,aLn,'-',c='r') # print(f' PRF H = {FUNCTIONALScalc.calculate_simplified_volavg(x,n)[0]:.3f}') @@ -52,13 +52,13 @@ print("\t* Took: " + IOtools.getTimeDifference(timeBeginning)) axs[0, 0].plot(x, T, "-", c="g", label="Piecewise linear gradient") -aLT = CALCtools.produceGradient(torch.from_numpy(x), torch.from_numpy(T)).cpu().numpy() +aLT = CALCtools.derivation_into_Lx(torch.from_numpy(x), torch.from_numpy(T)).cpu().numpy() axs[1, 0].plot(x, aLT, "-", c="g") print(f" PRF L = {FUNCTIONALScalc.calculate_simplified_volavg(x,T)[0]:.3f}keV") axs[0, 1].plot(x, n, "-", c="g") -aLn = CALCtools.produceGradient(torch.from_numpy(x), torch.from_numpy(n)).cpu().numpy() +aLn = CALCtools.derivation_into_Lx(torch.from_numpy(x), torch.from_numpy(n)).cpu().numpy() axs[1, 1].plot(x, aLn, "-", c="g") print(f" PRF L = {FUNCTIONALScalc.calculate_simplified_volavg(x,n)[0]:.3f}") From f76d03b86529169d0bab7a4f5571494af9b08286 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 21 Jun 2025 20:45:06 +0200 Subject: [PATCH 089/385] misc --- src/mitim_modules/powertorch/physics_models/parameterizers.py | 4 ++-- src/mitim_modules/powertorch/utils/CALCtools.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/parameterizers.py b/src/mitim_modules/powertorch/physics_models/parameterizers.py index 0f99e5d8..56a4dfa6 100644 --- a/src/mitim_modules/powertorch/physics_models/parameterizers.py +++ b/src/mitim_modules/powertorch/physics_models/parameterizers.py @@ -106,7 +106,7 @@ def profile_constructor_middle(x, y, multiplier=multiplier_quantity): Reason why something like this is not used for the full profile is because derivative of this will not be as original, which is needed to match TGYRO """ - yCPs = CALCtools.Interp1d()(aLy_coarse[:, 0][:-1].repeat((y.shape[0], 1)), y, x) + yCPs = CALCtools.Interp1d_torch()(aLy_coarse[:, 0][:-1].repeat((y.shape[0], 1)), y, x) return x, integrator_function(x, yCPs, y_bc_real) / multiplier def profile_constructor_fine(x, y, multiplier=multiplier_quantity): @@ -124,7 +124,7 @@ def profile_constructor_fine(x, y, multiplier=multiplier_quantity): y = torch.cat((y, aLy_coarse[-1][-1].repeat((y.shape[0], 1))), dim=1) # Model curve (basically, what happens in between points) - yBS = CALCtools.Interp1d()(x.repeat(y.shape[0], 1), y, x_notrail.repeat(y.shape[0], 1)) + yBS = CALCtools.Interp1d_torch()(x.repeat(y.shape[0], 1), y, x_notrail.repeat(y.shape[0], 1)) """ --------------------------------------------------------------------------------------------------------- diff --git a/src/mitim_modules/powertorch/utils/CALCtools.py b/src/mitim_modules/powertorch/utils/CALCtools.py index 7a910d8c..fc0e6a65 100644 --- a/src/mitim_modules/powertorch/utils/CALCtools.py +++ b/src/mitim_modules/powertorch/utils/CALCtools.py @@ -129,7 +129,7 @@ def volume_integration(p, r, volp): ---------------------------------------------------------------------------------------------------------------- """ -class Interp1d(torch.autograd.Function): +class Interp1d_torch(torch.autograd.Function): def __call__(self, x, y, xnew, out=None): return self.forward(x, y, xnew, out) From 6f1b31d35a6c1c12ffbb847c02078b31c772f93e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 22 Jun 2025 15:02:43 +0200 Subject: [PATCH 090/385] Hook method for before and after decorator --- src/mitim_tools/misc_tools/IOtools.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index ac74c0ac..d417f6a7 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -94,6 +94,19 @@ def wrapper_timer(*args, **kwargs): return wrapper_timer return decorator_timer +# Decorator to hook methods before and after execution +def hook_method(before=None, after=None): + def decorator(func): + def wrapper(self, *args, **kwargs): + if before: + before(self) + result = func(self, *args, **kwargs) + if after: + after(self) + return result + return wrapper + return decorator + def clipstr(txt, chars=40): if not isinstance(txt, str): txt = f"{txt}" From 37eddc8ad4d77720dc07e344f383933fa4285af6 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 22 Jun 2025 15:02:59 +0200 Subject: [PATCH 091/385] Ensure variables in profiles as part of decorator --- src/mitim_tools/gacode_tools/PROFILEStools.py | 73 +++++-------------- 1 file changed, 19 insertions(+), 54 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index ab9a64d4..37b1e3ee 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -1,13 +1,13 @@ import copy import numpy as np from collections import OrderedDict -from mitim_tools.plasmastate_tools.MITIMstate import mitim_state +from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.gs_tools import GEQtools -from mitim_tools.misc_tools import MATHtools +from mitim_tools.misc_tools import MATHtools, IOtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -class gacode_state(mitim_state): +class gacode_state(MITIMstate.mitim_state): ''' Class to read and manipulate GACODE profiles files (input.gacode). It inherits from the main MITIMstate class, which provides basic @@ -18,21 +18,24 @@ class gacode_state(mitim_state): ''' # ------------------------------------------------------------------ - # Reading and interpreting + # Reading and interpreting input.gacode files # ------------------------------------------------------------------ - def __init__(self, file, calculateDerived=True, mi_ref=None): + def __init__(self, file, derive_quantities=True, mi_ref=None): + # Initialize the base class and tell it the type of file super().__init__(type_file='input.gacode') + # Read the input file and store the raw data self.file = file - self._read_inputgacocde() + # Derive quantities if requested if self.file is not None: # Derive (Depending on resolution, derived can be expensive, so I mmay not do it every time) - self.derive_quantities(mi_ref=mi_ref, calculateDerived=calculateDerived) + self.derive_quantities(mi_ref=mi_ref, derive_quantities=derive_quantities) + @IOtools.hook_method(after=MITIMstate.ensure_variables_existence) def _read_inputgacocde(self): self.titles_singleNum = ["nexp", "nion", "shot", "name", "type", "time"] @@ -46,7 +49,15 @@ def _read_inputgacocde(self): # Read file and store raw data self._read_header() self._read_profiles() - self._ensure_existence() + + # Ensure correctness (wrong names in older input.gacode files) + if "qmom(Nm)" in self.profiles: + self.profiles["qmom(N/m^2)"] = self.profiles.pop("qmom(Nm)") + if "qpar_beam(MW/m^3)" in self.profiles: + self.profiles["qpar_beam(1/m^3/s)"] = self.profiles.pop("qpar_beam(MW/m^3)") + if "qpar_wall(MW/m^3)" in self.profiles: + self.profiles["qpar_wall(1/m^3/s)"] = self.profiles.pop("qpar_wall(MW/m^3)") + def _read_header(self): for i in range(len(self.lines)): @@ -110,52 +121,6 @@ def _read_profiles(self): self.profiles["w0(rad/s)"] = self.profiles["omega0(rad/s)"] del self.profiles["omega0(rad/s)"] - def _ensure_existence(self): - # Calculate necessary quantities - - if "qpar_beam(MW/m^3)" in self.profiles: - self.varqpar, self.varqpar2 = "qpar_beam(MW/m^3)", "qpar_wall(MW/m^3)" - else: - self.varqpar, self.varqpar2 = "qpar_beam(1/m^3/s)", "qpar_wall(1/m^3/s)" - - if "qmom(Nm)" in self.profiles: - self.varqmom = "qmom(Nm)" # Old, wrong one. But Candy fixed it as of 02/24/2023 - else: - self.varqmom = "qmom(N/m^2)" # CORRECT ONE - - # ------------------------------------------------------------------------------------------------------------------- - # Insert zeros in those cases whose column are not there - # ------------------------------------------------------------------------------------------------------------------- - - some_times_are_not_here = [ - "qei(MW/m^3)", - "qohme(MW/m^3)", - "johm(MA/m^2)", - "jbs(MA/m^2)", - "jbstor(MA/m^2)", - "w0(rad/s)", - "ptot(Pa)", # e.g. if I haven't written that info from ASTRA - "zeta(-)", # e.g. if TGYRO is run with zeta=0, it won't write this column in .new - "zmag(m)", - "qsync(MW/m^3)", - "qbrem(MW/m^3)", - "qline(MW/m^3)", - self.varqpar, - self.varqpar2, - "shape_cos0(-)", - self.varqmom, - ] - - num_moments = 7 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros - for i in range(num_moments): - some_times_are_not_here.append(f"shape_cos{i + 1}(-)") - if i > 1: - some_times_are_not_here.append(f"shape_sin{i + 1}(-)") - - for ikey in some_times_are_not_here: - if ikey not in self.profiles.keys(): - self.profiles[ikey] = copy.deepcopy(self.profiles["rmin(m)"]) * 0.0 - # ------------------------------------------------------------------ # Derivation (different from MITIMstate) # ------------------------------------------------------------------ From 2bfd160a8c950f98649e2c76334975147409b29f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 22 Jun 2025 15:03:11 +0200 Subject: [PATCH 092/385] Clearer naming --- .../portals/utils/PORTALSanalysis.py | 4 +- .../powertorch/utils/CALCtools.py | 1 - src/mitim_tools/gacode_tools/NEOtools.py | 2 +- src/mitim_tools/gacode_tools/TGYROtools.py | 6 +- .../plasmastate_tools/MITIMstate.py | 106 +++++++++++++++--- .../plasmastate_tools/utils/VMECtools.py | 4 +- 6 files changed, 97 insertions(+), 26 deletions(-) diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 12b1e2a1..65c58660 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -183,12 +183,12 @@ def prep_metrics(self, ilast=None): file = self.opt_fun.folder / "Execution" / f"Evaluation.{x_train_num}" / "model_complete" / "input.gacode_unmodified" if file.exists(): print("\t\t- Reading next profile to evaluate (from folder)") - self.profiles_next = PROFILEStools.gacode_state(file, calculateDerived=False) + self.profiles_next = PROFILEStools.gacode_state(file, derive_quantities=False) file = self.opt_fun.folder / "Execution" / f"Evaluation.{x_train_num}" / "model_complete" / "input.gacode.new" if file.exists(): self.profiles_next_new = PROFILEStools.gacode_state( - file, calculateDerived=False + file, derive_quantities=False ) self.profiles_next_new.printInfo(label="NEXT") else: diff --git a/src/mitim_modules/powertorch/utils/CALCtools.py b/src/mitim_modules/powertorch/utils/CALCtools.py index fc0e6a65..6dc0d8ea 100644 --- a/src/mitim_modules/powertorch/utils/CALCtools.py +++ b/src/mitim_modules/powertorch/utils/CALCtools.py @@ -81,7 +81,6 @@ def derivation_into_dxdr(r, p): # Volume calculations # ******************************************************************************************************************** - def _to_2d(x, xp): """Ensure shape (batch, N) for either NumPy or Torch.""" if xp is np: diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index c5f1d15a..990470e4 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -68,7 +68,7 @@ def run_vgen(self, subfolder="vgen1", vgenOptions={}, cold_start=False): from mitim_tools.gacode_tools import PROFILEStools self.inputgacode_vgen = PROFILEStools.gacode_state( - file_new, calculateDerived=True, mi_ref=self.inputgacode.mi_ref + file_new, derive_quantities=True, mi_ref=self.inputgacode.mi_ref ) diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index d2a9a152..5b7fc18c 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -1203,7 +1203,7 @@ def __init__(self, FolderTGYRO, profiles=None): if (profiles is None) and (FolderTGYRO / "input.gacode").exists(): from mitim_tools.gacode_tools import PROFILEStools - profiles = PROFILEStools.gacode_state(FolderTGYRO / f"input.gacode", calculateDerived=False) + profiles = PROFILEStools.gacode_state(FolderTGYRO / f"input.gacode", derive_quantities=False) self.profiles = profiles @@ -1215,9 +1215,9 @@ def __init__(self, FolderTGYRO, profiles=None): self.readNu() self.readProfiles() - calculateDerived = True + derive_quantities = True try: - self.profiles_final = PROFILEStools.gacode_state(self.FolderTGYRO / "input.gacode.new",calculateDerived=calculateDerived,) + self.profiles_final = PROFILEStools.gacode_state(self.FolderTGYRO / "input.gacode.new",derive_quantities=derive_quantities,) except: self.profiles_final = None diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 7fcef888..d53c29d1 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -14,6 +14,85 @@ from mitim_tools import __version__ from IPython import embed +def ensure_variables_existence(self): + # --------------------------------------------------------------------------- + # Determine minimal set of variables that should be present in the profiles + # --------------------------------------------------------------------------- + + # Kinetics + required_profiles = { + "te(keV)": 1, + "ti(keV)": 2, + "ne(10^19/m^3)": 1, + "ni(10^19/m^3)": 2, + "ptot(Pa)": 1, + "zeff(-)": 1, + "w0(rad/s)": 1, + } + + # Electromagnetics + required_profiles.update({ + "q(-)": 1, + "torfluxa(Wb/radian)": 1, + "polflux(Wb/radian)": 1, + "johm(MA/m^2)": 1, + "jbs(MA/m^2)": 1, + "jbstor(MA/m^2)": 1, + }) + + # Geometry + required_profiles.update({ + "rho(-)": 1, + "rmin(m)": 1, + "rmaj(m)": 1, + "zmag(m)": 1, + "rcentr(m)": 1, + "kappa(-)": 1, + "delta(-)": 1, + "zeta(-)": 1, + "shape_cos0(-)": 1, + }) + + num_moments = 7 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros + for i in range(num_moments): + required_profiles[f"shape_cos{i + 1}(-)"] = 1 + if i > 1: + required_profiles[f"shape_sin{i + 1}(-)"] = 1 + + # Sources and Sinks + required_profiles.update({ + "qohme(MW/m^3)": 1, + "qei(MW/m^3)": 1, + "qbeame(MW/m^3)": 1, + "qbeami(MW/m^3)": 1, + "qrfe(MW/m^3)": 1, + "qrfi(MW/m^3)": 1, + "qfuse(MW/m^3)": 1, + "qfusi(MW/m^3)": 1, + "qsync(MW/m^3)": 1, + "qbrem(MW/m^3)": 1, + "qline(MW/m^3)": 1, + "qpar_beam(1/m^3/s)": 1, + "qpar_wall(1/m^3/s)": 1, + "qmom(N/m^2)": 1, + }) + + # --------------------------------------------------------------------------- + # Insert zeros in those cases whose column are not there + # --------------------------------------------------------------------------- + + # Choose a template for dimensionality + template_key_1d = "rmin(m)" + template_key_2d = "ti(keV)" + + template_1d = copy.deepcopy(self.profiles[template_key_1d]) * 0.0 + template_2d = copy.deepcopy(self.profiles[template_key_2d]) * 0.0 + + # Ensure required keys exist + for key, dim in required_profiles.items(): + if key not in self.profiles: + self.profiles[key] = template_1d if dim == 1 else template_2d + class mitim_state: ''' Class to manipulate the plasma state in MITIM. @@ -23,9 +102,9 @@ def __init__(self, type_file = 'input.gacode'): self.type = type_file - def derive_quantities(self, mi_ref=None, calculateDerived=True, rederiveGeometry=True): + def derive_quantities(self, mi_ref=None, derive_quantities=True, rederiveGeometry=True): - # ------------------------------------------------------------------------------------------------------------------- + # ------------------------------------- self.readSpecies() self.mi_first = self.Species[0]["A"] self.DTplasma() @@ -54,7 +133,7 @@ def derive_quantities(self, mi_ref=None, calculateDerived=True, rederiveGeometry self.derived["aLni"] = np.transpose(np.array(self.derived["aLni"])) # ------------------------------------------------------------------------------------------------ - if calculateDerived: + if derive_quantities: self.derive_quantities_full(rederiveGeometry=rederiveGeometry) def derive_geometry(self, **kwargs): @@ -172,10 +251,6 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): deriving geometry is expensive, so if I'm just updating profiles it may not be needed """ - self.varqmom = "qmom(N/m^2)" - if self.varqmom not in self.profiles: - self.profiles[self.varqmom] = self.profiles["rho(-)"] * 0.0 - if "derived" not in self.__dict__: self.derived = {} @@ -315,7 +390,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["qi"] += qi_terms[i] * self.profiles[i] # Depends on GACODE version - ge_terms = {self.varqpar: 1, self.varqpar2: 1} + ge_terms = {"qpar_beam(1/m^3/s)": 1, "qpar_wall(1/m^3/s)": 1} self.derived["ge"] = np.zeros(len(self.profiles["rho(-)"])) for i in ge_terms: @@ -354,7 +429,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): # qmom self.derived["mt_Jmiller"] = CALCtools.volume_integration( - self.profiles[self.varqmom], r, volp + self.profiles["qmom(N/m^2)"], r, volp ) self.derived["mt_Jm2"] = self.derived["mt_Jmiller"] / (volp) @@ -1616,8 +1691,8 @@ def correct(self, options={}, write=False, new_file=None): # If I don't trust the negative particle flux in the core that comes from TRANSP... if ensurePostiveGamma: print("\t\t- Making particle flux always positive", typeMsg="i") - self.profiles[self.varqpar] = self.profiles[self.varqpar].clip(0) - self.profiles[self.varqpar2] = self.profiles[self.varqpar2].clip(0) + self.profiles["qpar_beam(1/m^3/s)"] = self.profiles["qpar_beam(1/m^3/s)"].clip(0) + self.profiles["qpar_wall(1/m^3/s)"] = self.profiles["qpar_wall(1/m^3/s)"].clip(0) # Mach if ensureMachNumber is not None: @@ -2042,11 +2117,8 @@ def plot( GRAPHICStools.autoscale_y(ax) ax = ax01b - if "varqmom" not in self.__dict__: - self.varqmom = "qmom(N/m^2)" - self.profiles[self.varqmom] = self.profiles["rho(-)"] * 0.0 - ax.plot(rho, self.profiles[self.varqmom], lw=lw, ls="-", c=color) + ax.plot(rho, self.profiles["qmom(N/m^2)"], lw=lw, ls="-", c=color) ax.set_xlim([0, 1]) ax.set_xlabel("$\\rho$") ax.set_ylabel("$N/m^2$, $J/m^3$") @@ -2177,9 +2249,9 @@ def plot( ax = ax11b cont = 0 - var = self.profiles[self.varqpar] * 1e-20 + var = self.profiles["qpar_beam(1/m^3/s)"] * 1e-20 ax.plot(rho, var, lw=lw, ls=lines[0], c=color, label=extralab + "beam") - var = self.profiles[self.varqpar2] * 1e-20 + var = self.profiles["qpar_wall(1/m^3/s)"] * 1e-20 ax.plot(rho, var, lw=lw, ls=lines[1], c=color, label=extralab + "wall") ax.set_xlim([0, 1]) diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index 8480922d..206aa902 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -10,7 +10,7 @@ class vmec_state(mitim_state): # Reading and interpreting # ------------------------------------------------------------------ - def __init__(self, file, calculateDerived=True, mi_ref=None): + def __init__(self, file, derive_quantities=True, mi_ref=None): super().__init__(type_file='vmec') @@ -20,7 +20,7 @@ def __init__(self, file, calculateDerived=True, mi_ref=None): if self.file is not None: # Derive (Depending on resolution, derived can be expensive, so I mmay not do it every time) - self.derive_quantities(mi_ref=mi_ref, calculateDerived=calculateDerived) + self.derive_quantities(mi_ref=mi_ref, derive_quantities=derive_quantities) def _read_vmec(self): pass From b525850547c2bc413a9c89b1a622ded7bd618f87 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 22 Jun 2025 17:54:45 +0200 Subject: [PATCH 093/385] Extracted plotting routines from MITIMstate and made clearer division --- .../maestro/utils/MAESTROplot.py | 5 +- .../portals/utils/PORTALSanalysis.py | 8 +- .../portals/utils/PORTALSplot.py | 14 +- src/mitim_modules/powertorch/STATEtools.py | 4 +- .../powertorch/utils/POWERplot.py | 5 +- .../powertorch/utils/TRANSFORMtools.py | 4 +- src/mitim_tools/gacode_tools/PROFILEStools.py | 44 +- .../gacode_tools/scripts/compare_MXH3.py | 2 +- .../gacode_tools/scripts/read_gacode.py | 4 +- src/mitim_tools/gs_tools/GEQtools.py | 2 +- src/mitim_tools/gs_tools/scripts/mxh.py | 2 +- .../plasmastate_tools/MITIMstate.py | 1653 +---------------- .../plasmastate_tools/utils/state_plotting.py | 1445 ++++++++++++++ src/mitim_tools/popcon_tools/RAPIDStools.py | 2 +- .../transp_tools/utils/TRANSPhelpers.py | 2 +- 15 files changed, 1543 insertions(+), 1653 deletions(-) create mode 100644 src/mitim_tools/plasmastate_tools/utils/state_plotting.py diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index 41be692c..e7548aa3 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -3,6 +3,7 @@ from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.misc_tools import LOGtools, GRAPHICStools from mitim_tools.plasmastate_tools import MITIMstate +from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_tools.gs_tools import GEQtools from pathlib import Path from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -97,10 +98,10 @@ def plot_results(self, fn): maxPlot = 5 if len(ps) > 0: # Plot profiles - figs = MITIMstate.add_figures(fn,fnlab_pre = 'MAESTRO - ') + figs = state_plotting.add_figures(fn,fnlab_pre = 'MAESTRO - ') log_file = self.folder_logs/'plot_maestro.log' if (not self.terminal_outputs) else None with LOGtools.conditional_log_to_file(log_file=log_file): - MITIMstate.plotAll(ps[-maxPlot:], extralabs=ps_lab[-maxPlot:], figs=figs) + state_plotting.plotAll(ps[-maxPlot:], extralabs=ps_lab[-maxPlot:], figs=figs) for p,pl in zip(ps,ps_lab): p.printInfo(label = pl) diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 65c58660..8770371c 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -1010,7 +1010,7 @@ def plotMetrics(self, extra_lab="", **kwargs): self.powerstates[i].plot(axs=axs, c=colors[i], label=f"#{i}") # Add profiles too - self.powerstates[i].profiles.plotGradients( + self.powerstates[i].profiles.plot_gradients( axsGrads_extra, color=colors[i], plotImpurity=self.powerstates[-1].impurityPosition if 'nZ' in self.powerstates[-1].ProfilesPredicted else None, @@ -1025,7 +1025,7 @@ def plotMetrics(self, extra_lab="", **kwargs): # Add next profile if len(self.profiles) > len(self.powerstates): - self.profiles[-1].plotGradients( + self.profiles[-1].plot_gradients( axsGrads_extra, color=colors[i+1], plotImpurity=self.powerstates[-1].impurityPosition_transport if 'nZ' in self.powerstates[-1].ProfilesPredicted else None, @@ -1053,7 +1053,7 @@ def plotMetrics(self, extra_lab="", **kwargs): for i in range(2): axsGrads.append(figG.add_subplot(grid[i, j])) for i, p in enumerate(self.powerstates): - p.profiles.plotGradients( + p.profiles.plot_gradients( axsGrads, color=colors[i], plotImpurity=p.impurityPosition if 'nZ' in p.ProfilesPredicted else None, @@ -1065,7 +1065,7 @@ def plotMetrics(self, extra_lab="", **kwargs): if len(self.profiles) > len(self.powerstates): prof = self.profiles[-1] - prof.plotGradients( + prof.plot_gradients( axsGrads, color=colors[i+1], plotImpurity=p.impurityPosition_transport if 'nZ' in p.ProfilesPredicted else None, diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index 728c8639..ae618ba4 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -1,3 +1,4 @@ +from mitim_tools.plasmastate_tools.utils import state_plotting import torch import copy import numpy as np @@ -5,6 +6,7 @@ from mitim_tools.misc_tools import GRAPHICStools from mitim_modules.portals import PORTALStools from mitim_modules.powertorch import STATEtools +from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_modules.powertorch.utils import POWERplot from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -1826,10 +1828,10 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): # Plot PROFILES # ------------------------------------------------------- - figs = MITIMstate.add_figures(fn,fnlab_pre = "PROFILES - ") + figs = state_plotting.add_figures(fn,fnlab_pre = "PROFILES - ") if indecesPlot[0] < len(self.powerstates): - _ = MITIMstate.plotAll( + _ = state_plotting.plotAll( [ self.powerstates[indecesPlot[1]].profiles, self.powerstates[indecesPlot[0]].profiles, @@ -1887,7 +1889,7 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): [0.2, 1.0, 1.0, 1.0], ) ): - profiles.plotGradients( + profiles.plot_gradients( axs4, color=colors[i], label=label, @@ -1955,7 +1957,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): ms = 0 p = self.mitim_runs[self.i0]["powerstate"].profiles - p.plotGradients( + p.plot_gradients( axsR, color="b", lastRho=self.MODELparameters["RhoLocations"][-1], @@ -1972,7 +1974,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): break p = self.mitim_runs[ikey]["powerstate"].profiles - p.plotGradients( + p.plot_gradients( axsR, color="r", lastRho=self.MODELparameters["RhoLocations"][-1], @@ -1984,7 +1986,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): ) p = self.mitim_runs[self.ibest]["powerstate"].profiles - p.plotGradients( + p.plot_gradients( axsR, color="g", lastRho=self.MODELparameters["RhoLocations"][-1], diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index e3e9440f..c116508a 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -7,7 +7,7 @@ import dill as pickle from mitim_tools.misc_tools import PLASMAtools, IOtools from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.plasmastate_tools import MITIMstate +from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_modules.powertorch.utils import TRANSFORMtools, POWERplot from mitim_tools.opt_tools.optimizers import optim from mitim_modules.powertorch.utils import TARGETStools, CALCtools, TRANSPORTtools @@ -450,7 +450,7 @@ def plot(self, axs=None, axsRes=None, axsMetrics=None, figs=None, fn=None,c="r", axsRes.append(figOpt.add_subplot(grid[j, i+1])) # Profiles - figs = MITIMstate.add_figures(fn, tab_color='b') + figs = state_plotting.add_figures(fn, tab_color='b') axs, axsMetrics = add_axes_powerstate_plot(figMain, num_kp = len(self.ProfilesPredicted)) diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index 6d87f477..f97b7fab 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -1,7 +1,8 @@ -from mitim_tools.plasmastate_tools import MITIMstate +from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_tools.misc_tools import GRAPHICStools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +from mitim_tools.plasmastate_tools.utils import state_plotting def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, compare_to_state=None, c_orig = "b"): @@ -15,7 +16,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co profiles_new = self.from_powerstate(insert_highres_powers=True) # Plot the inserted profiles together with the original ones - _ = MITIMstate.plotAll([self.profiles, profiles_new], figs=figs) + _ = state_plotting.plotAll([self.profiles, profiles_new], figs=figs) # ----------------------------------------------------------------------------------------------------------- # ---- Plot plasma state diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 60398143..d45a5773 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -4,7 +4,7 @@ import numpy as np import pandas as pd from mitim_tools.misc_tools import LOGtools, IOtools -from mitim_tools.plasmastate_tools import MITIMstate +from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_modules.powertorch.physics_models import targets_analytic, parameterizers from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __mitimroot__ @@ -548,7 +548,7 @@ def debug_transformation(p, p_new, s): print(f'Profile mean error: {np.mean(err_prof):.2f}%', typeMsg='i' if np.mean(err_prof) < 1e-0 else 'w') print(f'Gradient mean error (ignoring 0.0): {np.mean(err_grad):.2f}%', typeMsg='i' if np.mean(err_grad) < 1e-0 else 'w') - fn = MITIMstate.plotAll([p,p_new],extralabs=['Original','New'],lastRhoGradients=rho[-1].item()+0.01) + fn = state_plotting.plotAll([p,p_new],extralabs=['Original','New'],lastRhoGradients=rho[-1].item()+0.01) axs = fn.figure_handles[3].figure.axes diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 37b1e3ee..09badca0 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -57,6 +57,33 @@ def _read_inputgacocde(self): self.profiles["qpar_beam(1/m^3/s)"] = self.profiles.pop("qpar_beam(MW/m^3)") if "qpar_wall(MW/m^3)" in self.profiles: self.profiles["qpar_wall(1/m^3/s)"] = self.profiles.pop("qpar_wall(MW/m^3)") + """ + Note that in prgen_map_plasmastate, that variable: + expro_qpar_beam(i) = plst_sn_trans(i-1)/dvol + + Note that in prgen_read_plasmastate, that variable: + ! Particle source + err = nf90_inq_varid(ncid,trim('sn_trans'),varid) + err = nf90_get_var(ncid,varid,plst_sn_trans(1:nx-1)) + plst_sn_trans(nx) = 0.0 + + Note that in the plasmastate file, the variable "sn_trans": + + long_name: particle transport (loss) + units: #/sec + component: PLASMA + section: STATE_PROFILES + specification: R|units=#/sec|step*dV sn_trans(~nrho,0:nspec_th) + + So, this means that expro_qpar_beam is in units of #/sec/m^3, meaning that + it is a particle flux DENSITY. It therefore requires volume integral and + divide by surface to produce a flux. + + The units of this qpar_beam column is NOT MW/m^3. In the gacode source codes + they also say that those units are wrong. + + """ + def _read_header(self): @@ -183,19 +210,26 @@ def derive_geometry(self, n_theta_geo=1001): self.profiles["zmag(m)"], cn, sn) - self.derived["R_surface"],self.derived["Z_surface"] = flux_surfaces.R, flux_surfaces.Z + self.derived["R_surface"],self.derived["Z_surface"] = np.array([flux_surfaces.R]), np.array([flux_surfaces.Z]) + + # R and Z have [toroidal, radius, point], to allow for non-axisymmetric cases # ----------------------------------------------- #cross-sectional area of each flux surface - self.derived["surfXS"] = xsec_area_RZ(self.derived["R_surface"],self.derived["Z_surface"]) + self.derived["surfXS"] = xsec_area_RZ(self.derived["R_surface"][0,...],self.derived["Z_surface"][0,...]) - self.derived["R_LF"] = self.derived["R_surface"].max(axis=1) # self.profiles['rmaj(m)'][0]+self.profiles['rmin(m)'] + self.derived["R_LF"] = self.derived["R_surface"][0,...].max(axis=-1) # self.profiles['rmaj(m)'][0]+self.profiles['rmin(m)'] # For Synchrotron self.derived["B_ref"] = np.abs(self.derived["B_unit"] * self.derived["bt_geo"]) - - + """ + surf_geo is truly surface area, but because of the GACODE definitions of flux, + Surf = V' <|grad r|> + Surf_GACODE = V' + """ + self.derived["surfGACODE_geo"] = (self.derived["surf_geo"] / self.derived["gradr_geo"]) + self.derived["surfGACODE_geo"][np.isnan(self.derived["surfGACODE_geo"])] = 0 def calculateGeometricFactors(profiles, n_theta=1001): diff --git a/src/mitim_tools/gacode_tools/scripts/compare_MXH3.py b/src/mitim_tools/gacode_tools/scripts/compare_MXH3.py index 26a01adb..b8928a99 100644 --- a/src/mitim_tools/gacode_tools/scripts/compare_MXH3.py +++ b/src/mitim_tools/gacode_tools/scripts/compare_MXH3.py @@ -25,7 +25,7 @@ g.plotFluxSurfaces(ax=ax, fluxes=ff, rhoPol=False, sqrt=True, color="r", plot1=False) -p.plotGeometry(ax=ax, surfaces_rho=ff, color="b") +p.plot_state_flux_surfaces(ax=ax, surfaces_rho=ff, color="b") ax.set_xlabel("R (m)") ax.set_ylabel("Z (m)") diff --git a/src/mitim_tools/gacode_tools/scripts/read_gacode.py b/src/mitim_tools/gacode_tools/scripts/read_gacode.py index c668a9ea..86c2e508 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_gacode.py +++ b/src/mitim_tools/gacode_tools/scripts/read_gacode.py @@ -1,5 +1,5 @@ import argparse -from mitim_tools.plasmastate_tools import MITIMstate +from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_tools.gacode_tools import PROFILEStools """ @@ -32,7 +32,7 @@ def main(): if not print_only: - fn = MITIMstate.plotAll(profs, lastRhoGradients=rho) + fn = state_plotting.plotAll(profs, lastRhoGradients=rho) fn.show() diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index e27774f7..8d67b764 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -438,7 +438,7 @@ def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH fig, ax = plt.subplots() ff = np.linspace(0, 1, 11) self.plotFluxSurfaces(ax=ax, fluxes=ff, rhoPol=False, sqrt=True, color="r", plot1=False) - p.plotGeometry(ax=ax, surfaces_rho=ff, color="b") + p.plot_state_flux_surfaces(ax=ax, surfaces_rho=ff, color="b") plt.show() return p diff --git a/src/mitim_tools/gs_tools/scripts/mxh.py b/src/mitim_tools/gs_tools/scripts/mxh.py index dd7a0289..0bb969aa 100644 --- a/src/mitim_tools/gs_tools/scripts/mxh.py +++ b/src/mitim_tools/gs_tools/scripts/mxh.py @@ -22,7 +22,7 @@ ff = np.linspace(0, 1, 11) for i, (coeffs_MXH, p) in enumerate(pc.items()): - p.plotGeometry(ax=ax[i], surfaces_rho=ff, color="b") + p.plot_state_flux_surfaces(ax=ax[i], surfaces_rho=ff, color="b") g.plotFluxSurfaces(ax=ax[i], fluxes=ff, rhoPol=False, sqrt=True, color="r", plot1=False) ax[i].set_title(f'coeffs_MXH = {coeffs_MXH}') diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index d53c29d1..8170b053 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -10,6 +10,7 @@ from mitim_tools.gacode_tools.utils import GACODEdefaults from mitim_tools.transp_tools import CDFtools from mitim_tools.transp_tools.utils import TRANSPhelpers +from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __version__ from IPython import embed @@ -25,15 +26,14 @@ def ensure_variables_existence(self): "ti(keV)": 2, "ne(10^19/m^3)": 1, "ni(10^19/m^3)": 2, + "w0(rad/s)": 1, "ptot(Pa)": 1, "zeff(-)": 1, - "w0(rad/s)": 1, } # Electromagnetics required_profiles.update({ "q(-)": 1, - "torfluxa(Wb/radian)": 1, "polflux(Wb/radian)": 1, "johm(MA/m^2)": 1, "jbs(MA/m^2)": 1, @@ -46,7 +46,6 @@ def ensure_variables_existence(self): "rmin(m)": 1, "rmaj(m)": 1, "zmag(m)": 1, - "rcentr(m)": 1, "kappa(-)": 1, "delta(-)": 1, "zeta(-)": 1, @@ -82,16 +81,12 @@ def ensure_variables_existence(self): # --------------------------------------------------------------------------- # Choose a template for dimensionality - template_key_1d = "rmin(m)" - template_key_2d = "ti(keV)" - - template_1d = copy.deepcopy(self.profiles[template_key_1d]) * 0.0 - template_2d = copy.deepcopy(self.profiles[template_key_2d]) * 0.0 + template_key_1d, template_key_2d = "rmin(m)", "ti(keV)" # Ensure required keys exist for key, dim in required_profiles.items(): if key not in self.profiles: - self.profiles[key] = template_1d if dim == 1 else template_2d + self.profiles[key] = copy.deepcopy(self.profiles[template_key_1d]) * 0.0 if dim == 1 else copy.deepcopy(self.profiles[template_key_2d]) * 0.0 class mitim_state: ''' @@ -102,6 +97,7 @@ def __init__(self, type_file = 'input.gacode'): self.type = type_file + @IOtools.hook_method(before=ensure_variables_existence) def derive_quantities(self, mi_ref=None, derive_quantities=True, rederiveGeometry=True): # ------------------------------------- @@ -320,15 +316,6 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): # --------------------------------------------------------------------------------------------------------------------- - """ - surf_geo is truly surface area, but because of the GACODE definitions of flux, - Surf = V' <|grad r|> - Surf_GACODE = V' - """ - - self.derived["surfGACODE_geo"] = (self.derived["surf_geo"] / self.derived["gradr_geo"]) - - self.derived["surfGACODE_geo"][np.isnan(self.derived["surfGACODE_geo"])] = 0 self.derived["c_s"] = PLASMAtools.c_s(self.profiles["te(keV)"], self.derived["mi_ref"]) self.derived["rho_s"] = PLASMAtools.rho_s(self.profiles["te(keV)"], self.derived["mi_ref"], self.derived["B_unit"]) @@ -1782,24 +1769,12 @@ def introduceRotationProfile(self, Mach_LF=1.0, new_file=None): if new_file is not None: self.writeCurrentStatus(file=new_file) + def plot( self, - axs1=None, - axs2=None, - axs3=None, - axs4=None, - axsFlows=None, - axs6=None, - axsImps=None, - color="b", - legYN=True, - extralab="", - fn=None, - fnlab="", - lsFlows="-", - legFlows=True, - showtexts=True, - lastRhoGradients=0.89, + fn=None,fnlab="", + axs1=None, axs2=None, axs3=None, axs4=None, axsFlows=None, axs6=None, axsImps=None, + color="b",legYN=True,extralab="",lsFlows="-",legFlows=True,showtexts=True,lastRhoGradients=0.89, ): if axs1 is None: if fn is None: @@ -1807,1463 +1782,22 @@ def plot( self.fn = FigureNotebook("PROFILES Notebook", geometry="1600x1000") - fig, fig2, fig3, fig4, fig5, fig6, fig7 = add_figures(self.fn, fnlab=fnlab) - - grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) - axs1 = [ - fig.add_subplot(grid[0, 0]), - fig.add_subplot(grid[1, 0]), - fig.add_subplot(grid[2, 0]), - fig.add_subplot(grid[0, 1]), - fig.add_subplot(grid[1, 1]), - fig.add_subplot(grid[2, 1]), - fig.add_subplot(grid[0, 2]), - fig.add_subplot(grid[1, 2]), - fig.add_subplot(grid[2, 2]), - ] - - - grid = plt.GridSpec(3, 2, hspace=0.3, wspace=0.3) - axs2 = [ - fig2.add_subplot(grid[0, 0]), - fig2.add_subplot(grid[0, 1]), - fig2.add_subplot(grid[1, 0]), - fig2.add_subplot(grid[1, 1]), - fig2.add_subplot(grid[2, 0]), - fig2.add_subplot(grid[2, 1]), - ] - - - grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.5) - ax00c = fig3.add_subplot(grid[0, 0]) - axs3 = [ - ax00c, - fig3.add_subplot(grid[1, 0], sharex=ax00c), - fig3.add_subplot(grid[2, 0], sharex=ax00c), - fig3.add_subplot(grid[0, 1], sharex=ax00c), - fig3.add_subplot(grid[1, 1], sharex=ax00c), - fig3.add_subplot(grid[2, 1], sharex=ax00c), - fig3.add_subplot(grid[0, 2], sharex=ax00c), - fig3.add_subplot(grid[1, 2], sharex=ax00c), - fig3.add_subplot(grid[2, 2], sharex=ax00c), - fig3.add_subplot(grid[0, 3], sharex=ax00c), - fig3.add_subplot(grid[1, 3], sharex=ax00c), - fig3.add_subplot(grid[2, 3], sharex=ax00c), - ] - - - grid = plt.GridSpec(2, 3, 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]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - - axsFlows = [ - fig5.add_subplot(grid[0, 0]), - fig5.add_subplot(grid[1, 0]), - fig5.add_subplot(grid[0, 1]), - fig5.add_subplot(grid[0, 2]), - fig5.add_subplot(grid[1, 1]), - fig5.add_subplot(grid[1, 2]), - ] - - - grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) - axs6 = [ - fig6.add_subplot(grid[0, 0]), - fig6.add_subplot(grid[:, 1]), - fig6.add_subplot(grid[0, 2]), - fig6.add_subplot(grid[1, 0]), - fig6.add_subplot(grid[1, 2]), - fig6.add_subplot(grid[0, 3]), - fig6.add_subplot(grid[1, 3]), - ] - - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsImps = [ - fig7.add_subplot(grid[0, 0]), - fig7.add_subplot(grid[0, 1]), - fig7.add_subplot(grid[0, 2]), - fig7.add_subplot(grid[1, 0]), - fig7.add_subplot(grid[1, 1]), - fig7.add_subplot(grid[1, 2]), - ] - - [ax00, ax10, ax20, ax01, ax11, ax21, ax02, ax12, ax22] = axs1 - [ax00b, ax01b, ax10b, ax11b, ax20b, ax21b] = axs2 - [ - ax00c, - ax10c, - ax20c, - ax01c, - ax11c, - ax21c, - ax02c, - ax12c, - ax22c, - ax03c, - ax13c, - ax23c, - ] = axs3 - - lw = 1 - fs = 6 - rho = self.profiles["rho(-)"] - - lines = ["-", "--", "-.", ":", "-", "--", "-."] - - self.plot_temps(ax=ax00, leg=legYN, col=color, lw=lw, fs=fs, extralab=extralab) - self.plot_dens(ax=ax01, leg=legYN, col=color, lw=lw, fs=fs, extralab=extralab) - - ax = ax10 - cont = 0 - for i in range(len(self.Species)): - if self.Species[i]["S"] == "therm": - var = self.profiles["ti(keV)"][:, i] - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - varL = "Thermal $T_i$ (keV)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax20 - cont = 0 - for i in range(len(self.Species)): - if self.Species[i]["S"] == "fast": - var = self.profiles["ti(keV)"][:, i] - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - varL = "Fast $T_i$ (keV)" - ax.plot( - rho, - self.profiles["ti(keV)"][:, 0], - lw=0.5, - ls="-", - alpha=0.5, - c=color, - label=extralab + "$T_{i,1}$", - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax11 - cont = 0 - for i in range(len(self.Species)): - if self.Species[i]["S"] == "therm": - var = self.profiles["ni(10^19/m^3)"][:, i] * 1e-1 - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - varL = "Thermal $n_i$ ($10^{20}/m^3$)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax21 - cont = 0 - for i in range(len(self.Species)): - if self.Species[i]["S"] == "fast": - var = self.profiles["ni(10^19/m^3)"][:, i] * 1e-1 * 1e5 - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - varL = "Fast $n_i$ ($10^{15}/m^3$)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN and cont>0: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax02 - var = self.profiles["w0(rad/s)"] - ax.plot(rho, var, lw=lw, ls="-", c=color) - varL = "$\\omega_{0}$ (rad/s)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax12 - var = self.profiles["ptot(Pa)"] * 1e-6 - ax.plot(rho, var, lw=lw, ls="-", c=color, label=extralab + "ptot") - if "ptot_manual" in self.derived: - ax.plot( - rho, - self.derived["ptot_manual"], - lw=lw, - ls="--", - c=color, - label=extralab + "check", - ) - # ax.plot(rho,np.abs(var-self.derived['ptot_manual']),lw=lw,ls='-.',c=color,label=extralab+'diff') - - ax.plot( - rho, - self.derived["pthr_manual"], - lw=lw, - ls="-.", - c=color, - label=extralab + "check, thrm", - ) - - - varL = "$p$ (MPa)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - # ax.set_ylim(bottom=0) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax00b - varL = "$MW/m^3$" - cont = 0 - var = -self.profiles["qei(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "i->e", c=color) - cont += 1 - if "qrfe(MW/m^3)" in self.profiles: - var = self.profiles["qrfe(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "rf", c=color) - cont += 1 - if "qfuse(MW/m^3)" in self.profiles: - var = self.profiles["qfuse(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "fus", c=color) - cont += 1 - if "qbeame(MW/m^3)" in self.profiles: - var = self.profiles["qbeame(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "beam", c=color) - cont += 1 - if "qione(MW/m^3)" in self.profiles: - var = self.profiles["qione(MW/m^3)"] - ax.plot( - rho, var, lw=lw / 2, ls=lines[cont], label=extralab + "extra", c=color - ) - cont += 1 - if "qohme(MW/m^3)" in self.profiles: - var = self.profiles["qohme(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "ohmic", c=color) - cont += 1 - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - ax.set_title("Electron Power Density") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax01b - - ax.plot(rho, self.profiles["qmom(N/m^2)"], lw=lw, ls="-", c=color) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$N/m^2$, $J/m^3$") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - ax.set_title("Momentum Source Density") - - ax = ax21b - ax.plot( - rho, self.derived["qe_MWm2"], lw=lw, ls="-", label=extralab + "qe", c=color - ) - ax.plot( - rho, self.derived["qi_MWm2"], lw=lw, ls="--", label=extralab + "qi", c=color - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("Heat Flux ($MW/m^2$)") - if legYN: - ax.legend(loc="lower left", fontsize=fs) - ax.set_title("Flux per unit area (gacode: P/V')") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax21b.twinx() - ax.plot( - rho, - self.derived["ge_10E20m2"], - lw=lw, - ls="-.", - label=extralab + "$\\Gamma_e$", - c=color, - ) - ax.set_ylabel("Particle Flux ($10^{20}/m^2/s$)") - if legYN: - ax.legend(loc="lower right", fontsize=fs) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax20b - varL = "$Q_{rad}$ ($MW/m^3$)" - if "qbrem(MW/m^3)" in self.profiles: - var = self.profiles["qbrem(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls="-", label=extralab + "brem", c=color) - if "qline(MW/m^3)" in self.profiles: - var = self.profiles["qline(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls="--", label=extralab + "line", c=color) - if "qsync(MW/m^3)" in self.profiles: - var = self.profiles["qsync(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=":", label=extralab + "sync", c=color) - - var = self.derived["qrad"] - ax.plot(rho, var, lw=lw * 1.5, ls="-", label=extralab + "Total", c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - ax.set_title("Radiation Contributions") - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax10b - varL = "$MW/m^3$" - cont = 0 - var = self.profiles["qei(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "e->i", c=color) - cont += 1 - if "qrfi(MW/m^3)" in self.profiles: - var = self.profiles["qrfi(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "rf", c=color) - cont += 1 - if "qfusi(MW/m^3)" in self.profiles: - var = self.profiles["qfusi(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "fus", c=color) - cont += 1 - if "qbeami(MW/m^3)" in self.profiles: - var = self.profiles["qbeami(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "beam", c=color) - cont += 1 - if "qioni(MW/m^3)" in self.profiles: - var = self.profiles["qioni(MW/m^3)"] - ax.plot( - rho, var, lw=lw / 2, ls=lines[cont], label=extralab + "extra", c=color - ) - cont += 1 - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - ax.set_title("Ion Power Density") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - """ - Note that in prgen_map_plasmastate, that variable: - expro_qpar_beam(i) = plst_sn_trans(i-1)/dvol - - Note that in prgen_read_plasmastate, that variable: - ! Particle source - err = nf90_inq_varid(ncid,trim('sn_trans'),varid) - err = nf90_get_var(ncid,varid,plst_sn_trans(1:nx-1)) - plst_sn_trans(nx) = 0.0 - - Note that in the plasmastate file, the variable "sn_trans": - - long_name: particle transport (loss) - units: #/sec - component: PLASMA - section: STATE_PROFILES - specification: R|units=#/sec|step*dV sn_trans(~nrho,0:nspec_th) - - So, this means that expro_qpar_beam is in units of #/sec/m^3, meaning that - it is a particle flux DENSITY. It therefore requires volume integral and - divide by surface to produce a flux. - - The units of this qpar_beam column is NOT MW/m^3. In the gacode source codes - they also say that those units are wrong. - - """ - - ax = ax11b - cont = 0 - var = self.profiles["qpar_beam(1/m^3/s)"] * 1e-20 - ax.plot(rho, var, lw=lw, ls=lines[0], c=color, label=extralab + "beam") - var = self.profiles["qpar_wall(1/m^3/s)"] * 1e-20 - ax.plot(rho, var, lw=lw, ls=lines[1], c=color, label=extralab + "wall") - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.axhline(y=0, lw=0.5, ls="--", c="k") - ax.set_ylabel("$10^{20}m^{-3}s^{-1}$") - ax.set_title("Particle Source Density") - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) + figs = state_plotting.add_figures(self.fn, fnlab=fnlab) + axs1, axs2, axs3, axs4, axsFlows, axs6, axsImps = state_plotting.add_axes(figs) - ax = ax00c - varL = "cos Shape Params" - yl = 0 - cont = 0 - - for i, s in enumerate(self.shape_cos): - if s is not None: - valmax = np.abs(s).max() - if valmax > 1e-10: - lab = f"c{i}" - ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) - cont += 1 - - yl = np.max([yl, valmax]) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) + lw, fs = 1, 6 + state_plotting.plot_profiles(self,axs1, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) + state_plotting.plot_powers(self,axs2, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) + state_plotting.plot_geometry(self,axs3, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) + state_plotting.plot_gradients(self,axs4, color=color, lw=lw, lastRho=lastRhoGradients, label=extralab) + if axsFlows is not None: + state_plotting.plot_flows(self, axsFlows, ls=lsFlows, leg=legFlows, showtexts=showtexts) + state_plotting.plot_other(self,axs6, color=color, lw=lw, extralab=extralab, fs=fs) + state_plotting.plot_ions(self,axsImps, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - if legYN: - ax.legend(loc="best", fontsize=fs) - - ax = ax01c - varL = "sin Shape Params" - cont = 0 - for i, s in enumerate(self.shape_sin): - if s is not None: - valmax = np.abs(s).max() - if valmax > 1e-10: - lab = f"s{i}" - ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) - cont += 1 - - yl = np.max([yl, valmax]) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax02c - var = self.profiles["q(-)"] - ax.plot(rho, var, lw=lw, ls="-", c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0) - ax.set_ylabel("q") - - ax.axhline(y=1, ls="--", c="k", lw=1) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0.0) - - - ax = ax12c - var = self.profiles["polflux(Wb/radian)"] - ax.plot(rho, var, lw=lw, ls="-", c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("Poloidal $\\psi$ ($Wb/rad$)") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax10c - - var = self.profiles["rho(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0) - ax.set_ylabel("$\\rho$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax11c - - var = self.profiles["rmin(m)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylim(bottom=0) - ax.set_ylabel("$r_{min}$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax20c - - var = self.profiles["rmaj(m)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$R_{maj}$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax21c - - var = self.profiles["zmag(m)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - yl = np.max([0.1, np.max(np.abs(var))]) - ax.set_ylim([-yl, yl]) - ax.set_ylabel("$Z_{maj}$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax22c - - var = self.profiles["kappa(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$\\kappa$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=1) - - ax = ax03c - - var = self.profiles["delta(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$\\delta$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax13c - - var = self.profiles["zeta(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("zeta") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax23c - - var = self.profiles["johm(MA/m^2)"] - ax.plot(rho, var, "-", lw=lw, c=color, label=extralab + "$J_{OH}$") - var = self.profiles["jbs(MA/m^2)"] - ax.plot(rho, var, "--", lw=lw, c=color, label=extralab + "$J_{BS,par}$") - var = self.profiles["jbstor(MA/m^2)"] - ax.plot(rho, var, "-.", lw=lw, c=color, label=extralab + "$J_{BS,tor}$") - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylim(bottom=0) - ax.set_ylabel("J ($MA/m^2$)") - if legYN: - ax.legend(loc="best", prop={"size": 7}) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - # Derived - self.plotGradients( - axs4, color=color, lw=lw, lastRho=lastRhoGradients, label=extralab - ) - - # Others - ax = axs6[0] - ax.plot(self.profiles["rho(-)"], self.derived["dw0dr"] * 1e-5, c=color, lw=lw) - ax.set_ylabel("$-d\\omega_0/dr$ (krad/s/cm)") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - ax.axhline(y=0, lw=1.0, c="k", ls="--") - - ax = axs6[2] - ax.plot(self.profiles["rho(-)"], self.derived["q_fus"], c=color, lw=lw) - ax.set_ylabel("$q_{fus}$ ($MW/m^3$)") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axs6[3] - ax.plot(self.profiles["rho(-)"], self.derived["q_fus_MW"], c=color, lw=lw) - ax.set_ylabel("$P_{fus}$ ($MW$)") - ax.set_xlim([0, 1]) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axs6[4] - ax.plot(self.profiles["rho(-)"], self.derived["tite"], c=color, lw=lw) - ax.set_ylabel("$T_i/T_e$") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - ax.axhline(y=1, ls="--", lw=1.0, c="k") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = axs6[5] - if "MachNum" in self.derived: - ax.plot(self.profiles["rho(-)"], self.derived["MachNum"], c=color, lw=lw) - ax.set_ylabel("Mach Number") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - ax.axhline(y=0, ls="--", c="k", lw=0.5) - ax.axhline(y=1, ls="--", c="k", lw=0.5) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = axs6[6] - safe_division = np.divide( - self.derived["qi_MWm2"], - self.derived["qe_MWm2"], - where=self.derived["qe_MWm2"] != 0, - out=np.full_like(self.derived["qi_MWm2"], np.nan), - ) - ax.plot( - self.profiles["rho(-)"], - safe_division, - c=color, - lw=lw, - label=extralab + "$Q_i/Q_e$", - ) - safe_division = np.divide( - self.derived["qi_aux_MW"], - self.derived["qe_aux_MW"], - where=self.derived["qe_aux_MW"] != 0, - out=np.full_like(self.derived["qi_aux_MW"], np.nan), - ) - ax.plot( - self.profiles["rho(-)"], - safe_division, - c=color, - lw=lw, - ls="--", - label=extralab + "$P_i/P_e$", - ) - ax.set_ylabel("Power ratios") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - ax.axhline(y=1.0, ls="--", c="k", lw=1.0) - GRAPHICStools.addDenseAxis(ax) - # GRAPHICStools.autoscale_y(ax,bottomy=0) - ax.set_ylim(bottom=0) - ax.legend(loc="best", fontsize=fs) - - # Final - if axsFlows is not None: - self.plotBalance( - axs=axsFlows, ls=lsFlows, leg=legFlows, showtexts=showtexts - ) - - # Geometry - ax = axs6[1] - self.plotGeometry(ax=ax, color=color) - - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - GRAPHICStools.addDenseAxis(ax) - - # Impurities - ax = axsImps[0] - for i in range(len(self.Species)): - var = ( - self.profiles["ni(10^19/m^3)"][:, i] - / self.profiles["ni(10^19/m^3)"][0, i] - ) - ax.plot( - rho, - var, - lw=lw, - ls=lines[i], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - varL = "$n_i/n_{i,0}$" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axsImps[1] - for i in range(len(self.Species)): - var = self.derived["fi"][:, i] - ax.plot( - rho, - var, - lw=lw, - ls=lines[i], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - varL = "$f_i$" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - ax.set_ylim([0, 1]) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axsImps[2] - - lastRho = 0.9 - - ix = np.argmin(np.abs(self.profiles["rho(-)"] - lastRho)) + 1 - ax.plot( - rho[:ix], self.derived["aLne"][:ix], lw=lw * 3, ls="-", c=color, label="e" - ) - for i in range(len(self.Species)): - var = self.derived["aLni"][:, i] - ax.plot( - rho[:ix], - var[:ix], - lw=lw, - ls=lines[i], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - varL = "$a/L_{ni}$" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axsImps[5] - - ax = axsImps[3] - ax.plot(self.profiles["rho(-)"], self.derived["Zeff"], c=color, lw=lw) - ax.set_ylabel("$Z_{eff}$") - ax.set_xlabel("$\\rho$") - ax.set_xlim([0, 1]) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axsImps[4] - cont = 0 - if "vtor(m/s)" in self.profiles: - for i in range(len(self.Species)): - try: # REMOVE FOR FUTURE - var = self.profiles["vtor(m/s)"][:, i] * 1e-3 - ax.plot( - rho, - var, - lw=lw, - ls=lines[cont], - c=color, - label=extralab + f"{i + 1} = {self.profiles['name'][i]}", - ) - cont += 1 - except: - break - varL = "$V_{tor}$ (km/s)" - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if "vtor(m/s)" in self.profiles and legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - def plotGradients( - self, - axs4, - color="b", - lw=1.0, - label="", - ls="-o", - lastRho=0.89, - ms=2, - alpha=1.0, - useRoa=False, - RhoLocationsPlot=None, - plotImpurity=None, - plotRotation=False, - autoscale=True, - ): - - if RhoLocationsPlot is None: RhoLocationsPlot=[] - - if axs4 is None: - plt.ion() - fig, axs = plt.subplots( - ncols=3 + int(plotImpurity is not None) + int(plotRotation), - nrows=2, - figsize=(12, 5), - ) - - axs4 = [] - for i in range(axs.shape[-1]): - axs4.append(axs[0, i]) - axs4.append(axs[1, i]) - - ix = np.argmin(np.abs(self.profiles["rho(-)"] - lastRho)) + 1 - - xcoord = self.profiles["rho(-)"] if (not useRoa) else self.derived["roa"] - labelx = "$\\rho$" if (not useRoa) else "$r/a$" - - ax = axs4[0] - ax.plot( - xcoord, - self.profiles["te(keV)"], - ls, - c=color, - lw=lw, - label=label, - markersize=ms, - alpha=alpha, - ) - ax = axs4[2] - ax.plot( - xcoord, - self.profiles["ti(keV)"][:, 0], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - ax = axs4[4] - ax.plot( - xcoord, - self.profiles["ne(10^19/m^3)"] * 1e-1, - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - - if "derived" in self.__dict__: - ax = axs4[1] - ax.plot( - xcoord[:ix], - self.derived["aLTe"][:ix], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - ax = axs4[3] - ax.plot( - xcoord[:ix], - self.derived["aLTi"][:ix, 0], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - ax = axs4[5] - ax.plot( - xcoord[:ix], - self.derived["aLne"][:ix], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - - for ax in axs4: - ax.set_xlim([0, 1]) - - ax = axs4[0] - ax.set_ylabel("$T_e$ (keV)") - ax.set_xlabel(labelx) - if autoscale: - 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) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax = axs4[4] - ax.set_ylabel("$n_e$ ($10^{20}m^{-3}$)") - ax.set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axs4[1] - ax.set_ylabel("$a/L_{Te}$") - ax.set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax = axs4[3] - ax.set_ylabel("$a/L_{Ti}$") - ax.set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax = axs4[5] - ax.set_ylabel("$a/L_{ne}$") - ax.axhline(y=0, ls="--", lw=0.5, c="k") - ax.set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - - cont = 0 - if plotImpurity is not None: - axs4[6 + cont].plot( - xcoord, - self.profiles["ni(10^19/m^3)"][:, plotImpurity] * 1e-1, - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - axs4[6 + cont].set_ylabel("$n_Z$ ($10^{20}m^{-3}$)") - axs4[6].set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - if "derived" in self.__dict__: - axs4[7 + cont].plot( - xcoord[:ix], - self.derived["aLni"][:ix, plotImpurity], - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - axs4[7 + cont].set_ylabel("$a/L_{nZ}$") - axs4[7 + cont].axhline(y=0, ls="--", lw=0.5, c="k") - axs4[7 + cont].set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - cont += 2 - - if plotRotation: - axs4[6 + cont].plot( - xcoord, - self.profiles["w0(rad/s)"] * 1e-3, - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - axs4[6 + cont].set_ylabel("$w_0$ (krad/s)") - axs4[6 + cont].set_xlabel(labelx) - if "derived" in self.__dict__: - axs4[7 + cont].plot( - xcoord[:ix], - self.derived["dw0dr"][:ix] * 1e-5, - ls, - c=color, - lw=lw, - markersize=ms, - alpha=alpha, - ) - axs4[7 + cont].set_ylabel("-$d\\omega_0/dr$ (krad/s/cm)") - axs4[7 + cont].axhline(y=0, ls="--", lw=0.5, c="k") - axs4[7 + cont].set_xlabel(labelx) - if autoscale: - GRAPHICStools.autoscale_y(ax, bottomy=0) - cont += 2 - - for x0 in RhoLocationsPlot: - ix = np.argmin(np.abs(self.profiles["rho(-)"] - x0)) - for ax in axs4: - ax.axvline(x=xcoord[ix], ls="--", lw=0.5, c=color) - - for i in range(len(axs4)): - ax = axs4[i] - GRAPHICStools.addDenseAxis(ax) - - def plotBalance(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): - if axs is None: - fig1 = plt.figure() - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - - axs = [ - fig1.add_subplot(grid[0, 0]), - fig1.add_subplot(grid[1, 0]), - fig1.add_subplot(grid[0, 1]), - fig1.add_subplot(grid[0, 2]), - fig1.add_subplot(grid[1, 1]), - fig1.add_subplot(grid[1, 2]), - ] - - # Profiles - - ax = axs[0] - axT = axs[1] - roa = self.profiles["rmin(m)"] / self.profiles["rmin(m)"][-1] - Te = self.profiles["te(keV)"] - ne = self.profiles["ne(10^19/m^3)"] * 1e-1 - ni = self.profiles["ni(10^19/m^3)"] * 1e-1 - niT = np.sum(ni, axis=1) - Ti = self.profiles["ti(keV)"][:, 0] - ax.plot(roa, Te, lw=2, c="r", label="$T_e$" if leg else "", ls=ls) - ax.plot(roa, Ti, lw=2, c="b", label="$T_i$" if leg else "", ls=ls) - axT.plot(roa, ne, lw=2, c="m", label="$n_e$" if leg else "", ls=ls) - axT.plot(roa, niT, lw=2, c="c", label="$\\sum n_i$" if leg else "", ls=ls) - if limits is not None: - [roa_first, roa_last] = limits - ax.plot(roa_last, np.interp(roa_last, roa, Te), "s", c="r", markersize=3) - ax.plot(roa_first, np.interp(roa_first, roa, Te), "s", c="r", markersize=3) - ax.plot(roa_last, np.interp(roa_last, roa, Ti), "s", c="b", markersize=3) - ax.plot(roa_first, np.interp(roa_first, roa, Ti), "s", c="b", markersize=3) - axT.plot(roa_last, np.interp(roa_last, roa, ne), "s", c="m", markersize=3) - axT.plot(roa_first, np.interp(roa_first, roa, ne), "s", c="m", markersize=3) - - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - axT.set_xlabel("r/a") - axT.set_xlim([0, 1]) - ax.set_ylabel("$T$ (keV)") - ax.set_ylim(bottom=0) - axT.set_ylabel("$n$ ($10^{20}m^{-3}$)") - axT.set_ylim(bottom=0) - # axT.set_ylim([0,np.max(ne)*1.5]) - ax.legend() - axT.legend() - ax.set_title("Final Temperature profiles") - axT.set_title("Final Density profiles") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - GRAPHICStools.addDenseAxis(axT) - GRAPHICStools.autoscale_y(axT, bottomy=0) - - if showtexts: - if self.derived["Q"] > 0.005: - ax.text( - 0.05, - 0.05, - f"Pfus = {self.derived['Pfus']:.1f}MW, Q = {self.derived['Q']:.2f}", - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - transform=ax.transAxes, - ) - - axT.text( - 0.05, - 0.4, - "ne_20 = {0:.1f} (fG = {1:.2f}), Zeff = {2:.1f}".format( - self.derived["ne_vol20"], - self.derived["fG"], - self.derived["Zeff_vol"], - ), - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - transform=axT.transAxes, - ) - - # F - ax = axs[2] - P = ( - self.derived["qe_fus_MW"] - + self.derived["qe_aux_MW"] - + -self.derived["qe_rad_MW"] - + -self.derived["qe_exc_MW"] - ) - - ax.plot( - roa, - -self.derived["qe_MW"], - c="g", - lw=2, - label="$P_{e}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qe_fus_MW"], - c="r", - lw=2, - label="$P_{fus,e}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qe_aux_MW"], - c="b", - lw=2, - label="$P_{aux,e}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - -self.derived["qe_exc_MW"], - c="m", - lw=2, - label="$P_{exc,e}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - -self.derived["qe_rad_MW"], - c="c", - lw=2, - label="$P_{rad,e}$" if leg else "", - ls=ls, - ) - ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) - - # Pe = self.profiles['te(keV)']*1E3*e_J*self.profiles['ne(10^19/m^3)']*1E-1*1E20 *1E-6 - # ax.plot(roa,Pe,ls='-',lw=3,alpha=0.1,c='k',label='$W_e$ (MJ/m^3)') - - ax.plot( - roa, - -self.derived["ce_MW"], - c="k", - lw=1, - label="($P_{conv,e}$)" if leg else "", - ) - - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - ax.set_ylabel("$P$ (MW)") - # ax.set_ylim(bottom=0) - ax.set_title("Electron Thermal Flows") - ax.axhline(y=0.0, lw=0.5, ls="--", c="k") - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" - ) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = axs[3] - P = ( - self.derived["qi_fus_MW"] - + self.derived["qi_aux_MW"] - + self.derived["qe_exc_MW"] - ) - - ax.plot( - roa, - -self.derived["qi_MW"], - c="g", - lw=2, - label="$P_{i}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qi_fus_MW"], - c="r", - lw=2, - label="$P_{fus,i}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qi_aux_MW"], - c="b", - lw=2, - label="$P_{aux,i}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qe_exc_MW"], - c="m", - lw=2, - label="$P_{exc,i}$" if leg else "", - ls=ls, - ) - ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) - - # Pi = self.profiles['ti(keV)'][:,0]*1E3*e_J*self.profiles['ni(10^19/m^3)'][:,0]*1E-1*1E20 *1E-6 - # ax.plot(roa,Pi,ls='-',lw=3,alpha=0.1,c='k',label='$W_$ (MJ/m^3)') - - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - ax.set_ylabel("$P$ (MW)") - # ax.set_ylim(bottom=0) - ax.set_title("Ion Thermal Flows") - ax.axhline(y=0.0, lw=0.5, ls="--", c="k") - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" - ) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - # F - ax = axs[4] - - ax.plot( - roa, - self.derived["ge_10E20"], - c="g", - lw=2, - label="$\\Gamma_{e}$" if leg else "", - ls=ls, - ) - # ax.plot(roa,self.profiles['ne(10^19/m^3)']*1E-1,lw=3,alpha=0.1,c='k',label='$n_e$ ($10^{20}/m^3$)' if leg else '',ls=ls) - - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - ax.set_ylabel("$\\Gamma$ ($10^{20}/s$)") - ax.set_title("Particle Flows") - ax.axhline(y=0.0, lw=0.5, ls="--", c="k") - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" - ) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - # TOTAL - ax = axs[5] - P = ( - self.derived["qOhm_MW"] - + self.derived["qRF_MW"] - + self.derived["qFus_MW"] - + -self.derived["qe_rad_MW"] - + self.derived["qz_MW"] - + self.derived["qBEAM_MW"] - ) - - ax.plot( - roa, - -self.derived["q_MW"], - c="g", - lw=2, - label="$P$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qOhm_MW"], - c="k", - lw=2, - label="$P_{Oh}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qRF_MW"], - c="b", - lw=2, - label="$P_{RF}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qBEAM_MW"], - c="pink", - lw=2, - label="$P_{NBI}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qFus_MW"], - c="r", - lw=2, - label="$P_{fus}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - -self.derived["qe_rad_MW"], - c="c", - lw=2, - label="$P_{rad}$" if leg else "", - ls=ls, - ) - ax.plot( - roa, - self.derived["qz_MW"], - c="orange", - lw=1, - label="$P_{ionz.}$" if leg else "", - ls=ls, - ) - - # P = Pe+Pi - # ax.plot(roa,P,ls='-',lw=3,alpha=0.1,c='k',label='$W$ (MJ)') - - ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) - ax.set_xlabel("r/a") - ax.set_xlim([0, 1]) - ax.set_ylabel("$P$ (MW)") - # ax.set_ylim(bottom=0) - ax.set_title("Total Thermal Flows") - - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" - ) - - ax.axhline(y=0.0, lw=0.5, ls="--", c="k") - # GRAPHICStools.drawLineWithTxt(ax,0.0,label='',orientation='vertical',color='k',lw=1,ls='--',alpha=1.0,fontsize=10,fromtop=0.85,fontweight='normal', - # verticalalignment='bottom',horizontalalignment='left',separation=0) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - def plot_temps(self, ax=None, leg=False, col="b", lw=2, extralab="", fs=10): - if ax is None: - fig, ax = plt.subplots() - - rho = self.profiles["rho(-)"] - - var = self.profiles["te(keV)"] - varL = "$T_e$ , $T_i$ (keV)" - if leg: - lab = extralab + "e" - else: - lab = "" - ax.plot(rho, var, lw=lw, ls="-", label=lab, c=col) - var = self.profiles["ti(keV)"][:, 0] - if leg: - lab = extralab + "i" - else: - lab = "" - ax.plot(rho, var, lw=lw, ls="--", label=lab, c=col) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - - def plot_dens(self, ax=None, leg=False, col="b", lw=2, extralab="", fs=10): - if ax is None: - fig, ax = plt.subplots() - - rho = self.profiles["rho(-)"] - - var = self.profiles["ne(10^19/m^3)"] * 1e-1 - varL = "$n_e$ ($10^{20}/m^3$)" - if leg: - lab = extralab + "e" - else: - lab = "" - ax.plot(rho, var, lw=lw, ls="-", label=lab, c=col) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - - def plotGeometry(self, ax=None, surfaces_rho=np.linspace(0, 1, 11), color="b", label = '', lw=1.0, lw1=2.0): + def plot_state_flux_surfaces(self, ax=None, surfaces_rho=np.linspace(0, 1, 11), color="b", label = '', lw=1.0, lw1=2.0): if ("R_surface" in self.derived) and (self.derived["R_surface"] is not None): if ax is None: plt.ion() @@ -3275,13 +1809,14 @@ def plotGeometry(self, ax=None, surfaces_rho=np.linspace(0, 1, 11), color="b", l for rho in surfaces_rho: ir = np.argmin(np.abs(self.profiles["rho(-)"] - rho)) - ax.plot( - self.derived["R_surface"][ir, :], - self.derived["Z_surface"][ir, :], - "-", - lw=lw if rho<1.0 else lw1, - c=color, - ) + for i_toroidal in range(self.derived["R_surface"].shape[0]): + ax.plot( + self.derived["R_surface"][i_toroidal,ir, :], + self.derived["Z_surface"][i_toroidal,ir, :], + "-", + lw=lw if rho<1.0 else lw1, + c=color, + ) ax.axhline(y=0, ls="--", lw=0.2, c="k") ax.plot( @@ -3387,7 +1922,7 @@ def plotRelevant(self, axs = None, color = 'b', label ='', lw = 1, ms = 1): ax = axs[0] rho = np.linspace(0, 1, 21) - self.plotGeometry(ax=ax, surfaces_rho=rho, label=label, color=color, lw=lw, lw1=lw*3) + self.plot_state_flux_surfaces(ax=ax, surfaces_rho=rho, label=label, color=color, lw=lw, lw1=lw*3) ax.set_xlabel("R (m)") ax.set_ylabel("Z (m)") @@ -3527,7 +2062,6 @@ def plotRelevant(self, axs = None, color = 'b', label ='', lw = 1, ms = 1): GRAPHICStools.addDenseAxis(ax) ax.set_title("Dynamic Targets") - def csv(self, file="input.gacode.xlsx"): dictExcel = IOtools.OrderedDict() @@ -3910,121 +2444,6 @@ def export_to_csv(self, filename, title=None): for row in self.data: writer.writerow(row) -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: - from mitim_tools.misc_tools.GUItools import FigureNotebook - - fn = FigureNotebook("Profiles", geometry="1800x900") - figProf_1, figProf_2, figProf_3, figProf_4, figFlows, figProf_6, fig7 = add_figures(fn) - - grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) - axsProf_1 = [ - figProf_1.add_subplot(grid[0, 0]), - figProf_1.add_subplot(grid[1, 0]), - figProf_1.add_subplot(grid[2, 0]), - figProf_1.add_subplot(grid[0, 1]), - figProf_1.add_subplot(grid[1, 1]), - figProf_1.add_subplot(grid[2, 1]), - figProf_1.add_subplot(grid[0, 2]), - figProf_1.add_subplot(grid[1, 2]), - figProf_1.add_subplot(grid[2, 2]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsProf_2 = [ - figProf_2.add_subplot(grid[0, 0]), - figProf_2.add_subplot(grid[0, 1]), - figProf_2.add_subplot(grid[1, 0]), - figProf_2.add_subplot(grid[1, 1]), - figProf_2.add_subplot(grid[0, 2]), - figProf_2.add_subplot(grid[1, 2]), - ] - grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.3) - ax00c = figProf_3.add_subplot(grid[0, 0]) - axsProf_3 = [ - ax00c, - figProf_3.add_subplot(grid[1, 0], sharex=ax00c), - figProf_3.add_subplot(grid[2, 0]), - figProf_3.add_subplot(grid[0, 1]), - figProf_3.add_subplot(grid[1, 1]), - figProf_3.add_subplot(grid[2, 1]), - figProf_3.add_subplot(grid[0, 2]), - figProf_3.add_subplot(grid[1, 2]), - figProf_3.add_subplot(grid[2, 2]), - figProf_3.add_subplot(grid[0, 3]), - figProf_3.add_subplot(grid[1, 3]), - figProf_3.add_subplot(grid[2, 3]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsProf_4 = [ - figProf_4.add_subplot(grid[0, 0]), - figProf_4.add_subplot(grid[1, 0]), - figProf_4.add_subplot(grid[0, 1]), - figProf_4.add_subplot(grid[1, 1]), - figProf_4.add_subplot(grid[0, 2]), - figProf_4.add_subplot(grid[1, 2]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsFlows = [ - figFlows.add_subplot(grid[0, 0]), - figFlows.add_subplot(grid[1, 0]), - figFlows.add_subplot(grid[0, 1]), - figFlows.add_subplot(grid[0, 2]), - figFlows.add_subplot(grid[1, 1]), - figFlows.add_subplot(grid[1, 2]), - ] - - grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) - axsProf_6 = [ - figProf_6.add_subplot(grid[0, 0]), - figProf_6.add_subplot(grid[:, 1]), - figProf_6.add_subplot(grid[0, 2]), - figProf_6.add_subplot(grid[1, 0]), - figProf_6.add_subplot(grid[1, 2]), - figProf_6.add_subplot(grid[0, 3]), - figProf_6.add_subplot(grid[1, 3]), - ] - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsImps = [ - fig7.add_subplot(grid[0, 0]), - fig7.add_subplot(grid[0, 1]), - fig7.add_subplot(grid[0, 2]), - fig7.add_subplot(grid[1, 0]), - fig7.add_subplot(grid[1, 1]), - fig7.add_subplot(grid[1, 2]), - ] - - ls = GRAPHICStools.listLS() - colors = GRAPHICStools.listColors() - for i, profiles in enumerate(profiles_list): - if extralabs is None: - extralab = f"#{i}, " - else: - extralab = f"{extralabs[i]}, " - profiles.plot( - axs1=axsProf_1, - axs2=axsProf_2, - axs3=axsProf_3, - axs4=axsProf_4, - axsFlows=axsFlows, - axs6=axsProf_6, - axsImps=axsImps, - color=colors[i], - legYN=True, - extralab=extralab, - lsFlows=ls[i], - legFlows=i == 0, - showtexts=False, - lastRhoGradients=lastRhoGradients, - ) - - return fn - def readTGYRO_profile_extra(file, varLabel="B_unit (T)"): with open(file) as f: @@ -4176,18 +2595,6 @@ def gradientsMerger(p0, p_true, roa=0.46, blending=0.1): return p -def add_figures(fn, fnlab='', fnlab_pre='', tab_color=None): - - figProf_1 = fn.add_figure(label= fnlab_pre + "Profiles" + fnlab, tab_color=tab_color) - figProf_2 = fn.add_figure(label= fnlab_pre + "Powers" + fnlab, tab_color=tab_color) - figProf_3 = fn.add_figure(label= fnlab_pre + "Geometry" + fnlab, tab_color=tab_color) - figProf_4 = fn.add_figure(label= fnlab_pre + "Gradients" + fnlab, tab_color=tab_color) - figFlows = fn.add_figure(label= fnlab_pre + "Flows" + fnlab, tab_color=tab_color) - figProf_6 = fn.add_figure(label= fnlab_pre + "Other" + fnlab, tab_color=tab_color) - fig7 = fn.add_figure(label= fnlab_pre + "Impurities" + fnlab, tab_color=tab_color) - figs = [figProf_1, figProf_2, figProf_3, figProf_4, figFlows, figProf_6, fig7] - - return figs def impurity_location(profiles, impurity_of_interest): diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py new file mode 100644 index 00000000..60ed078d --- /dev/null +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -0,0 +1,1445 @@ +import numpy as np +import matplotlib.pyplot as plt +from mitim_tools.misc_tools import GRAPHICStools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from mitim_tools import __version__ +from IPython import embed + + +def add_figures(fn, fnlab='', fnlab_pre='', tab_color=None): + + fig1 = fn.add_figure(label= fnlab_pre + "Profiles" + fnlab, tab_color=tab_color) + fig2 = fn.add_figure(label= fnlab_pre + "Powers" + fnlab, tab_color=tab_color) + fig3 = fn.add_figure(label= fnlab_pre + "Geometry" + fnlab, tab_color=tab_color) + fig4 = fn.add_figure(label= fnlab_pre + "Gradients" + fnlab, tab_color=tab_color) + fig5 = fn.add_figure(label= fnlab_pre + "Flows" + fnlab, tab_color=tab_color) + fig6 = fn.add_figure(label= fnlab_pre + "Other" + fnlab, tab_color=tab_color) + fig7 = fn.add_figure(label= fnlab_pre + "Impurities" + fnlab, tab_color=tab_color) + figs = [fig1, fig2, fig3, fig4, fig5, fig6, fig7] + + return figs + + +def add_axes(figs): + + fig1, fig2, fig3, fig4, fig5, fig6, fig7 = figs + + grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) + axsProf_1 = [ + fig1.add_subplot(grid[0, 0]), + fig1.add_subplot(grid[1, 0]), + fig1.add_subplot(grid[2, 0]), + fig1.add_subplot(grid[0, 1]), + fig1.add_subplot(grid[1, 1]), + fig1.add_subplot(grid[2, 1]), + fig1.add_subplot(grid[0, 2]), + fig1.add_subplot(grid[1, 2]), + fig1.add_subplot(grid[2, 2]), + ] + + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + axsProf_2 = [ + fig2.add_subplot(grid[0, 0]), + fig2.add_subplot(grid[0, 1]), + fig2.add_subplot(grid[1, 0]), + fig2.add_subplot(grid[1, 1]), + fig2.add_subplot(grid[0, 2]), + fig2.add_subplot(grid[1, 2]), + ] + grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.3) + ax00c = fig3.add_subplot(grid[0, 0]) + axsProf_3 = [ + ax00c, + fig3.add_subplot(grid[1, 0], sharex=ax00c), + fig3.add_subplot(grid[2, 0]), + fig3.add_subplot(grid[0, 1]), + fig3.add_subplot(grid[1, 1]), + fig3.add_subplot(grid[2, 1]), + fig3.add_subplot(grid[0, 2]), + fig3.add_subplot(grid[1, 2]), + fig3.add_subplot(grid[2, 2]), + fig3.add_subplot(grid[:, 3]), + ] + + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + axsProf_4 = [ + 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]), + ] + + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + axsFlows = [ + fig5.add_subplot(grid[0, 0]), + fig5.add_subplot(grid[1, 0]), + fig5.add_subplot(grid[0, 1]), + fig5.add_subplot(grid[0, 2]), + fig5.add_subplot(grid[1, 1]), + fig5.add_subplot(grid[1, 2]), + ] + + grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) + axsProf_6 = [ + fig6.add_subplot(grid[0, 0]), + fig6.add_subplot(grid[0, 1]), + fig6.add_subplot(grid[0, 2]), + fig6.add_subplot(grid[1, 0]), + fig6.add_subplot(grid[1, 1]), + fig6.add_subplot(grid[1, 2]), + fig6.add_subplot(grid[0, 3]), + fig6.add_subplot(grid[1, 3]), + ] + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + axsImps = [ + fig7.add_subplot(grid[0, 0]), + fig7.add_subplot(grid[0, 1]), + fig7.add_subplot(grid[0, 2]), + fig7.add_subplot(grid[1, 0]), + fig7.add_subplot(grid[1, 1]), + fig7.add_subplot(grid[1, 2]), + ] + + return axsProf_1, axsProf_2, axsProf_3, axsProf_4, axsFlows, axsProf_6, axsImps + + +def plot_profiles(self, axs1, color="b", legYN=True, extralab="", lw=1, fs=6): + + [ax00, ax10, ax20, ax01, ax11, ax21, ax02, ax12, ax22] = axs1 + + rho = self.profiles["rho(-)"] + + lines = GRAPHICStools.listLS() + + ax=ax00 + var = self.profiles["te(keV)"] + varL = "$T_e$ , $T_i$ (keV)" + if legYN: + lab = extralab + "e" + else: + lab = "" + ax.plot(rho, var, lw=lw, ls="-", label=lab, c=color) + var = self.profiles["ti(keV)"][:, 0] + if legYN: + lab = extralab + "i" + else: + lab = "" + ax.plot(rho, var, lw=lw, ls="--", label=lab, c=color) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + + + ax=ax01 + var = self.profiles["ne(10^19/m^3)"] * 1e-1 + varL = "$n_e$ ($10^{20}/m^3$)" + if legYN: + lab = extralab + "e" + else: + lab = "" + ax.plot(rho, var, lw=lw, ls="-", label=lab, c=color) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + + ax = ax10 + cont = 0 + for i in range(len(self.Species)): + if self.Species[i]["S"] == "therm": + var = self.profiles["ti(keV)"][:, i] + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + varL = "Thermal $T_i$ (keV)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax20 + cont = 0 + for i in range(len(self.Species)): + if self.Species[i]["S"] == "fast": + var = self.profiles["ti(keV)"][:, i] + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + varL = "Fast $T_i$ (keV)" + ax.plot( + rho, + self.profiles["ti(keV)"][:, 0], + lw=0.5, + ls="-", + alpha=0.5, + c=color, + label=extralab + "$T_{i,1}$", + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax11 + cont = 0 + for i in range(len(self.Species)): + if self.Species[i]["S"] == "therm": + var = self.profiles["ni(10^19/m^3)"][:, i] * 1e-1 + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + varL = "Thermal $n_i$ ($10^{20}/m^3$)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax21 + cont = 0 + for i in range(len(self.Species)): + if self.Species[i]["S"] == "fast": + var = self.profiles["ni(10^19/m^3)"][:, i] * 1e-1 * 1e5 + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + varL = "Fast $n_i$ ($10^{15}/m^3$)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN and cont>0: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax02 + var = self.profiles["w0(rad/s)"] + ax.plot(rho, var, lw=lw, ls="-", c=color) + varL = "$\\omega_{0}$ (rad/s)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax12 + var = self.profiles["ptot(Pa)"] * 1e-6 + ax.plot(rho, var, lw=lw, ls="-", c=color, label=extralab + "ptot") + if "ptot_manual" in self.derived: + ax.plot( + rho, + self.derived["ptot_manual"], + lw=lw, + ls="--", + c=color, + label=extralab + "check", + ) + ax.plot( + rho, + self.derived["pthr_manual"], + lw=lw, + ls="-.", + c=color, + label=extralab + "check, thrm", + ) + + varL = "$p$ (MPa)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + # ax.set_ylim(bottom=0) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax22 + var = self.profiles["q(-)"] + ax.plot(rho, var, lw=lw, ls="-", c=color) + varL = "$q$ profile" + ax.axhline(y=1.0, lw=0.5, ls="--", c="k") + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + +def plot_powers(self, axs2, legYN=True, extralab="", color="b", lw=1, fs=6): + + [ax00b, ax01b, ax10b, ax11b, ax20b, ax21b] = axs2 + + rho = self.profiles["rho(-)"] + + lines = GRAPHICStools.listLS() + + ax = ax00b + varL = "$MW/m^3$" + cont = 0 + var = -self.profiles["qei(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "i->e", c=color) + cont += 1 + if "qrfe(MW/m^3)" in self.profiles: + var = self.profiles["qrfe(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "rf", c=color) + cont += 1 + if "qfuse(MW/m^3)" in self.profiles: + var = self.profiles["qfuse(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "fus", c=color) + cont += 1 + if "qbeame(MW/m^3)" in self.profiles: + var = self.profiles["qbeame(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "beam", c=color) + cont += 1 + if "qione(MW/m^3)" in self.profiles: + var = self.profiles["qione(MW/m^3)"] + ax.plot( + rho, var, lw=lw / 2, ls=lines[cont], label=extralab + "extra", c=color + ) + cont += 1 + if "qohme(MW/m^3)" in self.profiles: + var = self.profiles["qohme(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "ohmic", c=color) + cont += 1 + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + ax.set_title("Electron Power Density") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax01b + + ax.plot(rho, self.profiles["qmom(N/m^2)"], lw=lw, ls="-", c=color) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$N/m^2$, $J/m^3$") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + ax.set_title("Momentum Source Density") + + ax = ax21b + ax.plot( + rho, self.derived["qe_MWm2"], lw=lw, ls="-", label=extralab + "qe", c=color + ) + ax.plot( + rho, self.derived["qi_MWm2"], lw=lw, ls="--", label=extralab + "qi", c=color + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("Heat Flux ($MW/m^2$)") + if legYN: + ax.legend(loc="lower left", fontsize=fs) + ax.set_title("Flux per unit area (gacode: P/V')") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax21b.twinx() + ax.plot( + rho, + self.derived["ge_10E20m2"], + lw=lw, + ls="-.", + label=extralab + "$\\Gamma_e$", + c=color, + ) + ax.set_ylabel("Particle Flux ($10^{20}/m^2/s$)") + if legYN: + ax.legend(loc="lower right", fontsize=fs) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax20b + varL = "$Q_{rad}$ ($MW/m^3$)" + if "qbrem(MW/m^3)" in self.profiles: + var = self.profiles["qbrem(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls="-", label=extralab + "brem", c=color) + if "qline(MW/m^3)" in self.profiles: + var = self.profiles["qline(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls="--", label=extralab + "line", c=color) + if "qsync(MW/m^3)" in self.profiles: + var = self.profiles["qsync(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=":", label=extralab + "sync", c=color) + + var = self.derived["qrad"] + ax.plot(rho, var, lw=lw * 1.5, ls="-", label=extralab + "Total", c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + ax.set_title("Radiation Contributions") + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax10b + varL = "$MW/m^3$" + cont = 0 + var = self.profiles["qei(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "e->i", c=color) + cont += 1 + if "qrfi(MW/m^3)" in self.profiles: + var = self.profiles["qrfi(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "rf", c=color) + cont += 1 + if "qfusi(MW/m^3)" in self.profiles: + var = self.profiles["qfusi(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "fus", c=color) + cont += 1 + if "qbeami(MW/m^3)" in self.profiles: + var = self.profiles["qbeami(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=lines[cont], label=extralab + "beam", c=color) + cont += 1 + if "qioni(MW/m^3)" in self.profiles: + var = self.profiles["qioni(MW/m^3)"] + ax.plot( + rho, var, lw=lw / 2, ls=lines[cont], label=extralab + "extra", c=color + ) + cont += 1 + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + ax.set_title("Ion Power Density") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax11b + cont = 0 + var = self.profiles["qpar_beam(1/m^3/s)"] * 1e-20 + ax.plot(rho, var, lw=lw, ls=lines[0], c=color, label=extralab + "beam") + var = self.profiles["qpar_wall(1/m^3/s)"] * 1e-20 + ax.plot(rho, var, lw=lw, ls=lines[1], c=color, label=extralab + "wall") + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.axhline(y=0, lw=0.5, ls="--", c="k") + ax.set_ylabel("$10^{20}m^{-3}s^{-1}$") + ax.set_title("Particle Source Density") + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + +def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): + + [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,ax13c] = axs3 + + rho = self.profiles["rho(-)"] + lines = GRAPHICStools.listLS() + + ax = ax00c + varL = "cos Shape Params" + yl = 0 + cont = 0 + + for i, s in enumerate(self.shape_cos): + if s is not None: + valmax = np.abs(s).max() + if valmax > 1e-10: + lab = f"c{i}" + ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) + cont += 1 + + yl = np.max([yl, valmax]) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + + + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + if legYN: + ax.legend(loc="best", fontsize=fs) + + ax = ax01c + varL = "sin Shape Params" + cont = 0 + for i, s in enumerate(self.shape_sin): + if s is not None: + valmax = np.abs(s).max() + if valmax > 1e-10: + lab = f"s{i}" + ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) + cont += 1 + + yl = np.max([yl, valmax]) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax02c + var = self.profiles["polflux(Wb/radian)"] + ax.plot(rho, var, lw=lw, ls="-", c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("Poloidal $\\psi$ ($Wb/rad$)") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax10c + var = self.profiles["delta(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$\\delta$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax11c + + var = self.profiles["rmin(m)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylim(bottom=0) + ax.set_ylabel("$r_{min}$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax20c + + var = self.profiles["rmaj(m)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$R_{maj}$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax21c + + var = self.profiles["zmag(m)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + yl = np.max([0.1, np.max(np.abs(var))]) + ax.set_ylim([-yl, yl]) + ax.set_ylabel("$Z_{maj}$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax22c + + var = self.profiles["kappa(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$\\kappa$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=1) + + ax = ax12c + + var = self.profiles["zeta(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("zeta") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax13c + self.plot_state_flux_surfaces(ax=ax, color=color) + + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + GRAPHICStools.addDenseAxis(ax) + + +def plot_gradients( + self, + axs4, + color="b", + lw=1.0, + label="", + ls="-o", + lastRho=0.89, + ms=2, + alpha=1.0, + useRoa=False, + RhoLocationsPlot=None, + plotImpurity=None, + plotRotation=False, + autoscale=True, + ): + + if RhoLocationsPlot is None: RhoLocationsPlot=[] + + if axs4 is None: + plt.ion() + fig, axs = plt.subplots( + ncols=3 + int(plotImpurity is not None) + int(plotRotation), + nrows=2, + figsize=(12, 5), + ) + + axs4 = [] + for i in range(axs.shape[-1]): + axs4.append(axs[0, i]) + axs4.append(axs[1, i]) + + ix = np.argmin(np.abs(self.profiles["rho(-)"] - lastRho)) + 1 + + xcoord = self.profiles["rho(-)"] if (not useRoa) else self.derived["roa"] + labelx = "$\\rho$" if (not useRoa) else "$r/a$" + + ax = axs4[0] + ax.plot( + xcoord, + self.profiles["te(keV)"], + ls, + c=color, + lw=lw, + label=label, + markersize=ms, + alpha=alpha, + ) + ax = axs4[2] + ax.plot( + xcoord, + self.profiles["ti(keV)"][:, 0], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + ax = axs4[4] + ax.plot( + xcoord, + self.profiles["ne(10^19/m^3)"] * 1e-1, + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + + if "derived" in self.__dict__: + ax = axs4[1] + ax.plot( + xcoord[:ix], + self.derived["aLTe"][:ix], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + ax = axs4[3] + ax.plot( + xcoord[:ix], + self.derived["aLTi"][:ix, 0], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + ax = axs4[5] + ax.plot( + xcoord[:ix], + self.derived["aLne"][:ix], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + + for ax in axs4: + ax.set_xlim([0, 1]) + + ax = axs4[0] + ax.set_ylabel("$T_e$ (keV)") + ax.set_xlabel(labelx) + if autoscale: + 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) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = axs4[4] + ax.set_ylabel("$n_e$ ($10^{20}m^{-3}$)") + ax.set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axs4[1] + ax.set_ylabel("$a/L_{Te}$") + ax.set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = axs4[3] + ax.set_ylabel("$a/L_{Ti}$") + ax.set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = axs4[5] + ax.set_ylabel("$a/L_{ne}$") + ax.axhline(y=0, ls="--", lw=0.5, c="k") + ax.set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + + cont = 0 + if plotImpurity is not None: + axs4[6 + cont].plot( + xcoord, + self.profiles["ni(10^19/m^3)"][:, plotImpurity] * 1e-1, + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + axs4[6 + cont].set_ylabel("$n_Z$ ($10^{20}m^{-3}$)") + axs4[6].set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + if "derived" in self.__dict__: + axs4[7 + cont].plot( + xcoord[:ix], + self.derived["aLni"][:ix, plotImpurity], + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + axs4[7 + cont].set_ylabel("$a/L_{nZ}$") + axs4[7 + cont].axhline(y=0, ls="--", lw=0.5, c="k") + axs4[7 + cont].set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + cont += 2 + + if plotRotation: + axs4[6 + cont].plot( + xcoord, + self.profiles["w0(rad/s)"] * 1e-3, + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + axs4[6 + cont].set_ylabel("$w_0$ (krad/s)") + axs4[6 + cont].set_xlabel(labelx) + if "derived" in self.__dict__: + axs4[7 + cont].plot( + xcoord[:ix], + self.derived["dw0dr"][:ix] * 1e-5, + ls, + c=color, + lw=lw, + markersize=ms, + alpha=alpha, + ) + axs4[7 + cont].set_ylabel("-$d\\omega_0/dr$ (krad/s/cm)") + axs4[7 + cont].axhline(y=0, ls="--", lw=0.5, c="k") + axs4[7 + cont].set_xlabel(labelx) + if autoscale: + GRAPHICStools.autoscale_y(ax, bottomy=0) + cont += 2 + + for x0 in RhoLocationsPlot: + ix = np.argmin(np.abs(self.profiles["rho(-)"] - x0)) + for ax in axs4: + ax.axvline(x=xcoord[ix], ls="--", lw=0.5, c=color) + + for i in range(len(axs4)): + ax = axs4[i] + GRAPHICStools.addDenseAxis(ax) + +def plot_other(self, axs6, color="b", lw=1.0, extralab="", fs=6): + + rho = self.profiles["rho(-)"] + lines = GRAPHICStools.listLS() + + # Others + ax = axs6[0] + ax.plot(self.profiles["rho(-)"], self.derived["dw0dr"] * 1e-5, c=color, lw=lw) + ax.set_ylabel("$-d\\omega_0/dr$ (krad/s/cm)") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + ax.axhline(y=0, lw=1.0, c="k", ls="--") + + ax = axs6[2] + ax.plot(self.profiles["rho(-)"], self.derived["q_fus"], c=color, lw=lw) + ax.set_ylabel("$q_{fus}$ ($MW/m^3$)") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axs6[3] + ax.plot(self.profiles["rho(-)"], self.derived["q_fus_MW"], c=color, lw=lw) + ax.set_ylabel("$P_{fus}$ ($MW$)") + ax.set_xlim([0, 1]) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axs6[4] + ax.plot(self.profiles["rho(-)"], self.derived["tite"], c=color, lw=lw) + ax.set_ylabel("$T_i/T_e$") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + ax.axhline(y=1, ls="--", lw=1.0, c="k") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = axs6[5] + if "MachNum" in self.derived: + ax.plot(self.profiles["rho(-)"], self.derived["MachNum"], c=color, lw=lw) + ax.set_ylabel("Mach Number") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + ax.axhline(y=0, ls="--", c="k", lw=0.5) + ax.axhline(y=1, ls="--", c="k", lw=0.5) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = axs6[6] + safe_division = np.divide( + self.derived["qi_MWm2"], + self.derived["qe_MWm2"], + where=self.derived["qe_MWm2"] != 0, + out=np.full_like(self.derived["qi_MWm2"], np.nan), + ) + ax.plot( + self.profiles["rho(-)"], + safe_division, + c=color, + lw=lw, + label=extralab + "$Q_i/Q_e$", + ) + safe_division = np.divide( + self.derived["qi_aux_MW"], + self.derived["qe_aux_MW"], + where=self.derived["qe_aux_MW"] != 0, + out=np.full_like(self.derived["qi_aux_MW"], np.nan), + ) + ax.plot( + self.profiles["rho(-)"], + safe_division, + c=color, + lw=lw, + ls="--", + label=extralab + "$P_i/P_e$", + ) + ax.set_ylabel("Power ratios") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + ax.axhline(y=1.0, ls="--", c="k", lw=1.0) + GRAPHICStools.addDenseAxis(ax) + # GRAPHICStools.autoscale_y(ax,bottomy=0) + ax.set_ylim(bottom=0) + ax.legend(loc="best", fontsize=fs) + + # Currents + + ax = axs6[1] + + var = self.profiles["johm(MA/m^2)"] + ax.plot(rho, var, "-", lw=lw, c=color, label=extralab + "$J_{OH}$") + var = self.profiles["jbs(MA/m^2)"] + ax.plot(rho, var, "--", lw=lw, c=color, label=extralab + "$J_{BS,par}$") + var = self.profiles["jbstor(MA/m^2)"] + ax.plot(rho, var, "-.", lw=lw, c=color, label=extralab + "$J_{BS,tor}$") + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylim(bottom=0) + ax.set_ylabel("J ($MA/m^2$)") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axs6[7] + cont = 0 + if "vtor(m/s)" in self.profiles: + for i in range(len(self.Species)): + try: # REMOVE FOR FUTURE + var = self.profiles["vtor(m/s)"][:, i] * 1e-3 + ax.plot( + rho, + var, + lw=lw, + ls=lines[cont], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + cont += 1 + except: + break + varL = "$V_{tor}$ (km/s)" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + +def plot_flows(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): + if axs is None: + fig1 = plt.figure() + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + + axs = [ + fig1.add_subplot(grid[0, 0]), + fig1.add_subplot(grid[1, 0]), + fig1.add_subplot(grid[0, 1]), + fig1.add_subplot(grid[0, 2]), + fig1.add_subplot(grid[1, 1]), + fig1.add_subplot(grid[1, 2]), + ] + + # Profiles + + ax = axs[0] + axT = axs[1] + roa = self.profiles["rmin(m)"] / self.profiles["rmin(m)"][-1] + Te = self.profiles["te(keV)"] + ne = self.profiles["ne(10^19/m^3)"] * 1e-1 + ni = self.profiles["ni(10^19/m^3)"] * 1e-1 + niT = np.sum(ni, axis=1) + Ti = self.profiles["ti(keV)"][:, 0] + ax.plot(roa, Te, lw=2, c="r", label="$T_e$" if leg else "", ls=ls) + ax.plot(roa, Ti, lw=2, c="b", label="$T_i$" if leg else "", ls=ls) + axT.plot(roa, ne, lw=2, c="m", label="$n_e$" if leg else "", ls=ls) + axT.plot(roa, niT, lw=2, c="c", label="$\\sum n_i$" if leg else "", ls=ls) + if limits is not None: + [roa_first, roa_last] = limits + ax.plot(roa_last, np.interp(roa_last, roa, Te), "s", c="r", markersize=3) + ax.plot(roa_first, np.interp(roa_first, roa, Te), "s", c="r", markersize=3) + ax.plot(roa_last, np.interp(roa_last, roa, Ti), "s", c="b", markersize=3) + ax.plot(roa_first, np.interp(roa_first, roa, Ti), "s", c="b", markersize=3) + axT.plot(roa_last, np.interp(roa_last, roa, ne), "s", c="m", markersize=3) + axT.plot(roa_first, np.interp(roa_first, roa, ne), "s", c="m", markersize=3) + + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + axT.set_xlabel("r/a") + axT.set_xlim([0, 1]) + ax.set_ylabel("$T$ (keV)") + ax.set_ylim(bottom=0) + axT.set_ylabel("$n$ ($10^{20}m^{-3}$)") + axT.set_ylim(bottom=0) + # axT.set_ylim([0,np.max(ne)*1.5]) + ax.legend() + axT.legend() + ax.set_title("Final Temperature profiles") + axT.set_title("Final Density profiles") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + GRAPHICStools.addDenseAxis(axT) + GRAPHICStools.autoscale_y(axT, bottomy=0) + + if showtexts: + if self.derived["Q"] > 0.005: + ax.text( + 0.05, + 0.05, + f"Pfus = {self.derived['Pfus']:.1f}MW, Q = {self.derived['Q']:.2f}", + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + transform=ax.transAxes, + ) + + axT.text( + 0.05, + 0.4, + "ne_20 = {0:.1f} (fG = {1:.2f}), Zeff = {2:.1f}".format( + self.derived["ne_vol20"], + self.derived["fG"], + self.derived["Zeff_vol"], + ), + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + transform=axT.transAxes, + ) + + # F + ax = axs[2] + P = ( + self.derived["qe_fus_MW"] + + self.derived["qe_aux_MW"] + + -self.derived["qe_rad_MW"] + + -self.derived["qe_exc_MW"] + ) + + ax.plot( + roa, + -self.derived["qe_MW"], + c="g", + lw=2, + label="$P_{e}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qe_fus_MW"], + c="r", + lw=2, + label="$P_{fus,e}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qe_aux_MW"], + c="b", + lw=2, + label="$P_{aux,e}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + -self.derived["qe_exc_MW"], + c="m", + lw=2, + label="$P_{exc,e}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + -self.derived["qe_rad_MW"], + c="c", + lw=2, + label="$P_{rad,e}$" if leg else "", + ls=ls, + ) + ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) + + # Pe = self.profiles['te(keV)']*1E3*e_J*self.profiles['ne(10^19/m^3)']*1E-1*1E20 *1E-6 + # ax.plot(roa,Pe,ls='-',lw=3,alpha=0.1,c='k',label='$W_e$ (MJ/m^3)') + + ax.plot( + roa, + -self.derived["ce_MW"], + c="k", + lw=1, + label="($P_{conv,e}$)" if leg else "", + ) + + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + ax.set_ylabel("$P$ (MW)") + # ax.set_ylim(bottom=0) + ax.set_title("Electron Thermal Flows") + ax.axhline(y=0.0, lw=0.5, ls="--", c="k") + GRAPHICStools.addLegendApart( + ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" + ) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axs[3] + P = ( + self.derived["qi_fus_MW"] + + self.derived["qi_aux_MW"] + + self.derived["qe_exc_MW"] + ) + + ax.plot( + roa, + -self.derived["qi_MW"], + c="g", + lw=2, + label="$P_{i}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qi_fus_MW"], + c="r", + lw=2, + label="$P_{fus,i}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qi_aux_MW"], + c="b", + lw=2, + label="$P_{aux,i}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qe_exc_MW"], + c="m", + lw=2, + label="$P_{exc,i}$" if leg else "", + ls=ls, + ) + ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) + + # Pi = self.profiles['ti(keV)'][:,0]*1E3*e_J*self.profiles['ni(10^19/m^3)'][:,0]*1E-1*1E20 *1E-6 + # ax.plot(roa,Pi,ls='-',lw=3,alpha=0.1,c='k',label='$W_$ (MJ/m^3)') + + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + ax.set_ylabel("$P$ (MW)") + # ax.set_ylim(bottom=0) + ax.set_title("Ion Thermal Flows") + ax.axhline(y=0.0, lw=0.5, ls="--", c="k") + GRAPHICStools.addLegendApart( + ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" + ) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + # F + ax = axs[4] + + ax.plot( + roa, + self.derived["ge_10E20"], + c="g", + lw=2, + label="$\\Gamma_{e}$" if leg else "", + ls=ls, + ) + # ax.plot(roa,self.profiles['ne(10^19/m^3)']*1E-1,lw=3,alpha=0.1,c='k',label='$n_e$ ($10^{20}/m^3$)' if leg else '',ls=ls) + + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + ax.set_ylabel("$\\Gamma$ ($10^{20}/s$)") + ax.set_title("Particle Flows") + ax.axhline(y=0.0, lw=0.5, ls="--", c="k") + GRAPHICStools.addLegendApart( + ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" + ) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + # TOTAL + ax = axs[5] + P = ( + self.derived["qOhm_MW"] + + self.derived["qRF_MW"] + + self.derived["qFus_MW"] + + -self.derived["qe_rad_MW"] + + self.derived["qz_MW"] + + self.derived["qBEAM_MW"] + ) + + ax.plot( + roa, + -self.derived["q_MW"], + c="g", + lw=2, + label="$P$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qOhm_MW"], + c="k", + lw=2, + label="$P_{Oh}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qRF_MW"], + c="b", + lw=2, + label="$P_{RF}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qBEAM_MW"], + c="pink", + lw=2, + label="$P_{NBI}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qFus_MW"], + c="r", + lw=2, + label="$P_{fus}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + -self.derived["qe_rad_MW"], + c="c", + lw=2, + label="$P_{rad}$" if leg else "", + ls=ls, + ) + ax.plot( + roa, + self.derived["qz_MW"], + c="orange", + lw=1, + label="$P_{ionz.}$" if leg else "", + ls=ls, + ) + + # P = Pe+Pi + # ax.plot(roa,P,ls='-',lw=3,alpha=0.1,c='k',label='$W$ (MJ)') + + ax.plot(roa, -P, lw=1, c="y", label="sum" if leg else "", ls=ls) + ax.set_xlabel("r/a") + ax.set_xlim([0, 1]) + ax.set_ylabel("$P$ (MW)") + # ax.set_ylim(bottom=0) + ax.set_title("Total Thermal Flows") + + GRAPHICStools.addLegendApart( + ax, ratio=0.9, withleg=True, extraPad=0, size=None, loc="upper left" + ) + + ax.axhline(y=0.0, lw=0.5, ls="--", c="k") + # GRAPHICStools.drawLineWithTxt(ax,0.0,label='',orientation='vertical',color='k',lw=1,ls='--',alpha=1.0,fontsize=10,fromtop=0.85,fontweight='normal', + # verticalalignment='bottom',horizontalalignment='left',separation=0) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + + + +def plot_ions(self, axsImps, legYN=True, extralab="", color="b", lw=1, fs=6): + + rho = self.profiles["rho(-)"] + lines = GRAPHICStools.listLS() + + # Impurities + ax = axsImps[0] + for i in range(len(self.Species)): + var = ( + self.profiles["ni(10^19/m^3)"][:, i] + / self.profiles["ni(10^19/m^3)"][0, i] + ) + ax.plot( + rho, + var, + lw=lw, + ls=lines[i], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + varL = "$n_i/n_{i,0}$" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axsImps[1] + for i in range(len(self.Species)): + var = self.derived["fi"][:, i] + ax.plot( + rho, + var, + lw=lw, + ls=lines[i], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + varL = "$f_i$" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + ax.set_ylim([0, 1]) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axsImps[2] + + lastRho = 0.9 + + ix = np.argmin(np.abs(self.profiles["rho(-)"] - lastRho)) + 1 + ax.plot( + rho[:ix], self.derived["aLne"][:ix], lw=lw * 3, ls="-", c=color, label="e" + ) + for i in range(len(self.Species)): + var = self.derived["aLni"][:, i] + ax.plot( + rho[:ix], + var[:ix], + lw=lw, + ls=lines[i], + c=color, + label=extralab + f"{i + 1} = {self.profiles['name'][i]}", + ) + varL = "$a/L_{ni}$" + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = axsImps[5] + + ax = axsImps[3] + ax.plot(self.profiles["rho(-)"], self.derived["Zeff"], c=color, lw=lw) + ax.set_ylabel("$Z_{eff}$") + ax.set_xlabel("$\\rho$") + ax.set_xlim([0, 1]) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + +def plotAll(profiles_list, figs=None, extralabs=None, lastRhoGradients=0.89): + if figs is not None: + fn = None + else: + from mitim_tools.misc_tools.GUItools import FigureNotebook + fn = FigureNotebook("Profiles", geometry="1800x900") + figs = add_figures(fn) + + axsProf_1, axsProf_2, axsProf_3, axsProf_4, axsFlows, axsProf_6, axsImps = add_axes(figs) + + ls = GRAPHICStools.listLS() + colors = GRAPHICStools.listColors() + for i, profiles in enumerate(profiles_list): + if extralabs is None: + extralab = f"#{i}, " + else: + extralab = f"{extralabs[i]}, " + + profiles.plot( + axs1=axsProf_1,axs2=axsProf_2,axs3=axsProf_3,axs4=axsProf_4,axsFlows=axsFlows,axs6=axsProf_6,axsImps=axsImps, + color=colors[i],legYN=True,extralab=extralab,lsFlows=ls[i],legFlows=i == 0,showtexts=False,lastRhoGradients=lastRhoGradients, + ) + + return fn diff --git a/src/mitim_tools/popcon_tools/RAPIDStools.py b/src/mitim_tools/popcon_tools/RAPIDStools.py index d435555d..e6ffb963 100644 --- a/src/mitim_tools/popcon_tools/RAPIDStools.py +++ b/src/mitim_tools/popcon_tools/RAPIDStools.py @@ -316,7 +316,7 @@ def scan_density_additional(nn, p_base, nominal_parameters, core, r, param, para resultsS, ['r','b','g'], ): - results['profs'][0].plotGeometry(ax=ax, surfaces_rho=[1.0], color=c) + results['profs'][0].plot_state_flux_surfaces(ax=ax, surfaces_rho=[1.0], color=c) GRAPHICStools.addDenseAxis(ax) ax.set_xlabel("R (m)") diff --git a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py index afd675af..1ff37140 100644 --- a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py +++ b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py @@ -820,7 +820,7 @@ def _produce_geometry_profiles(self): # Separatrix # -------------------------------------------------------------- - self.geometry['R_sep'], self.geometry['Z_sep'] = self.p.derived["R_surface"][-1], self.p.derived["Z_surface"][-1] + self.geometry['R_sep'], self.geometry['Z_sep'] = self.p.derived["R_surface"][0,:,-1], self.p.derived["Z_surface"][0,:,-1] # -------------------------------------------------------------- # VV From 47bcda969806a29ce38d39c40ea2a2cc30d82ab1 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 22 Jun 2025 18:13:54 +0200 Subject: [PATCH 094/385] Geometry only plotted for input.gacode --- src/mitim_tools/gacode_tools/PROFILEStools.py | 167 +++++++++++++++++- .../plasmastate_tools/MITIMstate.py | 12 +- .../plasmastate_tools/utils/VMECtools.py | 7 +- .../plasmastate_tools/utils/state_plotting.py | 153 ---------------- 4 files changed, 173 insertions(+), 166 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 09badca0..8b2d3c72 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -3,7 +3,7 @@ from collections import OrderedDict from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.gs_tools import GEQtools -from mitim_tools.misc_tools import MATHtools, IOtools +from mitim_tools.misc_tools import MATHtools, IOtools, GRAPHICStools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -57,6 +57,7 @@ def _read_inputgacocde(self): self.profiles["qpar_beam(1/m^3/s)"] = self.profiles.pop("qpar_beam(MW/m^3)") if "qpar_wall(MW/m^3)" in self.profiles: self.profiles["qpar_wall(1/m^3/s)"] = self.profiles.pop("qpar_wall(MW/m^3)") + """ Note that in prgen_map_plasmastate, that variable: expro_qpar_beam(i) = plst_sn_trans(i-1)/dvol @@ -84,7 +85,15 @@ def _read_inputgacocde(self): """ - + # Ensure that we also have the shape coefficients + num_moments = 7 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros + if "shape_cos0(-)" not in self.profiles: + self.profiles["shape_cos0(-)"] = np.ones(self.profiles["rmaj(m)"].shape) + for i in range(num_moments): + if f"shape_cos{i + 1}(-)" not in self.profiles: + self.profiles[f"shape_cos{i + 1}(-)"] = np.zeros(self.profiles["rmaj(m)"].shape) + if f"shape_sin{i + 1}(-)" not in self.profiles and i > 1: + self.profiles[f"shape_sin{i + 1}(-)"] = np.zeros(self.profiles["rmaj(m)"].shape) def _read_header(self): for i in range(len(self.lines)): @@ -231,6 +240,160 @@ def derive_geometry(self, n_theta_geo=1001): self.derived["surfGACODE_geo"] = (self.derived["surf_geo"] / self.derived["gradr_geo"]) self.derived["surfGACODE_geo"][np.isnan(self.derived["surfGACODE_geo"])] = 0 + def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): + + [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,ax13c] = axs3 + + rho = self.profiles["rho(-)"] + lines = GRAPHICStools.listLS() + + ax = ax00c + varL = "cos Shape Params" + yl = 0 + cont = 0 + + for i, s in enumerate(self.shape_cos): + if s is not None: + valmax = np.abs(s).max() + if valmax > 1e-10: + lab = f"c{i}" + ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) + cont += 1 + + yl = np.max([yl, valmax]) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + if legYN: + ax.legend(loc="best", fontsize=fs) + + ax = ax01c + varL = "sin Shape Params" + cont = 0 + for i, s in enumerate(self.shape_sin): + if s is not None: + valmax = np.abs(s).max() + if valmax > 1e-10: + lab = f"s{i}" + ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) + cont += 1 + + yl = np.max([yl, valmax]) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax02c + var = self.profiles["polflux(Wb/radian)"] + ax.plot(rho, var, lw=lw, ls="-", c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("Poloidal $\\psi$ ($Wb/rad$)") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax10c + var = self.profiles["delta(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$\\delta$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax11c + + var = self.profiles["rmin(m)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylim(bottom=0) + ax.set_ylabel("$r_{min}$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax20c + + var = self.profiles["rmaj(m)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$R_{maj}$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax21c + + var = self.profiles["zmag(m)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + yl = np.max([0.1, np.max(np.abs(var))]) + ax.set_ylim([-yl, yl]) + ax.set_ylabel("$Z_{maj}$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax22c + + var = self.profiles["kappa(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$\\kappa$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=1) + + ax = ax12c + + var = self.profiles["zeta(-)"] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("zeta") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax13c + self.plot_state_flux_surfaces(ax=ax, color=color) + + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + GRAPHICStools.addDenseAxis(ax) + + + + def calculateGeometricFactors(profiles, n_theta=1001): # ---------------------------------------- diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 8170b053..139b5655 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -49,15 +49,9 @@ def ensure_variables_existence(self): "kappa(-)": 1, "delta(-)": 1, "zeta(-)": 1, - "shape_cos0(-)": 1, }) - num_moments = 7 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros - for i in range(num_moments): - required_profiles[f"shape_cos{i + 1}(-)"] = 1 - if i > 1: - required_profiles[f"shape_sin{i + 1}(-)"] = 1 - + # Sources and Sinks required_profiles.update({ "qohme(MW/m^3)": 1, @@ -1789,13 +1783,15 @@ def plot( state_plotting.plot_profiles(self,axs1, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) state_plotting.plot_powers(self,axs2, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) - state_plotting.plot_geometry(self,axs3, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) + self.plot_geometry(axs3, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) state_plotting.plot_gradients(self,axs4, color=color, lw=lw, lastRho=lastRhoGradients, label=extralab) if axsFlows is not None: state_plotting.plot_flows(self, axsFlows, ls=lsFlows, leg=legFlows, showtexts=showtexts) state_plotting.plot_other(self,axs6, color=color, lw=lw, extralab=extralab, fs=fs) state_plotting.plot_ions(self,axsImps, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) + def plot_geometry(self, *args, **kwargs): + pass def plot_state_flux_surfaces(self, ax=None, surfaces_rho=np.linspace(0, 1, 11), color="b", label = '', lw=1.0, lw1=2.0): if ("R_surface" in self.derived) and (self.derived["R_surface"] is not None): diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index 206aa902..591657da 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -1,7 +1,7 @@ +from mitim_tools.misc_tools import IOtools +from mitim_tools.plasmastate_tools import MITIMstate -from mitim_tools.plasmastate_tools.MITIMstate import mitim_state - -class vmec_state(mitim_state): +class vmec_state(MITIMstate.mitim_state): ''' Class to read and manipulate VMEC files ''' @@ -22,5 +22,6 @@ def __init__(self, file, derive_quantities=True, mi_ref=None): # Derive (Depending on resolution, derived can be expensive, so I mmay not do it every time) self.derive_quantities(mi_ref=mi_ref, derive_quantities=derive_quantities) + @IOtools.hook_method(after=MITIMstate.ensure_variables_existence) def _read_vmec(self): pass diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py index 60ed078d..0555b8a9 100644 --- a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -486,161 +486,8 @@ def plot_powers(self, axs2, legYN=True, extralab="", color="b", lw=1, fs=6): GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax) - - -def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): - - [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,ax13c] = axs3 - - rho = self.profiles["rho(-)"] - lines = GRAPHICStools.listLS() - - ax = ax00c - varL = "cos Shape Params" - yl = 0 - cont = 0 - - for i, s in enumerate(self.shape_cos): - if s is not None: - valmax = np.abs(s).max() - if valmax > 1e-10: - lab = f"c{i}" - ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) - cont += 1 - - yl = np.max([yl, valmax]) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - if legYN: - ax.legend(loc="best", fontsize=fs) - - ax = ax01c - varL = "sin Shape Params" - cont = 0 - for i, s in enumerate(self.shape_sin): - if s is not None: - valmax = np.abs(s).max() - if valmax > 1e-10: - lab = f"s{i}" - ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) - cont += 1 - - yl = np.max([yl, valmax]) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax02c - var = self.profiles["polflux(Wb/radian)"] - ax.plot(rho, var, lw=lw, ls="-", c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("Poloidal $\\psi$ ($Wb/rad$)") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax10c - var = self.profiles["delta(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$\\delta$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax11c - - var = self.profiles["rmin(m)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylim(bottom=0) - ax.set_ylabel("$r_{min}$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax20c - - var = self.profiles["rmaj(m)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$R_{maj}$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax21c - - var = self.profiles["zmag(m)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - yl = np.max([0.1, np.max(np.abs(var))]) - ax.set_ylim([-yl, yl]) - ax.set_ylabel("$Z_{maj}$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax22c - - var = self.profiles["kappa(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$\\kappa$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=1) - - ax = ax12c - - var = self.profiles["zeta(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("zeta") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - ax = ax13c - self.plot_state_flux_surfaces(ax=ax, color=color) - - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - GRAPHICStools.addDenseAxis(ax) - - def plot_gradients( self, axs4, From ac7308b72af0317216501cbc8cdac263927b1dc0 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 22 Jun 2025 18:40:37 +0200 Subject: [PATCH 095/385] misc --- src/mitim_modules/portals/utils/PORTALSplot.py | 1 - src/mitim_tools/plasmastate_tools/MITIMstate.py | 4 ++++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index ae618ba4..de23f1ec 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -8,7 +8,6 @@ from mitim_modules.powertorch import STATEtools from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_modules.powertorch.utils import POWERplot -from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 139b5655..5edd359d 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1790,6 +1790,10 @@ def plot( state_plotting.plot_other(self,axs6, color=color, lw=lw, extralab=extralab, fs=fs) state_plotting.plot_ions(self,axsImps, color=color, legYN=legYN, extralab=extralab, lw=lw, fs=fs) + # To allow this to be called from the object + def plot_gradients(self, *args, **kwargs): + return state_plotting.plot_gradients(self, *args, **kwargs) + def plot_geometry(self, *args, **kwargs): pass From 2735b595fca78ffe74266b1787c89bbe6861ad96 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 22 Jun 2025 23:53:48 +0200 Subject: [PATCH 096/385] Moved write_state (renamed writteCurrentStatus) to specific children --- src/mitim_modules/maestro/MAESTROmain.py | 4 +- .../maestro/scripts/run_maestro.py | 2 +- src/mitim_modules/maestro/utils/EPEDbeat.py | 4 +- .../maestro/utils/MAESTRObeat.py | 4 +- .../maestro/utils/PORTALSbeat.py | 8 +- src/mitim_modules/maestro/utils/TRANSPbeat.py | 6 +- .../portals/utils/PORTALSanalysis.py | 4 +- .../portals/utils/PORTALSinit.py | 2 +- .../powertorch/utils/TRANSFORMtools.py | 2 +- src/mitim_tools/astra_tools/ASTRAtools.py | 2 +- src/mitim_tools/gacode_tools/CGYROtools.py | 6 +- src/mitim_tools/gacode_tools/NEOtools.py | 2 +- src/mitim_tools/gacode_tools/PROFILEStools.py | 65 +++++++++++ src/mitim_tools/gacode_tools/TGLFtools.py | 6 +- src/mitim_tools/gacode_tools/TGYROtools.py | 10 +- src/mitim_tools/gs_tools/GEQtools.py | 2 +- .../plasmastate_tools/MITIMstate.py | 103 ++---------------- .../plasmastate_tools/utils/VMECtools.py | 16 ++- 18 files changed, 118 insertions(+), 130 deletions(-) diff --git a/src/mitim_modules/maestro/MAESTROmain.py b/src/mitim_modules/maestro/MAESTROmain.py index 9b3923b2..3b633d1e 100644 --- a/src/mitim_modules/maestro/MAESTROmain.py +++ b/src/mitim_modules/maestro/MAESTROmain.py @@ -259,7 +259,7 @@ def _freeze_parameters(self, profiles = None): print('\t\t- Freezing engineering parameters from MAESTRO') self.profiles_with_engineering_parameters = copy.deepcopy(profiles) - self.profiles_with_engineering_parameters.writeCurrentStatus(file= (self.folder_output / 'input.gacode_frozen')) + self.profiles_with_engineering_parameters.write_state(file= (self.folder_output / 'input.gacode_frozen')) @mitim_timer('\t\t* Finalizing', name_timer=None) def finalize(self): @@ -271,7 +271,7 @@ def finalize(self): final_file= (self.folder_output / 'input.gacode_final') - self.beat.profiles_output.writeCurrentStatus(file= final_file) + self.beat.profiles_output.write_state(file= final_file) print(f'\t\t- Final input.gacode saved to {IOtools.clipstr(final_file)}') diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index b1a14749..b9c0d5f6 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -144,7 +144,7 @@ def profiles_postprocessing_fun(file_profs): p.lumpImpurities() if enforce_same_density_gradients: p.enforce_same_density_gradients() - p.writeCurrentStatus(file=file_profs) + p.write_state(file=file_profs) beat_namelist['PORTALSparameters']['profiles_postprocessing_fun'] = profiles_postprocessing_fun else: diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index 63267924..8de195d6 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -285,7 +285,7 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F for key in eped_results: print(f'\t\t- {key}: {eped_results[key]}') - self.profiles_output.writeCurrentStatus(file=self.folder / 'input.gacode.eped') + self.profiles_output.write_state(file=self.folder / 'input.gacode.eped') return eped_results @@ -293,7 +293,7 @@ def finalize(self, **kwargs): self.profiles_output = PROFILEStools.gacode_state(self.folder / 'input.gacode.eped') - self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') + self.profiles_output.write_state(file=self.folder_output / 'input.gacode') def merge_parameters(self): # EPED beat does not modify the profiles grid or anything, so I can keep it fine diff --git a/src/mitim_modules/maestro/utils/MAESTRObeat.py b/src/mitim_modules/maestro/utils/MAESTRObeat.py index a4fa0e9c..bae64b52 100644 --- a/src/mitim_modules/maestro/utils/MAESTRObeat.py +++ b/src/mitim_modules/maestro/utils/MAESTRObeat.py @@ -112,7 +112,7 @@ def __call__(self, profiles_file = None, Vsurf = None, **kwargs_beat): # -------------------------------------------------------------------------------------------- # Write it to initialization folder - self.profiles_current.writeCurrentStatus(file=self.folder / 'input.gacode') + self.profiles_current.write_state(file=self.folder / 'input.gacode') # Pass the profiles to the beat instance self.beat_instance.profiles_current = self.profiles_current @@ -199,7 +199,7 @@ def __call__( p.profiles['q(-)'] *= kwargs_profiles['B_T'] / Bt_in_geqdsk # Write it to initialization folder - p.writeCurrentStatus(file=self.folder / 'input.geqdsk.gacode') + p.write_state(file=self.folder / 'input.geqdsk.gacode') # Copy original geqdsk for reference use shutil.copy2(geqdsk_file, self.folder / "input.geqdsk") diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 0f80b716..d4a1bae9 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -63,7 +63,7 @@ def prepare(self, self.profiles_current = profiles - self.profiles_current.writeCurrentStatus(file = self.fileGACODE) + self.profiles_current.write_state(file = self.fileGACODE) self.PORTALSparameters = PORTALSparameters self.MODELparameters = MODELparameters @@ -155,7 +155,7 @@ def finalize(self, **kwargs): print('\t\t- PORTALS probably converged in training, so analyzing a bit differently') self.profiles_output = portals_output.profiles[portals_output.opt_fun_full.res.best_absolute_index] - self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') + self.profiles_output.write_state(file=self.folder_output / 'input.gacode') def merge_parameters(self): ''' @@ -172,7 +172,7 @@ def merge_parameters(self): ''' # Write the pre-merge input.gacode before modifying it - self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode_pre_merge') + self.profiles_output.write_state(file=self.folder_output / 'input.gacode_pre_merge') # First, bring back to the resolution of the frozen p_frozen = self.maestro_instance.profiles_with_engineering_parameters @@ -231,7 +231,7 @@ def merge_parameters(self): # Write to final input.gacode self.profiles_output.derive_quantities() - self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') + self.profiles_output.write_state(file=self.folder_output / 'input.gacode') def grab_output(self, full = False): diff --git a/src/mitim_modules/maestro/utils/TRANSPbeat.py b/src/mitim_modules/maestro/utils/TRANSPbeat.py index 739e0ddc..29fb7217 100644 --- a/src/mitim_modules/maestro/utils/TRANSPbeat.py +++ b/src/mitim_modules/maestro/utils/TRANSPbeat.py @@ -174,7 +174,7 @@ def finalize(self, force_auxiliary_heating_at_output = {'Pe': None, 'Pi': None}, self._add_heating_profiles(force_auxiliary_heating_at_output) # Write profiles - self.profiles_output.writeCurrentStatus(file=self.folder_output / "input.gacode") + self.profiles_output.write_state(file=self.folder_output / "input.gacode") def _add_heating_profiles(self, force_auxiliary_heating_at_output = {'Pe': None, 'Pi': None}): ''' @@ -206,7 +206,7 @@ def merge_parameters(self): # Write the pre-merge input.gacode before modifying it profiles_output_pre_merge = copy.deepcopy(self.profiles_output) - profiles_output_pre_merge.writeCurrentStatus(file=self.folder_output / 'input.gacode_pre_merge') + profiles_output_pre_merge.write_state(file=self.folder_output / 'input.gacode_pre_merge') # First, bring back to the resolution of the frozen p_frozen = self.maestro_instance.profiles_with_engineering_parameters @@ -237,7 +237,7 @@ def merge_parameters(self): # Write to final input.gacode self.profiles_output.derive_quantities() - self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') + self.profiles_output.write_state(file=self.folder_output / 'input.gacode') def grab_output(self): diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 8770371c..35bf5377 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -498,7 +498,7 @@ def extractPORTALS(self, evaluation=None, folder=None, modified_profiles=False): # Start from the profiles of that step fileGACODE = folder / "input.gacode_transferred" p = self.extractProfiles(evaluation=evaluation, modified_profiles=modified_profiles) - p.writeCurrentStatus(file=fileGACODE) + p.write_state(file=fileGACODE) # New class from mitim_modules.portals.PORTALSmain import portals @@ -592,7 +592,7 @@ def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=F inputgacode = folder / "input.gacode.start" p = self.extractProfiles(evaluation=evaluation,modified_profiles=modified_profiles) - p.writeCurrentStatus(file=inputgacode) + p.write_state(file=inputgacode) tglf = TGLFtools.TGLF(rhos=rhos) _ = tglf.prep(folder, cold_start=cold_start, inputgacode=inputgacode) diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index b436fbb8..8c6d88b9 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -185,7 +185,7 @@ def initializeProblem( ) # After resolution and corrections, store. - profiles.writeCurrentStatus(file=FolderInitialization / "input.gacode_modified") + profiles.write_state(file=FolderInitialization / "input.gacode_modified") # *************************************************************************************************** # *************************************************************************************************** diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index d45a5773..f7f0070a 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -226,7 +226,7 @@ def to_gacode( write_input_gacode = Path(write_input_gacode) print(f"\t- Writing input.gacode file: {IOtools.clipstr(write_input_gacode)}") write_input_gacode.parent.mkdir(parents=True, exist_ok=True) - profiles.writeCurrentStatus(file=write_input_gacode) + profiles.write_state(file=write_input_gacode) # If corrections modify the ions set... it's better to re-read, otherwise powerstate will be confused if rederive_profiles: diff --git a/src/mitim_tools/astra_tools/ASTRAtools.py b/src/mitim_tools/astra_tools/ASTRAtools.py index ccfbfba1..79e602d7 100644 --- a/src/mitim_tools/astra_tools/ASTRAtools.py +++ b/src/mitim_tools/astra_tools/ASTRAtools.py @@ -348,7 +348,7 @@ def convert_ASTRA_to_gacode_from_transp_output(c, p.printInfo() if gacode_out is not None: - p.writeCurrentStatus(file=gacode_out) + p.write_state(file=gacode_out) if plot_result: p.plot() diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 37ce636d..1e735677 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -110,7 +110,7 @@ def _prerun( control_file = 'input.cgyro.controls' ) - inputCGYRO.writeCurrentStatus() + inputCGYRO.write_state() return input_cgyro_file, inputgacode_file_this @@ -301,7 +301,7 @@ def run_full( control_file = 'input.cgyro.controls' ) - input_cgyro_file_this.writeCurrentStatus() + input_cgyro_file_this.write_state() # Copy the input.gacode file in the subfolder inputgacode_file_this = folder_run / "input.gacode" @@ -743,7 +743,7 @@ def __init__(self, file=None): self.controls = GACODErun.buildDictFromInput(self.file_txt) - def writeCurrentStatus(self, file=None): + def write_state(self, file=None): print("\t- Writting CGYRO input file") if file is None: diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 990470e4..12561fee 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -49,7 +49,7 @@ def run_vgen(self, subfolder="vgen1", vgenOptions={}, cold_start=False): if runThisCase: IOtools.askNewFolder(self.folder_vgen, force=True) - self.inputgacode.writeCurrentStatus(file=(self.folder_vgen / f"input.gacode")) + self.inputgacode.write_state(file=(self.folder_vgen / f"input.gacode")) # ---- Run diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 8b2d3c72..801d59ff 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -240,6 +240,71 @@ def derive_geometry(self, n_theta_geo=1001): self.derived["surfGACODE_geo"] = (self.derived["surf_geo"] / self.derived["gradr_geo"]) self.derived["surfGACODE_geo"][np.isnan(self.derived["surfGACODE_geo"])] = 0 + def write_state(self, file=None, limitedNames=False): + print("\t- Writting input.gacode file") + + if file is None: + file = self.file + + with open(file, "w") as f: + for line in self.header: + f.write(line) + + for i in self.profiles: + if "(" not in i: + f.write(f"# {i}\n") + else: + f.write(f"# {i.split('(')[0]} | {i.split('(')[-1].split(')')[0]}\n") + + if i in self.titles_single: + if i == "name" and limitedNames: + newlist = [self.profiles[i][0]] + for k in self.profiles[i][1:]: + if k not in [ + "D", + "H", + "T", + "He4", + "he4", + "C", + "O", + "Ar", + "W", + ]: + newlist.append("C") + else: + newlist.append(k) + print( + f"\n\n!! Correcting ion names from {self.profiles[i]} to {newlist} to avoid TGYRO radiation error (to solve in future?)\n\n", + typeMsg="w", + ) + listWrite = newlist + else: + listWrite = self.profiles[i] + + if IOtools.isfloat(listWrite[0]): + listWrite = [f"{i:.7e}".rjust(14) for i in listWrite] + f.write(f"{''.join(listWrite)}\n") + else: + f.write(f"{' '.join(listWrite)}\n") + + else: + if len(self.profiles[i].shape) == 1: + for j, val in enumerate(self.profiles[i]): + pos = f"{j + 1}".rjust(3) + valt = f"{round(val,99):.7e}".rjust(15) + f.write(f"{pos}{valt}\n") + else: + for j, val in enumerate(self.profiles[i]): + pos = f"{j + 1}".rjust(3) + txt = "".join([f"{k:.7e}".rjust(15) for k in val]) + f.write(f"{pos}{txt}\n") + + print(f"\t\t~ File {IOtools.clipstr(file)} written") + + # Update file + self.file = file + def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,ax13c] = axs3 diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 4df28493..49991cc5 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -429,7 +429,7 @@ def prep_direct_tglf( for rho in self.inputsTGLF: self.inputsTGLF[rho].file = self.FolderGACODE / f'input.tglf_{rho:.4f}' - self.inputsTGLF[rho].writeCurrentStatus() + self.inputsTGLF[rho].write_state() """ ~~~~~ Create Normalizations ~~~~~ @@ -3854,7 +3854,7 @@ def changeANDwrite_TGLF( else: print('\t- Not applying corrections nor quasineutrality because "TGLFsettings" is None') - inputTGLF_rho.writeCurrentStatus(file=newfile) + inputTGLF_rho.write_state(file=newfile) modInputTGLF[rho] = inputTGLF_rho @@ -4126,7 +4126,7 @@ def addTraceSpecie(self, ZS, MASS, AS=1e-6, position=None, increaseNS=True, posi return position - def writeCurrentStatus(self, file=None): + def write_state(self, file=None): print("\t- Writting TGLF input file") maxSpeciesTGLF = 6 # TGLF cannot handle more than 6 species diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 5b7fc18c..b18e4eeb 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -374,7 +374,7 @@ def run( addControlFunction=GACODEdefaults.addTGLFcontrol, NS=self.loc_n_ion + 1, ) - inputclass_TGLF.writeCurrentStatus(file=self.FolderTGYRO_tmp / "input.tglf") + inputclass_TGLF.write_state(file=self.FolderTGYRO_tmp / "input.tglf") # ----------------------------------- # ------ Write input profiles @@ -391,7 +391,7 @@ def run( if "z_eff(-)" not in self.profiles.profiles: self.profiles.profiles["z_eff(-)"] = self.profiles.derived["Zeff"] - self.profiles.writeCurrentStatus(file=self.FolderTGYRO_tmp / f"{fil}") + self.profiles.write_state(file=self.FolderTGYRO_tmp / f"{fil}") # ----------------------------------- # ------ Create TGYRO file @@ -410,7 +410,7 @@ def run( special_radii=special_radii_mod, ) - inputclass_TGYRO.writeCurrentStatus(file=self.FolderTGYRO_tmp / "input.tgyro") + inputclass_TGYRO.write_state(file=self.FolderTGYRO_tmp / "input.tgyro") # ----------------------------------- # ------ Check density for problems @@ -499,7 +499,7 @@ def run( inputgacode_new.profiles["rho(-)"] * 0.0 ) - inputgacode_new.writeCurrentStatus() + inputgacode_new.write_state() # ------------------------------------------------------------------------------------------------------------------------ # Copy those files that I'm interested in, plus the extra file, into the main folder @@ -4586,7 +4586,7 @@ def __init__(self, input_profiles, file=None, onlyThermal=False, limitSpecies=10 ) self.loc_n_ion = spec["LOC_N_ION"] - def writeCurrentStatus(self, file=None): + def write_state(self, file=None): print("\t- Writting TGYRO input file") if file is None: diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 8d67b764..3c24cd5d 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -450,7 +450,7 @@ def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', ne0_20 folder.mkdir(parents=True, exist_ok=True) p = self.to_profiles(ne0_20 = ne0_20, Zeff = Zeff, PichT = PichT_MW) - p.writeCurrentStatus(folder / 'input.gacode') + p.write_state(folder / 'input.gacode') transp = p.to_transp(folder = folder, shot = shot, runid = runid, times = times, Vsurf = Vsurf) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 5edd359d..5a31db52 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -126,7 +126,10 @@ def derive_quantities(self, mi_ref=None, derive_quantities=True, rederiveGeometr if derive_quantities: self.derive_quantities_full(rederiveGeometry=rederiveGeometry) - def derive_geometry(self, **kwargs): + def derive_geometry(self, *args, **kwargs): + raise Exception('[MITIM] This method is not implemented in the base class. Please use a derived class that implements it.') + + def write_state(self, *args, **kwargs): raise Exception('[MITIM] This method is not implemented in the base class. Please use a derived class that implements it.') # ------------------------------------------------------------------------------------- @@ -198,7 +201,7 @@ def calculate_Er( self.derive_quantities() if write_new_file is not None: - self.writeCurrentStatus(file=write_new_file) + self.write_state(file=write_new_file) def readSpecies(self, maxSpecies=100): @@ -1192,98 +1195,6 @@ def toNumpyArrays(self): self.profiles.update({key: tensor.cpu().detach().cpu().numpy() for key, tensor in self.profiles.items() if isinstance(tensor, torch.Tensor)}) self.derived.update({key: tensor.cpu().detach().cpu().numpy() for key, tensor in self.derived.items() if isinstance(tensor, torch.Tensor)}) - def writeCurrentStatus(self, file=None, limitedNames=False): - print("\t- Writting input.gacode file") - - if file is None: - file = self.file - - with open(file, "w") as f: - for line in self.header: - f.write(line) - - for i in self.profiles: - if "(" not in i: - f.write(f"# {i}\n") - else: - f.write(f"# {i.split('(')[0]} | {i.split('(')[-1].split(')')[0]}\n") - - if i in self.titles_single: - if i == "name" and limitedNames: - newlist = [self.profiles[i][0]] - for k in self.profiles[i][1:]: - if k not in [ - "D", - "H", - "T", - "He4", - "he4", - "C", - "O", - "Ar", - "W", - ]: - newlist.append("C") - else: - newlist.append(k) - print( - f"\n\n!! Correcting ion names from {self.profiles[i]} to {newlist} to avoid TGYRO radiation error (to solve in future?)\n\n", - typeMsg="w", - ) - listWrite = newlist - else: - listWrite = self.profiles[i] - - if IOtools.isfloat(listWrite[0]): - listWrite = [f"{i:.7e}".rjust(14) for i in listWrite] - f.write(f"{''.join(listWrite)}\n") - else: - f.write(f"{' '.join(listWrite)}\n") - - else: - if len(self.profiles[i].shape) == 1: - for j, val in enumerate(self.profiles[i]): - pos = f"{j + 1}".rjust(3) - valt = f"{round(val,99):.7e}".rjust(15) - f.write(f"{pos}{valt}\n") - else: - for j, val in enumerate(self.profiles[i]): - pos = f"{j + 1}".rjust(3) - txt = "".join([f"{k:.7e}".rjust(15) for k in val]) - f.write(f"{pos}{txt}\n") - - print(f"\t\t~ File {IOtools.clipstr(file)} written") - - # Update file - self.file = file - - def writeMiminalKinetic(self, file): - setProfs = [ - "rho(-)", - "polflux(Wb/radian)", - "q(-)", - "te(keV)", - "ti(keV)", - "ne(10^19/m^3)", - ] - - with open(file, "w") as f: - for i in setProfs: - if "(" not in i: - f.write(f"# {i}\n") - else: - f.write(f"# {i.split('(')[0]} | {i.split('(')[-1].split(')')[0]}\n") - - if len(self.profiles[i].shape) > 1: - p = self.profiles[i][:, 0] - else: - p = self.profiles[i] - - for j, val in enumerate(p): - pos = f"{j + 1}".rjust(3) - valt = f"{val:.7e}".rjust(15) - f.write(f"{pos}{valt}\n") - def changeResolution(self, n=100, rho_new=None, interpolation_function=MATHtools.extrapolateCubicSpline): rho = copy.deepcopy(self.profiles["rho(-)"]) @@ -1689,7 +1600,7 @@ def correct(self, options={}, write=False, new_file=None): # Write # ---------------------------------------------------------------------- if write: - self.writeCurrentStatus(file=new_file) + self.write_state(file=new_file) self.printInfo() def enforce_same_density_gradients(self): @@ -1761,7 +1672,7 @@ def introduceRotationProfile(self, Mach_LF=1.0, new_file=None): self.derive_quantities() if new_file is not None: - self.writeCurrentStatus(file=new_file) + self.write_state(file=new_file) def plot( diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index 591657da..bb58a007 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -1,5 +1,8 @@ +import vmecpp from mitim_tools.misc_tools import IOtools from mitim_tools.plasmastate_tools import MITIMstate +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed class vmec_state(MITIMstate.mitim_state): ''' @@ -12,16 +15,25 @@ class vmec_state(MITIMstate.mitim_state): def __init__(self, file, derive_quantities=True, mi_ref=None): + # Initialize the base class and tell it the type of file super().__init__(type_file='vmec') + # Read the input file and store the raw data self.file = file + self._read_inputgacocde() - self._read_vmec() - + # Derive quantities if requested if self.file is not None: # Derive (Depending on resolution, derived can be expensive, so I mmay not do it every time) self.derive_quantities(mi_ref=mi_ref, derive_quantities=derive_quantities) + @IOtools.hook_method(after=MITIMstate.ensure_variables_existence) def _read_vmec(self): pass + + def derive_geometry(self, **kwargs): + pass + + def write_state(self, file=None, **kwargs): + pass From c9d0c4c70f6d672aba9302da70a4ff1d39bdd044 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 23 Jun 2025 15:18:45 +0200 Subject: [PATCH 097/385] misc --- src/mitim_tools/gacode_tools/PROFILEStools.py | 122 +++++++++--------- .../plasmastate_tools/MITIMstate.py | 1 - 2 files changed, 61 insertions(+), 62 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 801d59ff..37001307 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -27,73 +27,73 @@ def __init__(self, file, derive_quantities=True, mi_ref=None): super().__init__(type_file='input.gacode') # Read the input file and store the raw data - self.file = file - self._read_inputgacocde() + self.files = [file] + + self.titles_singleNum = ["nexp", "nion", "shot", "name", "type", "time"] + self.titles_singleArr = ["masse","mass","ze","z","torfluxa(Wb/radian)","rcentr(m)","bcentr(T)","current(MA)"] + self.titles_single = self.titles_singleNum + self.titles_singleArr + + if self.files is not None: - # Derive quantities if requested - if self.file is not None: + self._read_inputgacocde() + # Derive (Depending on resolution, derived can be expensive, so I mmay not do it every time) self.derive_quantities(mi_ref=mi_ref, derive_quantities=derive_quantities) @IOtools.hook_method(after=MITIMstate.ensure_variables_existence) def _read_inputgacocde(self): - self.titles_singleNum = ["nexp", "nion", "shot", "name", "type", "time"] - self.titles_singleArr = ["masse","mass","ze","z","torfluxa(Wb/radian)","rcentr(m)","bcentr(T)","current(MA)"] - self.titles_single = self.titles_singleNum + self.titles_singleArr + with open(self.files[0], "r") as f: + self.lines = f.readlines() + + # Read file and store raw data + self._read_header() + self._read_profiles() + + # Ensure correctness (wrong names in older input.gacode files) + if "qmom(Nm)" in self.profiles: + self.profiles["qmom(N/m^2)"] = self.profiles.pop("qmom(Nm)") + if "qpar_beam(MW/m^3)" in self.profiles: + self.profiles["qpar_beam(1/m^3/s)"] = self.profiles.pop("qpar_beam(MW/m^3)") + if "qpar_wall(MW/m^3)" in self.profiles: + self.profiles["qpar_wall(1/m^3/s)"] = self.profiles.pop("qpar_wall(MW/m^3)") - if self.file is not None: - with open(self.file, "r") as f: - self.lines = f.readlines() - - # Read file and store raw data - self._read_header() - self._read_profiles() - - # Ensure correctness (wrong names in older input.gacode files) - if "qmom(Nm)" in self.profiles: - self.profiles["qmom(N/m^2)"] = self.profiles.pop("qmom(Nm)") - if "qpar_beam(MW/m^3)" in self.profiles: - self.profiles["qpar_beam(1/m^3/s)"] = self.profiles.pop("qpar_beam(MW/m^3)") - if "qpar_wall(MW/m^3)" in self.profiles: - self.profiles["qpar_wall(1/m^3/s)"] = self.profiles.pop("qpar_wall(MW/m^3)") - - """ - Note that in prgen_map_plasmastate, that variable: - expro_qpar_beam(i) = plst_sn_trans(i-1)/dvol - - Note that in prgen_read_plasmastate, that variable: - ! Particle source - err = nf90_inq_varid(ncid,trim('sn_trans'),varid) - err = nf90_get_var(ncid,varid,plst_sn_trans(1:nx-1)) - plst_sn_trans(nx) = 0.0 - - Note that in the plasmastate file, the variable "sn_trans": - - long_name: particle transport (loss) - units: #/sec - component: PLASMA - section: STATE_PROFILES - specification: R|units=#/sec|step*dV sn_trans(~nrho,0:nspec_th) - - So, this means that expro_qpar_beam is in units of #/sec/m^3, meaning that - it is a particle flux DENSITY. It therefore requires volume integral and - divide by surface to produce a flux. - - The units of this qpar_beam column is NOT MW/m^3. In the gacode source codes - they also say that those units are wrong. - - """ - - # Ensure that we also have the shape coefficients - num_moments = 7 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros - if "shape_cos0(-)" not in self.profiles: - self.profiles["shape_cos0(-)"] = np.ones(self.profiles["rmaj(m)"].shape) - for i in range(num_moments): - if f"shape_cos{i + 1}(-)" not in self.profiles: - self.profiles[f"shape_cos{i + 1}(-)"] = np.zeros(self.profiles["rmaj(m)"].shape) - if f"shape_sin{i + 1}(-)" not in self.profiles and i > 1: - self.profiles[f"shape_sin{i + 1}(-)"] = np.zeros(self.profiles["rmaj(m)"].shape) + """ + Note that in prgen_map_plasmastate, that variable: + expro_qpar_beam(i) = plst_sn_trans(i-1)/dvol + + Note that in prgen_read_plasmastate, that variable: + ! Particle source + err = nf90_inq_varid(ncid,trim('sn_trans'),varid) + err = nf90_get_var(ncid,varid,plst_sn_trans(1:nx-1)) + plst_sn_trans(nx) = 0.0 + + Note that in the plasmastate file, the variable "sn_trans": + + long_name: particle transport (loss) + units: #/sec + component: PLASMA + section: STATE_PROFILES + specification: R|units=#/sec|step*dV sn_trans(~nrho,0:nspec_th) + + So, this means that expro_qpar_beam is in units of #/sec/m^3, meaning that + it is a particle flux DENSITY. It therefore requires volume integral and + divide by surface to produce a flux. + + The units of this qpar_beam column is NOT MW/m^3. In the gacode source codes + they also say that those units are wrong. + + """ + + # Ensure that we also have the shape coefficients + num_moments = 7 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros + if "shape_cos0(-)" not in self.profiles: + self.profiles["shape_cos0(-)"] = np.ones(self.profiles["rmaj(m)"].shape) + for i in range(num_moments): + if f"shape_cos{i + 1}(-)" not in self.profiles: + self.profiles[f"shape_cos{i + 1}(-)"] = np.zeros(self.profiles["rmaj(m)"].shape) + if f"shape_sin{i + 1}(-)" not in self.profiles and i > 1: + self.profiles[f"shape_sin{i + 1}(-)"] = np.zeros(self.profiles["rmaj(m)"].shape) def _read_header(self): for i in range(len(self.lines)): @@ -244,7 +244,7 @@ def write_state(self, file=None, limitedNames=False): print("\t- Writting input.gacode file") if file is None: - file = self.file + file = self.files[0] with open(file, "w") as f: for line in self.header: @@ -303,7 +303,7 @@ def write_state(self, file=None, limitedNames=False): print(f"\t\t~ File {IOtools.clipstr(file)} written") # Update file - self.file = file + self.files[0] = file def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 5a31db52..c1652e89 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -50,7 +50,6 @@ def ensure_variables_existence(self): "delta(-)": 1, "zeta(-)": 1, }) - # Sources and Sinks required_profiles.update({ From 1d55e0d7da9d1f237d29a2128b15dfd9c27c1b77 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 24 Jun 2025 15:23:32 +0200 Subject: [PATCH 098/385] Vmec class start --- .../plasmastate_tools/utils/VMECtools.py | 40 +++++++++++++++---- 1 file changed, 32 insertions(+), 8 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index bb58a007..9a1f1795 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -1,4 +1,6 @@ import vmecpp +import numpy as np +from pathlib import Path from mitim_tools.misc_tools import IOtools from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -13,27 +15,49 @@ class vmec_state(MITIMstate.mitim_state): # Reading and interpreting # ------------------------------------------------------------------ - def __init__(self, file, derive_quantities=True, mi_ref=None): + def __init__(self, file_vmec, file_profs, derive_quantities=True, mi_ref=None): # Initialize the base class and tell it the type of file super().__init__(type_file='vmec') # Read the input file and store the raw data - self.file = file - self._read_inputgacocde() - - # Derive quantities if requested - if self.file is not None: + self.files = [file_vmec, file_profs] + if self.files is not None: + self._read_vmec() # Derive (Depending on resolution, derived can be expensive, so I mmay not do it every time) self.derive_quantities(mi_ref=mi_ref, derive_quantities=derive_quantities) @IOtools.hook_method(after=MITIMstate.ensure_variables_existence) def _read_vmec(self): - pass + + # Read VMEC file + print("\t- Reading VMEC file") + self.wout = vmecpp.VmecWOut.from_wout_file(Path(self.files[0])) + + # Produce variables + + self.profiles["torfluxa(Wb/radian)"] = [self.wout.phipf[-1]] + def derive_geometry(self, **kwargs): - pass + + rho = np.linspace(0, 1, self.wout.ns) + + ds = rho[1] - rho[0] + half_grid_rho = rho - ds / 2 + + d_volume_d_rho = ( + (2 * np.pi) ** 2 + * np.array(self.wout.vp) + * 2 + * np.sqrt(half_grid_rho) + ) + + #self.derived["B_unit"] = self.profiles["torfluxa(Wb/radian)"] / (np.pi * self.wout.Aminor_p**2) def write_state(self, file=None, **kwargs): pass + + def plot_geometry(self, axs, color="b", legYN=True, extralab="", lw=1, fs=6): + pass \ No newline at end of file From e9943ab9d3b8612d3c0b99603493589977ec08c1 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 23 Jul 2025 23:30:28 -0400 Subject: [PATCH 099/385] Minimal example of Profiles for stellarators working --- src/mitim_tools/gacode_tools/PROFILEStools.py | 13 +- .../plasmastate_tools/MITIMstate.py | 14 +- .../plasmastate_tools/utils/VMECtools.py | 313 +++++++++++++++++- 3 files changed, 321 insertions(+), 19 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 37001307..6b52a002 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -240,6 +240,17 @@ def derive_geometry(self, n_theta_geo=1001): self.derived["surfGACODE_geo"] = (self.derived["surf_geo"] / self.derived["gradr_geo"]) self.derived["surfGACODE_geo"][np.isnan(self.derived["surfGACODE_geo"])] = 0 + + self.derived["kappa95"] = np.interp(0.95, self.derived["psi_pol_n"], self.profiles["kappa(-)"]) + + self.derived["kappa995"] = np.interp(0.995, self.derived["psi_pol_n"], self.profiles["kappa(-)"]) + + self.derived["delta95"] = np.interp(0.95, self.derived["psi_pol_n"], self.profiles["delta(-)"]) + + self.derived["delta995"] = np.interp(0.995, self.derived["psi_pol_n"], self.profiles["delta(-)"]) + + self.derived["kappa_a"] = self.derived["surfXS"][-1] / np.pi / self.derived["a"] ** 2 + def write_state(self, file=None, limitedNames=False): print("\t- Writting input.gacode file") @@ -457,8 +468,6 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): GRAPHICStools.addDenseAxis(ax) - - def calculateGeometricFactors(profiles, n_theta=1001): # ---------------------------------------- diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index c1652e89..7b75157f 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -74,7 +74,7 @@ def ensure_variables_existence(self): # --------------------------------------------------------------------------- # Choose a template for dimensionality - template_key_1d, template_key_2d = "rmin(m)", "ti(keV)" + template_key_1d, template_key_2d = "rho(-)", "ti(keV)" # Ensure required keys exist for key, dim in required_profiles.items(): @@ -242,7 +242,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): """ deriving geometry is expensive, so if I'm just updating profiles it may not be needed """ - + if "derived" not in self.__dict__: self.derived = {} @@ -297,16 +297,6 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): # --------- Important for scaling laws # --------------------------------------------------------------------------------------------------------------------- - self.derived["kappa95"] = np.interp(0.95, self.derived["psi_pol_n"], self.profiles["kappa(-)"]) - - self.derived["kappa995"] = np.interp(0.995, self.derived["psi_pol_n"], self.profiles["kappa(-)"]) - - self.derived["kappa_a"] = self.derived["surfXS"][-1] / np.pi / self.derived["a"] ** 2 - - self.derived["delta95"] = np.interp(0.95, self.derived["psi_pol_n"], self.profiles["delta(-)"]) - - self.derived["delta995"] = np.interp(0.995, self.derived["psi_pol_n"], self.profiles["delta(-)"]) - self.derived["Rgeo"] = float(self.profiles["rcentr(m)"][-1]) self.derived["B0"] = np.abs(float(self.profiles["bcentr(T)"][-1])) diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index 9a1f1795..137bcdba 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -1,6 +1,8 @@ import vmecpp import numpy as np +import matplotlib.pyplot as plt from pathlib import Path +from scipy.interpolate import interp1d from mitim_tools.misc_tools import IOtools from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -15,7 +17,13 @@ class vmec_state(MITIMstate.mitim_state): # Reading and interpreting # ------------------------------------------------------------------ - def __init__(self, file_vmec, file_profs, derive_quantities=True, mi_ref=None): + def __init__( + self, + file_vmec, + file_profs, + derive_quantities=True, + mi_ref=None + ): # Initialize the base class and tell it the type of file super().__init__(type_file='vmec') @@ -24,22 +32,48 @@ def __init__(self, file_vmec, file_profs, derive_quantities=True, mi_ref=None): self.files = [file_vmec, file_profs] if self.files is not None: self._read_vmec() - # Derive (Depending on resolution, derived can be expensive, so I mmay not do it every time) + + # Derive (Depending on resolution, derived can be expensive, so I may not do it every time) self.derive_quantities(mi_ref=mi_ref, derive_quantities=derive_quantities) @IOtools.hook_method(after=MITIMstate.ensure_variables_existence) def _read_vmec(self): + self.profiles = {} + # Read VMEC file print("\t- Reading VMEC file") self.wout = vmecpp.VmecWOut.from_wout_file(Path(self.files[0])) # Produce variables - + self.profiles["rho(-)"] = np.linspace(0, 1, self.wout.ns)**0.5 + self.profiles["rmin(m)"] = self.profiles["rho(-)"] + self.profiles["torfluxa(Wb/radian)"] = [self.wout.phipf[-1]] - + # Read Profiles + data = self._read_profiles(x_coord=self.profiles["rho(-)"]) + + self.profiles['te(keV)'] = data['Te'] + self.profiles['ne(10^19/m^3)'] = data['ne'] + self.profiles['ti(keV)'] = np.atleast_2d(data['Ti']).T + self.profiles['ni(10^19/m^3)'] = np.atleast_2d(data['ni']).T + + self.profiles['qbeami(MW/m^3)'] = data['Qi'] * 1e-6 # Convert from W/m^3 to MW/m^3 + self.profiles['qrfe(MW/m^3)'] = data['Qe'] * 1e-6 # Convert from W/m^3 to MW/m^3 + self.profiles['qpar_beam(1/m^3/s)'] = data['S'] + + self.profiles['nion'] = np.array([1]) + self.profiles['name'] = np.array(['D']) + self.profiles['type'] = np.array(['[therm]']) + self.profiles['mass'] = np.array([2.0]) + self.profiles['z'] = np.array([1.0]) + + self.profiles['rcentr(m)'] = np.array([self.wout.Rmajor_p]) + self.profiles['bcentr(T)'] = np.array([self.wout.rbtor/self.wout.Rmajor_p]) + self.profiles["current(MA)"] = np.array([0.0]) + def derive_geometry(self, **kwargs): rho = np.linspace(0, 1, self.wout.ns) @@ -53,11 +87,280 @@ def derive_geometry(self, **kwargs): * 2 * np.sqrt(half_grid_rho) ) + d_volume_d_rho[0] = 0.0 # Set the first element to zero to avoid division by zero #self.derived["B_unit"] = self.profiles["torfluxa(Wb/radian)"] / (np.pi * self.wout.Aminor_p**2) + + self.derived["volp_geo"] = d_volume_d_rho + + self.derived["kappa_a"] = 0.0 + self.derived["kappa95"] = 0.0 + self.derived["delta95"] = 0.0 + self.derived["kappa995"] = 0.0 + self.derived["delta995"] = 0.0 + self.derived["R_LF"] = np.zeros(self.profiles["rho(-)"].shape) + + self.derived["bp2_exp"] = np.zeros(self.profiles["rho(-)"].shape) + self.derived["bt2_exp"] = np.zeros(self.profiles["rho(-)"].shape) + self.derived["bp2_geo"] = np.zeros(self.profiles["rho(-)"].shape) + self.derived["bt2_geo"] = np.zeros(self.profiles["rho(-)"].shape) + def write_state(self, file=None, **kwargs): pass def plot_geometry(self, axs, color="b", legYN=True, extralab="", lw=1, fs=6): - pass \ No newline at end of file + + pass + #self.plot_plasma_boundary() + + def _read_profiles(self, x_coord=None, debug = False): + + filename = self.files[1] + + if x_coord is None: + # Create uniform coordinate array from 0 to 1 + x_coord = np.linspace(0, 1, 200) + + # Raw data storage + raw_data = { + 'Te_data': {'x': [], 'y': []}, + 'ne_data': {'x': [], 'y': []}, + 'Ti_data': {'x': [], 'y': []}, + 'ni_data': {'x': [], 'y': []}, + 'S_data': {'x': [], 'y': []}, + 'Qe_data': {'x': [], 'y': []}, + 'Qi_data': {'x': [], 'y': []} + } + current_section = None + + with open(filename, 'r') as f: + lines = f.readlines() + + for line in lines: + line = line.strip() + if not line: + continue + + # Check for section headers + if line.startswith('#'): + if 'Te' in line and 'ne' in line: + current_section = 'Te_ne' + elif 'Ti' in line and 'ni' in line: + current_section = 'Ti_ni' + elif 'Particle source' in line: + current_section = 'S' + elif 'Q W/m3' in line: + current_section = 'Qe' + elif 'NBI W/m3' in line: + current_section = 'Qi' + continue + + # Parse data lines + parts = line.split() + if len(parts) < 2: + continue + + x = float(parts[0]) # x/a or r/a coordinate (first column) + + if current_section == 'Te_ne' and len(parts) == 4: + # x/a, x/rho, Te, ne - use x/a coordinate + raw_data['Te_data']['x'].append(x) + raw_data['Te_data']['y'].append(float(parts[2])) + raw_data['ne_data']['x'].append(x) + raw_data['ne_data']['y'].append(float(parts[3])) + + elif current_section == 'Ti_ni' and len(parts) == 4: + # x/a, x/rho, Ti, ni - use x/a coordinate + raw_data['Ti_data']['x'].append(x) + raw_data['Ti_data']['y'].append(float(parts[2])) + raw_data['ni_data']['x'].append(x) + raw_data['ni_data']['y'].append(float(parts[3])) + + elif current_section == 'S' and len(parts) == 2: + raw_data['S_data']['x'].append(x) + raw_data['S_data']['y'].append(float(parts[1])) + + elif current_section == 'Qe' and len(parts) == 2: + raw_data['Qe_data']['x'].append(x) + raw_data['Qe_data']['y'].append(float(parts[1])) + + elif current_section == 'Qi' and len(parts) == 2: + raw_data['Qi_data']['x'].append(x) + raw_data['Qi_data']['y'].append(float(parts[1])) + + # Convert to numpy arrays + for profile_data in raw_data.values(): + profile_data['x'] = np.array(profile_data['x']) + profile_data['y'] = np.array(profile_data['y']) + + # Interpolate each profile to uniform grid + uniform_data = {'x_coord': x_coord} + + profile_map = { + 'Te': 'Te_data', 'ne': 'ne_data', 'Ti': 'Ti_data', + 'ni': 'ni_data', 'S': 'S_data', 'Qe': 'Qe_data', 'Qi': 'Qi_data' + } + + for profile_name, data_key in profile_map.items(): + x_data = raw_data[data_key]['x'] + y_data = raw_data[data_key]['y'] + + if len(x_data) > 0: + # Interpolate using actual coordinates from the data + f = interp1d(x_data, y_data, kind='linear', + bounds_error=False, fill_value=(y_data[0], y_data[-1])) + uniform_data[profile_name] = f(x_coord) + else: + uniform_data[profile_name] = None + + # Also store original data for plotting + uniform_data['raw_data'] = raw_data + + if debug: + plot_profiles(uniform_data) + embed() + + return uniform_data + + + def plot_plasma_boundary(self, ax=None): + + # The output object contains the Fourier coefficients of the geometry in R and Z + # as a function of the poloidal (theta) and toroidal (phi) angle-like coordinates + # for a number of discrete radial locations. + + # number of flux surfaces, i.e., final radial resolution + ns = self.wout.ns + + # poloidal mode numbers: m + xm = self.wout.xm + + # toroidal mode numbers: n * nfp + xn = self.wout.xn + + # stellarator-symmetric Fourier coefficients of flux surface geometry R ~ cos(m * theta - n * nfp * phi) + rmnc = self.wout.rmnc + + # stellarator-symmetric Fourier coefficients of flux surface geometry Z ~ sin(m * theta - n * nfp * phi) + zmns = self.wout.zmns + + # plot the outermost (last) flux surface, which is the plasma boundary + j = ns - 1 + + # resolution over the flux surface + num_theta = 101 + num_phi = 181 + + min_phi = 0.0 + max_phi = 2.0 * np.pi + + # grid in theta and phi along the flux surface + grid_theta = np.linspace(0.0, 2.0 * np.pi, num_theta, endpoint=True) + grid_phi = np.linspace(min_phi, max_phi, num_phi, endpoint=True) + + # compute Cartesian coordinates of flux surface geometry + x = np.zeros([num_theta, num_phi]) + y = np.zeros([num_theta, num_phi]) + z = np.zeros([num_theta, num_phi]) + for idx_theta, theta in enumerate(grid_theta): + for idx_phi, phi in enumerate(grid_phi): + kernel = xm * theta - xn * phi + r = np.dot(rmnc[:, j], np.cos(kernel)) + x[idx_theta, idx_phi] = r * np.cos(phi) + y[idx_theta, idx_phi] = r * np.sin(phi) + z[idx_theta, idx_phi] = np.dot(zmns[:, j], np.sin(kernel)) + + # actually make the 3D plot + if ax is None: + fig = plt.figure() + ax = fig.add_subplot(projection="3d") + + # Plot the surface + ax.plot_surface(x, y, z) + + # Set an equal aspect ratio + ax.set_aspect("equal") + + plt.show() + +def plot_profiles(data): + """ + Create plots of the plasma profiles + """ + fig, axes = plt.subplots(2, 3, figsize=(15, 10)) + + raw_data = data['raw_data'] + + # Temperature profiles + if len(raw_data['Te_data']['x']) > 0: + axes[0,0].plot(raw_data['Te_data']['x'], raw_data['Te_data']['y'], 'ro-', label='Te (original)', markersize=3) + if data['Te'] is not None: + axes[0,0].plot(data['x_coord'], data['Te'], 'r-', label='Te (interpolated)', linewidth=2) + + if len(raw_data['Ti_data']['x']) > 0: + axes[0,0].plot(raw_data['Ti_data']['x'], raw_data['Ti_data']['y'], 'bo-', label='Ti (original)', markersize=3) + if data['Ti'] is not None: + axes[0,0].plot(data['x_coord'], data['Ti'], 'b-', label='Ti (interpolated)', linewidth=2) + + axes[0,0].set_xlabel('x/a') + axes[0,0].set_ylabel('Temperature [keV]') + axes[0,0].legend() + axes[0,0].grid(True) + axes[0,0].set_title('Temperature Profiles') + + # Density profiles + if len(raw_data['ne_data']['x']) > 0: + axes[0,1].plot(raw_data['ne_data']['x'], raw_data['ne_data']['y'], 'ro-', label='ne (original)', markersize=3) + if data['ne'] is not None: + axes[0,1].plot(data['x_coord'], data['ne'], 'r-', label='ne (interpolated)', linewidth=2) + + if len(raw_data['ni_data']['x']) > 0: + axes[0,1].plot(raw_data['ni_data']['x'], raw_data['ni_data']['y'], 'bo-', label='ni (original)', markersize=3) + if data['ni'] is not None: + axes[0,1].plot(data['x_coord'], data['ni'], 'b-', label='ni (interpolated)', linewidth=2) + + axes[0,1].set_xlabel('x/a') + axes[0,1].set_ylabel('Density [10¹⁹ m⁻³]') + axes[0,1].legend() + axes[0,1].grid(True) + axes[0,1].set_title('Density Profiles') + + # Particle source + if len(raw_data['S_data']['x']) > 0: + axes[0,2].plot(raw_data['S_data']['x'], raw_data['S_data']['y'], 'go-', label='S (original)', markersize=3) + if data['S'] is not None: + axes[0,2].plot(data['x_coord'], data['S'], 'g-', label='S (interpolated)', linewidth=2) + axes[0,2].set_xlabel('x/a') + axes[0,2].set_ylabel('Particle Source [m⁻³s⁻¹]') + axes[0,2].legend() + axes[0,2].grid(True) + axes[0,2].set_title('Particle Source') + + # Electron heating + if len(raw_data['Qe_data']['x']) > 0: + axes[1,0].plot(raw_data['Qe_data']['x'], raw_data['Qe_data']['y'], 'ro-', label='Qe (original)', markersize=3) + if data['Qe'] is not None: + axes[1,0].plot(data['x_coord'], data['Qe'], 'r-', label='Qe (interpolated)', linewidth=2) + axes[1,0].set_xlabel('x/a') + axes[1,0].set_ylabel('Qe [W/m³]') + axes[1,0].legend() + axes[1,0].grid(True) + axes[1,0].set_title('Electron Heating') + + # Ion heating + if len(raw_data['Qi_data']['x']) > 0: + axes[1,1].plot(raw_data['Qi_data']['x'], raw_data['Qi_data']['y'], 'bo-', label='Qi (original)', markersize=3) + if data['Qi'] is not None: + axes[1,1].plot(data['x_coord'], data['Qi'], 'b-', label='Qi (interpolated)', linewidth=2) + axes[1,1].set_xlabel('x/a') + axes[1,1].set_ylabel('Qi [W/m³]') + axes[1,1].legend() + axes[1,1].grid(True) + axes[1,1].set_title('Ion Heating (NBI)') + + # Remove empty subplot + axes[1,2].remove() + + plt.tight_layout() + plt.show() From 65ac743cd3b2c03d8eaa6173c075ba0bc0363f0c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Jul 2025 10:24:02 -0400 Subject: [PATCH 100/385] Code organization and writter of vmec state --- src/mitim_tools/gacode_tools/PROFILEStools.py | 80 +--- .../plasmastate_tools/MITIMstate.py | 416 ++++++++++-------- .../plasmastate_tools/utils/VMECtools.py | 68 +-- 3 files changed, 267 insertions(+), 297 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 6b52a002..c2caa3f1 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -17,9 +17,9 @@ class gacode_state(MITIMstate.mitim_state): and writes them in the way that MITIMstate class expects. ''' - # ------------------------------------------------------------------ + # ************************************************************************************************************************************************ # Reading and interpreting input.gacode files - # ------------------------------------------------------------------ + # ************************************************************************************************************************************************ def __init__(self, file, derive_quantities=True, mi_ref=None): @@ -157,9 +157,9 @@ def _read_profiles(self): self.profiles["w0(rad/s)"] = self.profiles["omega0(rad/s)"] del self.profiles["omega0(rad/s)"] - # ------------------------------------------------------------------ + # ************************************************************************************************************************************************ # Derivation (different from MITIMstate) - # ------------------------------------------------------------------ + # ************************************************************************************************************************************************ def derive_quantities(self, **kwargs): @@ -187,11 +187,7 @@ def _produce_shape_lists(self): self.profiles["shape_sin6(-)"], ] - # ------------------------------------------------------------------ - # Geometry - # ------------------------------------------------------------------ - - def derive_geometry(self, n_theta_geo=1001): + def derive_geometry(self, n_theta_geo=1001, **kwargs): self._produce_shape_lists() @@ -251,71 +247,6 @@ def derive_geometry(self, n_theta_geo=1001): self.derived["kappa_a"] = self.derived["surfXS"][-1] / np.pi / self.derived["a"] ** 2 - def write_state(self, file=None, limitedNames=False): - print("\t- Writting input.gacode file") - - if file is None: - file = self.files[0] - - with open(file, "w") as f: - for line in self.header: - f.write(line) - - for i in self.profiles: - if "(" not in i: - f.write(f"# {i}\n") - else: - f.write(f"# {i.split('(')[0]} | {i.split('(')[-1].split(')')[0]}\n") - - if i in self.titles_single: - if i == "name" and limitedNames: - newlist = [self.profiles[i][0]] - for k in self.profiles[i][1:]: - if k not in [ - "D", - "H", - "T", - "He4", - "he4", - "C", - "O", - "Ar", - "W", - ]: - newlist.append("C") - else: - newlist.append(k) - print( - f"\n\n!! Correcting ion names from {self.profiles[i]} to {newlist} to avoid TGYRO radiation error (to solve in future?)\n\n", - typeMsg="w", - ) - listWrite = newlist - else: - listWrite = self.profiles[i] - - if IOtools.isfloat(listWrite[0]): - listWrite = [f"{i:.7e}".rjust(14) for i in listWrite] - f.write(f"{''.join(listWrite)}\n") - else: - f.write(f"{' '.join(listWrite)}\n") - - else: - if len(self.profiles[i].shape) == 1: - for j, val in enumerate(self.profiles[i]): - pos = f"{j + 1}".rjust(3) - valt = f"{round(val,99):.7e}".rjust(15) - f.write(f"{pos}{valt}\n") - else: - for j, val in enumerate(self.profiles[i]): - pos = f"{j + 1}".rjust(3) - txt = "".join([f"{k:.7e}".rjust(15) for k in val]) - f.write(f"{pos}{txt}\n") - - print(f"\t\t~ File {IOtools.clipstr(file)} written") - - # Update file - self.files[0] = file - def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,ax13c] = axs3 @@ -342,7 +273,6 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): ax.set_xlabel("$\\rho$") ax.set_ylabel(varL) - GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 7b75157f..e4f5ac06 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -15,6 +15,7 @@ from mitim_tools import __version__ from IPython import embed + def ensure_variables_existence(self): # --------------------------------------------------------------------------- # Determine minimal set of variables that should be present in the profiles @@ -81,6 +82,19 @@ def ensure_variables_existence(self): if key not in self.profiles: self.profiles[key] = copy.deepcopy(self.profiles[template_key_1d]) * 0.0 if dim == 1 else copy.deepcopy(self.profiles[template_key_2d]) * 0.0 + +''' +The mitim_state class is the base class for manipulating plasma states in MITIM. +Any class that inherits from this class should implement the methods: + + - derive_geometry: to derive the geometry of the plasma state. + + - write_state: to write the plasma state to a file. + + - plot_geometry: to plot the geometry of the plasma state. + +''' + class mitim_state: ''' Class to manipulate the plasma state in MITIM. @@ -90,6 +104,25 @@ def __init__(self, type_file = 'input.gacode'): self.type = type_file + @classmethod + def scratch(cls, profiles, label_header='', **kwargs_process): + # Method to write a scratch file + + instance = cls(None) + + # Header + instance.header = f''' +# Created from scratch with MITIM version {__version__} +# {label_header} +# +''' + # Add data to profiles + instance.profiles = profiles + + instance.process(**kwargs_process) + + return instance + @IOtools.hook_method(before=ensure_variables_existence) def derive_quantities(self, mi_ref=None, derive_quantities=True, rederiveGeometry=True): @@ -125,34 +158,62 @@ def derive_quantities(self, mi_ref=None, derive_quantities=True, rederiveGeometr if derive_quantities: self.derive_quantities_full(rederiveGeometry=rederiveGeometry) - def derive_geometry(self, *args, **kwargs): - raise Exception('[MITIM] This method is not implemented in the base class. Please use a derived class that implements it.') + def write_state(self, file=None): + print("\t- Writting input.gacode file") - def write_state(self, *args, **kwargs): - raise Exception('[MITIM] This method is not implemented in the base class. Please use a derived class that implements it.') + if file is None: + file = self.files[0] - # ------------------------------------------------------------------------------------- - # Method to write a scratch file - # ------------------------------------------------------------------------------------- + with open(file, "w") as f: + for line in self.header: + f.write(line) - @classmethod - def scratch(cls, profiles, label_header='', **kwargs_process): - instance = cls(None) + for i in self.profiles: + if "(" not in i: + f.write(f"# {i}\n") + else: + f.write(f"# {i.split('(')[0]} | {i.split('(')[-1].split(')')[0]}\n") - # Header - instance.header = f''' -# Created from scratch with MITIM version {__version__} -# {label_header} -# -''' - # Add data to profiles - instance.profiles = profiles + if i in self.titles_single: + listWrite = self.profiles[i] - instance.process(**kwargs_process) + if IOtools.isnum(listWrite[0]): + listWrite = [f"{i:.7e}".rjust(14) for i in listWrite] + f.write(f"{''.join(listWrite)}\n") + else: + f.write(f"{' '.join(listWrite)}\n") - return instance + else: + if len(self.profiles[i].shape) == 1: + for j, val in enumerate(self.profiles[i]): + pos = f"{j + 1}".rjust(3) + valt = f"{round(val,99):.7e}".rjust(15) + f.write(f"{pos}{valt}\n") + else: + for j, val in enumerate(self.profiles[i]): + pos = f"{j + 1}".rjust(3) + txt = "".join([f"{k:.7e}".rjust(15) for k in val]) + f.write(f"{pos}{txt}\n") + + print(f"\t\t~ File {IOtools.clipstr(file)} written") + + # Update file + self.files[0] = file + + # ************************************************************************************************************************************************ + # Derivation methods that children classes should implement + # ************************************************************************************************************************************************ - # ------------------------------------------------------------------------------------- + def derive_geometry(self, *args, **kwargs): + raise Exception('[MITIM] This method is not implemented in the base class. Please use a derived class that implements it.') + + def plot_geometry(self, *args, **kwargs): + print('[MITIM] Method plot_geometry() is not implemented in the base class. Please use a derived class that implements it.') + pass + + # ************************************************************************************************************************************************ + # Derivation methods + # ************************************************************************************************************************************************ def calculate_Er( self, @@ -202,7 +263,6 @@ def calculate_Er( if write_new_file is not None: self.write_state(file=write_new_file) - def readSpecies(self, maxSpecies=100): maxSpecies = int(self.profiles["nion"][0]) @@ -1663,6 +1723,144 @@ def introduceRotationProfile(self, Mach_LF=1.0, new_file=None): if new_file is not None: self.write_state(file=new_file) + def parabolizePlasma(self): + _, T = PLASMAtools.parabolicProfile( + Tbar=self.derived["Te_vol"], + nu=self.derived["Te_peaking"], + rho=self.profiles["rho(-)"], + Tedge=self.profiles["te(keV)"][-1], + ) + _, Ti = PLASMAtools.parabolicProfile( + Tbar=self.derived["Ti_vol"], + nu=self.derived["Ti_peaking"], + rho=self.profiles["rho(-)"], + Tedge=self.profiles["ti(keV)"][-1, 0], + ) + _, n = PLASMAtools.parabolicProfile( + Tbar=self.derived["ne_vol20"] * 1e1, + nu=self.derived["ne_peaking"], + rho=self.profiles["rho(-)"], + Tedge=self.profiles["ne(10^19/m^3)"][-1], + ) + + self.profiles["te(keV)"] = T + + self.profiles["ti(keV)"][:, 0] = Ti + self.makeAllThermalIonsHaveSameTemp(refIon=0) + + factor_n = n / self.profiles["ne(10^19/m^3)"] + self.profiles["ne(10^19/m^3)"] = n + self.scaleAllThermalDensities(scaleFactor=factor_n) + + self.derive_quantities() + + def changeRFpower(self, PrfMW=25.0): + """ + keeps same partition + """ + print(f"- Changing the RF power from {self.derived['qRF_MW'][-1]:.1f} MW to {PrfMW:.1f} MW",typeMsg="i",) + + if self.derived["qRF_MW"][-1] == 0.0: + raise Exception("No RF power in the input.gacode, cannot modify the RF power") + + for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: + self.profiles[i] = self.profiles[i] * PrfMW / self.derived["qRF_MW"][-1] + + self.derive_quantities() + + def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0.2): + + ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) + + self.profiles["te(keV)"] = self.profiles["te(keV)"] * TkeV / self.profiles["te(keV)"][ix] + + print(f"- Producing {typeEdge} boundary condition @ rho = {rho}, T = {TkeV} keV",typeMsg="i",) + + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "therm": + self.profiles["ti(keV)"][:, sp] = self.profiles["ti(keV)"][:, sp] * TkeV / self.profiles["ti(keV)"][ix, sp] + + if typeEdge == "linear": + self.profiles["te(keV)"][ix:] = np.linspace(TkeV, Tesep, len(self.profiles["rho(-)"][ix:])) + + for sp in range(len(self.Species)): + if self.Species[sp]["S"] == "therm": + self.profiles["ti(keV)"][ix:, sp] = np.linspace(TkeV, Tisep, len(self.profiles["rho(-)"][ix:])) + + elif typeEdge == "same": + pass + else: + raise Exception("no edge") + + def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5, isn20_edge=True): + ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) + + # Determine the factor to scale the density (either average or at rho) + if not isn20_edge: + print(f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") + factor = n20 / self.derived["ne_vol20"] + else: + print(f"- Changing the density at rho={rho} from {self.profiles['ne(10^19/m^3)'][ix]*1E-1:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") + factor = n20 / (self.profiles["ne(10^19/m^3)"][ix]*1E-1) + # ------------------------------------------------------------------ + + # Scale the density profiles + for i in ["ne(10^19/m^3)", "ni(10^19/m^3)"]: + self.profiles[i] = self.profiles[i] * factor + + # Apply the edge condition + if typeEdge == "linear": + factor_x = np.linspace(self.profiles["ne(10^19/m^3)"][ix],nedge20 * 1e1,len(self.profiles["rho(-)"][ix:]),)/ self.profiles["ne(10^19/m^3)"][ix:] + + self.profiles["ne(10^19/m^3)"][ix:] = self.profiles["ne(10^19/m^3)"][ix:] * factor_x + + for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): + self.profiles["ni(10^19/m^3)"][ix:, i] = self.profiles["ni(10^19/m^3)"][ix:, i] * factor_x + + elif typeEdge == "same": + pass + else: + raise Exception("no edge") + + def addSawtoothEffectOnOhmic(self, PohTot, mixRadius=None, plotYN=False): + """ + This will implement a flat profile inside the mixRadius to reduce the ohmic power by certain amount + """ + + if mixRadius is None: + mixRadius = self.profiles["rho(-)"][np.where(self.profiles["q(-)"] > 1)][0] + + print(f"\t- Original Ohmic power: {self.derived['qOhm_MW'][-1]:.2f}MW") + Ohmic_old = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) + + dvol = self.derived["volp_geo"] * np.append( + [0], np.diff(self.profiles["rmin(m)"]) + ) + + print( + f"\t- Will implement sawtooth ohmic power correction inside rho={mixRadius}" + ) + Psaw = CDFtools.profilePower( + self.profiles["rho(-)"], + dvol, + PohTot - self.derived["qOhm_MW"][-1], + mixRadius, + ) + self.profiles["qohme(MW/m^3)"] += Psaw + self.derive_quantities() + + print(f"\t- New Ohmic power: {self.derived['qOhm_MW'][-1]:.2f}MW") + Ohmic_new = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) + + if plotYN: + fig, ax = plt.subplots() + ax.plot(self.profiles["rho(-)"], Ohmic_old, "r", lw=2) + ax.plot(self.profiles["rho(-)"], Ohmic_new, "g", lw=2) + plt.show() + + # ************************************************************************************************************************************************ + # Plotting methods for the state class, which is used to plot the profiles, powers, geometry, gradients, flows, and other quantities. + # ************************************************************************************************************************************************ def plot( self, @@ -1974,146 +2172,9 @@ def csv(self, file="input.gacode.xlsx"): IOtools.writeExcel_fromDict(dictExcel, file, fromRow=1) - def parabolizePlasma(self): - _, T = PLASMAtools.parabolicProfile( - Tbar=self.derived["Te_vol"], - nu=self.derived["Te_peaking"], - rho=self.profiles["rho(-)"], - Tedge=self.profiles["te(keV)"][-1], - ) - _, Ti = PLASMAtools.parabolicProfile( - Tbar=self.derived["Ti_vol"], - nu=self.derived["Ti_peaking"], - rho=self.profiles["rho(-)"], - Tedge=self.profiles["ti(keV)"][-1, 0], - ) - _, n = PLASMAtools.parabolicProfile( - Tbar=self.derived["ne_vol20"] * 1e1, - nu=self.derived["ne_peaking"], - rho=self.profiles["rho(-)"], - Tedge=self.profiles["ne(10^19/m^3)"][-1], - ) - - self.profiles["te(keV)"] = T - - self.profiles["ti(keV)"][:, 0] = Ti - self.makeAllThermalIonsHaveSameTemp(refIon=0) - - factor_n = n / self.profiles["ne(10^19/m^3)"] - self.profiles["ne(10^19/m^3)"] = n - self.scaleAllThermalDensities(scaleFactor=factor_n) - - self.derive_quantities() - - - def changeRFpower(self, PrfMW=25.0): - """ - keeps same partition - """ - print(f"- Changing the RF power from {self.derived['qRF_MW'][-1]:.1f} MW to {PrfMW:.1f} MW",typeMsg="i",) - - if self.derived["qRF_MW"][-1] == 0.0: - raise Exception("No RF power in the input.gacode, cannot modify the RF power") - - for i in ["qrfe(MW/m^3)", "qrfi(MW/m^3)"]: - self.profiles[i] = self.profiles[i] * PrfMW / self.derived["qRF_MW"][-1] - - self.derive_quantities() - - def imposeBCtemps(self, TkeV=0.5, rho=0.9, typeEdge="linear", Tesep=0.1, Tisep=0.2): - - ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) - - self.profiles["te(keV)"] = self.profiles["te(keV)"] * TkeV / self.profiles["te(keV)"][ix] - - print(f"- Producing {typeEdge} boundary condition @ rho = {rho}, T = {TkeV} keV",typeMsg="i",) - - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm": - self.profiles["ti(keV)"][:, sp] = self.profiles["ti(keV)"][:, sp] * TkeV / self.profiles["ti(keV)"][ix, sp] - - if typeEdge == "linear": - self.profiles["te(keV)"][ix:] = np.linspace(TkeV, Tesep, len(self.profiles["rho(-)"][ix:])) - - for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm": - self.profiles["ti(keV)"][ix:, sp] = np.linspace(TkeV, Tisep, len(self.profiles["rho(-)"][ix:])) - - elif typeEdge == "same": - pass - else: - raise Exception("no edge") - - - def imposeBCdens(self, n20=2.0, rho=0.9, typeEdge="linear", nedge20=0.5, isn20_edge=True): - ix = np.argmin(np.abs(rho - self.profiles["rho(-)"])) - - # Determine the factor to scale the density (either average or at rho) - if not isn20_edge: - print(f"- Changing the initial average density from {self.derived['ne_vol20']:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") - factor = n20 / self.derived["ne_vol20"] - else: - print(f"- Changing the density at rho={rho} from {self.profiles['ne(10^19/m^3)'][ix]*1E-1:.1f} 1E20/m3 to {n20:.1f} 1E20/m3",typeMsg="i") - factor = n20 / (self.profiles["ne(10^19/m^3)"][ix]*1E-1) - # ------------------------------------------------------------------ - - # Scale the density profiles - for i in ["ne(10^19/m^3)", "ni(10^19/m^3)"]: - self.profiles[i] = self.profiles[i] * factor - - # Apply the edge condition - if typeEdge == "linear": - factor_x = np.linspace(self.profiles["ne(10^19/m^3)"][ix],nedge20 * 1e1,len(self.profiles["rho(-)"][ix:]),)/ self.profiles["ne(10^19/m^3)"][ix:] - - self.profiles["ne(10^19/m^3)"][ix:] = self.profiles["ne(10^19/m^3)"][ix:] * factor_x - - for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): - self.profiles["ni(10^19/m^3)"][ix:, i] = self.profiles["ni(10^19/m^3)"][ix:, i] * factor_x - - elif typeEdge == "same": - pass - else: - raise Exception("no edge") - - def addSawtoothEffectOnOhmic(self, PohTot, mixRadius=None, plotYN=False): - """ - This will implement a flat profile inside the mixRadius to reduce the ohmic power by certain amount - """ - - if mixRadius is None: - mixRadius = self.profiles["rho(-)"][np.where(self.profiles["q(-)"] > 1)][0] - - print(f"\t- Original Ohmic power: {self.derived['qOhm_MW'][-1]:.2f}MW") - Ohmic_old = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) - - dvol = self.derived["volp_geo"] * np.append( - [0], np.diff(self.profiles["rmin(m)"]) - ) - - print( - f"\t- Will implement sawtooth ohmic power correction inside rho={mixRadius}" - ) - Psaw = CDFtools.profilePower( - self.profiles["rho(-)"], - dvol, - PohTot - self.derived["qOhm_MW"][-1], - mixRadius, - ) - self.profiles["qohme(MW/m^3)"] += Psaw - self.derive_quantities() - - print(f"\t- New Ohmic power: {self.derived['qOhm_MW'][-1]:.2f}MW") - Ohmic_new = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) - - if plotYN: - fig, ax = plt.subplots() - ax.plot(self.profiles["rho(-)"], Ohmic_old, "r", lw=2) - ax.plot(self.profiles["rho(-)"], Ohmic_new, "g", lw=2) - plt.show() - - # --------------------------------------------------------------------------------------------------------------------------------------- + # ************************************************************************************************************************************************ # Code conversions - # --------------------------------------------------------------------------------------------------------------------------------------- + # ************************************************************************************************************************************************ def to_tglf(self, rhos=[0.5], TGLFsettings=1): @@ -2270,6 +2331,7 @@ def to_eped(self, ped_rho = 0.95): return eped_evaluation + class DataTable: def __init__(self, variables=None): @@ -2344,37 +2406,6 @@ def export_to_csv(self, filename, title=None): for row in self.data: writer.writerow(row) - -def readTGYRO_profile_extra(file, varLabel="B_unit (T)"): - with open(file) as f: - aux = f.readlines() - - lenn = int(aux[36].split()[-1]) - - i = 38 - allVec = [] - while i < len(aux): - vec = np.array([float(j) for j in aux[i : i + lenn]]) - i += lenn - allVec.append(vec) - allVec = np.array(allVec) - - dictL = OrderedDict() - for line in aux[2:35]: - lab = line.split("(:)")[-1].split("\n")[0] - try: - dictL[lab] = int(line.split()[1]) - except: - dictL[lab] = [int(j) for j in line.split()[1].split("-")] - - for i in dictL: - if i.strip(" ") == varLabel: - val = allVec[dictL[i] - 1] - break - - return val - - def aLT(r, p): return ( r[-1] @@ -2495,7 +2526,6 @@ def gradientsMerger(p0, p_true, roa=0.46, blending=0.1): return p - def impurity_location(profiles, impurity_of_interest): position_of_impurity = None diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index 137bcdba..9d13f8ac 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -13,9 +13,9 @@ class vmec_state(MITIMstate.mitim_state): Class to read and manipulate VMEC files ''' - # ------------------------------------------------------------------ + # ************************************************************************************************************************************************ # Reading and interpreting - # ------------------------------------------------------------------ + # ************************************************************************************************************************************************ def __init__( self, @@ -28,6 +28,12 @@ def __init__( # Initialize the base class and tell it the type of file super().__init__(type_file='vmec') + + self.header = ['# VMEC state file generated by MITIMtools\n'] + self.titles_singleNum = ["nexp", "nion", "shot", "name", "type", "time"] + self.titles_singleArr = ["masse","mass","ze","z","torfluxa(Wb/radian)","rcentr(m)","bcentr(T)"] #,"current(MA)"] + self.titles_single = self.titles_singleNum + self.titles_singleArr + # Read the input file and store the raw data self.files = [file_vmec, file_profs] if self.files is not None: @@ -36,22 +42,36 @@ def __init__( # Derive (Depending on resolution, derived can be expensive, so I may not do it every time) self.derive_quantities(mi_ref=mi_ref, derive_quantities=derive_quantities) - @IOtools.hook_method(after=MITIMstate.ensure_variables_existence) def _read_vmec(self): - - self.profiles = {} - + # Read VMEC file print("\t- Reading VMEC file") self.wout = vmecpp.VmecWOut.from_wout_file(Path(self.files[0])) + + # Initialize profiles dictionary + self.profiles = {} - # Produce variables - self.profiles["rho(-)"] = np.linspace(0, 1, self.wout.ns)**0.5 - self.profiles["rmin(m)"] = self.profiles["rho(-)"] + self.profiles['nion'] = np.array([1]) + self.profiles['name'] = np.array(['D']) + self.profiles['type'] = np.array(['[therm]']) + self.profiles['mass'] = np.array([2.0]) + self.profiles['z'] = np.array([1.0]) + self.profiles['rcentr(m)'] = np.array([self.wout.Rmajor_p]) + self.profiles['bcentr(T)'] = np.array([self.wout.rbtor/self.wout.Rmajor_p]) + self.profiles["current(MA)"] = np.array([0.0]) + self.profiles["torfluxa(Wb/radian)"] = [self.wout.phipf[-1]] + # Produce variables + self.profiles["rho(-)"] = (self.wout.phi/self.wout.phi[-1])**0.5 #np.linspace(0, 1, self.wout.ns)**0.5 + self.profiles["presf"] = self.wout.presf + + self.profiles["rmin(m)"] = self.profiles["rho(-)"] + + self.profiles["q(-)"] = self.wout.q_factor + # Read Profiles data = self._read_profiles(x_coord=self.profiles["rho(-)"]) @@ -60,20 +80,14 @@ def _read_vmec(self): self.profiles['ti(keV)'] = np.atleast_2d(data['Ti']).T self.profiles['ni(10^19/m^3)'] = np.atleast_2d(data['ni']).T - self.profiles['qbeami(MW/m^3)'] = data['Qi'] * 1e-6 # Convert from W/m^3 to MW/m^3 - self.profiles['qrfe(MW/m^3)'] = data['Qe'] * 1e-6 # Convert from W/m^3 to MW/m^3 + self.profiles['qbeami(MW/m^3)'] = data['Qi'] * 1e-6 # Convert from W/m^3 to MW/m^3 + self.profiles['qrfe(MW/m^3)'] = data['Qe'] * 1e-6 # Convert from W/m^3 to MW/m^3 self.profiles['qpar_beam(1/m^3/s)'] = data['S'] - self.profiles['nion'] = np.array([1]) - self.profiles['name'] = np.array(['D']) - self.profiles['type'] = np.array(['[therm]']) - self.profiles['mass'] = np.array([2.0]) - self.profiles['z'] = np.array([1.0]) - - self.profiles['rcentr(m)'] = np.array([self.wout.Rmajor_p]) - self.profiles['bcentr(T)'] = np.array([self.wout.rbtor/self.wout.Rmajor_p]) - self.profiles["current(MA)"] = np.array([0.0]) - + # ************************************************************************************************************************************************ + # Derivation (different from MITIMstate) + # ************************************************************************************************************************************************ + def derive_geometry(self, **kwargs): rho = np.linspace(0, 1, self.wout.ns) @@ -87,11 +101,11 @@ def derive_geometry(self, **kwargs): * 2 * np.sqrt(half_grid_rho) ) - d_volume_d_rho[0] = 0.0 # Set the first element to zero to avoid division by zero - + #self.derived["B_unit"] = self.profiles["torfluxa(Wb/radian)"] / (np.pi * self.wout.Aminor_p**2) - self.derived["volp_geo"] = d_volume_d_rho + self.derived["volp_geo"] = self.wout.vp #d_volume_d_rho + self.derived["volp_geo"][0] = 0.0 self.derived["kappa_a"] = 0.0 self.derived["kappa95"] = 0.0 @@ -105,10 +119,6 @@ def derive_geometry(self, **kwargs): self.derived["bp2_geo"] = np.zeros(self.profiles["rho(-)"].shape) self.derived["bt2_geo"] = np.zeros(self.profiles["rho(-)"].shape) - - def write_state(self, file=None, **kwargs): - pass - def plot_geometry(self, axs, color="b", legYN=True, extralab="", lw=1, fs=6): pass @@ -223,7 +233,6 @@ def _read_profiles(self, x_coord=None, debug = False): return uniform_data - def plot_plasma_boundary(self, ax=None): # The output object contains the Fourier coefficients of the geometry in R and Z @@ -364,3 +373,4 @@ def plot_profiles(data): plt.tight_layout() plt.show() +'' \ No newline at end of file From c1c49ab9a54aaad4f574a3d5cdf565a7544718a8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Jul 2025 11:12:17 -0400 Subject: [PATCH 101/385] Population of vmec variables and graphical plotting of geometry --- src/mitim_tools/gacode_tools/PROFILEStools.py | 10 +- .../plasmastate_tools/MITIMstate.py | 71 +++++++------- .../plasmastate_tools/utils/VMECtools.py | 93 +++++++++++++++++-- .../plasmastate_tools/utils/state_plotting.py | 3 +- 4 files changed, 132 insertions(+), 45 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index c2caa3f1..e7a9a742 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -163,9 +163,15 @@ def _read_profiles(self): def derive_quantities(self, **kwargs): + if "derived" not in self.__dict__: + self.derived = {} + self._produce_shape_lists() - super().derive_quantities(**kwargs) + # Define the minor radius used in all calculations (could be the half-width of the midplance intersect, or an effective minor radius) + self.derived["r"] = self.profiles["rmin(m)"] + + super().derive_quantities_base(**kwargs) def _produce_shape_lists(self): self.shape_cos = [ @@ -249,7 +255,7 @@ def derive_geometry(self, n_theta_geo=1001, **kwargs): def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): - [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,ax13c] = axs3 + [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,_,ax13c] = axs3 rho = self.profiles["rho(-)"] lines = GRAPHICStools.listLS() diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index e4f5ac06..00db28be 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -87,6 +87,8 @@ def ensure_variables_existence(self): The mitim_state class is the base class for manipulating plasma states in MITIM. Any class that inherits from this class should implement the methods: + - derive_quantities: to derive quantities from the plasma state (must at least define "r" and call the derive_quantities_base method). + - derive_geometry: to derive the geometry of the plasma state. - write_state: to write the plasma state to a file. @@ -124,7 +126,7 @@ def scratch(cls, profiles, label_header='', **kwargs_process): return instance @IOtools.hook_method(before=ensure_variables_existence) - def derive_quantities(self, mi_ref=None, derive_quantities=True, rederiveGeometry=True): + def derive_quantities_base(self, mi_ref=None, derive_quantities=True, rederiveGeometry=True): # ------------------------------------- self.readSpecies() @@ -144,14 +146,14 @@ def derive_quantities(self, mi_ref=None, derive_quantities=True, rederiveGeometr print(f"\t* Reference mass ({self.derived['mi_ref']}) from first ion",typeMsg="i") # Useful to have gradients in the basic ---------------------------------------------------------- - self.derived["aLTe"] = aLT(self.profiles["rmin(m)"], self.profiles["te(keV)"]) - self.derived["aLne"] = aLT(self.profiles["rmin(m)"], self.profiles["ne(10^19/m^3)"]) + self.derived["aLTe"] = aLT(self.derived["r"], self.profiles["te(keV)"]) + self.derived["aLne"] = aLT(self.derived["r"], self.profiles["ne(10^19/m^3)"]) self.derived["aLTi"] = self.profiles["ti(keV)"] * 0.0 self.derived["aLni"] = [] for i in range(self.profiles["ti(keV)"].shape[1]): - self.derived["aLTi"][:, i] = aLT(self.profiles["rmin(m)"], self.profiles["ti(keV)"][:, i]) - self.derived["aLni"].append(aLT(self.profiles["rmin(m)"], self.profiles["ni(10^19/m^3)"][:, i])) + self.derived["aLTi"][:, i] = aLT(self.derived["r"], self.profiles["ti(keV)"][:, i]) + self.derived["aLni"].append(aLT(self.derived["r"], self.profiles["ni(10^19/m^3)"][:, i])) self.derived["aLni"] = np.transpose(np.array(self.derived["aLni"])) # ------------------------------------------------------------------------------------------------ @@ -204,6 +206,9 @@ def write_state(self, file=None): # Derivation methods that children classes should implement # ************************************************************************************************************************************************ + def derive_quantities(self, *args, **kwargs): + raise Exception('[MITIM] derive_quantities method is not implemented in the base class (to define "r"). Please use a derived class that implements it.') + def derive_geometry(self, *args, **kwargs): raise Exception('[MITIM] This method is not implemented in the base class. Please use a derived class that implements it.') @@ -310,17 +315,17 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): # --------- MAIN (useful for STATEtools) # --------------------------------------------------------------------------------------------------------------------- - self.derived["a"] = self.profiles["rmin(m)"][-1] + self.derived["a"] = self.derived["r"][-1] # self.derived['epsX'] = self.profiles['rmaj(m)'] / self.profiles['rmin(m)'] # self.derived['eps'] = self.derived['epsX'][-1] - self.derived["eps"] = self.profiles["rmin(m)"][-1] / self.profiles["rmaj(m)"][-1] + self.derived["eps"] = self.derived["r"][-1] / self.profiles["rmaj(m)"][-1] - self.derived["roa"] = self.profiles["rmin(m)"] / self.derived["a"] + self.derived["roa"] = self.derived["r"] / self.derived["a"] self.derived["Rmajoa"] = self.profiles["rmaj(m)"] / self.derived["a"] self.derived["Zmagoa"] = self.profiles["zmag(m)"] / self.derived["a"] self.derived["torflux"] = float(self.profiles["torfluxa(Wb/radian)"][0])* 2* np.pi* self.profiles["rho(-)"] ** 2 # Wb - self.derived["B_unit"] = PLASMAtools.Bunit(self.derived["torflux"], self.profiles["rmin(m)"]) + self.derived["B_unit"] = PLASMAtools.Bunit(self.derived["torflux"], self.derived["r"]) self.derived["psi_pol_n"] = ( self.profiles["polflux(Wb/radian)"] - self.profiles["polflux(Wb/radian)"][0] @@ -371,7 +376,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.profiles["ne(10^19/m^3)"] * 1e-1, self.derived["mi_ref"], np.abs(self.derived["B_unit"]), - self.profiles["rmin(m)"][-1], + self.derived["r"][-1], ) """ @@ -437,7 +442,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): profiles_gen puts any missing power into the CX: qioni, qione """ - r = self.profiles["rmin(m)"] + r = self.derived["r"] volp = self.derived["volp_geo"] self.derived["qe_MW"] = CALCtools.volume_integration(self.derived["qe"], r, volp) @@ -467,7 +472,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["mt_Jm2"] = self.derived["mt_Jmiller"] / (volp) # Extras for plotting in TGYRO for comparison - P = np.zeros(len(self.profiles["rmin(m)"])) + P = np.zeros(len(self.derived["r"])) if "qsync(MW/m^3)" in self.profiles: P += self.profiles["qsync(MW/m^3)"] if "qbrem(MW/m^3)" in self.profiles: @@ -603,30 +608,30 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): """ Derivatives """ - self.derived["aLTe"] = aLT(self.profiles["rmin(m)"], self.profiles["te(keV)"]) + self.derived["aLTe"] = aLT(self.derived["r"], self.profiles["te(keV)"]) self.derived["aLTi"] = self.profiles["ti(keV)"] * 0.0 for i in range(self.profiles["ti(keV)"].shape[1]): self.derived["aLTi"][:, i] = aLT( - self.profiles["rmin(m)"], self.profiles["ti(keV)"][:, i] + self.derived["r"], self.profiles["ti(keV)"][:, i] ) self.derived["aLne"] = aLT( - self.profiles["rmin(m)"], self.profiles["ne(10^19/m^3)"] + self.derived["r"], self.profiles["ne(10^19/m^3)"] ) self.derived["aLni"] = [] for i in range(self.profiles["ni(10^19/m^3)"].shape[1]): self.derived["aLni"].append( - aLT(self.profiles["rmin(m)"], self.profiles["ni(10^19/m^3)"][:, i]) + aLT(self.derived["r"], self.profiles["ni(10^19/m^3)"][:, i]) ) self.derived["aLni"] = np.transpose(np.array(self.derived["aLni"])) if "w0(rad/s)" not in self.profiles: self.profiles["w0(rad/s)"] = self.profiles["rho(-)"] * 0.0 - self.derived["aLw0"] = aLT(self.profiles["rmin(m)"], self.profiles["w0(rad/s)"]) + self.derived["aLw0"] = aLT(self.derived["r"], self.profiles["w0(rad/s)"]) self.derived["dw0dr"] = -grad( - self.profiles["rmin(m)"], self.profiles["w0(rad/s)"] + self.derived["r"], self.profiles["w0(rad/s)"] ) - self.derived["dqdr"] = grad(self.profiles["rmin(m)"], self.profiles["q(-)"]) + self.derived["dqdr"] = grad(self.derived["r"], self.profiles["q(-)"]) """ Other, performance @@ -919,7 +924,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): nG = PLASMAtools.Greenwald_density( np.abs(float(self.profiles["current(MA)"][-1])), - float(self.profiles["rmin(m)"][-1]), + float(self.derived["r"][-1]), ) self.derived["fG"] = self.derived["ne_vol20"] / nG self.derived["fG_x"] = self.profiles["ne(10^19/m^3)"]* 0.1 / nG @@ -1008,7 +1013,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): def tglf_plasma(self): def deriv_gacode(y): - return grad(self.profiles["rmin(m)"],y).cpu().numpy() + return grad(self.derived["r"],y).cpu().numpy() self.derived["tite_all"] = self.profiles["ti(keV)"] / self.profiles["te(keV)"][:, np.newaxis] @@ -1029,18 +1034,18 @@ def deriv_gacode(y): self.derived["mi_ref"], self.derived["B_unit"]) - self.derived['pprime'] = 1E-7 * self.profiles["q(-)"]*self.derived['a']**2/self.profiles["rmin(m)"]/self.derived["B_unit"]**2*deriv_gacode(self.profiles["ptot(Pa)"]) + self.derived['pprime'] = 1E-7 * self.profiles["q(-)"]*self.derived['a']**2/self.derived["r"]/self.derived["B_unit"]**2*deriv_gacode(self.profiles["ptot(Pa)"]) self.derived['pprime'][0] = 0.0 - self.derived['drmin/dr'] = deriv_gacode(self.profiles["rmin(m)"]) + self.derived['drmin/dr'] = deriv_gacode(self.derived["r"]) self.derived['dRmaj/dr'] = deriv_gacode(self.profiles["rmaj(m)"]) self.derived['dZmaj/dr'] = deriv_gacode(self.profiles["zmag(m)"]) - self.derived['s_kappa'] = self.profiles["rmin(m)"] / self.profiles["kappa(-)"] * deriv_gacode(self.profiles["kappa(-)"]) - self.derived['s_delta'] = self.profiles["rmin(m)"] * deriv_gacode(self.profiles["delta(-)"]) - self.derived['s_zeta'] = self.profiles["rmin(m)"] * deriv_gacode(self.profiles["zeta(-)"]) + self.derived['s_kappa'] = self.derived["r"] / self.profiles["kappa(-)"] * deriv_gacode(self.profiles["kappa(-)"]) + self.derived['s_delta'] = self.derived["r"] * deriv_gacode(self.profiles["delta(-)"]) + self.derived['s_zeta'] = self.derived["r"] * deriv_gacode(self.profiles["zeta(-)"]) - s = self.profiles["rmin(m)"] / self.profiles["q(-)"]*deriv_gacode(self.profiles["q(-)"]) + s = self.derived["r"] / self.profiles["q(-)"]*deriv_gacode(self.profiles["q(-)"]) self.derived['s_q'] = np.concatenate([np.array([0.0]),(self.profiles["q(-)"][1:] / self.derived['roa'][1:])**2 * s[1:]]) # infinite in first location ''' @@ -1055,7 +1060,7 @@ def deriv_gacode(y): w0p = deriv_gacode(self.profiles["w0(rad/s)"]) gamma_p0 = -self.profiles["rmaj(m)"]*w0p - gamma_eb0 = -deriv_gacode(self.profiles["w0(rad/s)"]) * self.profiles["rmin(m)"]/self.profiles["q(-)"] + gamma_eb0 = -deriv_gacode(self.profiles["w0(rad/s)"]) * self.derived["r"]/self.profiles["q(-)"] self.derived['vexb_shear'] = gamma_eb0 * self.derived["a"]/self.derived['c_s'] self.derived['vpar_shear'] = gamma_p0 * self.derived["a"]/self.derived['c_s'] @@ -1101,7 +1106,7 @@ def deriveContentByVolumes(self, rhos=[0.5], impurityPosition=3): for j in range(self.profiles["te(keV)"].shape[0] - min_number_points): i = j + min_number_points We_x[i], Wi_x[i], Ne_x[i], _ = PLASMAtools.calculateContent( - np.expand_dims(self.profiles["rmin(m)"][:i], 0), + np.expand_dims(self.derived["r"][:i], 0), np.expand_dims(self.profiles["te(keV)"][:i], 0), np.expand_dims(np.transpose(self.profiles["ti(keV)"][:i]), 0), np.expand_dims(self.profiles["ne(10^19/m^3)"][:i] * 0.1, 0), @@ -1112,7 +1117,7 @@ def deriveContentByVolumes(self, rhos=[0.5], impurityPosition=3): ) _, _, Ni_x[i], _ = PLASMAtools.calculateContent( - np.expand_dims(self.profiles["rmin(m)"][:i], 0), + np.expand_dims(self.derived["r"][:i], 0), np.expand_dims(self.profiles["te(keV)"][:i], 0), np.expand_dims(np.transpose(self.profiles["ti(keV)"][:i]), 0), np.expand_dims( @@ -1363,7 +1368,7 @@ def lumpSpecies( fZ2 += self.Species[i - 1]["Z"] ** 2 * self.derived["fi"][:, i - 1] Zr = fZ2 / fZ1 - Zr_vol = CALCtools.volume_integration(Zr, self.profiles["rmin(m)"], self.derived["volp_geo"])[-1] / self.derived["volume"] + Zr_vol = CALCtools.volume_integration(Zr, self.derived["r"], self.derived["volp_geo"])[-1] / self.derived["volume"] print(f'\t\t\t* Original plasma had Zeff_vol={self.derived["Zeff_vol"]:.2f}, QN error={self.derived["QN_Error"]:.4f}') @@ -1834,7 +1839,7 @@ def addSawtoothEffectOnOhmic(self, PohTot, mixRadius=None, plotYN=False): Ohmic_old = copy.deepcopy(self.profiles["qohme(MW/m^3)"]) dvol = self.derived["volp_geo"] * np.append( - [0], np.diff(self.profiles["rmin(m)"]) + [0], np.diff(self.derived["r"]) ) print( @@ -1945,7 +1950,7 @@ def plotPeaking( print(f"\t- nu_eff = {nu_effCGYRO}, ne_peaking = {ne_peaking}") # Extra - r = self.profiles["rmin(m)"] + r = self.derived["r"] volp = self.derived["volp_geo"] ix = np.argmin(np.abs(self.profiles["rho(-)"] - 0.9)) diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index 9d13f8ac..2a5387b5 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -1,9 +1,10 @@ import vmecpp import numpy as np +from collections import OrderedDict import matplotlib.pyplot as plt from pathlib import Path from scipy.interpolate import interp1d -from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools import GRAPHICStools, IOtools from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -50,7 +51,7 @@ def _read_vmec(self): self.wout = vmecpp.VmecWOut.from_wout_file(Path(self.files[0])) # Initialize profiles dictionary - self.profiles = {} + self.profiles = OrderedDict() self.profiles['nion'] = np.array([1]) self.profiles['name'] = np.array(['D']) @@ -63,14 +64,12 @@ def _read_vmec(self): self.profiles["current(MA)"] = np.array([0.0]) self.profiles["torfluxa(Wb/radian)"] = [self.wout.phipf[-1]] - + # Produce variables self.profiles["rho(-)"] = (self.wout.phi/self.wout.phi[-1])**0.5 #np.linspace(0, 1, self.wout.ns)**0.5 self.profiles["presf"] = self.wout.presf - - self.profiles["rmin(m)"] = self.profiles["rho(-)"] - self.profiles["q(-)"] = self.wout.q_factor + self.profiles["polflux(Wb/radian)"] = self.wout.chi # Read Profiles data = self._read_profiles(x_coord=self.profiles["rho(-)"]) @@ -88,6 +87,16 @@ def _read_vmec(self): # Derivation (different from MITIMstate) # ************************************************************************************************************************************************ + def derive_quantities(self, **kwargs): + + if "derived" not in self.__dict__: + self.derived = {} + + # Define the minor radius used in all calculations (could be the half-width of the midplance intersect, or an effective minor radius) + self.derived["r"] = self.profiles["rho(-)"] + + super().derive_quantities_base(**kwargs) + def derive_geometry(self, **kwargs): rho = np.linspace(0, 1, self.wout.ns) @@ -121,8 +130,33 @@ def derive_geometry(self, **kwargs): def plot_geometry(self, axs, color="b", legYN=True, extralab="", lw=1, fs=6): - pass - #self.plot_plasma_boundary() + [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,axs_3d,axs_2d] = axs + + self.plot_plasma_boundary(ax=axs_3d) + + self.plot_state_flux_surfaces(ax=axs_2d, c=color) + + + def plot_state_flux_surfaces(self, ax=None, c='b'): + + rhos_plot = np.linspace(0.0, 1.0, 10) + phis_plot = [0.0, np.pi/2, np.pi, 3*np.pi/2] #np.linspace(0.0, 2.0 * np.pi, 5) + + ls = GRAPHICStools.listLS() + + for phi_cut, lsi in zip(phis_plot, ls): + + for i in range(len(rhos_plot)): + self.plot_flux_surface(ax = ax, phi_cut=phi_cut, rho=rhos_plot[i], c=c, lw = 0.5, ls = lsi) + self.plot_flux_surface(ax = ax, phi_cut=phi_cut, rho=1.0, c=c, lw = 2, ls = lsi, label = f"$\\phi={phi_cut*180/np.pi:.1f} deg$") + + ax.set_aspect('equal') + ax.set_xlabel('R [m]') + ax.set_ylabel('Z [m]') + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.addLegendApart(ax, ratio=0.9, size=6) + + ax.set_title(f'Poloidal Cross-section') def _read_profiles(self, x_coord=None, debug = False): @@ -290,8 +324,49 @@ def plot_plasma_boundary(self, ax=None): # Set an equal aspect ratio ax.set_aspect("equal") + + ax.set_title(f'Plasma Boundary') + + def plot_flux_surface(self, ax=None, phi_cut=0.0, rho=1.0, c='b', lw=1, ls='-', label = ''): + """ + Plot poloidal cross-section of the torus at a specified toroidal angle. + + Parameters: + ----------- + ax : matplotlib axes object, optional + Axes to plot on. If None, creates new figure and axes. + phi_cut : float, optional + Toroidal angle for the cross-section in radians. Default is 0.0. + rho : float, optional + Normalized flux surface coordinate (0 to 1). Default is 1.0 (boundary). + """ + + ns = self.wout.ns + xm = self.wout.xm + xn = self.wout.xn + rmnc = self.wout.rmnc + zmns = self.wout.zmns + + # Find closest flux surface index for given rho + rho_grid = self.profiles["rho(-)"] + j = np.argmin(np.abs(rho_grid - rho)) + + num_theta = 201 + grid_theta = np.linspace(0.0, 2.0 * np.pi, num_theta, endpoint=True) + + R = np.zeros(num_theta) + Z = np.zeros(num_theta) + + for idx_theta, theta in enumerate(grid_theta): + kernel = xm * theta - xn * phi_cut + R[idx_theta] = np.dot(rmnc[:, j], np.cos(kernel)) + Z[idx_theta] = np.dot(zmns[:, j], np.sin(kernel)) + + if ax is None: + plt.ion() + fig, ax = plt.subplots() - plt.show() + ax.plot(R, Z, ls=ls, color = c, linewidth=lw, label=label) def plot_profiles(data): """ diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py index 0555b8a9..0519bdc8 100644 --- a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -58,7 +58,8 @@ def add_axes(figs): fig3.add_subplot(grid[0, 2]), fig3.add_subplot(grid[1, 2]), fig3.add_subplot(grid[2, 2]), - fig3.add_subplot(grid[:, 3]), + fig3.add_subplot(grid[0, 3], projection="3d"), + fig3.add_subplot(grid[1:, 3]), ] grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) From 1694e53ffafd5c29ae68f62661fa6158f5ee937d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Jul 2025 19:58:39 -0400 Subject: [PATCH 102/385] Successful run of PORTALS with TGLF without TGYRO --- .../physics_models/transport_tglf.py | 236 ++++++++++++++++++ src/mitim_tools/gacode_tools/TGLFtools.py | 39 +-- src/mitim_tools/gacode_tools/TGYROtools.py | 3 +- .../plasmastate_tools/MITIMstate.py | 4 +- .../plasmastate_tools/utils/VMECtools.py | 2 +- 5 files changed, 265 insertions(+), 19 deletions(-) create mode 100644 src/mitim_modules/powertorch/physics_models/transport_tglf.py diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py new file mode 100644 index 00000000..5de5dc7e --- /dev/null +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -0,0 +1,236 @@ +import torch +import numpy as np +from mitim_tools.misc_tools import IOtools +from mitim_tools.gacode_tools import TGLFtools +from mitim_modules.powertorch.utils import TRANSPORTtools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +class tglf_model(TRANSPORTtools.power_transport): + def __init__(self, powerstate, **kwargs): + super().__init__(powerstate, **kwargs) + + def produce_profiles(self): + self._produce_profiles() + + def evaluate(self): + + tglf = self._evaluate_tglf() + + # ************************************************************************************ + # Private functions for the evaluation + # ************************************************************************************ + + def _evaluate_tglf(self): + + # ------------------------------------------------------------------------------------------------------------------------ + # Grab options from powerstate + # ------------------------------------------------------------------------------------------------------------------------ + + ModelOptions = self.powerstate.TransportOptions["ModelOptions"] + + MODELparameters = ModelOptions.get("MODELparameters",None) + includeFast = ModelOptions.get("includeFastInQi",False) + launchMODELviaSlurm = ModelOptions.get("launchMODELviaSlurm", False) + cold_start = ModelOptions.get("cold_start", False) + provideTurbulentExchange = ModelOptions.get("TurbulentExchange", False) + percentError = ModelOptions.get("percentError", [5, 1, 0.5]) + use_tglf_scan_trick = ModelOptions.get("use_tglf_scan_trick", None) + cores_per_tglf_instance = ModelOptions.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) + + # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) + impurityPosition = self.powerstate.impurityPosition_transport #ModelOptions.get("impurityPosition", 1) + + # ------------------------------------------------------------------------------------------------------------------------ + # Prepare TGLF object + # ------------------------------------------------------------------------------------------------------------------------ + + RadiisToRun = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] + + tglf = TGLFtools.TGLF(rhos=RadiisToRun) + + _ = tglf.prep_direct_tglf( + self.folder, + cold_start = cold_start, + onlyThermal_TGYRO = not includeFast, + recalculatePTOT = False, # Use what's in the input.gacode, which is what PORTALS TGYRO does + inputgacode=self.powerstate.profiles_transport, + ) + + # ------------------------------------------------------------------------------------------------------------------------ + # Run TGLF + # ------------------------------------------------------------------------------------------------------------------------ + + Flux_base, Flux_mean, Flux_std = tglf_scan_trick( + tglf, + RadiisToRun, + self.powerstate.ProfilesPredicted, + impurityPosition=impurityPosition, + includeFast=includeFast, + delta = use_tglf_scan_trick, + cold_start=cold_start, + extra_name=self.name, + cores_per_tglf_instance=cores_per_tglf_instance + ) + + + # ------------------------------------------------------------------------------------------------------------------------ + # Pass the information to POWERSTATE + # ------------------------------------------------------------------------------------------------------------------------ + + self.powerstate.plasma["QeMWm2_tr_turb"] = Flux_mean[0] + self.powerstate.plasma["QeMWm2_tr_turb_stds"] = Flux_std[0] + + self.powerstate.plasma["QiMWm2_tr_turb"] = Flux_mean[1] + self.powerstate.plasma["QiMWm2_tr_turb_stds"] = Flux_std[1] + + self.powerstate.plasma["Ce_tr_turb"] = Flux_mean[2] + self.powerstate.plasma["Ce_tr_turb_stds"] = Flux_std[2] + + self.powerstate.plasma["QeMWm2_tr"] = self.powerstate.plasma["QeMWm2_tr_turb"]+ 0.0 + self.powerstate.plasma["QiMWm2_tr"] = self.powerstate.plasma["QiMWm2_tr_turb"]+ 0.0 + self.powerstate.plasma["Ce_tr"] = self.powerstate.plasma["Ce_tr_turb"] + 0.0 + + for variable in ['QeMWm2', 'QiMWm2', 'Ce']: + for suffix in ['_tr','_tr_turb', '_tr_turb_stds']: + + # Make them tensors and add a batch dimension + self.powerstate.plasma[f"{variable}{suffix}"] = torch.Tensor(self.powerstate.plasma[f"{variable}{suffix}"]).to(self.powerstate.dfT).unsqueeze(0) + + # Pad with zeros at rho=0.0 + self.powerstate.plasma[f"{variable}{suffix}"] = torch.cat(( + torch.zeros((1, 1)), + self.powerstate.plasma[f"{variable}{suffix}"], + ), dim=1) + + return tglf + + +def tglf_scan_trick( + tglf, + RadiisToRun, + ProfilesPredicted, + impurityPosition=1, + includeFast=False, + delta=0.02, + minimum_abs_gradient=0.005, # This is 0.5% of aLx=1.0, to avoid extremely small scans when, for example, having aLn ~ 0.0 + cold_start=False, + extra_name="", + remove_folders_out = False, + cores_per_tglf_instance = 4 # e.g. 4 core per radius, since this is going to launch ~ Nr=5 x (Nv=6 x Nd=2 + 1) = 65 TGLFs at once + ): + + print(f"\t- Running TGLF standalone scans ({delta = }) to determine relative errors") + + # Prepare scan + variables_to_scan = [] + for i in ProfilesPredicted: + if i == 'te': variables_to_scan.append('RLTS_1') + if i == 'ti': variables_to_scan.append('RLTS_2') + if i == 'ne': variables_to_scan.append('RLNS_1') + if i == 'nZ': variables_to_scan.append(f'RLNS_{impurityPosition+2}') + if i == 'w0': variables_to_scan.append('VEXB_SHEAR') #TODO: is this correct? or VPAR_SHEAR? + + #TODO: Only if that parameter is changing at that location + if 'te' in ProfilesPredicted or 'ti' in ProfilesPredicted: + variables_to_scan.append('TAUS_2') + if 'te' in ProfilesPredicted or 'ne' in ProfilesPredicted: + variables_to_scan.append('XNUE') + if 'te' in ProfilesPredicted or 'ne' in ProfilesPredicted: + variables_to_scan.append('BETAE') + + relative_scan = [1-delta, 1+delta] + + # Enforce at least "minimum_abs_gradient" in gradient, to avoid zero gradient situations + minimum_delta_abs = {} + for ikey in variables_to_scan: + if 'RL' in ikey: + minimum_delta_abs[ikey] = minimum_abs_gradient + + name = 'turb_drives' + + tglf.rhos = RadiisToRun # To avoid the case in which TGYRO was run with an extra rho point + + # Estimate job minutes based on cases and cores (mostly IO I think at this moment, otherwise it should be independent on cases) + num_cases = len(RadiisToRun) * len(variables_to_scan) * len(relative_scan) + if cores_per_tglf_instance == 1: + minutes = 10 * (num_cases / 60) # Ad-hoc formula + else: + minutes = 1 * (num_cases / 60) # Ad-hoc formula + + # Enforce minimum minutes + minutes = max(2, minutes) + + tglf.runScanTurbulenceDrives( + subFolderTGLF = name, + variablesDrives = variables_to_scan, + varUpDown = relative_scan, + minimum_delta_abs = minimum_delta_abs, + TGLFsettings = None, + ApplyCorrections = False, + add_baseline_to = 'first', + cold_start=cold_start, + forceIfcold_start=True, + slurm_setup={ + "cores": cores_per_tglf_instance, + "minutes": minutes, + }, + extra_name = f'{extra_name}_{name}', + positionIon=impurityPosition+2, + attempts_execution=2, + only_minimal_files=True, # Since I only care about fluxes here, do not retrieve all the files + ) + + # Remove folders because they are heavy to carry many throughout + if remove_folders_out: + IOtools.shutil_rmtree(tglf.FolderGACODE) + + Qe = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + Qi = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + Ge = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + GZ = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + Mt = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + S = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + + cont = 0 + for vari in variables_to_scan: + jump = tglf.scans[f'{name}_{vari}']['Qe'].shape[-1] + + Qe[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qe'] + Qi[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qi'] + Ge[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Ge'] + GZ[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Gi'] + Mt[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Mt'] + S[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['S'] + cont += jump + + # Calculate the standard deviation of the scans, that's going to be the reported stds + + def calculate_mean_std(Q): + # Assumes Q is [radii, points], with [radii, 0] being the baseline + + Qm = np.mean(Q, axis=1) + Qstd = np.std(Q, axis=1) + + # Qm = Q[:,0] + # Qstd = np.std(Q, axis=1) + + # Qstd = ( Q.max(axis=1)-Q.min(axis=1) )/2 /2 # Such that the range is 2*std + # Qm = Q.min(axis=1) + Qstd*2 # Mean is at the middle of the range + + return Qm, Qstd + + Qe_point, Qe_std = calculate_mean_std(Qe) + Qi_point, Qi_std = calculate_mean_std(Qi) + Ge_point, Ge_std = calculate_mean_std(Ge) + GZ_point, GZ_std = calculate_mean_std(GZ) + Mt_point, Mt_std = calculate_mean_std(Mt) + S_point, S_std = calculate_mean_std(S) + + #TODO: Careful with fast particles + + Flux_base = [Qe[:,0], Qi[:,0], Ge[:,0], GZ[:,0], Mt[:,0], S[:,0]] + Flux_mean = [Qe_point, Qi_point, Ge_point, GZ_point, Mt_point, S_point] + Flux_std = [Qe_std, Qi_std, Ge_std, GZ_std, Mt_std, S_std] + + return Flux_base, Flux_mean, Flux_std diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 49991cc5..5a654803 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -233,12 +233,17 @@ def prep( # PROFILES class. - from mitim_tools.gacode_tools import PROFILEStools - profiles = ( - PROFILEStools.gacode_state(inputgacode) - if inputgacode is not None - else None - ) + if inputgacode is not None: + + if isinstance(inputgacode, str) or isinstance(inputgacode, Path): + + from mitim_tools.gacode_tools import PROFILEStools + profiles = PROFILEStools.gacode_state(inputgacode) + + else: + + # If inputgacode is already a PROFILEStools object, just use it + profiles = inputgacode # TGYRO class. It checks existence and creates input.profiles/input.gacode @@ -373,14 +378,19 @@ def prep_direct_tglf( # PROFILES class. - from mitim_tools.gacode_tools import PROFILEStools - self.profiles = ( - PROFILEStools.gacode_state(inputgacode) - if inputgacode is not None - else None - ) - - if self.profiles is None: + if inputgacode is not None: + + if isinstance(inputgacode, str) or isinstance(inputgacode, Path): + + from mitim_tools.gacode_tools import PROFILEStools + self.profiles = PROFILEStools.gacode_state(inputgacode) + + else: + + # If inputgacode is already a PROFILEStools object, just use it + self.profiles = inputgacode + + else: # TGYRO class. It checks existence and creates input.profiles/input.gacode @@ -413,6 +423,7 @@ def prep_direct_tglf( for rho in self.inputsTGLF: self.inputsTGLF[rho] = TGLFinput.initialize_in_memory(self.inputsTGLF[rho]) + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize by taking directly the inputs # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index b18e4eeb..658c400f 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -380,7 +380,6 @@ def run( # ------ Write input profiles # ----------------------------------- - print(f"\t\t- Using input.profiles from {IOtools.clipstr(self.profiles.file)}") fil = "input.gacode" if self.profiles.profiles['rho(-)'][0] > 0.0: @@ -1201,8 +1200,8 @@ class TGYROoutput: def __init__(self, FolderTGYRO, profiles=None): self.FolderTGYRO = FolderTGYRO + from mitim_tools.gacode_tools import PROFILEStools if (profiles is None) and (FolderTGYRO / "input.gacode").exists(): - from mitim_tools.gacode_tools import PROFILEStools profiles = PROFILEStools.gacode_state(FolderTGYRO / f"input.gacode", derive_quantities=False) self.profiles = profiles diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 00db28be..19130f7b 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2247,7 +2247,7 @@ def interpolator(y): # --------------------------------------------------------------------------------------------------------------------------------------- plasma = { - 'NS': len(species)+1, + 'NS': len(species), 'SIGN_BT': -1.0, 'SIGN_IT': -1.0, 'VEXB': 0.0, @@ -2299,7 +2299,7 @@ def interpolator(y): input_dict[f'{k}_{i+1}'] = species[i+1][k] inputsTGLF[rho] = input_dict - + return inputsTGLF def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', times = [0.0,1.0], Vsurf = 0.0): diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index 2a5387b5..a3f23da4 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -140,7 +140,7 @@ def plot_geometry(self, axs, color="b", legYN=True, extralab="", lw=1, fs=6): def plot_state_flux_surfaces(self, ax=None, c='b'): rhos_plot = np.linspace(0.0, 1.0, 10) - phis_plot = [0.0, np.pi/2, np.pi, 3*np.pi/2] #np.linspace(0.0, 2.0 * np.pi, 5) + phis_plot = [0.0, np.pi/2, np.pi, 3*np.pi/2] ls = GRAPHICStools.listLS() From 5b3caafe4d89c6b819f6e8ae6c9358b3e5ebe35e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Jul 2025 20:06:57 -0400 Subject: [PATCH 103/385] use actual tglf settings --- .../physics_models/transport_tglf.py | 18 +++++++++++++----- src/mitim_tools/gacode_tools/TGLFtools.py | 5 +---- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 5de5dc7e..7793824b 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -61,10 +61,12 @@ def _evaluate_tglf(self): # Run TGLF # ------------------------------------------------------------------------------------------------------------------------ - Flux_base, Flux_mean, Flux_std = tglf_scan_trick( + Flux_base, Flux_mean, Flux_std = _run_tglf_model( tglf, RadiisToRun, self.powerstate.ProfilesPredicted, + TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], + extraOptionsTGLF=MODELparameters["transport_model"]["extraOptionsTGLF"], impurityPosition=impurityPosition, includeFast=includeFast, delta = use_tglf_scan_trick, @@ -73,7 +75,6 @@ def _evaluate_tglf(self): cores_per_tglf_instance=cores_per_tglf_instance ) - # ------------------------------------------------------------------------------------------------------------------------ # Pass the information to POWERSTATE # ------------------------------------------------------------------------------------------------------------------------ @@ -87,10 +88,15 @@ def _evaluate_tglf(self): self.powerstate.plasma["Ce_tr_turb"] = Flux_mean[2] self.powerstate.plasma["Ce_tr_turb_stds"] = Flux_std[2] + # Turbulence + Neoclassical (#TODO: NEO is not implemented yet) self.powerstate.plasma["QeMWm2_tr"] = self.powerstate.plasma["QeMWm2_tr_turb"]+ 0.0 self.powerstate.plasma["QiMWm2_tr"] = self.powerstate.plasma["QiMWm2_tr_turb"]+ 0.0 self.powerstate.plasma["Ce_tr"] = self.powerstate.plasma["Ce_tr_turb"] + 0.0 + # ------------------------------------------------------------------------------------------------------------------------ + # Curate information for the powerstate (e.g. add batch dimension, rho=0.0, and tensorize) + # ------------------------------------------------------------------------------------------------------------------------ + for variable in ['QeMWm2', 'QiMWm2', 'Ce']: for suffix in ['_tr','_tr_turb', '_tr_turb_stds']: @@ -105,11 +111,12 @@ def _evaluate_tglf(self): return tglf - -def tglf_scan_trick( +def _run_tglf_model( tglf, RadiisToRun, ProfilesPredicted, + TGLFsettings=None, + extraOptionsTGLF=None, impurityPosition=1, includeFast=False, delta=0.02, @@ -166,7 +173,8 @@ def tglf_scan_trick( variablesDrives = variables_to_scan, varUpDown = relative_scan, minimum_delta_abs = minimum_delta_abs, - TGLFsettings = None, + TGLFsettings = TGLFsettings, + extraOptions = extraOptionsTGLF, ApplyCorrections = False, add_baseline_to = 'first', cold_start=cold_start, diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 5a654803..9b768eb9 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -458,10 +458,7 @@ def prep_direct_tglf( print("> Setting up normalizations") - print( - "\t- Using mass of deuterium to normalize things (not necesarily the first ion)", - typeMsg="w", - ) + print("\t- Using mass of deuterium to normalize things (not necesarily the first ion)",typeMsg="w",) self.profiles.derive_quantities(mi_ref=mi_D) self.NormalizationSets, cdf = NORMtools.normalizations( From 460513926002ab8fa3ff070efdba65dc9e033493 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 24 Jul 2025 22:36:47 -0400 Subject: [PATCH 104/385] Base case working, with no UQ tglf --- .../physics_models/transport_tglf.py | 113 +++++++++++++----- 1 file changed, 86 insertions(+), 27 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 7793824b..9515c776 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -1,3 +1,4 @@ +from sympy import per import torch import numpy as np from mitim_tools.misc_tools import IOtools @@ -16,6 +17,14 @@ def produce_profiles(self): def evaluate(self): tglf = self._evaluate_tglf() + neo = self._evaluate_neo() + + # Sum the turbulent and neoclassical contributions + variables = ['QeMWm2', 'QiMWm2', 'Ce'] + + for variable in variables: + # Add model suffixes + self.powerstate.plasma[f"{variable}_tr"] = self.powerstate.plasma[f"{variable}_tr_turb"] + self.powerstate.plasma[f"{variable}_tr_neo"] # ************************************************************************************ # Private functions for the evaluation @@ -61,19 +70,54 @@ def _evaluate_tglf(self): # Run TGLF # ------------------------------------------------------------------------------------------------------------------------ - Flux_base, Flux_mean, Flux_std = _run_tglf_model( - tglf, - RadiisToRun, - self.powerstate.ProfilesPredicted, - TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], - extraOptionsTGLF=MODELparameters["transport_model"]["extraOptionsTGLF"], - impurityPosition=impurityPosition, - includeFast=includeFast, - delta = use_tglf_scan_trick, - cold_start=cold_start, - extra_name=self.name, - cores_per_tglf_instance=cores_per_tglf_instance - ) + if use_tglf_scan_trick is None: + + tglf.run( + 'base', + TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], + extraOptions= MODELparameters["transport_model"]["extraOptionsTGLF"], + ApplyCorrections=False, + launchSlurm= launchMODELviaSlurm, + cold_start= cold_start, + forceIfcold_start=True, + extra_name= self.name, + anticipate_problems=True, + slurm_setup={ + "cores": cores_per_tglf_instance, + "minutes": 2, + }, + attempts_execution=2, + only_minimal_files=True, + ) + + tglf.read(label='base',require_all_files=False) + + + Qe = [tglf.results['base']['TGLFout'][i].Qe_unn for i in range(len(RadiisToRun))] + Qi = [tglf.results['base']['TGLFout'][i].Qi_unn for i in range(len(RadiisToRun))] + Ge = [tglf.results['base']['TGLFout'][i].Ge_unn for i in range(len(RadiisToRun))] + GZ = [tglf.results['base']['TGLFout'][i].GiAll_unn[impurityPosition] for i in range(len(RadiisToRun))] + Mt = [tglf.results['base']['TGLFout'][i].Mt_unn for i in range(len(RadiisToRun))] + S = [tglf.results['base']['TGLFout'][i].Se_unn for i in range(len(RadiisToRun))] + + Flux_mean = np.array([Qe, Qi, Ge, GZ, Mt, S]) + Flux_std = abs(Flux_mean)*percentError[0]/100.0 + + else: + + Flux_base, Flux_mean, Flux_std = _run_tglf_uncertainty_model( + tglf, + RadiisToRun, + self.powerstate.ProfilesPredicted, + TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], + extraOptionsTGLF=MODELparameters["transport_model"]["extraOptionsTGLF"], + impurityPosition=impurityPosition, + delta = use_tglf_scan_trick, + cold_start=cold_start, + extra_name=self.name, + cores_per_tglf_instance=cores_per_tglf_instance, + launchMODELviaSlurm=launchMODELviaSlurm, + ) # ------------------------------------------------------------------------------------------------------------------------ # Pass the information to POWERSTATE @@ -81,24 +125,29 @@ def _evaluate_tglf(self): self.powerstate.plasma["QeMWm2_tr_turb"] = Flux_mean[0] self.powerstate.plasma["QeMWm2_tr_turb_stds"] = Flux_std[0] - + self.powerstate.plasma["QiMWm2_tr_turb"] = Flux_mean[1] self.powerstate.plasma["QiMWm2_tr_turb_stds"] = Flux_std[1] - + self.powerstate.plasma["Ce_tr_turb"] = Flux_mean[2] - self.powerstate.plasma["Ce_tr_turb_stds"] = Flux_std[2] - - # Turbulence + Neoclassical (#TODO: NEO is not implemented yet) - self.powerstate.plasma["QeMWm2_tr"] = self.powerstate.plasma["QeMWm2_tr_turb"]+ 0.0 - self.powerstate.plasma["QiMWm2_tr"] = self.powerstate.plasma["QiMWm2_tr_turb"]+ 0.0 - self.powerstate.plasma["Ce_tr"] = self.powerstate.plasma["Ce_tr_turb"] + 0.0 + self.powerstate.plasma["Ce_tr_turb_stds"] = Flux_std[2] + # if provideTurbulentExchange: + # self.powerstate.plasma["PexchTurb"] = + # self.powerstate.plasma["PexchTurb_stds"] = + # else: + # self.powerstate.plasma["PexchTurb"] = self.powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + # self.powerstate.plasma["PexchTurb_stds"] = self.powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + + # ------------------------------------------------------------------------------------------------------------------------ - # Curate information for the powerstate (e.g. add batch dimension, rho=0.0, and tensorize) + # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) # ------------------------------------------------------------------------------------------------------------------------ - for variable in ['QeMWm2', 'QiMWm2', 'Ce']: - for suffix in ['_tr','_tr_turb', '_tr_turb_stds']: + variables = ['QeMWm2', 'QiMWm2', 'Ce'] + + for variable in variables: + for suffix in ['_tr_turb', '_tr_turb_stds']: # Make them tensors and add a batch dimension self.powerstate.plasma[f"{variable}{suffix}"] = torch.Tensor(self.powerstate.plasma[f"{variable}{suffix}"]).to(self.powerstate.dfT).unsqueeze(0) @@ -111,20 +160,29 @@ def _evaluate_tglf(self): return tglf -def _run_tglf_model( + def _evaluate_neo(self): + + self.powerstate.plasma["QeMWm2_tr_neo"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) + self.powerstate.plasma["QiMWm2_tr_neo"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) + self.powerstate.plasma["Ce_tr_neo"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) + + return None + + +def _run_tglf_uncertainty_model( tglf, RadiisToRun, ProfilesPredicted, TGLFsettings=None, extraOptionsTGLF=None, impurityPosition=1, - includeFast=False, delta=0.02, minimum_abs_gradient=0.005, # This is 0.5% of aLx=1.0, to avoid extremely small scans when, for example, having aLn ~ 0.0 cold_start=False, extra_name="", remove_folders_out = False, - cores_per_tglf_instance = 4 # e.g. 4 core per radius, since this is going to launch ~ Nr=5 x (Nv=6 x Nd=2 + 1) = 65 TGLFs at once + cores_per_tglf_instance = 4, # e.g. 4 core per radius, since this is going to launch ~ Nr=5 x (Nv=6 x Nd=2 + 1) = 65 TGLFs at once + launchMODELviaSlurm=False, ): print(f"\t- Running TGLF standalone scans ({delta = }) to determine relative errors") @@ -187,6 +245,7 @@ def _run_tglf_model( positionIon=impurityPosition+2, attempts_execution=2, only_minimal_files=True, # Since I only care about fluxes here, do not retrieve all the files + launchSlurm=launchMODELviaSlurm, ) # Remove folders because they are heavy to carry many throughout From a1f0fc58ef01c233b57d355b5dad2f78a8516ef9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Jul 2025 07:03:24 -0400 Subject: [PATCH 105/385] CHANGE IN CONVENTION: input.gacode always initialized with Deuterium ref mass now, no first ion --- src/mitim_tools/plasmastate_tools/MITIMstate.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 19130f7b..5f2d8d6d 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -142,8 +142,8 @@ def derive_quantities_base(self, mi_ref=None, derive_quantities=True, rederiveGe self.derived["mi_ref"] = mi_ref print(f"\t* Reference mass ({self.derived['mi_ref']:.2f}) to use was forced by class initialization",typeMsg="w") else: - self.derived["mi_ref"] = self.mi_first - print(f"\t* Reference mass ({self.derived['mi_ref']}) from first ion",typeMsg="i") + self.derived["mi_ref"] = 2.0 #self.mi_first + print(f"\t* Reference mass ({self.derived['mi_ref']}) from Deuterium, as convention in gacode",typeMsg="i") # Useful to have gradients in the basic ---------------------------------------------------------- self.derived["aLTe"] = aLT(self.derived["r"], self.profiles["te(keV)"]) @@ -2211,13 +2211,14 @@ def interpolator(y): #mass_ref = self.derived["mi_ref"] # input.gacode uses the deuterium mass as reference already (https://github.com/gafusion/gacode/issues/398), so this should be 2.0 mass_ref = 2.0 - - mass_e = 0.000272445 * mass_ref + + if mass_ref != self.derived["mi_ref"]: + print(f"\t- Warning: the mass reference in the input.gacode is {self.derived['mi_ref']}, but TGLF expects {mass_ref}. This may lead to problems with the TGLF input file.", typeMsg="q") species = { 1: { 'ZS': -1.0, - 'MASS': mass_e/mass_ref, + 'MASS': 0.000272445, 'RLNS': interpolator(self.derived['aLne']), 'RLTS': interpolator(self.derived['aLTe']), 'TAUS': 1.0, From 0012943b642f73edf895d720e4f634d533d6e2f7 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Jul 2025 08:36:20 -0400 Subject: [PATCH 106/385] Store input.gacode as well --- .../physics_models/transport_tglf.py | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 9515c776..c04a4d20 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -1,4 +1,5 @@ -from sympy import per + +import shutil import torch import numpy as np from mitim_tools.misc_tools import IOtools @@ -72,6 +73,8 @@ def _evaluate_tglf(self): if use_tglf_scan_trick is None: + # Just run TGLF once and apply an ad-hoc percent error to the results + tglf.run( 'base', TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], @@ -92,7 +95,6 @@ def _evaluate_tglf(self): tglf.read(label='base',require_all_files=False) - Qe = [tglf.results['base']['TGLFout'][i].Qe_unn for i in range(len(RadiisToRun))] Qi = [tglf.results['base']['TGLFout'][i].Qi_unn for i in range(len(RadiisToRun))] Ge = [tglf.results['base']['TGLFout'][i].Ge_unn for i in range(len(RadiisToRun))] @@ -105,6 +107,8 @@ def _evaluate_tglf(self): else: + # Run TGLF with scans to estimate the uncertainty + Flux_base, Flux_mean, Flux_std = _run_tglf_uncertainty_model( tglf, RadiisToRun, @@ -168,6 +172,20 @@ def _evaluate_neo(self): return None + def _profiles_to_store(self): + + if "extra_params" in self.powerstate.TransportOptions["ModelOptions"] and "folder" in self.powerstate.TransportOptions["ModelOptions"]["extra_params"]: + whereFolder = IOtools.expandPath(self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") + if not whereFolder.exists(): + IOtools.askNewFolder(whereFolder) + + fil = whereFolder / f"input.gacode.{self.evaluation_number}" + shutil.copy2(self.file_profs, fil) + shutil.copy2(self.file_profs_unmod, fil.parent / f"{fil.name}_unmodified") + print(f"\t- Copied profiles to {IOtools.clipstr(fil)}") + else: + print("\t- Could not move files", typeMsg="w") + def _run_tglf_uncertainty_model( tglf, From 707062b1c8be53914865cd250311a02814c43571 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Jul 2025 08:37:19 -0400 Subject: [PATCH 107/385] Bug zeff vs z_eff; add shape coefficients to_tglf --- src/mitim_tools/gacode_tools/TGLFtools.py | 20 ++++------- .../misc_tools/scripts/compare_namelist.py | 33 ++++++++++++------- .../plasmastate_tools/MITIMstate.py | 28 ++++++++++++---- 3 files changed, 50 insertions(+), 31 deletions(-) diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 9b768eb9..591d3646 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -4,6 +4,7 @@ import numpy as np import xarray as xr import matplotlib.pyplot as plt +from mitim_tools import __version__ as mitim_version from mitim_tools.gacode_tools import TGYROtools from mitim_tools.misc_tools import ( @@ -422,8 +423,7 @@ def prep_direct_tglf( for rho in self.inputsTGLF: self.inputsTGLF[rho] = TGLFinput.initialize_in_memory(self.inputsTGLF[rho]) - - + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize by taking directly the inputs # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -3900,7 +3900,7 @@ def reduceToControls(dict_all): ]: plasma[ikey] = dict_all[ikey] - elif (len(ikey.split("_")) > 1) and (ikey.split("_")[-1] in ["SA", "LOC"]): + elif (len(ikey.split("_")) > 1) and ( (ikey.split("_")[-1] in ["SA", "LOC"]) or (ikey.split("_")[0] in ["SHAPE"]) ): geom[ikey] = dict_all[ikey] else: @@ -4143,15 +4143,9 @@ def write_state(self, file=None): file = self.file with open(file, "w") as f: - f.write( - "#-------------------------------------------------------------------------\n" - ) - f.write( - "# TGLF input file modified by MITIM framework (Rodriguez-Fernandez, 2020)\n" - ) - f.write( - "#-------------------------------------------------------------------------" - ) + f.write("#-------------------------------------------------------------------------\n") + f.write(f"# TGLF input file modified by MITIM {mitim_version}\n") + f.write("#-------------------------------------------------------------------------") f.write("\n\n# Control parameters\n") f.write("# ------------------\n\n") @@ -4164,7 +4158,7 @@ def write_state(self, file=None): for ikey in self.geom: var = self.geom[ikey] f.write(f"{ikey.ljust(23)} = {var}\n") - + f.write("\n\n# Plasma parameters\n") f.write("# ------------------\n\n") for ikey in self.plasma: diff --git a/src/mitim_tools/misc_tools/scripts/compare_namelist.py b/src/mitim_tools/misc_tools/scripts/compare_namelist.py index 9ad08b77..c739e2d3 100644 --- a/src/mitim_tools/misc_tools/scripts/compare_namelist.py +++ b/src/mitim_tools/misc_tools/scripts/compare_namelist.py @@ -1,4 +1,4 @@ -import sys +import argparse import numpy as np from mitim_tools.misc_tools import IOtools from IPython import embed @@ -91,7 +91,10 @@ def compare_number(a,b,precision_of=None): else: decimal_places = 0 - b_rounded = round(b, decimal_places) + if isinstance(b, str): + b_rounded = b + else: + b_rounded = round(b, decimal_places) a_rounded = a elif precision_of == 2: @@ -150,9 +153,9 @@ def printTable(diff, warning_percent=1e-1): typeMsg="i" if perc > warning_percent else "", ) else: - print(f"{key:>15}{str(diff[key][0]):>25}{'':>25}") + print(f"{key:>15}{str(diff[key][0]):>25}{'':>25} (100%)", typeMsg="i") else: - print(f"{key:>15}{'':>25}{str(diff[key][1]):>25}") + print(f"{key:>15}{'':>25}{str(diff[key][1]):>25} (100%)", typeMsg="i") print( "--------------------------------------------------------------------------------" ) @@ -160,19 +163,25 @@ def printTable(diff, warning_percent=1e-1): def main(): - file1 = sys.argv[1] - file2 = sys.argv[2] + parser = argparse.ArgumentParser() + parser.add_argument("file1", type=str, help="First namelist file to compare") + parser.add_argument("file2", type=str, help="Second namelist file to compare") + parser.add_argument("--separator", type=str, required=False, default="=", + help="Separator used in the namelist files, default is '='") + parser.add_argument("--precision", type=int, required=False, default=None, + help="Precision for comparing numbers: 1 for decimal places, 2 for significant figures, None for exact comparison") + args = parser.parse_args() - try: - separator = sys.argv[3] - except: - separator = "=" + # Get arguments + file1 = args.file1 + file2 = args.file2 + separator = args.separator + precision = args.precision - diff = compareNML(file1, file2, separator=separator) + diff = compareNML(file1, file2, separator=separator, precision_of=precision) printTable(diff) print(f"Differences: {len(diff)}") - if __name__ == "__main__": main() diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 5f2d8d6d..47ff5e09 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -15,7 +15,6 @@ from mitim_tools import __version__ from IPython import embed - def ensure_variables_existence(self): # --------------------------------------------------------------------------- # Determine minimal set of variables that should be present in the profiles @@ -29,7 +28,7 @@ def ensure_variables_existence(self): "ni(10^19/m^3)": 2, "w0(rad/s)": 1, "ptot(Pa)": 1, - "zeff(-)": 1, + "z_eff(-)": 1, } # Electromagnetics @@ -268,7 +267,7 @@ def calculate_Er( if write_new_file is not None: self.write_state(file=write_new_file) - def readSpecies(self, maxSpecies=100): + def readSpecies(self, maxSpecies=100, correct_zeff = True): maxSpecies = int(self.profiles["nion"][0]) Species = [] @@ -287,7 +286,11 @@ def readSpecies(self, maxSpecies=100): Species.append(sp) self.Species = Species - + + # Correct Zeff if needed + if correct_zeff and ("z_eff(-)" in self.profiles): + self.profiles["z_eff(-)"] = np.sum(self.profiles["ni(10^19/m^3)"] * self.profiles["z"] ** 2, axis=1) / self.profiles["ne(10^19/m^3)"] + def sumFast(self): self.nFast = self.profiles["ne(10^19/m^3)"] * 0.0 self.nZFast = self.profiles["ne(10^19/m^3)"] * 0.0 @@ -1648,7 +1651,7 @@ def correct(self, options={}, write=False, new_file=None): # Re-derive # ---------------------------------------------------------------------- - self.derive_quantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) # ---------------------------------------------------------------------- # Write @@ -2280,6 +2283,19 @@ def interpolator(y): 'P_PRIME_LOC': self.derived['pprime'], 'Q_PRIME_LOC': self.derived['s_q'], } + + # Add MXH and derivatives (#TODO) + for ikey in self.profiles: + if 'shape_cos' in ikey or 'shape_sin' in ikey: + + # TGLF only accepts 6, as of July 2025 + if int(ikey[-4]) > 6: + continue + + key_mod = ikey.upper().split('(')[0] # Remove any function call like 'shape_cos(1)' + + parameters[key_mod] = self.profiles[ikey] + parameters[f"{key_mod.split('_')[0]}_S_{key_mod.split('_')[-1]}"] = self.profiles[ikey]*0.0 #TODO geom = {} for k in parameters: @@ -2288,7 +2304,7 @@ def interpolator(y): geom['BETA_LOC'] = 0.0 geom['KX0_LOC'] = 0.0 - + # --------------------------------------------------------------------------------------------------------------------------------------- # Merging # --------------------------------------------------------------------------------------------------------------------------------------- From fb76a76b6927c0bb722bb2ce9885f3b5fba30840 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Jul 2025 10:23:18 -0400 Subject: [PATCH 108/385] Better plotting, including 3D tokamak --- src/mitim_tools/gacode_tools/PROFILEStools.py | 85 ++++++++++++++++++- .../plasmastate_tools/MITIMstate.py | 39 --------- .../plasmastate_tools/utils/VMECtools.py | 15 ++-- .../plasmastate_tools/utils/state_plotting.py | 27 +++--- 4 files changed, 107 insertions(+), 59 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index e7a9a742..0ec1fd6b 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -1,5 +1,7 @@ import copy +from turtle import color import numpy as np +import matplotlib.pyplot as plt from collections import OrderedDict from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.gs_tools import GEQtools @@ -255,7 +257,7 @@ def derive_geometry(self, n_theta_geo=1001, **kwargs): def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): - [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,_,ax13c] = axs3 + [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,ax03c,ax13c] = axs3 rho = self.profiles["rho(-)"] lines = GRAPHICStools.listLS() @@ -402,6 +404,87 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): ax.set_xlabel("R (m)") ax.set_ylabel("Z (m)") GRAPHICStools.addDenseAxis(ax) + + ax = ax03c + self.plot_plasma_boundary(ax=ax, color=color) + + def plot_state_flux_surfaces(self, ax=None, surfaces_rho=np.linspace(0, 1, 11), color="b", label = '', lw=1.0, lw1=2.0): + + if ax is None: + plt.ion() + fig, ax = plt.subplots() + provided = False + else: + provided = True + + for rho in surfaces_rho: + ir = np.argmin(np.abs(self.profiles["rho(-)"] - rho)) + + for i_toroidal in range(self.derived["R_surface"].shape[0]): + ax.plot( + self.derived["R_surface"][i_toroidal,ir, :], + self.derived["Z_surface"][i_toroidal,ir, :], + "-", + lw=lw if rho<1.0 else lw1, + c=color, + ) + + ax.axhline(y=0, ls="--", lw=0.2, c="k") + ax.plot( + [self.profiles["rmaj(m)"][0]], + [self.profiles["zmag(m)"][0]], + "o", + markersize=2, + c=color, + label = label + ) + + if not provided: + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + ax.set_title("Surfaces @ rho=" + str(surfaces_rho), fontsize=8) + ax.set_aspect("equal") + + def plot_plasma_boundary(self, ax=None, color="b"): + """ + Plot the 3D plasma boundary by extruding the poloidal cross-section toroidally. + """ + + n_phi = 50 # Number of toroidal points for the surface mesh + + R = self.derived["R_surface"][0,-1,:] # Outermost flux surface R coordinates + Z = self.derived["Z_surface"][0,-1,:] # Outermost flux surface Z coordinates + + # Create toroidal angle array + phi = np.linspace(0, 2*np.pi, n_phi) + + # Create meshgrid for toroidal extrusion + PHI, THETA_POINTS = np.meshgrid(phi, range(len(R))) + R_mesh = R[THETA_POINTS] + Z_mesh = Z[THETA_POINTS] + + # Convert to Cartesian coordinates + x = R_mesh * np.cos(PHI) + y = R_mesh * np.sin(PHI) + z = Z_mesh * np.ones_like(PHI) # Z doesn't depend on phi for axisymmetric case + + # Create the 3D plot + if ax is None: + fig = plt.figure(figsize=(10, 8)) + ax = fig.add_subplot(projection="3d") + + # Plot the surface + ax.plot_surface(x, y, z, alpha=0.7, color=color) + + # Set labels and title + ax.set_xlabel('X (m)') + ax.set_ylabel('Y (m)') + ax.set_zlabel('Z (m)') + ax.set_title('Plasma Boundary (3D)') + + # Set equal aspect ratio + ax.set_aspect("equal") + def calculateGeometricFactors(profiles, n_theta=1001): diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 47ff5e09..7519fd5e 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1902,45 +1902,6 @@ def plot_gradients(self, *args, **kwargs): def plot_geometry(self, *args, **kwargs): pass - - def plot_state_flux_surfaces(self, ax=None, surfaces_rho=np.linspace(0, 1, 11), color="b", label = '', lw=1.0, lw1=2.0): - if ("R_surface" in self.derived) and (self.derived["R_surface"] is not None): - if ax is None: - plt.ion() - fig, ax = plt.subplots() - provided = False - else: - provided = True - - for rho in surfaces_rho: - ir = np.argmin(np.abs(self.profiles["rho(-)"] - rho)) - - for i_toroidal in range(self.derived["R_surface"].shape[0]): - ax.plot( - self.derived["R_surface"][i_toroidal,ir, :], - self.derived["Z_surface"][i_toroidal,ir, :], - "-", - lw=lw if rho<1.0 else lw1, - c=color, - ) - - ax.axhline(y=0, ls="--", lw=0.2, c="k") - ax.plot( - [self.profiles["rmaj(m)"][0]], - [self.profiles["zmag(m)"][0]], - "o", - markersize=2, - c=color, - label = label - ) - - if not provided: - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - ax.set_title("Surfaces @ rho=" + str(surfaces_rho), fontsize=8) - ax.set_aspect("equal") - else: - print("\t- Cannot plot flux surface geometry", typeMsg="w") def plotPeaking( self, ax, c="b", marker="*", label="", debugPlot=False, printVals=False diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index a3f23da4..75e7aaf5 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -68,8 +68,8 @@ def _read_vmec(self): # Produce variables self.profiles["rho(-)"] = (self.wout.phi/self.wout.phi[-1])**0.5 #np.linspace(0, 1, self.wout.ns)**0.5 self.profiles["presf"] = self.wout.presf - self.profiles["q(-)"] = self.wout.q_factor - self.profiles["polflux(Wb/radian)"] = self.wout.chi + #self.profiles["q(-)"] = self.wout.q_factor + #self.profiles["polflux(Wb/radian)"] = self.wout.chi # Read Profiles data = self._read_profiles(x_coord=self.profiles["rho(-)"]) @@ -132,7 +132,7 @@ def plot_geometry(self, axs, color="b", legYN=True, extralab="", lw=1, fs=6): [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,axs_3d,axs_2d] = axs - self.plot_plasma_boundary(ax=axs_3d) + self.plot_plasma_boundary(ax=axs_3d, color=color) self.plot_state_flux_surfaces(ax=axs_2d, c=color) @@ -148,13 +148,14 @@ def plot_state_flux_surfaces(self, ax=None, c='b'): for i in range(len(rhos_plot)): self.plot_flux_surface(ax = ax, phi_cut=phi_cut, rho=rhos_plot[i], c=c, lw = 0.5, ls = lsi) - self.plot_flux_surface(ax = ax, phi_cut=phi_cut, rho=1.0, c=c, lw = 2, ls = lsi, label = f"$\\phi={phi_cut*180/np.pi:.1f} deg$") + self.plot_flux_surface(ax = ax, phi_cut=phi_cut, rho=1.0, c=c, lw = 2, ls = lsi, label = f"{phi_cut*180/np.pi:.1f}°") ax.set_aspect('equal') ax.set_xlabel('R [m]') ax.set_ylabel('Z [m]') GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.addLegendApart(ax, ratio=0.9, size=6) + ax.legend(loc='best', fontsize=6) + #GRAPHICStools.addLegendApart(ax, ratio=0.9, size=6) ax.set_title(f'Poloidal Cross-section') @@ -267,7 +268,7 @@ def _read_profiles(self, x_coord=None, debug = False): return uniform_data - def plot_plasma_boundary(self, ax=None): + def plot_plasma_boundary(self, ax=None, color="b"): # The output object contains the Fourier coefficients of the geometry in R and Z # as a function of the poloidal (theta) and toroidal (phi) angle-like coordinates @@ -320,7 +321,7 @@ def plot_plasma_boundary(self, ax=None): ax = fig.add_subplot(projection="3d") # Plot the surface - ax.plot_surface(x, y, z) + ax.plot_surface(x, y, z, alpha=0.7, color=color) # Set an equal aspect ratio ax.set_aspect("equal") diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py index 0519bdc8..af8b1275 100644 --- a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -46,21 +46,24 @@ def add_axes(figs): fig2.add_subplot(grid[0, 2]), fig2.add_subplot(grid[1, 2]), ] - grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.3) - ax00c = fig3.add_subplot(grid[0, 0]) + + # GEOMETRY + grid = plt.GridSpec(6, 4, hspace=0.7, wspace=0.3) + ax00c = fig3.add_subplot(grid[0:2, 0]) axsProf_3 = [ ax00c, - fig3.add_subplot(grid[1, 0], sharex=ax00c), - fig3.add_subplot(grid[2, 0]), - fig3.add_subplot(grid[0, 1]), - fig3.add_subplot(grid[1, 1]), - fig3.add_subplot(grid[2, 1]), - fig3.add_subplot(grid[0, 2]), - fig3.add_subplot(grid[1, 2]), - fig3.add_subplot(grid[2, 2]), - fig3.add_subplot(grid[0, 3], projection="3d"), - fig3.add_subplot(grid[1:, 3]), + fig3.add_subplot(grid[2:4, 0], sharex=ax00c), + fig3.add_subplot(grid[4:, 0]), + fig3.add_subplot(grid[0:2, 1]), + fig3.add_subplot(grid[2:4, 1]), + fig3.add_subplot(grid[4:, 1]), + fig3.add_subplot(grid[0:2, 2]), + fig3.add_subplot(grid[2:4, 2]), + fig3.add_subplot(grid[4:, 2]), + fig3.add_subplot(grid[0:3, 3:], projection="3d"), + fig3.add_subplot(grid[3:, 3:]), ] + # ---- grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) axsProf_4 = [ From 6166ed4e5e6da60c9a006f5b49224e36fcea26c6 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 25 Jul 2025 11:27:25 -0400 Subject: [PATCH 109/385] Impose dr/dx=1.0 otherwise BUG in tglf --- src/mitim_tools/plasmastate_tools/MITIMstate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 7519fd5e..4aad93db 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2231,7 +2231,7 @@ def interpolator(y): 'RMIN_LOC': self.derived['roa'], 'RMAJ_LOC': self.derived['Rmajoa'], 'ZMAJ_LOC': self.derived["Zmagoa"], - 'DRMINDX_LOC': self.derived['drmin/dr'], + 'DRMINDX_LOC': np.ones(self.profiles["rho(-)"].shape), # Force 1.0 instead of self.derived['drmin/dr'] because of numerical issues in TGLF 'DRMAJDX_LOC': self.derived['dRmaj/dr'], 'DZMAJDX_LOC': self.derived['dZmaj/dr'], 'Q_LOC': self.profiles["q(-)"], From 653001ca29be1bafe0c02d6f05a86a8b8d455f3b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 28 Jul 2025 17:56:52 +0200 Subject: [PATCH 110/385] Method for targets evaluator --- src/mitim_modules/portals/PORTALSmain.py | 19 +++++++++++-------- src/mitim_modules/portals/PORTALStools.py | 2 +- .../portals/utils/PORTALSinit.py | 4 ++-- src/mitim_modules/powertorch/STATEtools.py | 4 ++-- .../physics_models/transport_tgyro.py | 2 +- .../powertorch/scripts/calculateTargets.py | 4 ++-- .../powertorch/utils/TRANSFORMtools.py | 4 ++-- 7 files changed, 21 insertions(+), 18 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index ea56f4c5..ed337d7a 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -203,7 +203,7 @@ def __init__( "percentError": [5,10,1], # (%) Error (std, in percent) of model evaluation [TGLF (treated as minimum if scan trick), NEO, TARGET] "transport_evaluator": transport_evaluator, "targets_evaluator": targets_evaluator, - "TargetCalc": "powerstate", # Method to calculate targets (tgyro or powerstate) + "targets_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) "useConvectiveFluxes": True, # If True, then convective flux for final metric (not fitting). If False, particle flux "includeFastInQi": False, # If True, and fast ions have been included, in seprateNEO, sum fast @@ -221,7 +221,7 @@ def __init__( "applyImpurityGammaTrick": True, # If True, fit model to GZ/nZ, valid on the trace limit "UseOriginalImpurityConcentrationAsWeight": 1.0, # If not None, using UseOriginalImpurityConcentrationAsWeight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis "fImp_orig": 1.0, - "fineTargetsResolution": 20, # If not None, calculate targets with this radial resolution (defaults TargetCalc to powerstate) + "fineTargetsResolution": 20, # If not None, calculate targets with this radial resolution (defaults targets_evaluator_method to powerstate) "hardCodedCGYRO": None, # If not None, use this hard-coded CGYRO evaluation "additional_params_in_surrogate": additional_params_in_surrogate, "use_tglf_scan_trick": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta @@ -466,13 +466,13 @@ def check_flags(self): # ---------------------------------------------------------------------------------- if self.PORTALSparameters["fineTargetsResolution"] is not None: - if self.PORTALSparameters["TargetCalc"] != "powerstate": + if self.PORTALSparameters["targets_evaluator_method"] != "powerstate": print("\t- Requested fineTargetsResolution, so running powerstate target calculations",typeMsg="w") - self.PORTALSparameters["TargetCalc"] = "powerstate" + self.PORTALSparameters["targets_evaluator_method"] = "powerstate" - if not issubclass(self.PORTALSparameters["transport_evaluator"], transport_tgyro.tgyro_model) and (self.PORTALSparameters["TargetCalc"] == "tgyro"): + if not issubclass(self.PORTALSparameters["transport_evaluator"], transport_tgyro.tgyro_model) and (self.PORTALSparameters["targets_evaluator_method"] == "tgyro"): print("\t- Requested TGYRO targets, but transport evaluator is not tgyro, so changing to powerstate",typeMsg="w") - self.PORTALSparameters["TargetCalc"] = "powerstate" + self.PORTALSparameters["targets_evaluator_method"] = "powerstate" if ("InputType" not in self.MODELparameters["Physics_options"]) or self.MODELparameters["Physics_options"]["InputType"] != 1: print("\t- In PORTALS TGYRO evaluations, we need to use exact profiles (InputType=1)",typeMsg="i") @@ -482,8 +482,11 @@ def check_flags(self): print("\t- In PORTALS TGYRO evaluations, we need to not recompute gradients (GradientsType=0)",typeMsg="i") self.MODELparameters["Physics_options"]["GradientsType"] = 0 - if self.PORTALSparameters["TargetCalc"] == "tgyro" and self.PORTALSparameters['profiles_postprocessing_fun'] is not None: - print("\t- Requested custom modification of postprocessing function but targets from tgyro... are you sure?",typeMsg="q") + if self.PORTALSparameters["targets_evaluator_method"] == "tgyro" and self.PORTALSparameters['profiles_postprocessing_fun'] is not None: + print("\t- Requested custom modification of postprocessing function but targets from TGYRO... are you sure?",typeMsg="q") + + if self.PORTALSparameters["targets_evaluator_method"] == "tgyro" and self.PORTALSparameters['transport_evaluator'] != transport_tgyro.tgyro_model: + print("\t- Requested TGYRO targets but transport evaluator is not TGYRO... are you sure?",typeMsg="q") key_rhos = "RoaLocations" if self.MODELparameters["RoaLocations"] is not None else "RhoLocations" diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 9668a31a..34f83954 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -357,7 +357,7 @@ def constructEvaluationProfiles(X, surrogate_parameters, recalculateTargets=Fals # Targets only if needed (for speed, GB doesn't need it) if recalculateTargets: - powerstate.TargetOptions["ModelOptions"]["TargetCalc"] = "powerstate" # For surrogate evaluation, always powerstate, logically. + powerstate.TargetOptions["ModelOptions"]["targets_evaluator_method"] = "powerstate" # For surrogate evaluation, always powerstate, logically. powerstate.calculateTargets() return powerstate diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 8c6d88b9..756814ee 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -179,7 +179,7 @@ def initializeProblem( "targets_evaluator": portals_fun.PORTALSparameters["targets_evaluator"], "ModelOptions": { "TypeTarget": portals_fun.MODELparameters["Physics_options"]["TypeTarget"], - "TargetCalc": portals_fun.PORTALSparameters["TargetCalc"]}, + "targets_evaluator_method": portals_fun.PORTALSparameters["targets_evaluator_method"]}, }, tensor_opts = tensor_opts ) @@ -231,7 +231,7 @@ def initializeProblem( "targets_evaluator": portals_fun.PORTALSparameters["targets_evaluator"], "ModelOptions": { "TypeTarget": portals_fun.MODELparameters["Physics_options"]["TypeTarget"], - "TargetCalc": portals_fun.PORTALSparameters["TargetCalc"]}, + "targets_evaluator_method": portals_fun.PORTALSparameters["targets_evaluator_method"]}, }, tensor_opts = tensor_opts ) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index c116508a..2b3e5424 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -54,7 +54,7 @@ def __init__( "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": 3, - "TargetCalc": "powerstate" + "targets_evaluator_method": "powerstate" }, } if tensor_opts is None: @@ -681,7 +681,7 @@ def calculateTargets(self, relative_error_assumed=1.0): """ # If no targets evaluator is given or the targets will come from TGYRO, assume them as zero - if (self.TargetOptions["targets_evaluator"] is None) or (self.TargetOptions["ModelOptions"]["TargetCalc"] == "tgyro"): + if (self.TargetOptions["targets_evaluator"] is None) or (self.TargetOptions["ModelOptions"]["targets_evaluator_method"] == "tgyro"): targets = TARGETStools.power_targets(self) else: targets = self.TargetOptions["targets_evaluator"](self) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 57ffb452..abad823a 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -162,7 +162,7 @@ def _postprocess_results(self, tgyro, label): OriginalFimp=OriginalFimp, forceZeroParticleFlux=forceZeroParticleFlux, provideTurbulentExchange=provideTurbulentExchange, - provideTargets=self.powerstate.TargetOptions['ModelOptions']['TargetCalc'] == "tgyro", + provideTargets=self.powerstate.TargetOptions['ModelOptions']['targets_evaluator_method'] == "tgyro", ) tgyro.results["use"] = tgyro.results[label] diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 3cf769de..368781f6 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -36,7 +36,7 @@ def calculator( "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": TypeTarget, - "TargetCalc": "tgyro"}, + "targets_evaluator_method": "tgyro"}, }, TransportOptions={ "transport_evaluator": transport_tgyro.tgyro_model, @@ -78,7 +78,7 @@ def calculator( "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": TypeTarget, - "TargetCalc": "powerstate"}, + "targets_evaluator_method": "powerstate"}, }, TransportOptions={ "transport_evaluator": None, diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index f7f0070a..f131b24a 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -368,13 +368,13 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": self.TargetOptions["ModelOptions"]["TypeTarget"], # Important to keep the same as in the original - "TargetCalc": "powerstate", + "targets_evaluator_method": "powerstate", } }, increase_profile_resol = False ) state_temp.calculateProfileFunctions() - state_temp.TargetOptions["ModelOptions"]["TargetCalc"] = "powerstate" + state_temp.TargetOptions["ModelOptions"]["targets_evaluator_method"] = "powerstate" state_temp.calculateTargets() # ------------------------------------------------------------------------------------------ From d59ccc6e86344af45a2dea87e4c67b2eea57f05d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 29 Jul 2025 10:05:54 +0200 Subject: [PATCH 111/385] better nml comparisons --- .../misc_tools/scripts/compare_namelist.py | 23 ++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/src/mitim_tools/misc_tools/scripts/compare_namelist.py b/src/mitim_tools/misc_tools/scripts/compare_namelist.py index c739e2d3..878a8216 100644 --- a/src/mitim_tools/misc_tools/scripts/compare_namelist.py +++ b/src/mitim_tools/misc_tools/scripts/compare_namelist.py @@ -13,7 +13,7 @@ """ -def compareNML(file1, file2, commentCommand="!", separator="=", precision_of=None): +def compareNML(file1, file2, commentCommand="!", separator="=", precision_of=None, close_enough=1e-7): d1 = IOtools.generateMITIMNamelist( file1, commentCommand=commentCommand, separator=separator ) @@ -24,7 +24,7 @@ def compareNML(file1, file2, commentCommand="!", separator="=", precision_of=Non d1 = separateArrays(d1) d2 = separateArrays(d2) - diff = compareDictionaries(d1, d2, precision_of = precision_of) + diff = compareDictionaries(d1, d2, precision_of=precision_of, close_enough=close_enough) diffo = cleanDifferences(diff) @@ -77,7 +77,7 @@ def cleanDifferences(d, tol_rel=1e-7): return d_new -def compare_number(a,b,precision_of=None): +def compare_number(a,b,precision_of=None, close_enough=1e-7): if precision_of is None: a_rounded = a @@ -110,11 +110,15 @@ def compare_number(a,b,precision_of=None): a_rounded = round(a, decimal_places) # Compare the two numbers - are_equal = (a_rounded == b_rounded) + if isinstance(a_rounded, str) or isinstance(b_rounded, str): + # If either is a string, we cannot compare numerically + are_equal = a_rounded == b_rounded + else: + are_equal = np.isclose(a_rounded, b_rounded, rtol=close_enough) return are_equal -def compareDictionaries(d1, d2, precision_of=None): +def compareDictionaries(d1, d2, precision_of=None, close_enough=1e-7): different = {} for key in d1: @@ -123,7 +127,7 @@ def compareDictionaries(d1, d2, precision_of=None): different[key] = [d1[key], None] # Values are different else: - if not compare_number(d1[key],d2[key],precision_of=precision_of): + if not compare_number(d1[key],d2[key],precision_of=precision_of, close_enough=close_enough): different[key] = [d1[key], d2[key]] for key in d2: @@ -134,7 +138,7 @@ def compareDictionaries(d1, d2, precision_of=None): return different -def printTable(diff, warning_percent=1e-1): +def printTable(diff, printing_percent = 1e-5, warning_percent=1e-1): for key in diff: if diff[key][0] is not None: @@ -170,6 +174,8 @@ def main(): help="Separator used in the namelist files, default is '='") parser.add_argument("--precision", type=int, required=False, default=None, help="Precision for comparing numbers: 1 for decimal places, 2 for significant figures, None for exact comparison") + parser.add_argument("--close_enough", type=float, required=False, default=1e-7, + help="Tolerance for comparing numbers, default is 1e-7") args = parser.parse_args() # Get arguments @@ -177,8 +183,9 @@ def main(): file2 = args.file2 separator = args.separator precision = args.precision + close_enough = args.close_enough - diff = compareNML(file1, file2, separator=separator, precision_of=precision) + diff = compareNML(file1, file2, separator=separator, precision_of=precision, close_enough=close_enough) printTable(diff) print(f"Differences: {len(diff)}") From ec0c5c881b97cfb42fbb700f8aa9486d38ac55a2 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 29 Jul 2025 10:49:38 +0200 Subject: [PATCH 112/385] Correct signs to_tglf --- src/mitim_tools/gacode_tools/TGLFtools.py | 99 ++----------------- .../gacode_tools/utils/GACODErun.py | 2 +- .../plasmastate_tools/MITIMstate.py | 27 +++-- 3 files changed, 27 insertions(+), 101 deletions(-) diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 591d3646..a50ca675 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -13,6 +13,7 @@ PLASMAtools, GUItools, ) +from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_tools.gacode_tools.utils import ( NORMtools, GACODEinterpret, @@ -1264,87 +1265,9 @@ def plot( axFluc02Sym = axFluc02.twinx() if plotGACODE: - grid = plt.GridSpec(3, 3, hspace=0.3, wspace=0.3) - axsProf_1 = [ - figProf_1.add_subplot(grid[0, 0]), - figProf_1.add_subplot(grid[1, 0]), - figProf_1.add_subplot(grid[2, 0]), - figProf_1.add_subplot(grid[0, 1]), - figProf_1.add_subplot(grid[1, 1]), - figProf_1.add_subplot(grid[2, 1]), - figProf_1.add_subplot(grid[0, 2]), - figProf_1.add_subplot(grid[1, 2]), - figProf_1.add_subplot(grid[2, 2]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsProf_2 = [ - figProf_2.add_subplot(grid[0, 0]), - figProf_2.add_subplot(grid[0, 1]), - figProf_2.add_subplot(grid[1, 0]), - figProf_2.add_subplot(grid[1, 1]), - figProf_2.add_subplot(grid[0, 2]), - figProf_2.add_subplot(grid[1, 2]), - ] - grid = plt.GridSpec(3, 4, hspace=0.3, wspace=0.5) - ax00c = figProf_3.add_subplot(grid[0, 0]) - axsProf_3 = [ - ax00c, - figProf_3.add_subplot(grid[1, 0], sharex=ax00c), - figProf_3.add_subplot(grid[2, 0], sharex=ax00c), - figProf_3.add_subplot(grid[0, 1], sharex=ax00c), - figProf_3.add_subplot(grid[1, 1], sharex=ax00c), - figProf_3.add_subplot(grid[2, 1], sharex=ax00c), - figProf_3.add_subplot(grid[0, 2], sharex=ax00c), - figProf_3.add_subplot(grid[1, 2], sharex=ax00c), - figProf_3.add_subplot(grid[2, 2], sharex=ax00c), - figProf_3.add_subplot(grid[0, 3], sharex=ax00c), - figProf_3.add_subplot(grid[1, 3], sharex=ax00c), - figProf_3.add_subplot(grid[2, 3], sharex=ax00c), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsProf_4 = [ - figProf_4.add_subplot(grid[0, 0]), - figProf_4.add_subplot(grid[0, 1]), - figProf_4.add_subplot(grid[0, 2]), - figProf_4.add_subplot(grid[1, 0]), - figProf_4.add_subplot(grid[1, 1]), - figProf_4.add_subplot(grid[1, 2]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - - axsProf_5 = [ - figProf_5.add_subplot(grid[0, 0]), - figProf_5.add_subplot(grid[1, 0]), - figProf_5.add_subplot(grid[0, 1]), - figProf_5.add_subplot(grid[0, 2]), - figProf_5.add_subplot(grid[1, 1]), - figProf_5.add_subplot(grid[1, 2]), - ] - - grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) - axsProf_6 = [ - figProf_6.add_subplot(grid[0, 0]), - figProf_6.add_subplot(grid[:, 1]), - figProf_6.add_subplot(grid[0, 2]), - figProf_6.add_subplot(grid[1, 0]), - figProf_6.add_subplot(grid[1, 2]), - figProf_6.add_subplot(grid[0, 3]), - figProf_6.add_subplot(grid[1, 3]), - ] - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - axsProf_7 = [ - figProf_7.add_subplot(grid[0, 0]), - figProf_7.add_subplot(grid[0, 1]), - figProf_7.add_subplot(grid[0, 2]), - figProf_7.add_subplot(grid[1, 0]), - figProf_7.add_subplot(grid[1, 1]), - figProf_7.add_subplot(grid[1, 2]), - ] - + axsProf_1, axsProf_2, axsProf_3, axsProf_4, axsProf_5, axsProf_6, axsProf_7 = state_plotting.add_axes([figProf_1, figProf_2, figProf_3, figProf_4, figProf_5, figProf_6, figProf_7]) + + grid = plt.GridSpec(2, 4, hspace=0.2, wspace=0.6) if plotNormalizations: @@ -3867,19 +3790,13 @@ def changeANDwrite_TGLF( modInputTGLF[rho] = inputTGLF_rho ns_max.append(inputs[rho].plasma["NS"]) - + # Convert back to a string because that's how runTGLFproduction operates inputFileTGLF = inputToVariable(FolderTGLF, rhos) if (np.diff(ns_max) > 0).any(): - print( - "> Each radial location has its own number of species... probably because of removal of fast or low density...", - typeMsg="w", - ) - print( - "\t * Reading of TGLF results will fail... consider doing something before launching run", - typeMsg="q", - ) + print("> Each radial location has its own number of species... probably because of removal of fast or low density...",typeMsg="w") + print("\t * Reading of TGLF results will fail... consider doing something before launching run",typeMsg="q") return inputFileTGLF, modInputTGLF @@ -4164,6 +4081,8 @@ def write_state(self, file=None): for ikey in self.plasma: if ikey == "NS": var = np.min([self.plasma[ikey], maxSpeciesTGLF]) + if var < self.plasma[ikey]: + print(f"\t- Maximum number of species in TGLF reached, not considering after {maxSpeciesTGLF} species",typeMsg="w",) else: var = self.plasma[ikey] f.write(f"{ikey.ljust(23)} = {var}\n") diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 34a6e908..9ec6e86e 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -951,7 +951,7 @@ def runTGLF( fileTGLF = folderTGLF_this / "input.tglf" with open(fileTGLF, "w") as f: f.write(tglf_executor[subFolderTGLF][rho]["inputs"]) - + # --------------------------------------------- # Prepare command # --------------------------------------------- diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 4aad93db..fd3f7c97 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -3,7 +3,6 @@ import csv import numpy as np import matplotlib.pyplot as plt -from collections import OrderedDict from mitim_tools.misc_tools import GRAPHICStools, MATHtools, PLASMAtools, IOtools from mitim_modules.powertorch.utils import CALCtools from mitim_tools.gacode_tools import NEOtools @@ -1017,6 +1016,9 @@ def tglf_plasma(self): def deriv_gacode(y): return grad(self.derived["r"],y).cpu().numpy() + + self.derived["sign_it"] = - np.sign(self.profiles["current(MA)"][-1]) + self.derived["sign_bt"] = - np.sign(self.profiles["bcentr(T)"][-1]) self.derived["tite_all"] = self.profiles["ti(keV)"] / self.profiles["te(keV)"][:, np.newaxis] @@ -1037,7 +1039,7 @@ def deriv_gacode(y): self.derived["mi_ref"], self.derived["B_unit"]) - self.derived['pprime'] = 1E-7 * self.profiles["q(-)"]*self.derived['a']**2/self.derived["r"]/self.derived["B_unit"]**2*deriv_gacode(self.profiles["ptot(Pa)"]) + self.derived['pprime'] = 1E-7 * abs(self.profiles["q(-)"])*self.derived['a']**2/self.derived["r"]/self.derived["B_unit"]**2*deriv_gacode(self.profiles["ptot(Pa)"]) self.derived['pprime'][0] = 0.0 self.derived['drmin/dr'] = deriv_gacode(self.derived["r"]) @@ -1063,11 +1065,11 @@ def deriv_gacode(y): w0p = deriv_gacode(self.profiles["w0(rad/s)"]) gamma_p0 = -self.profiles["rmaj(m)"]*w0p - gamma_eb0 = -deriv_gacode(self.profiles["w0(rad/s)"]) * self.derived["r"]/self.profiles["q(-)"] + gamma_eb0 = -deriv_gacode(self.profiles["w0(rad/s)"]) * self.derived["r"]/ np.abs(self.profiles["q(-)"]) - self.derived['vexb_shear'] = gamma_eb0 * self.derived["a"]/self.derived['c_s'] - self.derived['vpar_shear'] = gamma_p0 * self.derived["a"]/self.derived['c_s'] - self.derived['vpar'] = self.profiles["rmaj(m)"]*self.profiles["w0(rad/s)"]/self.derived['c_s'] + self.derived['vexb_shear'] = -self.derived['sign_it'] * gamma_eb0 * self.derived["a"]/self.derived['c_s'] + self.derived['vpar_shear'] = -self.derived['sign_it'] * gamma_p0 * self.derived["a"]/self.derived['c_s'] + self.derived['vpar'] = -self.derived['sign_it'] * self.profiles["rmaj(m)"]*self.profiles["w0(rad/s)"]/self.derived['c_s'] def calculateMass(self): self.derived["mbg"] = 0.0 @@ -2147,6 +2149,8 @@ def csv(self, file="input.gacode.xlsx"): def to_tglf(self, rhos=[0.5], TGLFsettings=1): + max_species_tglf = 6 # TGLF only accepts up to 6 species + # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function @@ -2192,8 +2196,11 @@ def interpolator(y): 'VNS_SHEAR': 0.0, 'VTS_SHEAR': 0.0}, } + + if len(self.Species) > max_species_tglf-1: + print(f"\t- Warning: TGLF only accepts {max_species_tglf} species, but there are {len(self.Species)} ions pecies in the GACODE input. The first {max_species_tglf-1} will be used.", typeMsg="w") - for i in range(len(self.Species)): + for i in range(min(len(self.Species), max_species_tglf-1)): species[i+2] = { 'ZS': self.Species[i]['Z'], 'MASS': self.Species[i]['A']/mass_ref, @@ -2213,8 +2220,8 @@ def interpolator(y): plasma = { 'NS': len(species), - 'SIGN_BT': -1.0, - 'SIGN_IT': -1.0, + 'SIGN_BT': self.derived['sign_bt'], + 'SIGN_IT': self.derived['sign_it'], 'VEXB': 0.0, 'VEXB_SHEAR': interpolator(self.derived['vexb_shear']), 'BETAE': interpolator(self.derived['betae']), @@ -2234,7 +2241,7 @@ def interpolator(y): 'DRMINDX_LOC': np.ones(self.profiles["rho(-)"].shape), # Force 1.0 instead of self.derived['drmin/dr'] because of numerical issues in TGLF 'DRMAJDX_LOC': self.derived['dRmaj/dr'], 'DZMAJDX_LOC': self.derived['dZmaj/dr'], - 'Q_LOC': self.profiles["q(-)"], + 'Q_LOC': np.abs(self.profiles["q(-)"]), 'KAPPA_LOC': self.profiles["kappa(-)"], 'S_KAPPA_LOC': self.derived['s_kappa'], 'DELTA_LOC': self.profiles["delta(-)"], From 08f0b5ad0e849b368b890acec3a6a97d42898242 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 29 Jul 2025 16:50:16 +0200 Subject: [PATCH 113/385] Grab deuterium mass from same source --- src/mitim_modules/powertorch/STATEtools.py | 4 +++- src/mitim_tools/gacode_tools/TGLFtools.py | 8 ++++---- src/mitim_tools/misc_tools/PLASMAtools.py | 12 ++++++++---- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 2b3e5424..99189092 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -15,6 +15,8 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +from mitim_tools.misc_tools.PLASMAtools import md_u + # ------------------------------------------------------------------ # POWERSTATE Class # ------------------------------------------------------------------ @@ -623,7 +625,7 @@ def _update_plasma_var(var_key, clamp_min=None, clamp_max=None): # Toolset for calculation # ------------------------------------------------------------------ - def calculateProfileFunctions(self, calculateRotationQuantities=True, mref=2.01355): + def calculateProfileFunctions(self, calculateRotationQuantities=True, mref=md_u): """ Update the normalizations of the current state Notes: diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index a50ca675..6b5166bc 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -24,7 +24,7 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -mi_D = 2.01355 +from mitim_tools.misc_tools.PLASMAtools import md_u MAX_TGLF_SPECIES = 6 @@ -345,7 +345,7 @@ def prep( print("> Setting up normalizations") print("\t- Using mass of deuterium to unnormalize TGLF (not necesarily the first ion)",typeMsg="i") - self.tgyro.profiles.derive_quantities(mi_ref=mi_D) + self.tgyro.profiles.derive_quantities(mi_ref=md_u) self.NormalizationSets, cdf = NORMtools.normalizations( self.tgyro.profiles, @@ -410,7 +410,7 @@ def prep_direct_tglf( self.profiles = self.tgyro.profiles - self.profiles.derive_quantities(mi_ref=mi_D) + self.profiles.derive_quantities(mi_ref=md_u) self.profiles.correct(options={'recompute_ptot':recalculatePTOT,'removeFast':onlyThermal_TGYRO}) @@ -460,7 +460,7 @@ def prep_direct_tglf( print("> Setting up normalizations") print("\t- Using mass of deuterium to normalize things (not necesarily the first ion)",typeMsg="w",) - self.profiles.derive_quantities(mi_ref=mi_D) + self.profiles.derive_quantities(mi_ref=md_u) self.NormalizationSets, cdf = NORMtools.normalizations( self.profiles, diff --git a/src/mitim_tools/misc_tools/PLASMAtools.py b/src/mitim_tools/misc_tools/PLASMAtools.py index 6ef18d70..b794eecd 100644 --- a/src/mitim_tools/misc_tools/PLASMAtools.py +++ b/src/mitim_tools/misc_tools/PLASMAtools.py @@ -42,6 +42,7 @@ md = 3.34358e-24 me_u = 5.4488741e-04 # as in input.gacode +md_u = md*1E-3 / u factor_convection = 3 / 2 # IMPORTANT @@ -282,15 +283,17 @@ def calculatePressure(Te, Ti, ne, ni): - It only works if the vectors contain the entire plasma (i.e. roa[-1]=1.0), otherwise it will miss that contribution. """ - p, peT, piT = [], [], [] + p, peT, piT, piTall = [], [], [], [] for it in range(Te.shape[0]): pe = (Te[it, :] * 1e3 * e_J) * (ne[it, :] * 1e20) * 1e-6 # MPa + piall = (Ti[it, :, :] * 1e3 * e_J) * (ni[it, :, :] * 1e20) * 1e-6 # MPa pi = np.zeros(Te.shape[1]) for i in range(ni.shape[1]): - pi += (Ti[it, i, :] * 1e3 * e_J) * (ni[it, i, :] * 1e20) * 1e-6 # MPa + pi += piall[i, :] # Sum over all ions peT.append(pe) piT.append(pi) + piTall.append(piall) # Total pressure press = pe + pi @@ -299,8 +302,9 @@ def calculatePressure(Te, Ti, ne, ni): p = np.array(p) pe = np.array(peT) pi = np.array(piT) + pi_all = np.array(piTall) - return p, pe, pi + return p, pe, pi, pi_all def calculateVolumeAverage(rmin, var, dVdr): @@ -338,7 +342,7 @@ def calculateContent(rmin, Te, Ti, ne, ni, dVdr): """ - p, pe, pi = calculatePressure(Te, Ti, ne, ni) + p, pe, pi, _ = calculatePressure(Te, Ti, ne, ni) We, Wi, Ne, Ni = [], [], [], [] for it in range(rmin.shape[0]): From 3650cde1797b24c44add735a1da5130cedec0d4e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 29 Jul 2025 16:50:51 +0200 Subject: [PATCH 114/385] refurbished to_tglf method --- .../plasmastate_tools/MITIMstate.py | 246 ++++++++++-------- 1 file changed, 139 insertions(+), 107 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index fd3f7c97..3242f1f6 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -14,6 +14,8 @@ from mitim_tools import __version__ from IPython import embed +from mitim_tools.misc_tools.PLASMAtools import md_u + def ensure_variables_existence(self): # --------------------------------------------------------------------------- # Determine minimal set of variables that should be present in the profiles @@ -126,6 +128,11 @@ def scratch(cls, profiles, label_header='', **kwargs_process): @IOtools.hook_method(before=ensure_variables_existence) def derive_quantities_base(self, mi_ref=None, derive_quantities=True, rederiveGeometry=True): + # Make sure the profiles have the required dimensions + if len(self.profiles["ni(10^19/m^3)"].shape) == 1: + self.profiles["ni(10^19/m^3)"] = self.profiles["ni(10^19/m^3)"].reshape(-1, 1) + self.profiles["ti(keV)"] = self.profiles["ti(keV)"].reshape(-1, 1) + # ------------------------------------- self.readSpecies() self.mi_first = self.Species[0]["A"] @@ -138,9 +145,9 @@ def derive_quantities_base(self, mi_ref=None, derive_quantities=True, rederiveGe if mi_ref is not None: self.derived["mi_ref"] = mi_ref - print(f"\t* Reference mass ({self.derived['mi_ref']:.2f}) to use was forced by class initialization",typeMsg="w") + print(f"\t* Reference mass ({self.derived['mi_ref']}) to use was forced by class initialization",typeMsg="w") else: - self.derived["mi_ref"] = 2.0 #self.mi_first + self.derived["mi_ref"] = md_u #2.0 #md_u #self.mi_first print(f"\t* Reference mass ({self.derived['mi_ref']}) from Deuterium, as convention in gacode",typeMsg="i") # Useful to have gradients in the basic ---------------------------------------------------------- @@ -660,11 +667,20 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["Prad_line"] = self.derived["qrad_line_MW"][-1] self.derived["Psol"] = self.derived["qHeat"] - self.derived["Prad"] + self.derived["Ti_thr"] = [] self.derived["ni_thr"] = [] for sp in range(len(self.Species)): if self.Species[sp]["S"] == "therm": self.derived["ni_thr"].append(self.profiles["ni(10^19/m^3)"][:, sp]) + self.derived["Ti_thr"].append(self.profiles["ti(keV)"][:, sp]) + self.derived["ni_thr"] = np.transpose(self.derived["ni_thr"]) + self.derived["Ti_thr"] = np.transpose(np.array(self.derived["Ti_thr"])) + + if len(self.derived["ni_thr"].shape) == 1: + self.derived["ni_thr"] = self.derived["ni_thr"].reshape(-1, 1) + self.derived["Ti_thr"] = self.derived["Ti_thr"].reshape(-1, 1) + self.derived["ni_thrAll"] = self.derived["ni_thr"].sum(axis=1) self.derived["ni_All"] = self.profiles["ni(10^19/m^3)"].sum(axis=1) @@ -674,25 +690,30 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["ptot_manual"], self.derived["pe"], self.derived["pi"], + self.derived["pi_all"], ) = PLASMAtools.calculatePressure( np.expand_dims(self.profiles["te(keV)"], 0), np.expand_dims(np.transpose(self.profiles["ti(keV)"]), 0), np.expand_dims(self.profiles["ne(10^19/m^3)"] * 0.1, 0), np.expand_dims(np.transpose(self.profiles["ni(10^19/m^3)"] * 0.1), 0), ) - self.derived["ptot_manual"], self.derived["pe"], self.derived["pi"] = ( - self.derived["ptot_manual"][0], - self.derived["pe"][0], - self.derived["pi"][0], + self.derived["ptot_manual"], self.derived["pe"], self.derived["pi"], self.derived["pi_all"] = ( + self.derived["ptot_manual"][0,...], + self.derived["pe"][0,...], + self.derived["pi"][0,...], + self.derived["pi_all"][0,...], ) + self.derived['pi_all'] = np.transpose(self.derived['pi_all']) # to have the same shape as ni_thr + ( self.derived["pthr_manual"], _, self.derived["pi_thr"], + _, ) = PLASMAtools.calculatePressure( np.expand_dims(self.profiles["te(keV)"], 0), - np.expand_dims(np.transpose(self.profiles["ti(keV)"]), 0), + np.expand_dims(np.transpose(self.derived["Ti_thr"]), 0), np.expand_dims(self.profiles["ne(10^19/m^3)"] * 0.1, 0), np.expand_dims(np.transpose(self.derived["ni_thr"] * 0.1), 0), ) @@ -701,6 +722,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["pi_thr"][0], ) + # ------- # Content # ------- @@ -713,7 +735,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): ) = PLASMAtools.calculateContent( np.expand_dims(r, 0), np.expand_dims(self.profiles["te(keV)"], 0), - np.expand_dims(np.transpose(self.profiles["ti(keV)"]), 0), + np.expand_dims(np.transpose(self.derived["Ti_thr"]), 0), np.expand_dims(self.profiles["ne(10^19/m^3)"] * 0.1, 0), np.expand_dims(np.transpose(self.derived["ni_thr"] * 0.1), 0), np.expand_dims(volp, 0), @@ -931,7 +953,8 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["fG"] = self.derived["ne_vol20"] / nG self.derived["fG_x"] = self.profiles["ne(10^19/m^3)"]* 0.1 / nG - self.derived["tite"] = self.profiles["ti(keV)"][:, 0] / self.profiles["te(keV)"] + self.derived["tite_all"] = self.profiles["ti(keV)"] / self.profiles["te(keV)"][:, np.newaxis] + self.derived["tite"] = self.derived["tite_all"][:, 0] self.derived["tite_vol"] = self.derived["Ti_vol"] / self.derived["Te_vol"] self.derived["LH_nmin"] = PLASMAtools.LHthreshold_nmin( @@ -1006,71 +1029,6 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): # ~~~~ Estimate upstream density self.derived['ne_lcfs_estimate'] = self.derived["ne_vol20"] * 0.6 - # ------------------------------------------------------- - # TGLF-relevant quantities - # ------------------------------------------------------- - - self.tglf_plasma() - - def tglf_plasma(self): - - def deriv_gacode(y): - return grad(self.derived["r"],y).cpu().numpy() - - self.derived["sign_it"] = - np.sign(self.profiles["current(MA)"][-1]) - self.derived["sign_bt"] = - np.sign(self.profiles["bcentr(T)"][-1]) - - self.derived["tite_all"] = self.profiles["ti(keV)"] / self.profiles["te(keV)"][:, np.newaxis] - - self.derived['betae'] = PLASMAtools.betae( - self.profiles['te(keV)'], - self.profiles['ne(10^19/m^3)']*0.1, - self.derived["B_unit"]) - - self.derived['xnue'] = PLASMAtools.xnue( - torch.from_numpy(self.profiles['te(keV)']).to(torch.double), - torch.from_numpy(self.profiles['ne(10^19/m^3)']*0.1).to(torch.double), - self.derived["a"], - mref_u=self.derived["mi_ref"]).cpu().numpy() - - self.derived['debye'] = PLASMAtools.debye( - self.profiles['te(keV)'], - self.profiles['ne(10^19/m^3)']*0.1, - self.derived["mi_ref"], - self.derived["B_unit"]) - - self.derived['pprime'] = 1E-7 * abs(self.profiles["q(-)"])*self.derived['a']**2/self.derived["r"]/self.derived["B_unit"]**2*deriv_gacode(self.profiles["ptot(Pa)"]) - self.derived['pprime'][0] = 0.0 - - self.derived['drmin/dr'] = deriv_gacode(self.derived["r"]) - self.derived['dRmaj/dr'] = deriv_gacode(self.profiles["rmaj(m)"]) - self.derived['dZmaj/dr'] = deriv_gacode(self.profiles["zmag(m)"]) - - self.derived['s_kappa'] = self.derived["r"] / self.profiles["kappa(-)"] * deriv_gacode(self.profiles["kappa(-)"]) - self.derived['s_delta'] = self.derived["r"] * deriv_gacode(self.profiles["delta(-)"]) - self.derived['s_zeta'] = self.derived["r"] * deriv_gacode(self.profiles["zeta(-)"]) - - s = self.derived["r"] / self.profiles["q(-)"]*deriv_gacode(self.profiles["q(-)"]) - self.derived['s_q'] = np.concatenate([np.array([0.0]),(self.profiles["q(-)"][1:] / self.derived['roa'][1:])**2 * s[1:]]) # infinite in first location - - ''' - Rotations - -------------------------------------------------------- - From TGYRO/TGLF definitions - w0p = expro_w0p(:)/100.0 - f_rot(:) = w0p(:)/w0_norm - gamma_p0 = -r_maj(i_r)*f_rot(i_r)*w0_norm - gamma_eb0 = gamma_p0*r(i_r)/(q_abs*r_maj(i_r)) - ''' - - w0p = deriv_gacode(self.profiles["w0(rad/s)"]) - gamma_p0 = -self.profiles["rmaj(m)"]*w0p - gamma_eb0 = -deriv_gacode(self.profiles["w0(rad/s)"]) * self.derived["r"]/ np.abs(self.profiles["q(-)"]) - - self.derived['vexb_shear'] = -self.derived['sign_it'] * gamma_eb0 * self.derived["a"]/self.derived['c_s'] - self.derived['vpar_shear'] = -self.derived['sign_it'] * gamma_p0 * self.derived["a"]/self.derived['c_s'] - self.derived['vpar'] = -self.derived['sign_it'] * self.profiles["rmaj(m)"]*self.profiles["w0(rad/s)"]/self.derived['c_s'] - def calculateMass(self): self.derived["mbg"] = 0.0 self.derived["fmain"] = 0.0 @@ -1314,7 +1272,7 @@ def remove(self, ions_list): ions_list.sort() print("\t\t- Removing ions in positions (of ions order, no zero): ",ions_list,typeMsg="i",) - ions_list = [i - 1 for i in ions_list] + ions_list = [(i - 1 if i >-1 else i) for i in ions_list] fail = False @@ -2149,11 +2107,95 @@ def csv(self, file="input.gacode.xlsx"): def to_tglf(self, rhos=[0.5], TGLFsettings=1): - max_species_tglf = 6 # TGLF only accepts up to 6 species + # Derivate function + def deriv_gacode(y): + return grad(self.derived["r"],y).cpu().numpy() # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function + # Determine the number of species to use in TGLF + max_species_tglf = 6 # TGLF only accepts up to 6 species + if len(self.Species) > max_species_tglf-1: + print(f"\t- Warning: TGLF only accepts {max_species_tglf} species, but there are {len(self.Species)} ions pecies in the GACODE input. The first {max_species_tglf-1} will be used.", typeMsg="w") + tglf_ions_num = max_species_tglf - 1 + else: + tglf_ions_num = len(self.Species) + + # Determinte the mass reference + mass_ref = 2.0 # TODO: This is the only way to make it consistent with TGYRO (derivations with mD_u but mass in tglf with 2.0... https://github.com/gafusion/gacode/issues/398 + + # ----------------------------------------------------------------------- + # Derived profiles + # ----------------------------------------------------------------------- + + sign_it = -np.sign(self.profiles["current(MA)"][-1]) + sign_bt = -np.sign(self.profiles["bcentr(T)"][-1]) + + betae = PLASMAtools.betae( + self.profiles['te(keV)'], + self.profiles['ne(10^19/m^3)']*0.1, + self.derived["B_unit"] + ) + + xnue = PLASMAtools.xnue( + torch.from_numpy(self.profiles['te(keV)']).to(torch.double), + torch.from_numpy(self.profiles['ne(10^19/m^3)']*0.1).to(torch.double), + self.derived["a"], + mref_u=self.derived["mi_ref"] + ).cpu().numpy() + + debye = PLASMAtools.debye( + self.profiles['te(keV)'], + self.profiles['ne(10^19/m^3)']*0.1, + self.derived["mi_ref"], + self.derived["B_unit"] + ) + + s_kappa = self.derived["r"] / self.profiles["kappa(-)"] * deriv_gacode(self.profiles["kappa(-)"]) + s_delta = self.derived["r"] * deriv_gacode(self.profiles["delta(-)"]) + s_zeta = self.derived["r"] * deriv_gacode(self.profiles["zeta(-)"]) + + s_hat = self.derived["r"]*deriv_gacode( np.log(abs(self.profiles["q(-)"])) ) + s_q = (self.profiles["q(-)"] / self.derived['roa'])**2 * s_hat + s_q[0] = 0.0 # infinite in first location + + ''' + Total pressure + -------------------------------------------------------- + Recompute pprime with those species that belong to this run #TODO not exact? + ''' + + adpedr = - self.derived['pe'] * (self.derived['aLTe'] + self.derived['aLne']) + adpjdr = - self.derived['pi_all'][:,:tglf_ions_num] * (self.derived['aLTi'][:,:tglf_ions_num] + self.derived['aLni'][:,:tglf_ions_num]) + + dpdr = ( adpedr + adpjdr.sum(axis=-1)) / self.derived['a'] * 1E6 + + pprime = 1E-7 * abs(self.profiles["q(-)"])*self.derived['a']**2/self.derived["r"]/self.derived["B_unit"]**2*dpdr + pprime[0] = 0 # infinite in first location + + ''' + Rotations + -------------------------------------------------------- + From TGYRO/TGLF definitions + w0p = expro_w0p(:)/100.0 + f_rot(:) = w0p(:)/w0_norm + gamma_p0 = -r_maj(i_r)*f_rot(i_r)*w0_norm + gamma_eb0 = gamma_p0*r(i_r)/(q_abs*r_maj(i_r)) + ''' + + w0p = deriv_gacode(self.profiles["w0(rad/s)"]) + gamma_p0 = -self.profiles["rmaj(m)"]*w0p + gamma_eb0 = -deriv_gacode(self.profiles["w0(rad/s)"]) * self.derived["r"]/ np.abs(self.profiles["q(-)"]) + + vexb_shear = -sign_it * gamma_eb0 * self.derived["a"]/self.derived['c_s'] + vpar_shear = -sign_it * gamma_p0 * self.derived["a"]/self.derived['c_s'] + vpar = -sign_it * self.profiles["rmaj(m)"]*self.profiles["w0(rad/s)"]/self.derived['c_s'] + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Prepare the inputs for TGLF + # --------------------------------------------------------------------------------------------------------------------------------------- + inputsTGLF = {} for rho in rhos: @@ -2165,7 +2207,7 @@ def interpolator(y): return interpolation_function(rho, self.profiles['rho(-)'],y).item() TGLFinput, TGLFoptions, label = GACODEdefaults.addTGLFcontrol(TGLFsettings) - + # --------------------------------------------------------------------------------------------------------------------------------------- # Controls come from options # --------------------------------------------------------------------------------------------------------------------------------------- @@ -2176,13 +2218,6 @@ def interpolator(y): # Species come from profiles # --------------------------------------------------------------------------------------------------------------------------------------- - #mass_ref = self.derived["mi_ref"] - # input.gacode uses the deuterium mass as reference already (https://github.com/gafusion/gacode/issues/398), so this should be 2.0 - mass_ref = 2.0 - - if mass_ref != self.derived["mi_ref"]: - print(f"\t- Warning: the mass reference in the input.gacode is {self.derived['mi_ref']}, but TGLF expects {mass_ref}. This may lead to problems with the TGLF input file.", typeMsg="q") - species = { 1: { 'ZS': -1.0, @@ -2191,14 +2226,11 @@ def interpolator(y): 'RLTS': interpolator(self.derived['aLTe']), 'TAUS': 1.0, 'AS': 1.0, - 'VPAR': interpolator(self.derived['vpar']), - 'VPAR_SHEAR': interpolator(self.derived['vpar_shear']), + 'VPAR': interpolator(vpar), + 'VPAR_SHEAR': interpolator(vpar_shear), 'VNS_SHEAR': 0.0, 'VTS_SHEAR': 0.0}, } - - if len(self.Species) > max_species_tglf-1: - print(f"\t- Warning: TGLF only accepts {max_species_tglf} species, but there are {len(self.Species)} ions pecies in the GACODE input. The first {max_species_tglf-1} will be used.", typeMsg="w") for i in range(min(len(self.Species), max_species_tglf-1)): species[i+2] = { @@ -2208,8 +2240,8 @@ def interpolator(y): 'RLTS': interpolator(self.derived['aLTi'][:,0] if self.Species[i]['S'] == 'therm' else self.derived["aLTi"][:,i]), 'TAUS': interpolator(self.derived["tite_all"][:,i]), 'AS': interpolator(self.derived['fi'][:,i]), - 'VPAR': interpolator(self.derived['vpar']), - 'VPAR_SHEAR': interpolator(self.derived['vpar_shear']), + 'VPAR': interpolator(vpar), + 'VPAR_SHEAR': interpolator(vpar_shear), 'VNS_SHEAR': 0.0, 'VTS_SHEAR': 0.0 } @@ -2220,16 +2252,16 @@ def interpolator(y): plasma = { 'NS': len(species), - 'SIGN_BT': self.derived['sign_bt'], - 'SIGN_IT': self.derived['sign_it'], + 'SIGN_BT': sign_bt, + 'SIGN_IT': sign_it, 'VEXB': 0.0, - 'VEXB_SHEAR': interpolator(self.derived['vexb_shear']), - 'BETAE': interpolator(self.derived['betae']), - 'XNUE': interpolator(self.derived['xnue']), + 'VEXB_SHEAR': interpolator(vexb_shear), + 'BETAE': interpolator(betae), + 'XNUE': interpolator(xnue), 'ZEFF': interpolator(self.derived['Zeff']), - 'DEBYE': interpolator(self.derived['debye']), + 'DEBYE':interpolator(debye), } - + # --------------------------------------------------------------------------------------------------------------------------------------- # Geometry comes from profiles # --------------------------------------------------------------------------------------------------------------------------------------- @@ -2238,18 +2270,18 @@ def interpolator(y): 'RMIN_LOC': self.derived['roa'], 'RMAJ_LOC': self.derived['Rmajoa'], 'ZMAJ_LOC': self.derived["Zmagoa"], - 'DRMINDX_LOC': np.ones(self.profiles["rho(-)"].shape), # Force 1.0 instead of self.derived['drmin/dr'] because of numerical issues in TGLF - 'DRMAJDX_LOC': self.derived['dRmaj/dr'], - 'DZMAJDX_LOC': self.derived['dZmaj/dr'], + 'DRMINDX_LOC': np.ones(self.profiles["rho(-)"].shape), # Force 1.0 because of numerical issues in TGLF + 'DRMAJDX_LOC': deriv_gacode(self.profiles["rmaj(m)"]), + 'DZMAJDX_LOC': deriv_gacode(self.profiles["zmag(m)"]), 'Q_LOC': np.abs(self.profiles["q(-)"]), 'KAPPA_LOC': self.profiles["kappa(-)"], - 'S_KAPPA_LOC': self.derived['s_kappa'], + 'S_KAPPA_LOC': s_kappa, 'DELTA_LOC': self.profiles["delta(-)"], - 'S_DELTA_LOC': self.derived['s_delta'], + 'S_DELTA_LOC': s_delta, 'ZETA_LOC': self.profiles["zeta(-)"], - 'S_ZETA_LOC': self.derived['s_zeta'], - 'P_PRIME_LOC': self.derived['pprime'], - 'Q_PRIME_LOC': self.derived['s_q'], + 'S_ZETA_LOC': s_zeta, + 'Q_PRIME_LOC': s_q, + 'P_PRIME_LOC': pprime, } # Add MXH and derivatives (#TODO) @@ -2272,7 +2304,7 @@ def interpolator(y): geom['BETA_LOC'] = 0.0 geom['KX0_LOC'] = 0.0 - + # --------------------------------------------------------------------------------------------------------------------------------------- # Merging # --------------------------------------------------------------------------------------------------------------------------------------- From daaa5724ec12d7f288faf35f82dca9dec853b494 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 29 Jul 2025 18:48:45 +0200 Subject: [PATCH 115/385] More complete vmec interpreation --- src/mitim_tools/gacode_tools/PROFILEStools.py | 39 ++--- .../plasmastate_tools/MITIMstate.py | 89 +++++----- .../plasmastate_tools/utils/VMECtools.py | 64 +++++--- .../plasmastate_tools/utils/state_plotting.py | 152 ++++++++++-------- 4 files changed, 198 insertions(+), 146 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 0ec1fd6b..9fdaa432 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -263,37 +263,32 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): lines = GRAPHICStools.listLS() ax = ax00c - varL = "cos Shape Params" - yl = 0 - cont = 0 + ax.plot(self.profiles["rho(-)"], self.derived['volp_geo'], color=color, lw=lw, label = extralab) + ax.set_xlabel('$\\rho_N$'); ax.set_xlim(0, 1) + ax.set_ylabel(f"$dV/dr$ ($m^3/[r]$)") + GRAPHICStools.addDenseAxis(ax) + + if legYN: + ax.legend(loc="best", fontsize=fs) + minShape = 1E-4 + + ax = ax01c + cont = 0 + yl = 0 for i, s in enumerate(self.shape_cos): if s is not None: valmax = np.abs(s).max() - if valmax > 1e-10: + if valmax > minShape: lab = f"c{i}" ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) cont += 1 yl = np.max([yl, valmax]) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) - - if legYN: - ax.legend(loc="best", fontsize=fs) - - ax = ax01c - varL = "sin Shape Params" - cont = 0 for i, s in enumerate(self.shape_sin): if s is not None: valmax = np.abs(s).max() - if valmax > 1e-10: + if valmax > minShape: lab = f"s{i}" ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) cont += 1 @@ -302,7 +297,7 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): ax.set_xlim([0, 1]) ax.set_xlabel("$\\rho$") - ax.set_ylabel(varL) + ax.set_ylabel(f"Shape Parameters (>{minShape})") if legYN: ax.legend(loc="best", fontsize=fs) @@ -337,13 +332,13 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): ax = ax11c - var = self.profiles["rmin(m)"] + var = self.derived['r'] ax.plot(rho, var, "-", lw=lw, c=color) ax.set_xlim([0, 1]) ax.set_xlabel("$\\rho$") ax.set_ylim(bottom=0) - ax.set_ylabel("$r_{min}$") + ax.set_ylabel("Effective $r$") GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax, bottomy=0) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 3242f1f6..ffa35eb6 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1029,6 +1029,37 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): # ~~~~ Estimate upstream density self.derived['ne_lcfs_estimate'] = self.derived["ne_vol20"] * 0.6 + # ------------------------------------------------------- + # Transport parameters + # ------------------------------------------------------- + + self.derived['betae'] = PLASMAtools.betae( + self.profiles['te(keV)'], + self.profiles['ne(10^19/m^3)']*0.1, + self.derived["B_unit"] + ) + + self.derived['xnue'] = PLASMAtools.xnue( + torch.from_numpy(self.profiles['te(keV)']).to(torch.double), + torch.from_numpy(self.profiles['ne(10^19/m^3)']*0.1).to(torch.double), + self.derived["a"], + mref_u=self.derived["mi_ref"] + ).cpu().numpy() + + self.derived['debye'] = PLASMAtools.debye( + self.profiles['te(keV)'], + self.profiles['ne(10^19/m^3)']*0.1, + self.derived["mi_ref"], + self.derived["B_unit"] + ) + s_hat = self.derived["r"]*self._deriv_gacode( np.log(abs(self.profiles["q(-)"])) ) + self.derived['s_q'] = (self.profiles["q(-)"] / self.derived['roa'])**2 * s_hat + self.derived['s_q'][0] = 0.0 # infinite in first location + + # Derivate function + def _deriv_gacode(self,y): + return grad(self.derived["r"],y).cpu().numpy() + def calculateMass(self): self.derived["mbg"] = 0.0 self.derived["fmain"] = 0.0 @@ -1108,6 +1139,11 @@ def deriveContentByVolumes(self, rhos=[0.5], impurityPosition=3): def printInfo(self, label="", reDeriveIfNotFound=True): + Prad_ratio = self.derived['Prad'] / self.derived['qHeat'] + Prad_ratio_brem = self.derived['Prad_brem']/self.derived['Prad'] + Prad_ratio_line = self.derived['Prad_line']/self.derived['Prad'] + Prad_ratio_sync = self.derived['Prad_sync']/self.derived['Prad'] + try: ImpurityText = "" for i in range(len(self.Species)): @@ -1129,7 +1165,7 @@ def printInfo(self, label="", reDeriveIfNotFound=True): print(f"\tnu_Ti = {self.derived['Ti_peaking']:.2f}") print(f"\tp_vol = {self.derived['ptot_manual_vol']:.2f} MPa ({self.derived['pfast_fraction']*100.0:.1f}% fast)") print(f"\tBetaN = {self.derived['BetaN']:.3f} (BetaN w/B0 = {self.derived['BetaN_engineering']:.3f})") - print(f"\tPrad = {self.derived['Prad']:.1f}MW ({self.derived['Prad'] / self.derived['qHeat'] * 100.0:.1f}% of total) ({self.derived['Prad_brem']/self.derived['Prad'] * 100.0:.1f}% brem, {self.derived['Prad_line']/self.derived['Prad'] * 100.0:.1f}% line, {self.derived['Prad_sync']/self.derived['Prad'] * 100.0:.1f}% sync)") + print(f"\tPrad = {self.derived['Prad']:.1f}MW ({Prad_ratio*100.0:.1f}% of total) ({Prad_ratio_brem*100.0:.1f}% brem, {Prad_ratio_line*100.0:.1f}% line, {Prad_ratio_sync*100.0:.1f}% sync)") print("\tPsol = {0:.1f}MW (fLH = {1:.2f})".format(self.derived["Psol"], self.derived["LHratio"])) print("Operational point ( [,] = [{0:.2f},{1:.2f}] ) and species:".format(self.derived["ne_vol20"], self.derived["Te_vol"])) print("\t = {0:.2f} keV (/ = {1:.2f}, Ti0/Te0 = {2:.2f})".format(self.derived["Ti_vol"],self.derived["tite_vol"],self.derived["tite"][0],)) @@ -2107,10 +2143,6 @@ def csv(self, file="input.gacode.xlsx"): def to_tglf(self, rhos=[0.5], TGLFsettings=1): - # Derivate function - def deriv_gacode(y): - return grad(self.derived["r"],y).cpu().numpy() - # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function @@ -2132,33 +2164,10 @@ def deriv_gacode(y): sign_it = -np.sign(self.profiles["current(MA)"][-1]) sign_bt = -np.sign(self.profiles["bcentr(T)"][-1]) - betae = PLASMAtools.betae( - self.profiles['te(keV)'], - self.profiles['ne(10^19/m^3)']*0.1, - self.derived["B_unit"] - ) - - xnue = PLASMAtools.xnue( - torch.from_numpy(self.profiles['te(keV)']).to(torch.double), - torch.from_numpy(self.profiles['ne(10^19/m^3)']*0.1).to(torch.double), - self.derived["a"], - mref_u=self.derived["mi_ref"] - ).cpu().numpy() - - debye = PLASMAtools.debye( - self.profiles['te(keV)'], - self.profiles['ne(10^19/m^3)']*0.1, - self.derived["mi_ref"], - self.derived["B_unit"] - ) - - s_kappa = self.derived["r"] / self.profiles["kappa(-)"] * deriv_gacode(self.profiles["kappa(-)"]) - s_delta = self.derived["r"] * deriv_gacode(self.profiles["delta(-)"]) - s_zeta = self.derived["r"] * deriv_gacode(self.profiles["zeta(-)"]) + s_kappa = self.derived["r"] / self.profiles["kappa(-)"] * self._deriv_gacode(self.profiles["kappa(-)"]) + s_delta = self.derived["r"] * self._deriv_gacode(self.profiles["delta(-)"]) + s_zeta = self.derived["r"] * self._deriv_gacode(self.profiles["zeta(-)"]) - s_hat = self.derived["r"]*deriv_gacode( np.log(abs(self.profiles["q(-)"])) ) - s_q = (self.profiles["q(-)"] / self.derived['roa'])**2 * s_hat - s_q[0] = 0.0 # infinite in first location ''' Total pressure @@ -2184,9 +2193,9 @@ def deriv_gacode(y): gamma_eb0 = gamma_p0*r(i_r)/(q_abs*r_maj(i_r)) ''' - w0p = deriv_gacode(self.profiles["w0(rad/s)"]) + w0p = self._deriv_gacode(self.profiles["w0(rad/s)"]) gamma_p0 = -self.profiles["rmaj(m)"]*w0p - gamma_eb0 = -deriv_gacode(self.profiles["w0(rad/s)"]) * self.derived["r"]/ np.abs(self.profiles["q(-)"]) + gamma_eb0 = -self._deriv_gacode(self.profiles["w0(rad/s)"]) * self.derived["r"]/ np.abs(self.profiles["q(-)"]) vexb_shear = -sign_it * gamma_eb0 * self.derived["a"]/self.derived['c_s'] vpar_shear = -sign_it * gamma_p0 * self.derived["a"]/self.derived['c_s'] @@ -2256,12 +2265,12 @@ def interpolator(y): 'SIGN_IT': sign_it, 'VEXB': 0.0, 'VEXB_SHEAR': interpolator(vexb_shear), - 'BETAE': interpolator(betae), - 'XNUE': interpolator(xnue), + 'XNUE': interpolator(self.derived['xnue']), 'ZEFF': interpolator(self.derived['Zeff']), - 'DEBYE':interpolator(debye), + 'DEBYE': interpolator(self.derived['debye']), + 'BETAE': interpolator(self.derived['betae']), } - + # --------------------------------------------------------------------------------------------------------------------------------------- # Geometry comes from profiles # --------------------------------------------------------------------------------------------------------------------------------------- @@ -2271,8 +2280,8 @@ def interpolator(y): 'RMAJ_LOC': self.derived['Rmajoa'], 'ZMAJ_LOC': self.derived["Zmagoa"], 'DRMINDX_LOC': np.ones(self.profiles["rho(-)"].shape), # Force 1.0 because of numerical issues in TGLF - 'DRMAJDX_LOC': deriv_gacode(self.profiles["rmaj(m)"]), - 'DZMAJDX_LOC': deriv_gacode(self.profiles["zmag(m)"]), + 'DRMAJDX_LOC': self._deriv_gacode(self.profiles["rmaj(m)"]), + 'DZMAJDX_LOC': self._deriv_gacode(self.profiles["zmag(m)"]), 'Q_LOC': np.abs(self.profiles["q(-)"]), 'KAPPA_LOC': self.profiles["kappa(-)"], 'S_KAPPA_LOC': s_kappa, @@ -2280,7 +2289,7 @@ def interpolator(y): 'S_DELTA_LOC': s_delta, 'ZETA_LOC': self.profiles["zeta(-)"], 'S_ZETA_LOC': s_zeta, - 'Q_PRIME_LOC': s_q, + 'Q_PRIME_LOC': self.derived['s_q'], 'P_PRIME_LOC': pprime, } diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index 75e7aaf5..b94c89fb 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -62,12 +62,13 @@ def _read_vmec(self): self.profiles['rcentr(m)'] = np.array([self.wout.Rmajor_p]) self.profiles['bcentr(T)'] = np.array([self.wout.rbtor/self.wout.Rmajor_p]) self.profiles["current(MA)"] = np.array([0.0]) - - self.profiles["torfluxa(Wb/radian)"] = [self.wout.phipf[-1]] - + + self.profiles["torfluxa(Wb/radian)"] = np.array([self.wout.phipf[-1]]) + # Produce variables self.profiles["rho(-)"] = (self.wout.phi/self.wout.phi[-1])**0.5 #np.linspace(0, 1, self.wout.ns)**0.5 - self.profiles["presf"] = self.wout.presf + self.profiles["ptot(Pa)"] = self.wout.presf + #self.profiles["q(-)"] = self.wout.q_factor #self.profiles["polflux(Wb/radian)"] = self.wout.chi @@ -99,39 +100,62 @@ def derive_quantities(self, **kwargs): def derive_geometry(self, **kwargs): - rho = np.linspace(0, 1, self.wout.ns) + r = self.derived["r"] - ds = rho[1] - rho[0] - half_grid_rho = rho - ds / 2 - - d_volume_d_rho = ( + half_grid_r = r - (r[1] - r[0]) / 2 + d_volume_d_r = ( (2 * np.pi) ** 2 * np.array(self.wout.vp) * 2 - * np.sqrt(half_grid_rho) - ) + * np.sqrt(half_grid_r) + ) - #self.derived["B_unit"] = self.profiles["torfluxa(Wb/radian)"] / (np.pi * self.wout.Aminor_p**2) - - self.derived["volp_geo"] = self.wout.vp #d_volume_d_rho - self.derived["volp_geo"][0] = 0.0 + self.derived["B_ref"] = np.ones(r.shape) * self.profiles["torfluxa(Wb/radian)"][-1] / (np.pi * self.wout.Aminor_p**2) + + self.derived["volp_geo"] = d_volume_d_r + self.derived["volp_geo"][0] = 1E-9 self.derived["kappa_a"] = 0.0 self.derived["kappa95"] = 0.0 self.derived["delta95"] = 0.0 self.derived["kappa995"] = 0.0 self.derived["delta995"] = 0.0 - self.derived["R_LF"] = np.zeros(self.profiles["rho(-)"].shape) + self.derived["R_LF"] = np.zeros(r.shape) - self.derived["bp2_exp"] = np.zeros(self.profiles["rho(-)"].shape) - self.derived["bt2_exp"] = np.zeros(self.profiles["rho(-)"].shape) - self.derived["bp2_geo"] = np.zeros(self.profiles["rho(-)"].shape) - self.derived["bt2_geo"] = np.zeros(self.profiles["rho(-)"].shape) + self.derived["bp2_exp"] = np.zeros(r.shape) + self.derived["bt2_exp"] = np.zeros(r.shape) + self.derived["bp2_geo"] = np.zeros(r.shape) + self.derived["bt2_geo"] = np.zeros(r.shape) def plot_geometry(self, axs, color="b", legYN=True, extralab="", lw=1, fs=6): [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,axs_3d,axs_2d] = axs + rho = self.profiles["rho(-)"] + + ax = ax00c + ax.plot(rho, self.derived['volp_geo'], color=color, lw=lw, label = extralab) + ax.set_xlabel('$\\rho$'); ax.set_xlim(0, 1) + ax.set_ylabel(f"$dV/d\\rho$ ($m^3$)") + GRAPHICStools.addDenseAxis(ax) + + if legYN: + ax.legend(loc="best", fontsize=fs) + + ax = ax11c + + var = self.derived['r'] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylim(bottom=0) + ax.set_ylabel("Effective $r$") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + self.plot_plasma_boundary(ax=axs_3d, color=color) self.plot_state_flux_surfaces(ax=axs_2d, c=color) diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py index af8b1275..d8e89b19 100644 --- a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -37,7 +37,7 @@ def add_axes(figs): fig1.add_subplot(grid[2, 2]), ] - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) axsProf_2 = [ fig2.add_subplot(grid[0, 0]), fig2.add_subplot(grid[0, 1]), @@ -45,6 +45,8 @@ def add_axes(figs): fig2.add_subplot(grid[1, 1]), fig2.add_subplot(grid[0, 2]), fig2.add_subplot(grid[1, 2]), + fig2.add_subplot(grid[0, 3]), + fig2.add_subplot(grid[1, 3]), ] # GEOMETRY @@ -85,7 +87,7 @@ def add_axes(figs): fig5.add_subplot(grid[1, 2]), ] - grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) + grid = plt.GridSpec(2, 5, hspace=0.3, wspace=0.3) axsProf_6 = [ fig6.add_subplot(grid[0, 0]), fig6.add_subplot(grid[0, 1]), @@ -95,6 +97,8 @@ def add_axes(figs): fig6.add_subplot(grid[1, 2]), fig6.add_subplot(grid[0, 3]), fig6.add_subplot(grid[1, 3]), + fig6.add_subplot(grid[0, 4]), + fig6.add_subplot(grid[1, 4]), ] grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) axsImps = [ @@ -323,7 +327,7 @@ def plot_profiles(self, axs1, color="b", legYN=True, extralab="", lw=1, fs=6): def plot_powers(self, axs2, legYN=True, extralab="", color="b", lw=1, fs=6): - [ax00b, ax01b, ax10b, ax11b, ax20b, ax21b] = axs2 + [ax00b, ax01b, ax10b, ax11b, ax20b, ax21b, ax30b, ax31b] = axs2 rho = self.profiles["rho(-)"] @@ -379,63 +383,6 @@ def plot_powers(self, axs2, legYN=True, extralab="", color="b", lw=1, fs=6): GRAPHICStools.autoscale_y(ax) ax.set_title("Momentum Source Density") - ax = ax21b - ax.plot( - rho, self.derived["qe_MWm2"], lw=lw, ls="-", label=extralab + "qe", c=color - ) - ax.plot( - rho, self.derived["qi_MWm2"], lw=lw, ls="--", label=extralab + "qi", c=color - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("Heat Flux ($MW/m^2$)") - if legYN: - ax.legend(loc="lower left", fontsize=fs) - ax.set_title("Flux per unit area (gacode: P/V')") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax21b.twinx() - ax.plot( - rho, - self.derived["ge_10E20m2"], - lw=lw, - ls="-.", - label=extralab + "$\\Gamma_e$", - c=color, - ) - ax.set_ylabel("Particle Flux ($10^{20}/m^2/s$)") - if legYN: - ax.legend(loc="lower right", fontsize=fs) - GRAPHICStools.autoscale_y(ax, bottomy=0) - - ax = ax20b - varL = "$Q_{rad}$ ($MW/m^3$)" - if "qbrem(MW/m^3)" in self.profiles: - var = self.profiles["qbrem(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls="-", label=extralab + "brem", c=color) - if "qline(MW/m^3)" in self.profiles: - var = self.profiles["qline(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls="--", label=extralab + "line", c=color) - if "qsync(MW/m^3)" in self.profiles: - var = self.profiles["qsync(MW/m^3)"] - ax.plot(rho, var, lw=lw, ls=":", label=extralab + "sync", c=color) - - var = self.derived["qrad"] - ax.plot(rho, var, lw=lw * 1.5, ls="-", label=extralab + "Total", c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - # ax.set_ylim(bottom=0); - ax.set_ylabel(varL) - if legYN: - ax.legend(loc="best", fontsize=fs) - ax.set_title("Radiation Contributions") - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - ax = ax10b varL = "$MW/m^3$" cont = 0 @@ -492,6 +439,69 @@ def plot_powers(self, axs2, legYN=True, extralab="", color="b", lw=1, fs=6): GRAPHICStools.autoscale_y(ax) + ax = ax20b + varL = "$Q_{rad}$ ($MW/m^3$)" + if "qbrem(MW/m^3)" in self.profiles: + var = self.profiles["qbrem(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls="-", label=extralab + "brem", c=color) + if "qline(MW/m^3)" in self.profiles: + var = self.profiles["qline(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls="--", label=extralab + "line", c=color) + if "qsync(MW/m^3)" in self.profiles: + var = self.profiles["qsync(MW/m^3)"] + ax.plot(rho, var, lw=lw, ls=":", label=extralab + "sync", c=color) + + var = self.derived["qrad"] + ax.plot(rho, var, lw=lw * 1.5, ls="-", label=extralab + "Total", c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + # ax.set_ylim(bottom=0); + ax.set_ylabel(varL) + if legYN: + ax.legend(loc="best", fontsize=fs) + ax.set_title("Radiation Contributions") + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + + + ax = ax30b + ax.plot(rho, self.derived["qe_MWm2"], lw=lw, ls="-", label=extralab + "qe", c=color) + ax.plot(rho, self.derived["qi_MWm2"], lw=lw, ls="--", label=extralab + "qi", c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("Heat Flux ($MW/m^2$)") + if legYN: + ax.legend(loc="lower left", fontsize=fs) + ax.set_title("Heat flux per unit area (gacode: P/V')") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + ax = ax31b + ax.plot( + rho, + self.derived["ge_10E20m2"], + lw=lw, + ls="-.", + label=extralab + "$\\Gamma_e$", + c=color, + ) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("Particle Flux ($10^{20}/m^2/s$)") + if legYN: + ax.legend(loc="lower left", fontsize=fs) + ax.set_title("Particle Flux per unit area (gacode: P/V')") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + def plot_gradients( self, axs4, @@ -788,8 +798,7 @@ def plot_other(self, axs6, color="b", lw=1.0, extralab="", fs=6): ax.set_xlim([0, 1]) ax.axhline(y=1.0, ls="--", c="k", lw=1.0) GRAPHICStools.addDenseAxis(ax) - # GRAPHICStools.autoscale_y(ax,bottomy=0) - ax.set_ylim(bottom=0) + GRAPHICStools.autoscale_y(ax,bottomy=0)#ax.set_ylim(bottom=0) ax.legend(loc="best", fontsize=fs) # Currents @@ -805,7 +814,7 @@ def plot_other(self, axs6, color="b", lw=1.0, extralab="", fs=6): ax.set_xlim([0, 1]) ax.set_xlabel("$\\rho$") - ax.set_ylim(bottom=0) + GRAPHICStools.autoscale_y(ax,bottomy=0)#ax.set_ylim(bottom=0) ax.set_ylabel("J ($MA/m^2$)") GRAPHICStools.addDenseAxis(ax) @@ -836,6 +845,21 @@ def plot_other(self, axs6, color="b", lw=1.0, extralab="", fs=6): GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = axs6[8] + + ax.plot(rho, self.derived["B_unit"], "-", lw=lw, c=color, label=extralab + "$B_{unit}$") + ax.plot(rho, self.derived["B_ref"], "--", lw=lw, c=color, label=extralab + "$B_{ref}$") + ax.axhline(y=self.profiles["bcentr(T)"][0], lw=lw, ls=":", c=color, label=extralab + "$B_{centr}$") + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("$B$ (T)") + + GRAPHICStools.addDenseAxis(ax) + + ax.legend(loc="best", fontsize=fs) + + + def plot_flows(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): if axs is None: @@ -855,7 +879,7 @@ def plot_flows(self, axs=None, limits=None, ls="-", leg=True, showtexts=True): ax = axs[0] axT = axs[1] - roa = self.profiles["rmin(m)"] / self.profiles["rmin(m)"][-1] + roa = self.derived['roa'] Te = self.profiles["te(keV)"] ne = self.profiles["ne(10^19/m^3)"] * 1e-1 ni = self.profiles["ni(10^19/m^3)"] * 1e-1 From 0d994f986f0b6fb057a9c261d2e6524bdc2045ee Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 30 Jul 2025 10:11:45 +0200 Subject: [PATCH 116/385] Removed deprecated portals functions --- src/mitim_modules/portals/PORTALSmain.py | 2 - src/mitim_modules/portals/PORTALStools.py | 56 +------------------ .../portals/utils/PORTALSinit.py | 2 - 3 files changed, 1 insertion(+), 59 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index ed337d7a..99137e69 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -207,8 +207,6 @@ def __init__( "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) "useConvectiveFluxes": True, # If True, then convective flux for final metric (not fitting). If False, particle flux "includeFastInQi": False, # If True, and fast ions have been included, in seprateNEO, sum fast - "useDiffusivities": False, # If True, use [chi_e,chi_i,D] instead of [Qe,Qi,Gamma] - "useFluxRatios": False, # If True, fit to [Qi,Qe/Qi,Ge/Qi] "portals_transformation_variables": portals_transformation_variables, # Physics-informed parameters to fit surrogates "portals_transformation_variables_trace": portals_transformation_variables_trace, # Physics-informed parameters to fit surrogates for trace impurities "Qi_criterion_stable": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 34f83954..e38fdedd 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -169,12 +169,10 @@ def output_transform_portals(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, output, powerstate) - # --- Ratio of fluxes (quasilinear) - factorRat = ratioFactor(X, surrogate_parameters, output, powerstate) # --- Specific to output factorImp = ImpurityGammaTrick(X, surrogate_parameters, output, powerstate) - compounded = factorGB * factorRat * factorImp + compounded = factorGB * factorImp """ 3. Go back to the original batching system @@ -260,58 +258,6 @@ def ImpurityGammaTrick(x, surrogate_parameters, output, powerstate): return factor - -def ratioFactor(X, surrogate_parameters, output, powerstate): - """ - This defines the vector to divide by. - - THIS IS BROKEN RIGHT NOW - """ - - 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 - - return v - - def constructEvaluationProfiles(X, surrogate_parameters, recalculateTargets=False): """ Prepare powerstate for another evaluation with batches diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 756814ee..904ee4f5 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -344,8 +344,6 @@ def initializeProblem( "transformationOutputs": PORTALStools.output_transform_portals, "powerstate": portals_fun.powerstate, "applyImpurityGammaTrick": portals_fun.PORTALSparameters["applyImpurityGammaTrick"], - "useFluxRatios": portals_fun.PORTALSparameters["useFluxRatios"], - "useDiffusivities": portals_fun.PORTALSparameters["useDiffusivities"], "surrogate_transformation_variables_alltimes": Variables, "surrogate_transformation_variables_lasttime": copy.deepcopy(Variables[list(Variables.keys())[-1]]), "parameters_combined": {}, From e3dd5a029487389c31869c3f2f1b3f19672b87b3 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 30 Jul 2025 10:32:49 +0200 Subject: [PATCH 117/385] Improved naming of turb exchange --- regressions/portals_regressions.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 6 +-- src/mitim_modules/portals/PORTALStools.py | 52 ++++++++----------- .../portals/utils/PORTALSinit.py | 4 +- .../physics_models/transport_cgyro.py | 2 +- .../physics_models/transport_tglf.py | 8 +-- .../physics_models/transport_tgyro.py | 8 +-- .../powertorch/utils/TRANSPORTtools.py | 2 +- src/mitim_tools/gacode_tools/TGYROtools.py | 4 +- 9 files changed, 39 insertions(+), 49 deletions(-) diff --git a/regressions/portals_regressions.py b/regressions/portals_regressions.py index 64f32a9a..61d6bac4 100644 --- a/regressions/portals_regressions.py +++ b/regressions/portals_regressions.py @@ -151,5 +151,5 @@ def conditions_regressions(variables): # # Checks # conditions_regressions([ # [mitim_bo.optimization_data.data['QeTurb_1'][5],0.0713711320661], - # [mitim_runs[5]['powerstate'].plasma['PexchTurb'][0,3].item(),-0.0009466626542564001] + # [mitim_runs[5]['powerstate'].plasma['QieMWm3_tr_turb'][0,3].item(),-0.0009466626542564001] # ]) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 99137e69..c45abbd0 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -685,10 +685,10 @@ def map_powerstate_to_portals(powerstate, dictOFs): Turbulent Exchange ------------------ """ - if 'PexchTurb_1' in dictOFs: + if 'QieMWm3_tr_turb_1' in dictOFs: for i in range(powerstate.plasma["rho"].shape[1] - 1): - dictOFs[f"PexchTurb_{i+1}"]["value"] = powerstate.plasma["PexchTurb"][0, i+1] - dictOFs[f"PexchTurb_{i+1}"]["error"] = powerstate.plasma["PexchTurb_stds"][0, i+1] + dictOFs[f"QieMWm3_tr_turb_{i+1}"]["value"] = powerstate.plasma["QieMWm3_tr_turb"][0, i+1] + dictOFs[f"QieMWm3_tr_turb_{i+1}"]["error"] = powerstate.plasma["QieMWm3_tr_turb_stds"][0, i+1] return dictOFs diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index e38fdedd..9502165d 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -185,7 +185,7 @@ def output_transform_portals(X, surrogate_parameters, output): return compounded -def computeTurbExchangeIndividual(PexchTurb, powerstate): +def computeTurbExchangeIndividual(QieMWm3_tr_turb, powerstate): """ Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added """ @@ -195,8 +195,8 @@ def computeTurbExchangeIndividual(PexchTurb, powerstate): ------------------------------------------------------------------ E.g.: (batch1,batch2,batch3,dimR) -> (batch1*batch2*batch3,dimR) """ - shape_orig = np.array(PexchTurb.shape) - PexchTurb = PexchTurb.view(np.prod(shape_orig[:-1]), shape_orig[-1]) + shape_orig = np.array(QieMWm3_tr_turb.shape) + QieMWm3_tr_turb = QieMWm3_tr_turb.view(np.prod(shape_orig[:-1]), shape_orig[-1]) """ 2. Integrate @@ -206,18 +206,18 @@ def computeTurbExchangeIndividual(PexchTurb, powerstate): """ # Add zeros at zero - qExch = torch.cat((torch.zeros(PexchTurb.shape).to(PexchTurb)[..., :1], PexchTurb), dim=-1) + qExch = torch.cat((torch.zeros(QieMWm3_tr_turb.shape).to(QieMWm3_tr_turb)[..., :1], QieMWm3_tr_turb), dim=-1) - PexchTurb_integrated = powerstate.volume_integrate(qExch, force_dim=qExch.shape[0])[..., 1:] + QieMWm3_tr_turb_integrated = powerstate.volume_integrate(qExch, force_dim=qExch.shape[0])[..., 1:] """ 3. Go back to the original batching system ------------------------------------------------------------------------ E.g.: (batch1*batch2*batch3,dimR) -> (batch1,batch2,batch3,dimR) """ - PexchTurb_integrated = PexchTurb_integrated.view(tuple(shape_orig)) + QieMWm3_tr_turb_integrated = QieMWm3_tr_turb_integrated.view(tuple(shape_orig)) - return PexchTurb_integrated + return QieMWm3_tr_turb_integrated def GBfromXnorm(x, output, powerstate): # Decide, depending on the output here, which to use as normalization and at what location @@ -383,7 +383,7 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): "Ge_tar": "Ce", "GZ_tar": "CZ", "Mt_tar": "MtJm2", - "PexchTurb": "PexchTurb" + "QieMWm3_tr_turb": "QieMWm3_tr_turb" } for ikey in mapper: @@ -400,11 +400,11 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): # ------------------------------------------------------------------------- if PORTALSparameters["surrogateForTurbExch"]: - PexchTurb_integrated = computeTurbExchangeIndividual( - var_dict["PexchTurb"], powerstate + QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual( + var_dict["QieMWm3_tr_turb"], powerstate ) else: - PexchTurb_integrated = torch.zeros(dfT.shape).to(dfT) + QieMWm3_tr_turb_integrated = torch.zeros(dfT.shape).to(dfT) # ------------------------------------------------------------------------ # Go through each profile that needs to be predicted, calculate components @@ -440,9 +440,9 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): ----------------------------------------------------------------------------------- """ if var == "Qe": - cal0 = var_dict[f"{var}_tar"] + PexchTurb_integrated + cal0 = var_dict[f"{var}_tar"] + QieMWm3_tr_turb_integrated elif var == "Qi": - cal0 = var_dict[f"{var}_tar"] - PexchTurb_integrated + cal0 = var_dict[f"{var}_tar"] - QieMWm3_tr_turb_integrated else: cal0 = var_dict[f"{var}_tar"] @@ -517,7 +517,7 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): "Ge_tar": "Ce", "GZ_tar": "CZ", "Mt_tar": "MtJm2", - "PexchTurb": "PexchTurb" + "QieMWm3_tr_turb": "QieMWm3_tr_turb" } var_dict = {} @@ -535,15 +535,11 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): # ------------------------------------------------------------------------- if PORTALSparameters["surrogateForTurbExch"]: - PexchTurb_integrated = computeTurbExchangeIndividual( - var_dict["PexchTurb"], powerstate - ) - PexchTurb_integrated_stds = computeTurbExchangeIndividual( - var_dict["PexchTurb_stds"], powerstate - ) + QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["QieMWm3_tr_turb"], powerstate) + QieMWm3_tr_turb_integrated_stds = computeTurbExchangeIndividual(var_dict["QieMWm3_tr_turb_stds"], powerstate) else: - PexchTurb_integrated = torch.zeros(dfT.shape).to(dfT) - PexchTurb_integrated_stds = torch.zeros(dfT.shape).to(dfT) + QieMWm3_tr_turb_integrated = torch.zeros(dfT.shape).to(dfT) + QieMWm3_tr_turb_integrated_stds = torch.zeros(dfT.shape).to(dfT) # ------------------------------------------------------------------------ # Go through each profile that needs to be predicted, calculate components @@ -577,15 +573,11 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): ----------------------------------------------------------------------------------- """ if var == "Qe": - cal0 = var_dict[f"{var}_tar"] + PexchTurb_integrated - cal0E = ( - var_dict[f"{var}_tar_stds"] ** 2 + PexchTurb_integrated_stds**2 - ) ** 0.5 + cal0 = var_dict[f"{var}_tar"] + QieMWm3_tr_turb_integrated + cal0E = (var_dict[f"{var}_tar_stds"] ** 2 + QieMWm3_tr_turb_integrated_stds**2) ** 0.5 elif var == "Qi": - cal0 = var_dict[f"{var}_tar"] - PexchTurb_integrated - cal0E = ( - var_dict[f"{var}_tar_stds"] ** 2 + PexchTurb_integrated_stds**2 - ) ** 0.5 + cal0 = var_dict[f"{var}_tar"] - QieMWm3_tr_turb_integrated + cal0E = (var_dict[f"{var}_tar_stds"] ** 2 + QieMWm3_tr_turb_integrated_stds**2) ** 0.5 else: cal0 = var_dict[f"{var}_tar"] cal0E = var_dict[f"{var}_tar_stds"] diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 904ee4f5..7a88a739 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -303,7 +303,7 @@ def initializeProblem( if portals_fun.PORTALSparameters["surrogateForTurbExch"]: for i in range(len(portals_fun.MODELparameters["RhoLocations"])): - ofs.append(f"PexchTurb_{i+1}") + ofs.append(f"QieMWm3_tr_turb_{i+1}") name_transformed_ofs = [] for of in ofs: @@ -373,7 +373,7 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue "Ge", "Ge_tr_turb", "Ge_tr_neo", - "PexchTurb", + "QieMWm3_tr_turb", "Mt", "Mt_tr_turb", "Mt_tr_neo", diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index f7067722..6e46984b 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -25,7 +25,7 @@ def evaluate(self): # Some checks print("\t- Checking model modifications:") - for r in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "MtJm2_tr_turb"]: #, "PexchTurb"]: #TODO: FIX + for r in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "MtJm2_tr_turb"]: #, "QieMWm3_tr_turb"]: #TODO: FIX print(f"\t\t{r}(tglf) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(powerstate_orig.plasma[r][0][1:],powerstate_orig.plasma[r+'_stds'][0][1:]) ])}") print(f"\t\t{r}(cgyro) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(self.powerstate.plasma[r][0][1:],self.powerstate.plasma[r+'_stds'][0][1:]) ])}") diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index c04a4d20..f0e92579 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -137,11 +137,11 @@ def _evaluate_tglf(self): self.powerstate.plasma["Ce_tr_turb_stds"] = Flux_std[2] # if provideTurbulentExchange: - # self.powerstate.plasma["PexchTurb"] = - # self.powerstate.plasma["PexchTurb_stds"] = + # self.powerstate.plasma["QieMWm3_tr_turb"] = + # self.powerstate.plasma["QieMWm3_tr_turb_stds"] = # else: - # self.powerstate.plasma["PexchTurb"] = self.powerstate.plasma["QeMWm2_tr_turb"] * 0.0 - # self.powerstate.plasma["PexchTurb_stds"] = self.powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + # self.powerstate.plasma["QieMWm3_tr_turb"] = self.powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + # self.powerstate.plasma["QieMWm3_tr_turb_stds"] = self.powerstate.plasma["QeMWm2_tr_turb"] * 0.0 # ------------------------------------------------------------------------------------------------------------------------ diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index abad823a..d066fda1 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -1052,11 +1052,11 @@ def tgyro_to_powerstate(TGYROresults, # ********************************** if provideTurbulentExchange: - powerstate.plasma["PexchTurb"] = torch.Tensor(TGYROresults.EXe_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["PexchTurb_stds"] = torch.Tensor(TGYROresults.EXe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QieMWm3_tr_turb"] = torch.Tensor(TGYROresults.EXe_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QieMWm3_tr_turb_stds"] = torch.Tensor(TGYROresults.EXe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None else: - powerstate.plasma["PexchTurb"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 - powerstate.plasma["PexchTurb_stds"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + powerstate.plasma["QieMWm3_tr_turb"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + powerstate.plasma["QieMWm3_tr_turb_stds"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 # ********************************** # *********** Traget extra diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index dd94010b..eab921e7 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -31,7 +31,7 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ self.variables += [f'{i}_stds' for i in self.variables] # There is also turbulent exchange - self.variables += ['PexchTurb', 'PexchTurb_stds'] + self.variables += ['QieMWm3_tr_turb', 'QieMWm3_tr_turb_stds'] # And total transport flux self.variables += [f'{i}_tr' for i in self.quantities] diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 658c400f..556c0de3 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -1047,9 +1047,7 @@ def plot(self, fn=None, labels=["tgyro1"], doNotShow=False, fn_color=None): ax.plot(res.roa[-1], res.Qi_sim[-1], "-o", c="b", markersize=3) axE.plot(res.roa[-1], res.Qi_res[-1], "-o", c="b", markersize=3) - ax.plot( - res.roa[-1], res.Ce_tar[-1], "--o", c="m", label="Qconv", markersize=3 - ) + ax.plot(res.roa[-1], res.Ce_tar[-1], "--o", c="m", label="Qconv", markersize=3) ax.plot(res.roa[-1], res.Ce_sim[-1], "-o", c="m", markersize=3) axE.plot(res.roa[-1], np.abs(res.Ce_res[-1]), "-o", c="m", markersize=3) From e647fa2edd56ff7df262144b4e995b11410fb12c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 30 Jul 2025 10:35:05 +0200 Subject: [PATCH 118/385] Changed tr_neo to tr_neoc to clarify is not only NEO, and to have same chars as tr_turb --- src/mitim_modules/portals/PORTALSmain.py | 4 +- src/mitim_modules/portals/PORTALStools.py | 30 ++--- .../portals/utils/PORTALSinit.py | 12 +- .../portals/utils/PORTALSplot.py | 48 +++---- .../physics_models/transport_analytic.py | 8 +- .../physics_models/transport_cgyro.py | 2 +- .../physics_models/transport_tglf.py | 8 +- .../physics_models/transport_tgyro.py | 118 +++++++++--------- .../powertorch/utils/TRANSPORTtools.py | 4 +- 9 files changed, 117 insertions(+), 117 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index c45abbd0..9fb37b41 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -669,8 +669,8 @@ def map_powerstate_to_portals(powerstate, dictOFs): dictOFs[f"{var0}_tr_turb_{i+1}"]["value"] = powerstate.plasma[f"{var1}_tr_turb"][0, i+1] dictOFs[f"{var0}_tr_turb_{i+1}"]["error"] = powerstate.plasma[f"{var1}_tr_turb_stds"][0, i+1] - dictOFs[f"{var0}_tr_neo_{i+1}"]["value"] = powerstate.plasma[f"{var1}_tr_neo"][0, i+1] - dictOFs[f"{var0}_tr_neo_{i+1}"]["error"] = powerstate.plasma[f"{var1}_tr_neo_stds"][0, i+1] + dictOFs[f"{var0}_tr_neoc_{i+1}"]["value"] = powerstate.plasma[f"{var1}_tr_neoc"][0, i+1] + dictOFs[f"{var0}_tr_neoc_{i+1}"]["error"] = powerstate.plasma[f"{var1}_tr_neoc_stds"][0, i+1] """ TARGET calculation diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 9502165d..13ba4175 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -373,11 +373,11 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): "Ge_tr_turb": "Ce_tr_turb", "GZ_tr_turb": "CZ_tr_turb", "Mt_tr_turb": "MtJm2_tr_turb", - "Qe_tr_neo": "QeMWm2_tr_neo", - "Qi_tr_neo": "QiMWm2_tr_neo", - "Ge_tr_neo": "Ce_tr_neo", - "GZ_tr_neo": "CZ_tr_neo", - "Mt_tr_neo": "MtJm2_tr_neo", + "Qe_tr_neoc": "QeMWm2_tr_neoc", + "Qi_tr_neoc": "QiMWm2_tr_neoc", + "Ge_tr_neoc": "Ce_tr_neoc", + "GZ_tr_neoc": "CZ_tr_neoc", + "Mt_tr_neoc": "MtJm2_tr_neoc", "Qe_tar": "QeMWm2", "Qi_tar": "QiMWm2", "Ge_tar": "Ce", @@ -429,10 +429,10 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): """ ----------------------------------------------------------------------------------- - Transport (_tr_turb+_tr_neo) + Transport (_tr_turb+_tr_neoc) ----------------------------------------------------------------------------------- """ - of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neo"] + of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neoc"] """ ----------------------------------------------------------------------------------- @@ -507,11 +507,11 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): "Ge_tr_turb": "Ce_tr_turb", "GZ_tr_turb": "CZ_tr_turb", "Mt_tr_turb": "MtJm2_tr_turb", - "Qe_tr_neo": "QeMWm2_tr_neo", - "Qi_tr_neo": "QiMWm2_tr_neo", - "Ge_tr_neo": "Ce_tr_neo", - "GZ_tr_neo": "CZ_tr_neo", - "Mt_tr_neo": "MtJm2_tr_neo", + "Qe_tr_neoc": "QeMWm2_tr_neoc", + "Qi_tr_neoc": "QiMWm2_tr_neoc", + "Ge_tr_neoc": "Ce_tr_neoc", + "GZ_tr_neoc": "CZ_tr_neoc", + "Mt_tr_neoc": "MtJm2_tr_neoc", "Qe_tar": "QeMWm2", "Qi_tar": "QiMWm2", "Ge_tar": "Ce", @@ -561,11 +561,11 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): """ ----------------------------------------------------------------------------------- - Transport (_tr_turb+_tr_neo) + Transport (_tr_turb+_tr_neoc) ----------------------------------------------------------------------------------- """ - of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neo"] - of0E = (var_dict[f"{var}_tr_turb_stds"] ** 2 + var_dict[f"{var}_tr_neo_stds"] ** 2) ** 0.5 + of0 = var_dict[f"{var}_tr_turb"] + var_dict[f"{var}_tr_neoc"] + of0E = (var_dict[f"{var}_tr_turb_stds"] ** 2 + var_dict[f"{var}_tr_neoc_stds"] ** 2) ** 0.5 """ ----------------------------------------------------------------------------------- diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 7a88a739..eb921e71 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -295,7 +295,7 @@ def initializeProblem( for i in range(len(portals_fun.MODELparameters["RhoLocations"])): ofs.append(f"{var}_tr_turb_{i+1}") - ofs.append(f"{var}_tr_neo_{i+1}") + ofs.append(f"{var}_tr_neoc_{i+1}") ofs.append(f"{var}_tar_{i+1}") @@ -366,17 +366,17 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue if typ in [ "Qe", "Qe_tr_turb", - "Qe_tr_neo", + "Qe_tr_neoc", "Qi", "Qi_tr_turb", - "Qi_tr_neo", + "Qi_tr_neoc", "Ge", "Ge_tr_turb", - "Ge_tr_neo", + "Ge_tr_neoc", "QieMWm3_tr_turb", "Mt", "Mt_tr_turb", - "Mt_tr_neo", + "Mt_tr_neoc", ]: if doNotFitOnFixedValues: isAbsValFixed = pos == ( @@ -414,7 +414,7 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue if useThisOne: Variables[output].append(ikey) - elif typ in ["GZ", "GZ_tr_turb", "GZ_tr_neo"]: + elif typ in ["GZ", "GZ_tr_turb", "GZ_tr_neoc"]: if doNotFitOnFixedValues: isAbsValFixed = pos == portals_fun.powerstate.plasma["rho"].shape[-1] - 1 else: diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index de23f1ec..f7f52333 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -218,7 +218,7 @@ def PORTALSanalyzer_plotMetrics( if axTe_f is not None: axTe_f.plot( rho, - power.plasma['QeMWm2_tr_turb'].cpu().numpy() + power.plasma['QeMWm2_tr_neo'].cpu().numpy(), + power.plasma['QeMWm2_tr_turb'].cpu().numpy() + power.plasma['QeMWm2_tr_neoc'].cpu().numpy(), "-", c=col, lw=lwt, @@ -228,7 +228,7 @@ def PORTALSanalyzer_plotMetrics( if axTi_f is not None: axTi_f.plot( rho, - power.plasma['QiMWm2_tr_turb'].cpu().numpy() + power.plasma['QiMWm2_tr_neo'].cpu().numpy(), + power.plasma['QiMWm2_tr_turb'].cpu().numpy() + power.plasma['QiMWm2_tr_neoc'].cpu().numpy(), "-", c=col, lw=lwt, @@ -242,7 +242,7 @@ def PORTALSanalyzer_plotMetrics( axne_f.plot( rho, - power.plasma['Ge1E20sm2_tr_turb'].cpu().numpy()+power.plasma['Ge1E20sm2_tr_neo'].cpu().numpy(), + power.plasma['Ge1E20sm2_tr_turb'].cpu().numpy()+power.plasma['Ge1E20sm2_tr_neoc'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph) axne_f.plot( rho, @@ -255,13 +255,13 @@ def PORTALSanalyzer_plotMetrics( if axnZ_f is not None: - axnZ_f.plot(rho, power.plasma['CZ_raw_tr_turb'].cpu().numpy()+power.plasma['CZ_raw_tr_neo'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph) + axnZ_f.plot(rho, power.plasma['CZ_raw_tr_turb'].cpu().numpy()+power.plasma['CZ_raw_tr_neoc'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph) axnZ_f.plot(rho, power.plasma['CZ_raw'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) if axw0_f is not None: axw0_f.plot( rho, - power.plasma['MtJm2_tr_turb'].cpu().numpy() + power.plasma['MtJm2_tr_neo'].cpu().numpy(), + power.plasma['MtJm2_tr_turb'].cpu().numpy() + power.plasma['MtJm2_tr_neoc'].cpu().numpy(), "-", c=col, lw=lwt, @@ -2883,7 +2883,7 @@ def plotFluxComparison( if axTe_f is not None: axTe_f.plot( r[0][ixF:], - power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:], + power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neoc'].cpu().numpy()[0][ixF:], "-s", c=col, lw=2, @@ -2892,10 +2892,10 @@ def plotFluxComparison( alpha=alpha, ) - sigma = power.plasma['QeMWm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo_stds'].cpu().numpy()[0][ixF:] + sigma = power.plasma['QeMWm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neoc_stds'].cpu().numpy()[0][ixF:] - m_Qe, M_Qe = (power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:]) - stds * sigma, ( - power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:] + m_Qe, M_Qe = (power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neoc'].cpu().numpy()[0][ixF:]) - stds * sigma, ( + power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neoc'].cpu().numpy()[0][ixF:] ) + stds * sigma axTe_f.fill_between(r[0][ixF:], m_Qe, M_Qe, facecolor=col, alpha=alpha / 3) @@ -2906,7 +2906,7 @@ def plotFluxComparison( if axTi_f is not None: axTi_f.plot( r[0][ixF:], - power.plasma['QiMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neo'].cpu().numpy()[0][ixF:], + power.plasma['QiMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neoc'].cpu().numpy()[0][ixF:], "-s", markersize=msFlux, c=col, @@ -2916,13 +2916,13 @@ def plotFluxComparison( ) sigma = ( - power.plasma['QiMWm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neo_stds'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neoc_stds'].cpu().numpy()[0][ixF:] ) m_Qi, M_Qi = ( - power.plasma['QiMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neo'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neoc'].cpu().numpy()[0][ixF:] ) - stds * sigma, ( - power.plasma['QiMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neo'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QiMWm2_tr_neoc'].cpu().numpy()[0][ixF:] ) + stds * sigma axTi_f.fill_between(r[0][ixF:], m_Qi, M_Qi, facecolor=col, alpha=alpha / 3) @@ -2932,7 +2932,7 @@ def plotFluxComparison( if axne_f is not None: - Ge = power.plasma['Ge1E20sm2_tr_turb'].cpu().numpy() + power.plasma['Ge1E20sm2_tr_neo'].cpu().numpy() + Ge = power.plasma['Ge1E20sm2_tr_turb'].cpu().numpy() + power.plasma['Ge1E20sm2_tr_neoc'].cpu().numpy() axne_f.plot( r[0][ixF:], @@ -2945,7 +2945,7 @@ def plotFluxComparison( alpha=alpha, ) - sigma = power.plasma['Ge1E20sm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['Ge1E20sm2_tr_neo_stds'].cpu().numpy()[0][ixF:] + sigma = power.plasma['Ge1E20sm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['Ge1E20sm2_tr_neoc_stds'].cpu().numpy()[0][ixF:] m_Ge, M_Ge = Ge[0][ixF:] - stds * sigma, Ge[0][ixF:] + stds * sigma @@ -2956,7 +2956,7 @@ def plotFluxComparison( # ----------------------------------------------------------------------------------------------- if axnZ_f is not None: - GZ = power.plasma['CZ_raw_tr_turb'].cpu().numpy() + power.plasma['CZ_raw_tr_neo'].cpu().numpy() + GZ = power.plasma['CZ_raw_tr_turb'].cpu().numpy() + power.plasma['CZ_raw_tr_neoc'].cpu().numpy() axnZ_f.plot( r[0][ixF:], @@ -2969,7 +2969,7 @@ def plotFluxComparison( alpha=alpha, ) - sigma = power.plasma['CZ_raw_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['CZ_raw_tr_neo_stds'].cpu().numpy()[0][ixF:] + sigma = power.plasma['CZ_raw_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['CZ_raw_tr_neoc_stds'].cpu().numpy()[0][ixF:] m_Gi, M_Gi = ( GZ[0][ixF:] - stds * sigma, @@ -2984,7 +2984,7 @@ def plotFluxComparison( if axw0_f is not None: axw0_f.plot( r[0][ixF:], - power.plasma['MtJm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neo'].cpu().numpy()[0][ixF:], + power.plasma['MtJm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neoc'].cpu().numpy()[0][ixF:], "-s", markersize=msFlux, c=col, @@ -2993,10 +2993,10 @@ def plotFluxComparison( alpha=alpha, ) - sigma = power.plasma['MtJm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neo_stds'].cpu().numpy()[0][ixF:] + sigma = power.plasma['MtJm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neoc_stds'].cpu().numpy()[0][ixF:] - m_Mt, M_Mt = (power.plasma['MtJm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neo'].cpu().numpy()[0][ixF:]) - stds * sigma, ( - power.plasma['MtJm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neo'].cpu().numpy()[0][ixF:] + m_Mt, M_Mt = (power.plasma['MtJm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neoc'].cpu().numpy()[0][ixF:]) - stds * sigma, ( + power.plasma['MtJm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['MtJm2_tr_neoc'].cpu().numpy()[0][ixF:] ) + stds * sigma axw0_f.fill_between(r[0][ixF:], m_Mt, M_Mt, facecolor=col, alpha=alpha / 3) @@ -3127,7 +3127,7 @@ def plotFluxComparison( if axTe_f is not None: (l1,) = axTe_f.plot( r[0][ixF:], - power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:], + power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neoc'].cpu().numpy()[0][ixF:], "-", c="k", lw=2, @@ -3139,8 +3139,8 @@ def plotFluxComparison( ) l3 = axTe_f.fill_between( r[0][ixF:], - (power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:]) - stds, - (power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neo'].cpu().numpy()[0][ixF:]) + stds, + (power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neoc'].cpu().numpy()[0][ixF:]) - stds, + (power.plasma['QeMWm2_tr_turb'].cpu().numpy()[0][ixF:] + power.plasma['QeMWm2_tr_neoc'].cpu().numpy()[0][ixF:]) + stds, facecolor="k", alpha=0.3, ) diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py index d9002af3..a53ee114 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -46,11 +46,11 @@ def evaluate(self): self.powerstate.plasma["QeMWm2_tr_turb"] = Pe_tr * 2 / 3 self.powerstate.plasma["QiMWm2_tr_turb"] = Pi_tr * 2 / 3 - self.powerstate.plasma["QeMWm2_tr_neo"] = Pe_tr * 1 / 3 - self.powerstate.plasma["QiMWm2_tr_neo"] = Pi_tr * 1 / 3 + self.powerstate.plasma["QeMWm2_tr_neoc"] = Pe_tr * 1 / 3 + self.powerstate.plasma["QiMWm2_tr_neoc"] = Pi_tr * 1 / 3 - self.powerstate.plasma["QeMWm2_tr"] = self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["QeMWm2_tr_neo"] - self.powerstate.plasma["QiMWm2_tr"] = self.powerstate.plasma["QiMWm2_tr_turb"] + self.powerstate.plasma["QiMWm2_tr_neo"] + self.powerstate.plasma["QeMWm2_tr"] = self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["QeMWm2_tr_neoc"] + self.powerstate.plasma["QiMWm2_tr"] = self.powerstate.plasma["QiMWm2_tr_turb"] + self.powerstate.plasma["QiMWm2_tr_neoc"] # ------------------------------------------------------------------ # SURROGATE diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 6e46984b..99f59f45 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -126,7 +126,7 @@ def _print_info(self): ): txt += f"\n{var} = " for j in range(self.powerstate.plasma["rho"].shape[1] - 1): - txt += f"{self.powerstate.plasma[varn][0,j+1]-self.powerstate.plasma[f'{varn}_tr_neo'][0,j+1]:.4e} " + txt += f"{self.powerstate.plasma[varn][0,j+1]-self.powerstate.plasma[f'{varn}_tr_neoc'][0,j+1]:.4e} " return txt diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index f0e92579..0e6f6698 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -25,7 +25,7 @@ def evaluate(self): for variable in variables: # Add model suffixes - self.powerstate.plasma[f"{variable}_tr"] = self.powerstate.plasma[f"{variable}_tr_turb"] + self.powerstate.plasma[f"{variable}_tr_neo"] + self.powerstate.plasma[f"{variable}_tr"] = self.powerstate.plasma[f"{variable}_tr_turb"] + self.powerstate.plasma[f"{variable}_tr_neoc"] # ************************************************************************************ # Private functions for the evaluation @@ -166,9 +166,9 @@ def _evaluate_tglf(self): def _evaluate_neo(self): - self.powerstate.plasma["QeMWm2_tr_neo"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) - self.powerstate.plasma["QiMWm2_tr_neo"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) - self.powerstate.plasma["Ce_tr_neo"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) + self.powerstate.plasma["QeMWm2_tr_neoc"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) + self.powerstate.plasma["QiMWm2_tr_neoc"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) + self.powerstate.plasma["Ce_tr_neoc"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) return None diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index d066fda1..5408e139 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -448,23 +448,23 @@ def curateTGYROfiles( # Neo # ************************************************************************************************************************** - Qe_tr_neo = tgyro.Qe_sim_neo[0, 1:] + Qe_tr_neoc = tgyro.Qe_sim_neo[0, 1:] if includeFast: - Qi_tr_neo = tgyro.QiIons_sim_neo[0, 1:] + Qi_tr_neoc = tgyro.QiIons_sim_neo[0, 1:] else: - Qi_tr_neo = tgyro.QiIons_sim_neo_thr[0, 1:] - Ge_tr_neo = tgyro.Ge_sim_neo[0, 1:] - GZ_tr_neo = tgyro.Gi_sim_neo[impurityPosition, 0, 1:] - Mt_tr_neo = tgyro.Mt_sim_neo[0, 1:] + Qi_tr_neoc = tgyro.QiIons_sim_neo_thr[0, 1:] + Ge_tr_neoc = tgyro.Ge_sim_neo[0, 1:] + GZ_tr_neoc = tgyro.Gi_sim_neo[impurityPosition, 0, 1:] + Mt_tr_neoc = tgyro.Mt_sim_neo[0, 1:] - Qe_tr_neoE = abs(tgyro.Qe_sim_neo[0, 1:]) * relativeErrorNEO + Qe_tr_neocE = abs(tgyro.Qe_sim_neo[0, 1:]) * relativeErrorNEO if includeFast: - Qi_tr_neoE = abs(tgyro.QiIons_sim_neo[0, 1:]) * relativeErrorNEO + Qi_tr_neocE = abs(tgyro.QiIons_sim_neo[0, 1:]) * relativeErrorNEO else: - Qi_tr_neoE = abs(tgyro.QiIons_sim_neo_thr[0, 1:]) * relativeErrorNEO - Ge_tr_neoE = abs(tgyro.Ge_sim_neo[0, 1:]) * relativeErrorNEO - GZ_tr_neoE = abs(tgyro.Gi_sim_neo[impurityPosition, 0, 1:]) * relativeErrorNEO - Mt_tr_neoE = abs(tgyro.Mt_sim_neo[0, 1:]) * relativeErrorNEO + Qi_tr_neocE = abs(tgyro.QiIons_sim_neo_thr[0, 1:]) * relativeErrorNEO + Ge_tr_neocE = abs(tgyro.Ge_sim_neo[0, 1:]) * relativeErrorNEO + GZ_tr_neocE = abs(tgyro.Gi_sim_neo[impurityPosition, 0, 1:]) * relativeErrorNEO + Mt_tr_neocE = abs(tgyro.Mt_sim_neo[0, 1:]) * relativeErrorNEO # Merge @@ -477,11 +477,11 @@ def curateTGYROfiles( GZ, Mt, Pexch, - Qe_tr_neo=Qe_tr_neo, - Qi_tr_neo=Qi_tr_neo, - Ge_tr_neo=Ge_tr_neo, - GZ_tr_neo=GZ_tr_neo, - Mt_tr_neo=Mt_tr_neo, + Qe_tr_neoc=Qe_tr_neoc, + Qi_tr_neoc=Qi_tr_neoc, + Ge_tr_neoc=Ge_tr_neoc, + GZ_tr_neoc=GZ_tr_neoc, + Mt_tr_neoc=Mt_tr_neoc, impurityPosition=impurityPosition, ) @@ -494,11 +494,11 @@ def curateTGYROfiles( GZE, MtE, PexchE, - Qe_tr_neo=Qe_tr_neoE, - Qi_tr_neo=Qi_tr_neoE, - Ge_tr_neo=Ge_tr_neoE, - GZ_tr_neo=GZ_tr_neoE, - Mt_tr_neo=Mt_tr_neoE, + Qe_tr_neoc=Qe_tr_neocE, + Qi_tr_neoc=Qi_tr_neocE, + Ge_tr_neoc=Ge_tr_neocE, + GZ_tr_neoc=GZ_tr_neocE, + Mt_tr_neoc=Mt_tr_neocE, impurityPosition=impurityPosition, special_label="_stds", ) @@ -579,7 +579,7 @@ def modifyResults( tgyro, folder_tgyro, minErrorPercent=5.0, - percent_tr_neo=2.0, + percent_tr_neoc=2.0, useConvectiveFluxes=False, Qi_criterion_stable=0.0025, impurityPosition=3, @@ -732,11 +732,11 @@ def modifyFLUX( GZ, Mt, S, - Qe_tr_neo=None, - Qi_tr_neo=None, - Ge_tr_neo=None, - GZ_tr_neo=None, - Mt_tr_neo=None, + Qe_tr_neoc=None, + Qi_tr_neoc=None, + Ge_tr_neoc=None, + GZ_tr_neoc=None, + Mt_tr_neoc=None, impurityPosition=3, special_label=None, ): @@ -756,15 +756,15 @@ def modifyFLUX( # Particle flux: Update modTGYROfile(folder / "out.tgyro.flux_e", GeGB, pos=2, fileN_suffix=special_label) - if Ge_tr_neo is not None: - GeGB_neo = Ge_tr_neo / tgyro.Gamma_GB[-1, 1:] + if Ge_tr_neoc is not None: + GeGB_neo = Ge_tr_neoc / tgyro.Gamma_GB[-1, 1:] modTGYROfile(folder / "out.tgyro.flux_e", GeGB_neo, pos=1, fileN_suffix=special_label) # Energy flux: Update modTGYROfile(folder / "out.tgyro.flux_e", QeGB, pos=4, fileN_suffix=special_label) - if Qe_tr_neo is not None: - QeGB_neo = Qe_tr_neo / tgyro.Q_GB[-1, 1:] + if Qe_tr_neoc is not None: + QeGB_neo = Qe_tr_neoc / tgyro.Q_GB[-1, 1:] modTGYROfile(folder / "out.tgyro.flux_e", QeGB_neo, pos=3, fileN_suffix=special_label) # Rotation: Remove (it will be sum to the first ion) @@ -788,8 +788,8 @@ def modifyFLUX( modTGYROfile(folder / "out.tgyro.flux_i1", QiGB, pos=4, fileN_suffix=special_label) - if Qi_tr_neo is not None: - QiGB_neo = Qi_tr_neo / tgyro.Q_GB[-1, 1:] + if Qi_tr_neoc is not None: + QiGB_neo = Qi_tr_neoc / tgyro.Q_GB[-1, 1:] modTGYROfile(folder / "out.tgyro.flux_i1", QiGB_neo, pos=3, fileN_suffix=special_label) # Particle flux: Make ion particle fluxes zero, because I don't want to mistake TGLF with CGYRO when looking at tgyro results @@ -807,8 +807,8 @@ def modifyFLUX( modTGYROfile(folder / "out.tgyro.flux_i1", MtGB, pos=6, fileN_suffix=special_label) - if Mt_tr_neo is not None: - MtGB_neo = Mt_tr_neo / tgyro.Pi_GB[-1, 1:] + if Mt_tr_neoc is not None: + MtGB_neo = Mt_tr_neoc / tgyro.Pi_GB[-1, 1:] modTGYROfile(folder / "out.tgyro.flux_i1", MtGB_neo, pos=5, fileN_suffix=special_label) # Energy exchange: Remove (it will be the electrons one) @@ -831,8 +831,8 @@ def modifyFLUX( modTGYROfile(folder / f"out.tgyro.flux_i{i+2}",var,pos=pos,fileN_suffix=special_label) modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB,pos=2,fileN_suffix=special_label) - if GZ_tr_neo is not None: - GZGB_neo = GZ_tr_neo / tgyro.Gamma_GB[-1, 1:] + if GZ_tr_neoc is not None: + GZGB_neo = GZ_tr_neoc / tgyro.Gamma_GB[-1, 1:] modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB_neo,pos=1,fileN_suffix=special_label) @@ -922,10 +922,10 @@ def tgyro_to_powerstate(TGYROresults, # ********************************** powerstate.plasma["QeMWm2_tr_turb"] = torch.Tensor(TGYROresults.Qe_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QeMWm2_tr_neo"] = torch.Tensor(TGYROresults.Qe_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QeMWm2_tr_neoc"] = torch.Tensor(TGYROresults.Qe_sim_neo[:, :nr]).to(powerstate.dfT) powerstate.plasma["QeMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Qe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["QeMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Qe_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QeMWm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Qe_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: powerstate.plasma["QeMWm2"] = torch.Tensor(TGYROresults.Qe_tar[:, :nr]).to(powerstate.dfT) @@ -938,18 +938,18 @@ def tgyro_to_powerstate(TGYROresults, if includeFast: powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QiMWm2_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_tr_neoc"] = torch.Tensor(TGYROresults.QiIons_sim_neo[:, :nr]).to(powerstate.dfT) powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["QiMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QiMWm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None else: powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QiMWm2_tr_neo"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr[:, :nr]).to(powerstate.dfT) + powerstate.plasma["QiMWm2_tr_neoc"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr[:, :nr]).to(powerstate.dfT) powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["QiMWm2_tr_neo_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["QiMWm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: powerstate.plasma["QiMWm2"] = torch.Tensor(TGYROresults.Qi_tar[:, :nr]).to(powerstate.dfT) @@ -960,10 +960,10 @@ def tgyro_to_powerstate(TGYROresults, # ********************************** powerstate.plasma["MtJm2_tr_turb"] = torch.Tensor(TGYROresults.Mt_sim_turb[:, :nr]).to(powerstate.dfT) # So far, let's include fast in momentum - powerstate.plasma["MtJm2_tr_neo"] = torch.Tensor(TGYROresults.Mt_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["MtJm2_tr_neoc"] = torch.Tensor(TGYROresults.Mt_sim_neo[:, :nr]).to(powerstate.dfT) powerstate.plasma["MtJm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Mt_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["MtJm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Mt_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["MtJm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Mt_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: powerstate.plasma["MtJm2"] = torch.Tensor(TGYROresults.Mt_tar[:, :nr]).to(powerstate.dfT) @@ -975,10 +975,10 @@ def tgyro_to_powerstate(TGYROresults, # Store raw fluxes for better plotting later powerstate.plasma["Ge1E20sm2_tr_turb"] = torch.Tensor(TGYROresults.Ge_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ge1E20sm2_tr_neo"] = torch.Tensor(TGYROresults.Ge_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ge1E20sm2_tr_neoc"] = torch.Tensor(TGYROresults.Ge_sim_neo[:, :nr]).to(powerstate.dfT) powerstate.plasma["Ge1E20sm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Ge_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Ge1E20sm2_tr_neo_stds"] = torch.Tensor(TGYROresults.Ge_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ge1E20sm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ge_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: powerstate.plasma["Ge1E20sm2"] = torch.Tensor(TGYROresults.Ge_tar[:, :nr]).to(powerstate.dfT) @@ -987,10 +987,10 @@ def tgyro_to_powerstate(TGYROresults, if not useConvectiveFluxes: powerstate.plasma["Ce_tr_turb"] = powerstate.plasma["Ge1E20sm2_tr_turb"] - powerstate.plasma["Ce_tr_neo"] = powerstate.plasma["Ge1E20sm2_tr_neo"] + powerstate.plasma["Ce_tr_neoc"] = powerstate.plasma["Ge1E20sm2_tr_neoc"] powerstate.plasma["Ce_tr_turb_stds"] = powerstate.plasma["Ge1E20sm2_tr_turb_stds"] - powerstate.plasma["Ce_tr_neo_stds"] = powerstate.plasma["Ge1E20sm2_tr_neo_stds"] + powerstate.plasma["Ce_tr_neoc_stds"] = powerstate.plasma["Ge1E20sm2_tr_neoc_stds"] if provideTargets: powerstate.plasma["Ce"] = powerstate.plasma["Ge1E20sm2"] @@ -999,10 +999,10 @@ def tgyro_to_powerstate(TGYROresults, else: powerstate.plasma["Ce_tr_turb"] = torch.Tensor(TGYROresults.Ce_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_tr_neo"] = torch.Tensor(TGYROresults.Ce_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ce_tr_neoc"] = torch.Tensor(TGYROresults.Ce_sim_neo[:, :nr]).to(powerstate.dfT) powerstate.plasma["Ce_tr_turb_stds"] = torch.Tensor(TGYROresults.Ce_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Ce_tr_neo_stds"] = torch.Tensor(TGYROresults.Ce_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ce_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ce_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: powerstate.plasma["Ce"] = torch.Tensor(TGYROresults.Ce_tar[:, :nr]).to(powerstate.dfT) @@ -1014,10 +1014,10 @@ def tgyro_to_powerstate(TGYROresults, # Store raw fluxes for better plotting later powerstate.plasma["CZ_raw_tr_turb"] = torch.Tensor(TGYROresults.Gi_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) - powerstate.plasma["CZ_raw_tr_neo"] = torch.Tensor(TGYROresults.Gi_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) + powerstate.plasma["CZ_raw_tr_neoc"] = torch.Tensor(TGYROresults.Gi_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) powerstate.plasma["CZ_raw_tr_turb_stds"] = torch.Tensor(TGYROresults.Gi_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_raw_tr_neo_stds"] = torch.Tensor(TGYROresults.Gi_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["CZ_raw_tr_neoc_stds"] = torch.Tensor(TGYROresults.Gi_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: powerstate.plasma["CZ_raw"] = torch.Tensor(TGYROresults.Gi_tar[impurityPosition, :, :nr]).to(powerstate.dfT) @@ -1026,10 +1026,10 @@ def tgyro_to_powerstate(TGYROresults, if not useConvectiveFluxes: powerstate.plasma["CZ_tr_turb"] = powerstate.plasma["CZ_raw_tr_turb"] / OriginalFimp - powerstate.plasma["CZ_tr_neo"] = powerstate.plasma["CZ_raw_tr_neo"] / OriginalFimp + powerstate.plasma["CZ_tr_neoc"] = powerstate.plasma["CZ_raw_tr_neoc"] / OriginalFimp powerstate.plasma["CZ_tr_turb_stds"] = powerstate.plasma["CZ_raw_tr_turb_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_tr_neo_stds"] = powerstate.plasma["CZ_raw_tr_neo_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None + powerstate.plasma["CZ_tr_neoc_stds"] = powerstate.plasma["CZ_raw_tr_neoc_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None if provideTargets: powerstate.plasma["CZ"] = powerstate.plasma["CZ_raw"] / OriginalFimp @@ -1038,10 +1038,10 @@ def tgyro_to_powerstate(TGYROresults, else: powerstate.plasma["CZ_tr_turb"] = torch.Tensor(TGYROresults.Ci_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - powerstate.plasma["CZ_tr_neo"] = torch.Tensor(TGYROresults.Ci_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp + powerstate.plasma["CZ_tr_neoc"] = torch.Tensor(TGYROresults.Ci_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp powerstate.plasma["CZ_tr_turb_stds"] = torch.Tensor(TGYROresults.Ci_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_tr_neo_stds"] = torch.Tensor(TGYROresults.Ci_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None + powerstate.plasma["CZ_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ci_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None if provideTargets: powerstate.plasma["CZ"] = torch.Tensor(TGYROresults.Ci_tar[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp @@ -1072,7 +1072,7 @@ def tgyro_to_powerstate(TGYROresults, quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2', 'Ge1E20sm2', 'CZ_raw'] for ikey in quantities: - powerstate.plasma[ikey+"_tr"] = powerstate.plasma[ikey+"_tr_turb"] + powerstate.plasma[ikey+"_tr_neo"] + powerstate.plasma[ikey+"_tr"] = powerstate.plasma[ikey+"_tr_turb"] + powerstate.plasma[ikey+"_tr_neoc"] return powerstate diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index eab921e7..85735cc1 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -25,7 +25,7 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ self.quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2'] # Each flux has a turbulent and neoclassical component - self.variables = [f'{i}_tr_turb' for i in self.quantities] + [f'{i}_tr_neo' for i in self.quantities] + self.variables = [f'{i}_tr_turb' for i in self.quantities] + [f'{i}_tr_neoc' for i in self.quantities] # Each flux component has a standard deviation self.variables += [f'{i}_stds' for i in self.variables] @@ -129,7 +129,7 @@ def _modify_profiles(self): def evaluate(self): ''' This needs to populate the following in self.powerstate.plasma - - QeMWm2, QeMWm2_tr, QeMWm2_tr_turb, QeMWm2_tr_neo + - QeMWm2, QeMWm2_tr, QeMWm2_tr_turb, QeMWm2_tr_neoc Same for QiMWm2, Ce, CZ, MtJm2 and their respective standard deviations ''' From c545460a72830e859e571daef0b0453f730e3d8d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 30 Jul 2025 10:50:17 +0200 Subject: [PATCH 119/385] Removal of never-used functionality of using particle flux directly in metric (not convective) --- src/mitim_modules/portals/PORTALSmain.py | 1 - src/mitim_modules/portals/PORTALStools.py | 8 +- .../portals/utils/PORTALSanalysis.py | 1 - .../portals/utils/PORTALSinit.py | 3 - .../portals/utils/PORTALSoptimization.py | 1 - src/mitim_modules/powertorch/STATEtools.py | 3 - .../physics_models/transport_cgyro.py | 4 - .../physics_models/transport_tgyro.py | 73 +++++-------------- .../powertorch/utils/POWERplot.py | 28 ++----- .../powertorch/utils/TARGETStools.py | 16 ++-- .../powertorch/utils/TRANSPORTtools.py | 12 +-- 11 files changed, 39 insertions(+), 111 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 9fb37b41..ca913bc1 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -205,7 +205,6 @@ def __init__( "targets_evaluator": targets_evaluator, "targets_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) - "useConvectiveFluxes": True, # If True, then convective flux for final metric (not fitting). If False, particle flux "includeFastInQi": False, # If True, and fast ions have been included, in seprateNEO, sum fast "portals_transformation_variables": portals_transformation_variables, # Physics-informed parameters to fit surrogates "portals_transformation_variables_trace": portals_transformation_variables_trace, # Physics-informed parameters to fit surrogates for trace impurities diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 13ba4175..1ce38e20 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -232,9 +232,9 @@ def GBfromXnorm(x, output, powerstate): elif varFull[:2] == "Mt": quantity = "Pgb" elif varFull[:2] == "Ge": - quantity = "Ggb" if (not powerstate.useConvectiveFluxes) else "Qgb_convection" + quantity = "Qgb_convection" elif varFull[:2] == "GZ": - quantity = "Ggb" if (not powerstate.useConvectiveFluxes) else "Qgb_convection" + quantity = "Qgb_convection" elif varFull[:5] == "Pexch": quantity = "Sgb" @@ -400,9 +400,7 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): # ------------------------------------------------------------------------- if PORTALSparameters["surrogateForTurbExch"]: - QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual( - var_dict["QieMWm3_tr_turb"], powerstate - ) + QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["QieMWm3_tr_turb"], powerstate) else: QieMWm3_tr_turb_integrated = torch.zeros(dfT.shape).to(dfT) diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 35bf5377..5f4c7e52 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -160,7 +160,6 @@ def prep_metrics(self, ilast=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 diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index eb921e71..c553ce72 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -139,7 +139,6 @@ def initializeProblem( "profiles_postprocessing_fun" ], "impurityPosition": position_of_impurity, - "useConvectiveFluxes": portals_fun.PORTALSparameters["useConvectiveFluxes"], "UseFineGridTargets": portals_fun.PORTALSparameters["fineTargetsResolution"], "OriginalFimp": portals_fun.PORTALSparameters["fImp_orig"], "forceZeroParticleFlux": portals_fun.PORTALSparameters[ @@ -167,7 +166,6 @@ def initializeProblem( EvolutionOptions={ "ProfilePredicted": portals_fun.MODELparameters["ProfilesPredicted"], "rhoPredicted": xCPs, - "useConvectiveFluxes": portals_fun.PORTALSparameters["useConvectiveFluxes"], "impurityPosition": position_of_impurity, "fineTargetsResolution": portals_fun.PORTALSparameters["fineTargetsResolution"], }, @@ -223,7 +221,6 @@ def initializeProblem( EvolutionOptions={ "ProfilePredicted": portals_fun.MODELparameters["ProfilesPredicted"], "rhoPredicted": xCPs, - "useConvectiveFluxes": portals_fun.PORTALSparameters["useConvectiveFluxes"], "impurityPosition": position_of_impurity, "fineTargetsResolution": portals_fun.PORTALSparameters["fineTargetsResolution"], }, diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index ba2a9d7a..ef5f38d1 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -138,7 +138,6 @@ def flux_match_surrogate(step,profiles, plot_results=False, fn = None, file_writ EvolutionOptions={ "ProfilePredicted": step.surrogate_parameters["powerstate"].ProfilesPredicted, "rhoPredicted": step.surrogate_parameters["powerstate"].plasma["rho"][0,1:], - "useConvectiveFluxes": step.surrogate_parameters["powerstate"].useConvectiveFluxes, "impurityPosition": step.surrogate_parameters["powerstate"].impurityPosition, "fineTargetsResolution": step.surrogate_parameters["powerstate"].fineTargetsResolution, }, diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 99189092..f7744033 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -37,7 +37,6 @@ def __init__( - EvolutionOptions: - rhoPredicted: radial grid (MUST NOT CONTAIN ZERO, it will be added internally) - ProfilesPredicted: list of profiles to predict - - useConvectiveFluxes: boolean = whether to use convective fluxes instead of particle fluxes for FM - impurityPosition: int = position of the impurity in the ions set - fineTargetsResolution: int = resolution of the fine targets - TransportOptions: dictionary with transport_evaluator and ModelOptions @@ -76,7 +75,6 @@ def __init__( # Default options self.ProfilesPredicted = EvolutionOptions.get("ProfilePredicted", ["te", "ti", "ne"]) - self.useConvectiveFluxes = EvolutionOptions.get("useConvectiveFluxes", True) self.impurityPosition = EvolutionOptions.get("impurityPosition", 1) self.impurityPosition_transport = copy.deepcopy(self.impurityPosition) self.fineTargetsResolution = EvolutionOptions.get("fineTargetsResolution", None) @@ -705,7 +703,6 @@ def calculateTargets(self, relative_error_assumed=1.0): # Merge targets, calculate errors and normalize targets.postprocessing( relative_error_assumed=relative_error_assumed, - useConvectiveFluxes=self.useConvectiveFluxes, forceZeroParticleFlux=self.TransportOptions["ModelOptions"].get("forceZeroParticleFlux", False)) def calculateTransport( diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 99f59f45..41c7ece1 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -169,7 +169,6 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod minErrorPercent = PORTALSparameters["percentError_stable"] Qi_criterion_stable = PORTALSparameters["Qi_criterion_stable"] percentNeo = PORTALSparameters["percentError"][1] - useConvectiveFluxes = PORTALSparameters["useConvectiveFluxes"] try: impurityPosition = PROFILEStools.impurity_location(PROFILEStools.gacode_state(unmodified_profiles), PORTALSparameters["ImpurityOfInterest"]) @@ -189,7 +188,6 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod numPORTALS, minErrorPercent, Qi_criterion_stable, - useConvectiveFluxes, percentNeo, radii, OriginalFimp=OriginalFimp, @@ -221,7 +219,6 @@ def cgyroing( evaluations, minErrorPercent, Qi_criterion_stable, - useConvectiveFluxes, percentNeo, radii, OriginalFimp=1.0, @@ -267,7 +264,6 @@ def cgyroing( tgyro, FolderEvaluation, minErrorPercent=minErrorPercent, - useConvectiveFluxes=useConvectiveFluxes, Qi_criterion_stable=Qi_criterion_stable, percentNeo=percentNeo, impurityPosition=impurityPosition, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 5408e139..6862f8ef 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -142,7 +142,6 @@ def _postprocess_results(self, tgyro, label): ModelOptions = self.powerstate.TransportOptions["ModelOptions"] includeFast = ModelOptions.get("includeFastInQi",False) - useConvectiveFluxes = ModelOptions.get("useConvectiveFluxes", True) UseFineGridTargets = ModelOptions.get("UseFineGridTargets", False) provideTurbulentExchange = ModelOptions.get("TurbulentExchange", False) OriginalFimp = ModelOptions.get("OriginalFimp", 1.0) @@ -155,7 +154,6 @@ def _postprocess_results(self, tgyro, label): self.powerstate = tgyro_to_powerstate( tgyro.results[label], self.powerstate, - useConvectiveFluxes=useConvectiveFluxes, includeFast=includeFast, impurityPosition=impurityPosition, UseFineGridTargets=UseFineGridTargets, @@ -580,7 +578,6 @@ def modifyResults( folder_tgyro, minErrorPercent=5.0, percent_tr_neoc=2.0, - useConvectiveFluxes=False, Qi_criterion_stable=0.0025, impurityPosition=3, OriginalFimp=1.0, @@ -599,7 +596,6 @@ def modifyResults( Mt_target, ) = defineReferenceFluxes( tgyro, - useConvectiveFluxes=useConvectiveFluxes, impurityPosition=impurityPosition, ) @@ -859,7 +855,7 @@ def modTGYROfile(file, var, pos=0, fileN_suffix=None): f.write(line_new + "\n") def defineReferenceFluxes( - tgyro, factor_tauptauE=5, useConvectiveFluxes=False, impurityPosition=3 + tgyro, factor_tauptauE=5, impurityPosition=3 ): Qe_target = abs(tgyro.Qe_tar[0, 1:]) Qi_target = abs(tgyro.Qi_tar[0, 1:]) @@ -877,10 +873,9 @@ def defineReferenceFluxes( ) # tau_p in seconds Ge_target_special = (Ne / tau_special) / tgyro.dvoldr[0, 1:] # (1E20/seconds/m^2) - if useConvectiveFluxes: - Ge_target_special = PLASMAtools.convective_flux( - tgyro.Te[0, 1:], Ge_target_special - ) # (1E20/seconds/m^2) + Ge_target_special = PLASMAtools.convective_flux( + tgyro.Te[0, 1:], Ge_target_special + ) # (1E20/seconds/m^2) GZ_target_special = Ge_target_special * NZ / Ne @@ -894,7 +889,6 @@ def defineReferenceFluxes( def tgyro_to_powerstate(TGYROresults, powerstate, - useConvectiveFluxes=False, forceZeroParticleFlux=False, includeFast=False, impurityPosition=1, @@ -984,29 +978,16 @@ def tgyro_to_powerstate(TGYROresults, powerstate.plasma["Ge1E20sm2"] = torch.Tensor(TGYROresults.Ge_tar[:, :nr]).to(powerstate.dfT) powerstate.plasma["Ge1E20sm2_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - if not useConvectiveFluxes: - powerstate.plasma["Ce_tr_turb"] = powerstate.plasma["Ge1E20sm2_tr_turb"] - powerstate.plasma["Ce_tr_neoc"] = powerstate.plasma["Ge1E20sm2_tr_neoc"] + powerstate.plasma["Ce_tr_turb"] = torch.Tensor(TGYROresults.Ce_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ce_tr_neoc"] = torch.Tensor(TGYROresults.Ce_sim_neo[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_tr_turb_stds"] = powerstate.plasma["Ge1E20sm2_tr_turb_stds"] - powerstate.plasma["Ce_tr_neoc_stds"] = powerstate.plasma["Ge1E20sm2_tr_neoc_stds"] - - if provideTargets: - powerstate.plasma["Ce"] = powerstate.plasma["Ge1E20sm2"] - powerstate.plasma["Ce_stds"] = powerstate.plasma["Ge1E20sm2_stds"] - - else: - - powerstate.plasma["Ce_tr_turb"] = torch.Tensor(TGYROresults.Ce_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_tr_neoc"] = torch.Tensor(TGYROresults.Ce_sim_neo[:, :nr]).to(powerstate.dfT) - - powerstate.plasma["Ce_tr_turb_stds"] = torch.Tensor(TGYROresults.Ce_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Ce_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ce_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["Ce"] = torch.Tensor(TGYROresults.Ce_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_stds"] = torch.Tensor(TGYROresults.Ce_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ce_tr_turb_stds"] = torch.Tensor(TGYROresults.Ce_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ce_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ce_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["Ce"] = torch.Tensor(TGYROresults.Ce_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ce_stds"] = torch.Tensor(TGYROresults.Ce_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None # ********************************** # *********** Impurity Fluxes @@ -1023,29 +1004,15 @@ def tgyro_to_powerstate(TGYROresults, powerstate.plasma["CZ_raw"] = torch.Tensor(TGYROresults.Gi_tar[impurityPosition, :, :nr]).to(powerstate.dfT) powerstate.plasma["CZ_raw_stds"] = torch.Tensor(TGYROresults.Gi_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - if not useConvectiveFluxes: + powerstate.plasma["CZ_tr_turb"] = torch.Tensor(TGYROresults.Ci_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp + powerstate.plasma["CZ_tr_neoc"] = torch.Tensor(TGYROresults.Ci_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - powerstate.plasma["CZ_tr_turb"] = powerstate.plasma["CZ_raw_tr_turb"] / OriginalFimp - powerstate.plasma["CZ_tr_neoc"] = powerstate.plasma["CZ_raw_tr_neoc"] / OriginalFimp - - powerstate.plasma["CZ_tr_turb_stds"] = powerstate.plasma["CZ_raw_tr_turb_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_tr_neoc_stds"] = powerstate.plasma["CZ_raw_tr_neoc_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["CZ"] = powerstate.plasma["CZ_raw"] / OriginalFimp - powerstate.plasma["CZ_stds"] = powerstate.plasma["CZ_raw_stds"] / OriginalFimp if TGYROresults.tgyro_stds else None - - else: - - powerstate.plasma["CZ_tr_turb"] = torch.Tensor(TGYROresults.Ci_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - powerstate.plasma["CZ_tr_neoc"] = torch.Tensor(TGYROresults.Ci_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - - powerstate.plasma["CZ_tr_turb_stds"] = torch.Tensor(TGYROresults.Ci_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ci_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["CZ"] = torch.Tensor(TGYROresults.Ci_tar[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - powerstate.plasma["CZ_stds"] = torch.Tensor(TGYROresults.Ci_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None + powerstate.plasma["CZ_tr_turb_stds"] = torch.Tensor(TGYROresults.Ci_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None + powerstate.plasma["CZ_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ci_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None + + if provideTargets: + powerstate.plasma["CZ"] = torch.Tensor(TGYROresults.Ci_tar[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp + powerstate.plasma["CZ_stds"] = torch.Tensor(TGYROresults.Ci_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None # ********************************** # *********** Energy Exchange diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index f97b7fab..625292f0 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -43,16 +43,10 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co 'Electron Density','$n_e$ ($10^{20}m^{-3}$)','$a/Ln_e$','$\\Gamma_e$ (GB)','$\\Gamma_e$ ($10^{20}m^{-3}/s$)', 1E-1,"Ggb"]) else: - if self.useConvectiveFluxes: - set_plots.append( - [ 'ne', 'aLne', 'Ce_tr', 'Ce', - 'Electron Density','$n_e$ ($10^{20}m^{-3}$)','$a/Ln_e$','$Q_{conv,e}$ (GB)','$Q_{conv,e}$ ($MW/m^2$)', - 1E-1,"Qgb"]) - else: - set_plots.append( - [ 'ne', 'aLne', 'Ce_tr', 'Ce', - 'Electron Density','$n_e$ ($10^{20}m^{-3}$)','$a/Ln_e$','$\\Gamma_e$ (GB)','$\\Gamma_e$ ($10^{20}m^{-3}/s$)', - 1E-1,"Ggb"]) + set_plots.append( + [ 'ne', 'aLne', 'Ce_tr', 'Ce', + 'Electron Density','$n_e$ ($10^{20}m^{-3}$)','$a/Ln_e$','$Q_{conv,e}$ (GB)','$Q_{conv,e}$ ($MW/m^2$)', + 1E-1,"Qgb"]) if "nZ" in self.ProfilesPredicted: @@ -63,16 +57,10 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co 'Impurity Density','$n_Z$ ($10^{20}m^{-3}$)','$a/Ln_Z$','$\\Gamma_Z$ (GB)','$\\Gamma_Z$ ($10^{20}m^{-3}/s$)', 1E-1,"Ggb"]) else: - if self.useConvectiveFluxes: - set_plots.append( - [ 'nZ', 'aLnZ', 'CZ_tr', 'CZ', - 'Impurity Density','$n_Z$ ($10^{20}m^{-3}$)','$a/Ln_Z$','$\\widehat{Q}_{conv,Z}$ (GB)','$\\widehat{Q}_{conv,Z}$ ($MW/m^2$)', - 1E-1,"Qgb"]) - else: - set_plots.append( - [ 'nZ', 'aLnZ', 'CZ_tr', 'CZ', - 'Impurity Density','$n_Z$ ($10^{20}m^{-3}$)','$a/Ln_Z$','$\\Gamma_Z$ (GB)','$\\Gamma_Z$ ($10^{20}m^{-3}/s$)', - 1E-1,"Ggb"]) + set_plots.append( + [ 'nZ', 'aLnZ', 'CZ_tr', 'CZ', + 'Impurity Density','$n_Z$ ($10^{20}m^{-3}$)','$a/Ln_Z$','$\\widehat{Q}_{conv,Z}$ (GB)','$\\widehat{Q}_{conv,Z}$ ($MW/m^2$)', + 1E-1,"Qgb"]) if "w0" in self.ProfilesPredicted: set_plots.append( diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index ae420886..6f2502d9 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -110,7 +110,7 @@ def coarse_grid(self): for i in self.plasma_original: self.powerstate.plasma[i] = self.plasma_original[i] - def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, relative_error_assumed=1.0): + def postprocessing(self, forceZeroParticleFlux=False, relative_error_assumed=1.0): # ************************************************************************************************** # Plug-in targets that were fixed @@ -125,13 +125,9 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, if forceZeroParticleFlux: self.powerstate.plasma["Ge1E20sm2"] = self.powerstate.plasma["Ge1E20sm2"] * 0 - # Convective fluxes? - if useConvectiveFluxes: - self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["Ge1E20sm2"]) # MW/m^2 - self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["CZ_raw"]) # MW/m^2 - else: - self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ge1E20sm2"] - self.powerstate.plasma["CZ"] = self.powerstate.plasma["CZ_raw"] + # Convective fluxes + self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["Ge1E20sm2"]) # MW/m^2 + self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["CZ_raw"]) # MW/m^2 # ************************************************************************************************** # Error @@ -148,6 +144,6 @@ def postprocessing(self, useConvectiveFluxes=False, forceZeroParticleFlux=False, self.powerstate.plasma["QeGB"] = self.powerstate.plasma["QeMWm2"] / self.powerstate.plasma["Qgb"] self.powerstate.plasma["QiGB"] = self.powerstate.plasma["QiMWm2"] / self.powerstate.plasma["Qgb"] - self.powerstate.plasma["CeGB"] = self.powerstate.plasma["Ce"] / self.powerstate.plasma["Qgb" if useConvectiveFluxes else "Ggb"] - self.powerstate.plasma["CZGB"] = self.powerstate.plasma["CZ"] / self.powerstate.plasma["Qgb" if useConvectiveFluxes else "Ggb"] + self.powerstate.plasma["CeGB"] = self.powerstate.plasma["Ce"] / self.powerstate.plasma["Qgb"] + self.powerstate.plasma["CZGB"] = self.powerstate.plasma["CZ"] / self.powerstate.plasma["Qgb"] self.powerstate.plasma["MtGB"] = self.powerstate.plasma["MtJm2"] / self.powerstate.plasma["Pgb"] diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 85735cc1..fbf2ef42 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -53,16 +53,8 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ self.powerstate.labelsFluxes = { "te": "$Q_e$ ($MW/m^2$)", "ti": "$Q_i$ ($MW/m^2$)", - "ne": ( - "$Q_{conv}$ ($MW/m^2$)" - if self.powerstate.TransportOptions["ModelOptions"].get("useConvectiveFluxes", True) - else "$\\Gamma_e$ ($10^{20}/s/m^2$)" - ), - "nZ": ( - "$Q_{conv}$ $\\cdot f_{Z,0}$ ($MW/m^2$)" - if self.powerstate.TransportOptions["ModelOptions"].get("useConvectiveFluxes", True) - else "$\\Gamma_Z$ $\\cdot f_{Z,0}$ ($10^{20}/s/m^2$)" - ), + "ne": "$Q_{conv}$ ($MW/m^2$)", + "nZ": "$Q_{conv}$ $\\cdot f_{Z,0}$ ($MW/m^2$)", "w0": "$M_T$ ($J/m^2$)", } From 9423cf8dbf78d5a2c58b24ba80534f45394f9c2c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 31 Jul 2025 12:54:13 +0200 Subject: [PATCH 120/385] curation and naming of Ge and GZ --- .../portals/utils/PORTALSanalysis.py | 2 +- .../portals/utils/PORTALSplot.py | 10 +- .../physics_models/transport_tglf.py | 86 ++++++++---- .../physics_models/transport_tgyro.py | 14 +- .../powertorch/utils/POWERplot.py | 4 +- .../powertorch/utils/TARGETStools.py | 8 +- src/mitim_tools/gacode_tools/PROFILEStools.py | 129 ++++++++++-------- src/mitim_tools/opt_tools/scripts/read.py | 6 +- .../plasmastate_tools/utils/VMECtools.py | 88 ++++++++---- .../plasmastate_tools/utils/state_plotting.py | 8 +- 10 files changed, 224 insertions(+), 131 deletions(-) diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 5f4c7e52..5570e8f9 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -296,8 +296,8 @@ def prep_metrics(self, ilast=None): ] ) except: - embed() print("\t- Could not calculate Ricci metric", typeMsg="w") + embed() calculateRicci = None self.qR_Ricci, self.chiR_Ricci, self.points_Ricci = None, None, None diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index f7f52333..2440893c 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -242,11 +242,11 @@ def PORTALSanalyzer_plotMetrics( axne_f.plot( rho, - power.plasma['Ge1E20sm2_tr_turb'].cpu().numpy()+power.plasma['Ge1E20sm2_tr_neoc'].cpu().numpy(), + power.plasma['Ge1E20m2_tr_turb'].cpu().numpy()+power.plasma['Ge1E20m2_tr_neoc'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph) axne_f.plot( rho, - power.plasma['Ge1E20sm2'].cpu().numpy() * (1 - int(self.forceZeroParticleFlux)), + power.plasma['Ge1E20m2'].cpu().numpy() * (1 - int(self.forceZeroParticleFlux)), "--", c=col, lw=lwt, @@ -2932,7 +2932,7 @@ def plotFluxComparison( if axne_f is not None: - Ge = power.plasma['Ge1E20sm2_tr_turb'].cpu().numpy() + power.plasma['Ge1E20sm2_tr_neoc'].cpu().numpy() + Ge = power.plasma['Ge1E20m2_tr_turb'].cpu().numpy() + power.plasma['Ge1E20m2_tr_neoc'].cpu().numpy() axne_f.plot( r[0][ixF:], @@ -2945,7 +2945,7 @@ def plotFluxComparison( alpha=alpha, ) - sigma = power.plasma['Ge1E20sm2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['Ge1E20sm2_tr_neoc_stds'].cpu().numpy()[0][ixF:] + sigma = power.plasma['Ge1E20m2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['Ge1E20m2_tr_neoc_stds'].cpu().numpy()[0][ixF:] m_Ge, M_Ge = Ge[0][ixF:] - stds * sigma, Ge[0][ixF:] + stds * sigma @@ -3008,7 +3008,7 @@ def plotFluxComparison( Qe_tar = power.plasma['QeMWm2'].cpu().numpy()[0][ixF:] Qi_tar = power.plasma['QiMWm2'].cpu().numpy()[0][ixF:] - Ge_tar = power.plasma['Ge1E20sm2'].cpu().numpy()[0][ixF:] * (1-int(forceZeroParticleFlux)) + Ge_tar = power.plasma['Ge1E20m2'].cpu().numpy()[0][ixF:] * (1-int(forceZeroParticleFlux)) GZ_tar = power.plasma['CZ_raw'].cpu().numpy()[0][ixF:] Mt_tar = power.plasma['MtJm2'].cpu().numpy()[0][ixF:] diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 0e6f6698..8c5162b2 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -2,7 +2,7 @@ import shutil import torch import numpy as np -from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools import IOtools, PLASMAtools from mitim_tools.gacode_tools import TGLFtools from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -20,13 +20,8 @@ def evaluate(self): tglf = self._evaluate_tglf() neo = self._evaluate_neo() - # Sum the turbulent and neoclassical contributions - variables = ['QeMWm2', 'QiMWm2', 'Ce'] + self._postprocess() - for variable in variables: - # Add model suffixes - self.powerstate.plasma[f"{variable}_tr"] = self.powerstate.plasma[f"{variable}_tr_turb"] + self.powerstate.plasma[f"{variable}_tr_neoc"] - # ************************************************************************************ # Private functions for the evaluation # ************************************************************************************ @@ -133,25 +128,53 @@ def _evaluate_tglf(self): self.powerstate.plasma["QiMWm2_tr_turb"] = Flux_mean[1] self.powerstate.plasma["QiMWm2_tr_turb_stds"] = Flux_std[1] - self.powerstate.plasma["Ce_tr_turb"] = Flux_mean[2] - self.powerstate.plasma["Ce_tr_turb_stds"] = Flux_std[2] + self.powerstate.plasma["Ge1E20m2_tr_turb"] = Flux_mean[2] + self.powerstate.plasma["Ge1E20m2_tr_turb_stds"] = Flux_std[2] - # if provideTurbulentExchange: - # self.powerstate.plasma["QieMWm3_tr_turb"] = - # self.powerstate.plasma["QieMWm3_tr_turb_stds"] = - # else: - # self.powerstate.plasma["QieMWm3_tr_turb"] = self.powerstate.plasma["QeMWm2_tr_turb"] * 0.0 - # self.powerstate.plasma["QieMWm3_tr_turb_stds"] = self.powerstate.plasma["QeMWm2_tr_turb"] * 0.0 - + self.powerstate.plasma["GZ1E20m2_tr_turb"] = Flux_mean[3] + self.powerstate.plasma["GZ1E20m2_tr_turb_stds"] = Flux_std[3] + + self.powerstate.plasma["MtJm2_tr_turb"] = Flux_mean[4] + self.powerstate.plasma["MtJm2_tr_turb_stds"] = Flux_std[4] + + if provideTurbulentExchange: + self.powerstate.plasma["QieMWm3_tr_turb"] = Flux_mean[5] + self.powerstate.plasma["QieMWm3_tr_turb_stds"] = Flux_std[5] + else: + self.powerstate.plasma["QieMWm3_tr_turb"] = Flux_mean[5] * 0.0 + self.powerstate.plasma["QieMWm3_tr_turb_stds"] = Flux_std[5] * 0.0 + + return tglf + + def _evaluate_neo(self): + + self.powerstate.plasma["QeMWm2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["QiMWm2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["Ge1E20m2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["GZ1E20m2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["MtJm2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + + self.powerstate.plasma["QeMWm2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["QiMWm2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["Ge1E20m2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["GZ1E20m2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["MtJm2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + + self.powerstate.plasma["QieMWm3_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["QieMWm3_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + + return None + + def _postprocess(self): # ------------------------------------------------------------------------------------------------------------------------ # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) # ------------------------------------------------------------------------------------------------------------------------ - variables = ['QeMWm2', 'QiMWm2', 'Ce'] + variables = ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3'] for variable in variables: - for suffix in ['_tr_turb', '_tr_turb_stds']: + for suffix in ['_tr_turb', '_tr_turb_stds', '_tr_neoc', '_tr_neoc_stds']: # Make them tensors and add a batch dimension self.powerstate.plasma[f"{variable}{suffix}"] = torch.Tensor(self.powerstate.plasma[f"{variable}{suffix}"]).to(self.powerstate.dfT).unsqueeze(0) @@ -162,15 +185,30 @@ def _evaluate_tglf(self): self.powerstate.plasma[f"{variable}{suffix}"], ), dim=1) - return tglf + # ----------------------------------------------------------- + # Sum the turbulent and neoclassical contributions + # ----------------------------------------------------------- + + variables = ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2'] + + for variable in variables: + self.powerstate.plasma[f"{variable}_tr"] = self.powerstate.plasma[f"{variable}_tr_turb"] + self.powerstate.plasma[f"{variable}_tr_neoc"] - def _evaluate_neo(self): + # ----------------------------------------------------------- + # Convective fluxes + # ----------------------------------------------------------- - self.powerstate.plasma["QeMWm2_tr_neoc"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) - self.powerstate.plasma["QiMWm2_tr_neoc"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) - self.powerstate.plasma["Ce_tr_neoc"] = torch.zeros((1, len(self.powerstate.plasma["rho"][0, :]))) + mapper_convective = { + 'Ce': 'Ge1E20m2', + 'CZ': 'GZ1E20m2', + } - return None + for key in mapper_convective.keys(): + for tt in ['','_turb', '_turb_stds', '_neoc', '_neoc_stds']: + self.powerstate.plasma[f"{key}_tr{tt}"] = PLASMAtools.convective_flux( + self.powerstate.plasma["te"], + self.powerstate.plasma[f"{mapper_convective[key]}_tr{tt}"] + ) def _profiles_to_store(self): diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 6862f8ef..530cb5b5 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -968,15 +968,15 @@ def tgyro_to_powerstate(TGYROresults, # ********************************** # Store raw fluxes for better plotting later - powerstate.plasma["Ge1E20sm2_tr_turb"] = torch.Tensor(TGYROresults.Ge_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ge1E20sm2_tr_neoc"] = torch.Tensor(TGYROresults.Ge_sim_neo[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ge1E20m2_tr_turb"] = torch.Tensor(TGYROresults.Ge_sim_turb[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ge1E20m2_tr_neoc"] = torch.Tensor(TGYROresults.Ge_sim_neo[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ge1E20sm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Ge_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Ge1E20sm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ge_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ge1E20m2_tr_turb_stds"] = torch.Tensor(TGYROresults.Ge_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ge1E20m2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ge_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: - powerstate.plasma["Ge1E20sm2"] = torch.Tensor(TGYROresults.Ge_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ge1E20sm2_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["Ge1E20m2"] = torch.Tensor(TGYROresults.Ge_tar[:, :nr]).to(powerstate.dfT) + powerstate.plasma["Ge1E20m2_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None powerstate.plasma["Ce_tr_turb"] = torch.Tensor(TGYROresults.Ce_sim_turb[:, :nr]).to(powerstate.dfT) @@ -1037,7 +1037,7 @@ def tgyro_to_powerstate(TGYROresults, # Sum here turbulence and neoclassical, after modifications # ------------------------------------------------------------------------------------------------------------------------ - quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2', 'Ge1E20sm2', 'CZ_raw'] + quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2', 'Ge1E20m2', 'CZ_raw'] for ikey in quantities: powerstate.plasma[ikey+"_tr"] = powerstate.plasma[ikey+"_tr_turb"] + powerstate.plasma[ikey+"_tr_neoc"] diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index 625292f0..d8de6edd 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -37,9 +37,9 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co if "ne" in self.ProfilesPredicted: # If this model provides the raw particle flux, go for it - if 'Ge1E20sm2_tr' in self.plasma: + if 'Ge1E20m2_tr' in self.plasma: set_plots.append( - [ 'ne', 'aLne', 'Ge1E20sm2_tr', 'Ge1E20sm2', + [ 'ne', 'aLne', 'Ge1E20m2_tr', 'Ge1E20m2', 'Electron Density','$n_e$ ($10^{20}m^{-3}$)','$a/Ln_e$','$\\Gamma_e$ (GB)','$\\Gamma_e$ ($10^{20}m^{-3}/s$)', 1E-1,"Ggb"]) else: diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 6f2502d9..1cfa1f9e 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -118,22 +118,22 @@ def postprocessing(self, forceZeroParticleFlux=False, relative_error_assumed=1.0 self.powerstate.plasma["QeMWm2"] = self.powerstate.plasma["QeMWm2_fixedtargets"] + self.P[: self.P.shape[0]//2, :] # MW/m^2 self.powerstate.plasma["QiMWm2"] = self.powerstate.plasma["QiMWm2_fixedtargets"] + self.P[self.P.shape[0]//2 :, :] # MW/m^2 - self.powerstate.plasma["Ge1E20sm2"] = self.powerstate.plasma["Ge_fixedtargets"] # 1E20/s/m^2 + self.powerstate.plasma["Ge1E20m2"] = self.powerstate.plasma["Ge_fixedtargets"] # 1E20/s/m^2 self.powerstate.plasma["CZ_raw"] = self.powerstate.plasma["GZ_fixedtargets"] # 1E20/s/m^2 self.powerstate.plasma["MtJm2"] = self.powerstate.plasma["MtJm2_fixedtargets"] # J/m^2 if forceZeroParticleFlux: - self.powerstate.plasma["Ge1E20sm2"] = self.powerstate.plasma["Ge1E20sm2"] * 0 + self.powerstate.plasma["Ge1E20m2"] = self.powerstate.plasma["Ge1E20m2"] * 0 # Convective fluxes - self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["Ge1E20sm2"]) # MW/m^2 + self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["Ge1E20m2"]) # MW/m^2 self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["CZ_raw"]) # MW/m^2 # ************************************************************************************************** # Error # ************************************************************************************************** - variables_to_error = ["QeMWm2", "QiMWm2", "Ce", "CZ", "MtJm2", "Ge1E20sm2", "CZ_raw"] + variables_to_error = ["QeMWm2", "QiMWm2", "Ce", "CZ", "MtJm2", "Ge1E20m2", "CZ_raw"] for i in variables_to_error: self.powerstate.plasma[i + "_stds"] = abs(self.powerstate.plasma[i]) * relative_error_assumed / 100 diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 9fdaa432..2382682a 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -257,12 +257,27 @@ def derive_geometry(self, n_theta_geo=1001, **kwargs): def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): - [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,ax03c,ax13c] = axs3 + [ax00c,ax10c,ax20c,ax01c,ax11c,ax21c,ax02c,ax12c,ax22c,ax3D,ax2D] = axs3 rho = self.profiles["rho(-)"] lines = GRAPHICStools.listLS() + ax = ax00c + + var = self.derived['r'] + ax.plot(rho, var, "-", lw=lw, c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylim(bottom=0) + ax.set_ylabel("Effective radius ($r$)") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + + ax = ax01c ax.plot(self.profiles["rho(-)"], self.derived['volp_geo'], color=color, lw=lw, label = extralab) ax.set_xlabel('$\\rho_N$'); ax.set_xlim(0, 1) ax.set_ylabel(f"$dV/dr$ ($m^3/[r]$)") @@ -271,25 +286,55 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): if legYN: ax.legend(loc="best", fontsize=fs) - minShape = 1E-4 - ax = ax01c + ax = ax02c + var = self.profiles["polflux(Wb/radian)"] + ax.plot(rho, var, lw=lw, ls="-", c=color) + + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel("Poloidal $\\psi$ ($Wb/rad$)") + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax, bottomy=0) + + + # ---------------------------------------- + # Shaping params + # ---------------------------------------- + + minShape = 1E-3 + + ax = ax10c cont = 0 yl = 0 for i, s in enumerate(self.shape_cos): if s is not None: valmax = np.abs(s).max() if valmax > minShape: - lab = f"c{i}" + lab = f"s{i}" ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) cont += 1 yl = np.max([yl, valmax]) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\rho$") + ax.set_ylabel(f"sin-shape (>{minShape:.1e})") + if legYN: + ax.legend(loc="best", fontsize=fs) + GRAPHICStools.gradientSPAN(ax, -minShape, +minShape, color='k', startingalpha = 0.2, endingalpha = 0.2, orientation='horizontal') + + GRAPHICStools.addDenseAxis(ax) + GRAPHICStools.autoscale_y(ax) + + ax = ax11c + cont = 0 + yl = 0 for i, s in enumerate(self.shape_sin): if s is not None: valmax = np.abs(s).max() if valmax > minShape: - lab = f"s{i}" + lab = f"c{i}" ax.plot(rho, s, lw=lw, ls=lines[cont], label=lab, c=color) cont += 1 @@ -297,53 +342,48 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): ax.set_xlim([0, 1]) ax.set_xlabel("$\\rho$") - ax.set_ylabel(f"Shape Parameters (>{minShape})") + ax.set_ylabel(f"cos-shape (>{minShape:.1e})") if legYN: ax.legend(loc="best", fontsize=fs) + GRAPHICStools.gradientSPAN(ax, -minShape, +minShape, color='k', startingalpha = 0.2, endingalpha = 0.2, orientation='horizontal') + GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax) - ax = ax02c - var = self.profiles["polflux(Wb/radian)"] - ax.plot(rho, var, lw=lw, ls="-", c=color) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("Poloidal $\\psi$ ($Wb/rad$)") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = ax12c - ax = ax10c - var = self.profiles["delta(-)"] + var = self.profiles["kappa(-)"] ax.plot(rho, var, "-", lw=lw, c=color) ax.set_xlim([0, 1]) ax.set_xlabel("$\\rho$") - ax.set_ylabel("$\\delta$") + ax.set_ylabel("$\\kappa$") GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) - + GRAPHICStools.autoscale_y(ax, bottomy=1) - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = ax20c + var = self.profiles["delta(-)"] + ax.plot(rho, var, "-", lw=lw, c=color, label = extralab + ', $\\delta$') - ax = ax11c + var = self.profiles["zeta(-)"] + ax.plot(rho, var, "--", lw=lw, c=color, label = extralab + ', $\\zeta$') - var = self.derived['r'] - ax.plot(rho, var, "-", lw=lw, c=color) ax.set_xlim([0, 1]) ax.set_xlabel("$\\rho$") - ax.set_ylim(bottom=0) - ax.set_ylabel("Effective $r$") + ax.set_ylabel("$\\delta$ and $\\zeta$") + GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=0) + GRAPHICStools.autoscale_y(ax) + + if legYN: + ax.legend(loc="best", fontsize=fs) - ax = ax20c + ax = ax21c var = self.profiles["rmaj(m)"] ax.plot(rho, var, "-", lw=lw, c=color) @@ -355,7 +395,7 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax) - ax = ax21c + ax = ax22c var = self.profiles["zmag(m)"] ax.plot(rho, var, "-", lw=lw, c=color) @@ -369,38 +409,19 @@ def plot_geometry(self, axs3, color="b", legYN=True, extralab="", lw=1, fs=6): GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax) - ax = ax22c - var = self.profiles["kappa(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("$\\kappa$") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax, bottomy=1) - - ax = ax12c - - var = self.profiles["zeta(-)"] - ax.plot(rho, var, "-", lw=lw, c=color) - - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\rho$") - ax.set_ylabel("zeta") - - GRAPHICStools.addDenseAxis(ax) - GRAPHICStools.autoscale_y(ax) + # ------------------------------ + # 3D and 2D plots + # ------------------------------ - ax = ax13c + ax = ax2D self.plot_state_flux_surfaces(ax=ax, color=color) ax.set_xlabel("R (m)") ax.set_ylabel("Z (m)") GRAPHICStools.addDenseAxis(ax) - ax = ax03c + ax = ax3D self.plot_plasma_boundary(ax=ax, color=color) def plot_state_flux_surfaces(self, ax=None, surfaces_rho=np.linspace(0, 1, 11), color="b", label = '', lw=1.0, lw1=2.0): diff --git a/src/mitim_tools/opt_tools/scripts/read.py b/src/mitim_tools/opt_tools/scripts/read.py index f3001a08..f2edeec9 100644 --- a/src/mitim_tools/opt_tools/scripts/read.py +++ b/src/mitim_tools/opt_tools/scripts/read.py @@ -264,11 +264,7 @@ def main(): ] txt += f"\n\t...From remote folder {remote_parent}\n" - print( - "\n" - + txt - + "***************************************************************************" - ) + print("\n"+ txt+ "***************************************************************************") print(f"(Analysis level {analysis_level})\n") if len(folders_complete) == 1: diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index b94c89fb..eaf5f7e7 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -94,7 +94,7 @@ def derive_quantities(self, **kwargs): self.derived = {} # Define the minor radius used in all calculations (could be the half-width of the midplance intersect, or an effective minor radius) - self.derived["r"] = self.profiles["rho(-)"] + self.derived["r"] = self.profiles["rho(-)"] # Assume that r = rho so r/a = rho too super().derive_quantities_base(**kwargs) @@ -134,15 +134,6 @@ def plot_geometry(self, axs, color="b", legYN=True, extralab="", lw=1, fs=6): rho = self.profiles["rho(-)"] ax = ax00c - ax.plot(rho, self.derived['volp_geo'], color=color, lw=lw, label = extralab) - ax.set_xlabel('$\\rho$'); ax.set_xlim(0, 1) - ax.set_ylabel(f"$dV/d\\rho$ ($m^3$)") - GRAPHICStools.addDenseAxis(ax) - - if legYN: - ax.legend(loc="best", fontsize=fs) - - ax = ax11c var = self.derived['r'] ax.plot(rho, var, "-", lw=lw, c=color) @@ -150,21 +141,30 @@ def plot_geometry(self, axs, color="b", legYN=True, extralab="", lw=1, fs=6): ax.set_xlim([0, 1]) ax.set_xlabel("$\\rho$") ax.set_ylim(bottom=0) - ax.set_ylabel("Effective $r$") + ax.set_ylabel("Effective radius ($r$)") GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax, bottomy=0) + ax = ax01c + ax.plot(rho, self.derived['volp_geo'], color=color, lw=lw, label = extralab) + ax.set_xlabel('$\\rho$'); ax.set_xlim(0, 1) + ax.set_ylabel(f"$dV/d\\rho$ ($m^3$)") + GRAPHICStools.addDenseAxis(ax) + + if legYN: + ax.legend(loc="best", fontsize=fs) - self.plot_plasma_boundary(ax=axs_3d, color=color) + # ---- + phis_plot = [0.0, np.pi/2, np.pi, 3*np.pi/2] - self.plot_state_flux_surfaces(ax=axs_2d, c=color) + self.plot_plasma_boundary(ax=axs_3d, color=color, phi_cuts=phis_plot) + self.plot_state_flux_surfaces(ax=axs_2d, c=color, phis_plot=phis_plot) - - def plot_state_flux_surfaces(self, ax=None, c='b'): + def plot_state_flux_surfaces(self, ax=None, c='b', phis_plot=[0.0]): rhos_plot = np.linspace(0.0, 1.0, 10) - phis_plot = [0.0, np.pi/2, np.pi, 3*np.pi/2] + ls = GRAPHICStools.listLS() @@ -172,7 +172,7 @@ def plot_state_flux_surfaces(self, ax=None, c='b'): for i in range(len(rhos_plot)): self.plot_flux_surface(ax = ax, phi_cut=phi_cut, rho=rhos_plot[i], c=c, lw = 0.5, ls = lsi) - self.plot_flux_surface(ax = ax, phi_cut=phi_cut, rho=1.0, c=c, lw = 2, ls = lsi, label = f"{phi_cut*180/np.pi:.1f}°") + self.plot_flux_surface(ax = ax, phi_cut=phi_cut, rho=1.0, c=c, lw = 4, ls = lsi, label = f"{phi_cut*180/np.pi:.1f}°") ax.set_aspect('equal') ax.set_xlabel('R [m]') @@ -181,7 +181,7 @@ def plot_state_flux_surfaces(self, ax=None, c='b'): ax.legend(loc='best', fontsize=6) #GRAPHICStools.addLegendApart(ax, ratio=0.9, size=6) - ax.set_title(f'Poloidal Cross-section') + ax.set_title(f'Poloidal cross-sections') def _read_profiles(self, x_coord=None, debug = False): @@ -292,7 +292,7 @@ def _read_profiles(self, x_coord=None, debug = False): return uniform_data - def plot_plasma_boundary(self, ax=None, color="b"): + def plot_plasma_boundary(self, ax=None, color="b", phi_cuts=[]): # The output object contains the Fourier coefficients of the geometry in R and Z # as a function of the poloidal (theta) and toroidal (phi) angle-like coordinates @@ -345,12 +345,16 @@ def plot_plasma_boundary(self, ax=None, color="b"): ax = fig.add_subplot(projection="3d") # Plot the surface - ax.plot_surface(x, y, z, alpha=0.7, color=color) + ax.plot_surface(x, y, z, alpha=0.3 if len(phi_cuts)>0 else 0.7, color=color) + + # Add cutting planes at specific toroidal angles + for phi_cut in phi_cuts: + self._add_cutting_plane(ax, phi_cut, j, xm, xn, rmnc, zmns, color) # Set an equal aspect ratio ax.set_aspect("equal") - ax.set_title(f'Plasma Boundary') + ax.set_title(f'3D plasma boundary') def plot_flux_surface(self, ax=None, phi_cut=0.0, rho=1.0, c='b', lw=1, ls='-', label = ''): """ @@ -393,6 +397,45 @@ def plot_flux_surface(self, ax=None, phi_cut=0.0, rho=1.0, c='b', lw=1, ls='-', ax.plot(R, Z, ls=ls, color = c, linewidth=lw, label=label) + def _add_cutting_plane(self, ax, phi_cut, j, xm, xn, rmnc, zmns, plane_color): + """ + Add a cutting plane at a specific toroidal angle to the 3D plot. + + Parameters: + ----------- + ax : matplotlib 3D axes + The 3D axes to plot on + phi_cut : float + Toroidal angle for the cutting plane in radians + j : int + Flux surface index (typically ns-1 for boundary) + xm, xn : array + Poloidal and toroidal mode numbers + rmnc, zmns : array + Fourier coefficients for R and Z + plane_color : str + Color for the cutting plane + """ + num_theta = 101 + grid_theta = np.linspace(0.0, 2.0 * np.pi, num_theta, endpoint=True) + + R = np.zeros(num_theta) + Z = np.zeros(num_theta) + X = np.zeros(num_theta) + Y = np.zeros(num_theta) + + for idx_theta, theta in enumerate(grid_theta): + kernel = xm * theta - xn * phi_cut + r = np.dot(rmnc[:, j], np.cos(kernel)) + R[idx_theta] = r + Z[idx_theta] = np.dot(zmns[:, j], np.sin(kernel)) + X[idx_theta] = r * np.cos(phi_cut) + Y[idx_theta] = r * np.sin(phi_cut) + + # Plot the cutting plane as a line in 3D + ax.plot(X, Y, Z, color=plane_color, linewidth=2, + label=f'φ = {phi_cut*180/np.pi:.0f}°') + def plot_profiles(data): """ Create plots of the plasma profiles @@ -472,5 +515,4 @@ def plot_profiles(data): axes[1,2].remove() plt.tight_layout() - plt.show() -'' \ No newline at end of file + plt.show() \ No newline at end of file diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py index d8e89b19..277247bb 100644 --- a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -50,7 +50,7 @@ def add_axes(figs): ] # GEOMETRY - grid = plt.GridSpec(6, 4, hspace=0.7, wspace=0.3) + grid = plt.GridSpec(6, 5, hspace=0.8, wspace=0.4) ax00c = fig3.add_subplot(grid[0:2, 0]) axsProf_3 = [ ax00c, @@ -100,14 +100,12 @@ def add_axes(figs): fig6.add_subplot(grid[0, 4]), fig6.add_subplot(grid[1, 4]), ] - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + grid = plt.GridSpec(2, 2, hspace=0.3, wspace=0.3) axsImps = [ fig7.add_subplot(grid[0, 0]), fig7.add_subplot(grid[0, 1]), - fig7.add_subplot(grid[0, 2]), fig7.add_subplot(grid[1, 0]), fig7.add_subplot(grid[1, 1]), - fig7.add_subplot(grid[1, 2]), ] return axsProf_1, axsProf_2, axsProf_3, axsProf_4, axsFlows, axsProf_6, axsImps @@ -1282,8 +1280,6 @@ def plot_ions(self, axsImps, legYN=True, extralab="", color="b", lw=1, fs=6): GRAPHICStools.addDenseAxis(ax) GRAPHICStools.autoscale_y(ax, bottomy=0) - ax = axsImps[5] - ax = axsImps[3] ax.plot(self.profiles["rho(-)"], self.derived["Zeff"], c=color, lw=lw) ax.set_ylabel("$Z_{eff}$") From 413fdbd283274e640064c963ec634ab2fb93bd4e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 1 Aug 2025 10:45:21 +0200 Subject: [PATCH 121/385] Refinements Qie and CZ naming --- src/mitim_modules/portals/PORTALSmain.py | 6 +++--- src/mitim_modules/portals/PORTALStools.py | 10 +++++----- .../portals/utils/PORTALSinit.py | 6 +++--- .../portals/utils/PORTALSplot.py | 6 +++--- .../physics_models/transport_tglf.py | 19 ++++++++++++------- src/mitim_tools/gacode_tools/TGLFtools.py | 7 +++++++ 6 files changed, 33 insertions(+), 21 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index ca913bc1..54561fec 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -684,10 +684,10 @@ def map_powerstate_to_portals(powerstate, dictOFs): Turbulent Exchange ------------------ """ - if 'QieMWm3_tr_turb_1' in dictOFs: + if 'Qie_tr_turb_1' in dictOFs: for i in range(powerstate.plasma["rho"].shape[1] - 1): - dictOFs[f"QieMWm3_tr_turb_{i+1}"]["value"] = powerstate.plasma["QieMWm3_tr_turb"][0, i+1] - dictOFs[f"QieMWm3_tr_turb_{i+1}"]["error"] = powerstate.plasma["QieMWm3_tr_turb_stds"][0, i+1] + dictOFs[f"Qie_tr_turb_{i+1}"]["value"] = powerstate.plasma["QieMWm3_tr_turb"][0, i+1] + dictOFs[f"Qie_tr_turb_{i+1}"]["error"] = powerstate.plasma["QieMWm3_tr_turb_stds"][0, i+1] return dictOFs diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 1ce38e20..5f25103e 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -383,7 +383,7 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): "Ge_tar": "Ce", "GZ_tar": "CZ", "Mt_tar": "MtJm2", - "QieMWm3_tr_turb": "QieMWm3_tr_turb" + "Qie_tr_turb": "QieMWm3_tr_turb" } for ikey in mapper: @@ -400,7 +400,7 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): # ------------------------------------------------------------------------- if PORTALSparameters["surrogateForTurbExch"]: - QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["QieMWm3_tr_turb"], powerstate) + QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) else: QieMWm3_tr_turb_integrated = torch.zeros(dfT.shape).to(dfT) @@ -515,7 +515,7 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): "Ge_tar": "Ce", "GZ_tar": "CZ", "Mt_tar": "MtJm2", - "QieMWm3_tr_turb": "QieMWm3_tr_turb" + "Qie_tr_turb": "QieMWm3_tr_turb" } var_dict = {} @@ -533,8 +533,8 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): # ------------------------------------------------------------------------- if PORTALSparameters["surrogateForTurbExch"]: - QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["QieMWm3_tr_turb"], powerstate) - QieMWm3_tr_turb_integrated_stds = computeTurbExchangeIndividual(var_dict["QieMWm3_tr_turb_stds"], powerstate) + QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) + QieMWm3_tr_turb_integrated_stds = computeTurbExchangeIndividual(var_dict["Qie_tr_turb_stds"], powerstate) else: QieMWm3_tr_turb_integrated = torch.zeros(dfT.shape).to(dfT) QieMWm3_tr_turb_integrated_stds = torch.zeros(dfT.shape).to(dfT) diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index c553ce72..b263db21 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -98,7 +98,7 @@ def initializeProblem( if portals_fun.PORTALSparameters["UseOriginalImpurityConcentrationAsWeight"] is not None and portals_fun.PORTALSparameters["ImpurityOfInterest"] is not None: f0 = profiles.Species[position_of_impurity]["n0"] / profiles.profiles['ne(10^19/m^3)'][0] portals_fun.PORTALSparameters["fImp_orig"] = f0/portals_fun.PORTALSparameters["UseOriginalImpurityConcentrationAsWeight"] - print(f'\t- Ion {portals_fun.PORTALSparameters["ImpurityOfInterest"]} has original central concentration of {f0:.2e}, using its inverse multiplied by {portals_fun.PORTALSparameters["UseOriginalImpurityConcentrationAsWeight"]} as scaling factor of GZ -> {portals_fun.PORTALSparameters["fImp_orig"]}',typeMsg="i") + print(f'\t- Ion {portals_fun.PORTALSparameters["ImpurityOfInterest"]} has original central concentration of {f0:.2e}, using its inverse multiplied by {portals_fun.PORTALSparameters["UseOriginalImpurityConcentrationAsWeight"]} as scaling factor of GZ -> {portals_fun.PORTALSparameters["fImp_orig"]:.2e}',typeMsg="i") else: portals_fun.PORTALSparameters["fImp_orig"] = 1.0 @@ -300,7 +300,7 @@ def initializeProblem( if portals_fun.PORTALSparameters["surrogateForTurbExch"]: for i in range(len(portals_fun.MODELparameters["RhoLocations"])): - ofs.append(f"QieMWm3_tr_turb_{i+1}") + ofs.append(f"Qie_tr_turb_{i+1}") name_transformed_ofs = [] for of in ofs: @@ -370,7 +370,7 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue "Ge", "Ge_tr_turb", "Ge_tr_neoc", - "QieMWm3_tr_turb", + "Qie_tr_turb", "Mt", "Mt_tr_turb", "Mt_tr_neoc", diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index 2440893c..b89a759a 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -255,7 +255,7 @@ def PORTALSanalyzer_plotMetrics( if axnZ_f is not None: - axnZ_f.plot(rho, power.plasma['CZ_raw_tr_turb'].cpu().numpy()+power.plasma['CZ_raw_tr_neoc'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph) + axnZ_f.plot(rho, power.plasma['GZ1E20m2_tr_turb'].cpu().numpy()+power.plasma['GZ1E20m2_tr_neoc'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph) axnZ_f.plot(rho, power.plasma['CZ_raw'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) if axw0_f is not None: @@ -2956,7 +2956,7 @@ def plotFluxComparison( # ----------------------------------------------------------------------------------------------- if axnZ_f is not None: - GZ = power.plasma['CZ_raw_tr_turb'].cpu().numpy() + power.plasma['CZ_raw_tr_neoc'].cpu().numpy() + GZ = power.plasma['GZ1E20m2_tr_turb'].cpu().numpy() + power.plasma['GZ1E20m2_tr_neoc'].cpu().numpy() axnZ_f.plot( r[0][ixF:], @@ -2969,7 +2969,7 @@ def plotFluxComparison( alpha=alpha, ) - sigma = power.plasma['CZ_raw_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['CZ_raw_tr_neoc_stds'].cpu().numpy()[0][ixF:] + sigma = power.plasma['GZ1E20m2_tr_turb_stds'].cpu().numpy()[0][ixF:] + power.plasma['GZ1E20m2_tr_neoc_stds'].cpu().numpy()[0][ixF:] m_Gi, M_Gi = ( GZ[0][ixF:] - stds * sigma, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 8c5162b2..7b1acbd4 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -42,10 +42,10 @@ def _evaluate_tglf(self): percentError = ModelOptions.get("percentError", [5, 1, 0.5]) use_tglf_scan_trick = ModelOptions.get("use_tglf_scan_trick", None) cores_per_tglf_instance = ModelOptions.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) - + # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) impurityPosition = self.powerstate.impurityPosition_transport #ModelOptions.get("impurityPosition", 1) - + # ------------------------------------------------------------------------------------------------------------------------ # Prepare TGLF object # ------------------------------------------------------------------------------------------------------------------------ @@ -131,8 +131,8 @@ def _evaluate_tglf(self): self.powerstate.plasma["Ge1E20m2_tr_turb"] = Flux_mean[2] self.powerstate.plasma["Ge1E20m2_tr_turb_stds"] = Flux_std[2] - self.powerstate.plasma["GZ1E20m2_tr_turb"] = Flux_mean[3] - self.powerstate.plasma["GZ1E20m2_tr_turb_stds"] = Flux_std[3] + self.powerstate.plasma["GZ1E20m2_tr_turb"] = Flux_mean[3] + self.powerstate.plasma["GZ1E20m2_tr_turb_stds"] = Flux_std[3] self.powerstate.plasma["MtJm2_tr_turb"] = Flux_mean[4] self.powerstate.plasma["MtJm2_tr_turb_stds"] = Flux_std[4] @@ -167,6 +167,8 @@ def _evaluate_neo(self): def _postprocess(self): + OriginalFimp = self.powerstate.TransportOptions["ModelOptions"].get("OriginalFimp", 1.0) + # ------------------------------------------------------------------------------------------------------------------------ # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) # ------------------------------------------------------------------------------------------------------------------------ @@ -195,7 +197,7 @@ def _postprocess(self): self.powerstate.plasma[f"{variable}_tr"] = self.powerstate.plasma[f"{variable}_tr_turb"] + self.powerstate.plasma[f"{variable}_tr_neoc"] # ----------------------------------------------------------- - # Convective fluxes + # Convective fluxes (& Re-scale the GZ flux by the original impurity concentration) # ----------------------------------------------------------- mapper_convective = { @@ -205,11 +207,14 @@ def _postprocess(self): for key in mapper_convective.keys(): for tt in ['','_turb', '_turb_stds', '_neoc', '_neoc_stds']: + + mult = 1.0 if key == 'Ce' else 1/OriginalFimp + self.powerstate.plasma[f"{key}_tr{tt}"] = PLASMAtools.convective_flux( self.powerstate.plasma["te"], self.powerstate.plasma[f"{mapper_convective[key]}_tr{tt}"] - ) - + ) * mult + def _profiles_to_store(self): if "extra_params" in self.powerstate.TransportOptions["ModelOptions"] and "folder" in self.powerstate.TransportOptions["ModelOptions"]["extra_params"]: diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 6b5166bc..52a3b655 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -2396,6 +2396,9 @@ def _prepare_scan( def readScan( self, label="scan1", subFolderTGLF=None, variable="RLTS_1", positionIon=2 ): + ''' + positionIon is the index in the input.tglf file... so if you want for ion RLNS_5, positionIon=5 + ''' if subFolderTGLF is None: subFolderTGLF = self.subFolderTGLF_scan @@ -3088,6 +3091,10 @@ def runScanTurbulenceDrives( positionIon=2, **kwargs_TGLFrun, ): + + ''' + positionIon is the index in the input.tglf file... so if you want for ion RLNS_5, positionIon=5 + ''' self.variablesDrives = variablesDrives From 83afb53feaaca97ec197eb14947d8a2c754cbebd Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 4 Aug 2025 15:41:47 -0400 Subject: [PATCH 122/385] Unable to merge or rebase development, having to manually modify the scripts to accommodate misc changes and the following development commits: 0 Comments were switched 1 CGYRO plotting including ions 2 Bring back cross-phases that were removed before 3 Added mean QL flux retrieval in linear scans 4 check for linear convergence (make it easy to average in post) 5 added cpus_per_task for GPU submission 6 (cgyro) fixed queue specification and mem for job arrays 7 Added ion cross phases 8 Cleaned up animation function 9 if tmin provided for linear runs, use it (useful for unconverged cases) * 11 bug fix print ballooning 12 Correct mass when changing Z 13 sort folders and allow linear and nonlinear together * 15 Add EM fields by default 16 Calculate if signal explodes and adjust mean and plotting of fluctuations accordingly 17 Revert "added memory and queue arguments for job arrays" 18 Revert "plotting linear stability" 19 added memory and queue arguments for job arrays 20 plotting linear stability 21 added s_hat as variable 22 Bug fix infinite correlation 23 Added radial correlation length calculation from pygacode 24 Fixed bug in changeZeff 25 misc stylist 26 Fixed bug portals plot when surrogate for exchange 27 added grabber of linear NT 28 CGYRO improvements 29 Corrected autocorrelation function 30 Fluctuations to real units 31 First implementation of cross-phases 32 Intensities and kx plotting 33 CGYRO plotting improments 34 Added colorbars 35 Added x-y turb plots 36 Fluctuation intensities 37 Capability for full EPED beat 38 Added file backup 39 Added capability to EMBED in MAESTRO 40 Added standalone EPED plotting capabilities 41 Added option to prescribed EPED pedestal for debugging/exploration 42 Fixed issue restarting PORTALS beat 43 Fixed issue restarting TRANSP beat 44 misc improvements to CGYRO plotting 45 Cleaner CGYRO out class 46 Fixed workflow 47 Large improvements to CGYRO plotting 48 CGYRO plotting of NL frequencies 49 Possibility for more than one tmin in CGYRO reader 50 Added capability for horizontal lines fill graph 51 CGYRO nl plotting fluxes nicer 52 CGYRO saturation now native to MITIM 53 Increased lengthscale constraint to 5% 54 Dyn relax based on abs(Y) 55 Dynamic relaxation tolerance is now relative 56 Capability for dynamic relaxation per channel 57 Bug fix for SR stopping criterion 58 Flux match first point using new Target Options 59 bug fix no minimal 60 Capability to change Zeff maintaining fmain (changing Z) 61 bug fix plot special and vertical lines at Portals 62 Timing plotting improvements, accounting for restarts 63 MAESTRO timings plot 64 Created timing json for tracking times in MAESTRO. 65 Overplotting of special quantities for MAESTRO runs 66 Ease remote folder specification for plot_maestro 67 Added capability to retrieve from a remote MAESTRO run just the minimal information 68 plot fast ion gradients in red if present 69 Added object variables * 140 To make it clear for the user, rename the PORTALS profiles .new to new_fromtgyro_modified --- pyproject.toml | 1 + src/mitim_modules/maestro/MAESTROmain.py | 62 +- .../maestro/scripts/plot_maestro.py | 111 +- .../maestro/scripts/run_maestro.py | 18 +- src/mitim_modules/maestro/utils/EPEDbeat.py | 126 +- .../maestro/utils/MAESTRObeat.py | 3 +- .../maestro/utils/MAESTROplot.py | 203 +- .../maestro/utils/PORTALSbeat.py | 12 +- src/mitim_modules/maestro/utils/TRANSPbeat.py | 29 +- src/mitim_modules/portals/PORTALStools.py | 2 +- .../portals/utils/PORTALSoptimization.py | 14 +- .../portals/utils/PORTALSplot.py | 154 +- src/mitim_modules/powertorch/STATEtools.py | 6 +- .../physics_models/transport_cgyro.py | 4 +- .../physics_models/transport_tgyro.py | 2 +- src/mitim_tools/eped_tools/EPEDtools.py | 19 +- .../eped_tools/scripts/plot_eped.py | 27 + src/mitim_tools/gacode_tools/CGYROtools.py | 1646 ++++++++++++++--- src/mitim_tools/gacode_tools/TGLFtools.py | 4 +- .../gacode_tools/scripts/read_cgyro.py | 27 +- .../gacode_tools/utils/CGYROutils.py | 928 ++++++---- .../gacode_tools/utils/GACODErun.py | 4 +- src/mitim_tools/misc_tools/FARMINGtools.py | 10 +- src/mitim_tools/misc_tools/GRAPHICStools.py | 99 +- src/mitim_tools/misc_tools/IOtools.py | 130 +- src/mitim_tools/misc_tools/LOGtools.py | 4 +- src/mitim_tools/opt_tools/SURROGATEtools.py | 23 +- .../opt_tools/optimizers/ROOTtools.py | 24 +- src/mitim_tools/opt_tools/optimizers/optim.py | 80 +- .../opt_tools/scripts/evaluate_model.py | 16 +- src/mitim_tools/opt_tools/utils/BOgraphics.py | 8 +- .../plasmastate_tools/MITIMstate.py | 103 +- .../plasmastate_tools/utils/state_plotting.py | 14 + src/mitim_tools/transp_tools/CDFtools.py | 45 +- templates/input.cgyro.controls | 3 + tests/CGYRO_workflow.py | 6 +- tests/EPED_workflow.py | 8 +- 37 files changed, 3013 insertions(+), 962 deletions(-) create mode 100644 src/mitim_tools/eped_tools/scripts/plot_eped.py diff --git a/pyproject.toml b/pyproject.toml index 5ffd4e92..d5ee2da7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,6 +78,7 @@ mitim_plot_tglf = "mitim_tools.gacode_tools.scripts.read_tglf:main" # [--su mitim_plot_cgyro = "mitim_tools.gacode_tools.scripts.read_cgyro:main" mitim_plot_eq = "mitim_tools.gs_tools.scripts.read_eq:main" mitim_plot_transp = "mitim_tools.transp_tools.scripts.read_transp:main" +mitim_plot_eped = "mitim_tools.eped_tools.scripts.plot_eped:main" mitim_run_tglf = "mitim_tools.gacode_tools.scripts.run_tglf:main" # (folder input.tglf) [--gacode input.gacode] [--scan RLTS_2] [--drives True] diff --git a/src/mitim_modules/maestro/MAESTROmain.py b/src/mitim_modules/maestro/MAESTROmain.py index 3b633d1e..ceabcab2 100644 --- a/src/mitim_modules/maestro/MAESTROmain.py +++ b/src/mitim_modules/maestro/MAESTROmain.py @@ -19,9 +19,10 @@ MAESTRO: Modular and Accelerated Engine for Simulation of Transport and Reactor Optimization (If MAESTRO is the orchestrator, then BEAT is each of the beats (steps) that MAESTRO orchestrates) - ''' +ENABLE_EMBED = False # If True, will enable IPython embed, useful for debugging + class maestro: def __init__( @@ -52,13 +53,15 @@ def __init__( self.folder_output = self.folder / "Outputs" self.folder_logs = self.folder_output / "Logs" + self.folder_performance = self.folder_output / "Performance" self.folder_beats = self.folder / "Beats" self.folder_logs.mkdir(parents=True, exist_ok=True) self.folder_beats.mkdir(parents=True, exist_ok=True) + self.folder_performance.mkdir(parents=True, exist_ok=True) # If terminal outputs, I also want to keep track of what has happened in a log file - if terminal_outputs and overall_log_file: + if terminal_outputs and overall_log_file and not ENABLE_EMBED: sys.stdout = LOGtools.Logger(logFile=self.folder_output / "maestro.log", writeAlsoTerminal=True) branch, commit_hash = IOtools.get_git_info(__mitimroot__) @@ -134,7 +137,8 @@ def define_creator(self, method, **kwargs_creator): # Beat operations # -------------------------------------------------------------------------------------------- - @mitim_timer('\t\t* Checker', name_timer=None) + @mitim_timer( + lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Checker') def check(self, beat_check = None, cold_start = False, **kwargs): ''' Note: @@ -149,7 +153,7 @@ def check(self, beat_check = None, cold_start = False, **kwargs): print('\t- Checking...') log_file = self.folder_logs / f'beat_{self.counter_current}_check.log' if (not self.terminal_outputs) else None - with LOGtools.conditional_log_to_file(log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): + with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): output_file = None if not cold_start: @@ -169,13 +173,15 @@ def check(self, beat_check = None, cold_start = False, **kwargs): return output_file is not None - @mitim_timer('\t\t* Initializer', name_timer=None) + @mitim_timer( + lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Initializer', + log_file = lambda self: self.folder_performance / "timing.jsonl") def initialize(self, *args, **kwargs): print('\t- Initializing...') if self.beat.run_flag: log_file = self.folder_logs / f'beat_{self.counter_current}_ini.log' if (not self.terminal_outputs) else None - with LOGtools.conditional_log_to_file(log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): + with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): # Initialize: produce self.profiles_current self.beat.initialize(*args, **kwargs) @@ -183,7 +189,7 @@ def initialize(self, *args, **kwargs): print('\t\t- Skipping beat initialization because this beat was already run', typeMsg = 'i') log_file = self.folder_logs / f'beat_{self.counter_current}_inform.log' if (not self.terminal_outputs) else None - with LOGtools.conditional_log_to_file(log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): + with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): # Initializer can also save important parameters self.beat.initialize._inform_save() @@ -195,13 +201,15 @@ def initialize(self, *args, **kwargs): # First initialization, freeze engineering parameters self._freeze_parameters(profiles = PROFILEStools.gacode_state(self.beat.initialize.folder / 'input.gacode')) - @mitim_timer('\t\t* Preparation', name_timer=None) + @mitim_timer( + lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Preparation', + log_file = lambda self: self.folder_performance / "timing.jsonl") def prepare(self, *args, **kwargs): print('\t- Preparing...') if self.beat.run_flag: log_file = self.folder_logs / f'beat_{self.counter_current}_prep.log' if (not self.terminal_outputs) else None - with LOGtools.conditional_log_to_file(log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): + with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): # Initialize if necessary if not self.beat.initialize_called: @@ -209,20 +217,22 @@ def prepare(self, *args, **kwargs): self.beat.initialize() # ----------------------------- - self.beat.profiles_current.derive_quantities() + self.beat.profiles_current.deriveQuantities() self.beat.prepare(*args, **kwargs) else: print('\t\t- Skipping beat preparation because this beat was already run', typeMsg = 'i') - @mitim_timer('\t\t* Run + finalization', name_timer=None) + @mitim_timer( + lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Run + Finalization', + log_file = lambda self: self.folder_performance / "timing.jsonl") def run(self, **kwargs): # Run print('\t- Running...') if self.beat.run_flag: log_file = self.folder_logs / f'beat_{self.counter_current}_run.log' if (not self.terminal_outputs) else None - with LOGtools.conditional_log_to_file(log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): + with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): self.beat.run(**kwargs) else: print('\t\t- Skipping beat run because this beat was already run', typeMsg = 'i') @@ -230,7 +240,7 @@ def run(self, **kwargs): # Finalize, merging and freezing should occur even if the run has not been performed because the results are already there print('\t- Finalizing beat...') log_file = self.folder_logs / f'beat_{self.counter_current}_finalize.log' if (not self.terminal_outputs) else None - with LOGtools.conditional_log_to_file(log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): + with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): # Finalize self.beat.finalize(**kwargs) @@ -244,7 +254,7 @@ def run(self, **kwargs): # Inform next beats log_file = self.folder_logs / f'beat_{self.counter_current}_inform.log' if (not self.terminal_outputs) else None - with LOGtools.conditional_log_to_file(log_file=log_file): + with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file): self.beat._inform_save() # To save space, we can remove the contents of the run_ folder, as everything needed is in the output folder @@ -259,19 +269,21 @@ def _freeze_parameters(self, profiles = None): print('\t\t- Freezing engineering parameters from MAESTRO') self.profiles_with_engineering_parameters = copy.deepcopy(profiles) - self.profiles_with_engineering_parameters.write_state(file= (self.folder_output / 'input.gacode_frozen')) + self.profiles_with_engineering_parameters.writeCurrentStatus(file= (self.folder_output / 'input.gacode_frozen')) - @mitim_timer('\t\t* Finalizing', name_timer=None) + @mitim_timer( + lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Finalizing', + log_file = lambda self: self.folder_performance / "timing.jsonl") def finalize(self): print(f'- MAESTRO finalizing ******************************* {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}') log_file = self.folder_output / 'beat_final' if (not self.terminal_outputs) else None - with LOGtools.conditional_log_to_file(log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): + with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file, msg = f'\t\t* Log info being saved to {IOtools.clipstr(log_file)}'): final_file= (self.folder_output / 'input.gacode_final') - self.beat.profiles_output.write_state(file= final_file) + self.beat.profiles_output.writeCurrentStatus(file= final_file) print(f'\t\t- Final input.gacode saved to {IOtools.clipstr(final_file)}') @@ -279,7 +291,8 @@ def finalize(self): # Plotting operations # -------------------------------------------------------------------------------------------- - @mitim_timer('\t\t* Plotting', name_timer=None) + @mitim_timer( + lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Plotting') def plot(self, fn = None, num_beats = 2, only_beats = None, full_plot = True): print('*** Plotting MAESTRO ******************************************************************** ') @@ -291,11 +304,14 @@ def plot(self, fn = None, num_beats = 2, only_beats = None, full_plot = True): wasProvided = True self.fn = fn - self._plot_beats(self.fn, num_beats = num_beats, only_beats = only_beats, full_plot = full_plot) - self._plot_results(self.fn) + if num_beats>0: + self._plot_beats(self.fn, num_beats = num_beats, only_beats = only_beats, full_plot = full_plot) + ps, ps_lab = self._plot_results(self.fn) if not wasProvided: self.fn.show() + + return ps, ps_lab def _plot_beats(self, fn, num_beats = 2, only_beats = None, full_plot = True): @@ -306,7 +322,7 @@ def _plot_beats(self, fn, num_beats = 2, only_beats = None, full_plot = True): print(f'\t- Plotting beat #{counter}...') log_file = self.folder_logs / f'plot_{counter}.log' if (not self.terminal_outputs) else None - with LOGtools.conditional_log_to_file(log_file=log_file): + with LOGtools.conditional_log_to_file(write_log=not ENABLE_EMBED,log_file=log_file): msg = beat.plot(fn = self.fn, counter = i, full_plot = full_plot) print(msg) @@ -314,7 +330,7 @@ def _plot_results(self, fn): print('\t- Plotting MAESTRO results...') - MAESTROplot.plot_results(self, fn) + return MAESTROplot.plot_results(self, fn) diff --git a/src/mitim_modules/maestro/scripts/plot_maestro.py b/src/mitim_modules/maestro/scripts/plot_maestro.py index 98c13499..75a50adb 100644 --- a/src/mitim_modules/maestro/scripts/plot_maestro.py +++ b/src/mitim_modules/maestro/scripts/plot_maestro.py @@ -1,8 +1,9 @@ import argparse from mitim_modules.maestro.utils import MAESTROplot -from mitim_tools.misc_tools import IOtools, GUItools, FARMINGtools +from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools, FARMINGtools from mitim_tools.opt_tools import STRATEGYtools from pathlib import Path +from IPython import embed """ Quick way to plot several input.gacode files together (assumes unix in remote) @@ -33,27 +34,61 @@ def fix_maestro(folders): def main(): parser = argparse.ArgumentParser() - parser.add_argument("folders", type=str, nargs="*") - parser.add_argument("--remote",type=str, required=False, default=None) - parser.add_argument("--remote_folders",type=str, nargs="*", required=False, default=None) - parser.add_argument("--beats", type=int, required=False, default=2) # Last beats to plot - parser.add_argument("--only", type=str, required=False, default=None) - parser.add_argument("--full", required=False, default=False, action="store_true") - parser.add_argument('--fix', required=False, default=False, action='store_true') + + # Standard options + parser.add_argument("folders", type=str, nargs="*", + help="Paths to the folders to read.") + parser.add_argument("--beats", type=int, required=False, default=2, + help="Number of beats to plot. If 0, it will not plot beat information.") + parser.add_argument("--only", type=str, required=False, default=None, + help="If provided, it will only plot the specified beats (e.g., transp)") + parser.add_argument("--full", required=False, default=False, action="store_true", + help="If set, it will plot the full beat information.") + + # Remote options + parser.add_argument("--remote",type=str, required=False, default=None, + help="Remote machine to retrieve the folders from. If not provided, it will read the local folders.") + parser.add_argument("--remote_folder_parent",type=str, required=False, default=None, + help="Parent folder in the remote machine where the folders are located. If not provided, it will use --remote_folders.") + parser.add_argument("--remote_folders",type=str, nargs="*", required=False, default=None, + help="List of folders in the remote machine to retrieve. If not provided, it will use the local folder structures.") + parser.add_argument("--remote_minimal", required=False, default=False, action="store_true", + help="If set, it will only retrieve the folder structure with a few files (input.gacode, input.gacode_final, initializer_geqdsk/input.gacode).") + parser.add_argument('--fix', required=False, default=False, action='store_true', + help="If set, it will fix the pkl optimization portals in the remote folders.") args = parser.parse_args() remote = args.remote folders = args.folders fix = args.fix + beats = args.beats + only = args.only + full = args.full + + if args.remote_folder_parent is not None: + folders_remote = [args.remote_folder_parent + '/' + folder.split('/')[-1] for folder in folders] + elif args.remote_folders is not None: + folders_remote = args.remote_folders + else: + folders_remote = folders + # Retrieve remote if remote is not None: - if args.remote_folders is not None: - folders_remote = args.remote_folders - else: - folders_remote = folders - _, folders = FARMINGtools.retrieve_files_from_remote(IOtools.expandPath('./'), remote, folders_remote = folders_remote, purge_tmp_files = True) + + only_folder_structure_with_files = None + if args.remote_minimal: + only_folder_structure_with_files = ["beat_results/input.gacode", "input.gacode_final","initializer_geqdsk/input.gacode", "timing.jsonl"] + + beats = 0 + + _, folders = FARMINGtools.retrieve_files_from_remote( + IOtools.expandPath('./'), + remote, + folders_remote = folders_remote, + purge_tmp_files = True, + only_folder_structure_with_files=only_folder_structure_with_files) # Fix pkl optimization portals in remote if fix: @@ -62,17 +97,55 @@ def main(): # ----- folders = [IOtools.expandPath(folder) for folder in folders] - beats = args.beats - only = args.only - full = args.full - + fn = GUItools.FigureNotebook("MAESTRO") + if len(folders) > 1: + fig = fn.add_figure(label='MAESTRO special ALL', tab_color=4) + + axsAll = fig.subplot_mosaic( + """ + ABGI + ABGI + AEGI + DEHJ + DFHJ + DFHJ + """ + ) + + fig = fn.add_figure(label='MAESTRO timings ALL', tab_color=4) + axsTiming = fig.subplot_mosaic(""" + A + B + """) + + colors = GRAPHICStools.listColors() + ms = [] - for folder in folders: - m = MAESTROplot.plotMAESTRO(folder, fn = fn, num_beats=beats, only_beats = only, full_plot = full) + x, scripts = [], [] + x0, scripts0 = [], [] + for i,folder in enumerate(folders): + m, ps, ps_lab = MAESTROplot.plotMAESTRO(folder, fn = fn, num_beats=beats, only_beats = only, full_plot = full) ms.append(m) + # Plot all special quantities together + if len(folders) > 1: + MAESTROplot.plot_special_quantities(ps, ps_lab, axsAll, color=colors[i], label = f'Case #{i}', legYN = i==0) + if (m.folder_performance / 'timing.jsonl').exists(): + x0, scripts0 = MAESTROplot.plot_timings(m.folder_performance / 'timing.jsonl', axs = axsTiming, label = f'Case #{i}', color=colors[i]) + + # Only keep the longest + if len(x0) > len(x): + x = x0 + scripts = scripts0 + + if len(folders) > 1: + for let in ['A','B']: + axsTiming[let].set_xlim(left=0) + axsTiming[let].set_ylim(bottom=0) + axsTiming[let].set_xticks(x, scripts, rotation=10, ha="right", fontsize=8) + fn.show() # Import IPython and embed an interactive session diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index b9c0d5f6..1eb2643a 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -144,7 +144,7 @@ def profiles_postprocessing_fun(file_profs): p.lumpImpurities() if enforce_same_density_gradients: p.enforce_same_density_gradients() - p.write_state(file=file_profs) + p.writeCurrentStatus(file=file_profs) beat_namelist['PORTALSparameters']['profiles_postprocessing_fun'] = profiles_postprocessing_fun else: @@ -156,7 +156,7 @@ def profiles_postprocessing_fun(file_profs): return parameters_engineering, parameters_initialize, geometry, beat_namelists, maestro_beats, seed -@mitim_timer('\t\t* MAESTRO') +@mitim_timer('MAESTRO') def run_maestro_local( parameters_engineering, parameters_initialize, @@ -178,7 +178,13 @@ def run_maestro_local( if folder is None: folder = IOtools.expandPath('./') - m = maestro(folder, master_seed = seed, terminal_outputs = terminal_outputs, master_cold_start = force_cold_start, keep_all_files = keep_all_files) + m = maestro( + folder, + master_seed = seed, + terminal_outputs = terminal_outputs, + overall_log_file = True, + master_cold_start = force_cold_start, + keep_all_files = keep_all_files) # ------------------------------------------------------------------------- # Loop through beats @@ -221,6 +227,8 @@ def run_maestro_local( run_namelist = {} if maestro_beats["beats"][0] in ["transp", "transp_soft"]: run_namelist = {'mpisettings' : {"trmpi": cpus, "toricmpi": cpus, "ptrmpi": 1}} + elif maestro_beats["beats"][0] in ["eped"]: + run_namelist = {'cold_start': force_cold_start, 'cpus': cpus} m.prepare(**beat_namelists[maestro_beats["beats"][0]]) m.run(**run_namelist) @@ -246,8 +254,8 @@ def main(): if not folder.exists(): folder.mkdir(parents=True, exist_ok=True) - shutil.copy2(file_path, folder / 'maestro_namelist.json') - + IOtools.recursive_backup(folder / 'maestro_namelist.json') + run_maestro_local(*parse_maestro_nml(file_path),folder=folder,cpus = cpus, terminal_outputs = terminal_outputs) diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index 8de195d6..e57509c0 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -5,6 +5,7 @@ import matplotlib.pyplot as plt from scipy.optimize import curve_fit from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.eped_tools import EPEDtools from mitim_tools.misc_tools import IOtools, GRAPHICStools, GUItools from mitim_tools.surrogate_tools import NNtools from mitim_tools.popcon_tools import FunctionalForms @@ -35,11 +36,21 @@ def prepare( **kwargs ): - self.nn = NNtools.eped_nn(type='tf') - nn_location = IOtools.expandPath(nn_location) - norm_location = IOtools.expandPath(norm_location) + if nn_location is not None: - self.nn.load(nn_location, norm=norm_location) + print(f'\t- Choice of EPED: NN from {IOtools.clipstr(nn_location)}', typeMsg='i') + + self.nn = NNtools.eped_nn(type='tf') + nn_location = IOtools.expandPath(nn_location) + norm_location = IOtools.expandPath(norm_location) + + self.nn.load(nn_location, norm=norm_location) + + else: + + print('\t- Choice of EPED: full', typeMsg='i') + + self.nn = None # Parameters to run EPED with instead of those from the profiles self.neped_20 = neped_20 @@ -63,7 +74,7 @@ def run(self, **kwargs): # Run the NN # ------------------------------------------------------- - eped_results = self._run(loopBetaN = 1, store_scan=True) + eped_results = self._run(loopBetaN = 1, store_scan=True, nproc_per_run=kwargs.get('cpus', 16), cold_start=kwargs.get('cold_start', False)) # ------------------------------------------------------- # Save stuff @@ -73,7 +84,7 @@ def run(self, **kwargs): self.rhotop = eped_results['rhotop'] - def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = False): + def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = False, nproc_per_run=64, cold_start=True): ''' minimum_relative_change_in_x: minimum relative change in x to streach the core, otherwise it will keep the old core ''' @@ -137,7 +148,8 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F # --- Sometimes we may need specific EPED inputs for key, value in self.corrections_set.items(): - self.current_evaluation[key] = value + if key not in ['ptop_kPa', 'wtop_psipol']: + self.current_evaluation[key] = value # ---------------------------------------------- print('\n\t- Running EPED with:') @@ -162,7 +174,7 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F for i in range(loopBetaN): print(f'\t\t- BetaN: {BetaN:.2f}') - inputs_to_nn = ( + inputs_to_eped = ( self.current_evaluation["Ip"], self.current_evaluation["Bt"], self.current_evaluation["R"], @@ -176,8 +188,32 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F self.current_evaluation["nesep_ratio"] ) - ptop_kPa, wtop_psipol = self.nn(*inputs_to_nn) - + # ------------------------------------------------------- + # Give the option to override the ptop_kPa and wtop_psipol + if 'ptop_kPa' in self.corrections_set: + print(f'\t\t- Overriding ptop_kPa: {self.corrections_set["ptop_kPa"]:.2f} kPa', typeMsg='w') + ptop_kPa = self.corrections_set["ptop_kPa"] + else: + ptop_kPa = None + + if 'wtop_psipol' in self.corrections_set: + print(f'\t\t- Overriding wtop_psipol: {self.corrections_set["wtop_psipol"]:.5f}', typeMsg='w') + wtop_psipol = self.corrections_set["wtop_psipol"] + else: + wtop_psipol = None + # ------------------------------------------------------- + + if ptop_kPa is None or wtop_psipol is None: + + if self.nn is not None: + ptop_kPa, wtop_psipol = self.nn(*inputs_to_eped) + else: + ptop_kPa, wtop_psipol = self._run_full_eped(self.folder,*inputs_to_eped, nproc_per_run=nproc_per_run, cold_start=cold_start) + + if store_scan: + store_scan = False + print('\t- Warning: store_scan is not available for full EPED runs yet, only for NN-based EPED') + print('\t- Raw EPED results:') print(f'\t\t- ptop_kPa: {ptop_kPa:.4f}') print(f'\t\t- wtop_psipol: {wtop_psipol:.4f}') @@ -251,10 +287,10 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F scan_results = {} for k,key in enumerate(scan_relative): - inputs_scan = list(copy.deepcopy(inputs_to_nn)) + inputs_scan = list(copy.deepcopy(inputs_to_eped)) scan_results[key] = {'ptop_kPa': [], 'wtop_psipol': [], 'value': []} for m in np.linspace(1-scan_relative[key],1+scan_relative[key],15): - inputs_scan[k] = inputs_to_nn[k]*m + inputs_scan[k] = inputs_to_eped[k]*m ptop_kPa0, wtop_psipol0 = self.nn(*inputs_scan) scan_results[key]['ptop_kPa'].append(ptop_kPa0) scan_results[key]['wtop_psipol'].append(wtop_psipol0) @@ -263,7 +299,7 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F scan_results[key]['wtop_psipol'] = np.array(scan_results[key]['wtop_psipol']) scan_results[key]['value'] = np.array(scan_results[key]['value']) - scan_results[key]['ptop_kPa_nominal'], scan_results[key]['wtop_psipol_nominal'] = self.nn(*inputs_to_nn) + scan_results[key]['ptop_kPa_nominal'], scan_results[key]['wtop_psipol_nominal'] = self.nn(*inputs_to_eped) # --------------------------------- # Store @@ -278,22 +314,58 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F 'nesep_20': nesep_20, 'rhotop': rhotop, 'Tesep_keV': Tesep_keV, - 'inputs_to_nn': inputs_to_nn, + 'inputs_to_eped': inputs_to_eped, 'scan_results': scan_results } for key in eped_results: print(f'\t\t- {key}: {eped_results[key]}') - self.profiles_output.write_state(file=self.folder / 'input.gacode.eped') + self.profiles_output.writeCurrentStatus(file=self.folder / 'input.gacode.eped') return eped_results + def _run_full_eped(self, folder, Ip, Bt, R, a, kappa995, delta995, neped19, BetaN, zeff, Tesep_eV, nesep_ratio, nproc_per_run=64, cold_start=True): + ''' + Run the full EPED code with the given inputs. + Returns ptop_kPa and wtop_psipol. + ''' + + eped = EPEDtools.EPED(folder=folder) + + input_params = { + 'ip': Ip, + 'bt': Bt, + 'r': R, + 'a': a, + 'kappa': kappa995, + 'delta': delta995, + 'neped': neped19, + 'betan': BetaN, + 'zeffped': zeff, + 'nesep': nesep_ratio * neped19, + 'tesep': Tesep_eV + } + + eped.run( + subfolder = 'case1', + input_params = input_params, + nproc_per_run = nproc_per_run, + cold_start = cold_start, + ) + + eped.read(subfolder='case1') + + ptop_kPa = float(eped.results['case1']['run1']['ptop']) + wtop_psipol = float(eped.results['case1']['run1']['wptop']) + + return ptop_kPa, wtop_psipol + def finalize(self, **kwargs): - self.profiles_output = PROFILEStools.gacode_state(self.folder / 'input.gacode.eped') + self.profiles_output = PROFILEStools.PROFILES_GACODE(self.folder / 'input.gacode.eped') - self.profiles_output.write_state(file=self.folder_output / 'input.gacode') + self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') def merge_parameters(self): # EPED beat does not modify the profiles grid or anything, so I can keep it fine @@ -307,7 +379,7 @@ def grab_output(self): loaded_results = np.load(self.folder_output / 'eped_results.npy', allow_pickle=True).item() - profiles = PROFILEStools.gacode_state(self.folder_output / 'input.gacode') if isitfinished else None + profiles = PROFILEStools.PROFILES_GACODE(self.folder_output / 'input.gacode') if isitfinished else None else: @@ -332,7 +404,7 @@ def plot(self, fn = None, counter = 0, full_plot = True): loaded_results, profiles = self.grab_output() - profiles_current = PROFILEStools.gacode_state(self.folder / 'input.gacode') + profiles_current = PROFILEStools.PROFILES_GACODE(self.folder / 'input.gacode') profiles_current.plotRelevant(axs = axs, color = 'b', label = 'orig') @@ -397,10 +469,10 @@ def _plot_scan(self, ikey, loaded_results = None, axs = None, color = 'b'): axs[i].plot(loaded_results['scan_results'][key]['value'], loaded_results['scan_results'][key][ikey], 's-', color=color, markersize=3) - axs[i].plot([loaded_results['inputs_to_nn'][i]], [loaded_results[ikey]], '^', color=color) - axs[i].plot([loaded_results['inputs_to_nn'][i]], [loaded_results['scan_results'][key][f'{ikey}_nominal']], 'o', color=color) + axs[i].plot([loaded_results['inputs_to_eped'][i]], [loaded_results[ikey]], '^', color=color) + axs[i].plot([loaded_results['inputs_to_eped'][i]], [loaded_results['scan_results'][key][f'{ikey}_nominal']], 'o', color=color) - axs[i].axvline(loaded_results['inputs_to_nn'][i], color=color, ls='--') + axs[i].axvline(loaded_results['inputs_to_eped'][i], color=color, ls='--') axs[i].axhline(loaded_results['scan_results'][key][f'{ikey}_nominal'], color=color, ls='-.') max_val = np.max([max_val,np.max(loaded_results['scan_results'][key][ikey])]) @@ -488,12 +560,12 @@ def scale_profile_by_stretching( x, y, xp, yp, xp_old, plotYN=False, label='', k print('\t\t\t* Keeping old aLT profile in the core-predicted region, using r/a for it') # Calculate gradient in entire region - aLy = CALCtools.derivation_into_Lx( torch.from_numpy(roa), torch.from_numpy(y) ) + aLy = CALCtools.produceGradient( torch.from_numpy(roa), torch.from_numpy(y) ) # I'm only interested in core region, plus one ghost point with the same gradient aLy = torch.cat( (aLy[:ibc+1], aLy[ibc].unsqueeze(0)) ) - y_mod = CALCtools.integration_Lx( torch.from_numpy(roa[:ibc+2]).unsqueeze(0), aLy.unsqueeze(0), torch.from_numpy(np.array(ynew[ibc+1])).unsqueeze(0) ).squeeze().numpy() + y_mod = CALCtools.integrateGradient( torch.from_numpy(roa[:ibc+2]).unsqueeze(0), aLy.unsqueeze(0), torch.from_numpy(np.array(ynew[ibc+1])).unsqueeze(0) ).squeeze().numpy() ynew[:ibc+2] = y_mod @@ -511,11 +583,11 @@ def scale_profile_by_stretching( x, y, xp, yp, xp_old, plotYN=False, label='', k ax.legend() ax = axs[1] - aLy = CALCtools.derivation_into_Lx( torch.from_numpy(roa), torch.from_numpy(y) ) + aLy = CALCtools.produceGradient( torch.from_numpy(roa), torch.from_numpy(y) ) ax.plot(x,aLy,'-o',color='b', label='old') ax.axvline(x=xp_old,color='b',ls='--') - aLy = CALCtools.derivation_into_Lx( torch.from_numpy(roa), torch.from_numpy(ynew) ) + aLy = CALCtools.produceGradient( torch.from_numpy(roa), torch.from_numpy(ynew) ) ax.plot(x,aLy,'-o',color='r', label='new') ax.axvline(x=xp,color='r',ls='--') @@ -599,6 +671,6 @@ def eped_profiler(profiles, xp_old, rhotop, Tetop_keV, Titop_keV, netop_20, mini # Re-derive # --------------------------------- - profiles_output.derive_quantities(rederiveGeometry=False) + profiles_output.deriveQuantities(rederiveGeometry=False) return profiles_output \ No newline at end of file diff --git a/src/mitim_modules/maestro/utils/MAESTRObeat.py b/src/mitim_modules/maestro/utils/MAESTRObeat.py index bae64b52..7b1c281d 100644 --- a/src/mitim_modules/maestro/utils/MAESTRObeat.py +++ b/src/mitim_modules/maestro/utils/MAESTRObeat.py @@ -461,7 +461,8 @@ def __call__(self): self.beat_eped.profiles_current = self.initialize_instance.profiles_current # Run EPED - eped_results = self.beat_eped._run(loopBetaN = 1) + nproc_per_run = 64 #TODO: make it a parameter to be received from MAESTRO namelist + eped_results = self.beat_eped._run(loopBetaN = 1, nproc_per_run=nproc_per_run, cold_start=True) # Assume always cold start for a creator # Potentially save variables np.save(self.beat_eped.folder_output / 'eped_results.npy', eped_results) diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index e7548aa3..76f6812f 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -2,13 +2,13 @@ from collections import OrderedDict from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.misc_tools import LOGtools, GRAPHICStools -from mitim_tools.plasmastate_tools import MITIMstate -from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_tools.gs_tools import GEQtools from pathlib import Path from mitim_tools.misc_tools.LOGtools import printMsg as print +import json, re +from pathlib import Path +import matplotlib.pyplot as plt from IPython import embed - from mitim_modules.maestro.utils.TRANSPbeat import transp_beat from mitim_modules.maestro.utils.PORTALSbeat import portals_beat from mitim_modules.maestro.utils.EPEDbeat import eped_beat @@ -53,9 +53,9 @@ def plotMAESTRO(folder, fn = None, num_beats = 2, only_beats = None, full_plot = m = grabMAESTRO(folder) # Plot - m.plot(fn = fn, num_beats=num_beats, only_beats = only_beats, full_plot = full_plot) + ps, ps_lab = m.plot(fn = fn, num_beats=num_beats, only_beats = only_beats, full_plot = full_plot) - return m + return m, ps, ps_lab def plot_results(self, fn): @@ -75,7 +75,11 @@ def plot_results(self, fn): for i,beat in enumerate(self.beats.values()): - _, profs = beat.grab_output() + # _, profs = beat.grab_output() + if (beat.folder_output / 'input.gacode').exists(): + profs = PROFILEStools.gacode_state(beat.folder_output / 'input.gacode') + else: + profs = None if isinstance(beat, transp_beat): key = f'TRANSP b#{i+1}' @@ -98,10 +102,10 @@ def plot_results(self, fn): maxPlot = 5 if len(ps) > 0: # Plot profiles - figs = state_plotting.add_figures(fn,fnlab_pre = 'MAESTRO - ') + figs = PROFILEStools.add_figures(fn,fnlab_pre = 'MAESTRO - ') log_file = self.folder_logs/'plot_maestro.log' if (not self.terminal_outputs) else None with LOGtools.conditional_log_to_file(log_file=log_file): - state_plotting.plotAll(ps[-maxPlot:], extralabs=ps_lab[-maxPlot:], figs=figs) + PROFILEStools.plotAll(ps[-maxPlot:], extralabs=ps_lab[-maxPlot:], figs=figs) for p,pl in zip(ps,ps_lab): p.printInfo(label = pl) @@ -195,7 +199,24 @@ def plot_results(self, fn): DFHJ """ ) + + plot_special_quantities(ps, ps_lab, axs) + + if (self.folder_performance / 'timing.jsonl').exists(): + # ******************************************************************************************************** + # Timings + # ******************************************************************************************************** + fig = fn.add_figure(label='MAESTRO timings', tab_color=3) + axs = fig.subplot_mosaic(""" + A + B + """) + plot_timings(self.folder_performance / 'timing.jsonl', axs = axs) + + return ps, ps_lab +def plot_special_quantities(ps, ps_lab, axs, color='b', label = '', legYN=True): + x, BetaN, Pfus, p_th, p_tot, Pin, Q, fG, nu_ne, q95, q0, xsaw,p90 = [], [], [], [], [], [], [], [], [], [], [], [], [] for p,pl in zip(ps,ps_lab): x.append(pl) @@ -212,24 +233,37 @@ def plot_results(self, fn): xsaw.append(p.derived['rho_saw']) p90.append(np.interp(0.9,p.profiles['rho(-)'],p.derived['pthr_manual'])) + def _special(ax,x): + for xi in x: + if 'portals' in xi.lower(): + if legYN: + ax.axvline(xi, color='y', linestyle='-', lw=5, alpha=0.2) + # ----------------------------------------------------------------- ax = axs['A'] - ax.plot(x, BetaN, '-s', markersize=7, lw = 1) + ax.plot(x, BetaN, '-s', color=color, markersize=7, lw = 1, label = label) ax.set_ylabel('$\\beta_N$ (engineering)') ax.set_title('Pressure Evolution') + if len(label) > 0: + ax.legend() GRAPHICStools.addDenseAxis(ax) ax.set_ylim(bottom = 0) + + _special(ax, x) ax.set_xticklabels([]) ax = axs['D'] - ax.plot(x, p_th, '-s', markersize=7, lw = 1, label='Thermal

') - ax.plot(x, p_tot, '-o', markersize=7, lw = 1, label='Total

') - ax.plot(x, p90, '-*', markersize=7, lw = 1, label='Total, p(rho=0.9)') + ax.plot(x, p_th, '-s', color=color, markersize=7, lw = 1, label='Thermal

') + ax.plot(x, p_tot, '-o', color=color, markersize=7, lw = 1, label='Total

') + ax.plot(x, p90, '-*', color=color, markersize=7, lw = 1, label='Total, p(rho=0.9)') ax.set_ylabel('$p$ (MPa)') GRAPHICStools.addDenseAxis(ax) ax.set_ylim(bottom = 0) - ax.legend() + if legYN: + ax.legend() + + _special(ax, x) rotation = 90 fontsize = 6 @@ -238,35 +272,41 @@ def plot_results(self, fn): # ----------------------------------------------------------------- ax = axs['B'] - ax.plot(x, Q, '-s', markersize=7, lw = 1) + ax.plot(x, Q, '-s', color=color, markersize=7, lw = 1) ax.set_ylabel('$Q$') ax.set_title('Performance Evolution') GRAPHICStools.addDenseAxis(ax) ax.set_ylim(bottom = 0) ax.set_xticklabels([]) + + _special(ax, x) ax = axs['E'] - ax.plot(x, Pfus, '-s', markersize=7, lw = 1) + ax.plot(x, Pfus, '-s', color=color, markersize=7, lw = 1) ax.set_ylabel('$P_{fus}$ (MW)') GRAPHICStools.addDenseAxis(ax) ax.set_ylim(bottom = 0) ax.set_xticklabels([]) + + _special(ax, x) ax = axs['F'] - ax.plot(x, Pin, '-s', markersize=7, lw = 1) + ax.plot(x, Pin, '-s', color=color, markersize=7, lw = 1) ax.set_ylabel('$P_{in}$ (MW)') GRAPHICStools.addDenseAxis(ax) ax.set_ylim(bottom = 0) ax.tick_params(axis='x', rotation=rotation, labelsize=fontsize) + + _special(ax, x) # ----------------------------------------------------------------- ax = axs['G'] - ax.plot(x, fG, '-s', markersize=7, lw = 1) + ax.plot(x, fG, '-s', color=color, markersize=7, lw = 1) ax.set_ylabel('$f_{G}$') ax.set_title('Density Evolution') ax.axhline(y=1, color = 'k', lw = 1, ls = '--') @@ -275,37 +315,46 @@ def plot_results(self, fn): ax.set_ylim([0,1.2]) ax.set_xticklabels([]) + + _special(ax, x) ax = axs['H'] - ax.plot(x, nu_ne, '-s', markersize=7, lw = 1) + ax.plot(x, nu_ne, '-s', color=color, markersize=7, lw = 1) ax.set_ylabel('$\\nu_{ne}$') GRAPHICStools.addDenseAxis(ax) ax.set_ylim(bottom = 0) ax.tick_params(axis='x', rotation=rotation, labelsize=fontsize) + + _special(ax, x) # ----------------------------------------------------------------- # ----------------------------------------------------------------- ax = axs['I'] - ax.plot(x, q95, '-s', markersize=7, lw = 1, label='q95') - ax.plot(x, q0, '-*', markersize=7, lw = 1, label='q0') + ax.plot(x, q95, '-s', color=color, markersize=7, lw = 1, label='q95') + ax.plot(x, q0, '-*', color=color, markersize=7, lw = 1, label='q0') ax.set_ylabel('$q$') ax.set_title('Current Evolution') GRAPHICStools.addDenseAxis(ax) ax.axhline(y=1, color = 'k', lw = 2, ls = '--') - ax.legend() + if legYN: + ax.legend() ax.set_ylim(bottom = 0) ax.set_xticklabels([]) + + _special(ax, x) ax = axs['J'] - ax.plot(x, xsaw, '-s', markersize=7, lw = 1) + ax.plot(x, xsaw, '-s', color=color, markersize=7, lw = 1) ax.set_ylabel('Inversion radius (rho)') GRAPHICStools.addDenseAxis(ax) ax.set_ylim([0,1]) ax.tick_params(axis='x', rotation=rotation, labelsize=fontsize) + + _special(ax, x) # ----------------------------------------------------------------- @@ -315,3 +364,113 @@ def plot_g_quantities(g, axs, color = 'b', lw = 1, ms = 0): g.plotFluxSurfaces(ax=axs[0], fluxes=np.linspace(0, 1, 21), rhoPol=False, sqrt=True, color=color,lwB=lw*3, lw = lw,label='Initial geqdsk') axs[3].plot(g.g['RHOVN'], g.g['PRES']*1E-6, '-o', markersize=ms, lw = lw, label='Initial geqdsk', color=color) axs[4].plot(g.g['RHOVN'], g.g['QPSI'], '-o', markersize=ms, lw = lw, label='Initial geqdsk', color=color) + +# --------------------------------------------------------------------------- +def plot_timings(jsonl_path, axs = None, unit: str = "min", color = "b", label= ''): + """ + Plot cumulative durations from a .jsonl timing ledger written by @mitim_timer, + with vertical lines when the beat number changes. + + Parameters + ---------- + jsonl_path : str | Path + File with one JSON record per line. + unit : {"s", "min", "h"} + Unit for the y-axis. + """ + multiplier = {"s": 1, "min": 1 / 60, "h": 1 / 3600}[unit] + + scripts, script_time, cumulative, beat_nums, script_restarts = [], [], [], [], [] + running = 0.0 + beat_pat = re.compile(r"Beat\s*#\s*(\d+)") + + # ── read the file ─────────────────────────────────────────────────────── + with Path(jsonl_path).expanduser().open() as f: + for line in f: + if not line.strip(): + continue + rec = json.loads(line) + + if rec["script"] not in scripts: + + scripts.append(rec["script"]) + script_time.append(rec["duration_s"] * multiplier) + running += rec["duration_s"]* multiplier + cumulative.append(running) + + m = beat_pat.search(rec["script"]) + beat_nums.append(int(m.group(1)) if m else None) + + script_restarts.append(0.0) + + else: + # If the script is already in the list, it means it was restarted + idx = scripts.index(rec["script"]) + script_restarts[idx] += rec["duration_s"] * multiplier + + if not scripts: + raise ValueError(f"No records found in {jsonl_path}") + + beat_nums = [0] + beat_nums # Start with zero beat + scripts = ['ini'] + scripts # Add initial beat + script_time = [0.0] + script_time # Start with zero time + cumulative = [0.0] + cumulative # Start with zero time + script_restarts = [0.0] + script_restarts # Start with zero restarts + + # ── plot ──────────────────────────────────────────────────────────────── + x = list(range(len(scripts))) + + if axs is None: + plt.ion() + fig = plt.figure() + axs = fig.subplot_mosaic(""" + A + B + """) + + + ax = axs['A'] + ax.plot(x, cumulative, "-s", markersize=8, color=color, label=label) + + # Add restarts as vertical lines + for i in range(len(script_restarts)): + if script_restarts[i] > 0: + ax.plot( + [x[i],x[i]], + [cumulative[i],cumulative[i]+script_restarts[i]], + "-.o", markersize=5, color=color) + + + for i in range(1, len(beat_nums)): + if beat_nums[i] != beat_nums[i - 1]: + ax.axvline(i - 0.5, color='k',linestyle="-.") + + #ax.set_xlim(left=0) + ax.set_ylabel(f"Cumulative time ({unit})"); #ax.set_ylim(bottom=0) + ax.set_xticks(x, scripts, rotation=10, ha="right", fontsize=8) + GRAPHICStools.addDenseAxis(ax) + ax.legend(loc='upper left', fontsize=8) + + + ax = axs['B'] + for i in range(len(scripts)-1): + ax.plot([x[i], x[i+1]], [0, script_time[i+1]], "-s", markersize=8, color=color) + + # Add restarts as vertical lines + for i in range(len(script_restarts)-1): + if script_restarts[i] > 0: + ax.plot( + [x[i+1],x[i+1]], + [script_time[i+1],script_time[i+1]+script_restarts[i+1]], + "-.o", markersize=5, color=color) + + for i in range(1, len(beat_nums)): + if beat_nums[i] != beat_nums[i - 1]: + ax.axvline(i - 0.5, color='k',linestyle="-.") + + #ax.set_xlim(left=0) + ax.set_ylabel(f"Time ({unit})"); #ax.set_ylim(bottom=0) + ax.set_xticks(x, scripts, rotation=10, ha="right", fontsize=8) + GRAPHICStools.addDenseAxis(ax) + + return x, scripts diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index d4a1bae9..4491c732 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -115,7 +115,7 @@ def run(self, **kwargs): def _flux_match_for_first_point(self): - print('\t- Running flux match for first point') + print('\n\t- Running flux match for first point') # Flux-match first folder_fm = self.folder / 'flux_match' @@ -123,7 +123,13 @@ def _flux_match_for_first_point(self): portals = PORTALSanalysis.PORTALSanalyzer.from_folder(self.folder_starting_point) p = portals.powerstates[portals.ibest].profiles - _ = PORTALSoptimization.flux_match_surrogate(portals.step,p,file_write_csv=folder_fm / 'optimization_data.csv') + _ = PORTALSoptimization.flux_match_surrogate( + portals.step, + p, + TargetOptions_use = self.mitim_bo.optimization_object.powerstate.TargetOptions, # Use the TargetOptions of the new run, not the old one (which may be with fixed targets if soft) + file_write_csv=folder_fm / 'optimization_data.csv' + ) + # Move files (self.folder / 'Outputs').mkdir(parents=True, exist_ok=True) @@ -367,7 +373,7 @@ def _inform_save(self): self.maestro_instance.parameters_trans_beat['portals_neg_residual_obj'] = max_value_neg_residual print(f'\t\t* Maximum value of negative residual saved for future beats: {max_value_neg_residual}') - fileTraining = stepSettings['folderOutputs'] / 'surrogate_data.csv' + fileTraining = self.folder / 'Outputs/' / 'surrogate_data.csv' self.maestro_instance.parameters_trans_beat['portals_last_run_folder'] = self.folder self.maestro_instance.parameters_trans_beat['portals_surrogate_data_file'] = fileTraining diff --git a/src/mitim_modules/maestro/utils/TRANSPbeat.py b/src/mitim_modules/maestro/utils/TRANSPbeat.py index 29fb7217..74ed7245 100644 --- a/src/mitim_modules/maestro/utils/TRANSPbeat.py +++ b/src/mitim_modules/maestro/utils/TRANSPbeat.py @@ -104,7 +104,7 @@ def prepare( self._additional_operations_add_initialization() # ICRF on - PichT_MW = self.profiles_current.derived['qRF_MW'][-1] + PichT_MW = self.profiles_current.derived['qRF_MWmiller'][-1] if freq_ICH is None: @@ -156,8 +156,15 @@ def finalize(self, force_auxiliary_heating_at_output = {'Pe': None, 'Pi': None}, print('\t\t- No TRANSP files in beat folder, assuming they may exist in the output folder (MAESTRO restart case)', typeMsg='w') # Find CDF name - files = [f for f in self.folder_output.iterdir() if f.is_file()] - cdf_prefix = next((file.stem for file in files if file.suffix.lower() == '.cdf'), None) + files = [f for f in self.folder.iterdir() if f.is_file()] + cdf_prefix = next( + (file.stem + for file in files + if file.suffix.lower() == ".cdf" # keep only .cdf files … + and not file.name.lower().endswith("ph.cdf")), # … but skip *.ph.cdf + None + ) + shutil.copy2(self.folder / f"{cdf_prefix}TR.DAT", self.folder_output / f"{self.shot}{self.runid}TR.DAT") shutil.copy2(self.folder / f"{cdf_prefix}.CDF", self.folder_output / f"{self.shot}{self.runid}.CDF") shutil.copy2(self.folder / f"{cdf_prefix}tr.log", self.folder_output / f"{self.shot}{self.runid}tr.log") @@ -174,18 +181,18 @@ def finalize(self, force_auxiliary_heating_at_output = {'Pe': None, 'Pi': None}, self._add_heating_profiles(force_auxiliary_heating_at_output) # Write profiles - self.profiles_output.write_state(file=self.folder_output / "input.gacode") + self.profiles_output.writeCurrentStatus(file=self.folder_output / "input.gacode") def _add_heating_profiles(self, force_auxiliary_heating_at_output = {'Pe': None, 'Pi': None}): ''' force_auxiliary_heating_at_output['Pe'] has the shaping function (takes rho) and the integrated value ''' - for key, pkey, ikey in zip(['Pe','Pi'], ['qrfe(MW/m^3)', 'qrfi(MW/m^3)'], ['qRFe_MW', 'qRFi_MW']): + for key, pkey, ikey in zip(['Pe','Pi'], ['qrfe(MW/m^3)', 'qrfi(MW/m^3)'], ['qRFe_MWmiller', 'qRFi_MWmiller']): if force_auxiliary_heating_at_output[key] is not None: self.profiles_output.profiles[pkey] = force_auxiliary_heating_at_output[key][0](self.profiles_output.profiles['rho(-)']) - self.profiles_output.derive_quantities() + self.profiles_output.deriveQuantities() self.profiles_output.profiles[pkey] = self.profiles_output.profiles[pkey] * force_auxiliary_heating_at_output[key][1]/self.profiles_output.derived[ikey][-1] def merge_parameters(self): @@ -206,7 +213,7 @@ def merge_parameters(self): # Write the pre-merge input.gacode before modifying it profiles_output_pre_merge = copy.deepcopy(self.profiles_output) - profiles_output_pre_merge.write_state(file=self.folder_output / 'input.gacode_pre_merge') + profiles_output_pre_merge.writeCurrentStatus(file=self.folder_output / 'input.gacode_pre_merge') # First, bring back to the resolution of the frozen p_frozen = self.maestro_instance.profiles_with_engineering_parameters @@ -230,14 +237,14 @@ def merge_parameters(self): self.profiles_output.profiles[key] = p_frozen.profiles[key] # Power scale - self.profiles_output.profiles['qrfe(MW/m^3)'] *= p_frozen.derived['qRF_MW'][-1] / self.profiles_output.derived['qRF_MW'][-1] - self.profiles_output.profiles['qrfi(MW/m^3)'] *= p_frozen.derived['qRF_MW'][-1] / self.profiles_output.derived['qRF_MW'][-1] + self.profiles_output.profiles['qrfe(MW/m^3)'] *= p_frozen.derived['qRF_MWmiller'][-1] / self.profiles_output.derived['qRF_MWmiller'][-1] + self.profiles_output.profiles['qrfi(MW/m^3)'] *= p_frozen.derived['qRF_MWmiller'][-1] / self.profiles_output.derived['qRF_MWmiller'][-1] # -------------------------------------------------------------------------------------------- # Write to final input.gacode - self.profiles_output.derive_quantities() - self.profiles_output.write_state(file=self.folder_output / 'input.gacode') + self.profiles_output.deriveQuantities() + self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') def grab_output(self): diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 5f25103e..bf32e92c 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -23,7 +23,7 @@ def surrogate_selection_portals(output, surrogate_options, CGYROrun=False): surrogate_options["TypeKernel"] = 1 # RBF surrogate_options["additional_constraints"] = { - 'lenghtscale_constraint': gpytorch.constraints.constraints.GreaterThan(0.01) # inputs normalized to [0,1], this is 1% lengthscale + 'lenghtscale_constraint': gpytorch.constraints.constraints.GreaterThan(0.05) # inputs normalized to [0,1], this is 5% lengthscale } return surrogate_options diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index ef5f38d1..e9168ee1 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -92,7 +92,17 @@ def initialization_simple_relax(self): """ -def flux_match_surrogate(step,profiles, plot_results=False, fn = None, file_write_csv=None, algorithm = None, solver_options = None, keep_within_bounds = True): +def flux_match_surrogate( + step, + profiles, + plot_results=False, + fn = None, + file_write_csv=None, + algorithm = None, + solver_options = None, + keep_within_bounds = True, + TargetOptions_use = None, + ): ''' Technique to reutilize flux surrogates to predict new conditions ---------------------------------------------------------------- @@ -142,7 +152,7 @@ def flux_match_surrogate(step,profiles, plot_results=False, fn = None, file_writ "fineTargetsResolution": step.surrogate_parameters["powerstate"].fineTargetsResolution, }, TransportOptions=TransportOptions, - TargetOptions=step.surrogate_parameters["powerstate"].TargetOptions, + TargetOptions= step.surrogate_parameters["powerstate"].TargetOptions if TargetOptions_use is None else TargetOptions_use, tensor_opts = { "dtype": step.surrogate_parameters["powerstate"].dfT.dtype, "device": step.surrogate_parameters["powerstate"].dfT.device}, diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index b89a759a..592ebff2 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -2027,93 +2027,97 @@ def PORTALSanalyzer_plotModelComparison( metrics = {} # 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" - metrics["Qe"] = plotModelComparison_quantity( - self, - axs[cont], - 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, - ) + if 'te' in self.ProfilesPredicted: + 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" + metrics["Qe"] = plotModelComparison_quantity( + self, + axs[cont], + 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[cont].set_xscale("log") - axs[cont].set_yscale("log") + axs[cont].set_xscale("log") + axs[cont].set_yscale("log") - cont += 1 + cont += 1 # 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" - metrics["Qi"] = plotModelComparison_quantity( - self, - axs[cont], - 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, - ) + if 'ti' in self.ProfilesPredicted: + 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" + metrics["Qi"] = plotModelComparison_quantity( + self, + axs[cont], + 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[cont].set_xscale("log") - axs[cont].set_yscale("log") + axs[cont].set_xscale("log") + axs[cont].set_yscale("log") - cont += 1 + cont += 1 # 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" - metrics["Ge"] = plotModelComparison_quantity( - self, - axs[cont], - 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 'ne' in self.ProfilesPredicted: + 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" + metrics["Ge"] = plotModelComparison_quantity( + self, + axs[cont], + 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]["powerstate"].model_results.__dict__[quantityX][0, 1:] - else: - val_calc = np.array( - [ - self.tglf_full.results["ev0"]["TGLFout"][j].__dict__[ - quantityX.replace("[TGLF]", "") + if UseTGLFfull_x is None: + val_calc = self.mitim_runs[0]["powerstate"].model_results.__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)) ] - for j in range(len(self.rhos)) - ] - ) + ) - try: - thre = 10 ** round(np.log10(np.abs(val_calc).min())) - axs[cont].set_xscale("symlog", linthresh=thre) - axs[cont].set_yscale("symlog", linthresh=thre) - # axs[2].tick_params(axis="both", which="major", labelsize=8) - except OverflowError: - pass + try: + thre = 10 ** round(np.log10(np.abs(val_calc).min())) + axs[cont].set_xscale("symlog", linthresh=thre) + axs[cont].set_yscale("symlog", linthresh=thre) + # axs[2].tick_params(axis="both", which="major", labelsize=8) + except OverflowError: + pass + + cont += 1 - cont += 1 if "nZ" in self.ProfilesPredicted: diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index f7744033..d742839b 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -333,7 +333,7 @@ def modify(self, X): self.update_var(i) - def flux_match(self, algorithm="root", solver_options=None, bounds=None): + def flux_match(self, algorithm="root", solver_options=None, bounds=None, debugYN=False): self.FluxMatch_plasma_orig = copy.deepcopy(self.plasma) self.bounds_current = bounds @@ -423,6 +423,10 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): print("**********************************************************************************************") print(f"\t- Flux matching of powerstate finished, and took {IOtools.getTimeDifference(timeBeginning)}\n") + if debugYN: + self.plot() + embed() + # ------------------------------------------------------------------ # Plotting tools # ------------------------------------------------------------------ diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 41c7ece1..6ce070e1 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -196,7 +196,7 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod file=file_cgyro, ) ) - print(f"\t- Suggested function call for mitim evaluation {numPORTALS} (lambda for cgyroing):",typeMsg="i") + print(f"\t- Suggested function call for PORTALS evaluation {numPORTALS} (lambda for cgyroing):",typeMsg="i") cgyropath = IOtools.expandPath(folder, ensurePathValid=True) / 'Outputs' / 'cgyro_results' / f'cgyro_it_{numPORTALS}.txt' print(f"\tcgyroing_file('{cgyropath}')") @@ -265,7 +265,7 @@ def cgyroing( FolderEvaluation, minErrorPercent=minErrorPercent, Qi_criterion_stable=Qi_criterion_stable, - percentNeo=percentNeo, + percent_tr_neo=percentNeo, impurityPosition=impurityPosition, OriginalFimp=OriginalFimp, ) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 530cb5b5..34d1ed25 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -189,7 +189,7 @@ def _profiles_to_store(self): fil = whereFolder / f"input.gacode.{self.evaluation_number}" shutil.copy2(self.file_profs, fil) shutil.copy2(self.file_profs_unmod, fil.parent / f"{fil.name}_unmodified") - shutil.copy2(self.file_profs_targets, fil.parent / f"{fil.name}.new") + shutil.copy2(self.file_profs_targets, fil.parent / f"{fil.name}.new_fromtgyro_modified") print(f"\t- Copied profiles to {IOtools.clipstr(fil)}") else: print("\t- Could not move files", typeMsg="w") diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index 61cd0ed5..c5607723 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -5,7 +5,7 @@ import matplotlib.pyplot as plt import f90nml from pathlib import Path -from mitim_tools.misc_tools import FARMINGtools, GRAPHICStools, IOtools +from mitim_tools.misc_tools import FARMINGtools, GRAPHICStools, IOtools, GUItools import numpy as np import pandas as pd import xarray as xr @@ -18,9 +18,10 @@ def __init__( folder ): - self.folder = Path(folder) + self.folder = Path(folder) if folder is not None else None # None for just reading - self.folder.mkdir(parents=True, exist_ok=True) + if self.folder is not None: + self.folder.mkdir(parents=True, exist_ok=True) self.results = {} @@ -200,8 +201,10 @@ def read( ): self.results[label if label is not None else subfolder] = {} - - output_files = sorted(list((self.folder / subfolder).glob("*.nc"))) + + where_is_this = self.folder / subfolder if self.folder is not None else Path(subfolder) + + output_files = sorted(list(where_is_this.glob("*.nc"))) for output_file in output_files: @@ -236,9 +239,13 @@ def print(self,label,sublabel): def plot( self, labels = ['run1'], + axs = None, ): - plt.ion(); fig, axs = plt.subplots(2, 1, figsize=(10, 6)) + if axs is None: + self.fn = GUItools.FigureNotebook("EPED", geometry="900x900") + fig = self.fn.add_figure(label="Pedestal Top") + axs = fig.subplots(2, 1) colors = GRAPHICStools.listColors() diff --git a/src/mitim_tools/eped_tools/scripts/plot_eped.py b/src/mitim_tools/eped_tools/scripts/plot_eped.py new file mode 100644 index 00000000..42ce68ec --- /dev/null +++ b/src/mitim_tools/eped_tools/scripts/plot_eped.py @@ -0,0 +1,27 @@ +import argparse +from mitim_tools.misc_tools import IOtools +from mitim_tools.eped_tools import EPEDtools +from IPython import embed + +def main(): + + parser = argparse.ArgumentParser() + parser.add_argument("folders", type=str, nargs="*") + + args = parser.parse_args() + + folders = [IOtools.expandPath(folder) for folder in args.folders] + + eped = EPEDtools.EPED(folder=None) + + for i, folder in enumerate(folders): + eped.read(subfolder=folder, label=f"run{i}") + + eped.plot(labels=[f"run{i}" for i in range(len(folders))]) + + eped.fn.show() + + embed() + +if __name__ == "__main__": + main() diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 1e735677..bd262b8b 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -1,5 +1,4 @@ import os -from re import sub import shutil import datetime import time @@ -10,11 +9,8 @@ from mitim_tools.misc_tools import IOtools, GRAPHICStools, FARMINGtools from mitim_tools.gacode_tools.utils import GACODEplotting from mitim_tools.misc_tools.LOGtools import printMsg as print -from pygacode.cgyro.data_plot import cgyrodata_plot -from pygacode import gacodefuncs from IPython import embed - class CGYRO: def __init__(self): @@ -34,6 +30,8 @@ def __init__(self): "bin.cgyro.kxky_e", "bin.cgyro.kxky_n", "bin.cgyro.kxky_phi", + "bin.cgyro.kxky_apar", + "bin.cgyro.kxky_bpar", "bin.cgyro.kxky_v", "bin.cgyro.ky_cflux", "bin.cgyro.ky_flux", @@ -110,7 +108,7 @@ def _prerun( control_file = 'input.cgyro.controls' ) - inputCGYRO.write_state() + inputCGYRO.writeCurrentStatus() return input_cgyro_file, inputgacode_file_this @@ -179,6 +177,9 @@ def run_full( minutes = 5, n = 16, nomp = 1, + cpuspertask=None, # if None, will default to 1 + queue=None, #if blank will default to the one in settings + mem=None, # in MB submit_via_qsub=True, #TODO fix this, works only at NERSC? no scans? clean_folder_going_in=True, # Make sure the scratch folder is removed before running (unless I want a restart!) submit_run=True, # False if I just want to check and fetch the job that was already submitted (e.g. via qsub or slurm) @@ -210,7 +211,10 @@ def run_full( subfolder = "scan0" queue = "-queue " + self.cgyro_job.machineSettings['slurm']['partition'] if "partition" in self.cgyro_job.machineSettings['slurm'] else "" - CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} {queue} -w 0:{minutes}:00 -s' + if mem is not None: + CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} -mem {mem} {queue} -w 0:{minutes}:00 -s' + else: + CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} {queue} -w 0:{minutes}:00 -s' if "account" in self.cgyro_job.machineSettings["slurm"] and self.cgyro_job.machineSettings["slurm"]["account"] is not None: CGYROcommand += f" -repo {self.cgyro_job.machineSettings['slurm']['account']}" @@ -258,11 +262,24 @@ def run_full( "minutes": minutes, "ntasks": n, "job_array": job_array, + # Validate n and nomp before assigning cpuspertask }, ) - if not self.cgyro_job.launchSlurm: - raise Exception(" Cannot run CGYRO scans without slurm") + if cpuspertask is not None: + if not isinstance(cpuspertask, int): + raise TypeError(" cpuspertask must be an integer") + self.cgyro_job.slurm_settings["cpuspertask"] = cpuspertask + + + if queue is not None: + self.cgyro_job.machineSettings['slurm']['partition'] = queue + + if mem is not None: + self.cgyro_job.slurm_settings['mem'] = mem + + # if not self.cgyro_job.launchSlurm: + # raise Exception(" Cannot run CGYRO scans without slurm") # Command to run cgyro CGYROcommand = f'cgyro -e {folder} -n {n} -nomp {nomp} -p {self.cgyro_job.folderExecution}' @@ -301,7 +318,7 @@ def run_full( control_file = 'input.cgyro.controls' ) - input_cgyro_file_this.write_state() + input_cgyro_file_this.writeCurrentStatus() # Copy the input.gacode file in the subfolder inputgacode_file_this = folder_run / "input.gacode" @@ -342,9 +359,7 @@ def check(self, every_n_minutes=5): while True: self.cgyro_job.check(file_output = self.slurm_output) - print( - f'\t- Current status (as of {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}): {self.cgyro_job.status} ({self.cgyro_job.infoSLURM["STATE"]})' - ) + print(f'\t- Current status (as of {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}): {self.cgyro_job.status} ({self.cgyro_job.infoSLURM["STATE"]})') if self.cgyro_job.status == 2: break else: @@ -386,12 +401,12 @@ def delete(self): # Reading and plotting # --------------------------------------------------------------------------------------------------------- - def read(self, label="cgyro1", folder=None, tmin = 0.0): + def read(self, label="cgyro1", folder=None, tmin = 0.0, minimal = False, last_tmin_for_linear = True): folder = IOtools.expandPath(folder) if folder is not None else self.folderCGYRO - folders = sorted(list((folder).glob("scan*"))) - + folders = sorted(list((folder).glob("scan*")), key=lambda f: int(''.join(filter(str.isdigit, f.name)))) + if len(folders) == 0: folders = [folder] attach_name = False @@ -401,92 +416,314 @@ def read(self, label="cgyro1", folder=None, tmin = 0.0): print(f"\t\t- {f.name}") attach_name = True + data = {} + labels = [] for folder in folders: - - original_dir = os.getcwd() - + if attach_name: - label_new = f"{label}_{folder.name}" + label1 = f"{label}_{folder.name}" else: - label_new = label - - try: - print(f"\t- Reading CGYRO data from {folder.resolve()}") - self.results[label_new] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") - except: - if print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): - os.chdir(folder) - os.system("cgyro -t") - self.results[label_new] = cgyrodata_plot(f"{folder.resolve()}{os.sep}") - - os.chdir(original_dir) - - try: - self._postprocess(label_new, folder, tmin=tmin) - except Exception as e: - print(f"\t- Error during postprocessing: {e}") - print("\t- Skipping postprocessing, results may be incomplete or linear run") - - def _postprocess(self, label, folder, tmin=0.0): - - # Extra postprocessing - self.results[label].electron_flag = np.where(self.results[label].z == -1)[0][0] - self.results[label].all_flags = np.arange(0, len(self.results[label].z), 1) - self.results[label].ions_flags = self.results[label].all_flags[self.results[label].all_flags != self.results[label].electron_flag] + label1 = label - self.results[label].all_names = [f"{gacodefuncs.specmap(self.results[label].mass[i],self.results[label].z[i])}({self.results[label].z[i]},{self.results[label].mass[i]:.1f})" for i in self.results[label].all_flags] + data[label1] = CGYROutils.CGYROout(folder, tmin=tmin, minimal=minimal, last_tmin_for_linear=last_tmin_for_linear) + labels.append(label1) - # ************************ - # Calculations - # ************************ - - cgyro = self.results[label] - cgyro.getflux() - cgyro.getnorm("elec") - - self.results[label].t = self.results[label].tnorm - - ys = np.sum(cgyro.ky_flux, axis=(2, 3)) - - self.results[label].Qe = ys[-1, 1, :] / cgyro.qc - self.results[label].Qi_all = ys[:-1, 1, :] / cgyro.qc - self.results[label].Qi = self.results[label].Qi_all.sum(axis=0) - self.results[label].Ge = ys[-1, 0, :] - - roa,alne,self.results[label].aLTi,alte,self.results[label].Qi_mean,self.results[label].Qi_std,self.results[label].Qe_mean,self.results[label].Qe_std,self.results[label].Ge_mean, self.results[label].Ge_std,m_gimp,std_gimp,m_mo,std_mo,m_tur,std_tur,qgb,ggb,pgb,sgb,tstart,nt = CGYROutils.grab_cgyro_nth(str(folder.resolve()), tmin, False, False) - - self.results[label].Qi_mean *= qgb - self.results[label].Qi_std *= qgb - self.results[label].Qe_mean *= qgb - self.results[label].Qe_std *= qgb - - - def derive_statistics(self,x,y,x_min=0.0): + self.results.update(data) + + if attach_name: + self.results[label] = CGYROutils.CGYROlinear_scan(labels, data) - return y.mean(), y.std() + def plot(self, labels=[""], include_2D=True, common_colorbar=True): + - def plot(self, labels=[""]): + # If it has scans, we need to correct the labels + labels_corrected = [] + for i in range(len(labels)): + if isinstance(self.results[labels[i]], CGYROutils.CGYROlinear_scan): + for scan_label in self.results[labels[i]].labels: + labels_corrected.append(scan_label) + else: + labels_corrected.append(labels[i]) + labels = labels_corrected + # ------------------------------------------------ + + from mitim_tools.misc_tools.GUItools import FigureNotebook - self.fn = FigureNotebook("CGYRO Notebook", geometry="1600x1000") - colors = GRAPHICStools.listColors() - - fig = self.fn.add_figure(label="Fluxes Time Traces") + fig = self.fn.add_figure(label="Fluxes (time)") axsFluxes_t = fig.subplot_mosaic( """ AC BD """ ) + fig = self.fn.add_figure(label="Fluxes (ky)") + axsFluxes_ky = fig.subplot_mosaic( + """ + AC + BD + """ + ) + fig = self.fn.add_figure(label="Intensities (time)") + axsIntensities = fig.subplot_mosaic( + """ + ACEG + BDFH + """ + ) + fig = self.fn.add_figure(label="Intensities (ky)") + axsIntensities_ky = fig.subplot_mosaic( + """ + ACEG + BDFH + """ + ) + fig = self.fn.add_figure(label="Intensities (kx)") + axsIntensities_kx = fig.subplot_mosaic( + """ + AC + BD + """ + ) + fig = self.fn.add_figure(label="Cross-phases (ky)") + axsCrossPhases = fig.subplot_mosaic( + """ + ACEG + BDFH + """ + ) + fig = self.fn.add_figure(label="Turbulence (linear)") + axsTurbulence = fig.subplot_mosaic( + """ + AC + BD + """ + ) + + create_ballooning = False + for label in labels: + if 'phi_ballooning' in self.results[label].__dict__: + create_ballooning = True + + if create_ballooning: + + fig = self.fn.add_figure(label="Ballooning") + axsBallooning = fig.subplot_mosaic( + """ + 135 + 246 + """ + ) + else: + axsBallooning = None + + + if include_2D: + axs2D = [] + for i in range(len(labels)): + fig = self.fn.add_figure(label="Turbulence (2D), " + labels[i]) + + mosaic = _2D_mosaic(4) # Plot 4 times by default + + axs2D.append(fig.subplot_mosaic(mosaic)) + + fig = self.fn.add_figure(label="Inputs") + axsInputs = fig.subplot_mosaic( + """ + A + B + """ + ) + + + colors = GRAPHICStools.listColors() + colorbars_all = [] # Store all colorbars for later use for j in range(len(labels)): + self.plot_fluxes( axs=axsFluxes_t, label=labels[j], c=colors[j], plotLegend=j == len(labels) - 1, ) + self.plot_fluxes_ky( + axs=axsFluxes_ky, + label=labels[j], + c=colors[j], + plotLegend=j == len(labels) - 1, + ) + self.plot_intensities_ky( + axs=axsIntensities_ky, + label=labels[j], + c=colors[j], + addText=j == len(labels) - 1, + ) + self.plot_intensities( + axs=axsIntensities, + label=labels[j], + c=colors[j], + addText=j == len(labels) - 1, # Add text only for the last label + ) + self.plot_intensities_kx( + axs=axsIntensities_kx, + label=labels[j], + c=colors[j], + addText=j == len(labels) - 1, # Add text only for the last label + ) + self.plot_turbulence( + axs=axsTurbulence, + label=labels[j], + c=colors[j], + ) + self.plot_cross_phases( + axs=axsCrossPhases, + label=labels[j], + c=colors[j], + ) + if create_ballooning: + self.plot_ballooning( + axs=axsBallooning, + label=labels[j], + c=colors[j], + ) + + if include_2D: + + colorbars = self.plot_2D( + axs=axs2D[j], + label=labels[j], + ) + + colorbars_all.append(colorbars) + + self.plot_inputs( + ax=axsInputs["A"], + label=labels[j], + c=colors[j], + ms= 10-j*0.5, # Decrease marker size for each label + normalization_label= labels[0], # Normalize to the first label + only_plot_differences=len(labels) > 1, # Only plot differences if there are multiple labels + ) + + self.plot_inputs( + ax=axsInputs["B"], + label=labels[j], + c=colors[j], + ms= 10-j*0.5, # Decrease marker size for each label + ) + + axsInputs["A"].axhline( + 1.0, + color="k", + ls="--", + lw=2.0 + ) + + GRAPHICStools.adjust_subplots(axs=axsInputs, vertical=0.4, horizontal=0.3) + + # Modify the colorbars to have a common range + if include_2D and common_colorbar and len(colorbars_all) > 0: + for var in ['phi', 'n', 'e']: + min_val = np.inf + max_val = -np.inf + for ilabel in range(len(colorbars_all)): + cb = colorbars_all[ilabel][0][var] + vals = cb.mappable.get_clim() + min_val = min(min_val, vals[0]) + max_val = max(max_val, vals[1]) + + for ilabel in range(len(colorbars_all)): + for it in range(len(colorbars_all[ilabel])): + cb = colorbars_all[ilabel][it][var] + cb.mappable.set_clim(min_val, max_val) + cb.update_ticks() + #cb.set_label(f"{var} (common range)") + + def _plot_trace(self, ax, label, variable, c="b", lw=1, ls="-", label_plot='', meanstd=True, var_meanstd= None): + + t = self.results[label].t + + if not isinstance(variable, str): + z = variable + if var_meanstd is not None: + z_mean = var_meanstd[0] + z_std = var_meanstd[1] + + else: + z = self.results[label].__dict__[variable] + if meanstd and (f'{variable}_mean' in self.results[label].__dict__): + z_mean = self.results[label].__dict__[variable + '_mean'] + z_std = self.results[label].__dict__[variable + '_std'] + else: + z_mean = None + z_std = None + + ax.plot( + t, + z, + ls=ls, + lw=lw, + c=c, + label=label_plot, + ) + + if meanstd and z_std>0.0: + GRAPHICStools.fillGraph( + ax, + t[t>self.results[label].tmin], + z_mean, + y_down=z_mean + - z_std, + y_up=z_mean + + z_std, + alpha=0.1, + color=c, + lw=0.5, + islwOnlyMean=True, + label=label_plot + f" $\\mathbf{{{z_mean:.2f} \\pm {z_std:.2f}}}$ (1$\\sigma$)", + ) + + def plot_inputs(self, ax = None, label="", c="b", ms = 10, normalization_label=None, only_plot_differences=False): + + if ax is None: + plt.ion() + fig, ax = plt.subplots(1, 1, figsize=(18, 9)) + + rel_tol = 1e-2 + + legadded = False + for i, ikey in enumerate(self.results[label].params1D): + + z = self.results[label].params1D[ikey] + + if normalization_label is not None: + z0 = self.results[normalization_label].params1D[ikey] + zp = z/z0 if z0 != 0 else 0 + label_plot = f"{label} / {normalization_label}" + else: + label_plot = label + zp = z + + if (not only_plot_differences) or (not np.isclose(z, z0, rtol=rel_tol)): + ax.plot(ikey,zp,'o',markersize=ms,color=c,label=label_plot if not legadded else '') + legadded = True + + if normalization_label is not None: + if only_plot_differences: + ylabel = f"Parameters (DIFFERENT by {rel_tol*100:.2f}%) relative to {normalization_label}" + else: + ylabel = f"Parameters relative to {normalization_label}" + else: + ylabel = "Parameters" + + ax.set_xlabel("Parameter") + ax.tick_params(axis='x', rotation=60) + ax.set_ylabel(ylabel) + GRAPHICStools.addDenseAxis(ax) + if legadded: + ax.legend(loc='best') def plot_fluxes(self, axs=None, label="", c="b", lw=1, plotLegend=True): if axs is None: @@ -504,230 +741,1128 @@ def plot_fluxes(self, axs=None, label="", c="b", lw=1, plotLegend=True): # Electron energy flux ax = axs["A"] - ax.plot( - self.results[label].t, - self.results[label].Qe, - ls=ls[0], - lw=lw, - c=c, - label=f"{label}, electron", - ) - ax.set_xlabel("$t$ ($a/c_s$)") + self._plot_trace(ax,label,"Qe",c=c,lw=lw,ls=ls[0],label_plot=f"{label}, Total") + self._plot_trace(ax,label,"Qe_EM",c=c,lw=lw,ls=ls[1],label_plot=f"{label}, EM ($A_\\parallel$+$A_\\perp$)", meanstd=False) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) ax.set_ylabel("$Q_e$ (GB)") GRAPHICStools.addDenseAxis(ax) ax.set_title('Electron energy flux') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) - # Ion energy fluxes + # Electron particle flux ax = axs["B"] - ax.plot( - self.results[label].t, - self.results[label].Qi, - ls=ls[0], - lw=lw, - c=c, - label=f"{label}", - ) + self._plot_trace(ax,label,"Ge",c=c,lw=lw,ls=ls[0],label_plot=f"{label}, Total") + self._plot_trace(ax,label,"Ge_EM",c=c,lw=lw,ls=ls[1],label_plot=f"{label}, EM ($A_\\parallel$+$A_\\perp$)", meanstd=False) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\Gamma_e$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron particle flux') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + + # Ion energy fluxes + ax = axs["C"] + self._plot_trace(ax,label,"Qi",c=c,lw=lw,ls=ls[0],label_plot=f"{label}, Total") + self._plot_trace(ax,label,"Qi_EM",c=c,lw=lw,ls=ls[1],label_plot=f"{label}, EM ($A_\\parallel$+$A_\\perp$)", meanstd=False) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) ax.set_ylabel("$Q_i$ (GB)") GRAPHICStools.addDenseAxis(ax) ax.set_title('Ion energy fluxes') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) - ax = axs["C"] + # Ion species energy fluxes + ax = axs["D"] for j, i in enumerate(self.results[label].ions_flags): - ax.plot( - self.results[label].t, - self.results[label].Qi_all[j], - ls=ls[j + 1], - lw=lw / 2, - c=c, - label=f"{label}, {self.results[label].all_names[i]}", - ) + self._plot_trace(ax,label,self.results[label].Qi_all[j],c=c,lw=lw,ls=ls[j],label_plot=f"{label}, {self.results[label].all_names[i]}", meanstd=False) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) ax.set_ylabel("$Q_i$ (GB)") GRAPHICStools.addDenseAxis(ax) ax.set_title('Ion energy fluxes (separate species)') - GRAPHICStools.addLegendApart(ax,ratio=0.95, withleg=True, size = 8) + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_fluxes_ky(self, axs=None, label="", c="b", lw=1, plotLegend=True): + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + AC + BD + """ + ) + + ls = GRAPHICStools.listLS() + + # Electron energy flux + ax = axs["A"] + ax.plot(self.results[label].ky, self.results[label].Qe_ky_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Qe_ky_mean-self.results[label].Qe_ky_std, self.results[label].Qe_ky_mean+self.results[label].Qe_ky_std, color=c, alpha=0.2) + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$Q_e$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron energy flux vs. $k_\\theta\\rho_s$') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) # Electron particle flux - ax = axs["D"] - ax.plot( - self.results[label].t, - self.results[label].Ge, - ls=ls[0], - lw=lw, - c=c, - label=f"{label}, electron", - ) + ax = axs["B"] + ax.plot(self.results[label].ky, self.results[label].Ge_ky_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Ge_ky_mean-self.results[label].Ge_ky_std, self.results[label].Ge_ky_mean+self.results[label].Ge_ky_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") ax.set_ylabel("$\\Gamma_e$ (GB)") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Electron particle flux') + ax.set_title('Electron particle flux vs. $k_\\theta\\rho_s$') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + # Ion energy flux + ax = axs["C"] + ax.plot(self.results[label].ky, self.results[label].Qi_ky_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Qi_ky_mean-self.results[label].Qi_ky_std, self.results[label].Qi_ky_mean+self.results[label].Qi_ky_std, color=c, alpha=0.2) + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$Q_i$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion energy fluxes vs. $k_\\theta\\rho_s$') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) - def plotLS(self, labels=["cgyro1"], fig=None): - colors = GRAPHICStools.listColors() + # Ion species energy fluxes + ax = axs["D"] + for j, i in enumerate(self.results[label].ions_flags): + ax.plot(self.results[label].ky, self.results[label].Qi_all_ky_mean[j],ls[j]+'o', markersize=5, color=c, label=f"{label}, {self.results[label].all_names[i]}") - if fig is None: - # fig = plt.figure(figsize=(15,9)) + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$Q_i$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion energy fluxes vs. $k_\\theta\\rho_s$(separate species)') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) - from mitim_tools.misc_tools.GUItools import FigureNotebook - self.fn = FigureNotebook( - "Linear CGYRO Notebook", - geometry="1600x1000", + def plot_intensities(self, axs = None, label= "cgyro1", c="b", addText=True): + + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + ACEG + BDFH + """ ) - fig1 = self.fn.add_figure(label="Linear Stability") - fig2 = self.fn.add_figure(label="Ballooning") - - grid = plt.GridSpec(2, 2, hspace=0.3, wspace=0.3) - ax00 = fig1.add_subplot(grid[0, 0]) - ax10 = fig1.add_subplot(grid[1, 0], sharex=ax00) - ax01 = fig1.add_subplot(grid[0, 1]) - ax11 = fig1.add_subplot(grid[1, 1], sharex=ax01) - - K, G, F = [], [], [] - for cont, label in enumerate(self.results): - c = self.results[label] - baseColor = colors[cont] - colorsC, _ = GRAPHICStools.colorTableFade( - len(c.ky), - startcolor=baseColor, - endcolor=baseColor, - alphalims=[1.0, 0.4], + + ls = GRAPHICStools.listLS() + + ax = axs["A"] + ax.plot(self.results[label].t, self.results[label].phi_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].phi_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].phi_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta \\phi/\\phi_0$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Potential intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta\phi/\phi_0|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + ax = axs["B"] + if 'apar' in self.results[label].__dict__: + ax.plot(self.results[label].t, self.results[label].apar_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}, $A_\\parallel$") + ax.plot(self.results[label].t, self.results[label].bpar_rms_sumnr_sumn*100.0, '--', c=c, lw=2, label=f"{label}, $B_\\parallel$") + ax.legend(loc='best', prop={'size': 8},) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta F_\\parallel/F_{\\parallel,0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('EM potential intensity fluctuations') + + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta F_\parallel/F_{\parallel,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + + ax = axs["C"] + ax.plot(self.results[label].t, self.results[label].ne_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].ne_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].ne_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta n_e/n_{e,0}/n_{e0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron Density intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta n_e/n_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + + ax = axs["D"] + ax.plot(self.results[label].t, self.results[label].Te_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].Te_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].Te_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta T_e/T_{e,0}/T_{e0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron Temperature intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta T_e/T_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + + + ax = axs["E"] + ax.plot(self.results[label].t, self.results[label].ni_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].ni_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].ni_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta n_i/n_{i,0}/n_{i0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion Density intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta n_i/n_{i,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + + ax = axs["F"] + ax.plot(self.results[label].t, self.results[label].Ti_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].Ti_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].Ti_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta T_i/T_{i,0}/T_{i0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion Temperature intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta T_i/T_{i,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + ax = axs["G"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].t, self.results[label].ni_all_rms_sumnr_sumn[ion]*100.0, ls[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]}") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta n_i/n_{i,0}/n_{i0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ions (all) Density intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + + ax = axs["H"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].t, self.results[label].Ti_all_rms_sumnr_sumn[ion]*100.0, ls[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]}") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta T_i/T_{i,0}/n_{i0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ions (all) Temperature intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_intensities_ky(self, axs=None, label="", c="b", addText=True): + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + ACEG + BDFH + """ ) + + ls = GRAPHICStools.listLS() - ax = ax00 - for ky in range(len(c.ky)): - ax.plot( - c.t, - c.freq[1, ky, :], - color=colorsC[ky], - label=f"$k_{{\\theta}}\\rho_s={np.abs(c.ky[ky]):.2f}$", - ) + # Potential intensity + ax = axs["A"] + ax.plot(self.results[label].ky, self.results[label].phi_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].phi_rms_sumnr_mean-self.results[label].phi_rms_sumnr_std, self.results[label].phi_rms_sumnr_mean+self.results[label].phi_rms_sumnr_std, color=c, alpha=0.2) - ax = ax10 - for ky in range(len(c.ky)): - ax.plot( - c.t, - c.freq[0, ky, :], - color=colorsC[ky], - label=f"$k_{{\\theta}}\\rho_s={np.abs(c.ky[ky]):.2f}$", - ) + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel(r"$\delta\phi/\phi_0$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Potential intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta\phi/\phi_0|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + # EM potential intensity + ax = axs["B"] + if 'apar' in self.results[label].__dict__: + ax.plot(self.results[label].ky, self.results[label].apar_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+', $A_\\parallel$ (mean)') + ax.fill_between(self.results[label].ky, self.results[label].apar_rms_sumnr_mean-self.results[label].apar_rms_sumnr_std, self.results[label].apar_rms_sumnr_mean+self.results[label].apar_rms_sumnr_std, color=c, alpha=0.2) + ax.plot(self.results[label].ky, self.results[label].bpar_rms_sumnr_mean, '--', markersize=5, color=c, label=label+', $B_\\parallel$ (mean)') + ax.fill_between(self.results[label].ky, self.results[label].bpar_rms_sumnr_mean-self.results[label].bpar_rms_sumnr_std, self.results[label].bpar_rms_sumnr_mean+self.results[label].bpar_rms_sumnr_std, color=c, alpha=0.2) - K.append(np.abs(c.ky[0])) - G.append(c.freq[1, 0, -1]) - F.append(c.freq[0, 0, -1]) - - GACODEplotting.plotTGLFspectrum( - [ax01, ax11], - K, - G, - freq=F, - coeff=0.0, - c=colors[0], - ls="-", - lw=1, - label="", - facecolors=colors[: len(K)], - markersize=50, - alpha=1.0, - titles=["Growth Rate", "Real Frequency"], - removeLow=1e-4, - ylabel=True, - ) + ax.legend(loc='best', prop={'size': 8},) - ax = ax00 - ax.set_xlabel("Time $(a/c_s)$") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - ax.set_ylabel("$\\gamma$ $(c_s/a)$") - ax.set_title("Growth Rate") - ax.set_xlim(left=0) - ax.legend() - ax = ax10 - ax.set_xlabel("Time $(a/c_s)$") - ax.set_ylabel("$\\omega$ $(c_s/a)$") - ax.set_title("Real Frequency") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - ax.set_xlim(left=0) + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel(r"$\delta F_\parallel/F_{\parallel,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('EM potential intensity vs. $k_\\theta\\rho_s$') + + ax.axhline(0.0, color='k', ls='--', lw=1) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta F_\parallel/F_{\parallel,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + # Electron particle intensity + ax = axs["C"] + ax.plot(self.results[label].ky, self.results[label].ne_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].ne_rms_sumnr_mean-self.results[label].ne_rms_sumnr_std, self.results[label].ne_rms_sumnr_mean+self.results[label].ne_rms_sumnr_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta n_e/n_{e,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron particle intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta n_e/n_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + # Electron temperature intensity + ax = axs["D"] + ax.plot(self.results[label].ky, self.results[label].Te_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Te_rms_sumnr_mean-self.results[label].Te_rms_sumnr_std, self.results[label].Te_rms_sumnr_mean+self.results[label].Te_rms_sumnr_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta T_e/T_{e,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron temperature intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta T_e/T_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + # Ion particle intensity + ax = axs["E"] + ax.plot(self.results[label].ky, self.results[label].ni_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].ni_rms_sumnr_mean-self.results[label].ni_rms_sumnr_std, self.results[label].ni_rms_sumnr_mean+self.results[label].ni_rms_sumnr_std, color=c, alpha=0.2) - ax = ax01 - #ax.set_xlim([5e-2, 50.0]) + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta n_i/n_{i,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion particle intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta n_i/n_{i,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + # Ion temperature intensity + ax = axs["F"] + ax.plot(self.results[label].ky, self.results[label].Ti_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Ti_rms_sumnr_mean-self.results[label].Ti_rms_sumnr_std, self.results[label].Ti_rms_sumnr_mean+self.results[label].Ti_rms_sumnr_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta T_i/T_{i,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion temperature intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta T_i/T_{i,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + # Ion particle intensity + ax = axs["G"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].ky, self.results[label].ni_all_rms_sumnr_mean[ion], ls[ion]+'o', markersize=5, color=c, label=f"{label}, {self.results[label].all_names[ion]} (mean)") - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - ax00 = fig2.add_subplot(grid[0, 0]) - ax01 = fig2.add_subplot(grid[0, 1], sharex=ax00, sharey=ax00) - ax02 = fig2.add_subplot(grid[0, 2], sharex=ax00, sharey=ax00) - ax10 = fig2.add_subplot(grid[1, 0], sharex=ax00, sharey=ax00) - ax11 = fig2.add_subplot(grid[1, 1], sharex=ax01, sharey=ax00) - ax12 = fig2.add_subplot(grid[1, 2], sharex=ax02, sharey=ax00) - it = -1 + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta n_i/n_{i,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ions (all) particle intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + + # Ion temperature intensity + ax = axs["H"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].ky, self.results[label].Ti_all_rms_sumnr_mean[ion], ls[ion]+'o', markersize=5, color=c, label=f"{label}, {self.results[label].all_names[ion]} (mean)") + - for cont, label in enumerate(self.results): - c = self.results[label] - baseColor = colors[cont] + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta T_i/T_{i,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ions (all) temperature intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_intensities_kx(self, axs=None, label="", c="b", addText=True): + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) - colorsC, _ = GRAPHICStools.colorTableFade( - len(c.ky), - startcolor=baseColor, - endcolor=baseColor, - alphalims=[1.0, 0.4], + axs = fig.subplot_mosaic( + """ + AC + BD + """ ) - ax = ax00 - for ky in range(len(c.ky)): - for var, axs, label in zip( - ["phib", "aparb", "bparb"], - [[ax00, ax10], [ax01, ax11], [ax02, ax12]], - ["phi", "abar", "aper"], - ): - try: - f = c.__dict__[var][0, :, it] + 1j * c.__dict__[var][1, :, it] - y1 = np.real(f) - y2 = np.imag(f) - x = c.thetab / np.pi - - ax = axs[0] - ax.plot( - x, - y1, - color=colorsC[ky], - ls="-", - label=f"$k_{{\\theta}}\\rho_s={np.abs(c.ky[ky]):.2f}$", - ) - ax = axs[1] - ax.plot(x, y2, color=colorsC[ky], ls="-") - except: - pass - - ax = ax00 - ax.set_xlabel("$\\theta/\\pi$") + # Potential intensity + ax = axs["A"] + ax.plot(self.results[label].kx, self.results[label].phi_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+' (mean)') + ax.plot(self.results[label].kx, self.results[label].phi_rms_n0_mean, '-.', markersize=0.5, lw=0.5, color=c, label=label+', $n=0$ (mean)') + ax.plot(self.results[label].kx, self.results[label].phi_rms_sumn1_mean, '--', markersize=0.5, lw=0.5, color=c, label=label+', $n>0$ (mean)') + + ax.set_xlabel("$k_{x}$") + ax.set_ylabel("$\\delta \\phi/\\phi_0$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Potential intensity vs kx') + ax.legend(loc='best', prop={'size': 8},) + ax.set_yscale('log') + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}|\delta\phi/\phi_0|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + # EM potential intensity + ax = axs["C"] + if 'apar' in self.results[label].__dict__: + ax.plot(self.results[label].kx, self.results[label].apar_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+', $A_\\parallel$ (mean)') + ax.plot(self.results[label].kx, self.results[label].bpar_rms_sumn_mean, '--', markersize=1.0, lw=1.0, color=c, label=label+', $B_\\parallel$ (mean)') + + ax.legend(loc='best', prop={'size': 8},) + + + ax.set_xlabel("$k_{x}$") + ax.set_ylabel("$\\delta F_\\parallel/F_{\\parallel,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('EM potential intensity vs kx') + ax.set_yscale('log') + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}|\delta F_\parallel/F_{\parallel,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + # Electron particle intensity + ax = axs["B"] + ax.plot(self.results[label].kx, self.results[label].ne_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+' (mean)') + ax.plot(self.results[label].kx, self.results[label].ne_rms_n0_mean, '-.', markersize=0.5, lw=0.5, color=c, label=label+', $n=0$ (mean)') + ax.plot(self.results[label].kx, self.results[label].ne_rms_sumn1_mean, '--', markersize=0.5, lw=0.5, color=c, label=label+', $n>0$ (mean)') + + ax.set_xlabel("$k_{x}$") + ax.set_ylabel("$\\delta n_e/n_{e,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron particle intensity vs kx') + ax.legend(loc='best', prop={'size': 8},) + ax.set_yscale('log') + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}|\delta n_e/n_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + # Electron temperature intensity + ax = axs["D"] + ax.plot(self.results[label].kx, self.results[label].Te_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+' (mean)') + ax.plot(self.results[label].kx, self.results[label].Te_rms_n0_mean, '-.', markersize=0.5, lw=0.5, color=c, label=label+', $n=0$ (mean)') + ax.plot(self.results[label].kx, self.results[label].Te_rms_sumn1_mean, '--', markersize=0.5, lw=0.5, color=c, label=label+', $n>0$ (mean)') + + ax.set_xlabel("$k_{x}$") + ax.set_ylabel("$\\delta T_e/T_{e,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron temperature intensity vs kx') + ax.legend(loc='best', prop={'size': 8},) + ax.set_yscale('log') + + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}|\delta T_e/T_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_turbulence(self, axs = None, label= "cgyro1", c="b", kys = None): + + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + AC + BD + """ + ) + + # Is no kys provided, select just 3: first, last and middle + if kys is None: + ikys = [0] + if len(self.results[label].ky) > 1: + ikys.append(-1) + if len(self.results[label].ky) > 2: + ikys.append(len(self.results[label].ky) // 2) + + ikys = np.unique(ikys) + else: + ikys = [self.results[label].ky.index(ky) for ky in kys if ky in self.results[label].ky] + + # Growth rate as function of time + ax = axs["A"] + for i,ky in enumerate(ikys): + self._plot_trace( + ax, + label, + self.results[label].g[ky, :], + c=c, + ls = GRAPHICStools.listLS()[i], + lw=1, + label_plot=f"$k_{{\\theta}}\\rho_s={np.abs(self.results[label].ky[ky]):.2f}$", + var_meanstd = [self.results[label].g_mean[ky], self.results[label].g_std[ky]], + ) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\gamma$ (norm.)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Growth rate vs time') + ax.legend(loc='best', prop={'size': 8},) + + # Frequency as function of time + ax = axs["B"] + for i,ky in enumerate(ikys): + self._plot_trace( + ax, + label, + self.results[label].f[ky, :], + c=c, + ls = GRAPHICStools.listLS()[i], + lw=1, + label_plot=f"$k_{{\\theta}}\\rho_s={np.abs(self.results[label].ky[ky]):.2f}$", + var_meanstd = [self.results[label].f_mean[ky], self.results[label].f_std[ky]], + ) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\omega$ (norm.)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Real Frequency vs time') + ax.legend(loc='best', prop={'size': 8},) + + # Mean+Std Growth rate as function of ky + ax = axs["C"] + ax.errorbar(self.results[label].ky, self.results[label].g_mean, yerr=self.results[label].g_std, fmt='-o', markersize=5, color=c, label=label+' (mean+std)') + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\gamma$ (norm.)") + ax.set_title('Saturated Growth Rate') + GRAPHICStools.addDenseAxis(ax) + ax.legend(loc='best', prop={'size': 8},) + + # Mean+Std Frequency as function of ky + ax = axs["D"] + ax.errorbar(self.results[label].ky, self.results[label].f_mean, yerr=self.results[label].f_std, fmt='-o', markersize=5, color=c, label=label+' (mean+std)') + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\omega$ (norm.)") + ax.set_title('Saturated Real Frequency') + GRAPHICStools.addDenseAxis(ax) + ax.legend(loc='best', prop={'size': 8},) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_cross_phases(self, axs = None, label= "cgyro1", c="b"): + + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + ACEG + BDFH + """ + ) + + ls = GRAPHICStools.listLS() + m = GRAPHICStools.listmarkers() + + ax = axs["A"] + ax.plot(self.results[label].ky, self.results[label].neTe_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].neTe_kx0_mean-self.results[label].neTe_kx0_std, self.results[label].neTe_kx0_mean+self.results[label].neTe_kx0_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$n_e-T_e$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) + GRAPHICStools.addDenseAxis(ax) + ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$n_e-T_e$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + + ax = axs["B"] + ax.plot(self.results[label].ky, self.results[label].niTi_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].niTi_kx0_mean-self.results[label].niTi_kx0_std, self.results[label].niTi_kx0_mean+self.results[label].niTi_kx0_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$n_i-T_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) + GRAPHICStools.addDenseAxis(ax) + ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$n_i-T_i$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + ax = axs["C"] + ax.plot(self.results[label].ky, self.results[label].phine_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].phine_kx0_mean-self.results[label].phine_kx0_std, self.results[label].phine_kx0_mean+self.results[label].phine_kx0_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\phi-n_e$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) + GRAPHICStools.addDenseAxis(ax) + ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$\\phi-n_e$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + ax = axs["D"] + ax.plot(self.results[label].ky, self.results[label].phini_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].phini_kx0_mean-self.results[label].phini_kx0_std, self.results[label].phini_kx0_mean+self.results[label].phini_kx0_std, color=c, alpha=0.2) + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\phi-n_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) + GRAPHICStools.addDenseAxis(ax) + ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$\\phi-n_i$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + + ax = axs["E"] + ax.plot(self.results[label].ky, self.results[label].phiTe_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].phiTe_kx0_mean-self.results[label].phiTe_kx0_std, self.results[label].phiTe_kx0_mean+self.results[label].phiTe_kx0_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\phi-T_e$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) + GRAPHICStools.addDenseAxis(ax) + ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$\\phi-T_e$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + + ax = axs["F"] + ax.plot(self.results[label].ky, self.results[label].phiTi_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].phiTi_kx0_mean-self.results[label].phiTi_kx0_std, self.results[label].phiTi_kx0_mean+self.results[label].phiTi_kx0_std, color=c, alpha=0.2) + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\phi-T_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) + GRAPHICStools.addDenseAxis(ax) + ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$\\phi-T_i$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + + ax = axs["G"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].ky, self.results[label].phiTi_all_kx0_mean[ion], ls[ion]+m[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]} (mean)", markersize=4) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\phi-T_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) + GRAPHICStools.addDenseAxis(ax) + ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$\\phi-T_i$ (all) cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + + ax = axs["H"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].ky, self.results[label].phini_all_kx0_mean[ion], ls[ion]+m[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]} (mean)", markersize=4) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\phi-n_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) + GRAPHICStools.addDenseAxis(ax) + ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$\\phi-n_i$ (all) cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + + + + + + def plot_ballooning(self, time = None, label="cgyro1", c="b", axs=None): + + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + 135 + 246 + """ + ) + + if time is None: + time = np.min([self.results[label].tmin, self.results[label].tmax_fluct]) + + it = np.argmin(np.abs(self.results[label].t - time)) + + colorsC, _ = GRAPHICStools.colorTableFade( + len(self.results[label].ky), + startcolor=c, + endcolor=c, + alphalims=[1.0, 0.4], + ) + + ax = axs['1'] + for ky in range(len(self.results[label].ky)): + for var, axsT in zip( + ["phi_ballooning", "apar_ballooning", "bpar_ballooning"], + [[axs['1'], axs['2']], [axs['3'], axs['4']], [axs['5'], axs['6']]], + ): + + f = self.results[label].__dict__[var][:, it] + y1 = np.real(f) + y2 = np.imag(f) + x = self.results[label].theta_ballooning / np.pi + + # Normalize + y1_max = np.max(np.abs(y1)) + y2_max = np.max(np.abs(y2)) + y1 /= y1_max + y2 /= y2_max + + ax = axsT[0] + ax.plot( + x, + y1, + color=colorsC[ky], + ls="-", + label=f"$k_{{\\theta}}\\rho_s={np.abs( self.results[label].ky[ky]):.2f}$ (max {y1_max:.2e})", + ) + ax = axsT[1] + ax.plot( + x, + y2, + color=colorsC[ky], + ls="-", + label=f"$k_{{\\theta}}\\rho_s={np.abs( self.results[label].ky[ky]):.2f}$ (max {y2_max:.2e})", + ) + + + ax = axs['1'] + ax.set_xlabel("$\\theta/\\pi$ (normalized to maximum)") ax.set_ylabel("Re($\\delta\\phi$)") ax.set_title("$\\delta\\phi$") - ax.legend(loc="best") + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) ax.set_xlim([-2 * np.pi, 2 * np.pi]) - ax = ax01 - ax.set_xlabel("$\\theta/\\pi$") + ax = axs['3'] + ax.set_xlabel("$\\theta/\\pi$ (normalized to maximum)") ax.set_ylabel("Re($\\delta A\\parallel$)") ax.set_title("$\\delta A\\parallel$") - ax = ax02 - ax.set_xlabel("$\\theta/\\pi$") + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + ax = axs['5'] + ax.set_xlabel("$\\theta/\\pi$ (normalized to maximum)") ax.set_ylabel("Re($\\delta B\\parallel$)") ax.set_title("$\\delta B\\parallel$") - ax = ax10 + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + ax = axs['2'] ax.set_xlabel("$\\theta/\\pi$") ax.set_ylabel("Im($\\delta\\phi$)") - ax = ax11 + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + ax = axs['4'] ax.set_xlabel("$\\theta/\\pi$") ax.set_ylabel("Im($\\delta A\\parallel$)") - ax = ax12 + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + ax = axs['6'] ax.set_xlabel("$\\theta/\\pi$") ax.set_ylabel("Im($\\delta B\\parallel$)") + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + - for ax in [ax00, ax01, ax02, ax10, ax11, ax12]: + for ax in [axs['1'], axs['3'], axs['5'], axs['2'], axs['4'], axs['6']]: ax.axvline(x=0, lw=0.5, ls="--", c="k") ax.axhline(y=0, lw=0.5, ls="--", c="k") + + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_2D(self, label="cgyro1", axs=None, times = None): + + if times is None: + times = [] + + number_times = len(axs)//3 if axs is not None else 4 + + times = [self.results[label].t[-1-i*10] for i in range(number_times)] + + if axs is None: + + mosaic = _2D_mosaic(len(times)) + + plt.ion() + fig = plt.figure(figsize=(18, 9)) + axs = fig.subplot_mosaic(mosaic) + + # Pre-calculate global min/max for each field type across all times + phi_values = [] + n_values = [] + e_values = [] + + for time in times: + it = np.argmin(np.abs(self.results[label].t - time)) + + # Get phi values + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_phi', it = it) + phi_values.append(fp) + + # Get n values + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_n',species = self.results[label].electron_flag, it = it) + n_values.append(fp) + + # Get e values + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_e',species = self.results[label].electron_flag, it = it) + e_values.append(fp) + + # Calculate global ranges + phi_max = np.max([np.max(np.abs(fp)) for fp in phi_values]) + phi_min, phi_max = -phi_max, +phi_max + + n_max = np.max([np.max(np.abs(fp)) for fp in n_values]) + n_min, n_max = -n_max, +n_max + + e_max = np.max([np.max(np.abs(fp)) for fp in e_values]) + e_min, e_max = -e_max, +e_max + + colorbars = [] # Store colorbar references + # Now plot with consistent colorbar ranges + for time_i, time in enumerate(times): + + print(f"\t- Plotting 2D turbulence for {label} at time {time}") + + it = np.argmin(np.abs(self.results[label].t - time)) + + cfig = axs[str(time_i+1)].get_figure() + + # Phi plot + ax = axs[str(time_i+1)] + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_phi', it = it) + + cs1 = ax.contourf(xp,yp,np.transpose(fp),levels=np.arange(phi_min,phi_max,(phi_max-phi_min)/256),cmap=plt.get_cmap('jet')) + cphi = cfig.colorbar(cs1, ax=ax) + + ax.set_xlabel("$x/\\rho_s$") + ax.set_ylabel("$y/\\rho_s$") + ax.set_title(f"$\\delta\\phi/\\phi_0$ (t={self.results[label].t[it]} $a/c_s$)") + ax.set_aspect('equal') + + # N plot + ax = axs[str(time_i+1+len(times))] + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_n',species = self.results[label].electron_flag, it = it) + + cs2 = ax.contourf(xp,yp,np.transpose(fp),levels=np.arange(n_min,n_max,(n_max-n_min)/256),cmap=plt.get_cmap('jet')) + cn = cfig.colorbar(cs2, ax=ax) + + ax.set_xlabel("$x/\\rho_s$") + ax.set_ylabel("$y/\\rho_s$") + ax.set_title(f"$\\delta n_e/n_{{e,0}}$ (t={self.results[label].t[it]} $a/c_s$)") + ax.set_aspect('equal') + + # E plot + ax = axs[str(time_i+1+len(times)*2)] + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_e',species = self.results[label].electron_flag, it = it) + + cs3 = ax.contourf(xp,yp,np.transpose(fp),levels=np.arange(e_min,e_max,(e_max-e_min)/256),cmap=plt.get_cmap('jet')) + ce = cfig.colorbar(cs3, ax=ax) + + ax.set_xlabel("$x/\\rho_s$") + ax.set_ylabel("$y/\\rho_s$") + ax.set_title(f"$\\delta E_e/E_{{e,0}}$ (t={self.results[label].t[it]} $a/c_s$)") + ax.set_aspect('equal') + + # Store the colorbar objects with their associated contour plots + colorbars.append({ + 'phi': cphi, + 'n': cn, + 'e': ce + }) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.4, horizontal=0.3) + + return colorbars + + + def _to_real_space(self, variable = 'kxky_phi', species = None, label="cgyro1", theta_plot = 0, it = -1): + + # from pygacode + def maptoreal_fft(nr,nn,nx,ny,c): + + d = np.zeros([nx,nn],dtype=complex) + for i in range(nr): + p = i-nr//2 + if -p < 0: + k = -p+nx + else: + k = -p + d[k,0:nn] = np.conj(c[i,0:nn]) + f = np.fft.irfft2(d,s=[nx,ny],norm='forward')*0.5 + + # Correct for half-sum + f = 2*f + + return f + + # Real space + nr = self.results[label].cgyrodata.n_radial + nn = self.results[label].cgyrodata.n_n + craw = self.results[label].cgyrodata.__dict__[variable] + + itheta = np.argmin(np.abs(self.results[label].theta_stored-theta_plot)) + if species is None: + c = craw[:,itheta,:,it] + else: + c = craw[:,itheta,species,:,it] + + nx = self.results[label].cgyrodata.__dict__[variable].shape[0] + ny = nx + + # Arrays + x = np.arange(nx)*2*np.pi/nx + y = np.arange(ny)*2*np.pi/ny + f = maptoreal_fft(nr,nn,nx,ny,c) + + # Physical maxima + ky1 = self.results[label].cgyrodata.ky[1] if len(self.results[label].cgyrodata.ky) > 1 else self.results[label].cgyrodata.ky[0] + xmax = self.results[label].cgyrodata.length + ymax = (2*np.pi)/np.abs(ky1) + xp = x/(2*np.pi)*xmax + yp = y/(2*np.pi)*ymax + + # Periodic extensions + xp = np.append(xp,xmax) + yp = np.append(yp,ymax) + fp = np.zeros([nx+1,ny+1]) + fp[0:nx,0:ny] = f[:,:] + fp[-1,:] = fp[0,:] + fp[:,-1] = fp[:,0] + + return xp, yp, fp + + def plot_quick_linear(self, labels=["cgyro1"], fig=None): + colors = GRAPHICStools.listColors() + + if fig is None: + fig = plt.figure(figsize=(15,9)) + + axs = fig.subplot_mosaic( + """ + 12 + 34 + """ + ) + + def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): + + for cont, label in enumerate(labels): + c = self.results[label] + baseColor = colors[cont+start_cont+1] + colorsC, _ = GRAPHICStools.colorTableFade( + len(c.ky), + startcolor=baseColor, + endcolor=baseColor, + alphalims=[1.0, 0.4], + ) + + ax = axs['1'] + for ky in range(len(c.ky)): + ax.plot( + c.t, + c.g[ky,:], + color=colorsC[ky], + label=f"$k_{{\\theta}}\\rho_s={np.abs(c.ky[ky]):.2f}$", + ) + + ax = axs['2'] + for ky in range(len(c.ky)): + ax.plot( + c.t, + c.f[ky,:], + color=colorsC[ky], + label=f"$k_{{\\theta}}\\rho_s={np.abs(c.ky[ky]):.2f}$", + ) + + GACODEplotting.plotTGLFspectrum( + [axs['3'], axs['4']], + self.results[label_base].ky, + self.results[label_base].g_mean, + freq=self.results[label_base].f_mean, + coeff=0.0, + c=col_lin, + ls="-", + lw=1, + label="", + facecolors=colors, + markersize=50, + alpha=1.0, + titles=["Growth Rate", "Real Frequency"], + removeLow=1e-4, + ylabel=True, + ) + + return cont + + co = -1 + for i,label0 in enumerate(labels): + if isinstance(self.results[label0], CGYROutils.CGYROlinear_scan): + co = _plot_linear_stability(axs, self.results[label0].labels, label0, start_cont=co, col_lin=colors[i]) + else: + co = _plot_linear_stability(axs, [label0], label0, start_cont=co, col_lin=colors[i]) + + ax = axs['1'] + ax.set_xlabel("Time $(a/c_s)$") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + ax.set_ylabel("$\\gamma$ $(c_s/a)$") + ax.set_title("Growth Rate") + ax.set_xlim(left=0) + ax.legend() + ax = axs['2'] + ax.set_xlabel("Time $(a/c_s)$") + ax.set_ylabel("$\\omega$ $(c_s/a)$") + ax.set_title("Real Frequency") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + ax.set_xlim(left=0) + class CGYROinput: @@ -743,7 +1878,7 @@ def __init__(self, file=None): self.controls = GACODErun.buildDictFromInput(self.file_txt) - def write_state(self, file=None): + def writeCurrentStatus(self, file=None): print("\t- Writting CGYRO input file") if file is None: @@ -765,3 +1900,20 @@ def write_state(self, file=None): for ikey in self.controls: var = self.controls[ikey] f.write(f"{ikey.ljust(23)} = {var}\n") + + +def _2D_mosaic(n_times): + + num_cols = n_times + + # Create the mosaic layout dynamically + mosaic = [] + counter = 1 + for _ in range(3): + row = [] + for _ in range(num_cols): + row.append(str(counter)) + counter += 1 + mosaic.append(row) + + return mosaic \ No newline at end of file diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 52a3b655..516bd2e7 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -1049,9 +1049,7 @@ def plot( max_fields.append(il) if fn is None: - self.fn = GUItools.FigureNotebook( - "TGLF MITIM Notebook", geometry="1700x900", vertical=True - ) + self.fn = GUItools.FigureNotebook("TGLF MITIM Notebook", geometry="1700x900", vertical=True) else: self.fn = fn diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 22fde9cd..779bdb77 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -1,6 +1,7 @@ import argparse +from xml.etree.ElementInclude import include +import matplotlib.pyplot as plt from IPython import embed -from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import CGYROtools """ @@ -11,27 +12,39 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("folders", type=str, nargs="*") - parser.add_argument("--linear", action="store_true", help="linear run") + parser.add_argument("--two", action="store_true", help="Include 2D plots") + parser.add_argument("--linear", action="store_true", help="Just a plot of the linear spectra") + parser.add_argument("--tmin", type=float, nargs="*", default=None, help="Minimum time to calculate mean and std") args = parser.parse_args() folders = args.folders linear = args.linear - + tmin = args.tmin + include_2D = args.two + + if tmin is None: + tmin = [0.0] * len(folders) + last_tmin_for_linear = True + else: + last_tmin_for_linear = False + # Read c = CGYROtools.CGYRO() labels = [] for i, folder in enumerate(folders): labels.append(f"case {i + 1}") - c.read(label=labels[-1], folder=folder) + c.read(label=labels[-1], folder=folder, tmin=tmin[i], last_tmin_for_linear=last_tmin_for_linear) if linear: # Plot linear spectrum - c.plotLS(labels=labels) + c.plot_quick_linear(labels=labels) + plt.show() else: - c.plot(labels=labels) + c.plot(labels=labels, include_2D=include_2D, common_colorbar=True) + c.fn.show() - c.fn.show() + embed() if __name__ == "__main__": diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index ab28cc0a..54389192 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -1,367 +1,607 @@ -''' -From NTH -''' - -import sys +import os +import scipy import numpy as np -import matplotlib.pyplot as plt -from gacodefuncs import * import statsmodels.api as sm -from cgyro.data import cgyrodata -import math -#from omfit_classes import omfit_gapy +import matplotlib.pyplot as plt +from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from pygacode.cgyro.data_plot import cgyrodata_plot +from pygacode import gacodefuncs +from IPython import embed +import quends as qnds +import pandas as pd -def grab_cgyro_nth(data_dir, tstart, plotflag, printflag, file = None): - - #Get all the relevant simulation quantities - data = cgyrodata(data_dir+'/') - print(dir(data)) - data.getflux() - rho_star=(data.vth_norm/((1.602e-19*data.b_unit)/(data.mass_norm*1e-27)))/data.a_meters - nt=data.n_time - t=data.t - dt=t[0] - ky=abs(data.ky) - n_n=data.n_n - n_spec=data.n_species - n_field=data.n_field - #print(n_field) - flux = data.ky_flux - #Shape should be (species,flux,field,ky,time) - #rint(dir(data)) - roa=data.rmin - #print data.__init__ - tmax=np.amax(t) - sflux = np.sum(flux,axis=3) - kflux=np.mean(flux,axis=4) - tkflux=np.sum(kflux,axis=2) - #Values of gradients and GB normalizations - alti=data.dlntdr[0] #technically ion 1 scale length - alte=data.dlntdr[n_spec-1] - alne=data.dlnndr[n_spec-1] - qgb=data.q_gb_norm - ggb=data.gamma_gb_norm - pgb=data.pi_gb_norm - sgb=qgb/data.a_meters - - #Total electron heat flux, sum over field - eflux_all=sflux[n_spec-1,1,:,:] - eflux=np.sum(eflux_all,axis=0) - - #ES electron heat flux, just electrostatic - efluxes=eflux_all[0,:] +class CGYROlinear_scan: + def __init__(self, labels, cgyro_data): + + self.labels = labels + + # Store the data in a structured way + self.aLTi = [] + self.ky = [] + self.g_mean = [] + self.f_mean = [] + + self.neTe_mean = [] + + self.Qe_mean = [] + self.Qi_mean = [] + + for label in labels: + self.ky.append(cgyro_data[label].ky[0]) + self.aLTi.append(cgyro_data[label].aLTi) + self.g_mean.append(cgyro_data[label].g_mean[0]) + self.f_mean.append(cgyro_data[label].f_mean[0]) + + self.Qe_mean.append(cgyro_data[label].Qe_mean) + self.Qi_mean.append(cgyro_data[label].Qi_mean) + + try: + self.neTe_mean.append(cgyro_data[label].neTe_kx0_mean[0]) + except: + self.neTe_mean.append(np.nan) + + self.ky = np.array(self.ky) + self.aLTi = np.array(self.aLTi) + self.g_mean = np.array(self.g_mean) + self.f_mean = np.array(self.f_mean) + self.neTe_mean = np.array(self.neTe_mean) + self.Qe_mean = np.array(self.Qe_mean) + self.Qi_mean = np.array(self.Qi_mean) + +class CGYROout: + def __init__(self, folder, tmin=0.0, minimal=False, last_tmin_for_linear=True): + + original_dir = os.getcwd() + + self.folder = folder + self.tmin = tmin + + try: + print(f"\t- Reading CGYRO data from {self.folder.resolve()}") + self.cgyrodata = cgyrodata_plot(f"{self.folder.resolve()}{os.sep}") + except FileNotFoundError: + raise Exception(f"[MITIM] Could not find CGYRO data in {self.folder.resolve()}. Please check the folder path or run CGYRO first.") + except Exception as e: + print(f"\t- Error reading CGYRO data: {e}") + if print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): + os.chdir(self.folder) + os.system("cgyro -t") + self.cgyrodata = cgyrodata_plot(f"{self.folder.resolve()}{os.sep}") + + os.chdir(original_dir) + + # -------------------------------------------------------------- + # Read inputs + # -------------------------------------------------------------- + + self.params1D = {} + for var in self.cgyrodata.__dict__: + par = self.cgyrodata.__dict__[var] + if isinstance(par, bool) or IOtools.isnum(par): + self.params1D[var] = par + elif isinstance(par, (list, np.ndarray)) and par.ndim==1 and len(par) <= 5: + for i, p in enumerate(par): + self.params1D[f"{var}_{i}"] = p - #EM electron heat flux, just A_|| - efluxem=eflux_all[1,:] + # -------------------------------------------------------------- + # Postprocess with MITIM-curated structures and variables + # -------------------------------------------------------------- + + # Check for linear run + if 'phib' in self.cgyrodata.__dict__ and last_tmin_for_linear: + print('\t- Forcing tmin to the last time point because this is a linear run', typeMsg='i') + self.tmin = self.cgyrodata.t[-1] + self.linear = True + else: + self.linear = False + + self.cgyrodata.getflux(cflux='auto') + self.cgyrodata.getnorm("elec") + self.cgyrodata.getgeo() + self.cgyrodata.getxflux() + + # Understand positions + self.electron_flag = np.where(self.cgyrodata.z == -1)[0][0] + self.all_flags = np.arange(0, len(self.cgyrodata.z), 1) + self.ions_flags = self.all_flags[self.all_flags != self.electron_flag] + + self.all_names = [f"{gacodefuncs.specmap(self.cgyrodata.mass[i],self.cgyrodata.z[i])}({self.cgyrodata.z[i]},{self.cgyrodata.mass[i]:.1f})" for i in self.all_flags] + + self.fields = np.arange(self.cgyrodata.n_field) + + self.aLTi = self.cgyrodata.dlntdr[0] + self.aLTe = self.cgyrodata.dlntdr[self.electron_flag] + self.aLne = self.cgyrodata.dlnndr[self.electron_flag] - #Total electron particle flux, sum over fields - epflux_all=sflux[n_spec-1,0,:,:] - epflux=np.sum(epflux_all,axis=0) - - #Total impurity particle flux, sum over fields (***NOTE THIS IS FOR SPECIES n_Spec -3******) - impflux_all=sflux[n_spec-3,0,:,:] - impflux=np.sum(impflux_all,axis=0) - - #Depending on if electrostatic of E&M output components - if n_field ==1: - epfluxp=epflux_all[0,:] - epfluxap=epflux_all[0,:]*0.0 #Zero out the A_par array - epfluxbp=epflux_all[0,:]*0.0 #Zero out the B_par array - - if n_field == 2: - epfluxp=epflux_all[0,:] - epfluxap=epflux_all[1,:] - epfluxbp=epflux_all[0,:]*0.0 #Zero out the B_par array - - if n_field ==3: - epfluxp=epflux_all[0,:] - epfluxap=epflux_all[1,:] - epfluxbp=epflux_all[2,:] + + # ************************ + # Normalization + # ************************ - - #Total (ES+EM) ion flux summed over ions - iflux_all=sflux[0:n_spec-1,1,:,:] - iflux=np.sum(iflux_all,axis=0) #Sum over ions - iflux=np.sum(iflux,axis=0) #Sum over fields + self.t = self.cgyrodata.tnorm + self.ky = self.cgyrodata.kynorm + self.kx = self.cgyrodata.kxnorm + self.theta = self.cgyrodata.theta + + if self.cgyrodata.theta_plot == 1: + self.theta_stored = np.array([0.0]) + else: + self.theta_stored = np.array([-1+2.0*i/self.cgyrodata.theta_plot for i in range(self.cgyrodata.theta_plot)]) + + self.Qgb = self.cgyrodata.q_gb_norm + self.Ggb = self.cgyrodata.gamma_gb_norm + + self.artificial_rhos_factor = self.cgyrodata.rho_star_norm / self.cgyrodata.rhonorm - #Sum the momentum flux (all species) - mflux_all=sflux[0:n_spec,2,:,:] - mflux=np.sum(mflux_all,axis=0) # Sum over species - mflux=np.sum(mflux,axis=0) # Sum over fields + self._process_linear() - #Total electron turbulent exchange, sm over all fields + if not minimal: # or not self.linear: + self.cgyrodata.getbigfield() - #Put in an option if turbulent exchange is not enabled - if sflux.shape[1] == 4: - turflux_all=sflux[n_spec-1,3,:,:] - else: - turflux_all=sflux[n_spec-1,2,:,:]*0.0 #Fill in with dummy values and 0 + if 'kxky_phi' in self.cgyrodata.__dict__: + self._process_fluctuations() + else: + print(f'\t- No fluctuations found in CGYRO data ({IOtools.clipstr(self.folder)}), skipping fluctuation processing and will not be able to plot default Notebook', typeMsg='w') + else: + print('\t- Minimal mode, skipping fluctuations processing', typeMsg='i') + + self._process_fluxes() + self._saturate_signals() + + def _process_linear(self): + + # check for convergence + self.linear_converged = False + info_file = f"{self.folder.resolve()}/out.cgyro.info" + if not os.path.exists(info_file): + raise FileNotFoundError(f"[MITIM] Could not find CGYRO info file at {info_file}. Please check the folder path or run CGYRO first.") + else: + with open(info_file, 'r') as f: + lines = f.readlines() + for line in lines: + if "EXIT: (CGYRO) Linear converged" in line: + self.linear_converged = True + break + + self.f = self.cgyrodata.fnorm[0,:,:] # (ky, time) + self.g = self.cgyrodata.fnorm[1,:,:] # (ky, time) + if self.g is np.nan or self.f is np.nan: + raise ValueError(f"[MITIM] Could not find f or g in CGYRO data at {info_file}. Please check the folder path or run CGYRO first.") + + # Ballooning Modes (complex eigenfunctions) + if 'phib' in self.cgyrodata.__dict__: + self.phi_ballooning = self.cgyrodata.phib # (ball, time) + self.apar_ballooning = self.cgyrodata.aparb # (ball, time) + self.bpar_ballooning = self.cgyrodata.bparb # (ball, time) + self.theta_ballooning = self.cgyrodata.thetab # (ball, time) + + def _process_fluctuations(self): + # Fluctuations (complex numbers) + + gbnorm = False + + theta = -1 + + moment, species, field = 'phi', None, 0 + self.phi, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) + if 'kxky_apar' in self.cgyrodata.__dict__: + field = 1 + self.apar, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) + field = 2 + self.bpar, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) + + self.tmax_fluct = _detect_exploiding_signal(self.t, self.phi**2) + + moment, species, field = 'n', self.electron_flag, 0 + self.ne, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) + + species = self.ions_flags + self.ni_all, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,nions,ntoroidal,time) + self.ni = self.ni_all.sum(axis=1) # [COMPLEX] (nradial,ntoroidal,time) + + moment, species, field = 'e', self.electron_flag, 0 + Ee, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) + + species = self.ions_flags + Ei_all, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,nions,ntoroidal,time) + Ei = Ei_all.sum(axis=1) + + # Transform to temperature + self.Te = 2/3 * Ee - self.ne + self.Ti_all = 2/3 * Ei_all - self.ni_all + self.Ti = 2/3 * Ei - self.ni + + # Sum over radial modes and divide between n=0 and n>0 modes, RMS + variables = ['phi', 'apar', 'bpar', 'ne', 'ni_all', 'ni', 'Te', 'Ti', 'Ti_all'] + for var in variables: + if var in self.__dict__: + + # Make sure I go to the real units for all of them ******************* + self.__dict__[var] = self.__dict__[var] * self.artificial_rhos_factor + # ******************************************************************** + + # Sum over radial modes + self.__dict__[var+'_rms_sumnr'] = (abs(self.__dict__[var][:,:,:])**2).sum(axis=(0))**0.5 # (ntoroidal, time) + + # Sum over radial modes AND separate n=0 and n>0 (sum) modes + self.__dict__[var+'_rms_sumnr_n0'] = (abs(self.__dict__[var][:,0,:])**2).sum(axis=0)**0.5 # (time) + self.__dict__[var+'_rms_sumnr_sumn1'] = (abs(self.__dict__[var][:,1:,:])**2).sum(axis=(0,1))**0.5 # (time) + + # Sum over radial modes and toroidal modes + self.__dict__[var+'_rms_sumnr_sumn'] = (abs(self.__dict__[var][:,:,:])**2).sum(axis=(0,1))**0.5 # (time) + + # Separate n=0, n>0 (sum) modes, and all n (sum) modes + self.__dict__[var+'_rms_n0'] = (abs(self.__dict__[var][:,0,:])**2)**0.5 # (nradial,time) + self.__dict__[var+'_rms_sumn1'] = (abs(self.__dict__[var][:,1:,:])**2).sum(axis=(1))**0.5 # (nradial,time) + self.__dict__[var+'_rms_sumn'] = (abs(self.__dict__[var][:,:,:])**2).sum(axis=(1))**0.5 # (nradial,time) + + # Cross-phases + self.neTe = _cross_phase(self.t, self.ne, self.Te) * 180/ np.pi # (nradial, ntoroidal, time) + self.neTe_kx0 = self.neTe[np.argmin(np.abs(self.kx)),:,:] # (ntoroidal, time) + + self.niTi = _cross_phase(self.t, self.ni, self.Ti) * 180/ np.pi # (nradial, ntoroidal, time) + self.niTi_kx0 = self.niTi[np.argmin(np.abs(self.kx)),:,:] + + self.phiTe = _cross_phase(self.t, self.phi, self.Te) * 180/ np.pi # (nradial, ntoroidal, time) + self.phiTe_kx0 = self.phiTe[np.argmin(np.abs(self.kx)),:,:] + + self.phiTi = _cross_phase(self.t, self.phi, self.Ti) * 180/ np.pi # (nradial, ntoroidal, time) + self.phiTi_kx0 = self.phiTi[np.argmin(np.abs(self.kx)),:,:] + + self.phiTi_all = [] + for ion in self.ions_flags: + self.phiTi_all.append(_cross_phase(self.t, self.phi, self.Ti_all[:,ion,:]) * 180/ np.pi) + self.phiTi_all = np.array(self.phiTi_all) + self.phiTi_all_kx0 = self.phiTi_all[:,np.argmin(np.abs(self.kx)),:,:] + + self.phine = _cross_phase(self.t, self.phi, self.ne) * 180/ np.pi # (nradial, ntoroidal, time) + self.phine_kx0 = self.phine[np.argmin(np.abs(self.kx)),:,:] + + self.phini = _cross_phase(self.t, self.phi, self.ni) * 180/ np.pi # (nradial, ntoroidal, time) + self.phini_kx0 = self.phini[np.argmin(np.abs(self.kx)),:,:] + + self.phini_all = [] + for ion in self.ions_flags: + self.phini_all.append(_cross_phase(self.t, self.phi, self.ni_all[:,ion,:]) * 180/ np.pi) + self.phini_all = np.array(self.phini_all) + self.phini_all_kx0 = self.phini_all[:,np.argmin(np.abs(self.kx)),:,:] + + # Correlation length + phi = (abs(self.phi[:,self.ky>0,:])).sum(axis=1) # Sum over toroidal modes n>0 + phim, _ = apply_ac(self.t,phi,tmin=self.tmin) + phim = np.append(0, phim) # Add n=0 mode + if np.isinf(phim).any() or np.isnan(phim).any(): + print(f"\t- Warning: Correlation length calculation failed due to infinite/nan values. Setting l_corr to NaN.", typeMsg='w') + self.l_corr = np.nan + else: + self.lr_corr = calculate_lcorr(phim, self.kx, self.cgyrodata.n_radial) + + def _process_fluxes(self): + + # ************************ + # Fluxes + # ************************ + + ky_flux = self.cgyrodata.ky_flux # (species, moments, fields, ntoroidal, time) + + # Electron energy flux + + i_species, i_moment = -1, 1 + i_fields = 0 + self.Qe_ES_ky = ky_flux[i_species, i_moment, i_fields, :, :] + i_fields = 1 + self.Qe_EM_apar_ky = ky_flux[i_species, i_moment, i_fields, :, :] + i_fields = 2 + self.Qe_EM_aper_ky = ky_flux[i_species, i_moment, i_fields, :, :] + + self.Qe_EM_ky = self.Qe_EM_apar_ky + self.Qe_EM_aper_ky + self.Qe_ky = self.Qe_ES_ky + self.Qe_EM_ky + + # Electron particle flux + + i_species, i_moment = -1, 0 + i_fields = 0 + self.Ge_ES_ky = ky_flux[i_species, i_moment, i_fields, :] + i_fields = 1 + self.Ge_EM_apar_ky = ky_flux[i_species, i_moment, i_fields, :] + i_fields = 2 + self.Ge_EM_aper_ky = ky_flux[i_species, i_moment, i_fields, :] + + self.Ge_EM_ky = self.Ge_EM_apar_ky + self.Ge_EM_aper_ky + self.Ge_ky = self.Ge_ES_ky + self.Ge_EM_ky + + # Ions energy flux + + i_species, i_moment = self.ions_flags, 1 + i_fields = 0 + self.Qi_all_ES_ky = ky_flux[i_species, i_moment, i_fields, :] + i_fields = 1 + self.Qi_all_EM_apar_ky = ky_flux[i_species, i_moment, i_fields, :] + i_fields = 2 + self.Qi_all_EM_aper_ky = ky_flux[i_species, i_moment, i_fields, :] + + self.Qi_all_EM_ky = self.Qi_all_EM_apar_ky + self.Qi_all_EM_aper_ky + self.Qi_all_ky = self.Qi_all_ES_ky + self.Qi_all_EM_ky - turflux=np.sum(turflux_all,axis=0) + self.Qi_ky = self.Qi_all_ky.sum(axis=0) + self.Qi_EM_ky = self.Qi_all_EM_ky.sum(axis=0) + self.Qi_EM_apar_ky = self.Qi_all_EM_apar_ky.sum(axis=0) + self.Qi_EM_aper_ky = self.Qi_all_EM_aper_ky.sum(axis=0) + self.Qi_ES_ky = self.Qi_all_ES_ky.sum(axis=0) + # ************************ + # Sum total + # ************************ + variables = ['Qe','Ge','Qi','Qi_all'] + for var in variables: + for i in ['', '_ES', '_EM_apar', '_EM_aper', '_EM']: + self.__dict__[var+i] = self.__dict__[var+i+'_ky'].sum(axis=-2) # (time) + + # Convert to MW/m^2 + self.QeMWm2 = self.Qe * self.Qgb + self.QiMWm2 = self.Qi * self.Qgb + self.Qi_allMWm2 = self.Qi_all * self.Qgb + + def _saturate_signals(self): + + # ************************ + # Saturated + # ************************ + + flags = [ + 'Qe', + 'QeMWm2', + 'Qe_ky', + 'Qi', + 'QiMWm2', + 'Qi_all', + 'Qi_allMWm2', + 'Qi_ky', + 'Qi_all_ky', + 'Ge', + 'Ge_ky', + 'Qe_ES', + 'Qi_ES', + 'Ge_ES', + 'Qe_EM', + 'Qi_EM', + 'Ge_EM', + 'g', + 'f', + ] + + flags_fluctuations = [ + 'phi_rms_sumnr', + 'apar_rms_sumnr', + 'bpar_rms_sumnr', + 'ne_rms_sumnr', + 'ni_rms_sumnr', + 'ni_all_rms_sumnr', + 'Te_rms_sumnr', + 'Ti_rms_sumnr', + 'Ti_all_rms_sumnr', + 'phi_rms_n0', + 'phi_rms_sumn1', + 'phi_rms_sumn', + 'apar_rms_n0', + 'apar_rms_sumn1', + 'apar_rms_sumn', + 'bpar_rms_n0', + 'bpar_rms_sumn1', + 'bpar_rms_sumn', + 'ne_rms_n0', + 'ne_rms_sumn1', + 'ne_rms_sumn', + 'ni_rms_n0', + 'ni_rms_sumn1', + 'ni_rms_sumn', + 'ni_all_rms_n0', + 'ni_all_rms_sumn1', + 'ni_all_rms_sumn', + 'Te_rms_n0', + 'Te_rms_sumn1', + 'Te_rms_sumn', + 'Ti_rms_n0', + 'Ti_rms_sumn1', + 'Ti_rms_sumn', + 'Ti_all_rms_n0', + 'Ti_all_rms_sumn1', + 'Ti_all_rms_sumn', + 'neTe_kx0', + 'niTi_kx0', + 'phiTe_kx0', + 'phine_kx0', + 'phini_kx0', + 'phiTi_kx0', + 'phini_all_kx0', + 'phiTi_all_kx0', + ] + + for iflag in flags: + if iflag in self.__dict__: + self.__dict__[iflag+'_mean'], self.__dict__[iflag+'_std'] = apply_ac( + self.t, + self.__dict__[iflag], + tmin=self.tmin, + label_print=iflag, + print_msg=iflag in ['Qi', 'Qe', 'Ge'], + ) + + for iflag in flags_fluctuations: + if iflag in self.__dict__: + self.__dict__[iflag+'_mean'], self.__dict__[iflag+'_std'] = apply_ac( + self.t, + self.__dict__[iflag], + tmin=self.tmin, + tmax=self.tmax_fluct, + label_print=iflag, + ) + +def _grab_ncorrelation(S, debug=False): + # Calculate the autocorrelation function + i_acf = sm.tsa.acf(S, nlags=len(S)) + + if i_acf.min() > 1/np.e: + print("Autocorrelation function does not reach 1/e, will use full length of time series for n_corr.", typeMsg='w') + + # Calculate how many time slices make the autocorrelation function is 1/e (conventional decorrelation level) + icor = np.abs(i_acf-1/np.e).argmin() - if np.amax(ky) > 1.0: - e_ind=np.where(ky > 1.0)[0] - eflux_elec_tmp1=flux[:,:,:,e_ind,:] - eflux_elec_tmp2=np.sum(eflux_elec_tmp1,axis=3) - eflux_elec_all=eflux_elec_tmp2[n_spec-1,1,:,:] - eflux_elec=np.sum(eflux_elec_all,axis=0) - - #Define max values for plot - imax=np.amax(iflux)*qgb - emax=np.amax(eflux)*qgb - smax=imax+emax - - #Determine if the time step has changed mid simulation - #Change the value of tstart internally if it has - tstart_save=tstart - tend=data.t[nt-1] - tstart=(np.abs(t - tstart)).argmin() + # Define number of samples + n_corr = len(S) / ( 3.0 * icor ) #Define "sample" as 3 x autocor time - #Take the mean values of the fluxes - m_qe= np.mean(eflux[int(tstart):int(nt)+1]) - m_qi= np.mean(iflux[int(tstart):int(nt)+1]) - m_ge= np.mean(epflux[int(tstart):int(nt)+1]) - m_qe_elec=np.mean(eflux_elec[int(tstart):int(nt)+1]) - m_ge_p=np.mean(epfluxp[int(tstart):int(nt)+1]) - m_ge_ap=np.mean(epfluxap[int(tstart):int(nt)+1]) - m_ge_bp=np.mean(epfluxbp[int(tstart):int(nt)+1]) - m_gimp=np.mean(impflux[int(tstart):int(nt)+1]) - m_mo=np.mean(mflux[int(tstart):int(nt)+1]) - m_tur=np.mean(turflux[int(tstart):int(nt)+1]) - print(nt) - - #Calculate the standard deviations of the fluxes based on autocorrelation - i_tmp=iflux[int(tstart):int(nt-1)] - e_tmp=eflux[int(tstart):int(nt-1)] - ep_tmp=epflux[int(tstart):int(nt-1)] - imp_tmp=impflux[int(tstart):int(nt-1)] - mo_tmp=mflux[int(tstart):int(nt-1)] - tur_tmp=turflux[int(tstart):int(nt-1)] - i_acf=sm.tsa.acf(i_tmp) - i_array=np.asarray(i_acf) - icor=(np.abs(i_array-0.36)).argmin() - e_acf=sm.tsa.acf(e_tmp) - e_array=np.asarray(e_acf) - ecor=(np.abs(e_array-0.36)).argmin() - ep_acf=sm.tsa.acf(ep_tmp) - ep_array=np.asarray(ep_acf) - epcor=(np.abs(ep_array-0.36)).argmin() - imp_acf=sm.tsa.acf(imp_tmp) - imp_array=np.asarray(imp_acf) - impcor=(np.abs(imp_array-0.36)).argmin() - mo_acf=sm.tsa.acf(mo_tmp) - mo_array=np.asarray(mo_acf) - mocor=(np.abs(mo_array-0.36)).argmin() - tur_acf=sm.tsa.acf(tur_tmp) - tur_array=np.asarray(tur_acf) - turcor=(np.abs(tur_array-0.36)).argmin() - - n_corr_i=(nt-tstart)/(3.0*icor) #Define "sample" as 3 x autocor time - n_corr_e=(nt-tstart)/(3.0*ecor) - n_corr_ep=(nt-tstart)/(3.0*epcor) - n_corr_imp=(nt-tstart)/(3.0*impcor) - n_corr_mo=(nt-tstart)/(3.0*mocor) - n_corr_tur=(nt-tstart)/(3.0*turcor) - print(n_corr_i) - print(n_corr_e) - print(n_corr_ep) - print(n_corr_imp) - print(n_corr_mo) - print(n_corr_tur) - std_qe=np.std(eflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_e) - std_qi=np.std(iflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_i) - std_ge=np.std(epflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_ep) - std_gimp=np.std(impflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_imp) - std_mo=np.std(mflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_mo) - std_tur=np.std(turflux[int(tstart):int(nt-1)])/np.sqrt(n_corr_tur) + if debug: + fig, ax = plt.subplots() + ax.plot(i_acf, '-o', label='ACF') + ax.axhline(1/np.e, color='r', linestyle='--', label='1/e') + ax.set_xlabel('Lags'); ax.set_xlim([0, icor+20]) + ax.set_ylabel('ACF') + ax.legend() + plt.show() + embed() - m_qees= np.mean(efluxes[int(tstart):int(nt-1)]) - m_qeem= np.mean(efluxem[int(tstart):int(nt-1)]) + return n_corr, icor - #Make some string values so they can be truncated - s_alti=str(alti) - s_alte=str(alte) - s_alne=str(alne) +def apply_ac(t, S, tmin = 0, tmax = None, label_print = '', print_msg = False, debug=False): + + it0 = np.argmin(np.abs(t - tmin)) + it1 = np.argmin(np.abs(t - tmax)) if tmax is not None else len(t) # If tmax is None, use the full length of t - #Print all the simulation information - print('') - print('========================================') - print('Time to start is:') - print(tstart) - print('Max simulation time is') - print(nt) - print('Simulation Gradients') - print('a/LTi = '+s_alti[:6]) - print('a/LTe = '+s_alte[:6]) - print('a/Lne = '+s_alne[:6]) - print('') - print('======================') - print('Heat Flux') - print('======================') - print('Q_gb is') - print(f"{qgb:.4f}") - print('Mean Qi (in GB)') - print(f"{m_qi:.4f}") - print('Qi Std deviation (in GB)') - print(f"{std_qi:.4f}") - print('Mean in MW/m^2') - print(f"{m_qi*qgb:.4f}") - print('----------------------') - print('Mean Qe (in GB)') - print(f"{m_qe:.4f}") - print('Qe Std deviation (in GB)') - print(f"{std_qe:.4f}") - print('Mean in MW/m^2') - print(f"{m_qe*qgb:.4f}") - print('Qi/Qe') - print(f"{m_qi/m_qe:.4f}") - print('High-k Qe (GB)') - print(f"{m_qe_elec:.4f}") - print('') - print('======================') - print('Particle Flux') - print('======================') - print('Gamma_gb is:') - print(f"{ggb:.4f}") - print('Mean Gamma_e (in GB)') - print(f"{m_ge:.4f}") - print('Gamma_e Std deviation (in GB)') - print(f"{std_ge:.4f}") - print('Mean in e19/m^2*s') - print(f"{m_ge*ggb:.4f}") - print('Mean Impurity Flux (GB)') - print(f"{m_gimp:.4e}") - print('Gamma_imp Std. deviation (in GB)') - print(f"{std_gimp:.4e}") - print('Electron Particle Flux Components(GB)') - print('Phi: 'f"{m_ge_p:.4f}") - print('A_||: 'f"{m_ge_ap:.4f}") - print('B_||: 'f"{m_ge_bp:.4f}") - print('') - print('======================') - print('Momentum Flux') - print('======================') - print('Pi_gb is:') - print(f"{pgb:.4f}") - print('Mean Pi (in GB)') - print(f"{m_mo:.4f}") - print('Momentum Flux Std deviation (in GB)') - print(f"{std_mo:.4f}") - print('Mean in J/m^2') - print(f"{m_mo*pgb:.4f}") - print('') - print('======================') - print('Turbulent Exchange') - print('======================') - print('S_gb is:') - print(f"{sgb:.4f}") - print('Mean Elec Turb Ex. (in GB)') - print(f"{m_tur:.4f}") - print('Turb Ex. Std deviation (in GB)') - print(f"{std_tur:.4f}") + if it1 <= it0: + it0 = it1 + + # Calculate the mean and std of the signal after tmin (last dimension is time) + S_mean = np.mean(S[..., it0:it1+1], axis=-1) + S_std = np.std(S[..., it0:it1+1], axis=-1) + + if S.ndim == 1: + # 1D case: single time series + n_corr, icor = _grab_ncorrelation(S[it0:it1+1], debug=debug) + S_std = S_std / np.sqrt(n_corr) + + if print_msg: + print(f"\t- {(label_print + ': a') if len(label_print)>0 else 'A'}utocorr time: {icor:.1f} -> {n_corr:.1f} samples -> {S_mean:.2e} +-{S_std:.2e}") + + else: + # Multi-dimensional case: flatten all dimensions except the last one + shape_orig = S.shape[:-1] # Original shape without time dimension + S_reshaped = S.reshape(-1, S.shape[-1]) # Flatten to (n_series, n_time) + + n_series = S_reshaped.shape[0] + n_corr = np.zeros(n_series) + icor = np.zeros(n_series) + + # Calculate correlation for each flattened time series + for i in range(n_series): + n_corr[i], icor[i] = _grab_ncorrelation(S_reshaped[i, it0:it1+1], debug=debug) + + # Reshape correlation arrays back to original shape (without time dimension) + n_corr = n_corr.reshape(shape_orig) + icor = icor.reshape(shape_orig) + + # Apply correlation correction to standard deviation + S_std = S_std / np.sqrt(n_corr) + + # Print results - handle different dimensionalities + if print_msg: + if S.ndim == 2: + # 2D case: print each series + for i in range(S.shape[0]): + print(f"\t- {(label_print + f'_{i}: a') if len(label_print)>0 else 'A'}utocorr: {icor[i]:.1f} -> {n_corr[i]:.1f} samples -> {S_mean[i]:.2e} +-{S_std[i]:.2e}") + else: + # Higher dimensional case: print summary statistics + print(f"\t- {(label_print + ': a') if len(label_print)>0 else 'A'}utocorr time: {icor.mean():.1f}±{icor.std():.1f} -> {n_corr.mean():.1f}±{n_corr.std():.1f} samples -> shape {S_mean.shape}") + + return S_mean, S_std + + +def _cross_phase(t, f1, f2): + """ + Calculate the cross-phase between two complex signals. - #Reset the starting time to correspond to the a/cs - tstart=tstart_save + Parameters: + f1, f2 : np.ndarray + Complex signals (e.g., fluctuations). + + Returns: + np.ndarray + Cross-phase in radians. + """ + + return np.angle(f1 * np.conj(f2)) + +def _detect_exploiding_signal(t,f1): + + try: + idx = np.where(np.isnan(f1.sum(axis=(0,1))) | np.isinf(f1.sum(axis=(0,1))))[0][0] + max_t = t[idx] + if print(f"\t- Warning: Exploding signal detected at t>={max_t:.2f}", typeMsg='w'): + return max_t + else: + return t[-1] + except IndexError: + return t[-1] # No exploding signal detected, return last time point + +def calculate_lcorr(phim, kx, nx, debug=False): + """Calculate the correlation length in the radial direction. + + Completely based on pygacode + """ + + ave = np.roll(phim,-nx//2) + ave[0] = 0.0 + corr = np.fft.fft(ave,nx) + corr = np.fft.fftshift(corr) + corr /= np.max(np.abs(corr)) + corr = corr.real + delta_r = np.fft.fftfreq(nx) + delta_r = np.fft.fftshift(delta_r) + Lx = 2*np.pi/(kx[1]-kx[0]) + delta_r *= Lx - #Print the results to results file - if printflag ==1: - file1=open(file,"a") - file1.write('r/a = {0:4f} \t a/Lne = {1:4f} \t a/LTi = {2:4f} \t a/LTe = {3:4f} \t Qi = {4:4f} \t Qi_std = {5:4f} \t Qe = {6:4f} \t Qe_std = {7:4f} \t Ge = {8:4f} \t Ge_std = {9:4f} \t Gimp = {10:.4e} \t Gimp_std = {11:.4e} \t Pi = {12:4f} \t Pi_std = {13:4f} \t S = {14:4f} \t S_std = {15:4f} \t Q_gb = {16:4f} \t G_gb = {17:4f} \t Pi_gb = {18:4f} \t S_gb = {19:4f} \t tstart = {20:4f} \t tend = {21:4f} \n'.format(roa,alne,alti,alte,m_qi,std_qi,m_qe,std_qe,m_ge,std_ge,m_gimp,std_gimp,m_mo,std_mo,m_tur,std_tur,qgb,ggb,pgb,sgb,tstart,nt)) - file1.close() - - - - #Plot the results - if plotflag == 1: - - summax=2.0*(m_qe*qgb+m_qi*qgb) - maxf=np.amax([np.amax(iflux)*qgb,np.amax(eflux)*qgb]) - pmax=np.amin([2.0*(m_qe*qgb+m_qi*qgb),maxf]) - gmax=np.amax(epflux*ggb) - gmin=np.amin(epflux*ggb) - - # Setting the style for the plots - plt.rc('axes', labelsize=25) - plt.rc('xtick', labelsize=25) - plt.rc('ytick', labelsize=25) - - plt.rcParams['font.family'] = 'serif' - - # Create a figure with 3 subplots vertically stacked (3 rows, 1 column) - fig, ax = plt.subplots(2, 1, figsize=(10.0, 9.0)) - - # First plot (Q vs Time) - ax[0].plot(t, iflux * qgb, 'b-') - ax[0].plot([tstart, tend], [m_qi * qgb, m_qi * qgb], 'b-', linewidth=4.0) - ax[0].plot([tstart, tend], [(m_qi + std_qi) * qgb, (m_qi + std_qi) * qgb], 'b-', linewidth=2.0) - ax[0].plot([tstart, tend], [(m_qi - std_qi) * qgb, (m_qi - std_qi) * qgb], 'b-', linewidth=2.0) - ax[0].set_title('Q$_e$ and Q$_i$ vs Time',fontsize=18) - ax[0].set_xlabel('a/c$_s$',fontsize=15) - ax[0].set_ylabel('MW/m$^2$',fontsize=15) - ax[0].axis([0.0, tmax, 0.0, pmax]) - - # Second plot (Qe vs Time) - ax[0].plot(t, eflux * qgb, 'r-') - ax[0].plot([tstart, tend], [m_qe * qgb, m_qe * qgb], 'r-', linewidth=4.0) - ax[0].plot([tstart, tend], [(m_qe + std_qe) * qgb, (m_qe + std_qe) * qgb], 'r-', linewidth=2.0) - ax[0].plot([tstart, tend], [(m_qe - std_qe) * qgb, (m_qe - std_qe) * qgb], 'r-', linewidth=2.0) - - ax[0].fill_between([tstart,tend], [(m_qe - std_qe) * qgb, (m_qe - std_qe) * qgb], [(m_qe + std_qe) * qgb, (m_qe + std_qe) * qgb], color='red', alpha=0.2) - ax[0].fill_between([tstart,tend], [(m_qi - std_qi) * qgb, (m_qi - std_qi) * qgb], [(m_qi + std_qi) * qgb, (m_qi + std_qi) * qgb], color='blue', alpha=0.2) - - #ax[0].plot(t, eflux * qgb, 'g-') # You can customize this line further - ax[0].tick_params(axis='x', labelsize=12) # x-axis tick labels font size reduced - ax[0].tick_params(axis='y', labelsize=12) - factor=10 ** 3 - factor2=10 ** 2 - num=math.floor(m_qe*qgb*factor)/factor - num2=math.floor(std_qe*qgb*factor)/factor - ax[0].text(10,pmax*0.8,'$Q_e$='+str(num)+' +/- '+str(num2),fontsize=18,color='red') - - num=math.floor(m_qi*qgb*factor)/factor - num2=math.floor(std_qi*qgb*factor)/factor - num3=math.floor(m_qi/m_qe*factor2)/factor2 - ax[0].text(10,pmax*0.88,'$Q_i$='+str(num)+' +/- '+str(num2),fontsize=18,color='blue') - ax[0].text(10,pmax*0.72,'$Q_i/Q_e$='+str(num3),fontsize=18) - - # Third plot (epflux vs Time) - ax[1].plot(t, epflux * ggb, 'g-') # Plotting epflux on the third subplot - ax[1].set_title('$\Gamma_e$ vs Time',fontsize=18) - ax[1].set_xlabel('a/c$_s$',fontsize=15) - ax[1].set_ylabel('1e19/m$^2$s',fontsize=15) - ax[1].tick_params(axis='x', labelsize=12) # x-axis tick labels font size reduced - ax[1].plot([tstart, tend], [m_ge * ggb, m_ge * ggb], 'g-', linewidth=4.0) - ax[1].plot([tstart, tend], [(m_ge + std_ge) * ggb, (m_ge + std_ge) * ggb], 'g-', linewidth=2.0) - ax[1].plot([tstart, tend], [(m_ge - std_ge) * ggb, (m_ge - std_ge) * ggb], 'g-', linewidth=2.0) - ax[1].fill_between([tstart,tend], [(m_ge - std_ge) * ggb, (m_ge - std_ge) * ggb], [(m_ge + std_ge) * ggb, (m_ge + std_ge) * ggb], color='green', alpha=0.2) - - ax[1].plot(t, epflux * ggb, 'g-') - ax[1].plot(t,t*0.0,'k',linestyle='--') - ax[1].tick_params(axis='y', labelsize=12) - ax[1].axis([0.0, tmax, gmin, gmax]) - - factor=10 ** 3 - num=math.floor(m_ge*ggb*factor)/factor - num2=math.floor(std_ge*ggb*factor)/factor - ax[1].text(10,gmax*0.8,'$\Gamma_e$='+str(num)+' +/- '+str(num2),fontsize=18) - ax[1].text(10,gmin*0.95,str(int(tend))+' a/c$_s$',fontsize=18) - # Adjust spacing between subplots - plt.tight_layout() - - # Show the plot - plt.show() + corr_hilbert = scipy.signal.hilbert(corr) + corr_env = np.abs(corr_hilbert) + def absexp(x,tau): + return np.exp(-np.abs(x)/tau) + l_corr, _ = scipy.optimize.curve_fit(absexp, delta_r, corr_env, p0=10.0) - return roa,alne,alti,alte,m_qi,std_qi,m_qe,std_qe,m_ge,std_ge,m_gimp,std_gimp,m_mo,std_mo,m_tur,std_tur,qgb,ggb,pgb,sgb,tstart,nt + if debug: + fig, ax = plt.subplots() + ax.plot(delta_r,0*delta_r,color='k',ls='--') + ax.plot(delta_r,corr,color='m') + ax.plot(delta_r,corr_env,color='b') + ax.plot(delta_r,absexp(delta_r,l_corr),color='k',ls='-.') + ax.set_xlim([np.min(delta_r),np.max(delta_r)]) + ax.set_ylim(-1,1) + plt.show() + embed() + return l_corr[0] # Return the correlation length in the radial direction -if __name__ == "__main__": - # Example usage - data_dir = sys.argv[1] # Directory containing the cgyro data - tstart = float(sys.argv[2]) # Start time for analysis - plotflag = int(sys.argv[3]) # Flag to indicate if plotting is required - printflag = int(sys.argv[4]) # Flag to indicate if printing results is required - file = sys.argv[5] # Print to this file - grab_cgyro_nth(data_dir, tstart, plotflag, printflag, file = file) +def quends_analysis(t, S, debug = False): + + time_dependent_data = {'time': t, 'signal': S} + df = pd.DataFrame(time_dependent_data, index = pd.RangeIndex(len(t))) + + dst = qnds.DataStream(df) + window_size = 10 + + trimmed_df = dst.trim(column_name="signal", method="std") #, window_size=10) + + mean = trimmed_df.mean(window_size=window_size)['signal'] + std = trimmed_df.mean_uncertainty(window_size=window_size)['signal'] + + stats = trimmed_df.compute_statistics(window_size=window_size) + + if debug: + plotter = qnds.Plotter() + plotter.steady_state_automatic_plot(dst, ["signal"]) + plotter.plot_acf(trimmed_df) + print(stats) + plt.show() + embed() + + return mean, std, stats diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 9ec6e86e..9ef6c8f8 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -1022,7 +1022,7 @@ def runTGLF( else: - # Job array + # Standard job if total_cores_required < max_cores_per_node: print(f"\t- TGLF will be executed in SLURM as standard job (cpus: {total_cores_required})",typeMsg="i") @@ -1040,7 +1040,7 @@ def runTGLF( ntasks = total_tglf_executions cpuspertask = cores_tglf - # Standard job + # Job array else: #raise Exception("TGLF array not implemented yet") print(f"\t- TGLF will be executed in SLURM as job array due to its size (cpus: {total_cores_required})",typeMsg="i") diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index c9ee8ad0..2fdfd2e6 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -1246,7 +1246,15 @@ def perform_quick_remote_execution( job.run(check_if_files_received=check_if_files_received) -def retrieve_files_from_remote(folder_local, machine, files_remote = [], folders_remote = [], purge_tmp_files = False, ensure_files = True): +def retrieve_files_from_remote( + folder_local, + machine, + files_remote = [], + folders_remote = [], + only_folder_structure_with_files = None, # If not None, only the folder structure is retrieved, with files in the list + purge_tmp_files = False, + ensure_files = True + ): ''' Quick routine for file retrieval from remote machine (assumes remote machine is linux) diff --git a/src/mitim_tools/misc_tools/GRAPHICStools.py b/src/mitim_tools/misc_tools/GRAPHICStools.py index 04be8a6c..5a8c9c79 100644 --- a/src/mitim_tools/misc_tools/GRAPHICStools.py +++ b/src/mitim_tools/misc_tools/GRAPHICStools.py @@ -666,6 +666,14 @@ def fillGraph( ms=None, ls="-", ): + + if IOtools.isnum(y): + y = np.array([y] * len(x)) + if IOtools.isnum(y_down): + y_down = np.array([y_down] * len(x)) + if IOtools.isnum(y_up): + y_up = np.array([y_up] * len(x)) + if y_up is not None: l = ax.fill_between(x, y, y_up, facecolor=color, alpha=alpha, label=label) if y_down is not None: @@ -690,6 +698,19 @@ def fillGraph( return l +def adjust_subplots(fig = None, axs=None, vertical=0.4, horizontal=0.3): + + if fig is None and axs is None: + raise ValueError("Either fig or axs must be provided") + + if axs is not None: + fig = next(iter(axs.values())).get_figure() + + fig.subplots_adjust( + hspace=vertical, # vertical spacing between rows + wspace=horizontal # horizontal spacing between columns + ) + def listColors(): col = [ @@ -1072,44 +1093,6 @@ def colorTableFade(num, startcolor="b", endcolor="r", alphalims=[1.0, 1.0]): return cn, cpick -def createAnimation( - fig, FunctionToAnimate, framesCalc, FramesPerSecond, BITrate, MovieFile, DPIs -): - plt.rcParams["animation.ffmpeg_path"] = "/usr/local/bin/ffmpeg" - if "mfews" in socket.gethostname(): - plt.rcParams["animation.ffmpeg_path"] = "/usr/bin/ffmpeg" - ani = animation.FuncAnimation( - fig, FunctionToAnimate, frames=framesCalc, repeat=True - ) - - Writer = animation.writers["ffmpeg"] - writer = Writer(fps=FramesPerSecond, metadata=dict(artist="PRF"), bitrate=BITrate) - ani.save(writer=writer, filename=MovieFile, dpi=DPIs) - - -def animageFunction( - plottingFunction, - axs, - fig, - MovieFile, - HowManyFrames, - framePS=50, - BITrate=1200, - DPIs=150, -): - if type(axs) not in [np.ndarray, list]: - axs = [axs] - - def animate(i): - for j in range(len(axs)): - axs[j].clear() - plottingFunction(axs, i) - print(f"\t~~ Frame {i + 1}/{HowManyFrames}") - - print(" --> Creating animation") - createAnimation(fig, animate, HowManyFrames, framePS, BITrate, MovieFile, DPIs) - - def reduceVariable(var, howmanytimes, t=None, trange=[0, 100]): if t is not None: var = var[np.argmin(np.abs(t - trange[0])) : np.argmin(np.abs(t - trange[1]))] @@ -1493,4 +1476,42 @@ def PSFCcolors(): colors["Heated"] = "#F25757" colors["Orange Edge"] = "#FFA630" - return colors \ No newline at end of file + return colors + + +''' +******************************************************************** +Capabilities to create animations using matplotlib's FuncAnimation. +******************************************************************** +''' + +def animateFunction( + plottingFunction, # Function that plots the data, must receive axs and frame index as arguments + axs, # List of axes to plot on + fig, # Figure object to create the animation in + MovieFile, # Output filename for the movie + HowManyFrames: int, # Total number of frames in the animation + framePS: int = 50, # Frames per second for the animation + BITrate: int = 1200, # Bitrate for the output video + DPIs: int = 150, # Dots per inch for the output video + ffmpeg_path = None # Path to ffmpeg executable, if None it uses the default. e.g. /usr/local/bin/ffmpeg, /usr/bin/ffmpeg +): + if type(axs) not in [np.ndarray, list]: + axs = [axs] + + def animate(i): + for j in range(len(axs)): + axs[j].clear() + plottingFunction(axs, i) + print(f"\t~~ Frame {i + 1}/{HowManyFrames}") + + print(" --> Creating animation") + + if ffmpeg_path is not None: + plt.rcParams["animation.ffmpeg_path"] = ffmpeg_path + + ani = animation.FuncAnimation(fig, animate, frames=HowManyFrames, repeat=True) + + Writer = animation.writers["ffmpeg"] + writer = Writer(fps=framePS, metadata=dict(artist="PRF"), bitrate=BITrate) + ani.save(writer=writer, filename=MovieFile, dpi=DPIs) diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index d417f6a7..08c63c65 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -59,41 +59,118 @@ def _get_time(self): self.timeDiff = getTimeDifference(self.timeBeginning, niceText=False) self.profiler.dump_stats(self.file) - print( - f'Script took {createTimeTXT(self.timeDiff)}, profiler stats dumped to {self.file} (open with "python3 -m snakeviz {self.file}")' - ) - -class timer(object): - - def __init__(self, name="\t* Script", name_timer = '\t* Start time: '): - self.name = name - self.name_timer = name_timer + print(f'Script took {createTimeTXT(self.timeDiff)}, profiler stats dumped to {self.file} (open with "python3 -m snakeviz {self.file}")') +class timer: + ''' + Context manager to time a script or function execution. + ''' + # ──────────────────────────────────────────────────────────────────── + def __init__(self, + name: str = "Script", # Name of the script for printing, visualization + print_at_entering: str | None = None, # Prefix printed right before the timer starts + log_file: Path | None = None): # File to log the timing information in JSON format + self.name = name + self.print_at_entering = print_at_entering + self.log_file = log_file + + # ──────────────────────────────────────────────────────────────────── def __enter__(self): - self.timeBeginning = datetime.datetime.now() - if self.name_timer is not None: print(f'{self.name_timer}{self.timeBeginning.strftime("%Y-%m-%d %H:%M:%S")}') - return self - - def __exit__(self, *args): - self._get_time() + # high-resolution timer + wall-clock stamp + + self.t0_wall = time.perf_counter() + self.t0 = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - def _get_time(self): + if self.print_at_entering: + print(f'{self.print_at_entering}{self.t0}') + return self - self.timeDiff = getTimeDifference(self.timeBeginning, niceText=False) + # ──────────────────────────────────────────────────────────────────── + def __exit__(self, exc_type, exc, tb): + self._finish() + return False # propagate any exception - print(f'{self.name} took {createTimeTXT(self.timeDiff)}') + # ──────────────────────────────────────────────────────────────────── + def _finish(self): + + dt = time.perf_counter() - self.t0_wall + t1 = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + + print(f'\t\t* {self.name} took {createTimeTXT(dt)}') + + if self.log_file: + record = { + "script" : self.name, + "t_start" : self.t0, + "ts_end" : t1, + "duration_s" : dt, + } + with Path(self.log_file).open("a", buffering=1) as f: + f.write(json.dumps(record) + "\n") # Decorator to time functions -def mitim_timer(name="\t* Script",name_timer = '\t* Start time: '): +def mitim_timer( + name: str | None = None, + print_at_entering: str | None = None, + log_file: str | Path | Callable[[object], str | Path] | None = None + ): + """ + Decorator that times a function / method and optionally appends one JSON + line to *log_file* after the call finishes. + + Parameters + ---------- + name : str | None + Human-readable beat name. If None, defaults to the wrapped function's __name__. + print_at_entering : str + Prefix printed right before the timer starts + log_file : str | Path | callable(self) -> str | Path | None + • str / Path → literal path written every time the beat finishes + • callable → called **at call time** with the bound instance + (`self`) and must return the path to use + • None → no file is written, only console timing is printed + + Notes + ----- + *When* the wrapper runs it has access to the bound instance (`self`), so + callable argument values let you access self variables. + """ + def decorator_timer(func): + script_name = name or func.__name__ + @functools.wraps(func) def wrapper_timer(*args, **kwargs): - with timer(name,name_timer=name_timer): + # -------------------- resolve name -------------------------- + if callable(script_name): + # assume first positional arg is `self` for bound methods + instance = args[0] if args else None + chosen_script_name = script_name(instance) + else: + chosen_script_name = script_name + # --------------------------------------------------------------- + # -------------------- resolve log_file -------------------------- + if callable(log_file): + # assume first positional arg is `self` for bound methods + instance = args[0] if args else None + chosen_log_file = log_file(instance) + else: + chosen_log_file = log_file + # --------------------------------------------------------------- + + # Your original context-manager timer class: + with timer(chosen_script_name, + print_at_entering=print_at_entering, + log_file=chosen_log_file): return func(*args, **kwargs) + return wrapper_timer + return decorator_timer +# ------------------------------------ + # Decorator to hook methods before and after execution def hook_method(before=None, after=None): def decorator(func): @@ -1889,6 +1966,19 @@ def shutil_rmtree(item): shutil.move(item, new_item) print(f"> Folder {clipstr(item)} could not be removed. Renamed to {clipstr(new_item)}",typeMsg='w') +def recursive_backup(file, extension='bak'): + + index = 0 + file_new = file.with_suffix(f".{extension}.{index}") + + while file_new.exists(): + index += 1 + file_new = file.with_suffix(f".{extension}.{index}") + + shutil.copy2(file, file_new) + print(f"> File {clipstr(file)} backed up to {clipstr(file_new)}", typeMsg='i') + + def unpickle_mitim(file): with open(str(file), "rb") as handle: diff --git a/src/mitim_tools/misc_tools/LOGtools.py b/src/mitim_tools/misc_tools/LOGtools.py index 48b7ba02..f89b8387 100644 --- a/src/mitim_tools/misc_tools/LOGtools.py +++ b/src/mitim_tools/misc_tools/LOGtools.py @@ -164,9 +164,9 @@ def __exit__(self, exc_type, exc_val, exc_tb): ''' @contextlib.contextmanager -def conditional_log_to_file(log_file=None, msg=None): +def conditional_log_to_file(log_file=None, msg=None, write_log=True): - if log_file is not None: + if log_file is not None and write_log: with log_to_file(log_file, msg) as logger: yield logger else: diff --git a/src/mitim_tools/opt_tools/SURROGATEtools.py b/src/mitim_tools/opt_tools/SURROGATEtools.py index 23f4e0a9..38d5fdfd 100644 --- a/src/mitim_tools/opt_tools/SURROGATEtools.py +++ b/src/mitim_tools/opt_tools/SURROGATEtools.py @@ -102,9 +102,7 @@ def __init__( # Points to be added from file if ("extrapointsFile" in self.surrogate_options) and (self.surrogate_options["extrapointsFile"] is not None) and (self.output is not None) and (self.output in self.surrogate_options["extrapointsModels"]): - print( - f"\t* Requested extension of training set by points in file {self.surrogate_options['extrapointsFile']}" - ) + print(f"\t* Requested extension of training set by points in file {self.surrogate_options['extrapointsFile']}") df = pd.read_csv(self.surrogate_options["extrapointsFile"]) df_model = df[df['Model'] == self.output] @@ -114,6 +112,7 @@ def __init__( continueAdding = False else: continueAdding = True + print(f"\t\t- Found {len(df_model)} points for this output in the file, adding them to the training set", typeMsg="i") else: continueAdding = False @@ -334,17 +333,13 @@ def normalization_pass( outcome_transform_normalization._is_trained = torch.tensor(True) def fit(self): - print( - f"\t- Fitting model to {self.train_X.shape[0]+self.train_X_added.shape[0]} points" - ) + print(f"\t- Fitting model to {self.train_X.shape[0]+self.train_X_added.shape[0]} points") # --------------------------------------------------------------------------------------------------- # Define loss Function to minimize # --------------------------------------------------------------------------------------------------- - mll = gpytorch.mlls.ExactMarginalLogLikelihood( - self.gpmodel.likelihood, self.gpmodel - ) + mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.gpmodel.likelihood, self.gpmodel) # --------------------------------------------------------------------------------------------------- # Prepare for training @@ -397,17 +392,11 @@ def perform_model_fit(self, mll): (train_x,) = mll.model.train_inputs approx_mll = len(train_x) > 2000 if approx_mll: - print( - f"\t* Using approximate MLL because x has {len(train_x)} elements", - ) + print(f"\t* Using approximate MLL because x has {len(train_x)} elements") # -------------------------------------------------- # Store first MLL value - track_fval = [ - -mll.forward(mll.model(*mll.model.train_inputs), mll.model.train_targets) - .detach() - .item() - ] + track_fval = [-mll.forward(mll.model(*mll.model.train_inputs), mll.model.train_targets).detach().item()] def callback(x, y, mll=mll): track_fval.append(y.fval) diff --git a/src/mitim_tools/opt_tools/optimizers/ROOTtools.py b/src/mitim_tools/opt_tools/optimizers/ROOTtools.py index 5c6fee9a..9d448ace 100644 --- a/src/mitim_tools/opt_tools/optimizers/ROOTtools.py +++ b/src/mitim_tools/opt_tools/optimizers/ROOTtools.py @@ -1,12 +1,14 @@ import torch import copy import numpy as np +import matplotlib.pyplot as plt from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools.opt_tools.optimizers import optim from mitim_tools.opt_tools.utils import TESTtools +from mitim_tools.misc_tools import GRAPHICStools from IPython import embed -def optimize_function(fun, optimization_params = {}, writeTrajectory=False, method = 'scipy_root'): +def optimize_function(fun, optimization_params = {}, writeTrajectory=False, method = 'scipy_root', debugYN=False): np.random.seed(fun.seed) @@ -104,9 +106,27 @@ def flux_residual_evaluator(X, y_history=None, x_history=None, metric_history=No # -------------------------------------------------------------------------------------------------------- print("************************************************************************************************") - x_res, _, _, acq_evaluated = solver_fun(flux_residual_evaluator,xGuesses,solver_options=solver_options,bounds=bounds) + x_res, y_history, x_history, acq_evaluated = solver_fun(flux_residual_evaluator,xGuesses,solver_options=solver_options,bounds=bounds) print("************************************************************************************************") + if debugYN: + fig, axs = plt.subplots(nrows= 2, figsize=(6, 10)) + ax = axs[0] + ax.plot(np.array(x_history)) + ax.set_xlabel("Iteration") + ax.set_ylabel("X") + GRAPHICStools.addDenseAxis(ax) + + ax = axs[1] + ax.plot(np.abs(np.array(y_history))) + ax.set_xlabel("Iteration") + ax.set_ylabel("|Y|") + ax.set_yscale("log") + GRAPHICStools.addDenseAxis(ax) + + plt.show() + embed() + # -------------------------------------------------------------------------------------------------------- # Post-process # -------------------------------------------------------------------------------------------------------- diff --git a/src/mitim_tools/opt_tools/optimizers/optim.py b/src/mitim_tools/opt_tools/optimizers/optim.py index f8ca9d3c..c662c347 100644 --- a/src/mitim_tools/opt_tools/optimizers/optim.py +++ b/src/mitim_tools/opt_tools/optimizers/optim.py @@ -127,7 +127,7 @@ def function_for_optimizer(x, dfT1=torch.zeros(1).to(x_initial)): # Perform optimization # -------------------------------------------------------------------------------------------------------- - with IOtools.timer(name="\t- SCIPY.ROOT multi-variate root finding method"): + with IOtools.timer(name="SCIPY.ROOT multi-variate root finding method"): sol = root(function_for_optimizer, x_initial0, jac=jac_ad, method=solver, tol=tol, options=algorithm_options) # -------------------------------------------------------------------------------------------------------- @@ -193,7 +193,7 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o relax_dyn = solver_options.get("relax_dyn", False) # Dynamic relax, decreases relax if residual is not decreasing relax_dyn_decrease = solver_options.get("relax_dyn_decrease", 5) # Decrease relax by this factor relax_dyn_num = solver_options.get("relax_dyn_num", 100) # Number of iterations to average over - relax_dyn_tol = solver_options.get("relax_dyn_tol", 1e-4) # Tolerance to consider that the residual is not decreasing + relax_dyn_tol_rel = solver_options.get("relax_dyn_tol_rel", 5e-2) # Tolerance to consider that the residual is not decreasing (relative, 0.1 -> 10% minimum change) print_each = solver_options.get("print_each", 1e2) @@ -258,9 +258,8 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o break if relax_dyn and (i-its_since_last_dyn_relax > relax_dyn_num): - relax, changed, hardbreak = _dynamic_relaxation(relax, relax_dyn_decrease, metric_history, relax_dyn_num, relax_dyn_tol,i+1) - if changed: - its_since_last_dyn_relax = i + relax, hardbreak = _dynamic_relaxation(relax, relax_dyn_decrease, y_history, relax_dyn_num, relax_dyn_tol_rel,i+1) + its_since_last_dyn_relax = i if hardbreak: break @@ -288,38 +287,67 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o # The best candidate, regardless of the restarts x_best = x_history[index_best,:].unsqueeze(0) - + return x_best, y_history, x_history, metric_history -def _dynamic_relaxation(relax, relax_dyn_decrease, metric_history, relax_dyn_num, relax_dyn_tol, it, min_relax=1e-6): +def _dynamic_relaxation(relax, relax_dyn_decrease, y_history, relax_dyn_num, relax_dyn_tol_rel, it, min_relax=1e-6): ''' Logic: If the metric is not improving enough, decrease the relax parameter. To determine if the metric is improving enough, I will fit a line to the last relax_dyn_num points and check if the slope is small enough. If it is, I will decrease the relax parameter. ''' - metric_history_considered = torch.Tensor(metric_history[-relax_dyn_num:]) - - # Linear fit to the time series - x = np.arange(len(metric_history_considered)) - y = metric_history_considered - slope, intercept = np.polyfit(x, y, 1) - metric0 = intercept - metric1 = slope * len(metric_history_considered) + intercept - change_in_metric = metric1 - metric0 - - if (change_in_metric < relax_dyn_tol): - if relax.all() > min_relax: - print(f"\t\t\t<> Metric not improving enough (@{it}), decreasing relax from {relax.max():.1e} to {relax.max()/relax_dyn_decrease:.1e}") - relax = relax / relax_dyn_decrease - return relax, True, False - else: + # Only consider a number of last iterations + y_history = torch.stack(y_history) + y_history_considered = y_history[-relax_dyn_num:].abs() + + # --------------------------------------------------------- + # Calculate improvement in each dimension + # --------------------------------------------------------- + + # Linear fit to the time series for each radius + x_fit = np.arange(len(y_history_considered)) + n_radii = y_history_considered.shape[1] # Number of radius points + + # Initialize arrays to store results for each radius + change_in_metric = torch.zeros(n_radii) + + # Fit line to each radius dimension separately + for i_radius in range(n_radii): + + # Fit a line that fits all the considered history + y_fit = y_history_considered[:, i_radius].cpu().numpy() + slope, intercept = np.polyfit(x_fit, y_fit, 1) + + # Calculate the relative change in metric + metric0 = slope * x_fit[0] + intercept + metric1 = slope * x_fit[-1] + intercept + + change_in_metric[i_radius] = np.abs((metric1 - metric0) / metric0) + + # --------------------------------------------------------- + # Determine which dimensions will need a reduction in relax + # --------------------------------------------------------- + + mask_reduction = change_in_metric < relax_dyn_tol_rel + + if mask_reduction.any(): + + if (relax[:,mask_reduction] < min_relax).all(): print(f"\t\t\t<> Metric not improving enough (@{it}), relax already at minimum of {min_relax:.1e}, not worth continuing", typeMsg="i") - return relax, False, True + return relax, True + + print(f"\t\t\t<> Metric not improving enough (@{it}), decreasing relax for {mask_reduction.sum()} out of {n_radii} channels") + relax[:,mask_reduction] = relax[:,mask_reduction] / relax_dyn_decrease + print(f"\t\t\t\t- New relax values: from {relax.min():.1e} to {relax.max():.1e}") + + return relax, False else: - return relax, False, False - + print(f"\t\t\t<> Metric improving enough (@{it}), relax remains at {relax.min():.1e} to {relax.max():.1e}") + + return relax, False + def _simple_relax_iteration(x, Q, QT, relax, dx_max, dx_max_abs = None, dx_min_abs = None, threshold_zero_flux_issue=1e-10): # Calculate step in gradient (if target > transport, dx>0 because I want to increase gradients) dx = relax * (QT - Q) / (Q**2 + QT**2).clamp(min=threshold_zero_flux_issue) ** 0.5 diff --git a/src/mitim_tools/opt_tools/scripts/evaluate_model.py b/src/mitim_tools/opt_tools/scripts/evaluate_model.py index da2b417d..331c0804 100644 --- a/src/mitim_tools/opt_tools/scripts/evaluate_model.py +++ b/src/mitim_tools/opt_tools/scripts/evaluate_model.py @@ -2,7 +2,7 @@ import argparse import numpy as np import matplotlib.pyplot as plt -from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools import IOtools, GRAPHICStools from mitim_tools.opt_tools import STRATEGYtools """ @@ -10,7 +10,7 @@ This way, you can try plot, re-ft, find best parameters, etc. It calculates speed, and generates profile file to look at bottlenecks e.g. - evaluate_model.py --folder run1/ --output Qi_tr_turb_5 --input aLti_5 --around -3 + evaluate_model.py --folder run1/ --output Qi_tr_turb_5 --inputs aLti_5 --around -3 evaluate_model.py --folder run1/ --step -1 --output Qi_tr_turb_5 --file figure.eps """ @@ -20,7 +20,7 @@ parser.add_argument("--folder", required=True, type=str) parser.add_argument("--step", type=int, required=False, default=-1) parser.add_argument("--output", required=False, type=str, default="Qi_tr_turb_1") -parser.add_argument("--input", required=False, type=str, default="aLti_1") +parser.add_argument("--inputs", required=False, type=str,nargs='*', default=["aLti_1"]) parser.add_argument("--around", type=int, required=False, default=-1) parser.add_argument("--xrange", type=float, required=False, default=0.5) parser.add_argument("--file", type=str, required=False, default=None) # File to save .eps @@ -31,7 +31,7 @@ folderWork = IOtools.expandPath(args.folder) step_num = args.step output_label = args.output -input_label = args.input +input_labels = args.inputs file = args.file plotYN = args.plot around = args.around @@ -48,12 +48,18 @@ # ***************** Plot +cols = GRAPHICStools.listColors() + if plotYN: gp.plot() if file is not None: plt.savefig(file, transparent=True, dpi=300) - gp.localBehavior_scan(gpA.train_X[around, :], dimension_label=input_label,xrange=xrange) + fig, axs = plt.subplots(nrows=2, figsize=(6, 9)) + for i,input_label in enumerate(input_labels): + gp.localBehavior_scan(gpA.train_X[around, :], dimension_label=input_label,xrange=xrange, axs=axs, c=cols[i], label=input_label) + + axs[0].legend() # gp.plot(plotFundamental=False) # gp.plotTraining() diff --git a/src/mitim_tools/opt_tools/utils/BOgraphics.py b/src/mitim_tools/opt_tools/utils/BOgraphics.py index 98969f6b..d63d5fce 100644 --- a/src/mitim_tools/opt_tools/utils/BOgraphics.py +++ b/src/mitim_tools/opt_tools/utils/BOgraphics.py @@ -763,7 +763,7 @@ def ev(X): def localBehavior_scan_surrogate_model( - self, x, numP=50, dimension_label=None, plotYN=True, axs=None, c="b", xrange=0.5 + self, x, numP=50, dimension_label=None, plotYN=True, axs=None, c="b", xrange=0.5, label='' ): """ This works only for individual models @@ -801,15 +801,17 @@ def ev(X): xlabel = xlabels[x_dim_chosen] ax = axs[0] - ax.plot(Jx, Y, "-o", color=c, lw=1.0) + ax.plot(Jx, Y, "-o", color=c, lw=1.0, label=label, markersize=3) ax.set_xlabel(xlabel) ax.set_ylabel(f"{self.output}") + GRAPHICStools.addDenseAxis(ax) ax = axs[1] - ax.plot(Jx, J, "-o", color=c, lw=1.0) + ax.plot(Jx, J, "-o", color=c, lw=1.0, label=label, markersize=3) ax.set_xlabel(xlabel) ax.set_ylabel(f"d({self.output})/d({xlabel})") ax.set_title("Scan of local Jacobian") + GRAPHICStools.addDenseAxis(ax) # ---------------------------------------------------------------------------------------------------- diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index ffa35eb6..a5ba5b1a 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1422,18 +1422,32 @@ def lumpDT(self): self.moveSpecie(pos=len(self.Species), pos_new=1) - def changeZeff(self, Zeff, ion_pos=2, quasineutral_ions=None, enforceSameGradients=False): - """ - if (D,Z1,Z2), pos 1 -> change Z1 - """ + def changeZeff( + self, + Zeff, + ion_pos = 2, # Position of ion to change (if (D,Z1,Z2), pos 1 -> change Z1) + keep_fmain = False, # If True, it will keep fmain and change Z of ion in position ion_pos. If False, it will change the content of ion in position ion_pos and the content of quasineutral ions to achieve Zeff + fmain_force = None, # If keep_fmain is True, it will force fmain to this value. If None, it will use the current fmain + enforceSameGradients = False # If True, it will scale all thermal densities to have the same gradients after changing Zeff + ): - if quasineutral_ions is None: - if self.DTplasmaBool: - quasineutral_ions = [self.Dion, self.Tion] - else: - quasineutral_ions = [self.Mion] + if not keep_fmain and fmain_force is not None: + raise ValueError("[MITIM] fmain_force can only be used if keep_fmain is True") + + if fmain_force is not None: + fmain_factor = fmain_force / self.derived["fmain"] + else: + fmain_factor = 1.0 - print(f'\t\t- Changing Zeff (from {self.derived["Zeff_vol"]:.3f} to {Zeff=:.3f}) by changing content of ion in position {ion_pos} {self.Species[ion_pos]["N"],self.Species[ion_pos]["Z"]}, quasineutralized by ions {quasineutral_ions}',typeMsg="i",) + if self.DTplasmaBool: + quasineutral_ions = [self.Dion, self.Tion] + else: + quasineutral_ions = [self.Mion] + + if not keep_fmain: + print(f'\t\t- Changing Zeff (from {self.derived["Zeff_vol"]:.3f} to {Zeff=:.3f}) by changing content of ion in position {ion_pos} {self.Species[ion_pos]["N"],self.Species[ion_pos]["Z"]}, quasineutralized by ions {quasineutral_ions}',typeMsg="i") + else: + print(f'\t\t- Changing Zeff (from {self.derived["Zeff_vol"]:.3f} to {Zeff=:.3f}) by changing content and Z of ion in position {ion_pos} {self.Species[ion_pos]["N"],self.Species[ion_pos]["Z"]}, quasineutralized by ions {quasineutral_ions} and keeping fmain={self.derived["fmain"]*fmain_factor:.3f}',typeMsg="i") # Plasma needs to be in quasineutrality to start with self.enforceQuasineutrality() @@ -1443,47 +1457,82 @@ def changeZeff(self, Zeff, ion_pos=2, quasineutral_ions=None, enforceSameGradien # ------------------------------------------------------ Zq = np.zeros(self.derived["fi"].shape[0]) Zq2 = np.zeros(self.derived["fi"].shape[0]) + fZq = np.zeros(self.derived["fi"].shape[0]) + fZq2 = np.zeros(self.derived["fi"].shape[0]) fZj = np.zeros(self.derived["fi"].shape[0]) fZj2 = np.zeros(self.derived["fi"].shape[0]) for i in range(len(self.Species)): + + # Ions for quasineutrality (main ones) if i in quasineutral_ions: Zq += self.Species[i]["Z"] Zq2 += self.Species[i]["Z"] ** 2 + + fZq += self.Species[i]["Z"] * self.derived["fi"][:, i] * fmain_factor + fZq2 += self.Species[i]["Z"] ** 2 * self.derived["fi"][:, i] * fmain_factor + # Non-quasineutral and not the ion to change elif i != ion_pos: fZj += self.Species[i]["Z"] * self.derived["fi"][:, i] fZj2 += self.Species[i]["Z"] ** 2 * self.derived["fi"][:, i] + # Ion to change else: Zk = self.Species[i]["Z"] - # ------------------------------------------------------ - # Find free parameters (fk and fq) - # ------------------------------------------------------ + fi_orig = self.derived["fi"][:, ion_pos] + Zi_orig = self.Species[ion_pos]["Z"] + Ai_orig = self.Species[ion_pos]["A"] - fk = ( Zeff - (1-fZj)*Zq2/Zq - fZj2 ) / ( Zk**2 - Zk*Zq2/Zq) - fq = ( 1 - fZj - fk*Zk ) / Zq + if not keep_fmain: + # ------------------------------------------------------ + # Find free parameters (fk and fq) + # ------------------------------------------------------ - if (fq<0).any(): - raise ValueError(f"Zeff cannot be reduced by changing ion #{ion_pos} because it would require negative densities for quasineutral ions") + fk = ( Zeff - (1-fZj)*Zq2/Zq - fZj2 ) / ( Zk**2 - Zk*Zq2/Zq) + fq = ( 1 - fZj - fk*Zk ) / Zq - # ------------------------------------------------------ - # Insert - # ------------------------------------------------------ + if (fq<0).any(): + raise ValueError(f"Zeff cannot be reduced by changing ion #{ion_pos} because it would require negative densities for quasineutral ions") - fi_orig = self.derived["fi"][:, ion_pos] + # ------------------------------------------------------ + # Insert + # ------------------------------------------------------ + + self.profiles["ni(10^19/m^3)"][:, ion_pos] = fk * self.profiles["ne(10^19/m^3)"] + for i in quasineutral_ions: + self.profiles["ni(10^19/m^3)"][:, i] = fq * self.profiles["ne(10^19/m^3)"] + else: + # ------------------------------------------------------ + # Find free parameters (fk and Zk) + # ------------------------------------------------------ - self.profiles["ni(10^19/m^3)"][:, ion_pos] = fk * self.profiles["ne(10^19/m^3)"] - for i in quasineutral_ions: - self.profiles["ni(10^19/m^3)"][:, i] = fq * self.profiles["ne(10^19/m^3)"] + Zk = (Zeff - fZq2 - fZj2) / (1 - fZq - fZj) + + # I need a single value + Zk_ave = CALCtools.integrateFS(Zk, self.profiles["rmin(m)"], self.derived["volp_miller"])[-1] / self.derived["volume"] + + fk = (1 - fZq - fZj) / Zk_ave + + # ------------------------------------------------------ + # Insert + # ------------------------------------------------------ + + self.profiles['z'][ion_pos] = Zk_ave + self.profiles['mass'][ion_pos] = Zk_ave * 2 + self.profiles["ni(10^19/m^3)"][:, ion_pos] = fk * self.profiles["ne(10^19/m^3)"] + + if fmain_force is not None: + for i in quasineutral_ions: + self.profiles["ni(10^19/m^3)"][:, i] *= fmain_factor self.readSpecies() - self.derive_quantities(rederiveGeometry=False) + self.deriveQuantities(rederiveGeometry=False) if enforceSameGradients: self.scaleAllThermalDensities() - self.derive_quantities(rederiveGeometry=False) + self.deriveQuantities(rederiveGeometry=False) - print(f'\t\t\t* Dilution changed from {fi_orig.mean():.2e} (vol avg) to { self.derived["fi"][:, ion_pos].mean():.2e} to achieve Zeff={self.derived["Zeff_vol"]:.3f} (fDT={self.derived["fmain"]:.3f}) [quasineutrality error = {self.derived["QN_Error"]:.1e}]') + print(f'\t\t\t* Dilution changed from {fi_orig.mean():.2e} (vol avg) of ion [{Zi_orig:.2f},{Ai_orig:.2f}] to { self.derived["fi"][:, ion_pos].mean():.2e} of ion [{self.profiles["z"][ion_pos]:.2f}, {self.profiles["mass"][ion_pos]:.2f}] to achieve Zeff={self.derived["Zeff_vol"]:.3f} (fDT={self.derived["fmain"]:.3f}) [quasineutrality error = {self.derived["QN_Error"]:.1e}]') def moveSpecie(self, pos=2, pos_new=1): """ diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py index 277247bb..6c9a97de 100644 --- a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -504,6 +504,7 @@ def plot_gradients( self, axs4, color="b", + fast_color='r', lw=1.0, label="", ls="-o", @@ -590,6 +591,19 @@ def plot_gradients( markersize=ms, alpha=alpha, ) + for i in range(len(self.Species)): + if self.Species[i]["S"] != "therm": + ax.plot( + xcoord[:ix], + self.derived["aLTi"][:ix, i], + ls, + c=fast_color, + lw=lw, + markersize=ms, + alpha=alpha, + label=self.Species[i]["N"], + ) + ax.legend(loc="best", fontsize=7) ax = axs4[5] ax.plot( xcoord[:ix], diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index 0502824b..cbf3421f 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -15,7 +15,7 @@ GRAPHICStools, ) from mitim_tools.transp_tools import UFILEStools -from mitim_tools.gacode_tools import TGLFtools, TGYROtools +from mitim_tools.gacode_tools import TGLFtools, TGYROtools, PROFILEStools from mitim_tools.gacode_tools.utils import GACODEplotting, GACODErun, TRANSPgacode from mitim_tools.transp_tools.utils import ( FBMtools, @@ -7682,12 +7682,14 @@ def plotICRF_t(self, fig=None): if fig is None: fig = plt.figure() - grid = plt.GridSpec(2, 2, hspace=0.3, wspace=0.2) + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.4) ax1 = fig.add_subplot(grid[0, 0]) ax2 = fig.add_subplot(grid[0, 1], sharex=ax1) ax3 = fig.add_subplot(grid[1, 0], sharex=ax1) ax4 = fig.add_subplot(grid[1, 1], sharex=ax1) + ax5 = fig.add_subplot(grid[0, 2], sharex=ax1) + ax6 = fig.add_subplot(grid[1, 2], sharex=ax1) # ELECTRONS ax = ax1 @@ -7774,10 +7776,10 @@ def plotICRF_t(self, fig=None): # ax.plot(self.x_lw,tot,lw=3,label='$P_{ICH}$') ax.plot(self.t, self.PichT, c="r", lw=3, label="$P_{ICH}$") - ax.plot(self.t, self.PichT_min, lw=2, label="$P_{ICH->min}$") ax.plot(self.t, self.PiichT_dir, lw=2, label="$P_{ICH->i}$") ax.plot(self.t, self.PeichT_dir, lw=2, label="$P_{ICH->e}$") ax.plot(self.t, self.PfichT_dir, lw=2, label="$P_{ICH->fast}$") + ax.plot(self.t, self.PichT_min, lw=2, label="$P_{ICH->min}$") ax.plot(self.t, self.PichT_check, lw=2, ls="--", c="y", label="check (sum)") ax.set_title("Total Balance") @@ -7788,6 +7790,26 @@ def plotICRF_t(self, fig=None): GRAPHICStools.addLegendApart(ax, ratio=0.85, size=self.mainLegendSize) GRAPHICStools.addDenseAxis(ax) + # TOTAL + ax = ax6 + + ax.plot(self.t, self.PichT, c="r", lw=3, label="$P_{ICH}$") + ax.plot(self.t, self.PiichT, lw=2, label="$P_{ICH,i}$") + ax.plot(self.t, self.PeichT, lw=2, label="$P_{ICH,e}$") + ax.plot(self.t, self.PfichT_dir, lw=2, label="$P_{ICH,fast}$") + ax.plot(self.t, self.GainminT, lw=2, label="$dW_{min}/dt$") + P = self.PeichT + self.PiichT + self.PfichT_dir + self.GainminT + ax.plot(self.t, P, lw=2, ls="--", c="y", label="check (sum)") + + ax.set_title("Total Balance (after thermalization)") + ax.set_ylabel("Power (MW)") + ax.set_xlabel("Time (s)") + ax.set_ylim(bottom=0) + + GRAPHICStools.addLegendApart(ax, ratio=0.85, size=self.mainLegendSize) + GRAPHICStools.addDenseAxis(ax) + + def plotRelevantResonances(self, ax, Fich, time=None, legendYN=False, lw=3): if time is None: i1 = self.ind_saw @@ -8195,13 +8217,11 @@ def plotSeparateSystems(self, fig=None): if np.sum(self.PichT) > 1.0e-5: for i in range(len(self.PichT_ant)): ax.plot(self.t, self.PichT_ant[i], lw=2, label=f"{i + 1}") - ax.plot( - self.t, self.PeichT + self.PiichT, c="y", ls="--", label="to species" - ) + ax.plot(self.t, self.PeichT + self.PiichT + self.PfichT_dir + self.GainminT, c="y", ls="--", label="to species (e+i+f+dWmin/dt)") timeb = 0.25 it1 = np.argmin(np.abs(self.t - (self.t[-1] - timeb))) - mean = np.mean(self.PeichT[it1:] + self.PiichT[it1:]) + mean = np.mean(self.PeichT[it1:] + self.PiichT[it1:] + self.PfichT_dir[it1:] + self.GainminT[it1:]) ax.axhline( y=mean, alpha=0.5, @@ -8213,7 +8233,7 @@ def plotSeparateSystems(self, fig=None): ax.plot( self.t, - self.PeichT + self.PiichT + self.PichTOrbLoss, + self.PeichT + self.PiichT + self.PfichT_dir + self.GainminT + self.PichTOrbLoss, c="c", ls="--", label="+ orb losses", @@ -8462,7 +8482,7 @@ def plotHeating(self, fig=None): if fig is None: fig = plt.figure() - grid = plt.GridSpec(nrows=2, ncols=4, hspace=0.3, wspace=0.2) + grid = plt.GridSpec(nrows=2, ncols=4, hspace=0.3, wspace=0.4) ax1 = fig.add_subplot(grid[0, 0]) ax2 = fig.add_subplot(grid[0, 1]) @@ -8478,9 +8498,7 @@ def plotHeating(self, fig=None): ax1.plot(self.t, self.PichT, "r", ls="-", lw=2, label="$P_{ICH}$") ax1.plot(self.t, self.PeichT, "b", ls="-", lw=1, label="$P_{ICH,e}$") ax1.plot(self.t, self.PiichT, "g", ls="-", lw=1, label="$P_{ICH,i}$") - ax1.plot( - self.t, self.PeichT + self.PiichT, "y", ls="--", lw=1, label="$P_{ICH,e+i}$" - ) + ax1.plot(self.t, self.PeichT + self.PiichT, "y", ls="--", lw=1, label="$P_{ICH,e+i}$") ax1.plot(self.t, self.PichT_min, "r", ls="--", lw=1, label="$P_{min}$") ax1.plot(self.t, self.PeichT_dir, "r", ls="-.", lw=1, label="$P_{dir,e}$") @@ -15425,8 +15443,7 @@ def grid_interpolation_method_to_zero(x,y): for key in ['ne(10^19/m^3)', 'ni(10^19/m^3)', 'te(keV)', 'ti(keV)', 'rmin(m)']: profiles[key] = profiles[key].clip(min=minimum) - from mitim_tools.gacode_tools import PROFILEStools - p = PROFILEStools.gacode_state.scratch(profiles) + p = PROFILEStools.PROFILES_GACODE.scratch(profiles) return p diff --git a/templates/input.cgyro.controls b/templates/input.cgyro.controls index e8e84341..e2d85fa3 100644 --- a/templates/input.cgyro.controls +++ b/templates/input.cgyro.controls @@ -31,6 +31,9 @@ N_FIELD=3 #Data field output flag MOMENT_PRINT_FLAG=1 +#Make EM fields available as output +FIELD_PRINT_FLAG=1 + #Velocity Order #VELOCITY_ORDER=2 diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index 8ec4be0e..15ff7a8e 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -5,7 +5,7 @@ cold_start = True gacode_file = __mitimroot__ / "tests" / "data" / "input.gacode" -folder = __mitimroot__ / "tests" / "scratch" / "cgyro_test2" +folder = __mitimroot__ / "tests" / "scratch" / "cgyro_test" if cold_start and folder.exists(): os.system(f"rm -r {folder}") @@ -24,7 +24,7 @@ 'KY':0.3, 'MAX_TIME': 1E1, # Short, I just want to test the run }, - submit_via_qsub=True # NERSC: True #TODO change this + submit_via_qsub=False # NERSC: True #TODO change this ) cgyro.check(every_n_minutes=1) @@ -33,5 +33,5 @@ cgyro.read(label="cgyro1") -cgyro.plotLS() +cgyro.plot(labels=["cgyro1"]) cgyro.fn.show() diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py index be03ff77..93d39a4b 100644 --- a/tests/EPED_workflow.py +++ b/tests/EPED_workflow.py @@ -1,7 +1,13 @@ +import os from mitim_tools.eped_tools import EPEDtools from mitim_tools import __mitimroot__ -folder = __mitimroot__ / "tests" / "scratch" / "eped_test" +cold_start = True + +folder = __mitimroot__ / "tests" / "scratch" / "eped_test16" + +if cold_start and os.path.exists(folder): + os.system(f"rm -r {folder}") eped = EPEDtools.EPED(folder=folder) From 3bb6f843a1174c692b8c993636190fe6ad5059a5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 4 Aug 2025 16:02:40 -0400 Subject: [PATCH 123/385] vmecpp --- pyproject.toml | 1 + src/mitim_tools/misc_tools/IOtools.py | 1 + 2 files changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index d5ee2da7..e82ab84b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dependencies = [ "onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models "tensorflow", "f90nml", + "vmecpp" ] [project.optional-dependencies] diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index 08c63c65..4b6a01f6 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -2,6 +2,7 @@ import shutil import psutil import copy +from typing import Callable import dill as pickle_dill import pandas as pd from mitim_tools.misc_tools import GRAPHICStools From 5b597e9ee87406498ac072bae47ff297f68c6426 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 4 Aug 2025 16:23:57 -0400 Subject: [PATCH 124/385] Added quends --- pyproject.toml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e82ab84b..08dc56d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,8 @@ dependencies = [ "onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models "tensorflow", "f90nml", - "vmecpp" + "vmecpp", + "quends @ git+https://github.com/sandialabs/quends.git" ] [project.optional-dependencies] @@ -56,7 +57,7 @@ omfit = [ "omfit_classes>3.2024.19.2", # Otherwise, it will need an old version of matplotlib, matplotlib<3.6 "scipy<1.14.0", # As of 08/08/2024, because of https://github.com/gafusion/OMFIT-source/issues/7104 "numpy<2.0.0", # For the xarray requirement below to work - # "xarray==2022.6.0", # As of 08/08/2024, because https://github.com/gafusion/OMFIT-source/issues/7104. Froze this version because of the PLASMAstate xr reader (importlib_metadata issues) + "xarray==2022.6.0", # As of 08/08/2024, because https://github.com/gafusion/OMFIT-source/issues/7104. Froze this version because of the PLASMAstate xr reader (importlib_metadata issues) "omas", "fortranformat", "openpyxl", From 39d1a8bb880fc061beb69585ea78d78ab914bf15 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 4 Aug 2025 16:50:00 -0400 Subject: [PATCH 125/385] correct_zeff and avoid zero division warnings --- .../plasmastate_tools/MITIMstate.py | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index a5ba5b1a..da60d1ec 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -163,7 +163,10 @@ def derive_quantities_base(self, mi_ref=None, derive_quantities=True, rederiveGe # ------------------------------------------------------------------------------------------------ if derive_quantities: - self.derive_quantities_full(rederiveGeometry=rederiveGeometry) + + # Avoid division by zero warning by using np.errstate + with np.errstate(divide='ignore', invalid='ignore'): + self.derive_quantities_full(rederiveGeometry=rederiveGeometry) def write_state(self, file=None): print("\t- Writting input.gacode file") @@ -294,8 +297,12 @@ def readSpecies(self, maxSpecies=100, correct_zeff = True): self.Species = Species # Correct Zeff if needed - if correct_zeff and ("z_eff(-)" in self.profiles): - self.profiles["z_eff(-)"] = np.sum(self.profiles["ni(10^19/m^3)"] * self.profiles["z"] ** 2, axis=1) / self.profiles["ne(10^19/m^3)"] + if correct_zeff: + self.correct_zeff_array() + + def correct_zeff_array(self): + + self.profiles["z_eff(-)"] = np.sum(self.profiles["ni(10^19/m^3)"] * self.profiles["z"] ** 2, axis=1) / self.profiles["ne(10^19/m^3)"] def sumFast(self): self.nFast = self.profiles["ne(10^19/m^3)"] * 0.0 @@ -1405,9 +1412,7 @@ def lumpSpecies( self.remove(ions_list) # Contributions to dilution and to Zeff - print( - f'\t\t\t* New plasma has Zeff_vol={self.derived["Zeff_vol"]:.2f}, QN error={self.derived["QN_Error"]:.4f}' - ) + print(f'\t\t\t* New plasma has Zeff_vol={self.derived["Zeff_vol"]:.2f}, QN error={self.derived["QN_Error"]:.4f}') def lumpImpurities(self): @@ -1508,7 +1513,7 @@ def changeZeff( Zk = (Zeff - fZq2 - fZj2) / (1 - fZq - fZj) # I need a single value - Zk_ave = CALCtools.integrateFS(Zk, self.profiles["rmin(m)"], self.derived["volp_miller"])[-1] / self.derived["volume"] + Zk_ave = CALCtools.volume_integration(Zk, self.profiles["rmin(m)"], self.derived["volp_geo"])[-1] / self.derived["volume"] fk = (1 - fZq - fZj) / Zk_ave @@ -1526,11 +1531,11 @@ def changeZeff( self.readSpecies() - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) if enforceSameGradients: self.scaleAllThermalDensities() - self.deriveQuantities(rederiveGeometry=False) + self.derive_quantities(rederiveGeometry=False) print(f'\t\t\t* Dilution changed from {fi_orig.mean():.2e} (vol avg) of ion [{Zi_orig:.2f},{Ai_orig:.2f}] to { self.derived["fi"][:, ion_pos].mean():.2e} of ion [{self.profiles["z"][ion_pos]:.2f}, {self.profiles["mass"][ion_pos]:.2f}] to achieve Zeff={self.derived["Zeff_vol"]:.3f} (fDT={self.derived["fmain"]:.3f}) [quasineutrality error = {self.derived["QN_Error"]:.1e}]') From d919874a4747e83630defbcebc682b79d17b69fd Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 4 Aug 2025 17:22:03 -0400 Subject: [PATCH 126/385] Removal of deprecated neo error in CGYRO --- .../powertorch/physics_models/transport_cgyro.py | 4 ---- .../powertorch/physics_models/transport_tgyro.py | 1 - 2 files changed, 5 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 6ce070e1..cb43d248 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -168,7 +168,6 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod minErrorPercent = PORTALSparameters["percentError_stable"] Qi_criterion_stable = PORTALSparameters["Qi_criterion_stable"] - percentNeo = PORTALSparameters["percentError"][1] try: impurityPosition = PROFILEStools.impurity_location(PROFILEStools.gacode_state(unmodified_profiles), PORTALSparameters["ImpurityOfInterest"]) @@ -188,7 +187,6 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod numPORTALS, minErrorPercent, Qi_criterion_stable, - percentNeo, radii, OriginalFimp=OriginalFimp, evaluationsInFile=f"{numPORTALS_this}", @@ -219,7 +217,6 @@ def cgyroing( evaluations, minErrorPercent, Qi_criterion_stable, - percentNeo, radii, OriginalFimp=1.0, file=None, @@ -265,7 +262,6 @@ def cgyroing( FolderEvaluation, minErrorPercent=minErrorPercent, Qi_criterion_stable=Qi_criterion_stable, - percent_tr_neo=percentNeo, impurityPosition=impurityPosition, OriginalFimp=OriginalFimp, ) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 34d1ed25..482e362d 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -577,7 +577,6 @@ def modifyResults( tgyro, folder_tgyro, minErrorPercent=5.0, - percent_tr_neoc=2.0, Qi_criterion_stable=0.0025, impurityPosition=3, OriginalFimp=1.0, From eb50f5c0163109122d29066e53d8361255f947eb Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 4 Aug 2025 18:57:12 -0400 Subject: [PATCH 127/385] Allow reading VMEC state without profiles --- src/mitim_tools/plasmastate_tools/MITIMstate.py | 4 ++-- src/mitim_tools/plasmastate_tools/utils/VMECtools.py | 8 ++++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index da60d1ec..742b8628 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -75,12 +75,12 @@ def ensure_variables_existence(self): # --------------------------------------------------------------------------- # Choose a template for dimensionality - template_key_1d, template_key_2d = "rho(-)", "ti(keV)" + template_key_1d = "rho(-)" # Ensure required keys exist for key, dim in required_profiles.items(): if key not in self.profiles: - self.profiles[key] = copy.deepcopy(self.profiles[template_key_1d]) * 0.0 if dim == 1 else copy.deepcopy(self.profiles[template_key_2d]) * 0.0 + self.profiles[key] = copy.deepcopy(self.profiles[template_key_1d]) * 0.0 if dim == 1 else np.atleast_2d(copy.deepcopy(self.profiles[template_key_1d]) * 0.0).T ''' diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index eaf5f7e7..6a7ac4f1 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -21,7 +21,7 @@ class vmec_state(MITIMstate.mitim_state): def __init__( self, file_vmec, - file_profs, + file_profs=None, derive_quantities=True, mi_ref=None ): @@ -37,7 +37,7 @@ def __init__( # Read the input file and store the raw data self.files = [file_vmec, file_profs] - if self.files is not None: + if self.files[0] is not None: self._read_vmec() # Derive (Depending on resolution, derived can be expensive, so I may not do it every time) @@ -73,6 +73,10 @@ def _read_vmec(self): #self.profiles["polflux(Wb/radian)"] = self.wout.chi # Read Profiles + if self.files[1] is None: + print("\t- No profiles file provided, skipping profile reading") + return + data = self._read_profiles(x_coord=self.profiles["rho(-)"]) self.profiles['te(keV)'] = data['Te'] From 983532aa1835ad9dbcd7539b7e16a2e4e85d4ac8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 5 Aug 2025 11:09:54 -0400 Subject: [PATCH 128/385] Options renaming for clarity --- .../maestro/utils/PORTALSbeat.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 14 ++-- src/mitim_modules/portals/PORTALStools.py | 2 +- .../portals/utils/PORTALSinit.py | 36 ++++----- .../portals/utils/PORTALSoptimization.py | 16 ++-- src/mitim_modules/powertorch/STATEtools.py | 74 +++++++++---------- .../physics_models/targets_analytic.py | 4 +- .../physics_models/transport_analytic.py | 6 +- .../physics_models/transport_cgyro.py | 4 +- .../physics_models/transport_tglf.py | 8 +- .../physics_models/transport_tgyro.py | 10 +-- .../powertorch/scripts/calculateTargets.py | 12 +-- .../scripts/compareRadialResolution.py | 4 +- .../powertorch/scripts/compareWithTGYRO.py | 4 +- .../powertorch/utils/TARGETStools.py | 8 +- .../powertorch/utils/TRANSFORMtools.py | 16 ++-- .../powertorch/utils/TRANSPORTtools.py | 4 +- src/mitim_tools/opt_tools/STRATEGYtools.py | 8 +- src/mitim_tools/popcon_tools/RAPIDStools.py | 2 +- tests/POWERTORCH_workflow.py | 4 +- 20 files changed, 116 insertions(+), 122 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 4491c732..05787210 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -126,7 +126,7 @@ def _flux_match_for_first_point(self): _ = PORTALSoptimization.flux_match_surrogate( portals.step, p, - TargetOptions_use = self.mitim_bo.optimization_object.powerstate.TargetOptions, # Use the TargetOptions of the new run, not the old one (which may be with fixed targets if soft) + target_options_use = self.mitim_bo.optimization_object.powerstate.target_options, # Use the target_options of the new run, not the old one (which may be with fixed targets if soft) file_write_csv=folder_fm / 'optimization_data.csv' ) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 54561fec..58119fdb 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -82,7 +82,7 @@ def __init__( self, folder, # Folder where the PORTALS workflow will be run namelist=None, # If None, default namelist will be used. If not None, it will be read and used - tensor_opts = { + tensor_options = { "dtype": torch.double, "device": torch.device("cpu"), }, @@ -104,7 +104,7 @@ def __init__( super().__init__( folder, namelist=namelist, - tensor_opts=tensor_opts, + tensor_options=tensor_options, default_namelist_function=( partial(default_namelist, CGYROrun=CGYROrun) if (namelist is None) @@ -308,7 +308,7 @@ def prep( limitsAreRelative=limitsAreRelative, cold_start=cold_start, hardGradientLimits=hardGradientLimits, - tensor_opts = self.tensor_opts, + tensor_options = self.tensor_options, seedInitial=seedInitial, checkForSpecies=askQuestions, ModelOptions=ModelOptions, @@ -551,10 +551,10 @@ def reuseTrainingTabular( self_copy = copy.deepcopy(self) if reevaluateTargets == 1: - self_copy.powerstate.TransportOptions["transport_evaluator"] = None - self_copy.powerstate.TargetOptions["ModelOptions"]["TypeTarget"] = "powerstate" + self_copy.powerstate.transport_options["transport_evaluator"] = None + self_copy.powerstate.target_options["ModelOptions"]["TypeTarget"] = "powerstate" else: - self_copy.powerstate.TransportOptions["transport_evaluator"] = transport_tgyro.tgyro_model + self_copy.powerstate.transport_options["transport_evaluator"] = transport_tgyro.tgyro_model _, dictOFs = runModelEvaluator( self_copy, @@ -623,7 +623,7 @@ def runModelEvaluator( # --------------------------------------------------------------------------------------------------- # In certain cases, I want to cold_start the model directly from the PORTALS call instead of powerstate - powerstate.TransportOptions["ModelOptions"]["cold_start"] = cold_start + powerstate.transport_options["ModelOptions"]["cold_start"] = cold_start # Evaluate X (DVs) through powerstate.calculate(). This will populate .plasma with the results powerstate.calculate(X, nameRun=name, folder=folder_model, evaluation_number=numPORTALS) diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index bf32e92c..a5ac3ce2 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -303,7 +303,7 @@ def constructEvaluationProfiles(X, surrogate_parameters, recalculateTargets=Fals # Targets only if needed (for speed, GB doesn't need it) if recalculateTargets: - powerstate.TargetOptions["ModelOptions"]["targets_evaluator_method"] = "powerstate" # For surrogate evaluation, always powerstate, logically. + powerstate.target_options["ModelOptions"]["targets_evaluator_method"] = "powerstate" # For surrogate evaluation, always powerstate, logically. powerstate.calculateTargets() return powerstate diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index b263db21..c819f251 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -30,7 +30,7 @@ def initializeProblem( ModelOptions=None, seedInitial=None, checkForSpecies=True, - tensor_opts = { + tensor_options = { "dtype": torch.double, "device": torch.device("cpu"), } @@ -43,7 +43,7 @@ def initializeProblem( - define_ranges_from_profiles must be PROFILES class """ - dfT = torch.randn((2, 2), **tensor_opts) + dfT = torch.randn((2, 2), **tensor_options) if seedInitial is not None: torch.manual_seed(seed=seedInitial) @@ -128,24 +128,18 @@ def initializeProblem( if ModelOptions is None: ModelOptions = { - "cold_start": False, - "launchMODELviaSlurm": portals_fun.PORTALSparameters[ - "launchEvaluationsAsSlurmJobs" - ], - "MODELparameters": portals_fun.MODELparameters, + "launchMODELviaSlurm": portals_fun.PORTALSparameters["launchEvaluationsAsSlurmJobs"], "includeFastInQi": portals_fun.PORTALSparameters["includeFastInQi"], "TurbulentExchange": portals_fun.PORTALSparameters["surrogateForTurbExch"], - "profiles_postprocessing_fun": portals_fun.PORTALSparameters[ - "profiles_postprocessing_fun" - ], - "impurityPosition": position_of_impurity, + "profiles_postprocessing_fun": portals_fun.PORTALSparameters["profiles_postprocessing_fun"], "UseFineGridTargets": portals_fun.PORTALSparameters["fineTargetsResolution"], "OriginalFimp": portals_fun.PORTALSparameters["fImp_orig"], - "forceZeroParticleFlux": portals_fun.PORTALSparameters[ - "forceZeroParticleFlux" - ], + "forceZeroParticleFlux": portals_fun.PORTALSparameters["forceZeroParticleFlux"], "percentError": portals_fun.PORTALSparameters["percentError"], "use_tglf_scan_trick": portals_fun.PORTALSparameters["use_tglf_scan_trick"], + "cold_start": False, + "MODELparameters": portals_fun.MODELparameters, + "impurityPosition": position_of_impurity, } if "extra_params" not in ModelOptions: @@ -163,23 +157,23 @@ def initializeProblem( portals_fun.powerstate = STATEtools.powerstate( profiles, - EvolutionOptions={ + evolution_options={ "ProfilePredicted": portals_fun.MODELparameters["ProfilesPredicted"], "rhoPredicted": xCPs, "impurityPosition": position_of_impurity, "fineTargetsResolution": portals_fun.PORTALSparameters["fineTargetsResolution"], }, - TransportOptions={ + transport_options={ "transport_evaluator": portals_fun.PORTALSparameters["transport_evaluator"], "ModelOptions": ModelOptions, }, - TargetOptions={ + target_options={ "targets_evaluator": portals_fun.PORTALSparameters["targets_evaluator"], "ModelOptions": { "TypeTarget": portals_fun.MODELparameters["Physics_options"]["TypeTarget"], "targets_evaluator_method": portals_fun.PORTALSparameters["targets_evaluator_method"]}, }, - tensor_opts = tensor_opts + tensor_options = tensor_options ) # After resolution and corrections, store. @@ -218,19 +212,19 @@ def initializeProblem( if define_ranges_from_profiles is not None: # If I want to define ranges from a different profile powerstate_extra = STATEtools.powerstate( define_ranges_from_profiles, - EvolutionOptions={ + evolution_options={ "ProfilePredicted": portals_fun.MODELparameters["ProfilesPredicted"], "rhoPredicted": xCPs, "impurityPosition": position_of_impurity, "fineTargetsResolution": portals_fun.PORTALSparameters["fineTargetsResolution"], }, - TargetOptions={ + target_options={ "targets_evaluator": portals_fun.PORTALSparameters["targets_evaluator"], "ModelOptions": { "TypeTarget": portals_fun.MODELparameters["Physics_options"]["TypeTarget"], "targets_evaluator_method": portals_fun.PORTALSparameters["targets_evaluator_method"]}, }, - tensor_opts = tensor_opts + tensor_options = tensor_options ) dictCPs_base_extra = {} diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index e9168ee1..a354f6d4 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -101,7 +101,7 @@ def flux_match_surrogate( algorithm = None, solver_options = None, keep_within_bounds = True, - TargetOptions_use = None, + target_options_use = None, ): ''' Technique to reutilize flux surrogates to predict new conditions @@ -136,24 +136,24 @@ def flux_match_surrogate( # Create powerstate with new profiles # ---------------------------------------------------- - TransportOptions = copy.deepcopy(step.surrogate_parameters["powerstate"].TransportOptions) + transport_options = copy.deepcopy(step.surrogate_parameters["powerstate"].transport_options) # Define transport calculation function as a surrogate model - TransportOptions['transport_evaluator'] = transport_analytic.surrogate - TransportOptions['ModelOptions'] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} + transport_options['transport_evaluator'] = transport_analytic.surrogate + transport_options['ModelOptions'] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} # Create powerstate with the same options as the original portals but with the new profiles powerstate = STATEtools.powerstate( profiles, - EvolutionOptions={ + evolution_options={ "ProfilePredicted": step.surrogate_parameters["powerstate"].ProfilesPredicted, "rhoPredicted": step.surrogate_parameters["powerstate"].plasma["rho"][0,1:], "impurityPosition": step.surrogate_parameters["powerstate"].impurityPosition, "fineTargetsResolution": step.surrogate_parameters["powerstate"].fineTargetsResolution, }, - TransportOptions=TransportOptions, - TargetOptions= step.surrogate_parameters["powerstate"].TargetOptions if TargetOptions_use is None else TargetOptions_use, - tensor_opts = { + transport_options=transport_options, + target_options= step.surrogate_parameters["powerstate"].target_options if target_options_use is None else target_options_use, + tensor_options = { "dtype": step.surrogate_parameters["powerstate"].dfT.dtype, "device": step.surrogate_parameters["powerstate"].dfT.device}, ) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index d742839b..9848a273 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -26,40 +26,40 @@ def __init__( self, profiles_object, increase_profile_resol=True, - EvolutionOptions=None, - TransportOptions=None, - TargetOptions=None, - tensor_opts=None, + evolution_options=None, + transport_options=None, + target_options=None, + tensor_options=None, ): ''' Inputs: - profiles_object: Object for gacode_state or others - - EvolutionOptions: + - evolution_options: - rhoPredicted: radial grid (MUST NOT CONTAIN ZERO, it will be added internally) - ProfilesPredicted: list of profiles to predict - impurityPosition: int = position of the impurity in the ions set - fineTargetsResolution: int = resolution of the fine targets - - TransportOptions: dictionary with transport_evaluator and ModelOptions - - TargetOptions: dictionary with targets_evaluator and ModelOptions + - transport_options: dictionary with transport_evaluator and ModelOptions + - target_options: dictionary with targets_evaluator and ModelOptions ''' - if EvolutionOptions is None: - EvolutionOptions = {} - if TransportOptions is None: - TransportOptions = { + if evolution_options is None: + evolution_options = {} + if transport_options is None: + transport_options = { "transport_evaluator": None, "ModelOptions": {} } - if TargetOptions is None: - TargetOptions = { + if target_options is None: + target_options = { "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": 3, "targets_evaluator_method": "powerstate" }, } - if tensor_opts is None: - tensor_opts = { + if tensor_options is None: + tensor_options = { "dtype": torch.double, "device": torch.device("cpu"), } @@ -70,16 +70,16 @@ def __init__( print('>> Creating powerstate object...') - self.TransportOptions = TransportOptions - self.TargetOptions = TargetOptions + self.transport_options = transport_options + self.target_options = target_options # Default options - self.ProfilesPredicted = EvolutionOptions.get("ProfilePredicted", ["te", "ti", "ne"]) - self.impurityPosition = EvolutionOptions.get("impurityPosition", 1) + self.ProfilesPredicted = evolution_options.get("ProfilePredicted", ["te", "ti", "ne"]) + self.impurityPosition = evolution_options.get("impurityPosition", 1) self.impurityPosition_transport = copy.deepcopy(self.impurityPosition) - self.fineTargetsResolution = EvolutionOptions.get("fineTargetsResolution", None) - self.scaleIonDensities = EvolutionOptions.get("scaleIonDensities", True) - rho_vec = EvolutionOptions.get("rhoPredicted", [0.2, 0.4, 0.6, 0.8]) + self.fineTargetsResolution = evolution_options.get("fineTargetsResolution", None) + self.scaleIonDensities = evolution_options.get("scaleIonDensities", True) + rho_vec = evolution_options.get("rhoPredicted", [0.2, 0.4, 0.6, 0.8]) if rho_vec[0] == 0: raise ValueError("[MITIM] The radial grid must not contain the initial zero") @@ -96,7 +96,7 @@ def _ensure_ne_before_nz(lst): self.ProfilesPredicted = _ensure_ne_before_nz(self.ProfilesPredicted) # Default type and device tensor - self.dfT = torch.randn((2, 2), **tensor_opts) + self.dfT = torch.randn((2, 2), **tensor_options) ''' Potential profiles to evolve (aLX) and their corresponding flux matching @@ -227,7 +227,7 @@ def save(self, file): pickle.dump(self, handle, protocol=4) def combine_states(self, states, includeTransport=True): - self.TransportOptions_set = [self.TransportOptions] + self.transport_options_set = [self.transport_options] self.profiles_stored_set = self.profiles_stored for state in states: @@ -236,17 +236,17 @@ def combine_states(self, states, includeTransport=True): self.plasma[key] ) - self.TransportOptions_set.append(state.TransportOptions) + self.transport_options_set.append(state.transport_options) self.profiles_stored_set += state.profiles_stored if includeTransport: for key in ["chi_e", "chi_i"]: - self.TransportOptions["ModelOptions"][key] = torch.cat( + self.transport_options["ModelOptions"][key] = torch.cat( ( - self.TransportOptions["ModelOptions"][key], - state.TransportOptions["ModelOptions"][key], + self.transport_options["ModelOptions"][key], + state.transport_options["ModelOptions"][key], ) - ).to(self.TransportOptions["ModelOptions"][key]) + ).to(self.transport_options["ModelOptions"][key]) def copy_state(self): @@ -294,7 +294,7 @@ def calculate( self.calculateProfileFunctions() # 3. Sources and sinks (populates components and Pe,Pi,...) - relative_error_assumed = self.TransportOptions["ModelOptions"].get("percentError", [5, 1, 0.5])[-1] + relative_error_assumed = self.transport_options["ModelOptions"].get("percentError", [5, 1, 0.5])[-1] self.calculateTargets(relative_error_assumed=relative_error_assumed) # Calculate targets based on powerstate functions (it may be overwritten in next step, if chosen) # 4. Turbulent and neoclassical transport (populates components and Pe_tr,Pi_tr,...) @@ -363,7 +363,7 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): if folder_main is not None: folder = IOtools.expandPath(folder_main) / f"{namingConvention}_{cont}" - if issubclass(self.TransportOptions["transport_evaluator"], TRANSPORTtools.power_transport): + if issubclass(self.transport_options["transport_evaluator"], TRANSPORTtools.power_transport): (folder / "model_complete").mkdir(parents=True, exist_ok=True) # *************************************************************************************************************** @@ -377,7 +377,7 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): # Save state so that I can check initializations if folder_main is not None: - if issubclass(self.TransportOptions["transport_evaluator"], TRANSPORTtools.power_transport): + if issubclass(self.transport_options["transport_evaluator"], TRANSPORTtools.power_transport): self.save(folder / "powerstate.pkl") shutil.copy2(folder_run / "input.gacode", folder) @@ -685,10 +685,10 @@ def calculateTargets(self, relative_error_assumed=1.0): """ # If no targets evaluator is given or the targets will come from TGYRO, assume them as zero - if (self.TargetOptions["targets_evaluator"] is None) or (self.TargetOptions["ModelOptions"]["targets_evaluator_method"] == "tgyro"): + if (self.target_options["targets_evaluator"] is None) or (self.target_options["ModelOptions"]["targets_evaluator_method"] == "tgyro"): targets = TARGETStools.power_targets(self) else: - targets = self.TargetOptions["targets_evaluator"](self) + targets = self.target_options["targets_evaluator"](self) # [Optional] Calculate local targets and integrals on a fine grid if self.fineTargetsResolution is not None: @@ -707,7 +707,7 @@ def calculateTargets(self, relative_error_assumed=1.0): # Merge targets, calculate errors and normalize targets.postprocessing( relative_error_assumed=relative_error_assumed, - forceZeroParticleFlux=self.TransportOptions["ModelOptions"].get("forceZeroParticleFlux", False)) + forceZeroParticleFlux=self.transport_options["ModelOptions"].get("forceZeroParticleFlux", False)) def calculateTransport( self, nameRun="test", folder="~/scratch/", evaluation_number=0): @@ -717,10 +717,10 @@ def calculateTransport( folder = IOtools.expandPath(folder) # Select transport evaluator - if self.TransportOptions["transport_evaluator"] is None: + if self.transport_options["transport_evaluator"] is None: transport = TRANSPORTtools.power_transport( self, name=nameRun, folder=folder, evaluation_number=evaluation_number ) else: - transport = self.TransportOptions["transport_evaluator"]( self, name=nameRun, folder=folder, evaluation_number=evaluation_number ) + transport = self.transport_options["transport_evaluator"]( self, name=nameRun, folder=folder, evaluation_number=evaluation_number ) # Produce profile object (for certain transport evaluators, this is necessary) transport.produce_profiles() diff --git a/src/mitim_modules/powertorch/physics_models/targets_analytic.py b/src/mitim_modules/powertorch/physics_models/targets_analytic.py index b36dd63d..6d7f8695 100644 --- a/src/mitim_modules/powertorch/physics_models/targets_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/targets_analytic.py @@ -28,10 +28,10 @@ def __init__(self,powerstate, **kwargs): def evaluate(self): - if self.powerstate.TargetOptions["ModelOptions"]["TypeTarget"] >= 2: + if self.powerstate.target_options["ModelOptions"]["TypeTarget"] >= 2: self._evaluate_energy_exchange() - if self.powerstate.TargetOptions["ModelOptions"]["TypeTarget"] == 3: + if self.powerstate.target_options["ModelOptions"]["TypeTarget"] == 3: self._evaluate_alpha_heating() self._evaluate_radiation() diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py index a53ee114..ad47a7e4 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -13,8 +13,8 @@ def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) # Ensure that the provided diffusivities include the zero location - self.chi_e = self.powerstate.TransportOptions["ModelOptions"]["chi_e"] - self.chi_i = self.powerstate.TransportOptions["ModelOptions"]["chi_i"] + self.chi_e = self.powerstate.transport_options["ModelOptions"]["chi_e"] + self.chi_i = self.powerstate.transport_options["ModelOptions"]["chi_i"] if self.chi_e.shape[0] < self.powerstate.plasma['rho'].shape[-1]: self.chi_e = torch.cat((torch.zeros(1), self.chi_e)) @@ -73,7 +73,7 @@ def evaluate(self): for prof in self.powerstate.ProfilesPredicted: X = torch.cat((X,self.powerstate.plasma['aL'+prof][:,1:]),axis=1) - _, Q, _, _ = self.powerstate.TransportOptions["ModelOptions"]["flux_fun"](X) + _, Q, _, _ = self.powerstate.transport_options["ModelOptions"]["flux_fun"](X) numeach = self.powerstate.plasma["rho"].shape[1] - 1 diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index cb43d248..95099042 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -92,8 +92,8 @@ def _cgyro_trick(self,FolderEvaluation_TGYRO): # ************************************************************************************************************************** evaluateCGYRO( - self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["PORTALSparameters"], - self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"], + self.powerstate.transport_options["ModelOptions"]["extra_params"]["PORTALSparameters"], + self.powerstate.transport_options["ModelOptions"]["extra_params"]["folder"], self.evaluation_number, FolderEvaluation_TGYRO, self.file_profs, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 7b1acbd4..4514a066 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -32,7 +32,7 @@ def _evaluate_tglf(self): # Grab options from powerstate # ------------------------------------------------------------------------------------------------------------------------ - ModelOptions = self.powerstate.TransportOptions["ModelOptions"] + ModelOptions = self.powerstate.transport_options["ModelOptions"] MODELparameters = ModelOptions.get("MODELparameters",None) includeFast = ModelOptions.get("includeFastInQi",False) @@ -167,7 +167,7 @@ def _evaluate_neo(self): def _postprocess(self): - OriginalFimp = self.powerstate.TransportOptions["ModelOptions"].get("OriginalFimp", 1.0) + OriginalFimp = self.powerstate.transport_options["ModelOptions"].get("OriginalFimp", 1.0) # ------------------------------------------------------------------------------------------------------------------------ # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) @@ -217,8 +217,8 @@ def _postprocess(self): def _profiles_to_store(self): - if "extra_params" in self.powerstate.TransportOptions["ModelOptions"] and "folder" in self.powerstate.TransportOptions["ModelOptions"]["extra_params"]: - whereFolder = IOtools.expandPath(self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") + if "extra_params" in self.powerstate.transport_options["ModelOptions"] and "folder" in self.powerstate.transport_options["ModelOptions"]["extra_params"]: + whereFolder = IOtools.expandPath(self.powerstate.transport_options["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") if not whereFolder.exists(): IOtools.askNewFolder(whereFolder) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 482e362d..1385c7f4 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -27,7 +27,7 @@ def evaluate(self): def _evaluate_tglf_neo(self): - ModelOptions = self.powerstate.TransportOptions["ModelOptions"] + ModelOptions = self.powerstate.transport_options["ModelOptions"] MODELparameters = ModelOptions.get("MODELparameters",None) includeFast = ModelOptions.get("includeFastInQi",False) @@ -139,7 +139,7 @@ def _evaluate_tglf_neo(self): def _postprocess_results(self, tgyro, label): - ModelOptions = self.powerstate.TransportOptions["ModelOptions"] + ModelOptions = self.powerstate.transport_options["ModelOptions"] includeFast = ModelOptions.get("includeFastInQi",False) UseFineGridTargets = ModelOptions.get("UseFineGridTargets", False) @@ -160,7 +160,7 @@ def _postprocess_results(self, tgyro, label): OriginalFimp=OriginalFimp, forceZeroParticleFlux=forceZeroParticleFlux, provideTurbulentExchange=provideTurbulentExchange, - provideTargets=self.powerstate.TargetOptions['ModelOptions']['targets_evaluator_method'] == "tgyro", + provideTargets=self.powerstate.target_options['ModelOptions']['targets_evaluator_method'] == "tgyro", ) tgyro.results["use"] = tgyro.results[label] @@ -181,8 +181,8 @@ def _postprocess_results(self, tgyro, label): def _profiles_to_store(self): - if "extra_params" in self.powerstate.TransportOptions["ModelOptions"] and "folder" in self.powerstate.TransportOptions["ModelOptions"]["extra_params"]: - whereFolder = IOtools.expandPath(self.powerstate.TransportOptions["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") + if "extra_params" in self.powerstate.transport_options["ModelOptions"] and "folder" in self.powerstate.transport_options["ModelOptions"]["extra_params"]: + whereFolder = IOtools.expandPath(self.powerstate.transport_options["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") if not whereFolder.exists(): IOtools.askNewFolder(whereFolder) diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 368781f6..b61dd2a2 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -28,17 +28,17 @@ def calculator( if typeCalculation == 1: p = STATEtools.powerstate( profiles, - EvolutionOptions={ + evolution_options={ "rhoPredicted": rho_vec, 'fineTargetsResolution': fineTargetsResolution, }, - TargetOptions={ + target_options={ "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": TypeTarget, "targets_evaluator_method": "tgyro"}, }, - TransportOptions={ + transport_options={ "transport_evaluator": transport_tgyro.tgyro_model, "ModelOptions": { "cold_start": cold_start, @@ -70,17 +70,17 @@ def calculator( elif typeCalculation == 2: p = STATEtools.powerstate( profiles, - EvolutionOptions={ + evolution_options={ "rhoPredicted": rho_vec, 'fineTargetsResolution': fineTargetsResolution, }, - TargetOptions={ + target_options={ "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { "TypeTarget": TypeTarget, "targets_evaluator_method": "powerstate"}, }, - TransportOptions={ + transport_options={ "transport_evaluator": None, "ModelOptions": {} }, diff --git a/src/mitim_modules/powertorch/scripts/compareRadialResolution.py b/src/mitim_modules/powertorch/scripts/compareRadialResolution.py index d5a6646f..bdf140d9 100644 --- a/src/mitim_modules/powertorch/scripts/compareRadialResolution.py +++ b/src/mitim_modules/powertorch/scripts/compareRadialResolution.py @@ -35,14 +35,14 @@ ls = "o-" -sC = STATEtools.powerstate(profiles,EvolutionOptions={"rhoPredicted": rho},) +sC = STATEtools.powerstate(profiles,evolution_options={"rhoPredicted": rho},) sC.calculateProfileFunctions() sC.calculateTargets() # Full state rho = np.linspace(rho[0], rho[-1], args.res) -sF = STATEtools.powerstate(profiles,EvolutionOptions={"rhoPredicted": rho}) +sF = STATEtools.powerstate(profiles,evolution_options={"rhoPredicted": rho}) sF.calculateProfileFunctions() sF.calculateTargets() diff --git a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py index 564e784c..403f849f 100644 --- a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py +++ b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py @@ -28,9 +28,9 @@ # STATE -s = STATEtools.powerstate(t.profiles, EvolutionOptions={"rhoPredicted": t.rho[0,1:]}) +s = STATEtools.powerstate(t.profiles, evolution_options={"rhoPredicted": t.rho[0,1:]}) s.calculateProfileFunctions() -# s.TargetOptions['ModelOptions']['TypeTarget'] = 1 +# s.target_options['ModelOptions']['TypeTarget'] = 1 s.calculateTargets() # diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 1cfa1f9e..1b1e0f9f 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -71,11 +71,11 @@ def flux_integrate(self): qe = self.powerstate.plasma["te"]*0.0 qi = self.powerstate.plasma["te"]*0.0 - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] >= 2: + if self.powerstate.target_options['ModelOptions']['TypeTarget'] >= 2: qe += -self.powerstate.plasma["qie"] qi += self.powerstate.plasma["qie"] - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: + if self.powerstate.target_options['ModelOptions']['TypeTarget'] == 3: qe += self.powerstate.plasma["qfuse"] - self.powerstate.plasma["qrad"] qi += self.powerstate.plasma["qfusi"] @@ -89,11 +89,11 @@ def coarse_grid(self): # ************************************************************************************************** # Interpolate results from fine to coarse (i.e. whole point is that it is better than integrate interpolated values) - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] >= 2: + if self.powerstate.target_options['ModelOptions']['TypeTarget'] >= 2: for i in ["qie"]: self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] - if self.powerstate.TargetOptions['ModelOptions']['TypeTarget'] == 3: + if self.powerstate.target_options['ModelOptions']['TypeTarget'] == 3: for i in [ "qfuse", "qfusi", diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index f131b24a..ef97693a 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -120,12 +120,12 @@ def gacode_to_powerstate(self, rho_vec=None): quantitites["GZ_fixedtargets"] = input_gacode.derived["ge_10E20"] * 0.0 quantitites["MtJm2_fixedtargets"] = input_gacode.derived["mt_Jmiller"] - if self.TargetOptions["ModelOptions"]["TypeTarget"] < 3: + if self.target_options["ModelOptions"]["TypeTarget"] < 3: # Fusion and radiation fixed if 1,2 quantitites["QeMWm2_fixedtargets"] += input_gacode.derived["qe_fus_MW"] - input_gacode.derived["qrad_MW"] quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qi_fus_MW"] - if self.TargetOptions["ModelOptions"]["TypeTarget"] < 2: + if self.target_options["ModelOptions"]["TypeTarget"] < 2: # Exchange fixed if 1 quantitites["QeMWm2_fixedtargets"] -= input_gacode.derived["qe_exc_MW"] quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qe_exc_MW"] @@ -363,26 +363,26 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): with LOGtools.HiddenPrints(): state_temp.__init__( profiles, - EvolutionOptions={"rhoPredicted": rhoy}, - TargetOptions={ + evolution_options={"rhoPredicted": rhoy}, + target_options={ "targets_evaluator": targets_analytic.analytical_model, "ModelOptions": { - "TypeTarget": self.TargetOptions["ModelOptions"]["TypeTarget"], # Important to keep the same as in the original + "TypeTarget": self.target_options["ModelOptions"]["TypeTarget"], # Important to keep the same as in the original "targets_evaluator_method": "powerstate", } }, increase_profile_resol = False ) state_temp.calculateProfileFunctions() - state_temp.TargetOptions["ModelOptions"]["targets_evaluator_method"] = "powerstate" + state_temp.target_options["ModelOptions"]["targets_evaluator_method"] = "powerstate" state_temp.calculateTargets() # ------------------------------------------------------------------------------------------ conversions = {} - if self.TargetOptions["ModelOptions"]["TypeTarget"] > 1: + if self.target_options["ModelOptions"]["TypeTarget"] > 1: conversions['qie'] = "qei(MW/m^3)" - if self.TargetOptions["ModelOptions"]["TypeTarget"] > 2: + if self.target_options["ModelOptions"]["TypeTarget"] > 2: conversions['qrad_bremms'] = "qbrem(MW/m^3)" conversions['qrad_sync'] = "qsync(MW/m^3)" conversions['qrad_line'] = "qline(MW/m^3)" diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index fbf2ef42..b6c9fbcc 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -65,7 +65,7 @@ def produce_profiles(self): def _produce_profiles(self,derive_quantities=True): - self.applyCorrections = self.powerstate.TransportOptions["ModelOptions"].get("MODELparameters", {}).get("applyCorrections", {}) + self.applyCorrections = self.powerstate.transport_options["ModelOptions"].get("MODELparameters", {}).get("applyCorrections", {}) # Write this updated profiles class (with parameterized profiles and target powers) self.file_profs = self.folder / "input.gacode" @@ -92,7 +92,7 @@ def _modify_profiles(self): self.file_profs_unmod = self.file_profs.parent / f"{self.file_profs.name}_unmodified" shutil.copy2(self.file_profs, self.file_profs_unmod) - profiles_postprocessing_fun = self.powerstate.TransportOptions["ModelOptions"].get("profiles_postprocessing_fun", None) + profiles_postprocessing_fun = self.powerstate.transport_options["ModelOptions"].get("profiles_postprocessing_fun", None) if profiles_postprocessing_fun is not None: print(f"\t- Modifying input.gacode to run transport calculations based on {profiles_postprocessing_fun}",typeMsg="i") diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 49c0d40b..1533cfd1 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -69,7 +69,7 @@ def __init__( folder, namelist=None, default_namelist_function=None, - tensor_opts = { + tensor_options = { "dtype": torch.double, "device": torch.device("cpu"), } @@ -78,7 +78,7 @@ def __init__( Namelist file can be provided and will be copied to the folder """ - self.tensor_opts = tensor_opts + self.tensor_options = tensor_options print("- Parent opt_evaluator function initialized") @@ -123,8 +123,8 @@ def __init__( } # Determine type of tensors to work with - torch.set_default_dtype(self.tensor_opts["dtype"]) # In case I forgot to specify a type explicitly, use as default (https://github.com/pytorch/botorch/discussions/1444) - self.dfT = torch.randn( (2, 2), **tensor_opts) + torch.set_default_dtype(self.tensor_options["dtype"]) # In case I forgot to specify a type explicitly, use as default (https://github.com/pytorch/botorch/discussions/1444) + self.dfT = torch.randn( (2, 2), **tensor_options) # Name of calibrated objectives (e.g. QiRes1 to represent the objective from Qi1-QiT1) self.name_objectives = None diff --git a/src/mitim_tools/popcon_tools/RAPIDStools.py b/src/mitim_tools/popcon_tools/RAPIDStools.py index e6ffb963..6ee0dd8f 100644 --- a/src/mitim_tools/popcon_tools/RAPIDStools.py +++ b/src/mitim_tools/popcon_tools/RAPIDStools.py @@ -135,7 +135,7 @@ def pedestal(p): raise Exception('BetaN error too high') # Power - power = STATEtools.powerstate(p,EvolutionOptions={"rhoPredicted": np.linspace(0.0, 0.9, 50)[1:]}) + power = STATEtools.powerstate(p,evolution_options={"rhoPredicted": np.linspace(0.0, 0.9, 50)[1:]}) power.calculate(None, folder='~/scratch/power/') profiles_new = power.from_powerstate(insert_highres_powers=True) diff --git a/tests/POWERTORCH_workflow.py b/tests/POWERTORCH_workflow.py index d33b34e2..cea71585 100644 --- a/tests/POWERTORCH_workflow.py +++ b/tests/POWERTORCH_workflow.py @@ -11,10 +11,10 @@ rho = torch.from_numpy(np.linspace(0.1,0.9,9)).to(dtype=torch.double) s = STATEtools.powerstate(inputgacode, - EvolutionOptions = { 'ProfilePredicted': ['te', 'ti'], + evolution_options = { 'ProfilePredicted': ['te', 'ti'], 'rhoPredicted': rho }, - TransportOptions = { 'transport_evaluator': TRANSPORTtools.diffusion_model, + transport_options = { 'transport_evaluator': TRANSPORTtools.diffusion_model, 'ModelOptions': { 'chi_e': torch.ones(rho.shape[0]).to(rho)*0.8, 'chi_i': torch.ones(rho.shape[0]).to(rho)*1.2 From 9387abc9e4865fca7cdcf564342d90d7075ebcdf Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 5 Aug 2025 11:16:11 -0400 Subject: [PATCH 129/385] Renaming of ModelOptions --- regressions/portals_regressions.py | 4 +- src/mitim_modules/portals/PORTALSmain.py | 28 +++++++------- src/mitim_modules/portals/PORTALStools.py | 2 +- .../portals/utils/PORTALSinit.py | 26 ++++++------- .../portals/utils/PORTALSoptimization.py | 2 +- src/mitim_modules/powertorch/STATEtools.py | 28 +++++++------- .../physics_models/targets_analytic.py | 4 +- .../physics_models/transport_analytic.py | 8 ++-- .../physics_models/transport_cgyro.py | 4 +- .../physics_models/transport_tglf.py | 28 +++++++------- .../physics_models/transport_tgyro.py | 38 +++++++++---------- .../powertorch/scripts/calculateTargets.py | 16 ++++---- .../powertorch/scripts/compareWithTGYRO.py | 2 +- .../powertorch/utils/TARGETStools.py | 8 ++-- .../powertorch/utils/TRANSFORMtools.py | 18 ++++----- .../powertorch/utils/TRANSPORTtools.py | 4 +- tests/POWERTORCH_workflow.py | 2 +- 17 files changed, 111 insertions(+), 111 deletions(-) diff --git a/regressions/portals_regressions.py b/regressions/portals_regressions.py index 61d6bac4..abd4f4ae 100644 --- a/regressions/portals_regressions.py +++ b/regressions/portals_regressions.py @@ -61,9 +61,9 @@ def conditions_regressions(variables): portals_fun.optimization_options["acquisition_options"]["optimizers"] = ["botorch"] portals_fun.PORTALSparameters["transport_evaluator"] = TRANSPORTtools.diffusion_model - ModelOptions = {'chi_e': torch.ones(5)*0.5,'chi_i': torch.ones(5)*2.0} + transport_evaluator_options = {'chi_e': torch.ones(5)*0.5,'chi_i': torch.ones(5)*2.0} - portals_fun.prep(inputgacode, folderWork, ModelOptions=ModelOptions) + portals_fun.prep(inputgacode, folderWork, transport_evaluator_options=transport_evaluator_options) mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=False, askQuestions=False) mitim_bo.run() diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 58119fdb..27b39ffc 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -197,13 +197,13 @@ def __init__( else: transport_evaluator = transport_tgyro.tgyro_model - targets_evaluator = targets_analytic.analytical_model + target_evaluator = targets_analytic.analytical_model self.PORTALSparameters = { "percentError": [5,10,1], # (%) Error (std, in percent) of model evaluation [TGLF (treated as minimum if scan trick), NEO, TARGET] "transport_evaluator": transport_evaluator, - "targets_evaluator": targets_evaluator, - "targets_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) + "target_evaluator": target_evaluator, + "target_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) "includeFastInQi": False, # If True, and fast ions have been included, in seprateNEO, sum fast "portals_transformation_variables": portals_transformation_variables, # Physics-informed parameters to fit surrogates @@ -218,7 +218,7 @@ def __init__( "applyImpurityGammaTrick": True, # If True, fit model to GZ/nZ, valid on the trace limit "UseOriginalImpurityConcentrationAsWeight": 1.0, # If not None, using UseOriginalImpurityConcentrationAsWeight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis "fImp_orig": 1.0, - "fineTargetsResolution": 20, # If not None, calculate targets with this radial resolution (defaults targets_evaluator_method to powerstate) + "fineTargetsResolution": 20, # If not None, calculate targets with this radial resolution (defaults target_evaluator_method to powerstate) "hardCodedCGYRO": None, # If not None, use this hard-coded CGYRO evaluation "additional_params_in_surrogate": additional_params_in_surrogate, "use_tglf_scan_trick": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta @@ -244,7 +244,7 @@ def prep( reevaluateTargets=0, seedInitial=None, askQuestions=True, - ModelOptions=None, + transport_evaluator_options=None, ): """ Notes: @@ -311,7 +311,7 @@ def prep( tensor_options = self.tensor_options, seedInitial=seedInitial, checkForSpecies=askQuestions, - ModelOptions=ModelOptions, + transport_evaluator_options=transport_evaluator_options, ) print(">> PORTALS initalization module (END)", typeMsg="i") @@ -463,13 +463,13 @@ def check_flags(self): # ---------------------------------------------------------------------------------- if self.PORTALSparameters["fineTargetsResolution"] is not None: - if self.PORTALSparameters["targets_evaluator_method"] != "powerstate": + if self.PORTALSparameters["target_evaluator_method"] != "powerstate": print("\t- Requested fineTargetsResolution, so running powerstate target calculations",typeMsg="w") - self.PORTALSparameters["targets_evaluator_method"] = "powerstate" + self.PORTALSparameters["target_evaluator_method"] = "powerstate" - if not issubclass(self.PORTALSparameters["transport_evaluator"], transport_tgyro.tgyro_model) and (self.PORTALSparameters["targets_evaluator_method"] == "tgyro"): + if not issubclass(self.PORTALSparameters["transport_evaluator"], transport_tgyro.tgyro_model) and (self.PORTALSparameters["target_evaluator_method"] == "tgyro"): print("\t- Requested TGYRO targets, but transport evaluator is not tgyro, so changing to powerstate",typeMsg="w") - self.PORTALSparameters["targets_evaluator_method"] = "powerstate" + self.PORTALSparameters["target_evaluator_method"] = "powerstate" if ("InputType" not in self.MODELparameters["Physics_options"]) or self.MODELparameters["Physics_options"]["InputType"] != 1: print("\t- In PORTALS TGYRO evaluations, we need to use exact profiles (InputType=1)",typeMsg="i") @@ -479,10 +479,10 @@ def check_flags(self): print("\t- In PORTALS TGYRO evaluations, we need to not recompute gradients (GradientsType=0)",typeMsg="i") self.MODELparameters["Physics_options"]["GradientsType"] = 0 - if self.PORTALSparameters["targets_evaluator_method"] == "tgyro" and self.PORTALSparameters['profiles_postprocessing_fun'] is not None: + if self.PORTALSparameters["target_evaluator_method"] == "tgyro" and self.PORTALSparameters['profiles_postprocessing_fun'] is not None: print("\t- Requested custom modification of postprocessing function but targets from TGYRO... are you sure?",typeMsg="q") - if self.PORTALSparameters["targets_evaluator_method"] == "tgyro" and self.PORTALSparameters['transport_evaluator'] != transport_tgyro.tgyro_model: + if self.PORTALSparameters["target_evaluator_method"] == "tgyro" and self.PORTALSparameters['transport_evaluator'] != transport_tgyro.tgyro_model: print("\t- Requested TGYRO targets but transport evaluator is not TGYRO... are you sure?",typeMsg="q") key_rhos = "RoaLocations" if self.MODELparameters["RoaLocations"] is not None else "RhoLocations" @@ -552,7 +552,7 @@ def reuseTrainingTabular( self_copy = copy.deepcopy(self) if reevaluateTargets == 1: self_copy.powerstate.transport_options["transport_evaluator"] = None - self_copy.powerstate.target_options["ModelOptions"]["TypeTarget"] = "powerstate" + self_copy.powerstate.target_options["target_evaluator_options"]["TypeTarget"] = "powerstate" else: self_copy.powerstate.transport_options["transport_evaluator"] = transport_tgyro.tgyro_model @@ -623,7 +623,7 @@ def runModelEvaluator( # --------------------------------------------------------------------------------------------------- # In certain cases, I want to cold_start the model directly from the PORTALS call instead of powerstate - powerstate.transport_options["ModelOptions"]["cold_start"] = cold_start + powerstate.transport_options["transport_evaluator_options"]["cold_start"] = cold_start # Evaluate X (DVs) through powerstate.calculate(). This will populate .plasma with the results powerstate.calculate(X, nameRun=name, folder=folder_model, evaluation_number=numPORTALS) diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index a5ac3ce2..4c2f66c8 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -303,7 +303,7 @@ def constructEvaluationProfiles(X, surrogate_parameters, recalculateTargets=Fals # Targets only if needed (for speed, GB doesn't need it) if recalculateTargets: - powerstate.target_options["ModelOptions"]["targets_evaluator_method"] = "powerstate" # For surrogate evaluation, always powerstate, logically. + powerstate.target_options["target_evaluator_options"]["target_evaluator_method"] = "powerstate" # For surrogate evaluation, always powerstate, logically. powerstate.calculateTargets() return powerstate diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index c819f251..daa2a780 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -27,7 +27,7 @@ def initializeProblem( dvs_fixed=None, start_from_folder=None, define_ranges_from_profiles=None, - ModelOptions=None, + transport_evaluator_options=None, seedInitial=None, checkForSpecies=True, tensor_options = { @@ -39,7 +39,7 @@ def initializeProblem( Notes: - Specification of points occur in rho coordinate, although internally the work is r/a cold_start = True if cold_start from beginning - - I can give ModelOptions directly (e.g. if I want chis or something) + - I can give transport_evaluator_options directly (e.g. if I want chis or something) - define_ranges_from_profiles must be PROFILES class """ @@ -126,8 +126,8 @@ def initializeProblem( xCPs = torch.from_numpy(np.array(portals_fun.MODELparameters["RhoLocations"])).to(dfT) - if ModelOptions is None: - ModelOptions = { + if transport_evaluator_options is None: + transport_evaluator_options = { "launchMODELviaSlurm": portals_fun.PORTALSparameters["launchEvaluationsAsSlurmJobs"], "includeFastInQi": portals_fun.PORTALSparameters["includeFastInQi"], "TurbulentExchange": portals_fun.PORTALSparameters["surrogateForTurbExch"], @@ -142,8 +142,8 @@ def initializeProblem( "impurityPosition": position_of_impurity, } - if "extra_params" not in ModelOptions: - ModelOptions["extra_params"] = { + if "extra_params" not in transport_evaluator_options: + transport_evaluator_options["extra_params"] = { "PORTALSparameters": portals_fun.PORTALSparameters, "folder": portals_fun.folder, } @@ -165,13 +165,13 @@ def initializeProblem( }, transport_options={ "transport_evaluator": portals_fun.PORTALSparameters["transport_evaluator"], - "ModelOptions": ModelOptions, + "transport_evaluator_options": transport_evaluator_options, }, target_options={ - "targets_evaluator": portals_fun.PORTALSparameters["targets_evaluator"], - "ModelOptions": { + "target_evaluator": portals_fun.PORTALSparameters["target_evaluator"], + "target_evaluator_options": { "TypeTarget": portals_fun.MODELparameters["Physics_options"]["TypeTarget"], - "targets_evaluator_method": portals_fun.PORTALSparameters["targets_evaluator_method"]}, + "target_evaluator_method": portals_fun.PORTALSparameters["target_evaluator_method"]}, }, tensor_options = tensor_options ) @@ -219,10 +219,10 @@ def initializeProblem( "fineTargetsResolution": portals_fun.PORTALSparameters["fineTargetsResolution"], }, target_options={ - "targets_evaluator": portals_fun.PORTALSparameters["targets_evaluator"], - "ModelOptions": { + "target_evaluator": portals_fun.PORTALSparameters["target_evaluator"], + "transport_evaluator_options": { "TypeTarget": portals_fun.MODELparameters["Physics_options"]["TypeTarget"], - "targets_evaluator_method": portals_fun.PORTALSparameters["targets_evaluator_method"]}, + "target_evaluator_method": portals_fun.PORTALSparameters["target_evaluator_method"]}, }, tensor_options = tensor_options ) diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index a354f6d4..af6065f6 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -140,7 +140,7 @@ def flux_match_surrogate( # Define transport calculation function as a surrogate model transport_options['transport_evaluator'] = transport_analytic.surrogate - transport_options['ModelOptions'] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} + transport_options['transport_evaluator_options'] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} # Create powerstate with the same options as the original portals but with the new profiles powerstate = STATEtools.powerstate( diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 9848a273..3b6c61b7 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -39,8 +39,8 @@ def __init__( - ProfilesPredicted: list of profiles to predict - impurityPosition: int = position of the impurity in the ions set - fineTargetsResolution: int = resolution of the fine targets - - transport_options: dictionary with transport_evaluator and ModelOptions - - target_options: dictionary with targets_evaluator and ModelOptions + - transport_options: dictionary with transport_evaluator and transport_evaluator_options + - target_options: dictionary with target_evaluator and target_evaluator_options ''' if evolution_options is None: @@ -48,14 +48,14 @@ def __init__( if transport_options is None: transport_options = { "transport_evaluator": None, - "ModelOptions": {} + "transport_evaluator_options": {} } if target_options is None: target_options = { - "targets_evaluator": targets_analytic.analytical_model, - "ModelOptions": { + "target_evaluator": targets_analytic.analytical_model, + "target_evaluator_options": { "TypeTarget": 3, - "targets_evaluator_method": "powerstate" + "target_evaluator_method": "powerstate" }, } if tensor_options is None: @@ -241,12 +241,12 @@ def combine_states(self, states, includeTransport=True): if includeTransport: for key in ["chi_e", "chi_i"]: - self.transport_options["ModelOptions"][key] = torch.cat( + self.transport_options["transport_evaluator_options"][key] = torch.cat( ( - self.transport_options["ModelOptions"][key], - state.transport_options["ModelOptions"][key], + self.transport_options["transport_evaluator_options"][key], + state.transport_options["transport_evaluator_options"][key], ) - ).to(self.transport_options["ModelOptions"][key]) + ).to(self.transport_options["transport_evaluator_options"][key]) def copy_state(self): @@ -294,7 +294,7 @@ def calculate( self.calculateProfileFunctions() # 3. Sources and sinks (populates components and Pe,Pi,...) - relative_error_assumed = self.transport_options["ModelOptions"].get("percentError", [5, 1, 0.5])[-1] + relative_error_assumed = self.transport_options["transport_evaluator_options"].get("percentError", [5, 1, 0.5])[-1] self.calculateTargets(relative_error_assumed=relative_error_assumed) # Calculate targets based on powerstate functions (it may be overwritten in next step, if chosen) # 4. Turbulent and neoclassical transport (populates components and Pe_tr,Pi_tr,...) @@ -685,10 +685,10 @@ def calculateTargets(self, relative_error_assumed=1.0): """ # If no targets evaluator is given or the targets will come from TGYRO, assume them as zero - if (self.target_options["targets_evaluator"] is None) or (self.target_options["ModelOptions"]["targets_evaluator_method"] == "tgyro"): + if (self.target_options["target_evaluator"] is None) or (self.target_options["target_evaluator_options"]["target_evaluator_method"] == "tgyro"): targets = TARGETStools.power_targets(self) else: - targets = self.target_options["targets_evaluator"](self) + targets = self.target_options["target_evaluator"](self) # [Optional] Calculate local targets and integrals on a fine grid if self.fineTargetsResolution is not None: @@ -707,7 +707,7 @@ def calculateTargets(self, relative_error_assumed=1.0): # Merge targets, calculate errors and normalize targets.postprocessing( relative_error_assumed=relative_error_assumed, - forceZeroParticleFlux=self.transport_options["ModelOptions"].get("forceZeroParticleFlux", False)) + forceZeroParticleFlux=self.transport_options["transport_evaluator_options"].get("forceZeroParticleFlux", False)) def calculateTransport( self, nameRun="test", folder="~/scratch/", evaluation_number=0): diff --git a/src/mitim_modules/powertorch/physics_models/targets_analytic.py b/src/mitim_modules/powertorch/physics_models/targets_analytic.py index 6d7f8695..6f3956ad 100644 --- a/src/mitim_modules/powertorch/physics_models/targets_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/targets_analytic.py @@ -28,10 +28,10 @@ def __init__(self,powerstate, **kwargs): def evaluate(self): - if self.powerstate.target_options["ModelOptions"]["TypeTarget"] >= 2: + if self.powerstate.target_options["target_evaluator_options"]["TypeTarget"] >= 2: self._evaluate_energy_exchange() - if self.powerstate.target_options["ModelOptions"]["TypeTarget"] == 3: + if self.powerstate.target_options["target_evaluator_options"]["TypeTarget"] == 3: self._evaluate_alpha_heating() self._evaluate_radiation() diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py index ad47a7e4..57ac1922 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -13,8 +13,8 @@ def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) # Ensure that the provided diffusivities include the zero location - self.chi_e = self.powerstate.transport_options["ModelOptions"]["chi_e"] - self.chi_i = self.powerstate.transport_options["ModelOptions"]["chi_i"] + self.chi_e = self.powerstate.transport_options["transport_evaluator_options"]["chi_e"] + self.chi_i = self.powerstate.transport_options["transport_evaluator_options"]["chi_i"] if self.chi_e.shape[0] < self.powerstate.plasma['rho'].shape[-1]: self.chi_e = torch.cat((torch.zeros(1), self.chi_e)) @@ -66,14 +66,14 @@ def produce_profiles(self): def evaluate(self): """ - flux_fun as given in ModelOptions must produce Q and Qtargets in order of te,ti,ne + flux_fun as given in transport_evaluator_options must produce Q and Qtargets in order of te,ti,ne """ X = torch.Tensor() for prof in self.powerstate.ProfilesPredicted: X = torch.cat((X,self.powerstate.plasma['aL'+prof][:,1:]),axis=1) - _, Q, _, _ = self.powerstate.transport_options["ModelOptions"]["flux_fun"](X) + _, Q, _, _ = self.powerstate.transport_options["transport_evaluator_options"]["flux_fun"](X) numeach = self.powerstate.plasma["rho"].shape[1] - 1 diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 95099042..d8cb7fb2 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -92,8 +92,8 @@ def _cgyro_trick(self,FolderEvaluation_TGYRO): # ************************************************************************************************************************** evaluateCGYRO( - self.powerstate.transport_options["ModelOptions"]["extra_params"]["PORTALSparameters"], - self.powerstate.transport_options["ModelOptions"]["extra_params"]["folder"], + self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]["PORTALSparameters"], + self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]["folder"], self.evaluation_number, FolderEvaluation_TGYRO, self.file_profs, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 4514a066..50fcf364 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -32,19 +32,19 @@ def _evaluate_tglf(self): # Grab options from powerstate # ------------------------------------------------------------------------------------------------------------------------ - ModelOptions = self.powerstate.transport_options["ModelOptions"] - - MODELparameters = ModelOptions.get("MODELparameters",None) - includeFast = ModelOptions.get("includeFastInQi",False) - launchMODELviaSlurm = ModelOptions.get("launchMODELviaSlurm", False) - cold_start = ModelOptions.get("cold_start", False) - provideTurbulentExchange = ModelOptions.get("TurbulentExchange", False) - percentError = ModelOptions.get("percentError", [5, 1, 0.5]) - use_tglf_scan_trick = ModelOptions.get("use_tglf_scan_trick", None) - cores_per_tglf_instance = ModelOptions.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) + transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + + MODELparameters = transport_evaluator_options.get("MODELparameters",None) + includeFast = transport_evaluator_options.get("includeFastInQi",False) + launchMODELviaSlurm = transport_evaluator_options.get("launchMODELviaSlurm", False) + cold_start = transport_evaluator_options.get("cold_start", False) + provideTurbulentExchange = transport_evaluator_options.get("TurbulentExchange", False) + percentError = transport_evaluator_options.get("percentError", [5, 1, 0.5]) + use_tglf_scan_trick = transport_evaluator_options.get("use_tglf_scan_trick", None) + cores_per_tglf_instance = transport_evaluator_options.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) - impurityPosition = self.powerstate.impurityPosition_transport #ModelOptions.get("impurityPosition", 1) + impurityPosition = self.powerstate.impurityPosition_transport #transport_evaluator_options.get("impurityPosition", 1) # ------------------------------------------------------------------------------------------------------------------------ # Prepare TGLF object @@ -167,7 +167,7 @@ def _evaluate_neo(self): def _postprocess(self): - OriginalFimp = self.powerstate.transport_options["ModelOptions"].get("OriginalFimp", 1.0) + OriginalFimp = self.powerstate.transport_options["transport_evaluator_options"].get("OriginalFimp", 1.0) # ------------------------------------------------------------------------------------------------------------------------ # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) @@ -217,8 +217,8 @@ def _postprocess(self): def _profiles_to_store(self): - if "extra_params" in self.powerstate.transport_options["ModelOptions"] and "folder" in self.powerstate.transport_options["ModelOptions"]["extra_params"]: - whereFolder = IOtools.expandPath(self.powerstate.transport_options["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") + if "extra_params" in self.powerstate.transport_options["transport_evaluator_options"] and "folder" in self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]: + whereFolder = IOtools.expandPath(self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") if not whereFolder.exists(): IOtools.askNewFolder(whereFolder) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 1385c7f4..1b27e7ab 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -27,19 +27,19 @@ def evaluate(self): def _evaluate_tglf_neo(self): - ModelOptions = self.powerstate.transport_options["ModelOptions"] + transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] - MODELparameters = ModelOptions.get("MODELparameters",None) - includeFast = ModelOptions.get("includeFastInQi",False) - launchMODELviaSlurm = ModelOptions.get("launchMODELviaSlurm", False) - cold_start = ModelOptions.get("cold_start", False) - provideTurbulentExchange = ModelOptions.get("TurbulentExchange", False) - percentError = ModelOptions.get("percentError", [5, 1, 0.5]) - use_tglf_scan_trick = ModelOptions.get("use_tglf_scan_trick", None) - cores_per_tglf_instance = ModelOptions.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) + MODELparameters = transport_evaluator_options.get("MODELparameters",None) + includeFast = transport_evaluator_options.get("includeFastInQi",False) + launchMODELviaSlurm = transport_evaluator_options.get("launchMODELviaSlurm", False) + cold_start = transport_evaluator_options.get("cold_start", False) + provideTurbulentExchange = transport_evaluator_options.get("TurbulentExchange", False) + percentError = transport_evaluator_options.get("percentError", [5, 1, 0.5]) + use_tglf_scan_trick = transport_evaluator_options.get("use_tglf_scan_trick", None) + cores_per_tglf_instance = transport_evaluator_options.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) - impurityPosition = self.powerstate.impurityPosition_transport #ModelOptions.get("impurityPosition", 1) + impurityPosition = self.powerstate.impurityPosition_transport #transport_evaluator_options.get("impurityPosition", 1) # ------------------------------------------------------------------------------------------------------------------------ # tglf_neo_original: Run TGYRO workflow - TGLF + NEO in subfolder tglf_neo_original (original as in... without stds or merging) @@ -139,13 +139,13 @@ def _evaluate_tglf_neo(self): def _postprocess_results(self, tgyro, label): - ModelOptions = self.powerstate.transport_options["ModelOptions"] + transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] - includeFast = ModelOptions.get("includeFastInQi",False) - UseFineGridTargets = ModelOptions.get("UseFineGridTargets", False) - provideTurbulentExchange = ModelOptions.get("TurbulentExchange", False) - OriginalFimp = ModelOptions.get("OriginalFimp", 1.0) - forceZeroParticleFlux = ModelOptions.get("forceZeroParticleFlux", False) + includeFast = transport_evaluator_options.get("includeFastInQi",False) + UseFineGridTargets = transport_evaluator_options.get("UseFineGridTargets", False) + provideTurbulentExchange = transport_evaluator_options.get("TurbulentExchange", False) + OriginalFimp = transport_evaluator_options.get("OriginalFimp", 1.0) + forceZeroParticleFlux = transport_evaluator_options.get("forceZeroParticleFlux", False) # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) impurityPosition = self.powerstate.impurityPosition_transport @@ -160,7 +160,7 @@ def _postprocess_results(self, tgyro, label): OriginalFimp=OriginalFimp, forceZeroParticleFlux=forceZeroParticleFlux, provideTurbulentExchange=provideTurbulentExchange, - provideTargets=self.powerstate.target_options['ModelOptions']['targets_evaluator_method'] == "tgyro", + provideTargets=self.powerstate.target_options['target_evaluator_options']['target_evaluator_method'] == "tgyro", ) tgyro.results["use"] = tgyro.results[label] @@ -181,8 +181,8 @@ def _postprocess_results(self, tgyro, label): def _profiles_to_store(self): - if "extra_params" in self.powerstate.transport_options["ModelOptions"] and "folder" in self.powerstate.transport_options["ModelOptions"]["extra_params"]: - whereFolder = IOtools.expandPath(self.powerstate.transport_options["ModelOptions"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") + if "extra_params" in self.powerstate.transport_options["transport_evaluator_options"] and "folder" in self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]: + whereFolder = IOtools.expandPath(self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") if not whereFolder.exists(): IOtools.askNewFolder(whereFolder) diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index b61dd2a2..26a7c376 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -33,14 +33,14 @@ def calculator( 'fineTargetsResolution': fineTargetsResolution, }, target_options={ - "targets_evaluator": targets_analytic.analytical_model, - "ModelOptions": { + "target_evaluator": targets_analytic.analytical_model, + "target_evaluator_options": { "TypeTarget": TypeTarget, - "targets_evaluator_method": "tgyro"}, + "target_evaluator_method": "tgyro"}, }, transport_options={ "transport_evaluator": transport_tgyro.tgyro_model, - "ModelOptions": { + "transport_evaluator_options": { "cold_start": cold_start, "launchSlurm": True, "MODELparameters": { @@ -75,14 +75,14 @@ def calculator( 'fineTargetsResolution': fineTargetsResolution, }, target_options={ - "targets_evaluator": targets_analytic.analytical_model, - "ModelOptions": { + "target_evaluator": targets_analytic.analytical_model, + "target_evaluator_options": { "TypeTarget": TypeTarget, - "targets_evaluator_method": "powerstate"}, + "target_evaluator_method": "powerstate"}, }, transport_options={ "transport_evaluator": None, - "ModelOptions": {} + "transport_evaluator_options": {} }, ) diff --git a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py index 403f849f..dabae966 100644 --- a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py +++ b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py @@ -30,7 +30,7 @@ # STATE s = STATEtools.powerstate(t.profiles, evolution_options={"rhoPredicted": t.rho[0,1:]}) s.calculateProfileFunctions() -# s.target_options['ModelOptions']['TypeTarget'] = 1 +# s.target_options['target_evaluator_options']['TypeTarget'] = 1 s.calculateTargets() # diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 1b1e0f9f..2b0580d9 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -71,11 +71,11 @@ def flux_integrate(self): qe = self.powerstate.plasma["te"]*0.0 qi = self.powerstate.plasma["te"]*0.0 - if self.powerstate.target_options['ModelOptions']['TypeTarget'] >= 2: + if self.powerstate.target_options['target_evaluator_options']['TypeTarget'] >= 2: qe += -self.powerstate.plasma["qie"] qi += self.powerstate.plasma["qie"] - if self.powerstate.target_options['ModelOptions']['TypeTarget'] == 3: + if self.powerstate.target_options['target_evaluator_options']['TypeTarget'] == 3: qe += self.powerstate.plasma["qfuse"] - self.powerstate.plasma["qrad"] qi += self.powerstate.plasma["qfusi"] @@ -89,11 +89,11 @@ def coarse_grid(self): # ************************************************************************************************** # Interpolate results from fine to coarse (i.e. whole point is that it is better than integrate interpolated values) - if self.powerstate.target_options['ModelOptions']['TypeTarget'] >= 2: + if self.powerstate.target_options['target_evaluator_options']['TypeTarget'] >= 2: for i in ["qie"]: self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] - if self.powerstate.target_options['ModelOptions']['TypeTarget'] == 3: + if self.powerstate.target_options['target_evaluator_options']['TypeTarget'] == 3: for i in [ "qfuse", "qfusi", diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index ef97693a..f30d5a81 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -120,12 +120,12 @@ def gacode_to_powerstate(self, rho_vec=None): quantitites["GZ_fixedtargets"] = input_gacode.derived["ge_10E20"] * 0.0 quantitites["MtJm2_fixedtargets"] = input_gacode.derived["mt_Jmiller"] - if self.target_options["ModelOptions"]["TypeTarget"] < 3: + if self.target_options["target_evaluator_options"]["TypeTarget"] < 3: # Fusion and radiation fixed if 1,2 quantitites["QeMWm2_fixedtargets"] += input_gacode.derived["qe_fus_MW"] - input_gacode.derived["qrad_MW"] quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qi_fus_MW"] - if self.target_options["ModelOptions"]["TypeTarget"] < 2: + if self.target_options["target_evaluator_options"]["TypeTarget"] < 2: # Exchange fixed if 1 quantitites["QeMWm2_fixedtargets"] -= input_gacode.derived["qe_exc_MW"] quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qe_exc_MW"] @@ -365,24 +365,24 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): profiles, evolution_options={"rhoPredicted": rhoy}, target_options={ - "targets_evaluator": targets_analytic.analytical_model, - "ModelOptions": { - "TypeTarget": self.target_options["ModelOptions"]["TypeTarget"], # Important to keep the same as in the original - "targets_evaluator_method": "powerstate", + "target_evaluator": targets_analytic.analytical_model, + "target_evaluator_options": { + "TypeTarget": self.target_options["target_evaluator_options"]["TypeTarget"], # Important to keep the same as in the original + "target_evaluator_method": "powerstate", } }, increase_profile_resol = False ) state_temp.calculateProfileFunctions() - state_temp.target_options["ModelOptions"]["targets_evaluator_method"] = "powerstate" + state_temp.target_options["target_evaluator_options"]["target_evaluator_method"] = "powerstate" state_temp.calculateTargets() # ------------------------------------------------------------------------------------------ conversions = {} - if self.target_options["ModelOptions"]["TypeTarget"] > 1: + if self.target_options["target_evaluator_options"]["TypeTarget"] > 1: conversions['qie'] = "qei(MW/m^3)" - if self.target_options["ModelOptions"]["TypeTarget"] > 2: + if self.target_options["target_evaluator_options"]["TypeTarget"] > 2: conversions['qrad_bremms'] = "qbrem(MW/m^3)" conversions['qrad_sync'] = "qsync(MW/m^3)" conversions['qrad_line'] = "qline(MW/m^3)" diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index b6c9fbcc..270aeb7c 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -65,7 +65,7 @@ def produce_profiles(self): def _produce_profiles(self,derive_quantities=True): - self.applyCorrections = self.powerstate.transport_options["ModelOptions"].get("MODELparameters", {}).get("applyCorrections", {}) + self.applyCorrections = self.powerstate.transport_options["transport_evaluator_options"].get("MODELparameters", {}).get("applyCorrections", {}) # Write this updated profiles class (with parameterized profiles and target powers) self.file_profs = self.folder / "input.gacode" @@ -92,7 +92,7 @@ def _modify_profiles(self): self.file_profs_unmod = self.file_profs.parent / f"{self.file_profs.name}_unmodified" shutil.copy2(self.file_profs, self.file_profs_unmod) - profiles_postprocessing_fun = self.powerstate.transport_options["ModelOptions"].get("profiles_postprocessing_fun", None) + profiles_postprocessing_fun = self.powerstate.transport_options["transport_evaluator_options"].get("profiles_postprocessing_fun", None) if profiles_postprocessing_fun is not None: print(f"\t- Modifying input.gacode to run transport calculations based on {profiles_postprocessing_fun}",typeMsg="i") diff --git a/tests/POWERTORCH_workflow.py b/tests/POWERTORCH_workflow.py index cea71585..3eab027f 100644 --- a/tests/POWERTORCH_workflow.py +++ b/tests/POWERTORCH_workflow.py @@ -15,7 +15,7 @@ 'rhoPredicted': rho }, transport_options = { 'transport_evaluator': TRANSPORTtools.diffusion_model, - 'ModelOptions': { + 'transport_evaluator_options': { 'chi_e': torch.ones(rho.shape[0]).to(rho)*0.8, 'chi_i': torch.ones(rho.shape[0]).to(rho)*1.2 } From 2f10421e4388b14377aa48bc825e0a9eff53d1b0 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 5 Aug 2025 13:29:19 -0400 Subject: [PATCH 130/385] Clarified and fixed the treatment of Qifast in PORTALS --- regressions/portals_regressions.py | 8 +- .../maestro/tmp_tests/maestro_test1.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 14 +-- .../portals/utils/PORTALSanalysis.py | 2 +- .../portals/utils/PORTALSinit.py | 8 +- src/mitim_modules/powertorch/STATEtools.py | 3 + .../physics_models/transport_tglf.py | 93 +++++++++++++------ .../physics_models/transport_tgyro.py | 66 +++++++------ .../powertorch/scripts/calculateTargets.py | 6 +- .../powertorch/utils/TRANSFORMtools.py | 6 +- src/mitim_tools/gacode_tools/TGLFtools.py | 64 +++++++------ src/mitim_tools/gacode_tools/TGYROtools.py | 4 +- .../gacode_tools/scripts/run_tglf.py | 4 +- .../plasmastate_tools/MITIMstate.py | 20 ++-- src/mitim_tools/transp_tools/CDFtools.py | 6 +- templates/maestro_namelist.json | 2 +- tests/PORTALS_workflow.py | 4 +- tutorials/PORTALS_tutorial.py | 2 +- 18 files changed, 185 insertions(+), 129 deletions(-) diff --git a/regressions/portals_regressions.py b/regressions/portals_regressions.py index abd4f4ae..9890980f 100644 --- a/regressions/portals_regressions.py +++ b/regressions/portals_regressions.py @@ -55,7 +55,7 @@ def conditions_regressions(variables): portals_fun = PORTALSmain.portals(folderWork) portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - portals_fun.INITparameters["removeFast"] = True + portals_fun.INITparameters["remove_fast"] = True portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti"] portals_fun.optimization_options["acquisition_options"]["optimizers"] = ["botorch"] @@ -94,9 +94,9 @@ def conditions_regressions(variables): portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 portals_fun.MODELparameters["RhoLocations"] = [0.25, 0.45, 0.65, 0.85] - portals_fun.INITparameters["removeFast"] = True + portals_fun.INITparameters["remove_fast"] = True portals_fun.INITparameters["quasineutrality"] = True - portals_fun.INITparameters["sameDensityGradients"] = True + portals_fun.INITparameters["enforce_same_aLn"] = True portals_fun.MODELparameters["transport_model"]["TGLFsettings"] = 2 portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti", "ne"] @@ -134,7 +134,7 @@ def conditions_regressions(variables): # portals_fun = PORTALSmain.portals(folderWork) # portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 # portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - # portals_fun.INITparameters["removeFast"] = True + # portals_fun.INITparameters["remove_fast"] = True # portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti", "ne",'nZ','w0'] diff --git a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py b/src/mitim_modules/maestro/tmp_tests/maestro_test1.py index e1018d28..32846cb6 100644 --- a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py +++ b/src/mitim_modules/maestro/tmp_tests/maestro_test1.py @@ -47,7 +47,7 @@ "ProfilesPredicted": ["te", "ti", "ne"], "Physics_options": {"TypeTarget": 3}, "transport_model": {"TGLFsettings": 6, "extraOptionsTGLF": {'USE_BPER':True}}}, - "INITparameters": {"FastIsThermal": True, "removeIons": [5,6], "quasineutrality": True}, + "INITparameters": {"thermalize_fast": True, "removeIons": [5,6], "quasineutrality": True}, "optimization_options": { "convergence_options": { "maximum_iterations": 50, diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 27b39ffc..d0e98e40 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -128,14 +128,14 @@ def __init__( """ self.INITparameters = { - "recompute_ptot": True, # Recompute PTOT to match kinetic profiles (after removals) + "recalculate_ptot": True, # Recompute PTOT to match kinetic profiles (after removals) "quasineutrality": False, # Make sure things are quasineutral by changing the *MAIN* ion (D,T or both) (after removals) "removeIons": [], # Remove this ion from the input.gacode (if D,T,Z, eliminate T with [2]) - "removeFast": False, # Automatically detect which are fast ions and remove them - "FastIsThermal": False, # Do not remove fast, keep their diluiton effect but make them thermal - "sameDensityGradients": False, # Make all ion density gradients equal to electrons + "remove_fast": False, # Automatically detect which are fast ions and remove them + "thermalize_fast": False, # Do not remove fast, keep their diluiton effect but make them thermal + "enforce_same_aLn": False, # Make all ion density gradients equal to electrons "groupQIONE": False, - "ensurePostiveGamma": False, + "ensure_positive_Gamma": False, "ensureMachNumber": None, } @@ -166,7 +166,7 @@ def __init__( "applyCorrections": { "Ti_thermals": True, # Keep all thermal ion temperatures equal to the main Ti "ni_thermals": True, # Adjust for quasineutrality by modifying the thermal ion densities together with ne - "recompute_ptot": True, # Recompute PTOT to insert in input file each time + "recalculate_ptot": True, # Recompute PTOT to insert in input file each time "Tfast_ratio": False, # Keep the ratio of Tfast/Te constant throughout the Te evolution "ensureMachNumber": None, # Change w0 to match this Mach number when Ti varies }, @@ -205,7 +205,7 @@ def __init__( "target_evaluator": target_evaluator, "target_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) - "includeFastInQi": False, # If True, and fast ions have been included, in seprateNEO, sum fast + "Qi_includes_fast": False, # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) "portals_transformation_variables": portals_transformation_variables, # Physics-informed parameters to fit surrogates "portals_transformation_variables_trace": portals_transformation_variables_trace, # Physics-informed parameters to fit surrogates for trace impurities "Qi_criterion_stable": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 5570e8f9..765b825e 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -159,7 +159,7 @@ def prep_metrics(self, ilast=None): self.runWithImpurity = self.powerstate.impurityPosition if "nZ" in self.ProfilesPredicted else None self.runWithRotation = "w0" in self.ProfilesPredicted - self.includeFast = self.PORTALSparameters["includeFastInQi"] + self.includeFast = self.PORTALSparameters["Qi_includes_fast"] self.forceZeroParticleFlux = self.PORTALSparameters["forceZeroParticleFlux"] # Profiles and tgyro results diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index daa2a780..a158c974 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -83,10 +83,10 @@ def initializeProblem( if ( len(INITparameters["removeIons"]) > 0 - or INITparameters["removeFast"] + or INITparameters["remove_fast"] or INITparameters["quasineutrality"] - or INITparameters["sameDensityGradients"] - or INITparameters["recompute_ptot"] + or INITparameters["enforce_same_aLn"] + or INITparameters["recalculate_ptot"] ): profiles.correct(options=INITparameters) @@ -129,7 +129,7 @@ def initializeProblem( if transport_evaluator_options is None: transport_evaluator_options = { "launchMODELviaSlurm": portals_fun.PORTALSparameters["launchEvaluationsAsSlurmJobs"], - "includeFastInQi": portals_fun.PORTALSparameters["includeFastInQi"], + "Qi_includes_fast": portals_fun.PORTALSparameters["Qi_includes_fast"], "TurbulentExchange": portals_fun.PORTALSparameters["surrogateForTurbExch"], "profiles_postprocessing_fun": portals_fun.PORTALSparameters["profiles_postprocessing_fun"], "UseFineGridTargets": portals_fun.PORTALSparameters["fineTargetsResolution"], diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 3b6c61b7..8643170b 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -45,11 +45,13 @@ def __init__( if evolution_options is None: evolution_options = {} + if transport_options is None: transport_options = { "transport_evaluator": None, "transport_evaluator_options": {} } + if target_options is None: target_options = { "target_evaluator": targets_analytic.analytical_model, @@ -58,6 +60,7 @@ def __init__( "target_evaluator_method": "powerstate" }, } + if tensor_options is None: tensor_options = { "dtype": torch.double, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 50fcf364..bb547332 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -33,9 +33,11 @@ def _evaluate_tglf(self): # ------------------------------------------------------------------------------------------------------------------------ transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] - - MODELparameters = transport_evaluator_options.get("MODELparameters",None) - includeFast = transport_evaluator_options.get("includeFastInQi",False) + + TGLFsettings = transport_evaluator_options["MODELparameters"]["transport_model"]["TGLFsettings"] + extraOptions = transport_evaluator_options["MODELparameters"]["transport_model"]["extraOptionsTGLF"] + + Qi_includes_fast = transport_evaluator_options.get("Qi_includes_fast",False) launchMODELviaSlurm = transport_evaluator_options.get("launchMODELviaSlurm", False) cold_start = transport_evaluator_options.get("cold_start", False) provideTurbulentExchange = transport_evaluator_options.get("TurbulentExchange", False) @@ -44,21 +46,21 @@ def _evaluate_tglf(self): cores_per_tglf_instance = transport_evaluator_options.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) - impurityPosition = self.powerstate.impurityPosition_transport #transport_evaluator_options.get("impurityPosition", 1) + impurityPosition = self.powerstate.impurityPosition_transport # ------------------------------------------------------------------------------------------------------------------------ # Prepare TGLF object # ------------------------------------------------------------------------------------------------------------------------ - RadiisToRun = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] + rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] - tglf = TGLFtools.TGLF(rhos=RadiisToRun) + tglf = TGLFtools.TGLF(rhos=rho_locations) _ = tglf.prep_direct_tglf( self.folder, cold_start = cold_start, - onlyThermal_TGYRO = not includeFast, - recalculatePTOT = False, # Use what's in the input.gacode, which is what PORTALS TGYRO does + remove_fast = False, # Use what's in the input.gacode, the removal should happen at initialization + recalculate_ptot = False, # Use what's in the input.gacode, which is what PORTALS TGYRO does inputgacode=self.powerstate.profiles_transport, ) @@ -72,8 +74,8 @@ def _evaluate_tglf(self): tglf.run( 'base', - TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], - extraOptions= MODELparameters["transport_model"]["extraOptionsTGLF"], + TGLFsettings=TGLFsettings, + extraOptions=extraOptions, ApplyCorrections=False, launchSlurm= launchMODELviaSlurm, cold_start= cold_start, @@ -90,12 +92,20 @@ def _evaluate_tglf(self): tglf.read(label='base',require_all_files=False) - Qe = [tglf.results['base']['TGLFout'][i].Qe_unn for i in range(len(RadiisToRun))] - Qi = [tglf.results['base']['TGLFout'][i].Qi_unn for i in range(len(RadiisToRun))] - Ge = [tglf.results['base']['TGLFout'][i].Ge_unn for i in range(len(RadiisToRun))] - GZ = [tglf.results['base']['TGLFout'][i].GiAll_unn[impurityPosition] for i in range(len(RadiisToRun))] - Mt = [tglf.results['base']['TGLFout'][i].Mt_unn for i in range(len(RadiisToRun))] - S = [tglf.results['base']['TGLFout'][i].Se_unn for i in range(len(RadiisToRun))] + Qe = np.array([tglf.results['base']['TGLFout'][i].Qe_unn for i in range(len(rho_locations))]) + Qi = np.array([tglf.results['base']['TGLFout'][i].Qi_unn for i in range(len(rho_locations))]) + Ge = np.array([tglf.results['base']['TGLFout'][i].Ge_unn for i in range(len(rho_locations))]) + GZ = np.array([tglf.results['base']['TGLFout'][i].GiAll_unn[impurityPosition] for i in range(len(rho_locations))]) + Mt = np.array([tglf.results['base']['TGLFout'][i].Mt_unn for i in range(len(rho_locations))]) + S = np.array([tglf.results['base']['TGLFout'][i].Se_unn for i in range(len(rho_locations))]) + + if Qi_includes_fast: + + Qifast = [tglf.results['base']['TGLFout'][i].Qifast_unn for i in range(len(rho_locations))] + + if Qifast.sum() != 0.0: + print(f"\t- Qi includes fast ions, adding their contribution") + Qi += Qifast Flux_mean = np.array([Qe, Qi, Ge, GZ, Mt, S]) Flux_std = abs(Flux_mean)*percentError[0]/100.0 @@ -106,17 +116,37 @@ def _evaluate_tglf(self): Flux_base, Flux_mean, Flux_std = _run_tglf_uncertainty_model( tglf, - RadiisToRun, + rho_locations, self.powerstate.ProfilesPredicted, - TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], - extraOptionsTGLF=MODELparameters["transport_model"]["extraOptionsTGLF"], + TGLFsettings=TGLFsettings, + extraOptionsTGLF=extraOptions, impurityPosition=impurityPosition, delta = use_tglf_scan_trick, cold_start=cold_start, extra_name=self.name, cores_per_tglf_instance=cores_per_tglf_instance, launchMODELviaSlurm=launchMODELviaSlurm, + Qi_includes_fast=Qi_includes_fast, ) + + for i in range(len(tglf.profiles.Species)): + gacode_type = tglf.profiles.Species[i]['S'] + for rho in rho_locations: + tglf_type = tglf.inputsTGLF[0.25].ions_info[i+2]['type'] + + if gacode_type[:5] != tglf_type[:5]: + print(f"\t- For location {rho=:.2f}, ion specie #{i+1} ({tglf.profiles.Species[i]['N']}) is considered '{gacode_type}' by gacode but '{tglf_type}' by TGLF. Make sure this is consistent with your use case", typeMsg="w") + + if tglf_type == 'fast': + + if Qi_includes_fast: + print(f"\t\t\t* The fast ion considered by TGLF was summed into the Qi", typeMsg="i") + else: + print(f"\t\t\t* The fast ion considered by TGLF was NOT summed into the Qi", typeMsg="i") + + else: + + print(f"\t\t\t* The thermal ion considered by TGLF was summed into the Qi", typeMsg="i") # ------------------------------------------------------------------------------------------------------------------------ # Pass the information to POWERSTATE @@ -232,7 +262,7 @@ def _profiles_to_store(self): def _run_tglf_uncertainty_model( tglf, - RadiisToRun, + rho_locations, ProfilesPredicted, TGLFsettings=None, extraOptionsTGLF=None, @@ -244,6 +274,7 @@ def _run_tglf_uncertainty_model( remove_folders_out = False, cores_per_tglf_instance = 4, # e.g. 4 core per radius, since this is going to launch ~ Nr=5 x (Nv=6 x Nd=2 + 1) = 65 TGLFs at once launchMODELviaSlurm=False, + Qi_includes_fast=False, ): print(f"\t- Running TGLF standalone scans ({delta = }) to determine relative errors") @@ -275,10 +306,10 @@ def _run_tglf_uncertainty_model( name = 'turb_drives' - tglf.rhos = RadiisToRun # To avoid the case in which TGYRO was run with an extra rho point + tglf.rhos = rho_locations # To avoid the case in which TGYRO was run with an extra rho point # Estimate job minutes based on cases and cores (mostly IO I think at this moment, otherwise it should be independent on cases) - num_cases = len(RadiisToRun) * len(variables_to_scan) * len(relative_scan) + num_cases = len(rho_locations) * len(variables_to_scan) * len(relative_scan) if cores_per_tglf_instance == 1: minutes = 10 * (num_cases / 60) # Ad-hoc formula else: @@ -313,12 +344,13 @@ def _run_tglf_uncertainty_model( if remove_folders_out: IOtools.shutil_rmtree(tglf.FolderGACODE) - Qe = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - Qi = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - Ge = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - GZ = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - Mt = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - S = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + Qe = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + Qi = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + Qifast = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + Ge = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + GZ = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + Mt = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + S = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) cont = 0 for vari in variables_to_scan: @@ -326,11 +358,16 @@ def _run_tglf_uncertainty_model( Qe[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qe'] Qi[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qi'] + Qifast[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qifast'] Ge[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Ge'] GZ[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Gi'] Mt[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Mt'] S[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['S'] cont += jump + + if Qi_includes_fast: + print(f"\t- Qi includes fast ions, adding their contribution") + Qi += Qifast # Calculate the standard deviation of the scans, that's going to be the reported stds diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 1b27e7ab..0768e0ca 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -30,7 +30,7 @@ def _evaluate_tglf_neo(self): transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] MODELparameters = transport_evaluator_options.get("MODELparameters",None) - includeFast = transport_evaluator_options.get("includeFastInQi",False) + Qi_includes_fast = transport_evaluator_options.get("Qi_includes_fast",False) launchMODELviaSlurm = transport_evaluator_options.get("launchMODELviaSlurm", False) cold_start = transport_evaluator_options.get("cold_start", False) provideTurbulentExchange = transport_evaluator_options.get("TurbulentExchange", False) @@ -45,7 +45,7 @@ def _evaluate_tglf_neo(self): # tglf_neo_original: Run TGYRO workflow - TGLF + NEO in subfolder tglf_neo_original (original as in... without stds or merging) # ------------------------------------------------------------------------------------------------------------------------ - RadiisToRun = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] + rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] tgyro = TGYROtools.TGYRO(cdf=dummyCDF(self.folder, self.folder)) tgyro.prep(self.folder, profilesclass_custom=self.powerstate.profiles_transport) @@ -59,7 +59,7 @@ def _evaluate_tglf_neo(self): subFolderTGYRO="tglf_neo_original", cold_start=cold_start, forceIfcold_start=True, - special_radii=RadiisToRun, + special_radii=rho_locations, iterations=0, PredictionSet=[ int("te" in self.powerstate.ProfilesPredicted), @@ -92,12 +92,12 @@ def _evaluate_tglf_neo(self): curateTGYROfiles( tgyro, "tglf_neo_original", - RadiisToRun, + rho_locations, self.powerstate.ProfilesPredicted, self.folder / "tglf_neo", percentError, impurityPosition=impurityPosition, - includeFast=includeFast, + Qi_includes_fast=Qi_includes_fast, provideTurbulentExchange=provideTurbulentExchange, use_tglf_scan_trick = use_tglf_scan_trick, cold_start=cold_start, @@ -113,11 +113,11 @@ def _evaluate_tglf_neo(self): # $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ # from mitim_tools.gacode_tools import TGLFtools - # tglf = TGLFtools.TGLF(rhos=RadiisToRun) + # tglf = TGLFtools.TGLF(rhos=rho_locations) # _ = tglf.prep( # self.folder / 'stds', # inputgacode=self.file_profs, - # recalculatePTOT=False, # Use what's in the input.gacode, which is what PORTALS TGYRO does + # recalculate_ptot=False, # Use what's in the input.gacode, which is what PORTALS TGYRO does # cold_start=cold_start) # tglf.run( @@ -141,7 +141,7 @@ def _postprocess_results(self, tgyro, label): transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] - includeFast = transport_evaluator_options.get("includeFastInQi",False) + Qi_includes_fast = transport_evaluator_options.get("Qi_includes_fast",False) UseFineGridTargets = transport_evaluator_options.get("UseFineGridTargets", False) provideTurbulentExchange = transport_evaluator_options.get("TurbulentExchange", False) OriginalFimp = transport_evaluator_options.get("OriginalFimp", 1.0) @@ -154,7 +154,7 @@ def _postprocess_results(self, tgyro, label): self.powerstate = tgyro_to_powerstate( tgyro.results[label], self.powerstate, - includeFast=includeFast, + Qi_includes_fast=Qi_includes_fast, impurityPosition=impurityPosition, UseFineGridTargets=UseFineGridTargets, OriginalFimp=OriginalFimp, @@ -196,10 +196,10 @@ def _profiles_to_store(self): def tglf_scan_trick( tglf, - RadiisToRun, + rho_locations, ProfilesPredicted, impurityPosition=1, - includeFast=False, + Qi_includes_fast=False, delta=0.02, minimum_abs_gradient=0.005, # This is 0.5% of aLx=1.0, to avoid extremely small scans when, for example, having aLn ~ 0.0 cold_start=False, @@ -237,10 +237,10 @@ def tglf_scan_trick( name = 'turb_drives' - tglf.rhos = RadiisToRun # To avoid the case in which TGYRO was run with an extra rho point + tglf.rhos = rho_locations # To avoid the case in which TGYRO was run with an extra rho point # Estimate job minutes based on cases and cores (mostly IO I think at this moment, otherwise it should be independent on cases) - num_cases = len(RadiisToRun) * len(variables_to_scan) * len(relative_scan) + num_cases = len(rho_locations) * len(variables_to_scan) * len(relative_scan) if cores_per_tglf_instance == 1: minutes = 10 * (num_cases / 60) # Ad-hoc formula else: @@ -273,12 +273,13 @@ def tglf_scan_trick( if remove_folders_out: IOtools.shutil_rmtree(tglf.FolderGACODE) - Qe = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - Qi = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - Ge = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - GZ = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - Mt = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) - S = np.zeros((len(RadiisToRun), len(variables_to_scan)*len(relative_scan)+1 )) + Qe = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + Qi = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + Qifast = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + Ge = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + GZ = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + Mt = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + S = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) cont = 0 for vari in variables_to_scan: @@ -286,11 +287,16 @@ def tglf_scan_trick( Qe[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qe'] Qi[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qi'] + Qifast[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qifast'] Ge[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Ge'] GZ[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Gi'] Mt[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Mt'] S[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['S'] cont += jump + + if Qi_includes_fast: + print(f"\t- Qi includes fast ions, adding their contribution") + Qi += Qifast # Calculate the standard deviation of the scans, that's going to be the reported stds @@ -329,13 +335,13 @@ def calculate_mean_std(Q): def curateTGYROfiles( tgyroObject, label, - RadiisToRun, + rho_locations, ProfilesPredicted, folder, percentError, provideTurbulentExchange=False, impurityPosition=1, - includeFast=False, + Qi_includes_fast=False, use_tglf_scan_trick=None, cold_start=False, extra_name="", @@ -351,7 +357,7 @@ def curateTGYROfiles( # Grab fluxes from TGYRO Qe = tgyro.Qe_sim_turb[0, 1:] - Qi = tgyro.QiIons_sim_turb[0, 1:] if includeFast else tgyro.QiIons_sim_turb_thr[0, 1:] + Qi = tgyro.QiIons_sim_turb[0, 1:] if Qi_includes_fast else tgyro.QiIons_sim_turb_thr[0, 1:] Ge = tgyro.Ge_sim_turb[0, 1:] GZ = tgyro.Gi_sim_turb[impurityPosition, 0, 1:] Mt = tgyro.Mt_sim_turb[0, 1:] @@ -366,10 +372,10 @@ def curateTGYROfiles( # Run TGLF scan trick Flux_base, Flux_mean, Flux_std = tglf_scan_trick( tglfObject, - RadiisToRun, + rho_locations, ProfilesPredicted, impurityPosition=impurityPosition, - includeFast=includeFast, + Qi_includes_fast=Qi_includes_fast, delta = use_tglf_scan_trick, cold_start=cold_start, extra_name=extra_name, @@ -386,7 +392,7 @@ def curateTGYROfiles( # Grab fluxes from TGYRO Qe_tgyro = tgyro.Qe_sim_turb[0, 1:] - Qi_tgyro = tgyro.QiIons_sim_turb[0, 1:] if includeFast else tgyro.QiIons_sim_turb_thr[0, 1:] + Qi_tgyro = tgyro.QiIons_sim_turb[0, 1:] if Qi_includes_fast else tgyro.QiIons_sim_turb_thr[0, 1:] Ge_tgyro = tgyro.Ge_sim_turb[0, 1:] GZ_tgyro = tgyro.Gi_sim_turb[impurityPosition, 0, 1:] Mt_tgyro = tgyro.Mt_sim_turb[0, 1:] @@ -433,7 +439,7 @@ def curateTGYROfiles( # If simply a percentage error provided # -------------------------------------------------------------- - relativeErrorTGLF = [percentError[0] / 100.0]*len(RadiisToRun) + relativeErrorTGLF = [percentError[0] / 100.0]*len(rho_locations) QeE = abs(Qe) * relativeErrorTGLF QiE = abs(Qi) * relativeErrorTGLF @@ -447,7 +453,7 @@ def curateTGYROfiles( # ************************************************************************************************************************** Qe_tr_neoc = tgyro.Qe_sim_neo[0, 1:] - if includeFast: + if Qi_includes_fast: Qi_tr_neoc = tgyro.QiIons_sim_neo[0, 1:] else: Qi_tr_neoc = tgyro.QiIons_sim_neo_thr[0, 1:] @@ -456,7 +462,7 @@ def curateTGYROfiles( Mt_tr_neoc = tgyro.Mt_sim_neo[0, 1:] Qe_tr_neocE = abs(tgyro.Qe_sim_neo[0, 1:]) * relativeErrorNEO - if includeFast: + if Qi_includes_fast: Qi_tr_neocE = abs(tgyro.QiIons_sim_neo[0, 1:]) * relativeErrorNEO else: Qi_tr_neocE = abs(tgyro.QiIons_sim_neo_thr[0, 1:]) * relativeErrorNEO @@ -889,7 +895,7 @@ def defineReferenceFluxes( def tgyro_to_powerstate(TGYROresults, powerstate, forceZeroParticleFlux=False, - includeFast=False, + Qi_includes_fast=False, impurityPosition=1, UseFineGridTargets=False, OriginalFimp=1.0, @@ -928,7 +934,7 @@ def tgyro_to_powerstate(TGYROresults, # *********** Ion Energy Fluxes # ********************************** - if includeFast: + if Qi_includes_fast: powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb[:, :nr]).to(powerstate.dfT) powerstate.plasma["QiMWm2_tr_neoc"] = torch.Tensor(TGYROresults.QiIons_sim_neo[:, :nr]).to(powerstate.dfT) diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 26a7c376..f99b905e 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -57,11 +57,11 @@ def calculator( "Tfast_ratio": False, "Ti_thermals": True, "ni_thermals": True, - "recompute_ptot": False, + "recalculate_ptot": False, }, "transport_model": {"TGLFsettings": 5, "extraOptionsTGLF": {}}, }, - "includeFastInQi": False, + "Qi_includes_fast": False, }, }, ) @@ -119,7 +119,7 @@ def calculator( "Tfast_ratio": False, "Ti_thermals": False, "ni_thermals": False, - "recompute_ptot": False, + "recalculate_ptot": False, "ensureMachNumber": None, }, insert_highres_powers=True, diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index f30d5a81..60ccd39d 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -259,7 +259,7 @@ def powerstate_to_gacode( Tfast_ratio = postprocess_input_gacode.get("Tfast_ratio", True) Ti_thermals = postprocess_input_gacode.get("Ti_thermals", True) ni_thermals = postprocess_input_gacode.get("ni_thermals", True) - recompute_ptot = postprocess_input_gacode.get("recompute_ptot", True) + recalculate_ptot = postprocess_input_gacode.get("recalculate_ptot", True) ensureMachNumber = postprocess_input_gacode.get("ensureMachNumber", None) # ------------------------------------------------------------------------------------------ @@ -334,10 +334,10 @@ def powerstate_to_gacode( # Recalculate and change ptot to make it consistent? # ------------------------------------------------------------------------------------------ - if rederive or recompute_ptot: + if rederive or recalculate_ptot: profiles.derive_quantities(rederiveGeometry=False) - if recompute_ptot: + if recalculate_ptot: profiles.selfconsistentPTOT() if debugPlot: diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 516bd2e7..d0bf5057 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -217,8 +217,8 @@ def prep( self, FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) cold_start=False, # If True, do not use what it potentially inside the folder, run again - onlyThermal_TGYRO=False, # Ignore fast particles in TGYRO - recalculatePTOT=True, # Recalculate PTOT in TGYRO + remove_fast=False, # Ignore fast particles in TGYRO + recalculate_ptot=True, # Recalculate PTOT in TGYRO cdf_open=None, # Grab normalizations from CDF file that is open as transp_output class inputgacode=None, # *NOTE BELOW* specificInputs=None, # *NOTE BELOW* @@ -281,9 +281,7 @@ def prep( inp = TGLFinput(fii) exists = exists and not inp.onlyControl else: - print( - f"\t\t- Running scans because it does not exist file {IOtools.clipstr(fii)}" - ) + print(f"\t\t- Running scans because it does not exist file {IOtools.clipstr(fii)}") exists = False if exists: print( @@ -303,8 +301,8 @@ def prep( self.tgyro_results = self.tgyro.run_tglf_scan( rhos=self.rhos, cold_start=not exists, - onlyThermal=onlyThermal_TGYRO, - recalculatePTOT=recalculatePTOT, + onlyThermal=remove_fast, + recalculate_ptot=recalculate_ptot, donotrun=donotrun, ) @@ -362,8 +360,8 @@ def prep_direct_tglf( self, FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) cold_start=False, # If True, do not use what it potentially inside the folder, run again - onlyThermal_TGYRO=False, # Ignore fast particles in TGYRO - recalculatePTOT=True, # Recalculate PTOT in TGYRO + remove_fast=False, # Ignore fast particles in TGYRO + recalculate_ptot=True, # Recalculate PTOT in TGYRO cdf_open=None, # Grab normalizations from CDF file that is open as transp_output class inputgacode=None, # *NOTE BELOW* specificInputs=None, # *NOTE BELOW* @@ -412,7 +410,7 @@ def prep_direct_tglf( self.profiles.derive_quantities(mi_ref=md_u) - self.profiles.correct(options={'recompute_ptot':recalculatePTOT,'removeFast':onlyThermal_TGYRO}) + self.profiles.correct(options={'recalculate_ptot':recalculate_ptot,'remove_fast':remove_fast}) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize by preparing a tgyro class and running for -1 iterations @@ -473,8 +471,6 @@ def prep_direct_tglf( return cdf - - def prep_from_tglf( self, FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) @@ -2415,6 +2411,7 @@ def readScan( Qe_gb, Qi_gb, Ge_gb, Gi_gb, Mt_gb, S_gb = [],[],[],[],[],[] ky, g, f, eta1, eta2, itg, tem, etg = [],[],[],[],[],[],[],[] etalow_g, etalow_f, etalow_k = [], [], [] + Qifast, Qifast_gb = [],[] cont = 0 for ikey in self.results: isThisTheRightReadResults = (subFolderTGLF in ikey) and (variable== "_".join(ikey.split("_")[:-1]).split(subFolderTGLF + "_")[-1]) @@ -2427,6 +2424,7 @@ def readScan( Qe_gb0, Qi_gb0, Ge_gb0, Gi_gb0, Mt_gb0, S_gb0 = [],[],[],[],[],[] ky0, g0, f0, eta10, eta20, itg0, tem0, etg0 = [],[],[],[],[],[],[],[] etalow_g0, etalow_f0, etalow_k0 = [], [], [] + Qifast0, Qifast_gb0 = [],[] for irho_cont in range(len(self.rhos)): irho = np.where(self.results[ikey]["x"] == self.rhos[irho_cont])[0][0] @@ -2434,6 +2432,7 @@ def readScan( x0.append(self.results[ikey]["parsed"][irho][variable]) Qe_gb0.append(self.results[ikey]["TGLFout"][irho].Qe) Qi_gb0.append(self.results[ikey]["TGLFout"][irho].Qi) + Qifast_gb0.append(self.results[ikey]["TGLFout"][irho].Qifast) Ge_gb0.append(self.results[ikey]["TGLFout"][irho].Ge) Gi_gb0.append(self.results[ikey]["TGLFout"][irho].GiAll[self.positionIon_scan - 2]) Mt_gb0.append(self.results[ikey]["TGLFout"][irho].Mt) @@ -2453,6 +2452,7 @@ def readScan( if self.results[ikey]["TGLFout"][irho].unnormalization_successful: Qe0.append(self.results[ikey]["TGLFout"][irho].Qe_unn) Qi0.append(self.results[ikey]["TGLFout"][irho].Qi_unn) + Qifast0.append(self.results[ikey]["TGLFout"][irho].Qifast_unn) Ge0.append(self.results[ikey]["TGLFout"][irho].Ge_unn) Gi0.append(self.results[ikey]["TGLFout"][irho].GiAll_unn[self.positionIon_scan - 2]) Mt0.append(self.results[ikey]["TGLFout"][irho].Mt_unn) @@ -2463,9 +2463,11 @@ def readScan( x.append(x0) Qe.append(Qe0) Qi.append(Qi0) + Qifast.append(Qifast0) Ge.append(Ge0) Qe_gb.append(Qe_gb0) Qi_gb.append(Qi_gb0) + Qifast_gb.append(Qifast_gb0) Ge_gb.append(Ge_gb0) Gi_gb.append(Gi_gb0) Gi.append(Gi0) @@ -2491,12 +2493,14 @@ def readScan( self.scans[label]["xV"] = np.atleast_2d(np.transpose(x)) self.scans[label]["Qe_gb"] = np.atleast_2d(np.transpose(Qe_gb)) self.scans[label]["Qi_gb"] = np.atleast_2d(np.transpose(Qi_gb)) + self.scans[label]["Qifast_gb"] = np.atleast_2d(np.transpose(Qifast_gb)) self.scans[label]["Ge_gb"] = np.atleast_2d(np.transpose(Ge_gb)) self.scans[label]["Gi_gb"] = np.atleast_2d(np.transpose(Gi_gb)) self.scans[label]["Mt_gb"] = np.atleast_2d(np.transpose(Mt_gb)) self.scans[label]["S_gb"] = np.atleast_2d(np.transpose(S_gb)) self.scans[label]["Qe"] = np.atleast_2d(np.transpose(Qe)) self.scans[label]["Qi"] = np.atleast_2d(np.transpose(Qi)) + self.scans[label]["Qifast"] = np.atleast_2d(np.transpose(Qifast)) self.scans[label]["Ge"] = np.atleast_2d(np.transpose(Ge)) self.scans[label]["Gi"] = np.atleast_2d(np.transpose(Gi)) self.scans[label]["Mt"] = np.atleast_2d(np.transpose(Mt)) @@ -3782,7 +3786,7 @@ def changeANDwrite_TGLF( if ApplyCorrections: print("\t- Applying corrections") inputTGLF_rho.removeLowDensitySpecie() - inputTGLF_rho.removeFast() + inputTGLF_rho.remove_fast() # Ensure that plasma to run is quasineutral if Quasineutral: @@ -3910,6 +3914,7 @@ def processSpecies(self, MinMultiplierToBeFast=2.0): } thermal_indeces = [1, 2] + fast_indeces = [] for i in range(len(self.species) - 2): TiTe = self.species[3 + i]["TAUS"] if TiTe < thrTemperatureRatio: @@ -3917,17 +3922,17 @@ def processSpecies(self, MinMultiplierToBeFast=2.0): thermal_indeces.append(3 + i) else: self.ions_info[3 + i] = {"type": "fast"} + fast_indeces.append(3 + i) self.ions_info["thermal_list"] = thermal_indeces - self.ions_info["thermal_list_extras"] = thermal_indeces[ - 2: - ] # remove electrons and mains - + self.ions_info["thermal_list_extras"] = thermal_indeces[2:] # remove electrons and mains + + self.ions_info["fast_list"] = fast_indeces + self.onlyControl = False + else: - print( - "\t- No species in this input.tglf (it is either a controls-only file or there was a problem generating it)" - ) + print("\t- No species in this input.tglf (it is either a controls-only file or there was a problem generating it)") self.onlyControl = True def isThePlasmaDT(self): @@ -3943,7 +3948,7 @@ def isThePlasmaDT(self): return np.abs(mrat - 1.5) < 0.01 - def removeFast(self): + def remove_fast(self): self.processSpecies() i = 1 while i <= len(self.species): @@ -4483,12 +4488,10 @@ def read(self,require_all_files=True): # Ions to include? e.g. IncludeExtraIonsInQi = [2,3,4] -> This will sum to ion 1 # -------------------------------------------------------------------------------- - IncludeExtraIonsInQi = ( - [i - 1 for i in self.inputclass.ions_info["thermal_list_extras"]] - if self.inputclass is not None - else [] - ) + IncludeExtraIonsInQi = [i - 1 for i in self.inputclass.ions_info["thermal_list_extras"]] if self.inputclass is not None else [] self.ions_included = (1,) + tuple(IncludeExtraIonsInQi) + + self.fast_included = tuple(self.inputclass.ions_info["fast_list"]) if self.inputclass is not None else () # ------------------------------------------------------------------------ # Fluxes @@ -4515,6 +4518,13 @@ def read(self,require_all_files=True): print(f"\t\t- For Qi, summing contributions from ions {self.ions_included} (#0 is e-)",typeMsg="i",) self.Gi = data[0, self.ions_included].sum() self.Qi = data[1, self.ions_included].sum() + + if len(self.fast_included)>0: + print(f"\t\t- For Qifast, summing contributions from fast ions {self.fast_included} (#0 is e-)",typeMsg="i",) + self.Qifast = data[1, self.fast_included].sum() + else: + print(f"\t\t- No fast ions included",typeMsg="i",) + self.Qifast = 0.0 signMt = - self.inputclass.plasma['SIGN_IT'] # Following tgyro_flux.f90 print(f"\t\t- Sign of Mt given by toroidal current direction (SIGN_IT={-signMt}): {signMt}",typeMsg="i",) @@ -4997,6 +5007,8 @@ def unnormalize(self, normalization, rho=None, convolution_fun_fluct=None, facto self.Qe_unn = self.Qe * q_gb[ir] self.Qi_unn = self.Qi * q_gb[ir] + self.Qifast_unn = self.Qifast * q_gb[ir] + self.QiAll_unn = self.QiAll * q_gb[ir] self.Ge_unn = self.Ge * g_gb[ir] self.GiAll_unn = self.GiAll * g_gb[ir] diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 556c0de3..45305dbf 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -707,7 +707,7 @@ def run_tglf_scan( cold_start=False, label="tgyro1", donotrun=False, - recalculatePTOT=True, + recalculate_ptot=True, ): """ onlyThermal will remove from the TGYRO run the fast species, so the resulting input.tglf files will not have @@ -739,7 +739,7 @@ def run_tglf_scan( "onlyThermal": onlyThermal, "quasineutrality": quasineutrality, "neoclassical": 0, # Do not run or check NEOTGYRO canno - "PtotType": int(not recalculatePTOT), # Recalculate Ptot or use what's there + "PtotType": int(not recalculate_ptot), # Recalculate Ptot or use what's there } # ------------------------------------------------------------ diff --git a/src/mitim_tools/gacode_tools/scripts/run_tglf.py b/src/mitim_tools/gacode_tools/scripts/run_tglf.py index 0e5f3c55..b146cef5 100644 --- a/src/mitim_tools/gacode_tools/scripts/run_tglf.py +++ b/src/mitim_tools/gacode_tools/scripts/run_tglf.py @@ -28,9 +28,7 @@ def main(): 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, default=False, action="store_true") - parser.add_argument( - "--cold_start", "-r", required=False, default=False, action="store_true" - ) + parser.add_argument("--cold_start", "-r", required=False, default=False, action="store_true") args = parser.parse_args() diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 742b8628..169ddc92 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1613,15 +1613,15 @@ def correct(self, options={}, write=False, new_file=None): if name= T D LUMPED, and I want to eliminate D, removeIons = [2] """ - recompute_ptot = options.get("recompute_ptot", True) # Only done by default + recalculate_ptot = options.get("recalculate_ptot", True) # Only done by default removeIons = options.get("removeIons", []) - removeFast = options.get("removeFast", False) + remove_fast = options.get("remove_fast", False) quasineutrality = options.get("quasineutrality", False) - sameDensityGradients = options.get("sameDensityGradients", False) + enforce_same_aLn = options.get("enforce_same_aLn", False) groupQIONE = options.get("groupQIONE", False) - ensurePostiveGamma = options.get("ensurePostiveGamma", False) + ensure_positive_Gamma = options.get("ensure_positive_Gamma", False) ensureMachNumber = options.get("ensureMachNumber", None) - FastIsThermal = options.get("FastIsThermal", False) + thermalize_fast = options.get("thermalize_fast", False) print("\t- Custom correction of input.gacode file has been requested") @@ -1634,7 +1634,7 @@ def correct(self, options={}, write=False, new_file=None): self.remove(removeIons) # Remove fast - if removeFast: + if remove_fast: ions_fast = [] for sp in range(len(self.Species)): if self.Species[sp]["S"] != "therm": @@ -1645,7 +1645,7 @@ def correct(self, options={}, write=False, new_file=None): ) self.remove(ions_fast) # Fast as thermal - elif FastIsThermal: + elif thermalize_fast: self.make_fast_ions_thermal() # Correct LUMPED @@ -1673,7 +1673,7 @@ def correct(self, options={}, write=False, new_file=None): self.profiles["qione(MW/m^3)"] = self.profiles["qione(MW/m^3)"] * 0.0 # Make all thermal ions have the same gradient as the electron density, by keeping volume average constant - if sameDensityGradients: + if enforce_same_aLn: self.enforce_same_density_gradients() # Enforce quasineutrality @@ -1683,12 +1683,12 @@ def correct(self, options={}, write=False, new_file=None): print(f"\t\t\t* Quasineutrality error = {self.derived['QN_Error']:.1e}") # Recompute ptot - if recompute_ptot: + if recalculate_ptot: self.derive_quantities(rederiveGeometry=False) self.selfconsistentPTOT() # If I don't trust the negative particle flux in the core that comes from TRANSP... - if ensurePostiveGamma: + if ensure_positive_Gamma: print("\t\t- Making particle flux always positive", typeMsg="i") self.profiles["qpar_beam(1/m^3/s)"] = self.profiles["qpar_beam(1/m^3/s)"].clip(0) self.profiles["qpar_wall(1/m^3/s)"] = self.profiles["qpar_wall(1/m^3/s)"].clip(0) diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index cbf3421f..1529f7af 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -14041,13 +14041,13 @@ def runTGLFstandalone( cold_startPreparation=False, plotCompare=True, extraflag="", - onlyThermal_TGYRO=False, + remove_fast=False, forceIfcold_start=True, **kwargs_TGLFrun, ): """ Note: If this plasma had fast paricles but not at the time I'm running TGLF, then it will fail if I - set onlyThermal_TGYRO=False because at that time the particles are zero + set remove_fast=False because at that time the particles are zero """ if time is None: @@ -14078,7 +14078,7 @@ def runTGLFstandalone( cdf = self.TGLFstd[nameF].prep( folderGACODE, cold_start=cold_startPreparation, - onlyThermal_TGYRO=onlyThermal_TGYRO, + remove_fast=remove_fast, cdf_open=self, forceIfcold_start=forceIfcold_start, ) diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index ca8f857f..387454f2 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -75,7 +75,7 @@ } }, "INITparameters": { - "FastIsThermal": true, + "thermalize_fast": true, "quasineutrality": true }, "exploration_ranges": { diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index e6b3a6fe..9daf9cec 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -32,9 +32,9 @@ portals_fun.MODELparameters['ProfilesPredicted'] = ["te", "ti", "ne", "nZ", 'w0'] portals_fun.PORTALSparameters['ImpurityOfInterest'] = 'N' portals_fun.PORTALSparameters['surrogateForTurbExch'] = True -portals_fun.INITparameters["removeFast"] = True +portals_fun.INITparameters["remove_fast"] = True portals_fun.INITparameters["quasineutrality"] = True -portals_fun.INITparameters["sameDensityGradients"] = True +portals_fun.INITparameters["enforce_same_aLn"] = True portals_fun.MODELparameters["transport_model"]["TGLFsettings"] = 2 # Prepare run diff --git a/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index 9d4e0f82..9b6ca805 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -30,7 +30,7 @@ } # Plasma preparation: remove fast species, adjust quasineutrality -portals_fun.INITparameters["removeFast"] = True +portals_fun.INITparameters["remove_fast"] = True portals_fun.INITparameters["quasineutrality"] = True # Stopping criterion 1: 100x improvement in residual From 7391ef4492410b5eccee3af064b378858a31259b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 5 Aug 2025 13:33:26 -0400 Subject: [PATCH 131/385] Naming conventions --- regressions/portals_regressions.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 2 +- src/mitim_modules/portals/PORTALStools.py | 4 ++-- src/mitim_modules/portals/utils/PORTALSinit.py | 4 ++-- src/mitim_modules/portals/utils/PORTALSplot.py | 6 +++--- tests/PORTALS_workflow.py | 2 +- 6 files changed, 10 insertions(+), 10 deletions(-) diff --git a/regressions/portals_regressions.py b/regressions/portals_regressions.py index 9890980f..36ce13d6 100644 --- a/regressions/portals_regressions.py +++ b/regressions/portals_regressions.py @@ -139,7 +139,7 @@ def conditions_regressions(variables): # portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti", "ne",'nZ','w0'] # portals_fun.PORTALSparameters["ImpurityOfInterest"] = 'W' - # portals_fun.PORTALSparameters["surrogateForTurbExch"] = True + # portals_fun.PORTALSparameters["turbulent_exchange_as_surrogate"] = True # portals_fun.prep(inputgacode, folderWork) # mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=False, askQuestions=False) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index d0e98e40..e4c9702b 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -211,7 +211,7 @@ def __init__( "Qi_criterion_stable": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable "percentError_stable": 5.0, # (%) For CGYRO runs, minimum error based on target if case is considered stable "forceZeroParticleFlux": False, # If True, ignore particle flux profile and assume zero for all radii - "surrogateForTurbExch": False, # Run turbulent exchange as surrogate? + "turbulent_exchange_as_surrogate": False, # Run turbulent exchange as surrogate? "profiles_postprocessing_fun": None, # Function to post-process input.gacode only BEFORE passing to transport codes "Pseudo_multipliers": [1.0]*5, # [Qe,Qi,Ge] multipliers to calculate pseudo "ImpurityOfInterest": None, # Impurity to do flux-matching for if nZ enabled (name of first impurity instance AFTER postprocessing), e.g. "W" diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 4c2f66c8..a7a08fb9 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -399,7 +399,7 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added # ------------------------------------------------------------------------- - if PORTALSparameters["surrogateForTurbExch"]: + if PORTALSparameters["turbulent_exchange_as_surrogate"]: QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) else: QieMWm3_tr_turb_integrated = torch.zeros(dfT.shape).to(dfT) @@ -532,7 +532,7 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added # ------------------------------------------------------------------------- - if PORTALSparameters["surrogateForTurbExch"]: + if PORTALSparameters["turbulent_exchange_as_surrogate"]: QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) QieMWm3_tr_turb_integrated_stds = computeTurbExchangeIndividual(var_dict["Qie_tr_turb_stds"], powerstate) else: diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index a158c974..4cf5a56c 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -130,7 +130,7 @@ def initializeProblem( transport_evaluator_options = { "launchMODELviaSlurm": portals_fun.PORTALSparameters["launchEvaluationsAsSlurmJobs"], "Qi_includes_fast": portals_fun.PORTALSparameters["Qi_includes_fast"], - "TurbulentExchange": portals_fun.PORTALSparameters["surrogateForTurbExch"], + "TurbulentExchange": portals_fun.PORTALSparameters["turbulent_exchange_as_surrogate"], "profiles_postprocessing_fun": portals_fun.PORTALSparameters["profiles_postprocessing_fun"], "UseFineGridTargets": portals_fun.PORTALSparameters["fineTargetsResolution"], "OriginalFimp": portals_fun.PORTALSparameters["fImp_orig"], @@ -292,7 +292,7 @@ def initializeProblem( name_objectives.append(f"{var}Res_{i+1}") - if portals_fun.PORTALSparameters["surrogateForTurbExch"]: + if portals_fun.PORTALSparameters["turbulent_exchange_as_surrogate"]: for i in range(len(portals_fun.MODELparameters["RhoLocations"])): ofs.append(f"Qie_tr_turb_{i+1}") diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index 592ebff2..d9f0ebad 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -2011,10 +2011,10 @@ def PORTALSanalyzer_plotModelComparison( if (fig is None) and (axs is None): plt.ion() - fig = plt.figure(figsize=(15, 6 if len(self.ProfilesPredicted)+int(self.PORTALSparameters["surrogateForTurbExch"]) < 4 else 10)) + fig = plt.figure(figsize=(15, 6 if len(self.ProfilesPredicted)+int(self.PORTALSparameters["turbulent_exchange_as_surrogate"]) < 4 else 10)) if axs is None: - if len(self.ProfilesPredicted)+int(self.PORTALSparameters["surrogateForTurbExch"]) < 4: + if len(self.ProfilesPredicted)+int(self.PORTALSparameters["turbulent_exchange_as_surrogate"]) < 4: axs = fig.subplots(ncols=3) else: axs = fig.subplots(ncols=3, nrows=2) @@ -2201,7 +2201,7 @@ def PORTALSanalyzer_plotModelComparison( cont += 1 - if self.PORTALSparameters["surrogateForTurbExch"]: + if self.PORTALSparameters["turbulent_exchange_as_surrogate"]: if UseTGLFfull_x is not None: raise Exception("Turbulent exchange plot not implemented yet") # Sexch diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index 9daf9cec..e3bb56e6 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -31,7 +31,7 @@ portals_fun.MODELparameters["RhoLocations"] = [0.25, 0.45, 0.65, 0.85] portals_fun.MODELparameters['ProfilesPredicted'] = ["te", "ti", "ne", "nZ", 'w0'] portals_fun.PORTALSparameters['ImpurityOfInterest'] = 'N' -portals_fun.PORTALSparameters['surrogateForTurbExch'] = True +portals_fun.PORTALSparameters['turbulent_exchange_as_surrogate'] = True portals_fun.INITparameters["remove_fast"] = True portals_fun.INITparameters["quasineutrality"] = True portals_fun.INITparameters["enforce_same_aLn"] = True From 2e3fedcda7edc8f7bb72bba9ee83e6c95c23fe2b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 5 Aug 2025 14:19:52 -0400 Subject: [PATCH 132/385] Added shat in derived --- src/mitim_tools/plasmastate_tools/MITIMstate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 169ddc92..7180a044 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1059,8 +1059,8 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["mi_ref"], self.derived["B_unit"] ) - s_hat = self.derived["r"]*self._deriv_gacode( np.log(abs(self.profiles["q(-)"])) ) - self.derived['s_q'] = (self.profiles["q(-)"] / self.derived['roa'])**2 * s_hat + self.derived['s_hat'] = self.derived["r"]*self._deriv_gacode( np.log(abs(self.profiles["q(-)"])) ) + self.derived['s_q'] = (self.profiles["q(-)"] / self.derived['roa'])**2 * self.derived['s_hat'] self.derived['s_q'][0] = 0.0 # infinite in first location # Derivate function From 96be11bda6ecb70c189b507aca11adc01eea5698 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 5 Aug 2025 16:22:55 -0400 Subject: [PATCH 133/385] Fix partial import --- src/mitim_tools/plasmastate_tools/MITIMstate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 7180a044..41b831bc 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -7,8 +7,6 @@ from mitim_modules.powertorch.utils import CALCtools from mitim_tools.gacode_tools import NEOtools from mitim_tools.gacode_tools.utils import GACODEdefaults -from mitim_tools.transp_tools import CDFtools -from mitim_tools.transp_tools.utils import TRANSPhelpers from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __version__ @@ -1898,6 +1896,7 @@ def addSawtoothEffectOnOhmic(self, PohTot, mixRadius=None, plotYN=False): print( f"\t- Will implement sawtooth ohmic power correction inside rho={mixRadius}" ) + from mitim_tools.transp_tools import CDFtools Psaw = CDFtools.profilePower( self.profiles["rho(-)"], dvol, @@ -2388,6 +2387,7 @@ def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', times folder = IOtools.expandPath(folder) folder.mkdir(parents=True, exist_ok=True) + from mitim_tools.transp_tools.utils import TRANSPhelpers transp = TRANSPhelpers.transp_run(folder, shot, runid) for time in times: transp.populate_time.from_profiles(time,self, Vsurf = Vsurf) From 7a88e3208bfd0010db53e7619b32d424ce1b1620 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 5 Aug 2025 19:30:36 -0400 Subject: [PATCH 134/385] misc defaults deprecated --- src/mitim_tools/opt_tools/optimizers/ROOTtools.py | 5 ++--- src/mitim_tools/opt_tools/scripts/evaluate_model.py | 1 + 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/opt_tools/optimizers/ROOTtools.py b/src/mitim_tools/opt_tools/optimizers/ROOTtools.py index 9d448ace..5200bc2d 100644 --- a/src/mitim_tools/opt_tools/optimizers/ROOTtools.py +++ b/src/mitim_tools/opt_tools/optimizers/ROOTtools.py @@ -39,12 +39,11 @@ def optimize_function(fun, optimization_params = {}, writeTrajectory=False, meth print("\t- Implementation of simple relaxation method") solver_options = { - "tol": optimization_params.get("tol",-1e-6), "tol_rel": optimization_params.get("relative_improvement_for_stopping",1e-4), - "maxiter": optimization_params.get("maxiter",2000), + "maxiter": optimization_params.get("maxiter",1000), "relax": optimization_params.get("relax",0.1), "relax_dyn": optimization_params.get("relax_dyn",True), - "print_each": optimization_params.get("maxiter",2000)//20, + "print_each": optimization_params.get("maxiter",1000)//20, } solver_fun = optim.simple_relaxation numZ = 6 diff --git a/src/mitim_tools/opt_tools/scripts/evaluate_model.py b/src/mitim_tools/opt_tools/scripts/evaluate_model.py index 331c0804..95c7392d 100644 --- a/src/mitim_tools/opt_tools/scripts/evaluate_model.py +++ b/src/mitim_tools/opt_tools/scripts/evaluate_model.py @@ -60,6 +60,7 @@ gp.localBehavior_scan(gpA.train_X[around, :], dimension_label=input_label,xrange=xrange, axs=axs, c=cols[i], label=input_label) axs[0].legend() + axs[0].set_title("Full behavior (untransformed space)") # gp.plot(plotFundamental=False) # gp.plotTraining() From 10bfe6a0d41ac3302a994abded35f8e4b791b6e9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 5 Aug 2025 19:51:59 -0400 Subject: [PATCH 135/385] Bug fix on positive linear mean constraint --- src/mitim_tools/opt_tools/BOTORCHtools.py | 86 ++++++++--------------- 1 file changed, 30 insertions(+), 56 deletions(-) diff --git a/src/mitim_tools/opt_tools/BOTORCHtools.py b/src/mitim_tools/opt_tools/BOTORCHtools.py index f6b07cb6..f5c99fca 100644 --- a/src/mitim_tools/opt_tools/BOTORCHtools.py +++ b/src/mitim_tools/opt_tools/BOTORCHtools.py @@ -170,10 +170,6 @@ def __init__( self.mean_module = MITIM_LinearMeanGradients( batch_shape=self._aug_batch_shape, variables=variables, output=output ) - elif TypeMean == 3: - self.mean_module = MITIM_CriticalGradient( - batch_shape=self._aug_batch_shape, variables=variables - ) """ ----------------------------------------------------------------------- @@ -689,22 +685,22 @@ def __init__( else: mapping = { - 'Qe': 'aLte', - 'Qi': 'aLti', - 'Ge': 'aLne', - 'GZ': 'aLnZ', - 'Mt': 'dw0dr', - 'Pe': None # Referring to energy exchange + 'Qe_': 'aLte', + 'Qi_': 'aLti', + 'Ge_': 'aLne', + 'GZ_': 'aLnZ', + 'Mt_': 'dw0dr', + 'Qie': None # Referring to energy exchange } for i, variable in enumerate(variables): - if (mapping[output[:2]] is not None) and (mapping[output[:2]] == variable): + if (mapping[output[:3]] is not None) and (mapping[output[:3]] == variable): grad_vector.append(i) self.indeces_grad = tuple(grad_vector) # ---------------------------------------------------------------- - self.register_parameter(name="weights_lin",parameter=torch.nn.Parameter(torch.randn(*batch_shape, len(self.indeces_grad), 1)),) + self.register_parameter(name="raw_weights_lin",parameter=torch.nn.Parameter(torch.randn(*batch_shape, len(self.indeces_grad), 1)),) self.register_parameter(name="bias", parameter=torch.nn.Parameter(torch.randn(*batch_shape, 1))) # set the parameter constraint to be [0,1], when nothing is specified @@ -712,55 +708,33 @@ def __init__( # positive diffusion coefficient if only_diffusive: - self.register_constraint("weights_lin", diffusion_constraint) + self.register_constraint("raw_weights_lin", diffusion_constraint) def forward(self, x): - res = x[..., self.indeces_grad].matmul(self.weights_lin).squeeze(-1) + self.bias + weights_lin = self.weights_lin + res = x[..., self.indeces_grad].matmul(weights_lin).squeeze(-1) + self.bias return res + + # This follows the exact same pattern as in gpytorch's constant_mean.py + @property + def weights_lin(self): + return self._weights_lin_param(self) -class MITIM_CriticalGradient(gpytorch.means.mean.Mean): - def __init__(self, batch_shape=torch.Size(), variables=None, **kwargs): - super().__init__() - - # Indeces of variables that are gradient, so subject to CG behavior - grad_vector = [] - if variables is not None: - for i, variable in enumerate(variables): - if ("aL" in variable) or ("dw" in variable): - grad_vector.append(i) - self.indeces_grad = tuple(grad_vector) - # ---------------------------------------------------------------- - - self.register_parameter( - name="weights_lin", - parameter=torch.nn.Parameter( - torch.randn(*batch_shape, len(self.indeces_grad), 1) - ), - ) - self.register_parameter( - name="bias", parameter=torch.nn.Parameter(torch.randn(*batch_shape, 1)) - ) + @weights_lin.setter + def weights_lin(self, value): + self._weights_lin_closure(self, value) - self.NNfunc = ( - lambda x: x * (1 + torch.erf(x / 0.01)) / 2.0 - ) # https://paperswithcode.com/method/gelu + def _weights_lin_param(self, m): + if hasattr(m, "raw_weights_lin_constraint"): + return m.raw_weights_lin_constraint.transform(m.raw_weights_lin) + return m.raw_weights_lin - self.register_parameter( - name="relu_lin", - parameter=torch.nn.Parameter( - torch.randn(*batch_shape, len(self.indeces_grad), 1) - ), - ) - self.register_constraint( - "relu_lin", gpytorch.constraints.constraints.Interval(0, 1) - ) + def _weights_lin_closure(self, m, value): + if not torch.is_tensor(value): + value = torch.as_tensor(value).to(m.raw_weights_lin) - def forward(self, x): - res = ( - self.NNfunc(x[..., self.indeces_grad] - self.relu_lin.transpose(0, 1)) - .matmul(self.weights_lin) - .squeeze(-1) - + self.bias - ) - return res + if hasattr(m, "raw_weights_lin_constraint"): + m.initialize(raw_weights_lin=m.raw_weights_lin_constraint.inverse_transform(value)) + else: + m.initialize(raw_weights_lin=value) From 8b7cd6a6ba6432acb2b5b97cef0e46b2741afb56 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 5 Aug 2025 20:32:08 -0400 Subject: [PATCH 136/385] Transitioned to a oscillation-detection model for SR dynamic relaxation --- src/mitim_tools/opt_tools/optimizers/optim.py | 195 ++++++++++++------ 1 file changed, 136 insertions(+), 59 deletions(-) diff --git a/src/mitim_tools/opt_tools/optimizers/optim.py b/src/mitim_tools/opt_tools/optimizers/optim.py index c662c347..d1853cab 100644 --- a/src/mitim_tools/opt_tools/optimizers/optim.py +++ b/src/mitim_tools/opt_tools/optimizers/optim.py @@ -1,8 +1,10 @@ +from operator import index import torch import copy import numpy as np +import matplotlib.pyplot as plt from scipy.optimize import root -from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools import GRAPHICStools, IOtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -176,23 +178,27 @@ def function_for_optimizer(x, dfT1=torch.zeros(1).to(x_initial)): # Ready to go optimization tool: Simple Relax # -------------------------------------------------------------------------------------------------------- -def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_options=None ): +def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_options=None, debug=False ): """ See scipy_root for the inputs and outputs """ - tol = solver_options.get("tol", -1e-6) # Tolerance for the residual (negative because I want to maximize) - tol_rel = solver_options.get("tol_rel", None) # Relative tolerance for the residual (superseeds tol) + # ******************************************************************************************** + # Solver options + # ******************************************************************************************** + + tol = solver_options.get("tol", -1e-6) # Tolerance for the residual (negative because I want to maximize) + tol_rel = solver_options.get("tol_rel", None) # Relative tolerance for the residual (superseeds tol) maxiter = solver_options.get("maxiter", 1e5) - relax = solver_options.get("relax", 0.1) # Defines relationship between flux_residual_evaluator and gradient - dx_max = solver_options.get("dx_max", 0.1) # Maximum step size in gradient, relative (e.g. a/Lx can only increase by 10% each time) - dx_max_abs = solver_options.get("dx_max_abs", None) # Maximum step size in gradient, absolute (e.g. a/Lx can only increase by 0.1 each time) - dx_min_abs = solver_options.get("dx_min_abs", None) # Minimum step size in gradient, absolute (e.g. a/Lx must at least increase by 0.01 each time) + relax0 = solver_options.get("relax", 0.1) # Defines relationship between flux_residual_evaluator and gradient + dx_max = solver_options.get("dx_max", 0.1) # Maximum step size in gradient, relative (e.g. a/Lx can only increase by 10% each time) + dx_max_abs = solver_options.get("dx_max_abs", None) # Maximum step size in gradient, absolute (e.g. a/Lx can only increase by 0.1 each time) + dx_min_abs = solver_options.get("dx_min_abs", 1E-5) # Minimum step size in gradient, absolute (e.g. a/Lx must at least increase by 0.01 each time) relax_dyn = solver_options.get("relax_dyn", False) # Dynamic relax, decreases relax if residual is not decreasing relax_dyn_decrease = solver_options.get("relax_dyn_decrease", 5) # Decrease relax by this factor - relax_dyn_num = solver_options.get("relax_dyn_num", 100) # Number of iterations to average over + relax_dyn_num = solver_options.get("relax_dyn_num", 100) # Number of iterations to average over and check if the residual is decreasing relax_dyn_tol_rel = solver_options.get("relax_dyn_tol_rel", 5e-2) # Tolerance to consider that the residual is not decreasing (relative, 0.1 -> 10% minimum change) print_each = solver_options.get("print_each", 1e2) @@ -202,6 +208,15 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o thr_bounds = 1e-4 # To avoid being exactly in the bounds (relative -> 0.01%) + x_initial = x_initial[0,:].unsqueeze(0) + + # Convert relax to tensor of the same dimensions as x, such that it can be dynamically changed per channel + relax = torch.ones_like(x_initial) * relax0 + + # ******************************************************************************************** + # Initial condition + # ******************************************************************************************** + x = copy.deepcopy(x_initial) Q, QT, M = flux_residual_evaluator( x, @@ -215,20 +230,19 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o tol = tol_rel * M.max().item() print(f"\t* Relative tolerance of {tol_rel:.1e} will be used, resulting in an absolute tolerance of {tol:.1e}") + print(f"\t* Flux-grad relationship of {relax0*100.0:.1f}% and maximum gradient jump of {dx_max*100.0:.1f}%,{f' to achieve residual of {tol:.1e}' if tol is not None else ''} in maximum of {maxiter:.0f} iterations") - print(f"\t* Flux-grad relationship of {relax*100.0:.1f}% and maximum gradient jump of {dx_max*100.0:.1f}%,{f' to achieve residual of {tol:.1e}' if tol is not None else ''} in maximum of {maxiter:.0f} iterations") - - # Convert relax to tensor of the same dimensions as x, such that it can be dynamically changed per channel - relax = torch.ones_like(x) * relax + # ******************************************************************************************** + # Iterative strategy + # ******************************************************************************************** - its_since_last_dyn_relax = 0 - i = 0 + relax_history = [] + step_history = [] + its_since_last_dyn_relax, i = 0, 0 for i in range(int(maxiter) - 1): - # -------------------------------------------------------------------------------------------------------- - # Iterative Strategy - # -------------------------------------------------------------------------------------------------------- - x_new = _simple_relax_iteration(x, Q, QT, relax, dx_max, dx_max_abs = dx_max_abs, dx_min_abs = dx_min_abs) + # Make a step in the gradient direction + x_new, x_step = _simple_relax_iteration(x, Q, QT, relax, dx_max, dx_max_abs = dx_max_abs, dx_min_abs = dx_min_abs) # Clamp to bounds if bounds is not None: @@ -237,7 +251,7 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o x = x_new.clone() - # -------------------------------------------------------------------------------------------------------- + # Evaluate new residual Q, QT, M = flux_residual_evaluator( x, y_history = y_history if write_trajectory else None, @@ -257,15 +271,24 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o print(f"\t* Converged in {i+1} iterations with metric of {metric_best:.2e} > {tol:.2e}",typeMsg="i") break + # Update the dynamic relax if needed if relax_dyn and (i-its_since_last_dyn_relax > relax_dyn_num): - relax, hardbreak = _dynamic_relaxation(relax, relax_dyn_decrease, y_history, relax_dyn_num, relax_dyn_tol_rel,i+1) + relax, hardbreak = _dynamic_relaxation(relax, relax_dyn_decrease, x_history, y_history, relax_dyn_num, relax_dyn_tol_rel,i+1) its_since_last_dyn_relax = i if hardbreak: break + + # For debugging + if debug: + step_history.append(x_step[0,:].detach().clone()) + relax_history.append(relax[0,:].clone()) if i == int(maxiter) - 2: print(f"\t* Did not converge in {maxiter} iterations",typeMsg="i") + # ******************************************************************************************** + # Debugging, storing and plotting + # ******************************************************************************************** if write_trajectory: try: y_history = torch.stack(y_history) @@ -279,8 +302,51 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o metric_history = torch.stack(metric_history) except(TypeError,RuntimeError): metric_history = torch.Tensor(metric_history) + if debug: + relax_history = torch.stack(relax_history) + step_history = torch.stack(step_history) else: - y_history, x_history, metric_history = torch.Tensor(), torch.Tensor(), torch.Tensor() + y_history, x_history, metric_history, relax_history, step_history = torch.Tensor(), torch.Tensor(), torch.Tensor(), torch.Tensor(), torch.Tensor() + + + if debug: + fig, axs = plt.subplots(nrows=2, ncols=3, figsize=(15, 10), sharex=True) + + colors = GRAPHICStools.listColors()[:x_history.shape[1]] + + axs = axs.flatten() + + xvals = np.arange(x_history.shape[0]) + x = x_history.cpu().numpy() + y = y_history.cpu().numpy() + r = relax_history.cpu().numpy() + m = metric_history.cpu().numpy() + s = step_history.cpu().numpy() + + plot_ranges = range(x.shape[1]) + + for k in plot_ranges: + axs[0].plot(xvals, x[:,k], '-o', markersize=0.5, lw=1.0, label=f"x{k}", color=colors[k]) + axs[1].plot(xvals, y[:,k], '-o', markersize=0.5, lw=1.0,color=colors[k]) + axs[2].plot(xvals[1:r.shape[0]+1], r[:,k], '-o', markersize=0.5, lw=1.0,color=colors[k]) + axs[3].plot(xvals[1:r.shape[0]+1], s[:,k], '-o', markersize=0.5, lw=1.0,color=colors[k]) + axs[5].plot(xvals, m, '-o', markersize=0.5, lw=1.0) + + for i in range(len(axs)): + GRAPHICStools.addDenseAxis(axs[i]) + axs[i].set_xlabel("Iteration") + + axs[0].set_title("x history"); axs[0].legend() + axs[1].set_title("y history") + axs[2].set_title("Relax history"); axs[2].set_yscale('log') + axs[3].set_title("Step history") + axs[5].set_title("Metric history") + + plt.tight_layout() + + plt.show() + + embed() index_best = metric_history.argmax() print(f"\t* Best metric: {metric_history[index_best].mean().item():.2e} at iteration {index_best}",typeMsg="i") @@ -291,60 +357,71 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o return x_best, y_history, x_history, metric_history -def _dynamic_relaxation(relax, relax_dyn_decrease, y_history, relax_dyn_num, relax_dyn_tol_rel, it, min_relax=1e-6): - ''' - Logic: If the metric is not improving enough, decrease the relax parameter. To determine - if the metric is improving enough, I will fit a line to the last relax_dyn_num points and - check if the slope is small enough. If it is, I will decrease the relax parameter. - ''' +def _check_oscillation(signal): - # Only consider a number of last iterations - y_history = torch.stack(y_history) - y_history_considered = y_history[-relax_dyn_num:].abs() + """Check for oscillations using FFT to detect dominant frequencies""" - # --------------------------------------------------------- - # Calculate improvement in each dimension - # --------------------------------------------------------- + oscillating_dims = torch.zeros(signal.shape[1], dtype=torch.bool) - # Linear fit to the time series for each radius - x_fit = np.arange(len(y_history_considered)) - n_radii = y_history_considered.shape[1] # Number of radius points + # fig, axs = plt.subplots(nrows=2, figsize=(6, 6)) + # colors = GRAPHICStools.listColors() - # Initialize arrays to store results for each radius - change_in_metric = torch.zeros(n_radii) - - # Fit line to each radius dimension separately - for i_radius in range(n_radii): + for i in range(signal.shape[1]): + y_vals = signal[:, i].cpu().numpy() + + # Remove DC component and apply FFT + y_detrended = y_vals - np.mean(y_vals) + fft_vals = np.fft.fft(y_detrended) + power_spectrum = np.abs(fft_vals[1:len(fft_vals)//2+1]) # Exclude DC and negative frequencies - # Fit a line that fits all the considered history - y_fit = y_history_considered[:, i_radius].cpu().numpy() - slope, intercept = np.polyfit(x_fit, y_fit, 1) + # Check if there's a dominant frequency + max_power = np.max(power_spectrum[1:]) # Exclude lowest frequency + total_power = np.sum(power_spectrum) - # Calculate the relative change in metric - metric0 = slope * x_fit[0] + intercept - metric1 = slope * x_fit[-1] + intercept + # If a single frequency dominates (30%), it might be oscillating + single_frequency_power = max_power / total_power + single_frequency_dominance = bool(single_frequency_power > 0.3) - change_in_metric[i_radius] = np.abs((metric1 - metric0) / metric0) + # If more than 50% of the power comes from high frequencies (>1/3), consider it oscillating + index_high_freq = len(power_spectrum) // 3 + high_frequency_power = np.sum(power_spectrum[index_high_freq:]) / total_power + high_frequency_dominance = bool(high_frequency_power > 0.5) - # --------------------------------------------------------- - # Determine which dimensions will need a reduction in relax - # --------------------------------------------------------- + oscillating_dims[i] = single_frequency_dominance or high_frequency_dominance + + # axs[0].plot(y_vals, color=colors[i], ls='-' if oscillating_dims[i] else '--') + # axs[1].plot(power_spectrum/max_power, label = f"{single_frequency_power:.3f}, {high_frequency_power:.3f}", color=colors[i], ls='-' if oscillating_dims[i] else '--') + # axs[1].legend(loc='best',prop={'size': 6}) + # plt.show() + # embed() - mask_reduction = change_in_metric < relax_dyn_tol_rel + return oscillating_dims + +def _dynamic_relaxation(relax, relax_dyn_decrease, x_history, y_history, relax_dyn_num, relax_dyn_tol_rel, it, min_relax=1e-6): + ''' + Logic: If the metric is not improving enough, decrease the relax parameter. To determine + if the metric is improving enough, I will fit a line to the last relax_dyn_num points and + check if the slope is small enough. If it is, I will decrease the relax parameter. + ''' + + # Only consider a number of last iterations + x_history_considered = torch.stack(x_history)[-relax_dyn_num:] + y_history_considered = torch.stack(y_history)[-relax_dyn_num:] + + mask_reduction = _check_oscillation(x_history_considered) if mask_reduction.any(): - if (relax[:,mask_reduction] < min_relax).all(): - print(f"\t\t\t<> Metric not improving enough (@{it}), relax already at minimum of {min_relax:.1e}, not worth continuing", typeMsg="i") + if (relax < min_relax).all(): + print(f"\t\t\t<> Oscillatory behavior detected (@{it}), relax already at minimum of {min_relax:.1e}, not worth continuing", typeMsg="i") return relax, True - - print(f"\t\t\t<> Metric not improving enough (@{it}), decreasing relax for {mask_reduction.sum()} out of {n_radii} channels") + + print(f"\t\t\t<> Oscillatory behavior detected (@{it}), decreasing relax for {mask_reduction.sum()} out of {y_history_considered.shape[1]} channels") relax[:,mask_reduction] = relax[:,mask_reduction] / relax_dyn_decrease print(f"\t\t\t\t- New relax values: from {relax.min():.1e} to {relax.max():.1e}") return relax, False else: - print(f"\t\t\t<> Metric improving enough (@{it}), relax remains at {relax.min():.1e} to {relax.max():.1e}") return relax, False @@ -372,7 +449,7 @@ def _simple_relax_iteration(x, Q, QT, relax, dx_max, dx_max_abs = None, dx_min_a # Update x_new = x + x_step - return x_new + return x_new, x_step ''' ********************************************************************************************************************************** From 9dfa701353928df7db38140aa8216c985dd9671a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 11:38:15 -0400 Subject: [PATCH 137/385] Optimization tools renaming and improvement of SR, batching --- src/mitim_modules/powertorch/STATEtools.py | 27 +- src/mitim_tools/opt_tools/OPTtools.py | 6 +- src/mitim_tools/opt_tools/STRATEGYtools.py | 10 +- .../{BOTORCHoptim.py => botorch_tools.py} | 0 .../{GAtools.py => evolutionary.py} | 0 .../{ROOTtools.py => multivariate.py} | 45 +-- .../{optim.py => multivariate_tools.py} | 341 ++++++++++-------- .../scripts/evaluate_optimizer_root.py | 4 +- 8 files changed, 226 insertions(+), 207 deletions(-) rename src/mitim_tools/opt_tools/optimizers/{BOTORCHoptim.py => botorch_tools.py} (100%) rename src/mitim_tools/opt_tools/optimizers/{GAtools.py => evolutionary.py} (100%) rename src/mitim_tools/opt_tools/optimizers/{ROOTtools.py => multivariate.py} (81%) rename src/mitim_tools/opt_tools/optimizers/{optim.py => multivariate_tools.py} (72%) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 8643170b..a6a9b8a0 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -9,7 +9,7 @@ from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_modules.powertorch.utils import TRANSFORMtools, POWERplot -from mitim_tools.opt_tools.optimizers import optim +from mitim_tools.opt_tools.optimizers import multivariate_tools from mitim_modules.powertorch.utils import TARGETStools, CALCtools, TRANSPORTtools from mitim_modules.powertorch.physics_models import targets_analytic from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -345,9 +345,9 @@ def flux_match(self, algorithm="root", solver_options=None, bounds=None, debugYN timeBeginning = datetime.datetime.now() if algorithm == "root": - solver_fun = optim.scipy_root + solver_fun = multivariate_tools.scipy_root elif algorithm == "simple_relax": - solver_fun = optim.simple_relaxation + solver_fun = multivariate_tools.simple_relaxation else: raise ValueError(f"[MITIM] Algorithm {algorithm} not recognized") @@ -390,20 +390,15 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): # Residual is the difference between the target and the transport yRes = (QTarget - QTransport).abs() + # Metric is the mean of the absolute value of the residual - yMetric = -yRes.mean(axis=-1) - # Best in batch - best_candidate = yMetric.argmax().item() - # Only pass the best candidate - yRes = yRes[best_candidate, :].detach() - yMetric = yMetric[best_candidate].detach() - Xpass = X[best_candidate, :].detach() + yMetric = -yRes.mean(axis=-1).detach() # Store values if y_history is not None: - y_history.append(yRes) + y_history.append(yRes.detach()) if x_history is not None: - x_history.append(Xpass) + x_history.append(X.detach()) if metric_history is not None: metric_history.append(yMetric) @@ -418,10 +413,14 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): x0 = x0.view((self.plasma["rho"].shape[0],(self.plasma["rho"].shape[1] - 1) * len(self.ProfilesPredicted),)) # Optimize - _,Yopt, Xopt, metric_history = solver_fun(evaluator,x0, bounds=self.bounds_current,solver_options=solver_options) + x_best,Yopt, Xopt, metric_history = solver_fun(evaluator,x0, bounds=self.bounds_current,solver_options=solver_options) # For simplicity, return the trajectory of only the best candidate - self.FluxMatch_Yopt, self.FluxMatch_Xopt = Yopt, Xopt + + idx_flat = metric_history.argmax() + index_best = divmod(idx_flat.item(), metric_history.shape[1]) + + self.FluxMatch_Yopt, self.FluxMatch_Xopt = Yopt[:,index_best[1],:], Xopt[:,index_best[1],:] print("**********************************************************************************************") print(f"\t- Flux matching of powerstate finished, and took {IOtools.getTimeDifference(timeBeginning)}\n") diff --git a/src/mitim_tools/opt_tools/OPTtools.py b/src/mitim_tools/opt_tools/OPTtools.py index 6dec83f6..91af55ac 100644 --- a/src/mitim_tools/opt_tools/OPTtools.py +++ b/src/mitim_tools/opt_tools/OPTtools.py @@ -221,11 +221,11 @@ def acquire_next_points( # Prepare (run more now to find more solutions, more diversity, even if later best_points is 1) if optimizer == "ga": - from mitim_tools.opt_tools.optimizers.GAtools import optimize_function + from mitim_tools.opt_tools.optimizers.evolutionary import optimize_function elif optimizer == "botorch": - from mitim_tools.opt_tools.optimizers.BOTORCHoptim import optimize_function + from mitim_tools.opt_tools.optimizers.botorch_tools import optimize_function elif optimizer == "root" or optimizer == "sr": - from mitim_tools.opt_tools.optimizers.ROOTtools import optimize_function + from mitim_tools.opt_tools.optimizers.multivariate import optimize_function if optimizer == "root": optimize_function = partial(optimize_function, method="scipy_root") elif optimizer == "sr" : diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 1533cfd1..b3787f13 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -1583,11 +1583,16 @@ def plotAcquisitionOptimizationSummary(self, fn=None, step_from=0, step_to=-1): # Plot acquisition evolution for i in range(len(infoOPT)-1): #no cleanup stage y_acq = infoOPT[i]['info']['acq_evaluated'].cpu().numpy() - ax.plot(y_acq,'-o', c=colors[i], markersize=1, lw = 0.5, label=f'{infoOPT[i]["method"]} (max of batch)') + + if len(y_acq.shape)>1: + for j in range(y_acq.shape[1]): + ax.plot(y_acq[:,j],'-o', c=colors[i], markersize=0.5, lw = 0.3, label=f'{infoOPT[i]["method"]} (candidate #{j})') + else: + ax.plot(y_acq,'-o', c=colors[i], markersize=1, lw = 0.5, label=f'{infoOPT[i]["method"]}') # Plot max of guesses if len(y_acq)>0: - ax.axhline(y=y_acq[0], c=colors[i], ls='--', lw=1.0, label=f'{infoOPT[i]["method"]} (max of guesses)') + ax.axhline(y=y_acq.max(axis=1)[0], c=colors[i], ls='--', lw=1.0, label=f'{infoOPT[i]["method"]} (max of guesses)') ax.set_title(f'BO Step #{step}') ax.set_ylabel('$f_{acq}$ (to max)') @@ -1597,6 +1602,7 @@ def plotAcquisitionOptimizationSummary(self, fn=None, step_from=0, step_to=-1): GRAPHICStools.addDenseAxis(ax) + def plotModelStatus( self, fn=None, boStep=-1, plotsPerFigure=20, stds=2, tab_color=None ): diff --git a/src/mitim_tools/opt_tools/optimizers/BOTORCHoptim.py b/src/mitim_tools/opt_tools/optimizers/botorch_tools.py similarity index 100% rename from src/mitim_tools/opt_tools/optimizers/BOTORCHoptim.py rename to src/mitim_tools/opt_tools/optimizers/botorch_tools.py diff --git a/src/mitim_tools/opt_tools/optimizers/GAtools.py b/src/mitim_tools/opt_tools/optimizers/evolutionary.py similarity index 100% rename from src/mitim_tools/opt_tools/optimizers/GAtools.py rename to src/mitim_tools/opt_tools/optimizers/evolutionary.py diff --git a/src/mitim_tools/opt_tools/optimizers/ROOTtools.py b/src/mitim_tools/opt_tools/optimizers/multivariate.py similarity index 81% rename from src/mitim_tools/opt_tools/optimizers/ROOTtools.py rename to src/mitim_tools/opt_tools/optimizers/multivariate.py index 5200bc2d..b468e90e 100644 --- a/src/mitim_tools/opt_tools/optimizers/ROOTtools.py +++ b/src/mitim_tools/opt_tools/optimizers/multivariate.py @@ -1,14 +1,12 @@ import torch import copy import numpy as np -import matplotlib.pyplot as plt from mitim_tools.misc_tools.LOGtools import printMsg as print -from mitim_tools.opt_tools.optimizers import optim +from mitim_tools.opt_tools.optimizers import multivariate_tools from mitim_tools.opt_tools.utils import TESTtools -from mitim_tools.misc_tools import GRAPHICStools from IPython import embed -def optimize_function(fun, optimization_params = {}, writeTrajectory=False, method = 'scipy_root', debugYN=False): +def optimize_function(fun, optimization_params = {}, writeTrajectory=False, method = 'scipy_root'): np.random.seed(fun.seed) @@ -31,7 +29,7 @@ def optimize_function(fun, optimization_params = {}, writeTrajectory=False, meth 'solver': optimization_params.get("solver","lm"), 'write_trajectory': writeTrajectory } - solver_fun = optim.scipy_root + solver_fun = multivariate_tools.scipy_root numZ = 5 elif method == "sr": @@ -45,7 +43,7 @@ def optimize_function(fun, optimization_params = {}, writeTrajectory=False, meth "relax_dyn": optimization_params.get("relax_dyn",True), "print_each": optimization_params.get("maxiter",1000)//20, } - solver_fun = optim.simple_relaxation + solver_fun = multivariate_tools.simple_relaxation numZ = 6 # -------------------------------------------------------------------------------------------------------- @@ -57,24 +55,13 @@ def flux_residual_evaluator(X, y_history=None, x_history=None, metric_history=No # Evaluate source term yOut, y1, y2, _ = fun.evaluators["residual_function"](X, outputComponents=True) - # ----------------------------------------- - # Post-process - # ----------------------------------------- - - # Best in batch - best_candidate = yOut.argmax().item() - # Only pass the best candidate - yRes = (y2-y1)[best_candidate, :].detach() - yMetric = yOut[best_candidate].detach() - Xpass = X[best_candidate, :].detach() - # Store values if metric_history is not None: - metric_history.append(yMetric) + metric_history.append(yOut.detach()) if x_history is not None: - x_history.append(Xpass) + x_history.append(X.detach()) if y_history is not None: - y_history.append(yRes) + y_history.append((y2-y1).detach()) return y1, y2, yOut @@ -108,24 +95,6 @@ def flux_residual_evaluator(X, y_history=None, x_history=None, metric_history=No x_res, y_history, x_history, acq_evaluated = solver_fun(flux_residual_evaluator,xGuesses,solver_options=solver_options,bounds=bounds) print("************************************************************************************************") - if debugYN: - fig, axs = plt.subplots(nrows= 2, figsize=(6, 10)) - ax = axs[0] - ax.plot(np.array(x_history)) - ax.set_xlabel("Iteration") - ax.set_ylabel("X") - GRAPHICStools.addDenseAxis(ax) - - ax = axs[1] - ax.plot(np.abs(np.array(y_history))) - ax.set_xlabel("Iteration") - ax.set_ylabel("|Y|") - ax.set_yscale("log") - GRAPHICStools.addDenseAxis(ax) - - plt.show() - embed() - # -------------------------------------------------------------------------------------------------------- # Post-process # -------------------------------------------------------------------------------------------------------- diff --git a/src/mitim_tools/opt_tools/optimizers/optim.py b/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py similarity index 72% rename from src/mitim_tools/opt_tools/optimizers/optim.py rename to src/mitim_tools/opt_tools/optimizers/multivariate_tools.py index d1853cab..bba219ba 100644 --- a/src/mitim_tools/opt_tools/optimizers/optim.py +++ b/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py @@ -19,10 +19,10 @@ def scipy_root(flux_residual_evaluator, x_initial, bounds=None, solver_options=N - flux_residual_evaluator is a function that: - Takes X (batches,dimX) - Provides Y1: transport (batches,dimY), Y2: target (batches,dimY) and M: maximization metric (batches,1) - It must also take optional arguments, to capture the best in the batch: + It must also take optional arguments, to capture the evolution of the batch: x_history - y_history () - metric_history (to maximize, similar to acquisition definition, must be 1D, best in batch) + y_history + metric_history (to maximize, similar to acquisition definition) Outputs: - Optium vector x_sol with (batches,dimX) and the trajectory of the acquisition function evaluations (best per batch) Notes: @@ -173,7 +173,6 @@ def function_for_optimizer(x, dfT1=torch.zeros(1).to(x_initial)): return x_best, y_history, x_history, metric_history - # -------------------------------------------------------------------------------------------------------- # Ready to go optimization tool: Simple Relax # -------------------------------------------------------------------------------------------------------- @@ -199,23 +198,20 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o relax_dyn = solver_options.get("relax_dyn", False) # Dynamic relax, decreases relax if residual is not decreasing relax_dyn_decrease = solver_options.get("relax_dyn_decrease", 5) # Decrease relax by this factor relax_dyn_num = solver_options.get("relax_dyn_num", 100) # Number of iterations to average over and check if the residual is decreasing - relax_dyn_tol_rel = solver_options.get("relax_dyn_tol_rel", 5e-2) # Tolerance to consider that the residual is not decreasing (relative, 0.1 -> 10% minimum change) print_each = solver_options.get("print_each", 1e2) - write_trajectory = solver_options.get("write_trajectory", True) - x_history, y_history, metric_history = [], [], [] - + thr_bounds = 1e-4 # To avoid being exactly in the bounds (relative -> 0.01%) - x_initial = x_initial[0,:].unsqueeze(0) + # ******************************************************************************************** + # Initial condition + # ******************************************************************************************** # Convert relax to tensor of the same dimensions as x, such that it can be dynamically changed per channel relax = torch.ones_like(x_initial) * relax0 - # ******************************************************************************************** - # Initial condition - # ******************************************************************************************** + x_history, y_history, metric_history = [], [], [] x = copy.deepcopy(x_initial) Q, QT, M = flux_residual_evaluator( @@ -224,64 +220,79 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o x_history = x_history if write_trajectory else None, metric_history = metric_history if write_trajectory else None ) - print(f"\t* Starting residual: {(Q-QT).abs().mean(axis=1)[0].item():.4e}, will run {int(maxiter)-1} more evaluations, printing every {print_each} iteration:",typeMsg="i") + + print(f"\t* Starting best-candidate residual: {(Q-QT).abs().mean(axis=1).min().item():.4e}, will run {int(maxiter)-1} more evaluations, printing every {print_each} iteration",typeMsg="i") if tol_rel is not None: tol = tol_rel * M.max().item() print(f"\t* Relative tolerance of {tol_rel:.1e} will be used, resulting in an absolute tolerance of {tol:.1e}") - print(f"\t* Flux-grad relationship of {relax0*100.0:.1f}% and maximum gradient jump of {dx_max*100.0:.1f}%,{f' to achieve residual of {tol:.1e}' if tol is not None else ''} in maximum of {maxiter:.0f} iterations") + print(f"\t* Flux-grad relationship of {relax0} and maximum gradient jump of {dx_max}") # ******************************************************************************************** # Iterative strategy # ******************************************************************************************** - relax_history = [] - step_history = [] + hardbreak = False + relax_history, step_history = [], [] its_since_last_dyn_relax, i = 0, 0 + for i in range(int(maxiter) - 1): # Make a step in the gradient direction - x_new, x_step = _simple_relax_iteration(x, Q, QT, relax, dx_max, dx_max_abs = dx_max_abs, dx_min_abs = dx_min_abs) - - # Clamp to bounds - if bounds is not None: - bb = bounds[1,:]-bounds[0,:] - x_new = x_new.clamp(min=bounds[0,:]+thr_bounds*bb, max=bounds[1,:]-thr_bounds*bb) + x_new, x_step = _sr_step( + x, + Q, + QT, + relax, + dx_max, + dx_max_abs=dx_max_abs, + dx_min_abs=dx_min_abs, + bounds=bounds, + thr_bounds=thr_bounds + ) + # Make it the new point x = x_new.clone() # Evaluate new residual Q, QT, M = flux_residual_evaluator( x, - y_history = y_history if write_trajectory else None, - x_history = x_history if write_trajectory else None, - metric_history = metric_history - ) + y_history=y_history if write_trajectory else None, + x_history=x_history if write_trajectory else None, + metric_history=metric_history if write_trajectory else None + ) # Best metric of the batch - metric_best = M.max(axis=-1)[0].item() + metric_best = M.max().item() if (i + 1) % int(print_each) == 0: - print(f"\t\t- Metric (to maximize) @{i+1}: {metric_best:.2e}") + print(f"\t\t- Best metric (to maximize) @{i+1}: {metric_best:.2e}") # Stopping based on the best of the batch based on the metric - if tol is not None and M.max().item() > tol: - print(f"\t\t- Metric (to maximize) @{i+1}: {metric_best:.2e}",typeMsg="i") + if (tol is not None) and (M.max().item() > tol): print(f"\t* Converged in {i+1} iterations with metric of {metric_best:.2e} > {tol:.2e}",typeMsg="i") break # Update the dynamic relax if needed - if relax_dyn and (i-its_since_last_dyn_relax > relax_dyn_num): - relax, hardbreak = _dynamic_relaxation(relax, relax_dyn_decrease, x_history, y_history, relax_dyn_num, relax_dyn_tol_rel,i+1) - its_since_last_dyn_relax = i - if hardbreak: - break + if relax_dyn: + relax, its_since_last_dyn_relax, hardbreak = _dynamic_relax( + x_history, + y_history, + relax, + relax_dyn_decrease, + relax_dyn_num, + i, + its_since_last_dyn_relax + ) # For debugging if debug: - step_history.append(x_step[0,:].detach().clone()) - relax_history.append(relax[0,:].clone()) + step_history.append(x_step.detach().clone()) + relax_history.append(relax.clone()) + + if hardbreak: + break if i == int(maxiter) - 2: print(f"\t* Did not converge in {maxiter} iterations",typeMsg="i") @@ -289,6 +300,7 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o # ******************************************************************************************** # Debugging, storing and plotting # ******************************************************************************************** + if write_trajectory: try: y_history = torch.stack(y_history) @@ -300,132 +312,68 @@ def simple_relaxation( flux_residual_evaluator, x_initial, bounds=None, solver_o x_history = torch.Tensor(x_history) try: metric_history = torch.stack(metric_history) - except(TypeError,RuntimeError): + except (TypeError,RuntimeError): metric_history = torch.Tensor(metric_history) - if debug: - relax_history = torch.stack(relax_history) - step_history = torch.stack(step_history) else: - y_history, x_history, metric_history, relax_history, step_history = torch.Tensor(), torch.Tensor(), torch.Tensor(), torch.Tensor(), torch.Tensor() - + y_history, x_history, metric_history = torch.Tensor(), torch.Tensor(), torch.Tensor() if debug: - fig, axs = plt.subplots(nrows=2, ncols=3, figsize=(15, 10), sharex=True) - - colors = GRAPHICStools.listColors()[:x_history.shape[1]] + + relax_history = torch.stack(relax_history) + step_history = torch.stack(step_history) - axs = axs.flatten() - - xvals = np.arange(x_history.shape[0]) - x = x_history.cpu().numpy() - y = y_history.cpu().numpy() - r = relax_history.cpu().numpy() - m = metric_history.cpu().numpy() - s = step_history.cpu().numpy() - - plot_ranges = range(x.shape[1]) - - for k in plot_ranges: - axs[0].plot(xvals, x[:,k], '-o', markersize=0.5, lw=1.0, label=f"x{k}", color=colors[k]) - axs[1].plot(xvals, y[:,k], '-o', markersize=0.5, lw=1.0,color=colors[k]) - axs[2].plot(xvals[1:r.shape[0]+1], r[:,k], '-o', markersize=0.5, lw=1.0,color=colors[k]) - axs[3].plot(xvals[1:r.shape[0]+1], s[:,k], '-o', markersize=0.5, lw=1.0,color=colors[k]) - axs[5].plot(xvals, m, '-o', markersize=0.5, lw=1.0) - - for i in range(len(axs)): - GRAPHICStools.addDenseAxis(axs[i]) - axs[i].set_xlabel("Iteration") + for candidate in range(x_history.shape[1]): + + fig, axs = plt.subplots(nrows=2, ncols=3, figsize=(15, 10), sharex=True) + + axs = axs.flatten() + + x = x_history[:,candidate,:].cpu().numpy() + y = y_history[:,candidate,:].cpu().numpy() + r = relax_history[:,candidate,:].cpu().numpy() + m = metric_history[:,candidate].cpu().numpy() + s = step_history[:,candidate,:].cpu().numpy() - axs[0].set_title("x history"); axs[0].legend() - axs[1].set_title("y history") - axs[2].set_title("Relax history"); axs[2].set_yscale('log') - axs[3].set_title("Step history") - axs[5].set_title("Metric history") - - plt.tight_layout() + colors = GRAPHICStools.listColors()[:x.shape[-1]] + + xvals = np.arange(x.shape[0]) + plot_ranges = range(x.shape[1]) + + for k in plot_ranges: + axs[0].plot(xvals, x[:,k], '-o', markersize=0.5, lw=1.0, label=f"x{k}", color=colors[k]) + axs[1].plot(xvals, y[:,k], '-o', markersize=0.5, lw=1.0,color=colors[k]) + axs[2].plot(xvals[1:r.shape[0]+1], r[:,k], '-o', markersize=0.5, lw=1.0,color=colors[k]) + axs[3].plot(xvals[1:r.shape[0]+1], s[:,k], '-o', markersize=0.5, lw=1.0,color=colors[k]) + axs[5].plot(xvals, m, '-o', markersize=0.5, lw=1.0) + + for i in range(len(axs)): + GRAPHICStools.addDenseAxis(axs[i]) + axs[i].set_xlabel("Iteration") + + axs[0].set_title("x history"); axs[0].legend() + axs[1].set_title("y history") + axs[2].set_title("Relax history"); axs[2].set_yscale('log') + axs[3].set_title("Step history") + axs[5].set_title("Metric history") + + plt.tight_layout() plt.show() embed() - index_best = metric_history.argmax() - print(f"\t* Best metric: {metric_history[index_best].mean().item():.2e} at iteration {index_best}",typeMsg="i") - - # The best candidate, regardless of the restarts - x_best = x_history[index_best,:].unsqueeze(0) + # Find the best iteration of each candidate trajectory + index_bests = metric_history.argmax(dim=0) + x_best = x_history[index_bests, torch.arange(x_history.shape[1]), :] - return x_best, y_history, x_history, metric_history + idx_flat = metric_history.argmax() + index_best = divmod(idx_flat.item(), metric_history.shape[1]) + print(f"\t* Best metric: {metric_history[index_best].item():.2e} at iteration {index_best[0]} for candidate in position {index_best[1]}",typeMsg="i") + return x_best, y_history, x_history, metric_history -def _check_oscillation(signal): - - """Check for oscillations using FFT to detect dominant frequencies""" - - oscillating_dims = torch.zeros(signal.shape[1], dtype=torch.bool) - - # fig, axs = plt.subplots(nrows=2, figsize=(6, 6)) - # colors = GRAPHICStools.listColors() +def _sr_step(x, Q, QT, relax, dx_max, dx_max_abs = None, dx_min_abs = None, threshold_zero_flux_issue=1e-10, bounds=None, thr_bounds=1e-4): - for i in range(signal.shape[1]): - y_vals = signal[:, i].cpu().numpy() - - # Remove DC component and apply FFT - y_detrended = y_vals - np.mean(y_vals) - fft_vals = np.fft.fft(y_detrended) - power_spectrum = np.abs(fft_vals[1:len(fft_vals)//2+1]) # Exclude DC and negative frequencies - - # Check if there's a dominant frequency - max_power = np.max(power_spectrum[1:]) # Exclude lowest frequency - total_power = np.sum(power_spectrum) - - # If a single frequency dominates (30%), it might be oscillating - single_frequency_power = max_power / total_power - single_frequency_dominance = bool(single_frequency_power > 0.3) - - # If more than 50% of the power comes from high frequencies (>1/3), consider it oscillating - index_high_freq = len(power_spectrum) // 3 - high_frequency_power = np.sum(power_spectrum[index_high_freq:]) / total_power - high_frequency_dominance = bool(high_frequency_power > 0.5) - - oscillating_dims[i] = single_frequency_dominance or high_frequency_dominance - - # axs[0].plot(y_vals, color=colors[i], ls='-' if oscillating_dims[i] else '--') - # axs[1].plot(power_spectrum/max_power, label = f"{single_frequency_power:.3f}, {high_frequency_power:.3f}", color=colors[i], ls='-' if oscillating_dims[i] else '--') - # axs[1].legend(loc='best',prop={'size': 6}) - # plt.show() - # embed() - - return oscillating_dims - -def _dynamic_relaxation(relax, relax_dyn_decrease, x_history, y_history, relax_dyn_num, relax_dyn_tol_rel, it, min_relax=1e-6): - ''' - Logic: If the metric is not improving enough, decrease the relax parameter. To determine - if the metric is improving enough, I will fit a line to the last relax_dyn_num points and - check if the slope is small enough. If it is, I will decrease the relax parameter. - ''' - - # Only consider a number of last iterations - x_history_considered = torch.stack(x_history)[-relax_dyn_num:] - y_history_considered = torch.stack(y_history)[-relax_dyn_num:] - - mask_reduction = _check_oscillation(x_history_considered) - - if mask_reduction.any(): - - if (relax < min_relax).all(): - print(f"\t\t\t<> Oscillatory behavior detected (@{it}), relax already at minimum of {min_relax:.1e}, not worth continuing", typeMsg="i") - return relax, True - - print(f"\t\t\t<> Oscillatory behavior detected (@{it}), decreasing relax for {mask_reduction.sum()} out of {y_history_considered.shape[1]} channels") - relax[:,mask_reduction] = relax[:,mask_reduction] / relax_dyn_decrease - print(f"\t\t\t\t- New relax values: from {relax.min():.1e} to {relax.max():.1e}") - - return relax, False - else: - - return relax, False - -def _simple_relax_iteration(x, Q, QT, relax, dx_max, dx_max_abs = None, dx_min_abs = None, threshold_zero_flux_issue=1e-10): # Calculate step in gradient (if target > transport, dx>0 because I want to increase gradients) dx = relax * (QT - Q) / (Q**2 + QT**2).clamp(min=threshold_zero_flux_issue) ** 0.5 @@ -441,6 +389,7 @@ def _simple_relax_iteration(x, Q, QT, relax, dx_max, dx_max_abs = None, dx_min_a ix = x_step.abs() > dx_max_abs direction = torch.nan_to_num(x_step[ix] / x_step[ix].abs(), nan=1.0) x_step[ix] = dx_max_abs * direction + if dx_min_abs is not None: ix = x_step.abs() < dx_min_abs direction = torch.nan_to_num(x_step[ix] / x_step[ix].abs(), nan=1.0) @@ -449,8 +398,104 @@ def _simple_relax_iteration(x, Q, QT, relax, dx_max, dx_max_abs = None, dx_min_a # Update x_new = x + x_step + # Clamp to bounds + if bounds is not None: + thr_bounds_abs = ( bounds[1,:] - bounds[0,:]) * thr_bounds + x_new = x_new.clamp(min=bounds[0,:]+thr_bounds_abs, max=bounds[1,:]-thr_bounds_abs) + return x_new, x_step +def _dynamic_relax(x, y, relax, relax_dyn_decrease, relax_dyn_num, iteration_num, iteration_applied): + + min_relax = 1e-6 + + if iteration_num - iteration_applied > relax_dyn_num: + + mask_reduction = _check_oscillation(torch.stack(x), relax_dyn_num) + + if mask_reduction.any(): + + if (relax < min_relax).all(): + print(f"\t\t\t<> Oscillatory behavior detected (@{iteration_num}), all relax already at minimum of {min_relax:.1e}, not worth continuing", typeMsg="i") + return relax, iteration_applied, True + + print(f"\t\t\t<> Oscillatory behavior detected (@{iteration_num}), decreasing relax for {mask_reduction.sum()} out of {torch.stack(x).shape[1]*torch.stack(x).shape[2]} channels") + + relax[mask_reduction] = relax[mask_reduction] / relax_dyn_decrease + + print(f"\t\t\t\t- New relax values span from {relax.min():.1e} to {relax.max():.1e}") + + iteration_applied = iteration_num + + return relax, iteration_applied, False + +def _check_oscillation(signal_raw, relax_dyn_num): + + """Check for oscillations using FFT to detect dominant frequencies""" + + # Stack batch dimension (time, batch, dim) -> (time, batch*dim) + signal = signal_raw.reshape(signal_raw.shape[0], -1) + + oscillating_dims = torch.zeros(signal.shape[1:], dtype=torch.bool) + + # fig, axs = plt.subplots(nrows=2, figsize=(6, 6)) + # colors = GRAPHICStools.listColors() + + for i in range(signal.shape[1]): + + iterations_to_consider = relax_dyn_num + + # Only consider a number of last iterations + y_vals = signal[-iterations_to_consider:, i].cpu().numpy() + + # If the signal is not constant + if y_vals.std() > 0.0: + + # Remove DC component and apply FFT + y_detrended = y_vals - np.mean(y_vals) + fft_vals = np.fft.fft(y_detrended) + power_spectrum = np.abs(fft_vals[1:len(fft_vals)//2+1]) # Exclude DC and negative frequencies + + # Check if there's a dominant frequency + excl = 2 + p_around = 1 + argmax_power = np.argmax(power_spectrum[excl:]) # Exclude lowest frequencies + max_power = np.sum(power_spectrum[(argmax_power+excl) - p_around:(argmax_power+excl) + p_around]) + total_power = np.sum(power_spectrum) + + # If a single frequency dominates (30%), it might be oscillating (even if low frequency) + single_frequency_power = max_power / total_power + single_frequency_dominance = bool(single_frequency_power > 0.3) + + # If more than 50% of the power comes from high frequencies (>1/3), consider it oscillating + index_high_freq = len(power_spectrum) // 3 + high_frequency_power = np.sum(power_spectrum[index_high_freq:]) / total_power + high_frequency_dominance = bool(high_frequency_power > 0.5) + + # if signal completely flat, it's an indication that has hit the bounds, also consider it oscillating + signal_flat = bool(y_vals.std() < 1e-6) + + # If the signal is constant, consider it non-oscillating but flat + else: + single_frequency_dominance = False + high_frequency_dominance = False + signal_flat = True + + oscillating_dims[i] = single_frequency_dominance or high_frequency_dominance or signal_flat + + + # Back to the original shape + oscillating_dims = oscillating_dims.reshape(signal_raw.shape[1:]) + + # axs[0].plot(y_vals, color=colors[i], ls='-' if oscillating_dims[i] else '--') + # axs[1].plot(power_spectrum/max_power, label = f"{single_frequency_power:.3f}, {high_frequency_power:.3f}, {y_vals.std():.1e}", color=colors[i], ls='-' if oscillating_dims[i] else '--') + # axs[1].legend(loc='best',prop={'size': 6}) + # plt.show() + + return oscillating_dims + + + ''' ********************************************************************************************************************************** The original implementation of torch.autograd.functional.jacobian runs the function once and then computes the jacobian. diff --git a/src/mitim_tools/opt_tools/scripts/evaluate_optimizer_root.py b/src/mitim_tools/opt_tools/scripts/evaluate_optimizer_root.py index 1d54cc79..86580247 100644 --- a/src/mitim_tools/opt_tools/scripts/evaluate_optimizer_root.py +++ b/src/mitim_tools/opt_tools/scripts/evaluate_optimizer_root.py @@ -7,7 +7,7 @@ from mitim_tools.misc_tools import IOtools, GRAPHICStools from mitim_tools.opt_tools import STRATEGYtools, OPTtools from mitim_tools.opt_tools.utils import TESTtools -from mitim_tools.opt_tools.optimizers import ROOTtools +from mitim_tools.opt_tools.optimizers import multivariate from IPython import embed """ @@ -60,7 +60,7 @@ # for opt,lab in enumerate(['vectorize=True']): #,'vectorize=False']): for opt, lab in enumerate(["x0=0"]): # ,'x0=1.0']): #,'vectorize=False']): - logi = ROOTtools.logistic(l=bounds_logi[0, :], u=bounds_logi[1, :], k=0.5, x0=0) + logi = multivariate.logistic(l=bounds_logi[0, :], u=bounds_logi[1, :], k=0.5, x0=0) # *************************************************************************************************** # OPTIMIZER # *************************************************************************************************** From d16136608131fadc0966fa069e5161c9eeeee03c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 13:00:17 -0400 Subject: [PATCH 138/385] Changed the Ricci default --- src/mitim_modules/portals/PORTALSmain.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index e4c9702b..d8acafb3 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -52,10 +52,10 @@ def default_namelist(optimization_options, CGYROrun=False): optimization_options["convergence_options"]["maximum_iterations"] = 50 optimization_options['convergence_options']['stopping_criteria'] = PORTALStools.stopping_criteria_portals optimization_options['convergence_options']['stopping_criteria_parameters'] = { - "maximum_value": 5e-3, # Reducing residual by 200x is enough + "maximum_value": 5e-3, # Reducing residual by 200x is enough "maximum_value_is_rel": True, "minimum_dvs_variation": [10, 5, 0.1], # After iteration 10, Check if 5 consecutive DVs are varying less than 0.1% from the rest that has been evaluated - "ricci_value": 0.1, + "ricci_value": 0.05, "ricci_d0": 2.0, "ricci_lambda": 0.5, } From 696a6f8b22f16203ccd831fc7d87ba6a6ce88aa6 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 13:08:43 -0400 Subject: [PATCH 139/385] bug fix from mitimstate transition --- src/mitim_tools/gacode_tools/TGYROtools.py | 33 ++++--------------- .../plasmastate_tools/MITIMstate.py | 3 ++ 2 files changed, 9 insertions(+), 27 deletions(-) diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 45305dbf..247a7835 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -4338,7 +4338,7 @@ def plotBalance(self, fig=None): axs.append(fig.add_subplot(grid[1, 1])) axs.append(fig.add_subplot(grid[1, 2])) - self.profiles_final.plotBalance( + self.profiles_final.plot_flows( axs=axs, limits=[self.roa[-1, 1], self.roa[-1, -1]] ) @@ -4411,20 +4411,12 @@ def plotConvergence(self, fig1=None): ) GRAPHICStools.addDenseAxis(ax) - # GRAPHICStools.autoscale_y(ax) - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=False, extraPad=0, loc="center left", size=6 - ) - # GRAPHICStools.addLegendApart(ax2,ratio=0.9,withleg=False,extraPad=0,loc='center left',size=6) + GRAPHICStools.addLegendApart(ax, ratio=0.9, withleg=False, extraPad=0, loc="center left", size=6) ax = ax10 - colsE = ( - GRAPHICStools.listColors() - ) # GRAPHICStools.colorTableFade(self.radii-1,startcolor='b',endcolor='b',alphalims=[0.3,1.0]) - colsI = ( - GRAPHICStools.listColors() - ) # GRAPHICStools.colorTableFade(self.radii-1,startcolor='r',endcolor='r',alphalims=[0.3,1.0]) + colsE = GRAPHICStools.listColors() + colsI = GRAPHICStools.listColors() for i in range(self.radii - 1): label = f"r/a={self.roa[0, i + 1]:.4f}" @@ -4473,26 +4465,17 @@ def plotConvergence(self, fig1=None): ax.set_xlim(left=0) ax.set_ylabel("Individual Residuals (GB)") ax.set_yscale("log") - # ax.legend(loc='best',prop={'size':5}) - GRAPHICStools.addLegendApart( - ax, ratio=0.9, withleg=True, extraPad=0, loc="center left", size=6 - ) - # ax2 = GRAPHICStools.addXaxis(ax,self.iterations,self.calls_solver,label='Calls to transport solver',whichticks=whichticks) - + GRAPHICStools.addLegendApart(ax, ratio=0.9, withleg=True, extraPad=0, loc="center left", size=6) GRAPHICStools.addDenseAxis(ax) - # GRAPHICStools.autoscale_y(ax) ax = ax01 - ax.plot( - self.iterations, self.residual_manual_real, "-s", color="b", markersize=5 - ) + ax.plot(self.iterations, self.residual_manual_real, "-s", color="b", markersize=5) ax.set_xlabel("Iterations") ax.set_xlim(left=0) ax.set_ylabel("Residual (real)") ax.set_yscale("log") GRAPHICStools.addDenseAxis(ax) - # GRAPHICStools.autoscale_y(ax) _ = GRAPHICStools.addXaxis( ax, @@ -4536,11 +4519,7 @@ def plotConvergence(self, fig1=None): ax.set_ylabel("Individual Residuals (real)") ax.set_yscale("log") - # ax2 = GRAPHICStools.addXaxis(ax,self.iterations,self.calls_solver,label='Calls to transport solver',whichticks=whichticks) - GRAPHICStools.addDenseAxis(ax) - # GRAPHICStools.autoscale_y(ax) - def plotAll(TGYROoutputs, labels=None, fn=None): if fn is None: diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 41b831bc..7b136795 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1952,6 +1952,9 @@ def plot_gradients(self, *args, **kwargs): def plot_geometry(self, *args, **kwargs): pass + def plot_flows(self, *args, **kwargs): + return state_plotting.plot_flows(self, *args, **kwargs) + def plotPeaking( self, ax, c="b", marker="*", label="", debugPlot=False, printVals=False ): From c283077db48ccb6ec249adf937f11296d4e295d7 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 13:42:36 -0400 Subject: [PATCH 140/385] Misc informations --- src/mitim_modules/powertorch/STATEtools.py | 6 ++++-- .../powertorch/physics_models/targets_analytic.py | 4 +--- src/mitim_tools/opt_tools/STRATEGYtools.py | 4 +--- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index a6a9b8a0..7770c3cd 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -756,15 +756,17 @@ def volume_integrate(self, var, force_dim=None): """ If var in MW/m^3, this gives as output the MW/m^2 profile """ + + surface_used = self.plasma["volp"] # IMPORTANT Note: This is the GACODE definition, acknowledging that volp=dV/dr is not equal to the surface area if force_dim is None: return CALCtools.volume_integration( var, self.plasma["rmin"], self.plasma["volp"] - ) / self.plasma["volp"] + ) / surface_used else: return CALCtools.volume_integration( var, self.plasma["rmin"][0,:].repeat(force_dim,1), self.plasma["volp"][0,:].repeat(force_dim,1) - ) / self.plasma["volp"][0,:].repeat(force_dim,1) + ) / surface_used[0,:].repeat(force_dim,1) def add_axes_powerstate_plot(figMain, num_kp=3): diff --git a/src/mitim_modules/powertorch/physics_models/targets_analytic.py b/src/mitim_modules/powertorch/physics_models/targets_analytic.py index 6f3956ad..418065db 100644 --- a/src/mitim_modules/powertorch/physics_models/targets_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/targets_analytic.py @@ -96,9 +96,7 @@ def _evaluate_alpha_heating(self): for i in range(self.powerstate.plasma["ni"].shape[2]): c_a += (self.powerstate.plasma["ni"][..., i] / self.powerstate.plasma["ne"]) * self.powerstate.plasma["ions_set_Zi"][:,i].unsqueeze(-1) ** 2 * (Aalpha / self.powerstate.plasma["ions_set_mi"][:,i].unsqueeze(-1)) - W_crit = (self.powerstate.plasma["te"] * 1e3) * (4 * (Ae / Aalpha) ** 0.5 / (3 * pi**0.5 * c_a)) ** ( - -2.0 / 3.0 - ) # in eV + W_crit = (self.powerstate.plasma["te"] * 1e3) * (4 * (Ae / Aalpha) ** 0.5 / (3 * pi**0.5 * c_a)) ** (-2.0 / 3.0) # in eV frac_ai = sivukhin(Ealpha / W_crit) # This solves Eq 17 of Stix diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index b3787f13..d07518c0 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -1685,9 +1685,7 @@ def plotSurrogateOptimization(self, fig1=None, fig2=None, boStep=-1): info, boundsRaw = step.InfoOptimization, step.bounds bounds = torch.Tensor([boundsRaw[b] for b in boundsRaw]) - boundsThis = ( - info[0]["bounds"].cpu().numpy().transpose(1, 0) if "bounds" in info[0] else None - ) + boundsThis = info[0]["bounds"].cpu().numpy().transpose(1, 0) if "bounds" in info[0] else None # ---------------------------------------------------------------------- # Prep figures From 6c341fbbb88934962e7d77c1218914445c3a3dfa Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 13:51:16 -0400 Subject: [PATCH 141/385] Added capablity to limit job arrays --- src/mitim_tools/gacode_tools/utils/GACODErun.py | 5 +++++ src/mitim_tools/misc_tools/FARMINGtools.py | 4 +++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 9ef6c8f8..499e67d2 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -915,6 +915,7 @@ def runTGLF( name="", launchSlurm=True, attempts_execution=1, + max_jobs_at_once=None, ): """ launchSlurm = True -> Launch as a batch job in the machine chosen @@ -1015,6 +1016,7 @@ def runTGLF( # Slurm setup array_list = None + job_array_limit = None shellPreCommands = None shellPostCommands = None ntasks = total_cores_required @@ -1035,6 +1037,7 @@ def runTGLF( # Slurm setup array_list = None + job_array_limit = None shellPreCommands = None shellPostCommands = None ntasks = total_tglf_executions @@ -1064,6 +1067,7 @@ def runTGLF( array_list = ",".join(array_list) ntasks = 1 cpuspertask = cores_tglf + job_array_limit = max_jobs_at_once # Limit to this number at most running jobs at the same time # --------------------------------------------- # Execute @@ -1079,6 +1083,7 @@ def runTGLF( "name": name, "cpuspertask": cpuspertask, "job_array": array_list, + "job_array_limit": job_array_limit, #"nodes": 1, }, ) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 2fdfd2e6..dedb53ed 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -184,6 +184,7 @@ def run( self.folderExecution, modules_remote=self.machineSettings["modules"], job_array=self.slurm_settings["job_array"] if "job_array" in self.slurm_settings else None, + job_array_limit=self.slurm_settings["job_array_limit"] if "job_array_limit" in self.slurm_settings else None, folder_local=self.folder_local, shellPreCommands=self.shellPreCommands, shellPostCommands=self.shellPostCommands, @@ -1000,6 +1001,7 @@ def create_slurm_execution_files( cpuspertask=4, memory_req_by_job=None, job_array=None, + job_array_limit=None, # If job_array is not None, this is the limit of the array size at once nodes=None, label_log_files="", wait_until_sbatch=True, @@ -1081,7 +1083,7 @@ def create_slurm_execution_files( commandSBATCH.append(f"#SBATCH --time {time_com}") if job_array is not None: - commandSBATCH.append(f"#SBATCH --array={job_array}") + commandSBATCH.append(f"#SBATCH --array={job_array}{f'%{job_array_limit} ' if job_array_limit is not None else ''}") elif request_exclusive_node: commandSBATCH.append("#SBATCH --exclusive") From 75adf204398679b2825d216c1d6e259a3cac4f97 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 14:00:14 -0400 Subject: [PATCH 142/385] Naming convention for volume integrator --- src/mitim_modules/portals/PORTALStools.py | 32 +++++++++---------- src/mitim_modules/powertorch/STATEtools.py | 2 +- .../powertorch/scripts/calculateTargets.py | 4 +-- .../powertorch/scripts/compareWithTGYRO.py | 4 +-- .../powertorch/utils/POWERplot.py | 2 +- .../powertorch/utils/TARGETStools.py | 2 +- 6 files changed, 23 insertions(+), 23 deletions(-) diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index a7a08fb9..81389be2 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -202,22 +202,22 @@ def computeTurbExchangeIndividual(QieMWm3_tr_turb, powerstate): 2. Integrate ------------------------------------------------------------------------ qExch is in MW/m^3 - powerstate.volume_integrate produces in MW/m^2 + powerstate.from_density_to_flux produces in MW/m^2 """ # Add zeros at zero qExch = torch.cat((torch.zeros(QieMWm3_tr_turb.shape).to(QieMWm3_tr_turb)[..., :1], QieMWm3_tr_turb), dim=-1) - QieMWm3_tr_turb_integrated = powerstate.volume_integrate(qExch, force_dim=qExch.shape[0])[..., 1:] + QieMWm2_tr_turb = powerstate.from_density_to_flux(qExch, force_dim=qExch.shape[0])[..., 1:] """ 3. Go back to the original batching system ------------------------------------------------------------------------ E.g.: (batch1*batch2*batch3,dimR) -> (batch1,batch2,batch3,dimR) """ - QieMWm3_tr_turb_integrated = QieMWm3_tr_turb_integrated.view(tuple(shape_orig)) + QieMWm2_tr_turb = QieMWm2_tr_turb.view(tuple(shape_orig)) - return QieMWm3_tr_turb_integrated + return QieMWm2_tr_turb def GBfromXnorm(x, output, powerstate): # Decide, depending on the output here, which to use as normalization and at what location @@ -400,9 +400,9 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): # ------------------------------------------------------------------------- if PORTALSparameters["turbulent_exchange_as_surrogate"]: - QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) + QieMWm2_tr_turb = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) else: - QieMWm3_tr_turb_integrated = torch.zeros(dfT.shape).to(dfT) + QieMWm2_tr_turb = torch.zeros(dfT.shape).to(dfT) # ------------------------------------------------------------------------ # Go through each profile that needs to be predicted, calculate components @@ -438,9 +438,9 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): ----------------------------------------------------------------------------------- """ if var == "Qe": - cal0 = var_dict[f"{var}_tar"] + QieMWm3_tr_turb_integrated + cal0 = var_dict[f"{var}_tar"] + QieMWm2_tr_turb elif var == "Qi": - cal0 = var_dict[f"{var}_tar"] - QieMWm3_tr_turb_integrated + cal0 = var_dict[f"{var}_tar"] - QieMWm2_tr_turb else: cal0 = var_dict[f"{var}_tar"] @@ -533,11 +533,11 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): # ------------------------------------------------------------------------- if PORTALSparameters["turbulent_exchange_as_surrogate"]: - QieMWm3_tr_turb_integrated = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) - QieMWm3_tr_turb_integrated_stds = computeTurbExchangeIndividual(var_dict["Qie_tr_turb_stds"], powerstate) + QieMWm2_tr_turb = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) + QieMWm2_tr_turb_stds = computeTurbExchangeIndividual(var_dict["Qie_tr_turb_stds"], powerstate) else: - QieMWm3_tr_turb_integrated = torch.zeros(dfT.shape).to(dfT) - QieMWm3_tr_turb_integrated_stds = torch.zeros(dfT.shape).to(dfT) + QieMWm2_tr_turb = torch.zeros(dfT.shape).to(dfT) + QieMWm2_tr_turb_stds = torch.zeros(dfT.shape).to(dfT) # ------------------------------------------------------------------------ # Go through each profile that needs to be predicted, calculate components @@ -571,11 +571,11 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): ----------------------------------------------------------------------------------- """ if var == "Qe": - cal0 = var_dict[f"{var}_tar"] + QieMWm3_tr_turb_integrated - cal0E = (var_dict[f"{var}_tar_stds"] ** 2 + QieMWm3_tr_turb_integrated_stds**2) ** 0.5 + cal0 = var_dict[f"{var}_tar"] + QieMWm2_tr_turb + cal0E = (var_dict[f"{var}_tar_stds"] ** 2 + QieMWm2_tr_turb_stds**2) ** 0.5 elif var == "Qi": - cal0 = var_dict[f"{var}_tar"] - QieMWm3_tr_turb_integrated - cal0E = (var_dict[f"{var}_tar_stds"] ** 2 + QieMWm3_tr_turb_integrated_stds**2) ** 0.5 + cal0 = var_dict[f"{var}_tar"] - QieMWm2_tr_turb + cal0E = (var_dict[f"{var}_tar_stds"] ** 2 + QieMWm2_tr_turb_stds**2) ** 0.5 else: cal0 = var_dict[f"{var}_tar"] cal0E = var_dict[f"{var}_tar_stds"] diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 7770c3cd..e8dccaa9 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -752,7 +752,7 @@ def _concatenate_flux(plasma, profile_key, flux_key): self.plasma["S"] = self.plasma["P"] - self.plasma["P_tr"] self.plasma["residual"] = self.plasma["S"].abs().mean(axis=1, keepdim=True) - def volume_integrate(self, var, force_dim=None): + def from_density_to_flux(self, var, force_dim=None): """ If var in MW/m^3, this gives as output the MW/m^2 profile """ diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index f99b905e..f0370d58 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -101,13 +101,13 @@ def calculator( # ************************************ p.plasma["Pfus"] = ( - p.volume_integrate( + p.from_density_to_flux( (p.plasma["qfuse"] + p.plasma["qfusi"]) * 5.0 ) * p.plasma["volp"] )[..., -1] p.plasma["Prad"] = ( - p.volume_integrate(p.plasma["qrad"]) * p.plasma["volp"] + p.from_density_to_flux(p.plasma["qrad"]) * p.plasma["volp"] )[..., -1] p.profiles.derive_quantities() diff --git a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py index dabae966..ad798de8 100644 --- a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py +++ b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py @@ -99,7 +99,7 @@ label="TGYRO " + label, markersize=markersize, ) - P = s.volume_integrate(stateQuantity, dim=2) * s.plasma["volp"] + P = s.from_density_to_flux(stateQuantity, dim=2) * s.plasma["volp"] ax.plot( s.plasma["rho"][0], P[0], @@ -132,7 +132,7 @@ label="TGYRO " + label, markersize=markersize, ) - P = s.volume_integrate(stateQuantity, dim=2) * s.plasma["volp"] + P = s.from_density_to_flux(stateQuantity, dim=2) * s.plasma["volp"] ax.plot( s.plasma["rho"][0], P[0], diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index d8de6edd..458ac6e0 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -229,7 +229,7 @@ def plot_metrics_powerstates(axsM, powerstates, profiles=None, profiles_color='b x , y = [], [] for h in range(len(powerstates)): x.append(h) - Pfus = powerstates[h].volume_integrate( + Pfus = powerstates[h].from_density_to_flux( (powerstates[h].plasma["qfuse"] + powerstates[h].plasma["qfusi"]) * 5.0 ) * powerstates[h].plasma["volp"] y.append(Pfus[..., -1].item()) diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 2b0580d9..e53fd6ea 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -80,7 +80,7 @@ def flux_integrate(self): qi += self.powerstate.plasma["qfusi"] q = torch.cat((qe, qi)).to(qe) - self.P = self.powerstate.volume_integrate(q, force_dim=q.shape[0]) + self.P = self.powerstate.from_density_to_flux(q, force_dim=q.shape[0]) def coarse_grid(self): From e1928ace0a5b93ca14468c78968548dcfd908e8d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 15:29:58 -0400 Subject: [PATCH 143/385] bug fix --- src/mitim_tools/opt_tools/optimizers/botorch_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/opt_tools/optimizers/botorch_tools.py b/src/mitim_tools/opt_tools/optimizers/botorch_tools.py index efe757e3..4e8f6b6c 100644 --- a/src/mitim_tools/opt_tools/optimizers/botorch_tools.py +++ b/src/mitim_tools/opt_tools/optimizers/botorch_tools.py @@ -68,7 +68,7 @@ def __call__(self, x, *args, **kwargs): seq_message = f'({"sequential" if sequential_q else "joint"}) ' if q>1 else '' print(f"\t\t- Optimizing using optimize_acqf: {q = } {seq_message}, {num_restarts = }, {raw_samples = }") - with IOtools.timer(name = "\n\t- Optimization", name_timer = '\t\t- Time: '): + with IOtools.timer(name = "\n\t- Optimization"): x_opt, _ = botorch.optim.optimize_acqf( acq_function=fun_opt, bounds=fun.bounds_mod, From 497e250df7f8d6e92ab225c1056d03949dcedb0e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 16:50:19 -0400 Subject: [PATCH 144/385] Cleaned up handling of remote retrievals as common function --- .../maestro/scripts/plot_maestro.py | 56 +++---- .../maestro/scripts/run_maestro.py | 3 +- .../portals/scripts/read_portals.py | 97 ++++++----- src/mitim_tools/misc_tools/FARMINGtools.py | 14 +- src/mitim_tools/misc_tools/utils/__init__.py | 0 .../misc_tools/utils/remote_tools.py | 38 +++++ src/mitim_tools/opt_tools/scripts/read.py | 150 ++++-------------- 7 files changed, 162 insertions(+), 196 deletions(-) create mode 100644 src/mitim_tools/misc_tools/utils/__init__.py create mode 100644 src/mitim_tools/misc_tools/utils/remote_tools.py diff --git a/src/mitim_modules/maestro/scripts/plot_maestro.py b/src/mitim_modules/maestro/scripts/plot_maestro.py index 75a50adb..fb4c8177 100644 --- a/src/mitim_modules/maestro/scripts/plot_maestro.py +++ b/src/mitim_modules/maestro/scripts/plot_maestro.py @@ -1,7 +1,8 @@ import argparse from mitim_modules.maestro.utils import MAESTROplot -from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools, FARMINGtools +from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools from mitim_tools.opt_tools import STRATEGYtools +from mitim_tools.misc_tools.utils import remote_tools from pathlib import Path from IPython import embed @@ -38,6 +39,8 @@ def main(): # Standard options parser.add_argument("folders", type=str, nargs="*", help="Paths to the folders to read.") + + # MAESTRO specific options parser.add_argument("--beats", type=int, required=False, default=2, help="Number of beats to plot. If 0, it will not plot beat information.") parser.add_argument("--only", type=str, required=False, default=None, @@ -53,48 +56,39 @@ def main(): parser.add_argument("--remote_folders",type=str, nargs="*", required=False, default=None, help="List of folders in the remote machine to retrieve. If not provided, it will use the local folder structures.") parser.add_argument("--remote_minimal", required=False, default=False, action="store_true", - help="If set, it will only retrieve the folder structure with a few files (input.gacode, input.gacode_final, initializer_geqdsk/input.gacode).") + help="If set, it will only retrieve the folder structure with a few key files.") parser.add_argument('--fix', required=False, default=False, action='store_true', help="If set, it will fix the pkl optimization portals in the remote folders.") args = parser.parse_args() - remote = args.remote - folders = args.folders - fix = args.fix - beats = args.beats - only = args.only - full = args.full - if args.remote_folder_parent is not None: - folders_remote = [args.remote_folder_parent + '/' + folder.split('/')[-1] for folder in folders] - elif args.remote_folders is not None: - folders_remote = args.remote_folders - else: - folders_remote = folders - - - # Retrieve remote - if remote is not None: + # -------------------------------------------------------------------------------------------------------------------------------------------- + # Retrieve from remote + # -------------------------------------------------------------------------------------------------------------------------------------------- - only_folder_structure_with_files = None - if args.remote_minimal: - only_folder_structure_with_files = ["beat_results/input.gacode", "input.gacode_final","initializer_geqdsk/input.gacode", "timing.jsonl"] - - beats = 0 + only_folder_structure_with_files = None + if args.remote_minimal: + only_folder_structure_with_files = ["Outputs/optimization_data.csv","Outputs/optimization_extra.pkl","Outputs/optimization_object.pkl","Outputs/optimization_results.out"] - _, folders = FARMINGtools.retrieve_files_from_remote( - IOtools.expandPath('./'), - remote, - folders_remote = folders_remote, - purge_tmp_files = True, - only_folder_structure_with_files=only_folder_structure_with_files) + folders = remote_tools.retrieve_remote_folders(args.folders, args.remote, args.remote_folder_parent, args.remote_folders, only_folder_structure_with_files) + + # -------------------------------------------------------------------------------------------------------------------------------------------- # Fix pkl optimization portals in remote - if fix: + # -------------------------------------------------------------------------------------------------------------------------------------------- + + if args.fix: fix_maestro([Path(folder) for folder in folders]) - # ----- + # -------------------------------------------------------------------------------------------------------------------------------------------- + # Actual interpreting and plotting + # -------------------------------------------------------------------------------------------------------------------------------------------- + + beats = args.beats + only = args.only + full = args.full + folders = [IOtools.expandPath(folder) for folder in folders] diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index 1eb2643a..38b88f40 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -254,7 +254,8 @@ def main(): if not folder.exists(): folder.mkdir(parents=True, exist_ok=True) - IOtools.recursive_backup(folder / 'maestro_namelist.json') + if (folder / 'maestro_namelist.json').exists(): + IOtools.recursive_backup(folder / 'maestro_namelist.json') run_maestro_local(*parse_maestro_nml(file_path),folder=folder,cpus = cpus, terminal_outputs = terminal_outputs) diff --git a/src/mitim_modules/portals/scripts/read_portals.py b/src/mitim_modules/portals/scripts/read_portals.py index 01b535ad..3d669f7a 100644 --- a/src/mitim_modules/portals/scripts/read_portals.py +++ b/src/mitim_modules/portals/scripts/read_portals.py @@ -1,64 +1,74 @@ import argparse import matplotlib.pyplot as plt -from mitim_tools.misc_tools import IOtools from mitim_modules.portals.utils import PORTALSanalysis +from mitim_tools.misc_tools import IOtools +from mitim_tools.opt_tools import STRATEGYtools +from mitim_tools.misc_tools.utils import remote_tools from IPython import embed -""" -This script is to plot only the convergence figure, not the rest of surrogates that takes long. -It also does it on a separate figure, so easy to manage (e.g. for saving as .eps) -""" def main(): parser = argparse.ArgumentParser() - parser.add_argument("folders", type=str, nargs="*") - parser.add_argument("--remote", "-r", type=str, required=False, default=None) + + # Standard options + parser.add_argument("folders", type=str, nargs="*", + help="Paths to the folders to read.") - 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) + # PORTALS specific options + 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("--indeces_extra", type=int, required=False, default=[], nargs="*") - parser.add_argument( - "--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" - ) + parser.add_argument("--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") + + # Remote options + parser.add_argument("--remote",type=str, required=False, default=None, + help="Remote machine to retrieve the folders from. If not provided, it will read the local folders.") + parser.add_argument("--remote_folder_parent",type=str, required=False, default=None, + help="Parent folder in the remote machine where the folders are located. If not provided, it will use --remote_folders.") + parser.add_argument("--remote_folders",type=str, nargs="*", required=False, default=None, + help="List of folders in the remote machine to retrieve. If not provided, it will use the local folder structures.") + parser.add_argument("--remote_minimal", required=False, default=False, action="store_true", + help="If set, it will only retrieve the folder structure with a few key files.") + parser.add_argument('--fix', required=False, default=False, action='store_true', + help="If set, it will fix the pkl optimization portals in the remote folders.") args = parser.parse_args() - folders = [IOtools.expandPath(folder) for folder in args.folders] + # -------------------------------------------------------------------------------------------------------------------------------------------- + # Retrieve from remote + # -------------------------------------------------------------------------------------------------------------------------------------------- - portals_total = [] - for folderWork in folders: - folderRemote_reduced = args.remote - file = args.file - indexToMaximize = args.max - indeces_extra = args.indeces_extra - plotAllFluxes = args.all - complete = args.complete + only_folder_structure_with_files = None + if args.remote_minimal: + only_folder_structure_with_files = ["Outputs/optimization_data.csv","Outputs/optimization_extra.pkl","Outputs/optimization_object.pkl","Outputs/optimization_results.out"] + + folders = remote_tools.retrieve_remote_folders(args.folders, args.remote, args.remote_folder_parent, args.remote_folders, only_folder_structure_with_files) - if not folderWork.exists(): - folderWork.mkdir(parents=True, exist_ok=True) + # -------------------------------------------------------------------------------------------------------------------------------------------- + # Fix pkl optimization portals in remote + # -------------------------------------------------------------------------------------------------------------------------------------------- - folderRemote = ( - f"{folderRemote_reduced}/{IOtools.reducePathLevel(folderWork)[-1]}/" - if folderRemote_reduced is not None - else None - ) + if args.fix: + for folder in folders: + STRATEGYtools.clean_state(folder) - # Read PORTALS - portals = PORTALSanalysis.PORTALSanalyzer.from_folder( - folderWork, folderRemote=folderRemote - ) + # -------------------------------------------------------------------------------------------------------------------------------------------- + # PORTALS reading + # -------------------------------------------------------------------------------------------------------------------------------------------- + + portals_total = [PORTALSanalysis.PORTALSanalyzer.from_folder(folderWork) for folderWork in folders] - portals_total.append(portals) + # -------------------------------------------------------------------------------------------------------------------------------------------- + # Actual PORTALS plotting + # -------------------------------------------------------------------------------------------------------------------------------------------- - # PLOTTING + file = args.file + indexToMaximize = args.max + indeces_extra = args.indeces_extra + plotAllFluxes = args.all + complete = args.complete if not complete: size = 8 @@ -87,9 +97,7 @@ def main(): portals_total[i].fn = fn # Plot metrics - if (not complete) or ( - isinstance(portals_total[i], PORTALSanalysis.PORTALSinitializer) - ): + if (not complete) or isinstance(portals_total[i], PORTALSanalysis.PORTALSinitializer): if isinstance(portals_total[i], PORTALSanalysis.PORTALSinitializer): fig = None elif requiresFN: @@ -116,5 +124,6 @@ def main(): plt.show() embed() + if __name__ == "__main__": main() diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index dedb53ed..9a23c753 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -1283,7 +1283,19 @@ def retrieve_files_from_remote( output_files.append(file0) for folder in folders_remote: folder0 = f'{IOtools.expandPath(folder)}'.split('/')[-1] - command += f'cp -r {folder} {machineSettings["folderWork"]}/{folder0}\n' + + folder_source = folder + folder_destination = f'{machineSettings["folderWork"]}/{folder0}' + if only_folder_structure_with_files is None: + # Normal full copy + command += f'cp -r {folder_source} {folder_destination}\n' + else: + retrieve_files = '' + for file in only_folder_structure_with_files: + retrieve_files += f'-f"+ {file}" ' + # Only copy the folder structure with a few files + command += f'rsync -av -f"+ */" {retrieve_files}-f"- *" {folder_source}/ {folder_destination}/\n' + output_folders.append(folder0) # ------------------------------------------------ diff --git a/src/mitim_tools/misc_tools/utils/__init__.py b/src/mitim_tools/misc_tools/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/mitim_tools/misc_tools/utils/remote_tools.py b/src/mitim_tools/misc_tools/utils/remote_tools.py new file mode 100644 index 00000000..b50a74c5 --- /dev/null +++ b/src/mitim_tools/misc_tools/utils/remote_tools.py @@ -0,0 +1,38 @@ +import os +from mitim_tools.misc_tools import IOtools, FARMINGtools +from IPython import embed + +def retrieve_remote_folders(folders_local, remote, remote_folder_parent, remote_folders, only_folder_structure_with_files): + + # Make sure folders_local is a list of complete Paths + folders_local = [IOtools.expandPath(folder).resolve() for folder in folders_local] + + if remote_folder_parent is not None: + folders_remote = [remote_folder_parent + '/' + folder.split('/')[-1] for folder in folders_local] + elif remote_folders is not None: + folders_remote = remote_folders + else: + folders_remote = folders_local + + # Retrieve remote + if remote is not None: + + _, folders = FARMINGtools.retrieve_files_from_remote( + IOtools.expandPath('./'), + remote, + folders_remote = folders_remote, + purge_tmp_files = True, + only_folder_structure_with_files=only_folder_structure_with_files) + + # Renaming + for i in range(len(folders)): + folder = IOtools.expandPath(folders[i]) + folder_orig = IOtools.expandPath(folders_local[i]) + + if folder_orig.exists(): + IOtools.shutil_rmtree(folder_orig) + + os.rename(folder, folder_orig) + + + return folders_local \ No newline at end of file diff --git a/src/mitim_tools/opt_tools/scripts/read.py b/src/mitim_tools/opt_tools/scripts/read.py index f2edeec9..939ef3ba 100644 --- a/src/mitim_tools/opt_tools/scripts/read.py +++ b/src/mitim_tools/opt_tools/scripts/read.py @@ -5,6 +5,8 @@ from mitim_tools.misc_tools import IOtools, GRAPHICStools from mitim_tools.opt_tools import STRATEGYtools from mitim_tools.misc_tools.LOGtools import printMsg as print +from mitim_tools.misc_tools.utils import remote_tools + # These import are usually needed if they are called within the pickling object import torch @@ -56,91 +58,8 @@ def plotCompare(folders, plotMeanMax=[True, False]): ax3 = fig.add_subplot(grid[1, 1]) ax1i = fig.add_subplot(grid[2, 0], sharex=ax0) - types_ls = [ - "-", - "--", - "-.", - ":", - "-", - "--", - "-.", - ":", - "-", - "--", - "-.", - ":", - "-", - "--", - "-.", - ":", - "-", - "--", - "-.", - ":", - "-", - "--", - "-.", - ":", - "-", - "--", - "-.", - ":", - "-", - "--", - "-.", - ":", - "-", - "--", - "-.", - ":", - ] - types_m = [ - "o", - "s", - "^", - "v", - "*", - "o", - "s", - "^", - "v", - "*", - "o", - "s", - "^", - "v", - "*", - "o", - "s", - "^", - "v", - "*", - "o", - "s", - "^", - "v", - "*", - "o", - "s", - "^", - "v", - "*", - "o", - "s", - "^", - "v", - "*", - "o", - "s", - "^", - "v", - "*", - "o", - "s", - "^", - "v", - "*", - ] + types_ls = GRAPHICStools.listLS() + types_m = GRAPHICStools.listmarkers maxEv = -np.inf yCummMeans = [] @@ -202,8 +121,6 @@ def plotCompare(folders, plotMeanMax=[True, False]): def main(): - - # ----- Inputs parser = argparse.ArgumentParser() @@ -211,43 +128,49 @@ def main(): "--type", type=int, required=False, default=4 ) # 0: Only ResultsOpt plotting, 1: Also pickle, 2: Also final analysis, 3: Others 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) parser.add_argument("--conv", type=float, required=False, default=-1e-2) parser.add_argument("--its", type=int, nargs="*", required=False, default=None) + # Remote options + parser.add_argument("--remote",type=str, required=False, default=None, + help="Remote machine to retrieve the folders from. If not provided, it will read the local folders.") + parser.add_argument("--remote_folder_parent",type=str, required=False, default=None, + help="Parent folder in the remote machine where the folders are located. If not provided, it will use --remote_folders.") + parser.add_argument("--remote_folders",type=str, nargs="*", required=False, default=None, + help="List of folders in the remote machine to retrieve. If not provided, it will use the local folder structures.") + # parser.add_argument("--remote_minimal", required=False, default=False, action="store_true", + # help="If set, it will only retrieve the folder structure with a few key files") + parser.add_argument('--fix', required=False, default=False, action='store_true', + help="If set, it will fix the pkl optimization portals in the remote folders.") + args = parser.parse_args() analysis_level = args.type - folders = args.folders - remote_parent = args.remote seeds = args.seeds resolution = args.resolution save_folder = args.save conv = args.conv rangePlot = args.its -# ----------------------------------------- + # -------------------------------------------------------------------------------------------------------------------------------------------- + # Retrieve from remote + # -------------------------------------------------------------------------------------------------------------------------------------------- - # ----- Folders (complete local path) - folders_complete = [] - for i in range(len(folders)): - if seeds is not None: - aux = [f"{folders[i]}_s{k}" for k in range(seeds)] - folders_complete.extend(aux) - else: - folders_complete.append(folders[i]) + folders = remote_tools.retrieve_remote_folders(args.folders, args.remote, args.remote_folder_parent, args.remote_folders, None) - txt = "***************************************************************************\n" - for i in range(len(folders_complete)): - folders_complete[i] = IOtools.expandPath(folders_complete[i]) - folders_complete[i].mkdir(parents=True, exist_ok=True) - txt += f"* Reading results in {folders_complete[i]}\n" - - # ----- Folders (reduced local path) - folders_reduced = [IOtools.reducePathLevel(folderWork)[-1] for folderWork in folders_complete] + # -------------------------------------------------------------------------------------------------------------------------------------------- + # Fix pkl optimization portals in remote + # -------------------------------------------------------------------------------------------------------------------------------------------- + + if args.fix: + for folder in folders: + STRATEGYtools.clean_state(folder) + + + folders_complete = folders if len(folders_complete) > 1: retrieval_level = copy.deepcopy(analysis_level) @@ -255,23 +178,13 @@ def main(): else: retrieval_level = analysis_level - if remote_parent is None: - folders_remote = [None] * len(folders_complete) - else: - folders_remote = [ - f"{remote_parent}/{reduced_folder}/" - for reduced_folder in folders_reduced - ] - txt += f"\n\t...From remote folder {remote_parent}\n" - print("\n"+ txt+ "***************************************************************************") print(f"(Analysis level {analysis_level})\n") if len(folders_complete) == 1: opt_fun = STRATEGYtools.opt_evaluator(folders_complete[0]) opt_fun.plot_optimization_results( analysis_level=analysis_level, - folderRemote=folders_remote[0], retrieval_level=retrieval_level, pointsEvaluateEachGPdimension=resolution, save_folder=save_folder, @@ -279,12 +192,11 @@ def main(): ) else: opt_funs = [] - for folderWork, folderRemote in zip(folders_complete, folders_remote): + for folderWork in folders_complete: opt_fun = STRATEGYtools.opt_evaluator(folderWork) try: opt_fun.plot_optimization_results( analysis_level=analysis_level, - folderRemote=folderRemote, retrieval_level=retrieval_level, save_folder=save_folder, rangesPlot=rangePlot, From 1f9999846671807b7da68659f2945d5a6c898db0 Mon Sep 17 00:00:00 2001 From: pabloprf Date: Wed, 6 Aug 2025 17:33:06 -0400 Subject: [PATCH 145/385] Bug fixes maestro after mitimstate transition --- src/mitim_modules/maestro/MAESTROmain.py | 6 +++--- .../maestro/scripts/run_maestro.py | 2 +- src/mitim_modules/maestro/utils/EPEDbeat.py | 14 +++++++------- src/mitim_modules/maestro/utils/TRANSPbeat.py | 18 +++++++++--------- src/mitim_tools/gacode_tools/PROFILEStools.py | 2 +- src/mitim_tools/gs_tools/GEQtools.py | 3 ++- .../plasmastate_tools/MITIMstate.py | 2 +- 7 files changed, 24 insertions(+), 23 deletions(-) diff --git a/src/mitim_modules/maestro/MAESTROmain.py b/src/mitim_modules/maestro/MAESTROmain.py index ceabcab2..899a6c33 100644 --- a/src/mitim_modules/maestro/MAESTROmain.py +++ b/src/mitim_modules/maestro/MAESTROmain.py @@ -217,7 +217,7 @@ def prepare(self, *args, **kwargs): self.beat.initialize() # ----------------------------- - self.beat.profiles_current.deriveQuantities() + self.beat.profiles_current.derive_quantities() self.beat.prepare(*args, **kwargs) else: @@ -269,7 +269,7 @@ def _freeze_parameters(self, profiles = None): print('\t\t- Freezing engineering parameters from MAESTRO') self.profiles_with_engineering_parameters = copy.deepcopy(profiles) - self.profiles_with_engineering_parameters.writeCurrentStatus(file= (self.folder_output / 'input.gacode_frozen')) + self.profiles_with_engineering_parameters.write_state(file= (self.folder_output / 'input.gacode_frozen')) @mitim_timer( lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Finalizing', @@ -283,7 +283,7 @@ def finalize(self): final_file= (self.folder_output / 'input.gacode_final') - self.beat.profiles_output.writeCurrentStatus(file= final_file) + self.beat.profiles_output.write_state(file= final_file) print(f'\t\t- Final input.gacode saved to {IOtools.clipstr(final_file)}') diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index 38b88f40..ab53fff4 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -144,7 +144,7 @@ def profiles_postprocessing_fun(file_profs): p.lumpImpurities() if enforce_same_density_gradients: p.enforce_same_density_gradients() - p.writeCurrentStatus(file=file_profs) + p.write_state(file=file_profs) beat_namelist['PORTALSparameters']['profiles_postprocessing_fun'] = profiles_postprocessing_fun else: diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index e57509c0..d6f1df59 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -321,7 +321,7 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F for key in eped_results: print(f'\t\t- {key}: {eped_results[key]}') - self.profiles_output.writeCurrentStatus(file=self.folder / 'input.gacode.eped') + self.profiles_output.write_state(file=self.folder / 'input.gacode.eped') return eped_results @@ -365,7 +365,7 @@ def finalize(self, **kwargs): self.profiles_output = PROFILEStools.PROFILES_GACODE(self.folder / 'input.gacode.eped') - self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') + self.profiles_output.write_state(file=self.folder_output / 'input.gacode') def merge_parameters(self): # EPED beat does not modify the profiles grid or anything, so I can keep it fine @@ -560,12 +560,12 @@ def scale_profile_by_stretching( x, y, xp, yp, xp_old, plotYN=False, label='', k print('\t\t\t* Keeping old aLT profile in the core-predicted region, using r/a for it') # Calculate gradient in entire region - aLy = CALCtools.produceGradient( torch.from_numpy(roa), torch.from_numpy(y) ) + aLy = CALCtools.derivation_into_Lx( torch.from_numpy(roa), torch.from_numpy(y) ) # I'm only interested in core region, plus one ghost point with the same gradient aLy = torch.cat( (aLy[:ibc+1], aLy[ibc].unsqueeze(0)) ) - y_mod = CALCtools.integrateGradient( torch.from_numpy(roa[:ibc+2]).unsqueeze(0), aLy.unsqueeze(0), torch.from_numpy(np.array(ynew[ibc+1])).unsqueeze(0) ).squeeze().numpy() + y_mod = CALCtools.integration_Lx( torch.from_numpy(roa[:ibc+2]).unsqueeze(0), aLy.unsqueeze(0), torch.from_numpy(np.array(ynew[ibc+1])).unsqueeze(0) ).squeeze().numpy() ynew[:ibc+2] = y_mod @@ -583,11 +583,11 @@ def scale_profile_by_stretching( x, y, xp, yp, xp_old, plotYN=False, label='', k ax.legend() ax = axs[1] - aLy = CALCtools.produceGradient( torch.from_numpy(roa), torch.from_numpy(y) ) + aLy = CALCtools.derivation_into_Lx( torch.from_numpy(roa), torch.from_numpy(y) ) ax.plot(x,aLy,'-o',color='b', label='old') ax.axvline(x=xp_old,color='b',ls='--') - aLy = CALCtools.produceGradient( torch.from_numpy(roa), torch.from_numpy(ynew) ) + aLy = CALCtools.derivation_into_Lx( torch.from_numpy(roa), torch.from_numpy(ynew) ) ax.plot(x,aLy,'-o',color='r', label='new') ax.axvline(x=xp,color='r',ls='--') @@ -671,6 +671,6 @@ def eped_profiler(profiles, xp_old, rhotop, Tetop_keV, Titop_keV, netop_20, mini # Re-derive # --------------------------------- - profiles_output.deriveQuantities(rederiveGeometry=False) + profiles_output.derive_quantities(rederiveGeometry=False) return profiles_output \ No newline at end of file diff --git a/src/mitim_modules/maestro/utils/TRANSPbeat.py b/src/mitim_modules/maestro/utils/TRANSPbeat.py index 74ed7245..bd873022 100644 --- a/src/mitim_modules/maestro/utils/TRANSPbeat.py +++ b/src/mitim_modules/maestro/utils/TRANSPbeat.py @@ -104,7 +104,7 @@ def prepare( self._additional_operations_add_initialization() # ICRF on - PichT_MW = self.profiles_current.derived['qRF_MWmiller'][-1] + PichT_MW = self.profiles_current.derived['qRF_MW'][-1] if freq_ICH is None: @@ -181,18 +181,18 @@ def finalize(self, force_auxiliary_heating_at_output = {'Pe': None, 'Pi': None}, self._add_heating_profiles(force_auxiliary_heating_at_output) # Write profiles - self.profiles_output.writeCurrentStatus(file=self.folder_output / "input.gacode") + self.profiles_output.write_state(file=self.folder_output / "input.gacode") def _add_heating_profiles(self, force_auxiliary_heating_at_output = {'Pe': None, 'Pi': None}): ''' force_auxiliary_heating_at_output['Pe'] has the shaping function (takes rho) and the integrated value ''' - for key, pkey, ikey in zip(['Pe','Pi'], ['qrfe(MW/m^3)', 'qrfi(MW/m^3)'], ['qRFe_MWmiller', 'qRFi_MWmiller']): + for key, pkey, ikey in zip(['Pe','Pi'], ['qrfe(MW/m^3)', 'qrfi(MW/m^3)'], ['qRFe_MW', 'qRFi_MW']): if force_auxiliary_heating_at_output[key] is not None: self.profiles_output.profiles[pkey] = force_auxiliary_heating_at_output[key][0](self.profiles_output.profiles['rho(-)']) - self.profiles_output.deriveQuantities() + self.profiles_output.derive_quantities() self.profiles_output.profiles[pkey] = self.profiles_output.profiles[pkey] * force_auxiliary_heating_at_output[key][1]/self.profiles_output.derived[ikey][-1] def merge_parameters(self): @@ -213,7 +213,7 @@ def merge_parameters(self): # Write the pre-merge input.gacode before modifying it profiles_output_pre_merge = copy.deepcopy(self.profiles_output) - profiles_output_pre_merge.writeCurrentStatus(file=self.folder_output / 'input.gacode_pre_merge') + profiles_output_pre_merge.write_state(file=self.folder_output / 'input.gacode_pre_merge') # First, bring back to the resolution of the frozen p_frozen = self.maestro_instance.profiles_with_engineering_parameters @@ -237,14 +237,14 @@ def merge_parameters(self): self.profiles_output.profiles[key] = p_frozen.profiles[key] # Power scale - self.profiles_output.profiles['qrfe(MW/m^3)'] *= p_frozen.derived['qRF_MWmiller'][-1] / self.profiles_output.derived['qRF_MWmiller'][-1] - self.profiles_output.profiles['qrfi(MW/m^3)'] *= p_frozen.derived['qRF_MWmiller'][-1] / self.profiles_output.derived['qRF_MWmiller'][-1] + self.profiles_output.profiles['qrfe(MW/m^3)'] *= p_frozen.derived['qRF_MW'][-1] / self.profiles_output.derived['qRF_MW'][-1] + self.profiles_output.profiles['qrfi(MW/m^3)'] *= p_frozen.derived['qRF_MW'][-1] / self.profiles_output.derived['qRF_MW'][-1] # -------------------------------------------------------------------------------------------- # Write to final input.gacode - self.profiles_output.deriveQuantities() - self.profiles_output.writeCurrentStatus(file=self.folder_output / 'input.gacode') + self.profiles_output.derive_quantities() + self.profiles_output.write_state(file=self.folder_output / 'input.gacode') def grab_output(self): diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 2382682a..8c870155 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -35,7 +35,7 @@ def __init__(self, file, derive_quantities=True, mi_ref=None): self.titles_singleArr = ["masse","mass","ze","z","torfluxa(Wb/radian)","rcentr(m)","bcentr(T)","current(MA)"] self.titles_single = self.titles_singleNum + self.titles_singleArr - if self.files is not None: + if self.files[0] is not None: self._read_inputgacocde() diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 3c24cd5d..24e76123 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -5,6 +5,7 @@ import numpy as np import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, PLASMAtools +from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.gs_tools.utils import GEQplotting from shapely.geometry import LineString @@ -418,7 +419,7 @@ def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH _, profiles["qrfe(MW/m^3)"] = PLASMAtools.parabolicProfile(Tbar=1.0,nu=5.0,rho=rhotor,Tedge=0.0) - p = MITIMstate.mitim_state.scratch(profiles) + p = PROFILEStools.gacode_state.scratch(profiles) p.profiles["qrfe(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] * PichT/p.derived['qRF_MW'][-1] /2 p.profiles["qrfi(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 7b136795..9c656ada 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -119,7 +119,7 @@ def scratch(cls, profiles, label_header='', **kwargs_process): # Add data to profiles instance.profiles = profiles - instance.process(**kwargs_process) + instance.derive_quantities(**kwargs_process) return instance From 11768682a608d681657eca1cfaf4b9418889df5f Mon Sep 17 00:00:00 2001 From: pabloprf Date: Wed, 6 Aug 2025 17:52:38 -0400 Subject: [PATCH 146/385] bug fix --- src/mitim_tools/gacode_tools/PROFILEStools.py | 8 +++++++- src/mitim_tools/gs_tools/GEQtools.py | 1 - 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 8c870155..862209ec 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -1,5 +1,4 @@ import copy -from turtle import color import numpy as np import matplotlib.pyplot as plt from collections import OrderedDict @@ -86,7 +85,11 @@ def _read_inputgacocde(self): they also say that those units are wrong. """ + + self._ensure_shaping_coeffs() + def _ensure_shaping_coeffs(self): + # Ensure that we also have the shape coefficients num_moments = 7 # This is the max number of moments I'll be considering. If I don't have that many (usually there are 5 or 3), it'll be populated with zeros if "shape_cos0(-)" not in self.profiles: @@ -176,6 +179,9 @@ def derive_quantities(self, **kwargs): super().derive_quantities_base(**kwargs) def _produce_shape_lists(self): + + self._ensure_shaping_coeffs() + self.shape_cos = [ self.profiles["shape_cos0(-)"], # tilt self.profiles["shape_cos1(-)"], diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 24e76123..5c71f12d 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -6,7 +6,6 @@ import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, PLASMAtools from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.plasmastate_tools import MITIMstate from mitim_tools.gs_tools.utils import GEQplotting from shapely.geometry import LineString from scipy.integrate import quad From b7b6217aaeb6d8519c17bccc16db8dc9ac515b4d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 17:54:09 -0400 Subject: [PATCH 147/385] maestro plot bug fix --- src/mitim_modules/maestro/utils/MAESTROplot.py | 5 +++-- .../misc_tools/utils/remote_tools.py | 18 +++++++++--------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index 76f6812f..82f07871 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -1,6 +1,7 @@ import numpy as np from collections import OrderedDict from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_tools.misc_tools import LOGtools, GRAPHICStools from mitim_tools.gs_tools import GEQtools from pathlib import Path @@ -102,10 +103,10 @@ def plot_results(self, fn): maxPlot = 5 if len(ps) > 0: # Plot profiles - figs = PROFILEStools.add_figures(fn,fnlab_pre = 'MAESTRO - ') + figs = state_plotting.add_figures(fn,fnlab_pre = 'MAESTRO - ') log_file = self.folder_logs/'plot_maestro.log' if (not self.terminal_outputs) else None with LOGtools.conditional_log_to_file(log_file=log_file): - PROFILEStools.plotAll(ps[-maxPlot:], extralabs=ps_lab[-maxPlot:], figs=figs) + state_plotting.plotAll(ps[-maxPlot:], extralabs=ps_lab[-maxPlot:], figs=figs) for p,pl in zip(ps,ps_lab): p.printInfo(label = pl) diff --git a/src/mitim_tools/misc_tools/utils/remote_tools.py b/src/mitim_tools/misc_tools/utils/remote_tools.py index b50a74c5..32826203 100644 --- a/src/mitim_tools/misc_tools/utils/remote_tools.py +++ b/src/mitim_tools/misc_tools/utils/remote_tools.py @@ -24,15 +24,15 @@ def retrieve_remote_folders(folders_local, remote, remote_folder_parent, remote_ purge_tmp_files = True, only_folder_structure_with_files=only_folder_structure_with_files) - # Renaming - for i in range(len(folders)): - folder = IOtools.expandPath(folders[i]) - folder_orig = IOtools.expandPath(folders_local[i]) - - if folder_orig.exists(): - IOtools.shutil_rmtree(folder_orig) - - os.rename(folder, folder_orig) + # Renaming + for i in range(len(folders)): + folder = IOtools.expandPath(folders[i]) + folder_orig = IOtools.expandPath(folders_local[i]) + if folder_orig.exists(): + IOtools.shutil_rmtree(folder_orig) + + os.rename(folder, folder_orig) + return folders_local \ No newline at end of file From e9bdffa7fe9bc0decbfb9eaee484011ea9e7b07d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 18:36:27 -0400 Subject: [PATCH 148/385] Moved logger from BO to MAESTRO timings and removal of historical Logger object --- .../maestro/scripts/plot_maestro.py | 2 +- .../maestro/utils/MAESTROplot.py | 112 +------- src/mitim_tools/misc_tools/GRAPHICStools.py | 1 + src/mitim_tools/misc_tools/IOtools.py | 116 ++++++++ src/mitim_tools/opt_tools/STRATEGYtools.py | 159 +++++------ src/mitim_tools/opt_tools/scripts/read.py | 10 +- src/mitim_tools/opt_tools/utils/BOgraphics.py | 248 +----------------- .../opt_tools/utils/EVALUATORtools.py | 1 - 8 files changed, 209 insertions(+), 440 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/plot_maestro.py b/src/mitim_modules/maestro/scripts/plot_maestro.py index fb4c8177..a41e30bc 100644 --- a/src/mitim_modules/maestro/scripts/plot_maestro.py +++ b/src/mitim_modules/maestro/scripts/plot_maestro.py @@ -127,7 +127,7 @@ def main(): if len(folders) > 1: MAESTROplot.plot_special_quantities(ps, ps_lab, axsAll, color=colors[i], label = f'Case #{i}', legYN = i==0) if (m.folder_performance / 'timing.jsonl').exists(): - x0, scripts0 = MAESTROplot.plot_timings(m.folder_performance / 'timing.jsonl', axs = axsTiming, label = f'Case #{i}', color=colors[i]) + x0, scripts0 = IOtools.plot_timings(m.folder_performance / 'timing.jsonl', axs = axsTiming, label = f'Case #{i}', color=colors[i]) # Only keep the longest if len(x0) > len(x): diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index 82f07871..6777741c 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -2,7 +2,7 @@ from collections import OrderedDict from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.plasmastate_tools.utils import state_plotting -from mitim_tools.misc_tools import LOGtools, GRAPHICStools +from mitim_tools.misc_tools import LOGtools, GRAPHICStools, IOtools from mitim_tools.gs_tools import GEQtools from pathlib import Path from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -212,7 +212,7 @@ def plot_results(self, fn): A B """) - plot_timings(self.folder_performance / 'timing.jsonl', axs = axs) + IOtools.plot_timings(self.folder_performance / 'timing.jsonl', axs = axs) return ps, ps_lab @@ -366,112 +366,4 @@ def plot_g_quantities(g, axs, color = 'b', lw = 1, ms = 0): axs[3].plot(g.g['RHOVN'], g.g['PRES']*1E-6, '-o', markersize=ms, lw = lw, label='Initial geqdsk', color=color) axs[4].plot(g.g['RHOVN'], g.g['QPSI'], '-o', markersize=ms, lw = lw, label='Initial geqdsk', color=color) -# --------------------------------------------------------------------------- -def plot_timings(jsonl_path, axs = None, unit: str = "min", color = "b", label= ''): - """ - Plot cumulative durations from a .jsonl timing ledger written by @mitim_timer, - with vertical lines when the beat number changes. - - Parameters - ---------- - jsonl_path : str | Path - File with one JSON record per line. - unit : {"s", "min", "h"} - Unit for the y-axis. - """ - multiplier = {"s": 1, "min": 1 / 60, "h": 1 / 3600}[unit] - - scripts, script_time, cumulative, beat_nums, script_restarts = [], [], [], [], [] - running = 0.0 - beat_pat = re.compile(r"Beat\s*#\s*(\d+)") - - # ── read the file ─────────────────────────────────────────────────────── - with Path(jsonl_path).expanduser().open() as f: - for line in f: - if not line.strip(): - continue - rec = json.loads(line) - - if rec["script"] not in scripts: - - scripts.append(rec["script"]) - script_time.append(rec["duration_s"] * multiplier) - running += rec["duration_s"]* multiplier - cumulative.append(running) - - m = beat_pat.search(rec["script"]) - beat_nums.append(int(m.group(1)) if m else None) - - script_restarts.append(0.0) - - else: - # If the script is already in the list, it means it was restarted - idx = scripts.index(rec["script"]) - script_restarts[idx] += rec["duration_s"] * multiplier - - if not scripts: - raise ValueError(f"No records found in {jsonl_path}") - - beat_nums = [0] + beat_nums # Start with zero beat - scripts = ['ini'] + scripts # Add initial beat - script_time = [0.0] + script_time # Start with zero time - cumulative = [0.0] + cumulative # Start with zero time - script_restarts = [0.0] + script_restarts # Start with zero restarts - - # ── plot ──────────────────────────────────────────────────────────────── - x = list(range(len(scripts))) - - if axs is None: - plt.ion() - fig = plt.figure() - axs = fig.subplot_mosaic(""" - A - B - """) - - - ax = axs['A'] - ax.plot(x, cumulative, "-s", markersize=8, color=color, label=label) - - # Add restarts as vertical lines - for i in range(len(script_restarts)): - if script_restarts[i] > 0: - ax.plot( - [x[i],x[i]], - [cumulative[i],cumulative[i]+script_restarts[i]], - "-.o", markersize=5, color=color) - - - for i in range(1, len(beat_nums)): - if beat_nums[i] != beat_nums[i - 1]: - ax.axvline(i - 0.5, color='k',linestyle="-.") - - #ax.set_xlim(left=0) - ax.set_ylabel(f"Cumulative time ({unit})"); #ax.set_ylim(bottom=0) - ax.set_xticks(x, scripts, rotation=10, ha="right", fontsize=8) - GRAPHICStools.addDenseAxis(ax) - ax.legend(loc='upper left', fontsize=8) - - ax = axs['B'] - for i in range(len(scripts)-1): - ax.plot([x[i], x[i+1]], [0, script_time[i+1]], "-s", markersize=8, color=color) - - # Add restarts as vertical lines - for i in range(len(script_restarts)-1): - if script_restarts[i] > 0: - ax.plot( - [x[i+1],x[i+1]], - [script_time[i+1],script_time[i+1]+script_restarts[i+1]], - "-.o", markersize=5, color=color) - - for i in range(1, len(beat_nums)): - if beat_nums[i] != beat_nums[i - 1]: - ax.axvline(i - 0.5, color='k',linestyle="-.") - - #ax.set_xlim(left=0) - ax.set_ylabel(f"Time ({unit})"); #ax.set_ylim(bottom=0) - ax.set_xticks(x, scripts, rotation=10, ha="right", fontsize=8) - GRAPHICStools.addDenseAxis(ax) - - return x, scripts diff --git a/src/mitim_tools/misc_tools/GRAPHICStools.py b/src/mitim_tools/misc_tools/GRAPHICStools.py index 5a8c9c79..ebc52a16 100644 --- a/src/mitim_tools/misc_tools/GRAPHICStools.py +++ b/src/mitim_tools/misc_tools/GRAPHICStools.py @@ -1515,3 +1515,4 @@ def animate(i): Writer = animation.writers["ffmpeg"] writer = Writer(fps=framePS, metadata=dict(artist="PRF"), bitrate=BITrate) ani.save(writer=writer, filename=MovieFile, dpi=DPIs) + diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index 4b6a01f6..0e8f0c91 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -1,4 +1,5 @@ import os +import re import shutil import psutil import copy @@ -170,6 +171,121 @@ def wrapper_timer(*args, **kwargs): return decorator_timer +# --------------------------------------------------------------------------- +def plot_timings(jsonl_path, axs = None, unit: str = "min", color = "b", label= ''): + """ + Plot cumulative durations from a .jsonl timing ledger written by @mitim_timer, + with vertical lines when the beat number changes. + + Parameters + ---------- + jsonl_path : str | Path + File with one JSON record per line. + unit : {"s", "min", "h"} + Unit for the y-axis. + """ + multiplier = {"s": 1, "min": 1 / 60, "h": 1 / 3600}[unit] + + scripts, script_time, cumulative, beat_nums, script_restarts = [], [], [], [], [] + running = 0.0 + beat_pat = re.compile(r"Beat\s*#\s*(\d+)") + + # ── read the file ─────────────────────────────────────────────────────── + with Path(jsonl_path).expanduser().open() as f: + for line in f: + if not line.strip(): + continue + rec = json.loads(line) + + if rec["script"] not in scripts: + + scripts.append(rec["script"]) + script_time.append(rec["duration_s"] * multiplier) + running += rec["duration_s"]* multiplier + cumulative.append(running) + + m = beat_pat.search(rec["script"]) + beat_nums.append(int(m.group(1)) if m else None) + + script_restarts.append(0.0) + + else: + # If the script is already in the list, it means it was restarted + idx = scripts.index(rec["script"]) + script_restarts[idx] += rec["duration_s"] * multiplier + + if not scripts: + raise ValueError(f"No records found in {jsonl_path}") + + beat_nums = [0] + beat_nums # Start with zero beat + scripts = ['ini'] + scripts # Add initial beat + script_time = [0.0] + script_time # Start with zero time + cumulative = [0.0] + cumulative # Start with zero time + script_restarts = [0.0] + script_restarts # Start with zero restarts + + # ── plot ──────────────────────────────────────────────────────────────── + x = list(range(len(scripts))) + + if axs is None: + plt.ion() + fig = plt.figure() + axs = fig.subplot_mosaic(""" + A + B + """) + + try: + axs = [ax for ax in axs.values()] + except: + pass + + ax = axs[0] + ax.plot(x, cumulative, "-s", markersize=8, color=color, label=label) + + # Add restarts as vertical lines + for i in range(len(script_restarts)): + if script_restarts[i] > 0: + ax.plot( + [x[i],x[i]], + [cumulative[i],cumulative[i]+script_restarts[i]], + "-.o", markersize=5, color=color) + + + for i in range(1, len(beat_nums)): + if beat_nums[i] != beat_nums[i - 1]: + ax.axvline(i - 0.5, color='k',linestyle="-.") + + #ax.set_xlim(left=0) + ax.set_ylabel(f"Cumulative time ({unit})"); #ax.set_ylim(bottom=0) + ax.set_xticks(x, scripts, rotation=10, ha="right", fontsize=8) + GRAPHICStools.addDenseAxis(ax) + ax.legend(loc='upper left', fontsize=8) + + + ax = axs[1] + for i in range(len(scripts)-1): + ax.plot([x[i], x[i+1]], [0, script_time[i+1]], "-s", markersize=8, color=color) + + # Add restarts as vertical lines + for i in range(len(script_restarts)-1): + if script_restarts[i] > 0: + ax.plot( + [x[i+1],x[i+1]], + [script_time[i+1],script_time[i+1]+script_restarts[i+1]], + "-.o", markersize=5, color=color) + + for i in range(1, len(beat_nums)): + if beat_nums[i] != beat_nums[i - 1]: + ax.axvline(i - 0.5, color='k',linestyle="-.") + + #ax.set_xlim(left=0) + ax.set_ylabel(f"Time ({unit})"); #ax.set_ylim(bottom=0) + ax.set_xticks(x, scripts, rotation=10, ha="right", fontsize=8) + GRAPHICStools.addDenseAxis(ax) + + return x, scripts + + # ------------------------------------ # Decorator to hook methods before and after execution diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index d07518c0..2fbc88da 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -1,3 +1,4 @@ +import sys import copy import datetime import array @@ -10,6 +11,7 @@ import numpy as np import matplotlib.pyplot as plt from mitim_tools.misc_tools import IOtools, GRAPHICStools, GUItools, LOGtools +from mitim_tools.misc_tools.IOtools import mitim_timer from mitim_tools.opt_tools import OPTtools, STEPtools from mitim_tools.opt_tools.utils import ( BOgraphics, @@ -438,10 +440,11 @@ def __init__( Grab variables ------------------------------------------------------------------------------ """ + + self.timings_file = self.folderOutputs / "timing.jsonl" # Logger - self.logFile = BOgraphics.LogFile(self.folderOutputs / "optimization_log.txt") - self.logFile.activate() + sys.stdout = LOGtools.Logger(logFile=self.folderOutputs / "optimization_log.txt", writeAlsoTerminal=True) # Meta self.numIterations = self.optimization_options["convergence_options"]["maximum_iterations"] @@ -733,47 +736,8 @@ def run(self): --------------------------------------------------------------------------------------- """ - train_Ystd = self.train_Ystd if (self.optimization_options["evaluation_options"]["train_Ystd"] is None) else self.optimization_options["evaluation_options"]["train_Ystd"] - - current_step = STEPtools.OPTstep( - self.train_X, - self.train_Y, - train_Ystd, - bounds=self.bounds, - stepSettings=self.stepSettings, - currentIteration=self.currentIteration, - strategy_options=self.strategy_options_use, - BOmetrics=self.BOmetrics, - surrogate_parameters=self.surrogate_parameters, - ) - - # Incorporate strategy_options for later retrieving - current_step.strategy_options_use = copy.deepcopy(self.strategy_options_use) - - self.steps.append(current_step) - - # Avoid points - avoidPoints = np.append(self.avoidPoints_failed, self.avoidPoints_outside) - self.avoidPoints = np.unique([int(j) for j in avoidPoints]) + self._step() - # ***** Fit - self.steps[-1].fit_step(avoidPoints=self.avoidPoints) - - # ***** Define evaluators - self.steps[-1].defineFunctions(self.scalarized_objective) - - # Store class with the model fitted and evaluators defined - if self.storeClass: - self.save() - - # ***** Optimize - if not self.hard_finish: - self.steps[-1].optimize( - position_best_so_far=self.BOmetrics["overall"]["indBest"], - seed=self.seed, - ) - else: - self.steps[-1].x_next = None # Pass the information about next step self.x_next = self.steps[-1].x_next @@ -898,6 +862,74 @@ def read(self, name="optimization_object.pkl", iteration=None, file=None, provid return aux if provideFullClass else step + # Convenient helper methods to track timings of components + + @mitim_timer(lambda self: f'Eval @ {self.currentIteration}', log_file=lambda self: self.timings_file) + def _evaluate(self): + + y_next, ystd_next, self.numEval = EVALUATORtools.fun( + self.optimization_object, + self.x_next, + self.folderExecution, + self.bounds, + self.outputs, + self.optimization_data, + parallel=self.parallel_evaluations, + cold_start=self.cold_start, + numEval=self.numEval, + ) + + return y_next, ystd_next + + @mitim_timer(lambda self: f'Surr @ {self.currentIteration}', log_file=lambda self: self.timings_file) + def _step(self): + + train_Ystd = self.train_Ystd if (self.optimization_options["evaluation_options"]["train_Ystd"] is None) else self.optimization_options["evaluation_options"]["train_Ystd"] + + current_step = STEPtools.OPTstep( + self.train_X, + self.train_Y, + train_Ystd, + bounds=self.bounds, + stepSettings=self.stepSettings, + currentIteration=self.currentIteration, + strategy_options=self.strategy_options_use, + BOmetrics=self.BOmetrics, + surrogate_parameters=self.surrogate_parameters, + ) + + + # Incorporate strategy_options for later retrieving + current_step.strategy_options_use = copy.deepcopy(self.strategy_options_use) + + self.steps.append(current_step) + + # Avoid points + avoidPoints = np.append(self.avoidPoints_failed, self.avoidPoints_outside) + self.avoidPoints = np.unique([int(j) for j in avoidPoints]) + + # ***** Fit + self.steps[-1].fit_step(avoidPoints=self.avoidPoints) + + # ***** Define evaluators + self.steps[-1].defineFunctions(self.scalarized_objective) + + # Store class with the model fitted and evaluators defined + if self.storeClass: + self.save() + + # ***** Optimize + if not self.hard_finish: + self.steps[-1].optimize( + position_best_so_far=self.BOmetrics["overall"]["indBest"], + seed=self.seed, + ) + else: + self.steps[-1].x_next = None + + # --------------------------------------------------------------------------------- + + def updateSet( self, strategy_options_use, isThisCorrected=False, ForceNotApplyCorrections=False ): @@ -941,17 +973,7 @@ def updateSet( # --- Evaluation time1 = datetime.datetime.now() - y_next, ystd_next, self.numEval = EVALUATORtools.fun( - self.optimization_object, - self.x_next, - self.folderExecution, - self.bounds, - self.outputs, - self.optimization_data, - parallel=self.parallel_evaluations, - cold_start=self.cold_start, - numEval=self.numEval, - ) + y_next, ystd_next = self._evaluate() txt_time = IOtools.getTimeDifference(time1) print(f"\t- Complete model update took {txt_time}") # ------------------ @@ -1054,6 +1076,7 @@ def updateSet( return y_next, ystd_next + @mitim_timer(lambda self: f'Init', log_file=lambda self: self.timings_file) def initializeOptimization(self): print("\n") print("------------------------------------------------------------") @@ -1097,10 +1120,7 @@ def initializeOptimization(self): if (not self.cold_start) and (self.optimization_data is not None): self.type_initialization = 3 - print( - "--> Since restart from a previous MITIM has been requested, forcing initialization type to 3 (read from optimization_data)", - typeMsg="i", - ) + print("--> Since restart from a previous MITIM has been requested, forcing initialization type to 3 (read from optimization_data)",typeMsg="i",) if self.type_initialization == 3: print("--> Initialization by reading tabular data...") @@ -1132,9 +1152,7 @@ def initializeOptimization(self): # cold_started run from previous. Grab DVs of initial set if readCasesFromTabular: try: - self.train_X, self.train_Y, self.train_Ystd = self.optimization_data.extract_points( - points=np.arange(self.initial_training) - ) + self.train_X, self.train_Y, self.train_Ystd = self.optimization_data.extract_points(points=np.arange(self.initial_training)) # It could be the case that those points in Tabular are outside the bounds that I want to apply to this optimization, remove outside points? @@ -1148,20 +1166,14 @@ def initializeOptimization(self): self.avoidPoints_outside.append(i) except: - flagger = print( - "Error reading Tabular. Do you want to continue without cold_start and do standard initialization instead?", - typeMsg="q", - ) + flagger = print("Error reading Tabular. Do you want to continue without cold_start and do standard initialization instead?",typeMsg="q",) self.type_initialization = 1 self.cold_start = True readCasesFromTabular = False if readCasesFromTabular and IOtools.isAnyNan(self.train_X): - flagger = print( - " --> cold_start requires non-nan DVs, doing normal initialization", - typeMsg="q", - ) + flagger = print(" --> cold_start requires non-nan DVs, doing normal initialization",typeMsg="q",) if not flagger: embed() @@ -1174,10 +1186,7 @@ def initializeOptimization(self): if not readCasesFromTabular: if self.type_initialization == 1 and self.optimization_options["problem_options"]["dvs_base"] is not None: self.initial_training = self.initial_training - 1 - print( - f"--> Baseline point has been requested with LHS initialization, reducing requested initial random set to {self.initial_training}", - typeMsg="i", - ) + print(f"--> Baseline point has been requested with LHS initialization, reducing requested initial random set to {self.initial_training}",typeMsg="i",) """ Initialization @@ -1528,12 +1537,8 @@ def plot( if plotoptimization_results: # Most current state of the optimization_results.out self.optimization_results.read() - if "logFile" in self.__dict__.keys(): - logFile = self.logFile - else: - logFile = None self.optimization_results.plot( - fn=fn, doNotShow=True, log=logFile, tab_color=tab_color + fn=fn, doNotShow=True, log=self.timings_file, tab_color=tab_color ) """ diff --git a/src/mitim_tools/opt_tools/scripts/read.py b/src/mitim_tools/opt_tools/scripts/read.py index 939ef3ba..170e6418 100644 --- a/src/mitim_tools/opt_tools/scripts/read.py +++ b/src/mitim_tools/opt_tools/scripts/read.py @@ -75,7 +75,7 @@ def plotCompare(folders, plotMeanMax=[True, False]): ) res.read() - log_class = BOgraphics.LogFile(folderWork / "Outputs" / "optimization_log.txt") + log_class = folderWork / "Outputs" / "timing.jsonl" try: log_class.interpret() @@ -98,12 +98,8 @@ def plotCompare(folders, plotMeanMax=[True, False]): #ax1.axhline(y=compared, ls="-.", lw=0.3, color=color) if log_class is not None: - log_class.plot( - axs=[ax2, ax3], - ls=types_ls[i], - lab=name, - marker=types_m[i], - color=colors[i], + IOtools.plot_timings( + folderWork / "Outputs" / "timing.jsonl", axs=[ax2, ax3], label=name, color=color ) yCummMeans.append(yCummMean) diff --git a/src/mitim_tools/opt_tools/utils/BOgraphics.py b/src/mitim_tools/opt_tools/utils/BOgraphics.py index d63d5fce..49ab8552 100644 --- a/src/mitim_tools/opt_tools/utils/BOgraphics.py +++ b/src/mitim_tools/opt_tools/utils/BOgraphics.py @@ -860,12 +860,7 @@ def retrieveResults( res.read() # ---------------- Read Logger - log = LogFile(folderWork / "Outputs" / "optimization_log.txt") - try: - log.interpret() - except: - print("Could not read log", typeMsg="w") - log = None + timings_file = folderWork / "Outputs" / "timings.json" # ---------------- Read Tabular if analysis_level >= 0: @@ -886,7 +881,6 @@ def retrieveResults( # ------------------- mitim_model.optimization_results = res - mitim_model.logFile = log if plotFN is not None: fn = mitim_model.plot( doNotShow=doNotShow, @@ -898,245 +892,11 @@ def retrieveResults( # If no pickle, plot only the contents of optimization_results else: if plotFN: - fn = res.plot(doNotShow=doNotShow, log=log, fn = plotFN) + fn = res.plot(doNotShow=doNotShow, log=timings_file, fn = plotFN) mitim_model = None return fn, res, mitim_model, log, data_df - - -class LogFile: - def __init__(self, file): - self.file = file - - def activate(self, writeAlsoTerminal=True): - sys.stdout = LOGtools.Logger( - logFile=self.file, writeAlsoTerminal=writeAlsoTerminal - ) - - branch, commit_hash = IOtools.get_git_info(__mitimroot__) - print(f"Log file from MITIM version {mitim_version} from {branch} branch and commit hash {commit_hash}") - - def interpret(self): - with open(self.file, "r") as f: - lines = f.readlines() - - self.steps = {} - for line in lines: - if "Starting MITIM Optimization" in line: - try: - self.steps["start"] = IOtools.getTimeFromString( - line.split(",")[0].strip() - ) - except: - self.steps["start"] = IOtools.getTimeFromString( - " ".join(line.split(",")[0].strip().split()[-2:]) - ) - self.steps["steps"] = {} - if "MITIM Step" in line: - aft = line.split("Step")[-1] - ikey = int(aft.split()[0]) - time_str = aft.split("(")[-1].split(")")[0] - self.steps["steps"][ikey] = { - "start": IOtools.getTimeFromString(time_str), - "optimization": {}, - } - if "Posterior Optimization" in line: - time_str = line.split(",")[-1][:-2].strip() - self.steps["steps"][ikey]["optimization"] = { - "start": IOtools.getTimeFromString(time_str), - "steps": {}, - } - cont = 0 - if "Optimization stage " in line: - aft = line.split("Step")[-1] - time_str = aft.split("(")[-1].split(")")[0] - self.steps["steps"][ikey]["optimization"]["steps"][cont] = { - "name": line.split()[4], - "start": IOtools.getTimeFromString(time_str), - } - cont += 1 - - self.process() - - def process(self): - for step in self.steps["steps"]: - time_start = self.steps["steps"][step]["start"] - - if "start" not in self.steps["steps"][step]["optimization"]: - break - time_end = self.steps["steps"][step]["optimization"]["start"] - timeF = IOtools.getTimeDifference( - time_start, newTime=time_end, niceText=False - ) - self.steps["steps"][step]["fitting"] = timeF - - if step + 1 in self.steps["steps"]: - time_end = self.steps["steps"][step + 1]["start"] - time = IOtools.getTimeDifference( - time_start, newTime=time_end, niceText=False - ) - self.steps["steps"][step]["time_s"] = time - - for opt_step in self.steps["steps"][step]["optimization"]["steps"]: - time_start = self.steps["steps"][step]["optimization"]["steps"][ - opt_step - ]["start"] - - if opt_step + 1 in self.steps["steps"][step]["optimization"]["steps"]: - time_end = self.steps["steps"][step]["optimization"]["steps"][ - opt_step + 1 - ]["start"] - time = IOtools.getTimeDifference( - time_start, newTime=time_end, niceText=False - ) - self.steps["steps"][step]["optimization"]["steps"][opt_step][ - "time_s" - ] = time - else: - if step + 1 in self.steps["steps"]: - time_end = time_end = self.steps["steps"][step + 1]["start"] - time = IOtools.getTimeDifference( - time_start, newTime=time_end, niceText=False - ) - self.steps["steps"][step]["optimization"]["steps"][opt_step][ - "time_s" - ] = time - - self.points = [ - 0, - IOtools.getTimeDifference( - self.steps["start"], - newTime=self.steps["steps"][0]["start"], - niceText=False, - ), - ] - self.types = ["b"] - - for step in self.steps["steps"]: - if "fitting" in self.steps["steps"][step]: - self.points.append( - self.steps["steps"][step]["fitting"] + self.points[-1] - ) - self.types.append("r") - - if "steps" not in self.steps["steps"][step]["optimization"]: - break - for opt_step in self.steps["steps"][step]["optimization"]["steps"]: - if ( - "time_s" - in self.steps["steps"][step]["optimization"]["steps"][opt_step] - ): - self.points.append( - self.steps["steps"][step]["optimization"]["steps"][opt_step][ - "time_s" - ] - + self.points[-1] - ) - self.types.append("g") - - self.points = np.array(self.points) - - self.its = np.linspace(0, len(self.points) - 1, len(self.points)) - - def plot( - self, - axs=None, - factor=60.0, - fullCumulative=False, - ls="-", - lab="", - marker="o", - color="b", - ): - if axs is None: - plt.ion() - fig, axs = plt.subplots(ncols=2) - - ax = axs[0] - subtractor = 0 - totals = {"ini": 0.0, "fit": 0.0, "opt": 0.0} - - for i in range(len(self.points) - 1): - if self.types[i] == "r": - ax.axvline(x=self.its[i], ls="--", c="k", lw=0.5) - if not fullCumulative: - subtractor = self.points[i] - - ps = [ - (self.points[i] - subtractor) / factor, - (self.points[i + 1] - subtractor) / factor, - ] - - if self.types[i] == "b": - totals["ini"] += ps[1] - ps[0] - if self.types[i] == "g": - totals["opt"] += ps[1] - ps[0] - if self.types[i] == "r": - totals["fit"] += ps[1] - ps[0] - - if i == 0: - labb = lab - else: - labb = "" - - ax.plot( - [self.its[i], self.its[i + 1]], - ps, - marker + ls, - c=self.types[i], - label=labb, - ) - - if factor == 60.0: - label = "minutes" - ax.axhline(y=60, ls="-.", lw=0.2) - elif factor == 3600.0: - label = "hours" - ax.axhline(y=1, ls="-.", lw=0.2) - else: - label = "seconds" - - # ax.set_xlabel('Workflow Steps') - ax.set_ylabel(f"Cumulated Time ({label})") - # ax.set_xlim(left=0) - ax.set_ylim(bottom=0) - - from matplotlib.lines import Line2D - - custom_lines = [ - Line2D([0], [0], color="b", lw=2), - Line2D([0], [0], color="r", lw=2), - Line2D([0], [0], color="g", lw=2), - ] - - legs = [ - "Initialization + Evaluation", - "Evaluation + Fitting", - "Optimization", - "Total", - ] - ax.legend(custom_lines, legs) - - ax = axs[1] - ax.bar( - legs, - [ - totals["ini"], - totals["fit"], - totals["opt"], - totals["ini"] + totals["fit"] + totals["opt"], - ], - 1 / 3, - alpha=0.5, - label=lab, - color=color, - ) # , label=equil_names[i],color=colors[i],align='edge') - - # ax.set_xlabel('Workflow') - ax.set_ylabel(f"Cumulated Time ({label})") - - class optimization_data: def __init__( self, @@ -1661,7 +1421,7 @@ def plot( fig4 = self.fn.add_figure(label="Improvement", tab_color=tab_color) if log is not None: figTimes = self.fn.add_figure(label="Times", tab_color=tab_color) - grid = plt.GridSpec(1, 2, hspace=0.3, wspace=0.3) + grid = plt.GridSpec(2, 1, hspace=0.3, wspace=0.3) axsTimes = [figTimes.add_subplot(grid[0]), figTimes.add_subplot(grid[1])] _ = self.plotComplete( @@ -1695,7 +1455,7 @@ def plot( _, _ = self.plotImprovement(axs=[ax0, ax1, ax2, ax3]) if log is not None: - log.plot(axs=[axsTimes[0], axsTimes[1]]) + IOtools.plot_timings(log, axs = [axsTimes[0], axsTimes[1]]) return self.fn diff --git a/src/mitim_tools/opt_tools/utils/EVALUATORtools.py b/src/mitim_tools/opt_tools/utils/EVALUATORtools.py index 15781cd0..9a19b6f1 100644 --- a/src/mitim_tools/opt_tools/utils/EVALUATORtools.py +++ b/src/mitim_tools/opt_tools/utils/EVALUATORtools.py @@ -37,7 +37,6 @@ def parallel_main(Params, cont): lock=lock, ) - def fun( optimization_object, x, From 5c43855b598d55fb2826b4dd0cfd77f4b1fdbc98 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 6 Aug 2025 18:40:14 -0400 Subject: [PATCH 149/385] misc --- src/mitim_tools/opt_tools/STRATEGYtools.py | 2 -- src/mitim_tools/opt_tools/utils/BOgraphics.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 2fbc88da..5adad9ed 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -195,7 +195,6 @@ def read_optimization_results( self.fn, self.res, self.mitim_model, - self.log, self.data, ) = BOgraphics.retrieveResults( self.folder, @@ -897,7 +896,6 @@ def _step(self): BOmetrics=self.BOmetrics, surrogate_parameters=self.surrogate_parameters, ) - # Incorporate strategy_options for later retrieving current_step.strategy_options_use = copy.deepcopy(self.strategy_options_use) diff --git a/src/mitim_tools/opt_tools/utils/BOgraphics.py b/src/mitim_tools/opt_tools/utils/BOgraphics.py index 49ab8552..ebad817e 100644 --- a/src/mitim_tools/opt_tools/utils/BOgraphics.py +++ b/src/mitim_tools/opt_tools/utils/BOgraphics.py @@ -895,7 +895,7 @@ def retrieveResults( fn = res.plot(doNotShow=doNotShow, log=timings_file, fn = plotFN) mitim_model = None - return fn, res, mitim_model, log, data_df + return fn, res, mitim_model, data_df class optimization_data: def __init__( From 5f082809ce50a2837ff1a9d7fc7ef53ed45ad964 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 7 Aug 2025 17:28:12 -0400 Subject: [PATCH 150/385] Bug fix MAESTRO in new development changes --- src/mitim_tools/transp_tools/utils/TRANSPhelpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py index 1ff37140..428e6483 100644 --- a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py +++ b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py @@ -820,7 +820,7 @@ def _produce_geometry_profiles(self): # Separatrix # -------------------------------------------------------------- - self.geometry['R_sep'], self.geometry['Z_sep'] = self.p.derived["R_surface"][0,:,-1], self.p.derived["Z_surface"][0,:,-1] + self.geometry['R_sep'], self.geometry['Z_sep'] = self.p.derived["R_surface"][0,-1,:], self.p.derived["Z_surface"][0,-1,:] # -------------------------------------------------------------- # VV From 73baf47340f03cc3a0e394d8253228080f5d26d2 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 7 Aug 2025 17:28:52 -0400 Subject: [PATCH 151/385] For now, until further fixes, do not require quends or vmecpp --- pyproject.toml | 4 ++-- src/mitim_tools/gacode_tools/utils/CGYROutils.py | 3 ++- src/mitim_tools/plasmastate_tools/utils/VMECtools.py | 3 ++- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 08dc56d4..22a4ec74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,8 +45,8 @@ dependencies = [ "onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models "tensorflow", "f90nml", - "vmecpp", - "quends @ git+https://github.com/sandialabs/quends.git" + # "vmecpp", + # "quends @ git+https://github.com/sandialabs/quends.git" ] [project.optional-dependencies] diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 54389192..d10657df 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -8,7 +8,6 @@ from pygacode.cgyro.data_plot import cgyrodata_plot from pygacode import gacodefuncs from IPython import embed -import quends as qnds import pandas as pd class CGYROlinear_scan: @@ -582,6 +581,8 @@ def absexp(x,tau): def quends_analysis(t, S, debug = False): + import quends as qnds + time_dependent_data = {'time': t, 'signal': S} df = pd.DataFrame(time_dependent_data, index = pd.RangeIndex(len(t))) diff --git a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py index 6a7ac4f1..36f8fa87 100644 --- a/src/mitim_tools/plasmastate_tools/utils/VMECtools.py +++ b/src/mitim_tools/plasmastate_tools/utils/VMECtools.py @@ -1,4 +1,3 @@ -import vmecpp import numpy as np from collections import OrderedDict import matplotlib.pyplot as plt @@ -45,6 +44,8 @@ def __init__( @IOtools.hook_method(after=MITIMstate.ensure_variables_existence) def _read_vmec(self): + + import vmecpp # Read VMEC file print("\t- Reading VMEC file") From 66045a8fccd5ff4e0d8b8fdd66cf7e24844112af Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 8 Aug 2025 09:03:00 -0400 Subject: [PATCH 152/385] Bug fix fast --- src/mitim_tools/gacode_tools/TGLFtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index d0bf5057..c1582f85 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -4491,7 +4491,7 @@ def read(self,require_all_files=True): IncludeExtraIonsInQi = [i - 1 for i in self.inputclass.ions_info["thermal_list_extras"]] if self.inputclass is not None else [] self.ions_included = (1,) + tuple(IncludeExtraIonsInQi) - self.fast_included = tuple(self.inputclass.ions_info["fast_list"]) if self.inputclass is not None else () + self.fast_included = tuple([i-1 for i in self.inputclass.ions_info["fast_list"]]) if self.inputclass is not None else () # ------------------------------------------------------------------------ # Fluxes From 1919d7888d6aeb2c518d68427679bd03b8fac505 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 8 Aug 2025 09:56:58 -0400 Subject: [PATCH 153/385] Bug fix MAESTRO gacode_state --- src/mitim_modules/maestro/utils/EPEDbeat.py | 6 +++--- src/mitim_tools/transp_tools/CDFtools.py | 2 +- tests/MAESTRO_workflow.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index d6f1df59..f3dfeb3a 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -363,7 +363,7 @@ def _run_full_eped(self, folder, Ip, Bt, R, a, kappa995, delta995, neped19, Beta def finalize(self, **kwargs): - self.profiles_output = PROFILEStools.PROFILES_GACODE(self.folder / 'input.gacode.eped') + self.profiles_output = PROFILEStools.gacode_state(self.folder / 'input.gacode.eped') self.profiles_output.write_state(file=self.folder_output / 'input.gacode') @@ -379,7 +379,7 @@ def grab_output(self): loaded_results = np.load(self.folder_output / 'eped_results.npy', allow_pickle=True).item() - profiles = PROFILEStools.PROFILES_GACODE(self.folder_output / 'input.gacode') if isitfinished else None + profiles = PROFILEStools.gacode_state(self.folder_output / 'input.gacode') if isitfinished else None else: @@ -404,7 +404,7 @@ def plot(self, fn = None, counter = 0, full_plot = True): loaded_results, profiles = self.grab_output() - profiles_current = PROFILEStools.PROFILES_GACODE(self.folder / 'input.gacode') + profiles_current = PROFILEStools.gacode_state(self.folder / 'input.gacode') profiles_current.plotRelevant(axs = axs, color = 'b', label = 'orig') diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index 1529f7af..e60e3ae8 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -15443,7 +15443,7 @@ def grid_interpolation_method_to_zero(x,y): for key in ['ne(10^19/m^3)', 'ni(10^19/m^3)', 'te(keV)', 'ti(keV)', 'rmin(m)']: profiles[key] = profiles[key].clip(min=minimum) - p = PROFILEStools.PROFILES_GACODE.scratch(profiles) + p = PROFILEStools.gacode_state.scratch(profiles) return p diff --git a/tests/MAESTRO_workflow.py b/tests/MAESTRO_workflow.py index b206c428..ccb6eb01 100644 --- a/tests/MAESTRO_workflow.py +++ b/tests/MAESTRO_workflow.py @@ -3,7 +3,7 @@ from mitim_tools import __mitimroot__ from mitim_modules.maestro.scripts import run_maestro -cold_start = True +cold_start = False folder = __mitimroot__ / "tests" / "scratch" / "maestro_test" template = __mitimroot__ / "templates" / "maestro_namelist.json" From 1c2015f5c9d6ccf1d8b2b3628632f28357107ff4 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 8 Aug 2025 11:07:17 -0400 Subject: [PATCH 154/385] Bug fixes MAESTRO in new development --- src/mitim_tools/misc_tools/IOtools.py | 14 ++++++++++++-- src/mitim_tools/transp_tools/CDFtools.py | 2 +- .../transp_tools/src/TRANSPsingularity.py | 2 +- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index 0e8f0c91..ee68aa6e 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -599,7 +599,7 @@ def opt_crit(*args,**kwargs): print('--------------------------------------------------') print('Convergence criteria') print('--------------------------------------------------') - v = unprint_fun(*args,**kwargs) + v = unprint_fun(*args, **kwargs) print('--------------------------------------------------\n') return v optimization_options['convergence_options']['stopping_criteria'] = opt_crit @@ -926,7 +926,7 @@ def getLocInfo(locFile, with_extension=False): def findFileByExtension( - folder, extension, prefix=" ", fixSpaces=False, ForceFirst=False, agnostic_to_case=False + folder, extension, prefix=" ", fixSpaces=False, ForceFirst=False, agnostic_to_case=False, do_not_consider_files=None ): """ Retrieves the file without folder and extension @@ -938,6 +938,16 @@ def findFileByExtension( if fpath.exists(): allfiles = findExistingFiles(fpath, extension, agnostic_to_case = agnostic_to_case) + # Filter out files that contain any of the strings in do_not_consider_files + if do_not_consider_files is not None: + filtered_files = [] + for file_path in allfiles: + file_name = file_path.name + should_exclude = any(exclude_str in file_name for exclude_str in do_not_consider_files) + if not should_exclude: + filtered_files.append(file_path) + allfiles = filtered_files + if len(allfiles) > 1: # print(allfiles) if not ForceFirst: diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index e60e3ae8..af3f97a0 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -130,7 +130,7 @@ def __init__( # Capability to provide folder and just find the CDF in there if self.LocationCDF.is_dir(): - self.LocationCDF = IOtools.findFileByExtension(self.LocationCDF, ".CDF", agnostic_to_case=True) + self.LocationCDF = IOtools.findFileByExtension(self.LocationCDF, ".CDF", agnostic_to_case=True,do_not_consider_files=['PH.CDF']) if self.LocationCDF is None: raise ValueError(f"[MITIM] Could not find a CDF file in {self.LocationCDF}") # ---------------------------- diff --git a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py index 7a182af3..6ebbabc5 100644 --- a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py +++ b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py @@ -568,7 +568,7 @@ def runSINGULARITY_finish(folderWork, runid, tok, job_name): if item.is_file(): shutil.copy2(item, folderWork) elif item.is_dir(): - shutil.copytree(item, folderWork / item.name) + shutil.copytree(item, folderWork / item.name, dirs_exist_ok=True) def runSINGULARITY_look(folderWork, folderTRANSP, runid, job_name, times_retry_look = 3): From 8628af8fe8525d177a95e3568fcf7334e120b0cd Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 8 Aug 2025 11:14:47 -0400 Subject: [PATCH 155/385] misc --- tests/MAESTRO_workflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/MAESTRO_workflow.py b/tests/MAESTRO_workflow.py index ccb6eb01..b206c428 100644 --- a/tests/MAESTRO_workflow.py +++ b/tests/MAESTRO_workflow.py @@ -3,7 +3,7 @@ from mitim_tools import __mitimroot__ from mitim_modules.maestro.scripts import run_maestro -cold_start = False +cold_start = True folder = __mitimroot__ / "tests" / "scratch" / "maestro_test" template = __mitimroot__ / "templates" / "maestro_namelist.json" From f8f4d8fd1e2b062da3805ff860654943d242a243 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 8 Aug 2025 11:33:12 -0400 Subject: [PATCH 156/385] Bug fixes plot maestro --- src/mitim_modules/maestro/scripts/plot_maestro.py | 5 ++--- src/mitim_modules/maestro/utils/MAESTROplot.py | 2 +- src/mitim_tools/misc_tools/IOtools.py | 4 +++- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/plot_maestro.py b/src/mitim_modules/maestro/scripts/plot_maestro.py index a41e30bc..ca89e6f0 100644 --- a/src/mitim_modules/maestro/scripts/plot_maestro.py +++ b/src/mitim_modules/maestro/scripts/plot_maestro.py @@ -69,7 +69,7 @@ def main(): only_folder_structure_with_files = None if args.remote_minimal: - only_folder_structure_with_files = ["Outputs/optimization_data.csv","Outputs/optimization_extra.pkl","Outputs/optimization_object.pkl","Outputs/optimization_results.out"] + only_folder_structure_with_files = ["beat_results/input.gacode", "input.gacode_final","initializer_geqdsk/input.gacode", "timing.jsonl"] folders = remote_tools.retrieve_remote_folders(args.folders, args.remote, args.remote_folder_parent, args.remote_folders, only_folder_structure_with_files) @@ -85,11 +85,10 @@ def main(): # Actual interpreting and plotting # -------------------------------------------------------------------------------------------------------------------------------------------- - beats = args.beats + beats = args.beats if not args.remote_minimal else 0 only = args.only full = args.full - folders = [IOtools.expandPath(folder) for folder in folders] fn = GUItools.FigureNotebook("MAESTRO") diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index 6777741c..be5ab6d0 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -212,7 +212,7 @@ def plot_results(self, fn): A B """) - IOtools.plot_timings(self.folder_performance / 'timing.jsonl', axs = axs) + IOtools.plot_timings(self.folder_performance / 'timing.jsonl', axs = axs, log=True) return ps, ps_lab diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index ee68aa6e..11411a03 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -172,7 +172,7 @@ def wrapper_timer(*args, **kwargs): return decorator_timer # --------------------------------------------------------------------------- -def plot_timings(jsonl_path, axs = None, unit: str = "min", color = "b", label= ''): +def plot_timings(jsonl_path, axs = None, unit: str = "min", color = "b", label= '', log=False): """ Plot cumulative durations from a .jsonl timing ledger written by @mitim_timer, with vertical lines when the beat number changes. @@ -282,6 +282,8 @@ def plot_timings(jsonl_path, axs = None, unit: str = "min", color = "b", label= ax.set_ylabel(f"Time ({unit})"); #ax.set_ylim(bottom=0) ax.set_xticks(x, scripts, rotation=10, ha="right", fontsize=8) GRAPHICStools.addDenseAxis(ax) + if log: + ax.set_yscale('log') return x, scripts From b51ce946dea6c30496d609abaf9237fa63905569 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 8 Aug 2025 12:13:23 -0400 Subject: [PATCH 157/385] better maestro plotting timings --- src/mitim_modules/maestro/scripts/plot_maestro.py | 2 +- src/mitim_modules/maestro/utils/MAESTROplot.py | 2 +- src/mitim_tools/misc_tools/IOtools.py | 6 +++++- src/mitim_tools/opt_tools/STRATEGYtools.py | 1 + src/mitim_tools/opt_tools/scripts/read.py | 2 +- src/mitim_tools/opt_tools/utils/BOgraphics.py | 4 +++- 6 files changed, 12 insertions(+), 5 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/plot_maestro.py b/src/mitim_modules/maestro/scripts/plot_maestro.py index ca89e6f0..5c53c1fd 100644 --- a/src/mitim_modules/maestro/scripts/plot_maestro.py +++ b/src/mitim_modules/maestro/scripts/plot_maestro.py @@ -111,7 +111,7 @@ def main(): axsTiming = fig.subplot_mosaic(""" A B - """) + """,sharex=True) colors = GRAPHICStools.listColors() diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index be5ab6d0..11595dd5 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -211,7 +211,7 @@ def plot_results(self, fn): axs = fig.subplot_mosaic(""" A B - """) + """,sharex=True) IOtools.plot_timings(self.folder_performance / 'timing.jsonl', axs = axs, log=True) return ps, ps_lab diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index 11411a03..46260070 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -213,6 +213,10 @@ def plot_timings(jsonl_path, axs = None, unit: str = "min", color = "b", label= # If the script is already in the list, it means it was restarted idx = scripts.index(rec["script"]) script_restarts[idx] += rec["duration_s"] * multiplier + + cumulative[-1] += script_restarts[idx] + running += script_restarts[idx] + if not scripts: raise ValueError(f"No records found in {jsonl_path}") @@ -247,7 +251,7 @@ def plot_timings(jsonl_path, axs = None, unit: str = "min", color = "b", label= if script_restarts[i] > 0: ax.plot( [x[i],x[i]], - [cumulative[i],cumulative[i]+script_restarts[i]], + [cumulative[i],cumulative[i]-script_restarts[i]], "-.o", markersize=5, color=color) diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 5adad9ed..5ab13698 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -1954,6 +1954,7 @@ def clean_state(folder): aux.optimization_options['convergence_options']['stopping_criteria'] = PORTALStools.stopping_criteria_portals aux.folderOutputs = folder / "Outputs" + aux.timings_file = aux.folderOutputs / "timing.jsonl" aux.save() diff --git a/src/mitim_tools/opt_tools/scripts/read.py b/src/mitim_tools/opt_tools/scripts/read.py index 170e6418..b44a0952 100644 --- a/src/mitim_tools/opt_tools/scripts/read.py +++ b/src/mitim_tools/opt_tools/scripts/read.py @@ -55,7 +55,7 @@ def plotCompare(folders, plotMeanMax=[True, False]): ax0 = fig.add_subplot(grid[0, 0]) ax1 = fig.add_subplot(grid[1, 0], sharex=ax0) ax2 = fig.add_subplot(grid[0, 1]) - ax3 = fig.add_subplot(grid[1, 1]) + ax3 = fig.add_subplot(grid[1, 1],sharex=ax2) ax1i = fig.add_subplot(grid[2, 0], sharex=ax0) types_ls = GRAPHICStools.listLS() diff --git a/src/mitim_tools/opt_tools/utils/BOgraphics.py b/src/mitim_tools/opt_tools/utils/BOgraphics.py index ebad817e..bd32b55c 100644 --- a/src/mitim_tools/opt_tools/utils/BOgraphics.py +++ b/src/mitim_tools/opt_tools/utils/BOgraphics.py @@ -1422,7 +1422,9 @@ def plot( if log is not None: figTimes = self.fn.add_figure(label="Times", tab_color=tab_color) grid = plt.GridSpec(2, 1, hspace=0.3, wspace=0.3) - axsTimes = [figTimes.add_subplot(grid[0]), figTimes.add_subplot(grid[1])] + axx0 = figTimes.add_subplot(grid[0]) + axx1 = figTimes.add_subplot(grid[1], sharex=axx0) + axsTimes = [axx0, axx1] _ = self.plotComplete( fig=fig1, From cbd2cae9c04d15da2bf7b85b639331496d5f8fdf Mon Sep 17 00:00:00 2001 From: audreysa Date: Wed, 13 Aug 2025 15:40:23 -0400 Subject: [PATCH 158/385] Quick patch so that query_yes_no automatically responds yes if running on a non-interactive terminal. --- src/mitim_tools/misc_tools/LOGtools.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/mitim_tools/misc_tools/LOGtools.py b/src/mitim_tools/misc_tools/LOGtools.py index f89b8387..1f3f7f80 100644 --- a/src/mitim_tools/misc_tools/LOGtools.py +++ b/src/mitim_tools/misc_tools/LOGtools.py @@ -119,6 +119,11 @@ def query_yes_no(question, extra=""): ''' From https://stackoverflow.com/questions/3041986/apt-command-line-interface-like-yes-no-input ''' + # Fix suggested by chatGPT 4.1 to address failure bc of non-interactive terminal + if not sys.stdin.isatty(): + printMsg(f"\t\t>> Non-interactive terminal detected, auto-confirming '{question}'") + return True + valid = {"y": True, "n": False, "e": None} prompt = " [y/n/e] (yes, no, exit)" From 48303e86ef1dc204c153f315b4f88993cf952490 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 18 Aug 2025 12:53:58 -0400 Subject: [PATCH 159/385] bug fix remote retrieval same name --- src/mitim_tools/misc_tools/utils/remote_tools.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/mitim_tools/misc_tools/utils/remote_tools.py b/src/mitim_tools/misc_tools/utils/remote_tools.py index 32826203..1e183044 100644 --- a/src/mitim_tools/misc_tools/utils/remote_tools.py +++ b/src/mitim_tools/misc_tools/utils/remote_tools.py @@ -29,6 +29,9 @@ def retrieve_remote_folders(folders_local, remote, remote_folder_parent, remote_ folder = IOtools.expandPath(folders[i]) folder_orig = IOtools.expandPath(folders_local[i]) + if folder == folder_orig: + continue + if folder_orig.exists(): IOtools.shutil_rmtree(folder_orig) From f1d28530db2cd7d1989bcb4c579aeaa1372b72aa Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 19 Aug 2025 12:58:28 -0400 Subject: [PATCH 160/385] misc --- .../opt_tools/scripts/evaluate_speed.py | 8 +++---- src/mitim_tools/opt_tools/scripts/read.py | 21 +++++-------------- 2 files changed, 8 insertions(+), 21 deletions(-) diff --git a/src/mitim_tools/opt_tools/scripts/evaluate_speed.py b/src/mitim_tools/opt_tools/scripts/evaluate_speed.py index bc917504..a84625d3 100644 --- a/src/mitim_tools/opt_tools/scripts/evaluate_speed.py +++ b/src/mitim_tools/opt_tools/scripts/evaluate_speed.py @@ -25,9 +25,7 @@ x = torch.rand(cases, step.train_X.shape[-1]) with IOtools.speeder(f"profiler{name}.prof") as s: - with torch.no_grad(): - mean, upper, lower, _ = step.GP["combined_model"].predict(x) + #with torch.no_grad(): + mean, upper, lower, _ = step.GP["combined_model"].predict(x) -print( - f"\nIt took {s.timeDiff:.3f}s to run {x.shape[0]:.1e} parallel evaluations (i.e. {s.timeDiff*1E6/cases:.3f}micro-s/member) of {mean.shape[-1]} GPs with {x.shape[-1]} raw input dimensions" -) +print(f"\nIt took {s.timeDiff:.3f}s to run {x.shape[0]:.1e} parallel evaluations (i.e. {s.timeDiff*1E6/cases:.3f}micro-s/member) of {mean.shape[-1]} GPs with {x.shape[-1]} raw input dimensions") diff --git a/src/mitim_tools/opt_tools/scripts/read.py b/src/mitim_tools/opt_tools/scripts/read.py index b44a0952..a5b86494 100644 --- a/src/mitim_tools/opt_tools/scripts/read.py +++ b/src/mitim_tools/opt_tools/scripts/read.py @@ -65,7 +65,6 @@ def plotCompare(folders, plotMeanMax=[True, False]): yCummMeans = [] xes = [] resS = [] - logS = [] for i, (color, name, folderWork) in enumerate(zip(colors, names, folderWorks)): res = BOgraphics.optimization_results( folderWork / "Outputs" / "optimization_results.out" @@ -75,14 +74,6 @@ def plotCompare(folders, plotMeanMax=[True, False]): ) res.read() - log_class = folderWork / "Outputs" / "timing.jsonl" - - try: - log_class.interpret() - except: - print("Could not read log", typeMsg="w") - log_class = None - plotAllmembers = len(folderWorks) <= 3 xe, yCummMean = res.plotImprovement( axs=[ax0, ax1, ax1i, None], @@ -97,22 +88,20 @@ def plotCompare(folders, plotMeanMax=[True, False]): #compared = -yCummMean[0] * conv if conv < 0 else conv #ax1.axhline(y=compared, ls="-.", lw=0.3, color=color) - if log_class is not None: - IOtools.plot_timings( - folderWork / "Outputs" / "timing.jsonl", axs=[ax2, ax3], label=name, color=color - ) + IOtools.plot_timings( + folderWork / "Outputs" / "timing.jsonl", axs=[ax2, ax3], label=name, color=color + ) yCummMeans.append(yCummMean) xes.append(xe) resS.append(res) - logS.append(log_class) ax0.set_xlim([0, maxEv]) ax2.legend(prop={"size": 6}) ax3.legend(prop={"size": 6}) - return yCummMeans, xes, resS, logS, fig + return yCummMeans, xes, resS, fig def main(): @@ -202,7 +191,7 @@ def main(): opt_funs.append(opt_fun) if analysis_level == -1: - yCummMeans, xes, resS, logS, fig = plotCompare( + yCummMeans, xes, resS, fig = plotCompare( folders_complete, plotMeanMax=[True, len(folders_complete) < 2] ) From ba26164c08ae09d9cb11b6534954ecc9b8640fe9 Mon Sep 17 00:00:00 2001 From: audreysa Date: Wed, 20 Aug 2025 16:02:00 -0400 Subject: [PATCH 161/385] Adjust patch for interactive response required on slurm system to throw an error instead of giving a warning --- src/mitim_tools/misc_tools/LOGtools.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/misc_tools/LOGtools.py b/src/mitim_tools/misc_tools/LOGtools.py index 1f3f7f80..6cb7c82f 100644 --- a/src/mitim_tools/misc_tools/LOGtools.py +++ b/src/mitim_tools/misc_tools/LOGtools.py @@ -119,10 +119,9 @@ def query_yes_no(question, extra=""): ''' From https://stackoverflow.com/questions/3041986/apt-command-line-interface-like-yes-no-input ''' - # Fix suggested by chatGPT 4.1 to address failure bc of non-interactive terminal + if not sys.stdin.isatty(): - printMsg(f"\t\t>> Non-interactive terminal detected, auto-confirming '{question}'") - return True + raise Exception("Interactive terminal response required - something is wrong with this run") valid = {"y": True, "n": False, "e": None} From db12da5f2e63e8a0ca4a41c4acf07e5690d85d6f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 21 Aug 2025 16:10:03 -0400 Subject: [PATCH 162/385] Fixed bug EPED local --- src/mitim_tools/eped_tools/EPEDtools.py | 15 +++++++++++---- tests/EPED_workflow.py | 4 +++- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index c5607723..fb087025 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -132,15 +132,22 @@ def run( # ------------------------------------- # Execute # ------------------------------------- - - # Command to execute by each job in the array - EPEDcommand = f'cd {self.eped_job.folderExecution}/run"$SLURM_ARRAY_TASK_ID" && export NPROC_EPED={nproc_per_run} && ips.py --config=eped.config --platform=psfc_cluster.conf' + + # Submit as a slurm job array + if self.eped_job.launchSlurm: + EPEDcommand = f'cd {self.eped_job.folderExecution}/run"$SLURM_ARRAY_TASK_ID" && export NPROC_EPED={nproc_per_run} && ips.py --config=eped.config --platform=psfc_cluster.conf' + # Submit locally in parallel + else: + EPEDcommand = "" + for i in job_array.split(','): + EPEDcommand += f'cd {self.eped_job.folderExecution}/run{i} && export NPROC_EPED={nproc_per_run} && ips.py --config=eped.config --platform=psfc_cluster.conf & \n' + EPEDcommand += 'wait\n' # Prepare the job script self.eped_job.prep(EPEDcommand,input_folders=folder_cases,output_files=copy.deepcopy(output_files),shellPreCommands=shellPreCommands) # Run the job - self.eped_job.run() #removeScratchFolders=False) + self.eped_job.run() # ------------------------------------- # Postprocessing diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py index 93d39a4b..c8740354 100644 --- a/tests/EPED_workflow.py +++ b/tests/EPED_workflow.py @@ -1,10 +1,11 @@ import os +import matplotlib.pyplot as plt from mitim_tools.eped_tools import EPEDtools from mitim_tools import __mitimroot__ cold_start = True -folder = __mitimroot__ / "tests" / "scratch" / "eped_test16" +folder = __mitimroot__ / "tests" / "scratch" / "eped_test" if cold_start and os.path.exists(folder): os.system(f"rm -r {folder}") @@ -35,3 +36,4 @@ eped.read(subfolder='case1') eped.plot(labels=['case1']) +plt.show() \ No newline at end of file From 479e2773553b00942d99db60934088d26b9dc59b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 19 Aug 2025 17:22:16 -0400 Subject: [PATCH 163/385] Removed deprecated options for direct tglf prep --- .../physics_models/transport_tglf.py | 4 +- src/mitim_tools/gacode_tools/TGLFtools.py | 142 +++++++----------- .../plasmastate_tools/MITIMstate.py | 19 ++- 3 files changed, 66 insertions(+), 99 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index bb547332..8ec3e5de 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -56,12 +56,12 @@ def _evaluate_tglf(self): tglf = TGLFtools.TGLF(rhos=rho_locations) - _ = tglf.prep_direct_tglf( + _ = tglf.prep_direct( self.folder, + self.powerstate.profiles_transport cold_start = cold_start, remove_fast = False, # Use what's in the input.gacode, the removal should happen at initialization recalculate_ptot = False, # Use what's in the input.gacode, which is what PORTALS TGYRO does - inputgacode=self.powerstate.profiles_transport, ) # ------------------------------------------------------------------------------------------------------------------------ diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index c1582f85..40af13b7 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -356,118 +356,54 @@ def prep( return cdf - def prep_direct_tglf( + def prep_direct( self, + mitim_state, # A MITIM state class FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) cold_start=False, # If True, do not use what it potentially inside the folder, run again remove_fast=False, # Ignore fast particles in TGYRO recalculate_ptot=True, # Recalculate PTOT in TGYRO - cdf_open=None, # Grab normalizations from CDF file that is open as transp_output class - inputgacode=None, # *NOTE BELOW* - specificInputs=None, # *NOTE BELOW* - tgyro_results=None, # *NOTE BELOW* forceIfcold_start=False, # Extra flag ): - """ - * Note on inputgacode, specificInputs and tgyro_results: - If I don't want to prepare, I can provide inputgacode and specificInputs, but I have to make sure they are consistent with one another! - Optionally, I can give tgyro_results for further info in such a case - """ - - print("> Preparation of TGLF run") - # PROFILES class. + print("> Preparation of TGLF run from input.gacode (direct conversion)") - if inputgacode is not None: - - if isinstance(inputgacode, str) or isinstance(inputgacode, Path): - - from mitim_tools.gacode_tools import PROFILEStools - self.profiles = PROFILEStools.gacode_state(inputgacode) - - else: - - # If inputgacode is already a PROFILEStools object, just use it - self.profiles = inputgacode - - else: + self.FolderGACODE = IOtools.expandPath(FolderGACODE) + + if cold_start or not self.FolderGACODE.exists(): + IOtools.askNewFolder(self.FolderGACODE, force=forceIfcold_start) - # TGYRO class. It checks existence and creates input.profiles/input.gacode - - self.tgyro = TGYROtools.TGYRO( - cdf=self.LocationCDF, time=self.time, avTime=self.avTime - ) - self.tgyro.prep( - FolderGACODE, - cold_start=cold_start, - remove_tmp=True, - subfolder="tmp_tgyro_prep", - profilesclass_custom=self.profiles, - forceIfcold_start=forceIfcold_start, - ) - - self.profiles = self.tgyro.profiles - - self.profiles.derive_quantities(mi_ref=md_u) - - self.profiles.correct(options={'recalculate_ptot':recalculate_ptot,'remove_fast':remove_fast}) - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize by preparing a tgyro class and running for -1 iterations + # Prepare state # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self.profiles = mitim_state - if specificInputs is None: + self.profiles.derive_quantities(mi_ref=md_u) - self.inputsTGLF = self.profiles.to_tglf(rhos=self.rhos) + self.profiles.correct(options={'recalculate_ptot':recalculate_ptot,'remove_fast':remove_fast}) - for rho in self.inputsTGLF: - self.inputsTGLF[rho] = TGLFinput.initialize_in_memory(self.inputsTGLF[rho]) - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize by taking directly the inputs + # Initialize from state # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - else: - self.inputsTGLF = specificInputs - self.tgyro_results = tgyro_results - - self.FolderGACODE = IOtools.expandPath(FolderGACODE) - - if cold_start or not self.FolderGACODE.exists(): - IOtools.askNewFolder(self.FolderGACODE, force=forceIfcold_start) + self.inputsTGLF = self.profiles.to_tglf(r=self.rhos, r_is_rho = True) for rho in self.inputsTGLF: + + # Initialize class + self.inputsTGLF[rho] = TGLFinput.initialize_in_memory(self.inputsTGLF[rho]) + + # Write input.tglf file self.inputsTGLF[rho].file = self.FolderGACODE / f'input.tglf_{rho:.4f}' self.inputsTGLF[rho].write_state() - """ - ~~~~~ Create Normalizations ~~~~~ - - Only input.gacode needed - - I can also give TRANSP CDF for complement. It is used in prep anyway, so good to store here - and have the values for plotting the experimental fluxes. - - I can also give TGYRO class for complement. It is used in prep anyway, so good to store here - for plotting and check grid conversions. - - Note about the TGLF normalization: - What matters is what's the mass used to normalized the MASS_X. - If TGYRO was used to generate the input.tglf file, then the normalization mass is deuterium and all - must be normalized to deuterium - """ + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Definining normalizations + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ print("> Setting up normalizations") - - print("\t- Using mass of deuterium to normalize things (not necesarily the first ion)",typeMsg="w",) - self.profiles.derive_quantities(mi_ref=md_u) - - self.NormalizationSets, cdf = NORMtools.normalizations( - self.profiles, - LocationCDF=self.LocationCDF, - time=self.time, - avTime=self.avTime, - cdf_open=cdf_open, - tgyro=self.tgyro_results, - ) + self.NormalizationSets, cdf = NORMtools.normalizations(self.profiles) return cdf @@ -4069,6 +4005,30 @@ def write_state(self, file=None): if file is None: file = self.file + # Local formatter: floats -> 6 significant figures in exponential (uppercase), + # ints stay as ints, bools as 0/1, sequences space-separated with same rule. + def _fmt_num(x): + import numpy as _np + if isinstance(x, (bool, _np.bool_)): + return "1" if x else "0" + if isinstance(x, (_np.floating, float)): + # 6 significant figures in exponential => 5 digits after decimal + return f"{float(x):.5E}" + if isinstance(x, (_np.integer, int)): + return f"{int(x)}" + return str(x) + + def _fmt_value(val): + import numpy as _np + if isinstance(val, (list, tuple, _np.ndarray)): + # Flatten numpy arrays but keep ordering; join with spaces + if isinstance(val, _np.ndarray): + flat = val.flatten().tolist() + else: + flat = list(val) + return " ".join(_fmt_num(v) for v in flat) + return _fmt_num(val) + with open(file, "w") as f: f.write("#-------------------------------------------------------------------------\n") f.write(f"# TGLF input file modified by MITIM {mitim_version}\n") @@ -4078,13 +4038,13 @@ def write_state(self, file=None): f.write("# ------------------\n\n") for ikey in self.controls: var = self.controls[ikey] - f.write(f"{ikey.ljust(23)} = {var}\n") + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") f.write("\n\n# Geometry parameters\n") f.write("# ------------------\n\n") for ikey in self.geom: var = self.geom[ikey] - f.write(f"{ikey.ljust(23)} = {var}\n") + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") f.write("\n\n# Plasma parameters\n") f.write("# ------------------\n\n") @@ -4095,7 +4055,7 @@ def write_state(self, file=None): print(f"\t- Maximum number of species in TGLF reached, not considering after {maxSpeciesTGLF} species",typeMsg="w",) else: var = self.plasma[ikey] - f.write(f"{ikey.ljust(23)} = {var}\n") + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") f.write("\n\n# Species\n") f.write("# -------\n") @@ -4117,7 +4077,7 @@ def write_state(self, file=None): f.write(f"\n# Specie #{ikey}{extralab}\n") for ivar in self.species[ikey]: ikar = f"{ivar}_{ikey}" - f.write(f"{ikar.ljust(12)} = {self.species[ikey][ivar]}\n") + f.write(f"{ikar.ljust(12)} = {_fmt_value(self.species[ikey][ivar])}\n") print(f"\t\t~ File {IOtools.clipstr(file)} written") diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 9c656ada..a558d040 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2197,11 +2197,17 @@ def csv(self, file="input.gacode.xlsx"): # Code conversions # ************************************************************************************************************************************************ - def to_tglf(self, rhos=[0.5], TGLFsettings=1): + def to_tglf(self, r=[0.5], TGLFsettings=1, r_is_rho = True): # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function + # Determine if the input radius is rho toroidal or r/a + if r_is_rho: + r_interpolation = self.profiles['rho(-)'] + else: + r_interpolation = self.derived['roa'] + # Determine the number of species to use in TGLF max_species_tglf = 6 # TGLF only accepts up to 6 species if len(self.Species) > max_species_tglf-1: @@ -2210,7 +2216,7 @@ def to_tglf(self, rhos=[0.5], TGLFsettings=1): else: tglf_ions_num = len(self.Species) - # Determinte the mass reference + # Determine the mass reference mass_ref = 2.0 # TODO: This is the only way to make it consistent with TGYRO (derivations with mD_u but mass in tglf with 2.0... https://github.com/gafusion/gacode/issues/398 # ----------------------------------------------------------------------- @@ -2262,14 +2268,14 @@ def to_tglf(self, rhos=[0.5], TGLFsettings=1): # --------------------------------------------------------------------------------------------------------------------------------------- inputsTGLF = {} - for rho in rhos: + for rho in r: # --------------------------------------------------------------------------------------------------------------------------------------- # Define interpolator at this rho # --------------------------------------------------------------------------------------------------------------------------------------- def interpolator(y): - return interpolation_function(rho, self.profiles['rho(-)'],y).item() + return interpolation_function(rho, r_interpolation,y).item() TGLFinput, TGLFoptions, label = GACODEdefaults.addTGLFcontrol(TGLFsettings) @@ -2327,6 +2333,7 @@ def interpolator(y): 'BETAE': interpolator(self.derived['betae']), } + # --------------------------------------------------------------------------------------------------------------------------------------- # Geometry comes from profiles # --------------------------------------------------------------------------------------------------------------------------------------- @@ -2357,10 +2364,10 @@ def interpolator(y): if int(ikey[-4]) > 6: continue - key_mod = ikey.upper().split('(')[0] # Remove any function call like 'shape_cos(1)' + key_mod = ikey.upper().split('(')[0] parameters[key_mod] = self.profiles[ikey] - parameters[f"{key_mod.split('_')[0]}_S_{key_mod.split('_')[-1]}"] = self.profiles[ikey]*0.0 #TODO + parameters[f"{key_mod.split('_')[0]}_S_{key_mod.split('_')[-1]}"] = self.derived["r"] * self._deriv_gacode(self.profiles[ikey]) geom = {} for k in parameters: From 82e6915d3e67f7511be82197a5f28b102239d230 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 19 Aug 2025 17:35:55 -0400 Subject: [PATCH 164/385] direct preparation should use what the input.gacode has --- src/mitim_modules/powertorch/physics_models/transport_tglf.py | 4 +--- src/mitim_tools/gacode_tools/TGLFtools.py | 4 ---- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 8ec3e5de..83ee3b20 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -58,10 +58,8 @@ def _evaluate_tglf(self): _ = tglf.prep_direct( self.folder, - self.powerstate.profiles_transport + self.powerstate.profiles_transport, cold_start = cold_start, - remove_fast = False, # Use what's in the input.gacode, the removal should happen at initialization - recalculate_ptot = False, # Use what's in the input.gacode, which is what PORTALS TGYRO does ) # ------------------------------------------------------------------------------------------------------------------------ diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 40af13b7..e970169c 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -361,8 +361,6 @@ def prep_direct( mitim_state, # A MITIM state class FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) cold_start=False, # If True, do not use what it potentially inside the folder, run again - remove_fast=False, # Ignore fast particles in TGYRO - recalculate_ptot=True, # Recalculate PTOT in TGYRO forceIfcold_start=False, # Extra flag ): @@ -381,8 +379,6 @@ def prep_direct( self.profiles.derive_quantities(mi_ref=md_u) - self.profiles.correct(options={'recalculate_ptot':recalculate_ptot,'remove_fast':remove_fast}) - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize from state # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ From 01f8c4c9d78dc2a4c8da57e08086d4b4c370a971 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 19 Aug 2025 17:38:14 -0400 Subject: [PATCH 165/385] Removal of deprecated options --- .../gacode_tools/utils/GACODEdefaults.py | 27 +++++++------------ src/mitim_tools/transp_tools/NMLtools.py | 6 ++--- 2 files changed, 11 insertions(+), 22 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index aa890fcd..8748f3b7 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -31,9 +31,7 @@ def addTGLFcontrol(TGLFsettings, NS=2, minimal=False): # Define every flag else: - TGLFoptions = IOtools.generateMITIMNamelist( - __mitimroot__ / "templates" / "input.tglf.controls", caseInsensitive=False - ) + TGLFoptions = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.tglf.controls", caseInsensitive=False) TGLFoptions["NMODES"] = NS + 2 """ @@ -48,29 +46,22 @@ def addTGLFcontrol(TGLFsettings, NS=2, minimal=False): 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 - # -------------------------------- - TGLFinput = [""] - for ikey in TGLFoptions: - TGLFinput.append(f"{ikey} = {TGLFoptions[ikey]}") - TGLFinput.append("") - TGLFinput.append("# -- Begin overlay") - TGLFinput.append("") + return TGLFoptions - return TGLFinput, TGLFoptions, label +def addNEOcontrol(): + NEOoptions = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.neo.controls", caseInsensitive=False) + + return NEOoptions def TGLFinTRANSP(TGLFsettings, NS=3): - _, TGLFoptions, label = addTGLFcontrol(TGLFsettings, NS=NS) + + TGLFoptions = addTGLFcontrol(TGLFsettings, NS=NS) """ ------------------------------------------------------------------------------------------------------ @@ -125,7 +116,7 @@ def TGLFinTRANSP(TGLFsettings, NS=3): # **** Other modifications TGLFoptions["UNITS"] = f"'{TGLFoptions['UNITS']}'" - return TGLFoptions, label + return TGLFoptions def addCGYROcontrol(Settings, rmin): diff --git a/src/mitim_tools/transp_tools/NMLtools.py b/src/mitim_tools/transp_tools/NMLtools.py index 0b8663d1..d53d92b1 100644 --- a/src/mitim_tools/transp_tools/NMLtools.py +++ b/src/mitim_tools/transp_tools/NMLtools.py @@ -1335,10 +1335,8 @@ def addGLF23(self): self.contents_ptr_glf23 = "\n".join(lines) + "\n" def addTGLF(self): - TGLFoptions, label = GACODEdefaults.TGLFinTRANSP(self.TGLFsettings) - print( - f"\t- Adding TGLF control parameters with TGLFsettings = {self.TGLFsettings} ({label})" - ) + TGLFoptions = GACODEdefaults.TGLFinTRANSP(self.TGLFsettings) + print(f"\t- Adding TGLF control parameters with TGLFsettings = {self.TGLFsettings}") lines = [ "!------ TGLF namelist", From daacd0f9e08b06f585b609482b1b8acd0426f50b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 19 Aug 2025 18:18:45 -0400 Subject: [PATCH 166/385] Minimal NEO object --- src/mitim_tools/gacode_tools/NEOtools.py | 150 ++++++++++++++++-- .../plasmastate_tools/MITIMstate.py | 144 ++++++++++++++++- templates/config_user_example.json | 4 +- templates/input.neo.controls | 18 +++ tests/NEO_workflow.py | 16 ++ 5 files changed, 317 insertions(+), 15 deletions(-) create mode 100644 templates/input.neo.controls create mode 100644 tests/NEO_workflow.py diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 12561fee..b5c194d3 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -1,8 +1,11 @@ +from pathlib import Path +from mitim_tools import __version__ as mitim_version from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools.utils import GACODErun from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +from mitim_tools.misc_tools.PLASMAtools import md_u class NEO: def __init__(self): @@ -14,6 +17,75 @@ def prep(self, inputgacode, folder): self.folder.mkdir(parents=True, exist_ok=True) + + def prep_direct( + self, + mitim_state, # A MITIM state class + FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) + cold_start=False, # If True, do not use what it potentially inside the folder, run again + forceIfcold_start=False, # Extra flag + ): + + print("> Preparation of NEO run from input.gacode (direct conversion)") + + self.FolderGACODE = IOtools.expandPath(FolderGACODE) + + if cold_start or not self.FolderGACODE.exists(): + IOtools.askNewFolder(self.FolderGACODE, force=forceIfcold_start) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare state + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self.profiles = mitim_state + + self.profiles.derive_quantities(mi_ref=md_u) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize from state + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self.inputsNEO = self.profiles.to_neo(r=self.rhos, r_is_rho = True) + + for rho in self.inputsNEO: + + # Initialize class + self.inputsNEO[rho] = NEOinput.initialize_in_memory(self.inputsNEO[rho]) + + # Write input.tglf file + self.inputsNEO[rho].file = self.FolderGACODE / f'input.neo_{rho:.4f}' + self.inputsNEO[rho].write_state() + + def run( + self, + ): + + pass + # tmp_folder = self.FolderGACODE / "tmp" + + # os.system(f'cp {self.FolderGACODE / f'input.neo_{rho:.4f}'}') + + # neo_job = FARMINGtools.mitim_job(self.FolderGACODE) + # neo_job.define_machine( + # "neo", + # f"mitim_neo" + # ) + + # neo_job.prep( + # 'neo -e .', + # input_folders=[self.FolderGACODE], + # output_folders=folders_red, + # check_files_in_folder=check_files_in_folder, + # shellPreCommands=shellPreCommands, + # shellPostCommands=shellPostCommands, + # ) + + # neo_job.run( + # removeScratchFolders=True, + # attempts_execution=attempts_execution + # ) + + def run_vgen(self, subfolder="vgen1", vgenOptions={}, cold_start=False): self.folder_vgen = self.folder / f"{subfolder}" @@ -41,10 +113,7 @@ def run_vgen(self, subfolder="vgen1", vgenOptions={}, cold_start=False): ) if (not runThisCase) and cold_start: - runThisCase = print( - "\t- Files found in folder, but cold_start requested. Are you sure?", - typeMsg="q", - ) + runThisCase = print("\t- Files found in folder, but cold_start requested. Are you sure?",typeMsg="q",) if runThisCase: IOtools.askNewFolder(self.folder_vgen, force=True) @@ -58,18 +127,13 @@ def run_vgen(self, subfolder="vgen1", vgenOptions={}, cold_start=False): self.folder_vgen, vgenOptions=vgenOptions, name_run=subfolder ) else: - print( - f"\t- Required files found in {subfolder}, not running VGEN", - typeMsg="i", - ) + print(f"\t- Required files found in {subfolder}, not running VGEN",typeMsg="i",) file_new = self.folder_vgen / f"vgen" / f"input.gacode" # ---- Postprocess from mitim_tools.gacode_tools import PROFILEStools - self.inputgacode_vgen = PROFILEStools.gacode_state( - file_new, derive_quantities=True, mi_ref=self.inputgacode.mi_ref - ) + self.inputgacode_vgen = PROFILEStools.gacode_state(file_new, derive_quantities=True, mi_ref=self.inputgacode.mi_ref) def check_if_files_exist(folder, list_files): @@ -83,3 +147,67 @@ def check_if_files_exist(folder, list_files): return False return True + + + +class NEOinput: + def __init__(self, file=None): + self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None + + if self.file is not None and self.file.exists(): + with open(self.file, "r") as f: + lines = f.readlines() + file_txt = "".join(lines) + else: + file_txt = "" + input_dict = GACODErun.buildDictFromInput(file_txt) + + self.process(input_dict) + + @classmethod + def initialize_in_memory(cls, input_dict): + instance = cls() + instance.process(input_dict) + return instance + + def process(self, input_dict): + #TODO + self.all = input_dict + + def write_state(self, file=None): + + if file is None: + file = self.file + + # Local formatter: floats -> 6 significant figures in exponential (uppercase), + # ints stay as ints, bools as 0/1, sequences space-separated with same rule. + def _fmt_num(x): + import numpy as _np + if isinstance(x, (bool, _np.bool_)): + return "1" if x else "0" + if isinstance(x, (_np.floating, float)): + # 6 significant figures in exponential => 5 digits after decimal + return f"{float(x):.5E}" + if isinstance(x, (_np.integer, int)): + return f"{int(x)}" + return str(x) + + def _fmt_value(val): + import numpy as _np + if isinstance(val, (list, tuple, _np.ndarray)): + # Flatten numpy arrays but keep ordering; join with spaces + if isinstance(val, _np.ndarray): + flat = val.flatten().tolist() + else: + flat = list(val) + return " ".join(_fmt_num(v) for v in flat) + return _fmt_num(val) + + with open(file, "w") as f: + f.write("#-------------------------------------------------------------------------\n") + f.write(f"# NEO input file modified by MITIM {mitim_version}\n") + f.write("#-------------------------------------------------------------------------") + + for ikey in self.all: + var = self.all[ikey] + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") \ No newline at end of file diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index a558d040..2b868094 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -384,6 +384,7 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["c_s"] = PLASMAtools.c_s(self.profiles["te(keV)"], self.derived["mi_ref"]) self.derived["rho_s"] = PLASMAtools.rho_s(self.profiles["te(keV)"], self.derived["mi_ref"], self.derived["B_unit"]) + self.derived["rho_sa"] = self.derived["rho_s"] / self.derived["a"] self.derived["q_gb"], self.derived["g_gb"], self.derived["pi_gb"], self.derived["s_gb"], _ = PLASMAtools.gyrobohmUnits( self.profiles["te(keV)"], @@ -2230,7 +2231,6 @@ def to_tglf(self, r=[0.5], TGLFsettings=1, r_is_rho = True): s_delta = self.derived["r"] * self._deriv_gacode(self.profiles["delta(-)"]) s_zeta = self.derived["r"] * self._deriv_gacode(self.profiles["zeta(-)"]) - ''' Total pressure -------------------------------------------------------- @@ -2277,7 +2277,7 @@ def to_tglf(self, r=[0.5], TGLFsettings=1, r_is_rho = True): def interpolator(y): return interpolation_function(rho, r_interpolation,y).item() - TGLFinput, TGLFoptions, label = GACODEdefaults.addTGLFcontrol(TGLFsettings) + TGLFoptions = GACODEdefaults.addTGLFcontrol(TGLFsettings) # --------------------------------------------------------------------------------------------------------------------------------------- # Controls come from options @@ -2356,7 +2356,7 @@ def interpolator(y): 'P_PRIME_LOC': pprime, } - # Add MXH and derivatives (#TODO) + # Add MXH and derivatives for ikey in self.profiles: if 'shape_cos' in ikey or 'shape_sin' in ikey: @@ -2391,6 +2391,144 @@ def interpolator(y): return inputsTGLF + def to_neo(self, r=[0.5], r_is_rho = True): + + # <> Function to interpolate a curve <> + from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function + + # Determine if the input radius is rho toroidal or r/a + if r_is_rho: + r_interpolation = self.profiles['rho(-)'] + else: + r_interpolation = self.derived['roa'] + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Prepare the inputs + # --------------------------------------------------------------------------------------------------------------------------------------- + + # Determine the mass reference + mass_ref = 2.0 + + sign_it = -np.sign(self.profiles["current(MA)"][-1]) + sign_bt = -np.sign(self.profiles["bcentr(T)"][-1]) + + s_kappa = self.derived["r"] / self.profiles["kappa(-)"] * self._deriv_gacode(self.profiles["kappa(-)"]) + s_delta = self.derived["r"] * self._deriv_gacode(self.profiles["delta(-)"]) + s_zeta = self.derived["r"] * self._deriv_gacode(self.profiles["zeta(-)"]) + + inputsNEO = {} + for rho in r: + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Define interpolator at this rho + # --------------------------------------------------------------------------------------------------------------------------------------- + + def interpolator(y): + return interpolation_function(rho, r_interpolation,y).item() + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Controls come from options + # --------------------------------------------------------------------------------------------------------------------------------------- + + controls = GACODEdefaults.addNEOcontrol() + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Species come from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + max_species_neo = 6 #TODO: True? + + species = { + 1: { + 'Z': -1.0, + 'MASS': 0.000272445, + 'DLNNDR': interpolator(self.derived['aLne']), + 'DLNTDR': interpolator(self.derived['aLTe']), + 'TEMP': 1.0, + 'DENS': 1.0, + } + } + + for i in range(min(len(self.Species), max_species_neo-1)): + species[i+2] = { + 'Z': self.Species[i]['Z'], + 'MASS': self.Species[i]['A']/mass_ref, + 'DLNNDR': interpolator(self.derived['aLni'][:,i]), + 'DLNTDR': interpolator(self.derived['aLTi'][:,0] if self.Species[i]['S'] == 'therm' else self.derived["aLTi"][:,i]), + 'TEMP': interpolator(self.derived["tite_all"][:,i]), + 'DENS': interpolator(self.derived['fi'][:,i]), + } + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Plasma comes from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + #TODO Does this work with no deuterium first ion? + factor_nu = species[2]['Z']**4 * species[2]['DENS'] * species[1]['MASS']**0.5 * species[2]['TEMP']**(-1.5) + + plasma = { + 'N_SPECIES': len(species), + 'IPCCW': sign_bt, + 'BTCCW': sign_it, + 'OMEGA_ROT': 0.0, #TODO + 'OMEGA_ROT_DERIV': 0.0, #TODO + 'NU_1': interpolator(self.derived['xnue'])* factor_nu, + 'RHO_STAR': interpolator(self.derived["rho_sa"]), + } + + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Geometry comes from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + parameters = { + 'RMIN_OVER_A': self.derived['roa'], + 'RMAJ_OVER_A': self.derived['Rmajoa'], + 'SHIFT': self._deriv_gacode(self.profiles["rmaj(m)"]), + 'ZMAG_OVER_A': self.derived["Zmagoa"], + 'S_ZMAG': self._deriv_gacode(self.profiles["zmag(m)"]), + 'Q': np.abs(self.profiles["q(-)"]), + 'SHEAR': self.derived["s_hat"], + 'KAPPA': self.profiles["kappa(-)"], + 'S_KAPPA': s_kappa, + 'DELTA': self.profiles["delta(-)"], + 'S_DELTA': s_delta, + 'ZETA': self.profiles["zeta(-)"], + 'S_ZETA': s_zeta, + } + + # Add MXH and derivatives + for ikey in self.profiles: + if 'shape_cos' in ikey or 'shape_sin' in ikey: + + # TGLF only accepts 6, as of July 2025 + if int(ikey[-4]) > 6: + continue + + key_mod = ikey.upper().split('(')[0] + + parameters[key_mod] = self.profiles[ikey] + parameters[f"{key_mod.split('_')[0]}_S_{key_mod.split('_')[-1]}"] = self.derived["r"] * self._deriv_gacode(self.profiles[ikey]) + + geom = {} + for k in parameters: + par = torch.nan_to_num(torch.from_numpy(parameters[k]) if type(parameters[k]) is np.ndarray else parameters[k], nan=0.0, posinf=1E10, neginf=-1E10) + geom[k] = interpolator(par) + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Merging + # --------------------------------------------------------------------------------------------------------------------------------------- + + input_dict = {**controls, **plasma, **geom} + + for i in range(len(species)): + for k in species[i+1]: + input_dict[f'{k}_{i+1}'] = species[i+1][k] + + inputsNEO[rho] = input_dict + + return inputsNEO + def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', times = [0.0,1.0], Vsurf = 0.0): print("\t- Converting to TRANSP") diff --git a/templates/config_user_example.json b/templates/config_user_example.json index 0aadcac0..df1223fe 100644 --- a/templates/config_user_example.json +++ b/templates/config_user_example.json @@ -6,6 +6,7 @@ "profiles_gen": "engaging", "tgyro": "engaging", "tglf": "engaging", + "neo": "engaging", "cgyro": "engaging", "astra": "engaging", "eq": "mfews", @@ -13,7 +14,8 @@ "ntcc": "mfews", "get_fbm": "mfews", "transp": "globus", - "idl": "mfews" + "idl": "mfews", + "eped": "engaging" }, "local": { "machine": "local", diff --git a/templates/input.neo.controls b/templates/input.neo.controls new file mode 100644 index 00000000..09673778 --- /dev/null +++ b/templates/input.neo.controls @@ -0,0 +1,18 @@ +#------------------------------------------------------------------------- +# Template input.neo file (controls-only) +#------------------------------------------------------------------------- + +# Resolution +N_ENERGY=5 +N_XI=11 +N_THETA=11 +N_RADIAL=1 + +# Geometry (Miller) +EQUILIBRIUM_MODEL=2 + +# Rotation (Sonic) +ROTATION_MODEL=2 + +# Collisions (FP) +COLLISION_MODEL=4 diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py new file mode 100644 index 00000000..73745b7a --- /dev/null +++ b/tests/NEO_workflow.py @@ -0,0 +1,16 @@ +import os +from mitim_tools.gacode_tools import NEOtools +from mitim_tools import __mitimroot__ + +cold_start = True + +(__mitimroot__ / 'tests' / 'scratch').mkdir(parents=True, exist_ok=True) + +folder = __mitimroot__ / "tests" / "scratch" / "neo_test" +input_gacode = __mitimroot__ / "tests" / "data" / "input.gacode" + +if cold_start and folder.exists(): + os.system(f"rm -r {folder.resolve()}") + +neo = NEOtools.NEO() +neo.prep_direct(folder, input_gacode) From c054051f929cbf90a39a490cf5c57c668db38fd9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 19 Aug 2025 18:49:58 -0400 Subject: [PATCH 167/385] Minimal local neo example running --- src/mitim_tools/gacode_tools/NEOtools.py | 75 ++++++++++++------- .../plasmastate_tools/MITIMstate.py | 4 +- tests/NEO_workflow.py | 10 ++- 3 files changed, 58 insertions(+), 31 deletions(-) diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index b5c194d3..698438e1 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -1,6 +1,7 @@ +import os from pathlib import Path from mitim_tools import __version__ as mitim_version -from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools import FARMINGtools, IOtools from mitim_tools.gacode_tools.utils import GACODErun from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -8,8 +9,12 @@ from mitim_tools.misc_tools.PLASMAtools import md_u class NEO: - def __init__(self): - pass + def __init__( + self, + rhos=[0.4, 0.6], # rho locations of interest + ): + + self.rhos = rhos def prep(self, inputgacode, folder): self.inputgacode = inputgacode @@ -58,33 +63,51 @@ def prep_direct( def run( self, + subfolder, + forceIfcold_start=False, ): + + # Create this run folder + + subfolder = Path(subfolder) - pass - # tmp_folder = self.FolderGACODE / "tmp" + FolderNEO = self.FolderGACODE / subfolder + IOtools.askNewFolder(FolderNEO, force=forceIfcold_start) + + folders, folders_red = [], [] + for rho in self.rhos: + # Create subfolder for each rho + FolderNEO_rho = FolderNEO / f"rho_{rho:.4f}" + IOtools.askNewFolder(FolderNEO_rho, force=forceIfcold_start) + + # Copy the file + os.system(f"cp {self.FolderGACODE / f'input.neo_{rho:.4f}'} {FolderNEO_rho / 'input.neo'}") + + folders.append(FolderNEO_rho) + folders_red.append(str(subfolder / f"rho_{rho:.4f}")) + + # Run NEO - # os.system(f'cp {self.FolderGACODE / f'input.neo_{rho:.4f}'}') + neo_job = FARMINGtools.mitim_job(self.FolderGACODE) + neo_job.define_machine_quick("neo",f"mitim_neo") - # neo_job = FARMINGtools.mitim_job(self.FolderGACODE) - # neo_job.define_machine( - # "neo", - # f"mitim_neo" - # ) - - # neo_job.prep( - # 'neo -e .', - # input_folders=[self.FolderGACODE], - # output_folders=folders_red, - # check_files_in_folder=check_files_in_folder, - # shellPreCommands=shellPreCommands, - # shellPostCommands=shellPostCommands, - # ) - - # neo_job.run( - # removeScratchFolders=True, - # attempts_execution=attempts_execution - # ) + NEOcommand = "" + for folder in folders_red: + NEOcommand += f"neo -e {folder} -p {neo_job.folderExecution} &\n" + NEOcommand += "wait\n" + + neo_job.define_machine("neo",f"mitim_neo") + + neo_job.prep( + NEOcommand, + input_folders=[FolderNEO], + output_folders=folders_red, + ) + + neo_job.run( + removeScratchFolders=True, + ) def run_vgen(self, subfolder="vgen1", vgenOptions={}, cold_start=False): @@ -206,7 +229,7 @@ def _fmt_value(val): with open(file, "w") as f: f.write("#-------------------------------------------------------------------------\n") f.write(f"# NEO input file modified by MITIM {mitim_version}\n") - f.write("#-------------------------------------------------------------------------") + f.write("#-------------------------------------------------------------------------\n") for ikey in self.all: var = self.all[ikey] diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 2b868094..8d5b8aea 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2409,8 +2409,8 @@ def to_neo(self, r=[0.5], r_is_rho = True): # Determine the mass reference mass_ref = 2.0 - sign_it = -np.sign(self.profiles["current(MA)"][-1]) - sign_bt = -np.sign(self.profiles["bcentr(T)"][-1]) + sign_it = int(-np.sign(self.profiles["current(MA)"][-1])) + sign_bt = int(-np.sign(self.profiles["bcentr(T)"][-1])) s_kappa = self.derived["r"] / self.profiles["kappa(-)"] * self._deriv_gacode(self.profiles["kappa(-)"]) s_delta = self.derived["r"] * self._deriv_gacode(self.profiles["delta(-)"]) diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 73745b7a..95a78c72 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -1,5 +1,5 @@ import os -from mitim_tools.gacode_tools import NEOtools +from mitim_tools.gacode_tools import NEOtools, PROFILEStools from mitim_tools import __mitimroot__ cold_start = True @@ -12,5 +12,9 @@ if cold_start and folder.exists(): os.system(f"rm -r {folder.resolve()}") -neo = NEOtools.NEO() -neo.prep_direct(folder, input_gacode) +neo = NEOtools.NEO( + rhos=[0.55] +) +neo.prep_direct(PROFILEStools.gacode_state(input_gacode), folder, ) + +neo.run('neo1') From 705e6dcc77b06b95af479c6336965fb0729e62c0 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 11:03:29 -0400 Subject: [PATCH 168/385] Generalization of gacode simulation --- .../physics_models/transport_tglf.py | 1 - src/mitim_tools/gacode_tools/NEOtools.py | 2 +- src/mitim_tools/gacode_tools/TGLFtools.py | 416 ++-------- .../gacode_tools/utils/GACODErun.py | 759 ++++++++++++------ tests/TGLF_workflow.py | 1 + tests/TGLFscan_workflow.py | 3 +- 6 files changed, 592 insertions(+), 590 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 83ee3b20..87ab6f6d 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -79,7 +79,6 @@ def _evaluate_tglf(self): cold_start= cold_start, forceIfcold_start=True, extra_name= self.name, - anticipate_problems=True, slurm_setup={ "cores": cores_per_tglf_instance, "minutes": 2, diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 698438e1..b4c1d958 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -207,7 +207,7 @@ def write_state(self, file=None): def _fmt_num(x): import numpy as _np if isinstance(x, (bool, _np.bool_)): - return "1" if x else "0" + return "True" if x else "False" if isinstance(x, (_np.floating, float)): # 6 significant figures in exponential => 5 digits after decimal return f"{float(x):.5E}" diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index e970169c..1ea3e048 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -28,7 +28,7 @@ MAX_TGLF_SPECIES = 6 -class TGLF: +class TGLF(GACODErun.gacode_simulation): def __init__( self, rhos=[0.4, 0.6], # rho locations of interest @@ -121,13 +121,18 @@ def __init__( ** Modify the class as wish, and do run,read, etc ** ** Because normalizations are stored in the prep phase, that's all ready ** """ - print( - "\n-----------------------------------------------------------------------------------------" - ) + + super().__init__(rhos=rhos) + + self.run_specifications = { + 'code': 'tglf', + 'input_file': 'input.tglf', + 'code_call': 'tglf -e', + 'control_function': GACODEdefaults.addTGLFcontrol + } + print("\n-----------------------------------------------------------------------------------------") print("\t\t\t TGLF class module") - print( - "-----------------------------------------------------------------------------------------\n" - ) + print("-----------------------------------------------------------------------------------------\n") if alreadyRun is not None: # For the case in which I have run TGLF somewhere else, not using to plot and modify the class @@ -135,6 +140,11 @@ def __init__( self.__dict__ = alreadyRun.__dict__ print("* Readying previously-run TGLF class", typeMsg="i") else: + + self.ResultsFiles_minimal = [ + "out.tglf.gbflux", + ] + self.ResultsFiles = [ "out.tglf.run", "out.tglf.gbflux", @@ -167,7 +177,6 @@ def __init__( else: self.nameRunid = "0" self.time, self.avTime = time, avTime - self.rhos = np.array(rhos) ( self.results, @@ -461,7 +470,6 @@ def run( cold_start=False, forceIfcold_start=False, extra_name="exe", - anticipate_problems=True, slurm_setup={ "cores": 4, "minutes": 5, @@ -476,7 +484,7 @@ def run( # Prepare inputs # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - tglf_executor, tglf_executor_full, folderlast = self._prepare_run_radii( + tglf_executor, tglf_executor_full, folderlast = self._prep( subFolderTGLF, tglf_executor={}, tglf_executor_full={}, @@ -493,7 +501,6 @@ def run( forceIfcold_start=forceIfcold_start, extra_name=extra_name, slurm_setup=slurm_setup, - anticipate_problems=anticipate_problems, attempts_execution=attempts_execution, only_minimal_files=only_minimal_files, ) @@ -520,6 +527,24 @@ def run( self.FolderTGLFlast = folderlast + def _prep( + self, + subFolder, + tglf_executor={}, + tglf_executor_full={}, + TGLFsettings=None, + **kwargs + ): + + return self._generic_prep( + subFolder, + code_executor=tglf_executor, + code_executor_full=tglf_executor_full, + code_settings=TGLFsettings, + addControlFunction=self.run_specifications['control_function'], + **kwargs + ) + def _run( self, tglf_executor, @@ -532,28 +557,11 @@ def _run( print("\n> Run TGLF") - if kwargs_TGLFrun.get("only_minimal_files", False): - filesToRetrieve = ["out.tglf.gbflux"] - else: - filesToRetrieve = self.ResultsFiles - - c = 0 - for subFolderTGLF in tglf_executor: - c += len(tglf_executor[subFolderTGLF]) - - if c > 0: - GACODErun.runTGLF( - self.FolderGACODE, - tglf_executor, - filesToRetrieve=filesToRetrieve, - minutes=kwargs_TGLFrun.get("slurm_setup", {}).get("minutes", 5), - cores_tglf=kwargs_TGLFrun.get("slurm_setup", {}).get("cores", 4), - name=f"tglf_{self.nameRunid}{kwargs_TGLFrun.get('extra_name', '')}", - launchSlurm=kwargs_TGLFrun.get("launchSlurm", True), - attempts_execution=kwargs_TGLFrun.get("attempts_execution", 1), - ) - else: - print("\t- TGLF not run because all results files found (please ensure consistency!)",typeMsg="i") + self._generic_run( + tglf_executor, + self.run_specifications, + **kwargs_TGLFrun + ) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Waveform if requested @@ -563,107 +571,6 @@ def _run( if "runWaveForms" in kwargs_TGLFrun and kwargs_TGLFrun["runWaveForms"] is not None and len(kwargs_TGLFrun["runWaveForms"]) > 0: self._run_wf(kwargs_TGLFrun["runWaveForms"], tglf_executor_full, **kwargs_TGLFrun) - def _prepare_run_radii( - self, - subFolderTGLF, # 'tglf1/', - rhos=None, - tglf_executor={}, - tglf_executor_full={}, - TGLFsettings=None, - extraOptions={}, - multipliers={}, - minimum_delta_abs={}, - ApplyCorrections=True, # Removing ions with too low density and that are fast species - Quasineutral=False, # Ensures quasineutrality. By default is False because I may want to run the file directly - launchSlurm=True, - cold_start=False, - forceIfcold_start=False, - anticipate_problems=True, - slurm_setup={ - "cores": 4, - "minutes": 5, - }, # Cores per TGLF call (so, when running nR radii -> nR*4) - only_minimal_files=False, - **kwargs): - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare for run - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if rhos is None: - rhos = self.rhos - - inputs = copy.deepcopy(self.inputsTGLF) - FolderTGLF = self.FolderGACODE / subFolderTGLF - - ResultsFiles_new = [] - for i in self.ResultsFiles: - if "mitim.out" not in i: - ResultsFiles_new.append(i) - self.ResultsFiles = ResultsFiles_new - - if only_minimal_files: - filesToRetrieve = ["out.tglf.gbflux"] - else: - filesToRetrieve = self.ResultsFiles - - # Do I need to run all radii? - rhosEvaluate = cold_start_checker( - rhos, - filesToRetrieve, - FolderTGLF, - cold_start=cold_start, - ) - - if len(rhosEvaluate) == len(rhos): - # All radii need to be evaluated - IOtools.askNewFolder(FolderTGLF, force=forceIfcold_start) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Change this specific run of TGLF - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - ( - latest_inputsFileTGLF, - latest_inputsFileTGLFDict, - ) = changeANDwrite_TGLF( - rhos, - inputs, - FolderTGLF, - TGLFsettings=TGLFsettings, - extraOptions=extraOptions, - multipliers=multipliers, - minimum_delta_abs=minimum_delta_abs, - ApplyCorrections=ApplyCorrections, - Quasineutral=Quasineutral, - ) - - tglf_executor_full[subFolderTGLF] = {} - tglf_executor[subFolderTGLF] = {} - for irho in self.rhos: - tglf_executor_full[subFolderTGLF][irho] = { - "folder": FolderTGLF, - "dictionary": latest_inputsFileTGLFDict[irho], - "inputs": latest_inputsFileTGLF[irho], - "extraOptions": extraOptions, - "multipliers": multipliers, - } - if irho in rhosEvaluate: - tglf_executor[subFolderTGLF][irho] = tglf_executor_full[subFolderTGLF][ - irho - ] - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Stop if I expect problems - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if anticipate_problems: - anticipate_problems_func( - latest_inputsFileTGLFDict, rhosEvaluate, slurm_setup, launchSlurm - ) - - return tglf_executor, tglf_executor_full, FolderTGLF - def _run_wf(self, kys, tglf_executor, **kwargs_TGLFrun): """ extraOptions and multipliers are not being grabbed from kwargs_TGLFrun, but from tglf_executor @@ -697,13 +604,10 @@ def _run_wf(self, kys, tglf_executor, **kwargs_TGLFrun): ky_single_orig = copy.deepcopy(ky_single0) - FolderTGLF_old = tglf_executor[subFolderTGLF][ - list(tglf_executor[subFolderTGLF].keys())[0] - ]["folder"] + FolderTGLF_old = tglf_executor[subFolderTGLF][list(tglf_executor[subFolderTGLF].keys())[0]]["folder"] self.ky_single = None - self.read( - label=f"ky{ky_single0}", folder=FolderTGLF_old, cold_startWF = False) + self.read(label=f"ky{ky_single0}", folder=FolderTGLF_old, cold_startWF = False) self.ky_single = kys self.FoldersTGLF_WF[f"ky{ky_single0}"][ @@ -713,37 +617,18 @@ def _run_wf(self, kys, tglf_executor, **kwargs_TGLFrun): ky_singles = [] for i, ir in enumerate(self.rhos): # -------- Get the closest unstable mode to the one requested - if ( - kwargs_TGLFrun["forceClosestUnstableWF"] - if "forceClosestUnstableWF" in kwargs_TGLFrun - else True - ): + if kwargs_TGLFrun.get("forceClosestUnstableWF", True): + # Only unstable ones kys_n = [] - for j in range( - len(self.results[f"ky{ky_single0}"]["TGLFout"][i].ky) - ): - if ( - self.results[f"ky{ky_single0}"]["TGLFout"][i].g[ - 0, j - ] - > 0.0 - ): - kys_n.append( - self.results[f"ky{ky_single0}"]["TGLFout"][ - i - ].ky[j] - ) + for j in range(len(self.results[f"ky{ky_single0}"]["TGLFout"][i].ky)): + if self.results[f"ky{ky_single0}"]["TGLFout"][i].g[0, j] > 0.0: + kys_n.append(self.results[f"ky{ky_single0}"]["TGLFout"][i].ky[j]) kys_n = np.array(kys_n) # ---- - closest_ky = kys_n[ - np.argmin(np.abs(kys_n - ky_single_orig)) - ] - print( - f"\t- rho = {ir:.3f}, requested ky={ky_single_orig:.3f}, & closest unstable ky based on previous run: ky={closest_ky:.3f}", - typeMsg="i", - ) + closest_ky = kys_n[np.argmin(np.abs(kys_n - ky_single_orig))] + print(f"\t- rho = {ir:.3f}, requested ky={ky_single_orig:.3f}, & closest unstable ky based on previous run: ky={closest_ky:.3f}",typeMsg="i",) ky_single = closest_ky else: ky_single = ky_single0 @@ -759,21 +644,15 @@ def _run_wf(self, kys, tglf_executor, **kwargs_TGLFrun): else: extraOptions_WF = {} - extraOptions_WF = copy.deepcopy(tglf_executor[subFolderTGLF][ - list(tglf_executor[subFolderTGLF].keys())[0] - ]["extraOptions"]) - multipliers_WF = copy.deepcopy(tglf_executor[subFolderTGLF][ - list(tglf_executor[subFolderTGLF].keys())[0] - ]["multipliers"]) + extraOptions_WF = copy.deepcopy(tglf_executor[subFolderTGLF][list(tglf_executor[subFolderTGLF].keys())[0]]["extraOptions"]) + multipliers_WF = copy.deepcopy(tglf_executor[subFolderTGLF][list(tglf_executor[subFolderTGLF].keys())[0]]["multipliers"]) 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 - ) + extraOptions_WF["VEXB_SHEAR"] = 0.0 # See email from G. Staebler on 05/16/2021 - tglf_executorWF, _, _ = self._prepare_run_radii( + tglf_executorWF, _, _ = self._prep( (FolderTGLF_old / f"ky{ky_single0}").relative_to(FolderTGLF_old.parent), tglf_executor=tglf_executorWF, extraOptions=extraOptions_WF, @@ -2306,7 +2185,7 @@ def _prepare_scan( "forceIfcold_start" in kwargs_TGLFrun and kwargs_TGLFrun["forceIfcold_start"] ) - tglf_executor, tglf_executor_full, folderlast = self._prepare_run_radii( + tglf_executor, tglf_executor_full, folderlast = self._prep( f"{self.subFolderTGLF_scan}_{name}", tglf_executor=tglf_executor, tglf_executor_full=tglf_executor_full, @@ -3673,75 +3552,6 @@ def completeVariation(setVariations, species): return setVariations_new - -# ~~~~~~~~~ Input class - - -def changeANDwrite_TGLF( - rhos, - inputs0, - FolderTGLF, - TGLFsettings=None, - extraOptions={}, - multipliers={}, - minimum_delta_abs={}, - ApplyCorrections=True, - Quasineutral=False, -): - """ - Received inputs classes and gives text. - ApplyCorrections refer to removing ions with too low density and that are fast species - """ - - inputs = copy.deepcopy(inputs0) - - modInputTGLF = {} - ns_max = [] - for i, rho in enumerate(rhos): - print(f"\t- Changing input file for rho={rho:.4f}") - NS = inputs[rho].plasma["NS"] - inputTGLF_rho = GACODErun.modifyInputs( - inputs[rho], - Settings=TGLFsettings, - extraOptions=extraOptions, - multipliers=multipliers, - minimum_delta_abs=minimum_delta_abs, - position_change=i, - addControlFunction=GACODEdefaults.addTGLFcontrol, - NS=NS, - ) - - newfile = FolderTGLF / f"input.tglf_{rho:.4f}" - - if TGLFsettings is not None: - # Apply corrections - if ApplyCorrections: - print("\t- Applying corrections") - inputTGLF_rho.removeLowDensitySpecie() - inputTGLF_rho.remove_fast() - - # Ensure that plasma to run is quasineutral - if Quasineutral: - inputTGLF_rho.ensureQuasineutrality() - else: - print('\t- Not applying corrections nor quasineutrality because "TGLFsettings" is None') - - inputTGLF_rho.write_state(file=newfile) - - modInputTGLF[rho] = inputTGLF_rho - - ns_max.append(inputs[rho].plasma["NS"]) - - # Convert back to a string because that's how runTGLFproduction operates - inputFileTGLF = inputToVariable(FolderTGLF, rhos) - - if (np.diff(ns_max) > 0).any(): - print("> Each radial location has its own number of species... probably because of removal of fast or low density...",typeMsg="w") - print("\t * Reading of TGLF results will fail... consider doing something before launching run",typeMsg="q") - - return inputFileTGLF, modInputTGLF - - def reduceToControls(dict_all): controls, plasma, geom = {}, {}, {} for ikey in dict_all: @@ -3867,6 +3677,21 @@ def processSpecies(self, MinMultiplierToBeFast=2.0): print("\t- No species in this input.tglf (it is either a controls-only file or there was a problem generating it)") self.onlyControl = True + def anticipate_problems(self): + + threshold = 1e-10 + + minn = [] + for cont, ip in enumerate(self.species): + if (cont <= self.num_recorded) and ( + self.species[ip]["AS"] < threshold + ): + minn.append(ip) + + if len(minn) > 0: + print(f"* Ions in positions {ip} have a relative density lower than {threshold}, which can cause problems",typeMsg="q") + + def isThePlasmaDT(self): """ First two ions are D and T? @@ -4006,7 +3831,7 @@ def write_state(self, file=None): def _fmt_num(x): import numpy as _np if isinstance(x, (bool, _np.bool_)): - return "1" if x else "0" + return "True" if x else "False" if isinstance(x, (_np.floating, float)): # 6 significant figures in exponential => 5 digits after decimal return f"{float(x):.5E}" @@ -4338,25 +4163,6 @@ def identifySpecie(dict_species, dict_find): return found_index -# From file to dict - - -def inputToVariable(finalFolder, rhos): - """ - Entire text file to variable - """ - - inputFilesTGLF = {} - for cont, rho in enumerate(rhos): - fileN = finalFolder / f"input.tglf_{rho:.4f}" - - with open(fileN, "r") as f: - lines = f.readlines() - inputFilesTGLF[rho] = "".join(lines) - - return inputFilesTGLF - - # ~~~~~~~~~~~~~ Functions to handle results @@ -6314,83 +6120,3 @@ def createCombinedRuns(tglfs=(), new_names=(), results_names=(), isItScan=False) normalizations[new_names[i]] = tglfs[i].NormalizationSets return tglf, normalizations - - -def cold_start_checker( - rhos, - ResultsFiles, - FolderTGLF, - cold_start=False, - print_each_time=False, -): - """ - This function checks if the TGLF inputs are already in the folder. If they are, it returns True - """ - cont_each = 0 - if cold_start: - rhosEvaluate = rhos - else: - rhosEvaluate = [] - for ir in rhos: - existsRho = True - for j in ResultsFiles: - ffi = FolderTGLF / f"{j}_{ir:.4f}" - existsThis = ffi.exists() - existsRho = existsRho and existsThis - if not existsThis: - if print_each_time: - print(f"\t* {ffi} does not exist") - else: - cont_each += 1 - if not existsRho: - rhosEvaluate.append(ir) - - if not print_each_time and cont_each > 0: - print(f'\t* {cont_each} files from expected set are missing') - - if len(rhosEvaluate) < len(rhos) and len(rhosEvaluate) > 0: - print( - "~ Not all radii are found, but not removing folder and running only those that are needed", - typeMsg="i", - ) - - return rhosEvaluate - - -def anticipate_problems_func( - latest_inputsFileTGLFDict, rhosEvaluate, slurm_setup, launchSlurm -): - - # ----------------------------------- - # ------ Check density for problems - # ----------------------------------- - - threshold = 1e-10 - - minn = [] - for irho in latest_inputsFileTGLFDict: - for cont, ip in enumerate(latest_inputsFileTGLFDict[irho].species): - if (cont <= latest_inputsFileTGLFDict[irho].plasma["NS"]) and ( - latest_inputsFileTGLFDict[irho].species[ip]["AS"] < threshold - ): - minn.append([irho, ip]) - - if len(minn) > 0: - print( - f"* Ions in positions [rho,pos] {minn} have a relative density lower than {threshold}, which can cause problems", - typeMsg="q", - ) - - # ----------------------------------- - # ------ Check cores problem - # ----------------------------------- - - expected_allocated_cores = int(len(rhosEvaluate) * slurm_setup["cores"]) - - warning = 32 * 2 - - if launchSlurm: - print( - f'\t- Slurm job will be submitted with {expected_allocated_cores} cores ({len(rhosEvaluate)} radii x {slurm_setup["cores"]} cores/radius)', - typeMsg="" if expected_allocated_cores < warning else "q", - ) diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 499e67d2..36198bbf 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -1,5 +1,6 @@ import shutil import os +import copy import numpy as np import matplotlib.pyplot as plt from scipy.interpolate import interp1d @@ -9,6 +10,520 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +class gacode_simulation: + def __init__( + self, + rhos=[0.4, 0.6], # rho locations of interest + ): + self.rhos = np.array(rhos) if rhos is not None else None + + self.ResultsFiles = [] + self.ResultsFiles_minimal = [] + + def _generic_prep( + self, + subfolder_simulation, + rhos=None, + code_executor=None, + code_executor_full=None, + code_settings=None, + extraOptions={}, + multipliers={}, + cold_start=False, + forceIfcold_start=False, + only_minimal_files=False, + minimum_delta_abs={}, + ApplyCorrections=True, + Quasineutral=False, + launchSlurm=True, + slurm_setup={ + "cores": 4, + "minutes": 5, + }, # Cores per TGLF call (so, when running nR radii -> nR*4) + addControlFunction=None, + **kwargs + ): + + if code_executor is None: + code_executor = {} + if code_executor_full is None: + code_executor_full = {} + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare for run + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + if rhos is None: + rhos = self.rhos + + inputs = copy.deepcopy(self.inputsTGLF) + Folder_sim = self.FolderGACODE / subfolder_simulation + + ResultsFiles_new = [] + for i in self.ResultsFiles: + if "mitim.out" not in i: + ResultsFiles_new.append(i) + self.ResultsFiles = ResultsFiles_new + + if only_minimal_files: + filesToRetrieve = self.ResultsFiles_minimal + else: + filesToRetrieve = self.ResultsFiles + + # Do I need to run all radii? + rhosEvaluate = cold_start_checker( + rhos, + filesToRetrieve, + Folder_sim, + cold_start=cold_start, + ) + + if len(rhosEvaluate) == len(rhos): + # All radii need to be evaluated + IOtools.askNewFolder(Folder_sim, force=forceIfcold_start) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Change this specific run of TGLF + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + ( + latest_inputsFile, + latest_inputsFileDict, + ) = change_and_write_code( + rhos, + inputs, + Folder_sim, + code_settings=code_settings, + extraOptions=extraOptions, + multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, + ApplyCorrections=ApplyCorrections, + Quasineutral=Quasineutral, + addControlFunction=addControlFunction + ) + + code_executor_full[subfolder_simulation] = {} + code_executor[subfolder_simulation] = {} + for irho in self.rhos: + code_executor_full[subfolder_simulation][irho] = { + "folder": Folder_sim, + "dictionary": latest_inputsFileDict[irho], + "inputs": latest_inputsFile[irho], + "extraOptions": extraOptions, + "multipliers": multipliers, + } + if irho in rhosEvaluate: + code_executor[subfolder_simulation][irho] = code_executor_full[subfolder_simulation][ + irho + ] + + # Check input file problems + for irho in latest_inputsFileDict: + latest_inputsFileDict[irho].anticipate_problems() + + # Check cores problem + expected_allocated_cores = int(len(rhosEvaluate) * slurm_setup["cores"]) + warning = 32 * 2 + if launchSlurm: + print(f'\t- Slurm job will be submitted with {expected_allocated_cores} cores ({len(rhosEvaluate)} radii x {slurm_setup["cores"]} cores/radius)', + typeMsg="" if expected_allocated_cores < warning else "q",) + + return code_executor, code_executor_full, Folder_sim + + def _generic_run( + self, + code_executor, + run_specifications, + **kwargs_run + ): + + if kwargs_run.get("only_minimal_files", False): + filesToRetrieve = self.ResultsFiles_minimal + else: + filesToRetrieve = self.ResultsFiles + + c = 0 + for subfolder_simulation in code_executor: + c += len(code_executor[subfolder_simulation]) + + if c > 0: + run_gacode_simulation( + self.FolderGACODE, + code_executor, + run_specifications=run_specifications, + filesToRetrieve=filesToRetrieve, + minutes=kwargs_run.get("slurm_setup", {}).get("minutes", 5), + cores_simulation=kwargs_run.get("slurm_setup", {}).get("cores", 4), + name=f"{run_specifications['code']}_{self.nameRunid}{kwargs_run.get('extra_name', '')}", + launchSlurm=kwargs_run.get("launchSlurm", True), + attempts_execution=kwargs_run.get("attempts_execution", 1), + ) + else: + print(f"\t- {run_specifications['code'].upper()} not run because all results files found (please ensure consistency!)",typeMsg="i") + + +def change_and_write_code( + rhos, + inputs0, + Folder_sim, + code_settings=None, + extraOptions={}, + multipliers={}, + minimum_delta_abs={}, + ApplyCorrections=True, + Quasineutral=False, + addControlFunction=None, +): + """ + Received inputs classes and gives text. + ApplyCorrections refer to removing ions with too low density and that are fast species + """ + + inputs = copy.deepcopy(inputs0) + + mod_input_file = {} + ns_max = [] + for i, rho in enumerate(rhos): + print(f"\t- Changing input file for rho={rho:.4f}") + input_sim_rho = modifyInputs( + inputs[rho], + Settings=code_settings, + extraOptions=extraOptions, + multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, + position_change=i, + addControlFunction=addControlFunction, + NS=inputs[rho].num_recorded, + ) + + newfile = Folder_sim / f"input.tglf_{rho:.4f}" + + if code_settings is not None: + # Apply corrections + if ApplyCorrections: + print("\t- Applying corrections") + input_sim_rho.removeLowDensitySpecie() + input_sim_rho.remove_fast() + + # Ensure that plasma to run is quasineutral + if Quasineutral: + input_sim_rho.ensureQuasineutrality() + else: + print('\t- Not applying corrections because settings is None') + + input_sim_rho.write_state(file=newfile) + + mod_input_file[rho] = input_sim_rho + + ns_max.append(inputs[rho].num_recorded) + + # Convert back to a string because that's how runTGLFproduction operates + inputFile = inputToVariable(Folder_sim, rhos, file='input.tglf') + + if (np.diff(ns_max) > 0).any(): + print("> Each radial location has its own number of species... probably because of removal of fast or low density...",typeMsg="w") + print("\t * Reading of simulation results will fail... consider doing something before launching run",typeMsg="q") + + return inputFile, mod_input_file + +def inputToVariable(finalFolder, rhos, file='input.tglf'): + """ + Entire text file to variable + """ + + inputFilesTGLF = {} + for cont, rho in enumerate(rhos): + fileN = finalFolder / f"{file}_{rho:.4f}" + + with open(fileN, "r") as f: + lines = f.readlines() + inputFilesTGLF[rho] = "".join(lines) + + return inputFilesTGLF + + +def cold_start_checker( + rhos, + ResultsFiles, + Folder_sim, + cold_start=False, + print_each_time=False, +): + """ + This function checks if the TGLF inputs are already in the folder. If they are, it returns True + """ + cont_each = 0 + if cold_start: + rhosEvaluate = rhos + else: + rhosEvaluate = [] + for ir in rhos: + existsRho = True + for j in ResultsFiles: + ffi = Folder_sim / f"{j}_{ir:.4f}" + existsThis = ffi.exists() + existsRho = existsRho and existsThis + if not existsThis: + if print_each_time: + print(f"\t* {ffi} does not exist") + else: + cont_each += 1 + if not existsRho: + rhosEvaluate.append(ir) + + if not print_each_time and cont_each > 0: + print(f'\t* {cont_each} files from expected set are missing') + + if len(rhosEvaluate) < len(rhos) and len(rhosEvaluate) > 0: + print("~ Not all radii are found, but not removing folder and running only those that are needed",typeMsg="i",) + + return rhosEvaluate + + +def run_gacode_simulation( + FolderGACODE, + code_executor, + run_specifications = None, + minutes = 5, + cores_simulation = 4, + extraFlag = "", + filesToRetrieve = None, + name = "", + launchSlurm = True, + attempts_execution = 1, + max_jobs_at_once = None, +): + """ + launchSlurm = True -> Launch as a batch job in the machine chosen + launchSlurm = False -> Launch locally as a bash script + """ + + code = run_specifications.get('code', 'tglf') + input_file = run_specifications.get('input_file', 'input.tglf') + code_call = run_specifications.get('code_call', 'tglf -e') + + tmpFolder = FolderGACODE / f"tmp_{code}" + IOtools.askNewFolder(tmpFolder, force=True) + + gacode_job = FARMINGtools.mitim_job(tmpFolder) + + gacode_job.define_machine_quick(code,f"mitim_{name}") + + folders, folders_red = [], [] + for subfolder_sim in code_executor: + + rhos = list(code_executor[subfolder_sim].keys()) + + # --------------------------------------------- + # Prepare files and folders + # --------------------------------------------- + + for i, rho in enumerate(rhos): + print(f"\t- Preparing {code.upper()} execution ({subfolder_sim}) at rho={rho:.4f}") + + folder_sim_this = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" + folders.append(folder_sim_this) + + folder_sim_this_rel = folder_sim_this.relative_to(tmpFolder) + folders_red.append(folder_sim_this_rel.as_posix() if gacode_job.machineSettings['machine'] != 'local' else str(folder_sim_this_rel)) + + folder_sim_this.mkdir(parents=True, exist_ok=True) + + input_file_sim = folder_sim_this / input_file + with open(input_file_sim, "w") as f: + f.write(code_executor[subfolder_sim][rho]["inputs"]) + + # --------------------------------------------- + # Prepare command + # --------------------------------------------- + + # Grab machine local limits ------------------------------------------------- + max_cores_per_node = FARMINGtools.mitim_job.grab_machine_settings(code)["cores_per_node"] + + # If the run is local and not slurm, let's check the number of cores + if (FARMINGtools.mitim_job.grab_machine_settings(code)["machine"] == "local") and not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): + + cores_in_machine = int(os.cpu_count()) + cores_allocated = int(os.environ.get('SLURM_CPUS_PER_TASK')) if os.environ.get('SLURM_CPUS_PER_TASK') is not None else None + + if cores_allocated is not None: + if max_cores_per_node is None or (cores_allocated < max_cores_per_node): + print(f"\t - Detected {cores_allocated} cores allocated by SLURM, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node = cores_allocated + elif cores_in_machine is not None: + if max_cores_per_node is None or (cores_in_machine < max_cores_per_node): + print(f"\t - Detected {cores_in_machine} cores in machine, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node = cores_in_machine + else: + # Default to just 16 just in case + if max_cores_per_node is None: + max_cores_per_node = 16 + else: + # For remote execution, default to just 16 just in case + if max_cores_per_node is None: + max_cores_per_node = 16 + # --------------------------------------------------------------------------- + + # Grab the total number of cores of this job -------------------------------- + total_simulation_executions = len(rhos) * len(code_executor) + total_cores_required = int(cores_simulation) * total_simulation_executions + # --------------------------------------------------------------------------- + + # Simply bash, no slurm + if not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): + + max_parallel_execution = max_cores_per_node // cores_simulation # Make sure we don't overload the machine when running locally (assuming no farming trans-node) + + print(f"\t- {code.upper()} will be executed as bash script (total cores: {total_cores_required}, cores per simulation: {cores_simulation}). MITIM will launch {total_simulation_executions // max_parallel_execution+1} sequential executions",typeMsg="i") + + # Build the bash script with job control enabled and a loop to limit parallel jobs + GACODEcommand = "#!/usr/bin/env bash\n" + GACODEcommand += "set -m\n" # Enable job control even in non-interactive mode + GACODEcommand += f"max_parallel_execution={max_parallel_execution}\n\n" # Set the maximum number of parallel processes + + # Create a bash array of folders + GACODEcommand += "folders=(\n" + for folder in folders_red: + GACODEcommand += f' "{folder}"\n' + GACODEcommand += ")\n\n" + + # Loop over each folder and launch code, waiting if we've reached max_parallel_execution + GACODEcommand += "for folder in \"${folders[@]}\"; do\n" + GACODEcommand += f" {code_call} \"$folder\" -n {cores_simulation} -p {gacode_job.folderExecution} &\n" + GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" + GACODEcommand += "done\n\n" + GACODEcommand += "wait\n" + + # Slurm setup + array_list = None + job_array_limit = None + shellPreCommands = None + shellPostCommands = None + ntasks = total_cores_required + cpuspertask = cores_simulation + + else: + + # Standard job + if total_cores_required < max_cores_per_node: + + print(f"\t- {code.upper()} will be executed in SLURM as standard job (cpus: {total_cores_required})",typeMsg="i") + + # Code launches + GACODEcommand = "" + for folder in folders_red: + GACODEcommand += f"{code_call} {folder} -n {cores_simulation} -p {gacode_job.folderExecution} &\n" + GACODEcommand += "\nwait" # This is needed so that the script doesn't end before each job + + # Slurm setup + array_list = None + job_array_limit = None + shellPreCommands = None + shellPostCommands = None + ntasks = total_simulation_executions + cpuspertask = cores_simulation + + # Job array + else: + + print(f"\t- {code.upper()} will be executed in SLURM as job array due to its size (cpus: {total_cores_required})",typeMsg="i") + + # As a pre-command, organize all folders in a simpler way + shellPreCommands = [] + shellPostCommands = [] + array_list = [] + for i, folder in enumerate(folders_red): + array_list.append(f"{i}") + folder_temp_array = f"run{i}" + folder_actual = folder + shellPreCommands.append(f"mkdir {gacode_job.folderExecution}/{folder_temp_array}; cp {gacode_job.folderExecution}/{folder_actual}/* {gacode_job.folderExecution}/{folder_temp_array}/.") + shellPostCommands.append(f"cp {gacode_job.folderExecution}/{folder_temp_array}/* {gacode_job.folderExecution}/{folder_actual}/.; rm -r {gacode_job.folderExecution}/{folder_temp_array}") + + # Code launches + indexed_folder = 'run"$SLURM_ARRAY_TASK_ID"' + GACODEcommand = f'{code_call} {indexed_folder} -n {cores_simulation} -p {gacode_job.folderExecution} 1> {gacode_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {gacode_job.folderExecution}/{indexed_folder}/slurm_error.dat\n' + + # Slurm setup + array_list = ",".join(array_list) + ntasks = 1 + cpuspertask = cores_simulation + job_array_limit = max_jobs_at_once # Limit to this number at most running jobs at the same time + + # --------------------------------------------- + # Execute + # --------------------------------------------- + + gacode_job.define_machine( + code, + f"mitim_{name}", + launchSlurm=launchSlurm, + slurm_settings={ + "minutes": minutes, + "ntasks": ntasks, + "name": name, + "cpuspertask": cpuspertask, + "job_array": array_list, + "job_array_limit": job_array_limit, + #"nodes": 1, + }, + ) + + # I would like the mitim_job to check if the retrieved folders were complete + check_files_in_folder = {} + for folder in folders_red: + check_files_in_folder[folder] = filesToRetrieve + # --------------------------------------------- + + gacode_job.prep( + GACODEcommand, + input_folders=folders, + output_folders=folders_red, + check_files_in_folder=check_files_in_folder, + shellPreCommands=shellPreCommands, + shellPostCommands=shellPostCommands, + ) + + gacode_job.run( + removeScratchFolders=False, + attempts_execution=attempts_execution + ) + + # --------------------------------------------- + # Organize + # --------------------------------------------- + + print("\t- Retrieving files and changing names for storing") + fineall = True + for subfolder_sim in code_executor: + + for i, rho in enumerate(code_executor[subfolder_sim].keys()): + for file in filesToRetrieve: + original_file = f"{file}_{rho:.4f}{extraFlag}" + final_destination = ( + code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" + ) + final_destination.unlink(missing_ok=True) + + temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" + temp_file.replace(final_destination) + + fineall = fineall and final_destination.exists() + + if not final_destination.exists(): + print(f"\t!! file {file} ({original_file}) could not be retrived",typeMsg="w",) + + if fineall: + print("\t\t- All files were successfully retrieved") + + # Remove temporary folder + shutil.rmtree(tmpFolder) + + else: + print("\t\t- Some files were not retrieved", typeMsg="w") + + + + def runTGYRO( folderWork, outputFiles=None, @@ -112,10 +627,10 @@ def modifyInputs( # ------------------------------------------- if Settings is not None: - _, CodeOptions, label = addControlFunction(Settings, **kwargs_to_function) + CodeOptions = addControlFunction(Settings, **kwargs_to_function) # ~~~~~~~~~~ Change with presets - print(f" \t- Using presets Settings = {Settings} ({label})", typeMsg="i") + print(f" \t- Using presets Settings = {Settings}", typeMsg="i") input_class.controls = CodeOptions else: @@ -903,243 +1418,3 @@ def defineNewGrid( plt.show() return x[imin:imax], y[imin:imax] - - -def runTGLF( - FolderGACODE, - tglf_executor, - minutes=5, - cores_tglf=4, - extraFlag="", - filesToRetrieve=["out.tglf.gbflux"], - name="", - launchSlurm=True, - attempts_execution=1, - max_jobs_at_once=None, -): - """ - launchSlurm = True -> Launch as a batch job in the machine chosen - launchSlurm = False -> Launch locally as a bash script - """ - - tmpFolder = FolderGACODE / "tmp_tglf" - IOtools.askNewFolder(tmpFolder, force=True) - - tglf_job = FARMINGtools.mitim_job(tmpFolder) - - tglf_job.define_machine_quick("tglf",f"mitim_{name}") - - folders, folders_red = [], [] - for subFolderTGLF in tglf_executor: - - rhos = list(tglf_executor[subFolderTGLF].keys()) - - # --------------------------------------------- - # Prepare files and folders - # --------------------------------------------- - - for i, rho in enumerate(rhos): - print(f"\t- Preparing TGLF ({subFolderTGLF}) at rho={rho:.4f}") - - folderTGLF_this = tmpFolder / subFolderTGLF / f"rho_{rho:.4f}" - folders.append(folderTGLF_this) - - folderTGLF_this_rel = folderTGLF_this.relative_to(tmpFolder) - folders_red.append(folderTGLF_this_rel.as_posix() if tglf_job.machineSettings['machine'] != 'local' else str(folderTGLF_this_rel)) - - folderTGLF_this.mkdir(parents=True, exist_ok=True) - - fileTGLF = folderTGLF_this / "input.tglf" - with open(fileTGLF, "w") as f: - f.write(tglf_executor[subFolderTGLF][rho]["inputs"]) - - # --------------------------------------------- - # Prepare command - # --------------------------------------------- - - # Grab machine local limits ------------------------------------------------- - max_cores_per_node = FARMINGtools.mitim_job.grab_machine_settings("tglf")["cores_per_node"] - - # If the run is local and not slurm, let's check the number of cores - if (FARMINGtools.mitim_job.grab_machine_settings("tglf")["machine"] == "local") and not (launchSlurm and ("partition" in tglf_job.machineSettings["slurm"])): - - cores_in_machine = int(os.cpu_count()) - cores_allocated = int(os.environ.get('SLURM_CPUS_PER_TASK')) if os.environ.get('SLURM_CPUS_PER_TASK') is not None else None - - if cores_allocated is not None: - if max_cores_per_node is None or (cores_allocated < max_cores_per_node): - print(f"\t - Detected {cores_allocated} cores allocated by SLURM, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") - max_cores_per_node = cores_allocated - elif cores_in_machine is not None: - if max_cores_per_node is None or (cores_in_machine < max_cores_per_node): - print(f"\t - Detected {cores_in_machine} cores in machine, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") - max_cores_per_node = cores_in_machine - else: - # Default to just 16 just in case - if max_cores_per_node is None: - max_cores_per_node = 16 - else: - # For remote execution, default to just 16 just in case - if max_cores_per_node is None: - max_cores_per_node = 16 - # --------------------------------------------------------------------------- - - # Grab the total number of cores of this job -------------------------------- - total_tglf_executions = len(rhos) * len(tglf_executor) - total_cores_required = int(cores_tglf) * total_tglf_executions - # --------------------------------------------------------------------------- - - # Simply bash, no slurm - if not (launchSlurm and ("partition" in tglf_job.machineSettings["slurm"])): - - max_parallel_execution = max_cores_per_node // cores_tglf # Make sure we don't overload the machine when running locally (assuming no farming trans-node) - - print(f"\t- TGLF will be executed as bash script (total cores: {total_cores_required}, cores per TGLF: {cores_tglf}). MITIM will launch {total_tglf_executions // max_parallel_execution+1} sequential executions",typeMsg="i") - - # Build the bash script with job control enabled and a loop to limit parallel jobs - TGLFcommand = "#!/usr/bin/env bash\n" - TGLFcommand += "set -m\n" # Enable job control even in non-interactive mode - TGLFcommand += f"max_parallel_execution={max_parallel_execution}\n\n" # Set the maximum number of parallel processes - - # Create a bash array of folders - TGLFcommand += "folders=(\n" - for folder in folders_red: - TGLFcommand += f' "{folder}"\n' - TGLFcommand += ")\n\n" - - # Loop over each folder and launch tglf, waiting if we've reached max_parallel_execution - TGLFcommand += "for folder in \"${folders[@]}\"; do\n" - TGLFcommand += f" tglf -e \"$folder\" -n {cores_tglf} -p {tglf_job.folderExecution} &\n" - TGLFcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" - TGLFcommand += "done\n\n" - TGLFcommand += "wait\n" - - # Slurm setup - array_list = None - job_array_limit = None - shellPreCommands = None - shellPostCommands = None - ntasks = total_cores_required - cpuspertask = cores_tglf - - else: - - # Standard job - if total_cores_required < max_cores_per_node: - - print(f"\t- TGLF will be executed in SLURM as standard job (cpus: {total_cores_required})",typeMsg="i") - - # TGLF launches - TGLFcommand = "" - for folder in folders_red: - TGLFcommand += f"tglf -e {folder} -n {cores_tglf} -p {tglf_job.folderExecution} &\n" - TGLFcommand += "\nwait" # This is needed so that the script doesn't end before each job - - # Slurm setup - array_list = None - job_array_limit = None - shellPreCommands = None - shellPostCommands = None - ntasks = total_tglf_executions - cpuspertask = cores_tglf - - # Job array - else: - #raise Exception("TGLF array not implemented yet") - print(f"\t- TGLF will be executed in SLURM as job array due to its size (cpus: {total_cores_required})",typeMsg="i") - - # As a pre-command, organize all folders in a simpler way - shellPreCommands = [] - shellPostCommands = [] - array_list = [] - for i, folder in enumerate(folders_red): - array_list.append(f"{i}") - folder_temp_array = f"run{i}" - folder_actual = folder - shellPreCommands.append(f"mkdir {tglf_job.folderExecution}/{folder_temp_array}; cp {tglf_job.folderExecution}/{folder_actual}/* {tglf_job.folderExecution}/{folder_temp_array}/.") - shellPostCommands.append(f"cp {tglf_job.folderExecution}/{folder_temp_array}/* {tglf_job.folderExecution}/{folder_actual}/.; rm -r {tglf_job.folderExecution}/{folder_temp_array}") - - # TGLF launches - indexed_folder = 'run"$SLURM_ARRAY_TASK_ID"' - TGLFcommand = f'tglf -e {indexed_folder} -n {cores_tglf} -p {tglf_job.folderExecution} 1> {tglf_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {tglf_job.folderExecution}/{indexed_folder}/slurm_error.dat\n' - - # Slurm setup - array_list = ",".join(array_list) - ntasks = 1 - cpuspertask = cores_tglf - job_array_limit = max_jobs_at_once # Limit to this number at most running jobs at the same time - - # --------------------------------------------- - # Execute - # --------------------------------------------- - - tglf_job.define_machine( - "tglf", - f"mitim_{name}", - launchSlurm=launchSlurm, - slurm_settings={ - "minutes": minutes, - "ntasks": ntasks, - "name": name, - "cpuspertask": cpuspertask, - "job_array": array_list, - "job_array_limit": job_array_limit, - #"nodes": 1, - }, - ) - - # I would like the mitim_job to check if the retrieved folders were complete - check_files_in_folder = {} - for folder in folders_red: - check_files_in_folder[folder] = filesToRetrieve - # --------------------------------------------- - - tglf_job.prep( - TGLFcommand, - input_folders=folders, - output_folders=folders_red, - check_files_in_folder=check_files_in_folder, - shellPreCommands=shellPreCommands, - shellPostCommands=shellPostCommands, - ) - - tglf_job.run( - removeScratchFolders=True, - attempts_execution=attempts_execution - ) - - # --------------------------------------------- - # Organize - # --------------------------------------------- - - print("\t- Retrieving files and changing names for storing") - fineall = True - for subFolderTGLF in tglf_executor: - - for i, rho in enumerate(tglf_executor[subFolderTGLF].keys()): - for file in filesToRetrieve: - original_file = f"{file}_{rho:.4f}{extraFlag}" - final_destination = ( - tglf_executor[subFolderTGLF][rho]['folder'] / f"{original_file}" - ) - final_destination.unlink(missing_ok=True) - - temp_file = tmpFolder / subFolderTGLF / f"rho_{rho:.4f}" / f"{file}" - temp_file.replace(final_destination) - - fineall = fineall and final_destination.exists() - - if not final_destination.exists(): - print( - f"\t!! file {file} ({original_file}) could not be retrived", - typeMsg="w", - ) - - if fineall: - print("\t\t- All files were successfully retrieved") - - # Remove temporary folder - shutil.rmtree(tmpFolder) - - else: - print("\t\t- Some files were not retrieved", typeMsg="w") diff --git a/tests/TGLF_workflow.py b/tests/TGLF_workflow.py index 5c912513..34cdd186 100644 --- a/tests/TGLF_workflow.py +++ b/tests/TGLF_workflow.py @@ -19,6 +19,7 @@ subFolderTGLF="run1/", TGLFsettings=None, cold_start=cold_start, + runWaveForms = [0.67, 10.0], forceIfcold_start=True, extraOptions={"USE_BPER": False, "USE_BPAR": False}, slurm_setup={"cores": 4, "minutes": 1}, diff --git a/tests/TGLFscan_workflow.py b/tests/TGLFscan_workflow.py index 5cf22f23..cdac3c9d 100644 --- a/tests/TGLFscan_workflow.py +++ b/tests/TGLFscan_workflow.py @@ -21,7 +21,7 @@ cold_start = cold_start, runWaveForms = [0.67, 10.0], variable = 'RLTS_1', - varUpDown = np.linspace(0.5,1.5,16)) + varUpDown = np.linspace(0.5,1.5,4)) tglf.readScan(label='scan1',variable = 'RLTS_1') tglf.plotScan(labels=['scan1']) @@ -31,6 +31,7 @@ tglf.runScanTurbulenceDrives( subFolderTGLF = 'turb_drives', TGLFsettings = None, + resolutionPoints=3, cold_start = cold_start) tglf.plotScanTurbulenceDrives(label='turb_drives') From 3bd6548d2ac99ace2962178c8528d41f4ac31270 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 11:32:51 -0400 Subject: [PATCH 169/385] Generalization applied to NEO as well --- src/mitim_tools/gacode_tools/NEOtools.py | 220 ++++++++++++------ src/mitim_tools/gacode_tools/TGLFtools.py | 76 ++---- .../gacode_tools/utils/GACODEdefaults.py | 2 +- .../gacode_tools/utils/GACODErun.py | 74 +++++- .../optimizers/multivariate_tools.py | 2 +- tests/NEO_workflow.py | 2 +- tests/TGLFscan_workflow.py | 4 +- 7 files changed, 244 insertions(+), 136 deletions(-) diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index b4c1d958..336ffb84 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -2,26 +2,32 @@ from pathlib import Path from mitim_tools import __version__ as mitim_version from mitim_tools.misc_tools import FARMINGtools, IOtools -from mitim_tools.gacode_tools.utils import GACODErun +from mitim_tools.gacode_tools.utils import GACODErun, GACODEdefaults from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed from mitim_tools.misc_tools.PLASMAtools import md_u -class NEO: +class NEO(GACODErun.gacode_simulation): def __init__( self, rhos=[0.4, 0.6], # rho locations of interest ): - self.rhos = rhos - - def prep(self, inputgacode, folder): - self.inputgacode = inputgacode - self.folder = IOtools.expandPath(folder) - - self.folder.mkdir(parents=True, exist_ok=True) + super().__init__(rhos=rhos) + + self.run_specifications = { + 'code': 'neo', + 'input_file': 'input.neo', + 'code_call': 'neo -e', + 'control_function': GACODEdefaults.addNEOcontrol + } + + print("\n-----------------------------------------------------------------------------------------") + print("\t\t\t NEO class module") + print("-----------------------------------------------------------------------------------------\n") + self.ResultsFiles = self.ResultsFiles_minimal = ['out.neo.transport_flux'] def prep_direct( self, @@ -31,83 +37,148 @@ def prep_direct( forceIfcold_start=False, # Extra flag ): - print("> Preparation of NEO run from input.gacode (direct conversion)") - - self.FolderGACODE = IOtools.expandPath(FolderGACODE) - - if cold_start or not self.FolderGACODE.exists(): - IOtools.askNewFolder(self.FolderGACODE, force=forceIfcold_start) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare state - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - self.profiles = mitim_state - - self.profiles.derive_quantities(mi_ref=md_u) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize from state - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + print("> Preparation of TGLF run from input.gacode (direct conversion)") - self.inputsNEO = self.profiles.to_neo(r=self.rhos, r_is_rho = True) + cdf = self._prep_direct( + mitim_state, + FolderGACODE, + cold_start=cold_start, + forceIfcold_start=forceIfcold_start, + state_converter='to_neo', + input_class=NEOinput, + input_file='input.neo' + ) - for rho in self.inputsNEO: - - # Initialize class - self.inputsNEO[rho] = NEOinput.initialize_in_memory(self.inputsNEO[rho]) - - # Write input.tglf file - self.inputsNEO[rho].file = self.FolderGACODE / f'input.neo_{rho:.4f}' - self.inputsNEO[rho].write_state() + return cdf def run( self, - subfolder, + subFolderNEO, # 'neo1/', + NEOsettings=None, + extraOptions={}, + multipliers={}, + minimum_delta_abs={}, + # runWaveForms=None, # e.g. runWaveForms = [0.3,1.0] + # forceClosestUnstableWF=True, # Look at the growth rate spectrum and run exactly the ky of the closest unstable + ApplyCorrections=True, # Removing ions with too low density and that are fast species + Quasineutral=False, # Ensures quasineutrality. By default is False because I may want to run the file directly + launchSlurm=True, + cold_start=False, forceIfcold_start=False, - ): + extra_name="exe", + slurm_setup={ + "cores": 4, + "minutes": 5, + }, # Cores per NEO call (so, when running nR radii -> nR*4) + attempts_execution=1, + only_minimal_files=False, + ): - # Create this run folder - - subfolder = Path(subfolder) - - FolderNEO = self.FolderGACODE / subfolder - IOtools.askNewFolder(FolderNEO, force=forceIfcold_start) + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare inputs + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + neo_executor, neo_executor_full, folderlast = self.prep_run( + subFolderNEO, + neo_executor={}, + neo_executor_full={}, + NEOsettings=NEOsettings, + extraOptions=extraOptions, + multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, + # runWaveForms=runWaveForms, + # forceClosestUnstableWF=forceClosestUnstableWF, + ApplyCorrections=ApplyCorrections, + Quasineutral=Quasineutral, + launchSlurm=launchSlurm, + cold_start=cold_start, + forceIfcold_start=forceIfcold_start, + extra_name=extra_name, + slurm_setup=slurm_setup, + attempts_execution=attempts_execution, + only_minimal_files=only_minimal_files, + ) - folders, folders_red = [], [] - for rho in self.rhos: - # Create subfolder for each rho - FolderNEO_rho = FolderNEO / f"rho_{rho:.4f}" - IOtools.askNewFolder(FolderNEO_rho, force=forceIfcold_start) - - # Copy the file - os.system(f"cp {self.FolderGACODE / f'input.neo_{rho:.4f}'} {FolderNEO_rho / 'input.neo'}") + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Run NEO + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self._run( + neo_executor, + neo_executor_full=neo_executor_full, + NEOsettings=NEOsettings, + # runWaveForms=runWaveForms, + # forceClosestUnstableWF=forceClosestUnstableWF, + ApplyCorrections=ApplyCorrections, + Quasineutral=Quasineutral, + launchSlurm=launchSlurm, + cold_start=cold_start, + forceIfcold_start=forceIfcold_start, + extra_name=extra_name, + slurm_setup=slurm_setup, + only_minimal_files=only_minimal_files, + ) - folders.append(FolderNEO_rho) - folders_red.append(str(subfolder / f"rho_{rho:.4f}")) + self.FolderNEOlast = folderlast - # Run NEO - - neo_job = FARMINGtools.mitim_job(self.FolderGACODE) - neo_job.define_machine_quick("neo",f"mitim_neo") - - NEOcommand = "" + def prep_run( + self, + subFolder, + neo_executor={}, + neo_executor_full={}, + NEOsettings=None, + **kwargs + ): - for folder in folders_red: - NEOcommand += f"neo -e {folder} -p {neo_job.folderExecution} &\n" - NEOcommand += "wait\n" - - neo_job.define_machine("neo",f"mitim_neo") + return self._prep_run( + subFolder, + code_executor=neo_executor, + code_executor_full=neo_executor_full, + code_settings=NEOsettings, + addControlFunction=self.run_specifications['control_function'], + **kwargs + ) - neo_job.prep( - NEOcommand, - input_folders=[FolderNEO], - output_folders=folders_red, + def _run( + self, + neo_executor, + neo_executor_full={}, + **kwargs_NEOrun + ): + """ + extraOptions and multipliers are not being grabbed from kwargs_NEOrun, but from neo_executor for WF + """ + + print("\n> Run NEO") + + self._generic_run( + neo_executor, + self.run_specifications, + **kwargs_NEOrun ) - neo_job.run( - removeScratchFolders=True, - ) + + + + + + + + + + + + + + + + def prep(self, inputgacode, folder): + self.inputgacode = inputgacode + self.folder = IOtools.expandPath(folder) + + self.folder.mkdir(parents=True, exist_ok=True) + + def run_vgen(self, subfolder="vgen1", vgenOptions={}, cold_start=False): @@ -197,6 +268,11 @@ def process(self, input_dict): #TODO self.all = input_dict + self.num_recorded = 6 + + def anticipate_problems(self): + pass + def write_state(self, file=None): if file is None: diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 1ea3e048..0ffb1499 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -130,6 +130,7 @@ def __init__( 'code_call': 'tglf -e', 'control_function': GACODEdefaults.addTGLFcontrol } + print("\n-----------------------------------------------------------------------------------------") print("\t\t\t TGLF class module") print("-----------------------------------------------------------------------------------------\n") @@ -174,8 +175,6 @@ def __init__( self.LocationCDF = cdf if self.LocationCDF is not None: _, self.nameRunid = IOtools.getLocInfo(self.LocationCDF) - else: - self.nameRunid = "0" self.time, self.avTime = time, avTime ( @@ -321,18 +320,18 @@ def prep( print("\t- Creating dictionary with all input files generated by TGLF_scans") - self.inputsTGLF = {} + self.inputs_files = {} for cont, rho in enumerate(self.rhos): fileN = self.FolderGACODE / f"input.tglf_{rho:.4f}" inputclass = TGLFinput(file=fileN) - self.inputsTGLF[rho] = inputclass + self.inputs_files[rho] = inputclass # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize by taking directly the inputs # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ else: - self.inputsTGLF = specificInputs + self.inputs_files = specificInputs self.tgyro_results = tgyro_results """ @@ -374,41 +373,16 @@ def prep_direct( ): print("> Preparation of TGLF run from input.gacode (direct conversion)") - - self.FolderGACODE = IOtools.expandPath(FolderGACODE) - - if cold_start or not self.FolderGACODE.exists(): - IOtools.askNewFolder(self.FolderGACODE, force=forceIfcold_start) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare state - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - self.profiles = mitim_state - - self.profiles.derive_quantities(mi_ref=md_u) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize from state - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.inputsTGLF = self.profiles.to_tglf(r=self.rhos, r_is_rho = True) - - for rho in self.inputsTGLF: - - # Initialize class - self.inputsTGLF[rho] = TGLFinput.initialize_in_memory(self.inputsTGLF[rho]) - - # Write input.tglf file - self.inputsTGLF[rho].file = self.FolderGACODE / f'input.tglf_{rho:.4f}' - self.inputsTGLF[rho].write_state() - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Definining normalizations - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - print("> Setting up normalizations") - self.NormalizationSets, cdf = NORMtools.normalizations(self.profiles) + cdf = self._prep_direct( + mitim_state, + FolderGACODE, + cold_start=cold_start, + forceIfcold_start=forceIfcold_start, + state_converter='to_tglf', + input_class=TGLFinput, + input_file='input.tglf' + ) return cdf @@ -453,7 +427,7 @@ def prep_from_tglf( self.rhos = [rho] - self.inputsTGLF = {self.rhos[0]: inputclass} + self.inputs_files = {self.rhos[0]: inputclass} def run( self, @@ -484,7 +458,7 @@ def run( # Prepare inputs # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - tglf_executor, tglf_executor_full, folderlast = self._prep( + tglf_executor, tglf_executor_full, folderlast = self.prep_run( subFolderTGLF, tglf_executor={}, tglf_executor_full={}, @@ -527,7 +501,7 @@ def run( self.FolderTGLFlast = folderlast - def _prep( + def prep_run( self, subFolder, tglf_executor={}, @@ -536,7 +510,7 @@ def _prep( **kwargs ): - return self._generic_prep( + return self._prep_run( subFolder, code_executor=tglf_executor, code_executor_full=tglf_executor_full, @@ -652,7 +626,7 @@ def _run_wf(self, kys, tglf_executor, **kwargs_TGLFrun): extraOptions_WF["KY"] = ky_singles extraOptions_WF["VEXB_SHEAR"] = 0.0 # See email from G. Staebler on 05/16/2021 - tglf_executorWF, _, _ = self._prep( + tglf_executorWF, _, _ = self.prep_run( (FolderTGLF_old / f"ky{ky_single0}").relative_to(FolderTGLF_old.parent), tglf_executor=tglf_executorWF, extraOptions=extraOptions_WF, @@ -2167,7 +2141,7 @@ def _prepare_scan( name = f"{variable}_{mult}" - species = self.inputsTGLF[self.rhos[0]] # Any rho will do + species = self.inputs_files[self.rhos[0]] # Any rho will do multipliers_mod = completeVariation(multipliers_mod, species) @@ -2185,7 +2159,7 @@ def _prepare_scan( "forceIfcold_start" in kwargs_TGLFrun and kwargs_TGLFrun["forceIfcold_start"] ) - tglf_executor, tglf_executor_full, folderlast = self._prep( + tglf_executor, tglf_executor_full, folderlast = self.prep_run( f"{self.subFolderTGLF_scan}_{name}", tglf_executor=tglf_executor, tglf_executor_full=tglf_executor_full, @@ -3125,14 +3099,14 @@ def runAnalysis( print(f"*** Running D and V analysis for trace ({fimp:.1e}) species with Z={trace[0]:.1f}, A={trace[1]:.1f}") - self.inputsTGLF_orig = copy.deepcopy(self.inputsTGLF) + self.inputs_files_orig = copy.deepcopy(self.inputs_files) # ------------------------ # Add trace impurity # ------------------------ - for irho in self.inputsTGLF: - position = self.inputsTGLF[irho].addTraceSpecie(Z, A, AS=fimp) + for irho in self.inputs_files: + position = self.inputs_files[irho].addTraceSpecie(Z, A, AS=fimp) self.variable = f"RLNS_{position}" @@ -3177,7 +3151,7 @@ def runAnalysis( self.scans[label]["VoD"].append(V / D) # Back to original (not trace) - self.inputsTGLF = self.inputsTGLF_orig + self.inputs_files = self.inputs_files_orig def plotAnalysis(self, labels=["analysis1"], analysisType="chi_e", figs=None): if figs is None: @@ -4948,7 +4922,7 @@ def plotTGLF_Summary(self, c="b", label="", axs=None, irho_cont=0): try: gammaExB = np.abs( - self.inputsTGLF[self.rhos[irho_cont]].plasma["VEXB_SHEAR"] + self.inputs_files[self.rhos[irho_cont]].plasma["VEXB_SHEAR"] ) if gammaExB > 1e-5: GRAPHICStools.drawLineWithTxt( diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index 8748f3b7..972ac156 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -53,7 +53,7 @@ def addTGLFcontrol(TGLFsettings, NS=2, minimal=False): return TGLFoptions -def addNEOcontrol(): +def addNEOcontrol(*args, **kwargs): NEOoptions = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.neo.controls", caseInsensitive=False) diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 36198bbf..aae4849f 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -4,12 +4,14 @@ import numpy as np import matplotlib.pyplot as plt from scipy.interpolate import interp1d -from mitim_tools.gacode_tools.utils import GACODEdefaults +from mitim_tools.gacode_tools.utils import GACODEdefaults, NORMtools from mitim_tools.transp_tools.utils import NTCCtools from mitim_tools.misc_tools import FARMINGtools, IOtools, MATHtools, GRAPHICStools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +from mitim_tools.misc_tools.PLASMAtools import md_u + class gacode_simulation: def __init__( self, @@ -19,8 +21,62 @@ def __init__( self.ResultsFiles = [] self.ResultsFiles_minimal = [] + + self.nameRunid = "0" + + def _prep_direct( + self, + mitim_state, # A MITIM state class + FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) + cold_start=False, # If True, do not use what it potentially inside the folder, run again + forceIfcold_start=False, # Extra flag + state_converter='to_tglf', + input_class=None, + input_file='input.tglf' + ): + + print("> Preparation run from input.gacode (direct conversion)") + + self.FolderGACODE = IOtools.expandPath(FolderGACODE) + + if cold_start or not self.FolderGACODE.exists(): + IOtools.askNewFolder(self.FolderGACODE, force=forceIfcold_start) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare state + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self.profiles = mitim_state - def _generic_prep( + self.profiles.derive_quantities(mi_ref=md_u) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize from state + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + # Call the method dynamically based on state_converter + conversion_method = getattr(self.profiles, state_converter) + self.inputs_files = conversion_method(r=self.rhos, r_is_rho=True) + + for rho in self.inputs_files: + + # Initialize class + self.inputs_files[rho] = input_class.initialize_in_memory(self.inputs_files[rho]) + + # Write input.tglf file + self.inputs_files[rho].file = self.FolderGACODE / f'{input_file}_{rho:.4f}' + self.inputs_files[rho].write_state() + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Definining normalizations + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + print("> Setting up normalizations") + self.NormalizationSets, cdf = NORMtools.normalizations(self.profiles) + + return cdf + + def _prep_run( self, subfolder_simulation, rhos=None, @@ -39,7 +95,7 @@ def _generic_prep( slurm_setup={ "cores": 4, "minutes": 5, - }, # Cores per TGLF call (so, when running nR radii -> nR*4) + }, # Cores per call (so, when running nR radii -> nR*4) addControlFunction=None, **kwargs ): @@ -56,7 +112,7 @@ def _generic_prep( if rhos is None: rhos = self.rhos - inputs = copy.deepcopy(self.inputsTGLF) + inputs = copy.deepcopy(self.inputs_files) Folder_sim = self.FolderGACODE / subfolder_simulation ResultsFiles_new = [] @@ -83,7 +139,7 @@ def _generic_prep( IOtools.askNewFolder(Folder_sim, force=forceIfcold_start) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Change this specific run of TGLF + # Change this specific run # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ( @@ -196,7 +252,9 @@ def change_and_write_code( NS=inputs[rho].num_recorded, ) - newfile = Folder_sim / f"input.tglf_{rho:.4f}" + input_file = input_sim_rho.file.name.split('_')[0] + + newfile = Folder_sim / f"{input_file}_{rho:.4f}" if code_settings is not None: # Apply corrections @@ -217,8 +275,8 @@ def change_and_write_code( ns_max.append(inputs[rho].num_recorded) - # Convert back to a string because that's how runTGLFproduction operates - inputFile = inputToVariable(Folder_sim, rhos, file='input.tglf') + # Convert back to a string because that's how the run operates + inputFile = inputToVariable(Folder_sim, rhos, file=input_file) if (np.diff(ns_max) > 0).any(): print("> Each radial location has its own number of species... probably because of removal of fast or low density...",typeMsg="w") diff --git a/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py b/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py index bba219ba..6fc768a9 100644 --- a/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py +++ b/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py @@ -58,7 +58,7 @@ def scipy_root(flux_residual_evaluator, x_initial, bounds=None, solver_options=N # -------------------------------------------------------------------------------------------------------- x_history, y_history, metric_history = [], [], [] - def function_for_optimizer_prep(x, dimX=x_initial.shape[-1], flux_residual_evaluator=flux_residual_evaluator, bound_transform=bound_transform): + def function_for_optimizer_prep_run(x, dimX=x_initial.shape[-1], flux_residual_evaluator=flux_residual_evaluator, bound_transform=bound_transform): """ Notes: - x comes extended, batch*dim diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 95a78c72..05de225c 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -13,7 +13,7 @@ os.system(f"rm -r {folder.resolve()}") neo = NEOtools.NEO( - rhos=[0.55] + rhos=[0.5, 0.6, 0.7] ) neo.prep_direct(PROFILEStools.gacode_state(input_gacode), folder, ) diff --git a/tests/TGLFscan_workflow.py b/tests/TGLFscan_workflow.py index cdac3c9d..f07c842e 100644 --- a/tests/TGLFscan_workflow.py +++ b/tests/TGLFscan_workflow.py @@ -1,6 +1,6 @@ import os import numpy as np -from mitim_tools.gacode_tools import TGLFtools +from mitim_tools.gacode_tools import TGLFtools, PROFILEStools from mitim_tools import __mitimroot__ cold_start = True @@ -14,7 +14,7 @@ os.system(f"rm -r {folder.resolve()}") tglf = TGLFtools.TGLF(rhos=[0.5, 0.7]) -tglf.prep(folder, inputgacode=input_gacode, cold_start=cold_start) +tglf.prep_direct(PROFILEStools.gacode_state(input_gacode),folder, cold_start=cold_start) tglf.runScan( subFolderTGLF = 'scan1', TGLFsettings = None, From 95e23af1e563ed1a8a8516ce55dd2555759ed430 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 12:05:49 -0400 Subject: [PATCH 170/385] NEO module fully working --- src/mitim_tools/gacode_tools/NEOtools.py | 140 +++++++++++++++--- src/mitim_tools/gacode_tools/TGLFtools.py | 11 +- .../gacode_tools/utils/GACODErun.py | 16 +- tests/NEO_workflow.py | 14 +- 4 files changed, 145 insertions(+), 36 deletions(-) diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 336ffb84..0c520cd1 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -1,7 +1,8 @@ -import os +import numpy as np from pathlib import Path +import matplotlib.pyplot as plt from mitim_tools import __version__ as mitim_version -from mitim_tools.misc_tools import FARMINGtools, IOtools +from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools from mitim_tools.gacode_tools.utils import GACODErun, GACODEdefaults from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -20,7 +21,8 @@ def __init__( 'code': 'neo', 'input_file': 'input.neo', 'code_call': 'neo -e', - 'control_function': GACODEdefaults.addNEOcontrol + 'control_function': GACODEdefaults.addNEOcontrol, + 'controls_file': 'input.neo.controls' } print("\n-----------------------------------------------------------------------------------------") @@ -67,8 +69,8 @@ def run( forceIfcold_start=False, extra_name="exe", slurm_setup={ - "cores": 4, - "minutes": 5, + "cores": 1, + "minutes": 1, }, # Cores per NEO call (so, when running nR radii -> nR*4) attempts_execution=1, only_minimal_files=False, @@ -136,6 +138,7 @@ def prep_run( code_executor_full=neo_executor_full, code_settings=NEOsettings, addControlFunction=self.run_specifications['control_function'], + controls_file=self.run_specifications['controls_file'], **kwargs ) @@ -157,20 +160,75 @@ def _run( **kwargs_NEOrun ) + def read( + self, + label="neo1", + folder=None, # If None, search in the previously run folder + suffix=None, # If None, search with my standard _0.55 suffixes corresponding to rho of this TGLF class + require_all_files = True, # If False, I only need the fluxes + ): + print("> Reading NEO results") + + # If no specified folder, check the last one + if folder is None: + folder = self.FolderNEOlast + + self.results[label] = {'NEOout':[]} + for rho in self.rhos: + + NEOout = NEOoutput( + folder, + suffix=f"_{rho:.4f}" if suffix is None else suffix, + ) + self.results[label]['NEOout'].append(NEOout) - - - - - - - - - - - - + + def plot( + self, + fn=None, + labels=["neo1"], + extratitle="", + fn_color=None, + colors=None, + ): + + if fn is None: + self.fn = GUItools.FigureNotebook("NEO MITIM Notebook", geometry="1700x900", vertical=True) + else: + self.fn = fn + + fig1 = self.fn.add_figure(label=f"{extratitle}Summary", tab_color=fn_color) + + grid = plt.GridSpec(1, 3, hspace=0.7, wspace=0.2) + + if colors is None: + colors = GRAPHICStools.listColors() + + axQe = fig1.add_subplot(grid[0, 0]) + axQi = fig1.add_subplot(grid[0, 1]) + axGe = fig1.add_subplot(grid[0, 2]) + + for i,label in enumerate(labels): + roa, QeGB, QiGB, GeGB = [], [], [], [] + for irho in range(len(self.rhos)): + roa.append(self.results[label]['NEOout'][irho].roa) + QeGB.append(self.results[label]['NEOout'][irho].QeGB) + QiGB.append(self.results[label]['NEOout'][irho].QiGB) + GeGB.append(self.results[label]['NEOout'][irho].GeGB) + + axQe.plot(roa, QeGB, label=label, color=colors[i], marker='o', linestyle='-') + axQi.plot(roa, QiGB, label=label, color=colors[i], marker='o', linestyle='-') + axGe.plot(roa, GeGB, label=label, color=colors[i], marker='o', linestyle='-') + + for ax in [axQe, axQi, axGe]: + ax.set_xlabel("$r/a$"); ax.set_xlim([0,1]) + GRAPHICStools.addDenseAxis(ax) + ax.legend(loc="best") + + axQe.set_ylabel("$Q_e$ (GB)"); axQe.set_yscale('log') + axQi.set_ylabel("$Q_i$ (GB)"); axQi.set_yscale('log') + axGe.set_ylabel("$G_e$ (GB)"); #axGe.set_yscale('log') def prep(self, inputgacode, folder): self.inputgacode = inputgacode @@ -243,7 +301,6 @@ def check_if_files_exist(folder, list_files): return True - class NEOinput: def __init__(self, file=None): self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None @@ -266,7 +323,7 @@ def initialize_in_memory(cls, input_dict): def process(self, input_dict): #TODO - self.all = input_dict + self.controls = input_dict self.num_recorded = 6 @@ -307,6 +364,45 @@ def _fmt_value(val): f.write(f"# NEO input file modified by MITIM {mitim_version}\n") f.write("#-------------------------------------------------------------------------\n") - for ikey in self.all: - var = self.all[ikey] - f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") \ No newline at end of file + for ikey in self.controls: + var = self.controls[ikey] + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") + + +class NEOoutput: + def __init__(self, FolderGACODE, suffix=""): + self.FolderGACODE, self.suffix = FolderGACODE, suffix + + if suffix == "": + print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} without suffix") + else: + print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} with suffix {suffix}") + + self.inputclass = NEOinput(file=self.FolderGACODE / f"input.neo{self.suffix}") + + self.read() + + def read(self): + + with open(self.FolderGACODE / ("out.neo.transport_flux" + self.suffix), "r") as f: + lines = f.readlines() + + for i in range(len(lines)): + if '# Z pflux_tgyro eflux_tgyro mflux_tgyro' in lines[i]: + # Found the header line, now process the data + break + + line = lines[i+2] + self.GeGB, self.QeGB, self.MeGB = [float(x) for x in line.split()[1:]] + + self.GiAllGB, self.QiAllGB, self.MiAllGB = [], [], [] + for i in range(i+3, len(lines)): + line = lines[i] + self.GiAllGB.append(float(line.split()[1])) + self.QiAllGB.append(float(line.split()[2])) + self.MiAllGB.append(float(line.split()[3])) + + self.QiGB = np.sum(self.QiAllGB) + + self.roa = float(lines[0].split()[-1]) + diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 0ffb1499..2a936efe 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -128,7 +128,8 @@ def __init__( 'code': 'tglf', 'input_file': 'input.tglf', 'code_call': 'tglf -e', - 'control_function': GACODEdefaults.addTGLFcontrol + 'control_function': GACODEdefaults.addTGLFcontrol, + 'controls_file': 'input.tglf.controls' } print("\n-----------------------------------------------------------------------------------------") @@ -177,12 +178,7 @@ def __init__( _, self.nameRunid = IOtools.getLocInfo(self.LocationCDF) self.time, self.avTime = time, avTime - ( - self.results, - self.scans, - self.tgyro, - self.ky_single, - ) = ({}, {}, None, None) + self.tgyro,self.ky_single = None, None self.NormalizationSets = { "TRANSP": None, @@ -516,6 +512,7 @@ def prep_run( code_executor_full=tglf_executor_full, code_settings=TGLFsettings, addControlFunction=self.run_specifications['control_function'], + controls_file=self.run_specifications['controls_file'], **kwargs ) diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index aae4849f..0e522b2c 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -23,6 +23,8 @@ def __init__( self.ResultsFiles_minimal = [] self.nameRunid = "0" + + self.results, self.scans = {}, {} def _prep_direct( self, @@ -97,6 +99,7 @@ def _prep_run( "minutes": 5, }, # Cores per call (so, when running nR radii -> nR*4) addControlFunction=None, + controls_file='input.tglf.controls', **kwargs ): @@ -155,7 +158,8 @@ def _prep_run( minimum_delta_abs=minimum_delta_abs, ApplyCorrections=ApplyCorrections, Quasineutral=Quasineutral, - addControlFunction=addControlFunction + addControlFunction=addControlFunction, + controls_file=controls_file, ) code_executor_full[subfolder_simulation] = {} @@ -229,6 +233,7 @@ def change_and_write_code( ApplyCorrections=True, Quasineutral=False, addControlFunction=None, + controls_file='input.tglf.controls', ): """ Received inputs classes and gives text. @@ -249,6 +254,7 @@ def change_and_write_code( minimum_delta_abs=minimum_delta_abs, position_change=i, addControlFunction=addControlFunction, + controls_file=controls_file, NS=inputs[rho].num_recorded, ) @@ -542,7 +548,7 @@ def run_gacode_simulation( ) gacode_job.run( - removeScratchFolders=False, + removeScratchFolders=True, attempts_execution=attempts_execution ) @@ -675,13 +681,13 @@ def modifyInputs( minimum_delta_abs={}, position_change=0, addControlFunction=None, - control_file = 'input.tglf.controls', + controls_file = 'input.tglf.controls', **kwargs_to_function, ): # Check that those are valid flags - GACODEdefaults.review_controls(extraOptions, control = control_file) - GACODEdefaults.review_controls(multipliers, control = control_file) + GACODEdefaults.review_controls(extraOptions, control = controls_file) + GACODEdefaults.review_controls(multipliers, control = controls_file) # ------------------------------------------- if Settings is not None: diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 05de225c..8f8788a5 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -1,4 +1,5 @@ import os +import numpy as np from mitim_tools.gacode_tools import NEOtools, PROFILEStools from mitim_tools import __mitimroot__ @@ -13,8 +14,17 @@ os.system(f"rm -r {folder.resolve()}") neo = NEOtools.NEO( - rhos=[0.5, 0.6, 0.7] + rhos=np.linspace(0.1,0.95,20) ) neo.prep_direct(PROFILEStools.gacode_state(input_gacode), folder, ) -neo.run('neo1') +neo.run('neo1/') +neo.read('neo1') + +neo.run('neo2/', extraOptions={'N_XI': 17, 'N_THETA': 17}) +neo.read('neo2') + +neo.plot(labels=['neo1', 'neo2']) + +neo.fn.show() +neo.fn.close() \ No newline at end of file From 91f13864439d0de6145a76ae1e8e6c5941dcb676 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 12:51:04 -0400 Subject: [PATCH 171/385] PORTALS working fully with standalone TGLF + NEO --- ...transport_tglf.py => transport_tglfneo.py} | 68 +++++++++++++----- src/mitim_modules/vitals/VITALSmain.py | 2 +- src/mitim_tools/gacode_tools/NEOtools.py | 70 +++++++++++++++---- .../gacode_tools/utils/GACODErun.py | 8 ++- .../optimizers/multivariate_tools.py | 2 +- tests/NEO_workflow.py | 2 +- tests/TGLFscan_workflow.py | 2 +- 7 files changed, 119 insertions(+), 35 deletions(-) rename src/mitim_modules/powertorch/physics_models/{transport_tglf.py => transport_tglfneo.py} (88%) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py similarity index 88% rename from src/mitim_modules/powertorch/physics_models/transport_tglf.py rename to src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 87ab6f6d..7687f9fc 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -3,12 +3,12 @@ import torch import numpy as np from mitim_tools.misc_tools import IOtools, PLASMAtools -from mitim_tools.gacode_tools import TGLFtools +from mitim_tools.gacode_tools import TGLFtools, NEOtools from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -class tglf_model(TRANSPORTtools.power_transport): +class tglfneo_model(TRANSPORTtools.power_transport): def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) @@ -57,8 +57,8 @@ def _evaluate_tglf(self): tglf = TGLFtools.TGLF(rhos=rho_locations) _ = tglf.prep_direct( - self.folder, self.powerstate.profiles_transport, + self.folder, cold_start = cold_start, ) @@ -71,7 +71,7 @@ def _evaluate_tglf(self): # Just run TGLF once and apply an ad-hoc percent error to the results tglf.run( - 'base', + 'base_tglf', TGLFsettings=TGLFsettings, extraOptions=extraOptions, ApplyCorrections=False, @@ -129,7 +129,7 @@ def _evaluate_tglf(self): for i in range(len(tglf.profiles.Species)): gacode_type = tglf.profiles.Species[i]['S'] for rho in rho_locations: - tglf_type = tglf.inputsTGLF[0.25].ions_info[i+2]['type'] + tglf_type = tglf.inputs_files[0.25].ions_info[i+2]['type'] if gacode_type[:5] != tglf_type[:5]: print(f"\t- For location {rho=:.2f}, ion specie #{i+1} ({tglf.profiles.Species[i]['N']}) is considered '{gacode_type}' by gacode but '{tglf_type}' by TGLF. Make sure this is consistent with your use case", typeMsg="w") @@ -175,22 +175,54 @@ def _evaluate_tglf(self): def _evaluate_neo(self): - self.powerstate.plasma["QeMWm2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] - self.powerstate.plasma["QiMWm2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] - self.powerstate.plasma["Ge1E20m2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] - self.powerstate.plasma["GZ1E20m2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] - self.powerstate.plasma["MtJm2_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + # Options - self.powerstate.plasma["QeMWm2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] - self.powerstate.plasma["QiMWm2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] - self.powerstate.plasma["Ge1E20m2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] - self.powerstate.plasma["GZ1E20m2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] - self.powerstate.plasma["MtJm2_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] - self.powerstate.plasma["QieMWm3_tr_neoc"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] - self.powerstate.plasma["QieMWm3_tr_neoc_stds"] = 0.0 * self.powerstate.plasma["QeMWm2_tr_turb"] + cold_start = transport_evaluator_options.get("cold_start", False) + percentError = transport_evaluator_options.get("percentError", [5, 1, 0.5]) + impurityPosition = self.powerstate.impurityPosition_transport + + # Run + + rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] - return None + neo = NEOtools.NEO(rhos=rho_locations) + + _ = neo.prep_direct( + self.powerstate.profiles_transport, + self.folder, + cold_start = cold_start, + ) + + neo.run( + 'base_neo', + ) + + neo.read(label='base') + + Qe = np.array([neo.results['base']['NEOout'][i].Qe_unn for i in range(len(rho_locations))]) + Qi = np.array([neo.results['base']['NEOout'][i].Qi_unn for i in range(len(rho_locations))]) + Ge = np.array([neo.results['base']['NEOout'][i].Ge_unn for i in range(len(rho_locations))]) + GZ = np.array([neo.results['base']['NEOout'][i].GiAll_unn[impurityPosition] for i in range(len(rho_locations))]) + Mt = np.array([neo.results['base']['NEOout'][i].Mt_unn for i in range(len(rho_locations))]) + + self.powerstate.plasma["QeMWm2_tr_neoc"] = Qe + self.powerstate.plasma["QiMWm2_tr_neoc"] = Qi + self.powerstate.plasma["Ge1E20m2_tr_neoc"] = Ge + self.powerstate.plasma["GZ1E20m2_tr_neoc"] = GZ + self.powerstate.plasma["MtJm2_tr_neoc"] = Mt + + self.powerstate.plasma["QeMWm2_tr_neoc_stds"] = Qe * percentError[1]/100.0 + self.powerstate.plasma["QiMWm2_tr_neoc_stds"] = Qi * percentError[1]/100.0 + self.powerstate.plasma["Ge1E20m2_tr_neoc_stds"] = Ge * percentError[1]/100.0 + self.powerstate.plasma["GZ1E20m2_tr_neoc_stds"] = GZ * percentError[1]/100.0 + self.powerstate.plasma["MtJm2_tr_neoc_stds"] = Mt * percentError[1]/100.0 + + self.powerstate.plasma["QieMWm3_tr_neoc"] = Qe * 0.0 + self.powerstate.plasma["QieMWm3_tr_neoc_stds"] = Qe * 0.0 + + return neo def _postprocess(self): diff --git a/src/mitim_modules/vitals/VITALSmain.py b/src/mitim_modules/vitals/VITALSmain.py index fedb3d97..998547ec 100644 --- a/src/mitim_modules/vitals/VITALSmain.py +++ b/src/mitim_modules/vitals/VITALSmain.py @@ -287,7 +287,7 @@ def runTGLF( numSim = self.folder.name - variation = TGLFtools.completeVariation(variation, tglf.inputsTGLF[tglf.rhos[0]]) + variation = TGLFtools.completeVariation(variation, tglf.inputs_files[tglf.rhos[0]]) extraOptions = self.TGLFparameters["extraOptions"] multipliers = {} diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 0c520cd1..8452c794 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -180,6 +180,12 @@ def read( folder, suffix=f"_{rho:.4f}" if suffix is None else suffix, ) + + # Unnormalize + NEOout.unnormalize( + self.NormalizationSets["SELECTED"], + rho=rho, + ) self.results[label]['NEOout'].append(NEOout) @@ -213,9 +219,9 @@ def plot( roa, QeGB, QiGB, GeGB = [], [], [], [] for irho in range(len(self.rhos)): roa.append(self.results[label]['NEOout'][irho].roa) - QeGB.append(self.results[label]['NEOout'][irho].QeGB) - QiGB.append(self.results[label]['NEOout'][irho].QiGB) - GeGB.append(self.results[label]['NEOout'][irho].GeGB) + QeGB.append(self.results[label]['NEOout'][irho].Qe) + QiGB.append(self.results[label]['NEOout'][irho].Qi) + GeGB.append(self.results[label]['NEOout'][irho].Ge) axQe.plot(roa, QeGB, label=label, color=colors[i], marker='o', linestyle='-') axQi.plot(roa, QiGB, label=label, color=colors[i], marker='o', linestyle='-') @@ -226,9 +232,9 @@ def plot( GRAPHICStools.addDenseAxis(ax) ax.legend(loc="best") - axQe.set_ylabel("$Q_e$ (GB)"); axQe.set_yscale('log') - axQi.set_ylabel("$Q_i$ (GB)"); axQi.set_yscale('log') - axGe.set_ylabel("$G_e$ (GB)"); #axGe.set_yscale('log') + axQe.set_ylabel("$Q_e$ ($MW/m^2$)"); axQe.set_yscale('log') + axQi.set_ylabel("$Q_i$ ($MW/m^2$)"); axQi.set_yscale('log') + axGe.set_ylabel("$G_e$ ($1E20/s/m^2$)"); #axGe.set_yscale('log') def prep(self, inputgacode, folder): self.inputgacode = inputgacode @@ -393,16 +399,56 @@ def read(self): break line = lines[i+2] - self.GeGB, self.QeGB, self.MeGB = [float(x) for x in line.split()[1:]] + self.Ge, self.Qe, self.Me = [float(x) for x in line.split()[1:]] - self.GiAllGB, self.QiAllGB, self.MiAllGB = [], [], [] + self.GiAll, self.QiAll, self.MiAll = [], [], [] for i in range(i+3, len(lines)): line = lines[i] - self.GiAllGB.append(float(line.split()[1])) - self.QiAllGB.append(float(line.split()[2])) - self.MiAllGB.append(float(line.split()[3])) + self.GiAll.append(float(line.split()[1])) + self.QiAll.append(float(line.split()[2])) + self.MiAll.append(float(line.split()[3])) + + self.GiAll = np.array(self.GiAll) + self.QiAll = np.array(self.QiAll) + self.MiAll = np.array(self.MiAll) - self.QiGB = np.sum(self.QiAllGB) + self.Qi = self.QiAll.sum() + self.Mt = self.Me + self.MiAll.sum() self.roa = float(lines[0].split()[-1]) + def unnormalize(self, normalization, rho=None): + + if normalization is not None: + rho_x = normalization["rho"] + roa_x = normalization["roa"] + q_gb = normalization["q_gb"] + g_gb = normalization["g_gb"] + pi_gb = normalization["pi_gb"] + s_gb = normalization["s_gb"] + rho_s = normalization["rho_s"] + a = normalization["rmin"][-1] + + # ------------------------------------ + # Usage of normalization quantities + # ------------------------------------ + + if rho is None: + ir = np.argmin(np.abs(roa_x - self.roa)) + rho_eval = rho_x[ir] + else: + ir = np.argmin(np.abs(rho_x - rho)) + rho_eval = rho + + self.Qe_unn = self.Qe * q_gb[ir] + self.Qi_unn = self.Qi * q_gb[ir] + self.QiAll_unn = self.QiAll * q_gb[ir] + self.Ge_unn = self.Ge * g_gb[ir] + self.GiAll_unn = self.GiAll * g_gb[ir] + self.MiAll_unn = self.MiAll * g_gb[ir] + self.Mt_unn = self.Mt * s_gb[ir] + + self.unnormalization_successful = True + + else: + self.unnormalization_successful = False diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 0e522b2c..69cfcea8 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -2,8 +2,10 @@ import os import copy import numpy as np +from pathlib import Path import matplotlib.pyplot as plt from scipy.interpolate import interp1d +from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.gacode_tools.utils import GACODEdefaults, NORMtools from mitim_tools.transp_tools.utils import NTCCtools from mitim_tools.misc_tools import FARMINGtools, IOtools, MATHtools, GRAPHICStools @@ -48,7 +50,11 @@ def _prep_direct( # Prepare state # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.profiles = mitim_state + if isinstance(mitim_state, str) or isinstance(mitim_state, Path): + # If a string, assume it's a path to input.gacode + mitim_state = PROFILEStools.gacode_state(mitim_state) + else: + self.profiles = mitim_state self.profiles.derive_quantities(mi_ref=md_u) diff --git a/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py b/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py index 6fc768a9..bba219ba 100644 --- a/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py +++ b/src/mitim_tools/opt_tools/optimizers/multivariate_tools.py @@ -58,7 +58,7 @@ def scipy_root(flux_residual_evaluator, x_initial, bounds=None, solver_options=N # -------------------------------------------------------------------------------------------------------- x_history, y_history, metric_history = [], [], [] - def function_for_optimizer_prep_run(x, dimX=x_initial.shape[-1], flux_residual_evaluator=flux_residual_evaluator, bound_transform=bound_transform): + def function_for_optimizer_prep(x, dimX=x_initial.shape[-1], flux_residual_evaluator=flux_residual_evaluator, bound_transform=bound_transform): """ Notes: - x comes extended, batch*dim diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 8f8788a5..fe73956e 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -16,7 +16,7 @@ neo = NEOtools.NEO( rhos=np.linspace(0.1,0.95,20) ) -neo.prep_direct(PROFILEStools.gacode_state(input_gacode), folder, ) +neo.prep_direct(input_gacode, folder) neo.run('neo1/') neo.read('neo1') diff --git a/tests/TGLFscan_workflow.py b/tests/TGLFscan_workflow.py index f07c842e..5f5a7a70 100644 --- a/tests/TGLFscan_workflow.py +++ b/tests/TGLFscan_workflow.py @@ -14,7 +14,7 @@ os.system(f"rm -r {folder.resolve()}") tglf = TGLFtools.TGLF(rhos=[0.5, 0.7]) -tglf.prep_direct(PROFILEStools.gacode_state(input_gacode),folder, cold_start=cold_start) +tglf.prep_direct(input_gacode,folder, cold_start=cold_start) tglf.runScan( subFolderTGLF = 'scan1', TGLFsettings = None, From 0e9c3b020ecfb7db830bae4a85736930dd2d483c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 12:53:47 -0400 Subject: [PATCH 172/385] Solved circular import --- .../powertorch/physics_models/transport_tglfneo.py | 10 +++++----- src/mitim_tools/plasmastate_tools/MITIMstate.py | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 7687f9fc..d188d55c 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -213,11 +213,11 @@ def _evaluate_neo(self): self.powerstate.plasma["GZ1E20m2_tr_neoc"] = GZ self.powerstate.plasma["MtJm2_tr_neoc"] = Mt - self.powerstate.plasma["QeMWm2_tr_neoc_stds"] = Qe * percentError[1]/100.0 - self.powerstate.plasma["QiMWm2_tr_neoc_stds"] = Qi * percentError[1]/100.0 - self.powerstate.plasma["Ge1E20m2_tr_neoc_stds"] = Ge * percentError[1]/100.0 - self.powerstate.plasma["GZ1E20m2_tr_neoc_stds"] = GZ * percentError[1]/100.0 - self.powerstate.plasma["MtJm2_tr_neoc_stds"] = Mt * percentError[1]/100.0 + self.powerstate.plasma["QeMWm2_tr_neoc_stds"] = abs(Qe) * percentError[1]/100.0 + self.powerstate.plasma["QiMWm2_tr_neoc_stds"] = abs(Qi) * percentError[1]/100.0 + self.powerstate.plasma["Ge1E20m2_tr_neoc_stds"] = abs(Ge) * percentError[1]/100.0 + self.powerstate.plasma["GZ1E20m2_tr_neoc_stds"] = abs(GZ) * percentError[1]/100.0 + self.powerstate.plasma["MtJm2_tr_neoc_stds"] = abs(Mt) * percentError[1]/100.0 self.powerstate.plasma["QieMWm3_tr_neoc"] = Qe * 0.0 self.powerstate.plasma["QieMWm3_tr_neoc_stds"] = Qe * 0.0 diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 8d5b8aea..9e040599 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -5,7 +5,6 @@ import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, MATHtools, PLASMAtools, IOtools from mitim_modules.powertorch.utils import CALCtools -from mitim_tools.gacode_tools import NEOtools from mitim_tools.gacode_tools.utils import GACODEdefaults from mitim_tools.plasmastate_tools.utils import state_plotting from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -244,6 +243,7 @@ def calculate_Er( profiles.changeResolution(rho_new=rhos) resol_changed = True + from mitim_tools.gacode_tools import NEOtools self.neo = NEOtools.NEO() self.neo.prep(profiles, folder) self.neo.run_vgen(subfolder=name, vgenOptions=vgenOptions, cold_start=cold_start) From 2a193b67babe5fa399ac0832d0e3bfa7f81e61e9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 15:47:14 -0400 Subject: [PATCH 173/385] Rotation in to_neo working --- .../plasmastate_tools/MITIMstate.py | 36 ++++++++++--------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 9e040599..48e6f05b 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2416,6 +2416,9 @@ def to_neo(self, r=[0.5], r_is_rho = True): s_delta = self.derived["r"] * self._deriv_gacode(self.profiles["delta(-)"]) s_zeta = self.derived["r"] * self._deriv_gacode(self.profiles["zeta(-)"]) + omega_rot = self.profiles["w0(rad/s)"] / self.derived['c_s'] * self.derived['a'] + omega_rot_deriv = self._deriv_gacode(self.profiles["w0(rad/s)"])/ self.derived['c_s'] * self.derived['a']**2 + inputsNEO = {} for rho in r: @@ -2436,10 +2439,20 @@ def interpolator(y): # Species come from profiles # --------------------------------------------------------------------------------------------------------------------------------------- - max_species_neo = 6 #TODO: True? + species = {} + for i in range(len(self.Species)): + species[i+1] = { + 'Z': self.Species[i]['Z'], + 'MASS': self.Species[i]['A']/mass_ref, + 'DLNNDR': interpolator(self.derived['aLni'][:,i]), + 'DLNTDR': interpolator(self.derived['aLTi'][:,0] if self.Species[i]['S'] == 'therm' else self.derived["aLTi"][:,i]), + 'TEMP': interpolator(self.derived["tite_all"][:,i]), + 'DENS': interpolator(self.derived['fi'][:,i]), + } - species = { - 1: { + ie = i+2 + + species[ie] = { 'Z': -1.0, 'MASS': 0.000272445, 'DLNNDR': interpolator(self.derived['aLne']), @@ -2447,31 +2460,20 @@ def interpolator(y): 'TEMP': 1.0, 'DENS': 1.0, } - } - - for i in range(min(len(self.Species), max_species_neo-1)): - species[i+2] = { - 'Z': self.Species[i]['Z'], - 'MASS': self.Species[i]['A']/mass_ref, - 'DLNNDR': interpolator(self.derived['aLni'][:,i]), - 'DLNTDR': interpolator(self.derived['aLTi'][:,0] if self.Species[i]['S'] == 'therm' else self.derived["aLTi"][:,i]), - 'TEMP': interpolator(self.derived["tite_all"][:,i]), - 'DENS': interpolator(self.derived['fi'][:,i]), - } # --------------------------------------------------------------------------------------------------------------------------------------- # Plasma comes from profiles # --------------------------------------------------------------------------------------------------------------------------------------- #TODO Does this work with no deuterium first ion? - factor_nu = species[2]['Z']**4 * species[2]['DENS'] * species[1]['MASS']**0.5 * species[2]['TEMP']**(-1.5) + factor_nu = species[1]['Z']**4 * species[1]['DENS'] * (species[ie]['MASS']/species[1]['MASS'])**0.5 * species[1]['TEMP']**(-1.5) plasma = { 'N_SPECIES': len(species), 'IPCCW': sign_bt, 'BTCCW': sign_it, - 'OMEGA_ROT': 0.0, #TODO - 'OMEGA_ROT_DERIV': 0.0, #TODO + 'OMEGA_ROT': interpolator(omega_rot), + 'OMEGA_ROT_DERIV': interpolator(omega_rot_deriv), 'NU_1': interpolator(self.derived['xnue'])* factor_nu, 'RHO_STAR': interpolator(self.derived["rho_sa"]), } From 9957f6598cd59e3e3c5b18e7151d45c1fd5f8712 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 15:50:27 -0400 Subject: [PATCH 174/385] NEO module and portals ready to ship --- src/mitim_tools/gacode_tools/utils/GACODErun.py | 2 +- tests/NEO_workflow.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 69cfcea8..1ceaa18a 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -52,7 +52,7 @@ def _prep_direct( if isinstance(mitim_state, str) or isinstance(mitim_state, Path): # If a string, assume it's a path to input.gacode - mitim_state = PROFILEStools.gacode_state(mitim_state) + self.profiles = PROFILEStools.gacode_state(mitim_state) else: self.profiles = mitim_state diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index fe73956e..63a935ad 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -21,10 +21,10 @@ neo.run('neo1/') neo.read('neo1') -neo.run('neo2/', extraOptions={'N_XI': 17, 'N_THETA': 17}) -neo.read('neo2') +neo.run('neo2/', extraOptions={'N_ENERGY':6,'N_XI': 17, 'N_THETA': 17}) +neo.read('neo_highres') -neo.plot(labels=['neo1', 'neo2']) +neo.plot(labels=['neo1', 'neo_highres']) neo.fn.show() neo.fn.close() \ No newline at end of file From 103a227dcb571352c59238dd27c87da6517edeb1 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 15:56:59 -0400 Subject: [PATCH 175/385] misc --- src/mitim_tools/gacode_tools/NEOtools.py | 3 +-- src/mitim_tools/gacode_tools/utils/GACODErun.py | 13 +++++++------ tests/NEO_workflow.py | 12 ++++++++---- 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 8452c794..35609441 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -330,8 +330,7 @@ def initialize_in_memory(cls, input_dict): def process(self, input_dict): #TODO self.controls = input_dict - - self.num_recorded = 6 + self.num_recorded = 100 def anticipate_problems(self): pass diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 1ceaa18a..bf494fb3 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -296,6 +296,7 @@ def change_and_write_code( return inputFile, mod_input_file + def inputToVariable(finalFolder, rhos, file='input.tglf'): """ Entire text file to variable @@ -761,26 +762,26 @@ def modifyInputs( print("\t\t- Variables change:") for ikey in multipliers: # is a specie one? - if ikey.split("_")[0] in input_class.species[1]: + if "species" in input_class.__dict__.keys() and key.split("_")[0] in input_class.species[1]: specie = int(ikey.split("_")[-1]) varK = "_".join(ikey.split("_")[:-1]) var_orig = input_class.species[specie][varK] - var_new = multiplier_tglf_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.species[specie][varK] = var_new else: if ikey in input_class.controls: var_orig = input_class.controls[ikey] - var_new = multiplier_tglf_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.controls[ikey] = var_new elif ikey in input_class.geom: var_orig = input_class.geom[ikey] - var_new = multiplier_tglf_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.geom[ikey] = var_new elif ikey in input_class.plasma: var_orig = input_class.plasma[ikey] - var_new = multiplier_tglf_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.plasma[ikey] = var_new else: @@ -790,7 +791,7 @@ def modifyInputs( return input_class -def multiplier_tglf_input(var_orig, multiplier, minimum_delta_abs = None): +def multiplier_input(var_orig, multiplier, minimum_delta_abs = None): delta = var_orig * (multiplier - 1.0) diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 63a935ad..80796731 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -1,7 +1,8 @@ import os import numpy as np -from mitim_tools.gacode_tools import NEOtools, PROFILEStools +from mitim_tools.gacode_tools import NEOtools from mitim_tools import __mitimroot__ +from torch import mul cold_start = True @@ -19,12 +20,15 @@ neo.prep_direct(input_gacode, folder) neo.run('neo1/') -neo.read('neo1') +neo.read('NEO') neo.run('neo2/', extraOptions={'N_ENERGY':6,'N_XI': 17, 'N_THETA': 17}) -neo.read('neo_highres') +neo.read('NEO high res') -neo.plot(labels=['neo1', 'neo_highres']) +neo.run('neo3/', extraOptions={'N_ENERGY':6,'N_XI': 17, 'N_THETA': 17}, multipliers={'DLNTDR_1': 1.25}) +neo.read('NEO high res + 25% aLTe') + +neo.plot(labels=['NEO', 'NEO high res', 'NEO high res + 25% aLTe']) neo.fn.show() neo.fn.close() \ No newline at end of file From 795fea38d4acc5c62baf23be89d324da382bd192 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 21:59:10 -0400 Subject: [PATCH 176/385] NEO defaults --- templates/input.neo.controls | 28 +++++++++++++++++----------- tests/NEO_workflow.py | 12 ++++++------ 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/templates/input.neo.controls b/templates/input.neo.controls index 09673778..3dd4b5be 100644 --- a/templates/input.neo.controls +++ b/templates/input.neo.controls @@ -2,17 +2,23 @@ # Template input.neo file (controls-only) #------------------------------------------------------------------------- -# Resolution -N_ENERGY=5 -N_XI=11 -N_THETA=11 -N_RADIAL=1 +SILENT_FLAG=0 # 0: output files are written -# Geometry (Miller) -EQUILIBRIUM_MODEL=2 +# Setup +EQUILIBRIUM_MODEL=2 # 2: Miller +PROFILE_MODEL=1 # 1: local (one radius) +PROFILE_ERAD0_MODEL=1 # 1: Use the profile of the Erad0 parameter as specified in input.profiles -# Rotation (Sonic) -ROTATION_MODEL=2 +# Resolution +N_RADIAL=1 # Number of radial grid points +N_THETA=17 # Number of poloidal grid points +N_ENERGY=6 # Number of v polynomials +N_XI=17 # Number of xi polynomials -# Collisions (FP) -COLLISION_MODEL=4 +# Models +COLLISION_MODEL=4 # 4: Full linearized Fokker-Plank operator +SIM_MODEL=2 # 2: Numerical solution and analytic theory only +SPITZER_MODEL=0 # 0: solve the standard neoclassical transport problem +ROTATION_MODEL=2 # 2: Sonic rotation effects included (solves the Hinton-Wong generalized +THREED_MODEL=0 # 0: Toroidally axisymmetric limit (2D). +THREED_EXB_MODEL=0 # 0: higher-order drift velocity not included. \ No newline at end of file diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 80796731..56d229f6 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -20,15 +20,15 @@ neo.prep_direct(input_gacode, folder) neo.run('neo1/') -neo.read('NEO') +neo.read('NEO default') -neo.run('neo2/', extraOptions={'N_ENERGY':6,'N_XI': 17, 'N_THETA': 17}) -neo.read('NEO high res') +neo.run('neo2/', extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}) +neo.read('NEO low res') -neo.run('neo3/', extraOptions={'N_ENERGY':6,'N_XI': 17, 'N_THETA': 17}, multipliers={'DLNTDR_1': 1.25}) -neo.read('NEO high res + 25% aLTe') +neo.run('neo3/', extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}, multipliers={'DLNTDR_1': 1.25}) +neo.read('NEO low res + 25% aLTe') -neo.plot(labels=['NEO', 'NEO high res', 'NEO high res + 25% aLTe']) +neo.plot(labels=['NEO default', 'NEO low res', 'NEO low res + 25% aLTe']) neo.fn.show() neo.fn.close() \ No newline at end of file From 4d69a3eaaad75b2702a88c1c1d5df539c28d6b92 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 20 Aug 2025 22:30:15 -0400 Subject: [PATCH 177/385] misc --- src/mitim_tools/gacode_tools/NEOtools.py | 24 ++++++++++--------- src/mitim_tools/gacode_tools/TGLFtools.py | 20 +++++++++------- .../gacode_tools/utils/GACODErun.py | 9 +++++-- tests/NEO_workflow.py | 12 ++++------ 4 files changed, 35 insertions(+), 30 deletions(-) diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 35609441..3972115d 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -22,7 +22,9 @@ def __init__( 'input_file': 'input.neo', 'code_call': 'neo -e', 'control_function': GACODEdefaults.addNEOcontrol, - 'controls_file': 'input.neo.controls' + 'controls_file': 'input.neo.controls', + 'state_converter': 'to_neo', + 'input_class': NEOinput, } print("\n-----------------------------------------------------------------------------------------") @@ -33,22 +35,22 @@ def __init__( def prep_direct( self, - mitim_state, # A MITIM state class - FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) - cold_start=False, # If True, do not use what it potentially inside the folder, run again - forceIfcold_start=False, # Extra flag + mitim_state, # A MITIM state class + FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) + cold_start=False, # If True, do not use what it potentially inside the folder, run again + forceIfcold_start=False, # Extra flag ): - print("> Preparation of TGLF run from input.gacode (direct conversion)") + print("> Preparation of NEO run from input.gacode (direct conversion)") cdf = self._prep_direct( mitim_state, FolderGACODE, cold_start=cold_start, forceIfcold_start=forceIfcold_start, - state_converter='to_neo', - input_class=NEOinput, - input_file='input.neo' + state_converter=self.run_specifications['state_converter'], + input_class=self.run_specifications['input_class'], + input_file=self.run_specifications['input_file'] ) return cdf @@ -132,7 +134,7 @@ def prep_run( **kwargs ): - return self._prep_run( + return self._generic_run_prep( subFolder, code_executor=neo_executor, code_executor_full=neo_executor_full, @@ -234,7 +236,7 @@ def plot( axQe.set_ylabel("$Q_e$ ($MW/m^2$)"); axQe.set_yscale('log') axQi.set_ylabel("$Q_i$ ($MW/m^2$)"); axQi.set_yscale('log') - axGe.set_ylabel("$G_e$ ($1E20/s/m^2$)"); #axGe.set_yscale('log') + axGe.set_ylabel("$\\Gamma_e$ ($1E20/s/m^2$)"); #axGe.set_yscale('log') def prep(self, inputgacode, folder): self.inputgacode = inputgacode diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 2a936efe..aa1f99d3 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -129,7 +129,9 @@ def __init__( 'input_file': 'input.tglf', 'code_call': 'tglf -e', 'control_function': GACODEdefaults.addTGLFcontrol, - 'controls_file': 'input.tglf.controls' + 'controls_file': 'input.tglf.controls', + 'state_converter': 'to_tglf', + 'input_class': TGLFinput, } print("\n-----------------------------------------------------------------------------------------") @@ -362,10 +364,10 @@ def prep( def prep_direct( self, - mitim_state, # A MITIM state class - FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) - cold_start=False, # If True, do not use what it potentially inside the folder, run again - forceIfcold_start=False, # Extra flag + mitim_state, # A MITIM state class + FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) + cold_start=False, # If True, do not use what it potentially inside the folder, run again + forceIfcold_start=False, # Extra flag ): print("> Preparation of TGLF run from input.gacode (direct conversion)") @@ -375,9 +377,9 @@ def prep_direct( FolderGACODE, cold_start=cold_start, forceIfcold_start=forceIfcold_start, - state_converter='to_tglf', - input_class=TGLFinput, - input_file='input.tglf' + state_converter=self.run_specifications['state_converter'], + input_class=self.run_specifications['input_class'], + input_file=self.run_specifications['input_file'], ) return cdf @@ -506,7 +508,7 @@ def prep_run( **kwargs ): - return self._prep_run( + return self._generic_run_prep( subFolder, code_executor=tglf_executor, code_executor_full=tglf_executor_full, diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index bf494fb3..89993db6 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -14,7 +14,12 @@ from mitim_tools.misc_tools.PLASMAtools import md_u + + class gacode_simulation: + ''' + Main class for running GACODE simulations. + ''' def __init__( self, rhos=[0.4, 0.6], # rho locations of interest @@ -84,7 +89,7 @@ def _prep_direct( return cdf - def _prep_run( + def _generic_run_prep( self, subfolder_simulation, rhos=None, @@ -762,7 +767,7 @@ def modifyInputs( print("\t\t- Variables change:") for ikey in multipliers: # is a specie one? - if "species" in input_class.__dict__.keys() and key.split("_")[0] in input_class.species[1]: + if "species" in input_class.__dict__.keys() and ikey.split("_")[0] in input_class.species[1]: specie = int(ikey.split("_")[-1]) varK = "_".join(ikey.split("_")[:-1]) var_orig = input_class.species[specie][varK] diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 56d229f6..807d40a5 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -2,7 +2,6 @@ import numpy as np from mitim_tools.gacode_tools import NEOtools from mitim_tools import __mitimroot__ -from torch import mul cold_start = True @@ -14,9 +13,7 @@ if cold_start and folder.exists(): os.system(f"rm -r {folder.resolve()}") -neo = NEOtools.NEO( - rhos=np.linspace(0.1,0.95,20) -) +neo = NEOtools.NEO(rhos=np.linspace(0.1,0.95,20)) neo.prep_direct(input_gacode, folder) neo.run('neo1/') @@ -25,10 +22,9 @@ neo.run('neo2/', extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}) neo.read('NEO low res') -neo.run('neo3/', extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}, multipliers={'DLNTDR_1': 1.25}) -neo.read('NEO low res + 25% aLTe') - -neo.plot(labels=['NEO default', 'NEO low res', 'NEO low res + 25% aLTe']) +neo.run('neo3/', extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}, multipliers={'DLNTDR_1': 1.5}) +neo.read('NEO low res + 50% aLTe') +neo.plot(labels=['NEO default', 'NEO low res', 'NEO low res + 50% aLTe']) neo.fn.show() neo.fn.close() \ No newline at end of file From 780be5dda556da9188e5d8ee170eea1591c21b91 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 21 Aug 2025 15:45:33 -0400 Subject: [PATCH 178/385] changing prep naming and made run_scan general --- docs/capabilities/tglf_capabilities.rst | 4 +- .../portals/utils/PORTALSanalysis.py | 2 +- .../physics_models/transport_tglfneo.py | 4 +- .../physics_models/transport_tgyro.py | 2 +- src/mitim_tools/gacode_tools/CGYROtools.py | 4 +- src/mitim_tools/gacode_tools/NEOtools.py | 143 +---- src/mitim_tools/gacode_tools/TGLFtools.py | 570 ++++++------------ src/mitim_tools/gacode_tools/TGYROtools.py | 4 +- .../gacode_tools/scripts/read_tglf.py | 2 +- .../gacode_tools/scripts/run_tglf.py | 2 +- .../gacode_tools/utils/GACODEdefaults.py | 8 +- .../gacode_tools/utils/GACODErun.py | 330 ++++++++-- tests/NEO_workflow.py | 12 +- tests/TGLF_workflow.py | 10 +- tests/TGLFfull_workflow.py | 2 +- tests/TGLFscan_workflow.py | 4 +- tests/VITALS_workflow.py | 2 +- tests/data/FolderTRANSP/12345X01TR.DAT | 2 +- tutorials/TGLF_tutorial.py | 2 +- .../run_slurm_array_tutorial/test_launcher.py | 2 +- 20 files changed, 506 insertions(+), 605 deletions(-) diff --git a/docs/capabilities/tglf_capabilities.rst b/docs/capabilities/tglf_capabilities.rst index 852b1534..cfc34614 100644 --- a/docs/capabilities/tglf_capabilities.rst +++ b/docs/capabilities/tglf_capabilities.rst @@ -168,7 +168,7 @@ If you have a input.tglf file already, you can still use this script to run it. inputtglf_file = Path('MITIM-fusion/tests/data/input.tglf') tglf = TGLFtools.TGLF() - tglf.prep_from_tglf( folder, inputtglf_file, input_gacode = inputgacode_file ) + tglf.prep_from_file( folder, inputtglf_file, input_gacode = inputgacode_file ) The rest of the workflow is identical, including ``.run()``, ``.read()`` and ``.plot()``. @@ -190,7 +190,7 @@ The rest of the workflow is identical, including ``.run()``, ``.read()`` and ``. inputtglf_file = Path('MITIM-fusion/tests/data/input.tglf') tglf = TGLFtools.TGLF() - tglf.prep_from_tglf( folder, inputtglf_file ) + tglf.prep_from_file( folder, inputtglf_file ) tglf.read (folder = f'{folder}/', label = 'yes_em' ) tglf.plot( labels = ['yes_em'] ) diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 765b825e..64eb1efd 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -594,7 +594,7 @@ def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=F p.write_state(file=inputgacode) tglf = TGLFtools.TGLF(rhos=rhos) - _ = tglf.prep(folder, cold_start=cold_start, inputgacode=inputgacode) + _ = tglf.prep_using_tgyro(folder, cold_start=cold_start, inputgacode=inputgacode) TGLFsettings = self.MODELparameters["transport_model"]["TGLFsettings"] extraOptions = self.MODELparameters["transport_model"]["extraOptionsTGLF"] diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index d188d55c..7dabee02 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -56,7 +56,7 @@ def _evaluate_tglf(self): tglf = TGLFtools.TGLF(rhos=rho_locations) - _ = tglf.prep_direct( + _ = tglf.prep( self.powerstate.profiles_transport, self.folder, cold_start = cold_start, @@ -189,7 +189,7 @@ def _evaluate_neo(self): neo = NEOtools.NEO(rhos=rho_locations) - _ = neo.prep_direct( + _ = neo.prep( self.powerstate.profiles_transport, self.folder, cold_start = cold_start, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 0768e0ca..655be8b9 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -114,7 +114,7 @@ def _evaluate_tglf_neo(self): # from mitim_tools.gacode_tools import TGLFtools # tglf = TGLFtools.TGLF(rhos=rho_locations) - # _ = tglf.prep( + # _ = tglf.prep_using_tgyro( # self.folder / 'stds', # inputgacode=self.file_profs, # recalculate_ptot=False, # Use what's in the input.gacode, which is what PORTALS TGYRO does diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index bd262b8b..95877f4a 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -100,7 +100,7 @@ def _prerun( inputCGYRO = GACODErun.modifyInputs( inputCGYRO, - Settings=CGYROsettings, + code_settings=CGYROsettings, extraOptions=extraOptions, multipliers=multipliers, addControlFunction=GACODEdefaults.addCGYROcontrol, @@ -310,7 +310,7 @@ def run_full( inputCGYRO = CGYROinput(file=input_cgyro_file_this) input_cgyro_file_this = GACODErun.modifyInputs( inputCGYRO, - Settings=CGYROsettings, + code_settings=CGYROsettings, extraOptions=extraOptions_this, multipliers=multipliers, addControlFunction=GACODEdefaults.addCGYROcontrol, diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 3972115d..b6664ee3 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -25,6 +25,8 @@ def __init__( 'controls_file': 'input.neo.controls', 'state_converter': 'to_neo', 'input_class': NEOinput, + 'complete_variation': None, + 'default_cores': 1, # Default cores to use in the simulation } print("\n-----------------------------------------------------------------------------------------") @@ -33,147 +35,18 @@ def __init__( self.ResultsFiles = self.ResultsFiles_minimal = ['out.neo.transport_flux'] - def prep_direct( - self, - mitim_state, # A MITIM state class - FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) - cold_start=False, # If True, do not use what it potentially inside the folder, run again - forceIfcold_start=False, # Extra flag - ): - - print("> Preparation of NEO run from input.gacode (direct conversion)") - - cdf = self._prep_direct( - mitim_state, - FolderGACODE, - cold_start=cold_start, - forceIfcold_start=forceIfcold_start, - state_converter=self.run_specifications['state_converter'], - input_class=self.run_specifications['input_class'], - input_file=self.run_specifications['input_file'] - ) - - return cdf - - def run( - self, - subFolderNEO, # 'neo1/', - NEOsettings=None, - extraOptions={}, - multipliers={}, - minimum_delta_abs={}, - # runWaveForms=None, # e.g. runWaveForms = [0.3,1.0] - # forceClosestUnstableWF=True, # Look at the growth rate spectrum and run exactly the ky of the closest unstable - ApplyCorrections=True, # Removing ions with too low density and that are fast species - Quasineutral=False, # Ensures quasineutrality. By default is False because I may want to run the file directly - launchSlurm=True, - cold_start=False, - forceIfcold_start=False, - extra_name="exe", - slurm_setup={ - "cores": 1, - "minutes": 1, - }, # Cores per NEO call (so, when running nR radii -> nR*4) - attempts_execution=1, - only_minimal_files=False, - ): - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare inputs - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - neo_executor, neo_executor_full, folderlast = self.prep_run( - subFolderNEO, - neo_executor={}, - neo_executor_full={}, - NEOsettings=NEOsettings, - extraOptions=extraOptions, - multipliers=multipliers, - minimum_delta_abs=minimum_delta_abs, - # runWaveForms=runWaveForms, - # forceClosestUnstableWF=forceClosestUnstableWF, - ApplyCorrections=ApplyCorrections, - Quasineutral=Quasineutral, - launchSlurm=launchSlurm, - cold_start=cold_start, - forceIfcold_start=forceIfcold_start, - extra_name=extra_name, - slurm_setup=slurm_setup, - attempts_execution=attempts_execution, - only_minimal_files=only_minimal_files, - ) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Run NEO - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - self._run( - neo_executor, - neo_executor_full=neo_executor_full, - NEOsettings=NEOsettings, - # runWaveForms=runWaveForms, - # forceClosestUnstableWF=forceClosestUnstableWF, - ApplyCorrections=ApplyCorrections, - Quasineutral=Quasineutral, - launchSlurm=launchSlurm, - cold_start=cold_start, - forceIfcold_start=forceIfcold_start, - extra_name=extra_name, - slurm_setup=slurm_setup, - only_minimal_files=only_minimal_files, - ) - - self.FolderNEOlast = folderlast - - def prep_run( - self, - subFolder, - neo_executor={}, - neo_executor_full={}, - NEOsettings=None, - **kwargs - ): - - return self._generic_run_prep( - subFolder, - code_executor=neo_executor, - code_executor_full=neo_executor_full, - code_settings=NEOsettings, - addControlFunction=self.run_specifications['control_function'], - controls_file=self.run_specifications['controls_file'], - **kwargs - ) - - def _run( - self, - neo_executor, - neo_executor_full={}, - **kwargs_NEOrun - ): - """ - extraOptions and multipliers are not being grabbed from kwargs_NEOrun, but from neo_executor for WF - """ - - print("\n> Run NEO") - - self._generic_run( - neo_executor, - self.run_specifications, - **kwargs_NEOrun - ) - def read( self, label="neo1", folder=None, # If None, search in the previously run folder suffix=None, # If None, search with my standard _0.55 suffixes corresponding to rho of this TGLF class - require_all_files = True, # If False, I only need the fluxes + **kwargs ): print("> Reading NEO results") # If no specified folder, check the last one if folder is None: - folder = self.FolderNEOlast + folder = self.FolderSimLast self.results[label] = {'NEOout':[]} for rho in self.rhos: @@ -238,11 +111,11 @@ def plot( axQi.set_ylabel("$Q_i$ ($MW/m^2$)"); axQi.set_yscale('log') axGe.set_ylabel("$\\Gamma_e$ ($1E20/s/m^2$)"); #axGe.set_yscale('log') - def prep(self, inputgacode, folder): - self.inputgacode = inputgacode - self.folder = IOtools.expandPath(folder) + # def prep(self, inputgacode, folder): + # self.inputgacode = inputgacode + # self.folder = IOtools.expandPath(folder) - self.folder.mkdir(parents=True, exist_ok=True) + # self.folder.mkdir(parents=True, exist_ok=True) diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index aa1f99d3..10ed3462 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -71,7 +71,7 @@ def __init__( tglf = TGLF(cdf='~/testTGLF/12345B12.CDF',time=1.45,avTime=0.1,rhos=[0.4,0.6]) # Prepare TGLF (this will create input.tglf in the specified folder) - cdf = tglf.prep('~/testTGLF/') + cdf = tglf.prep_using_tgyro('~/testTGLF/') # Run standalone TGLF (this will find the input.tglf in the previous folder, # and then copy to this specify TGLF run, and run it there) @@ -91,7 +91,7 @@ def __init__( tglf = TGLF(cdf='~/testTGLF/12345B12.CDF',time=1.45,avTime=0.1,rhos=[0.4,0.6]) # Prepare TGLF (this will create input.tglf in the specified folder) - cdf = tglf.prep('~/testTGLF/') + cdf = tglf.prep_using_tgyro('~/testTGLF/') # Run tglf.runScan('scan1/',TGLFsettings=1,varUpDown=np.linspace(0.5,2.0,20),variable='RLTS_2') @@ -132,6 +132,8 @@ def __init__( 'controls_file': 'input.tglf.controls', 'state_converter': 'to_tglf', 'input_class': TGLFinput, + 'complete_variation': completeVariation_TGLF, + 'default_cores': 4, # Default cores to use in the simulation } print("\n-----------------------------------------------------------------------------------------") @@ -149,9 +151,8 @@ def __init__( "out.tglf.gbflux", ] - self.ResultsFiles = [ + self.ResultsFiles = self.ResultsFiles_minimal + [ "out.tglf.run", - "out.tglf.gbflux", "out.tglf.eigenvalue_spectrum", "out.tglf.sum_flux_spectrum", "out.tglf.ky_spectrum", @@ -191,6 +192,143 @@ def __init__( "SELECTED": None, } + def run( + self, + subFolderTGLF, + runWaveForms=None, # e.g. runWaveForms = [0.3,1.0] + forceClosestUnstableWF=True, # Look at the growth rate spectrum and run exactly the ky of the closest unstable + **kwargs_generic_run + ): + ''' + I need to redefine the run method for the TGLF class because it has the option of producing WaveForms + ''' + + code_executor_full = super().run(subFolderTGLF, **kwargs_generic_run) + + kwargs_generic_run['runWaveForms'] = runWaveForms + kwargs_generic_run['forceClosestUnstableWF'] = forceClosestUnstableWF + self._helper_wf(code_executor_full, **kwargs_generic_run) + + def run_scan( + self, + subFolderTGLF, + **kwargs, + ): + + code_executor_full = super().run_scan(subFolderTGLF,**kwargs) + + self._helper_wf(code_executor_full, **kwargs) + + # TOREMOVE #TODO + def runScan( + self, + subFolderTGLF, + **kwargs, + ): + + self.run_scan(subFolderTGLF, **kwargs) + + def _run_wf(self, kys, code_executor, forceClosestUnstableWF=True, **kwargs_TGLFrun): + """ + extraOptions and multipliers are not being grabbed from kwargs_TGLFrun, but from code_executor + """ + + if kwargs_TGLFrun.get("only_minimal_files", False): + raise Exception('[MITIM] Option to run WF with only minimal files is not available yet') + + if "runWaveForms" in kwargs_TGLFrun: + del kwargs_TGLFrun["runWaveForms"] + + # Grab these from code_executor + if "extraOptions" in kwargs_TGLFrun: + del kwargs_TGLFrun["extraOptions"] + if "multipliers" in kwargs_TGLFrun: + del kwargs_TGLFrun["multipliers"] + + self.ky_single = kys + ResultsFiles = copy.deepcopy(self.ResultsFiles) + self.ResultsFiles = copy.deepcopy(self.ResultsFiles_WF) + + self.FoldersTGLF_WF = {} + if self.ky_single is not None: + + code_executorWF = {} + for ky_single0 in self.ky_single: + print(f"> Running TGLF waveform analysis, ky~{ky_single0}") + + self.FoldersTGLF_WF[f"ky{ky_single0}"] = {} + for subFolderTGLF in code_executor: + + ky_single_orig = copy.deepcopy(ky_single0) + + FolderTGLF_old = code_executor[subFolderTGLF][list(code_executor[subFolderTGLF].keys())[0]]["folder"] + + self.ky_single = None + self.read(label=f"ky{ky_single0}", folder=FolderTGLF_old, cold_startWF = False) + self.ky_single = kys + + self.FoldersTGLF_WF[f"ky{ky_single0}"][ + FolderTGLF_old + ] = FolderTGLF_old / f"ky{ky_single0}" + + ky_singles = [] + for i, ir in enumerate(self.rhos): + # -------- Get the closest unstable mode to the one requested + if forceClosestUnstableWF: + + # Only unstable ones + kys_n = [] + for j in range(len(self.results[f"ky{ky_single0}"]["TGLFout"][i].ky)): + if self.results[f"ky{ky_single0}"]["TGLFout"][i].g[0, j] > 0.0: + kys_n.append(self.results[f"ky{ky_single0}"]["TGLFout"][i].ky[j]) + kys_n = np.array(kys_n) + # ---- + + closest_ky = kys_n[np.argmin(np.abs(kys_n - ky_single_orig))] + print(f"\t- rho = {ir:.3f}, requested ky={ky_single_orig:.3f}, & closest unstable ky based on previous run: ky={closest_ky:.3f}",typeMsg="i",) + ky_single = closest_ky + else: + ky_single = ky_single0 + + ky_singles.append(ky_single) + # ------------------------------------------------------------ + + kwargs_TGLFrun0 = copy.deepcopy(kwargs_TGLFrun) + if "extraOptions" in kwargs_TGLFrun: + extraOptions_WF = copy.deepcopy(kwargs_TGLFrun["extraOptions"]) + del kwargs_TGLFrun0["extraOptions"] + + else: + extraOptions_WF = {} + + extraOptions_WF = copy.deepcopy(code_executor[subFolderTGLF][list(code_executor[subFolderTGLF].keys())[0]]["extraOptions"]) + multipliers_WF = copy.deepcopy(code_executor[subFolderTGLF][list(code_executor[subFolderTGLF].keys())[0]]["multipliers"]) + + 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 + + code_executorWF, _ = self._run_prepare( + (FolderTGLF_old / f"ky{ky_single0}").relative_to(FolderTGLF_old.parent), + code_executor=code_executorWF, + extraOptions=extraOptions_WF, + multipliers=multipliers_WF, + **kwargs_TGLFrun0, + ) + + # Run them all + self._run( + code_executorWF, + runWaveForms=[], + **kwargs_TGLFrun0, + ) + + # Recover previous stuff + self.ResultsFiles_WF = copy.deepcopy(self.ResultsFiles) + self.ResultsFiles = ResultsFiles + # ----------- + def prepare_for_save_TGLF(self): """ This is a function that will be called when saving the class as pickle. @@ -219,7 +357,7 @@ def save_pkl(self, file): with open(file, "wb") as handle: pickle.dump(tglf_copy, handle, protocol=4) - def prep( + def prep_using_tgyro( self, FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) cold_start=False, # If True, do not use what it potentially inside the folder, run again @@ -290,10 +428,7 @@ def prep( print(f"\t\t- Running scans because it does not exist file {IOtools.clipstr(fii)}") exists = False if exists: - print( - "\t\t- All input files to TGLF exist, not running scans", - typeMsg="i", - ) + print("\t\t- All input files to TGLF exist, not running scans",typeMsg="i",) """ Sometimes, if I'm running TGLF only from input.tglf file, I may not need to run the entire TGYRO workflow @@ -362,29 +497,7 @@ def prep( return cdf - def prep_direct( - self, - mitim_state, # A MITIM state class - FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) - cold_start=False, # If True, do not use what it potentially inside the folder, run again - forceIfcold_start=False, # Extra flag - ): - - print("> Preparation of TGLF run from input.gacode (direct conversion)") - - cdf = self._prep_direct( - mitim_state, - FolderGACODE, - cold_start=cold_start, - forceIfcold_start=forceIfcold_start, - state_converter=self.run_specifications['state_converter'], - input_class=self.run_specifications['input_class'], - input_file=self.run_specifications['input_file'], - ) - - return cdf - - def prep_from_tglf( + def prep_from_file( self, FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) input_tglf_file, # input.tglf file to start with @@ -427,223 +540,6 @@ def prep_from_tglf( self.inputs_files = {self.rhos[0]: inputclass} - def run( - self, - subFolderTGLF, # 'tglf1/', - TGLFsettings=None, - extraOptions={}, - multipliers={}, - minimum_delta_abs={}, - runWaveForms=None, # e.g. runWaveForms = [0.3,1.0] - forceClosestUnstableWF=True, # Look at the growth rate spectrum and run exactly the ky of the closest unstable - ApplyCorrections=True, # Removing ions with too low density and that are fast species - Quasineutral=False, # Ensures quasineutrality. By default is False because I may want to run the file directly - launchSlurm=True, - cold_start=False, - forceIfcold_start=False, - extra_name="exe", - slurm_setup={ - "cores": 4, - "minutes": 5, - }, # Cores per TGLF call (so, when running nR radii -> nR*4) - attempts_execution=1, - only_minimal_files=False, - ): - - if runWaveForms is None: runWaveForms = [] - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare inputs - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - tglf_executor, tglf_executor_full, folderlast = self.prep_run( - subFolderTGLF, - tglf_executor={}, - tglf_executor_full={}, - TGLFsettings=TGLFsettings, - extraOptions=extraOptions, - multipliers=multipliers, - minimum_delta_abs=minimum_delta_abs, - runWaveForms=runWaveForms, - forceClosestUnstableWF=forceClosestUnstableWF, - ApplyCorrections=ApplyCorrections, - Quasineutral=Quasineutral, - launchSlurm=launchSlurm, - cold_start=cold_start, - forceIfcold_start=forceIfcold_start, - extra_name=extra_name, - slurm_setup=slurm_setup, - attempts_execution=attempts_execution, - only_minimal_files=only_minimal_files, - ) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Run TGLF - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - self._run( - tglf_executor, - tglf_executor_full=tglf_executor_full, - TGLFsettings=TGLFsettings, - runWaveForms=runWaveForms, - forceClosestUnstableWF=forceClosestUnstableWF, - ApplyCorrections=ApplyCorrections, - Quasineutral=Quasineutral, - launchSlurm=launchSlurm, - cold_start=cold_start, - forceIfcold_start=forceIfcold_start, - extra_name=extra_name, - slurm_setup=slurm_setup, - only_minimal_files=only_minimal_files, - ) - - self.FolderTGLFlast = folderlast - - def prep_run( - self, - subFolder, - tglf_executor={}, - tglf_executor_full={}, - TGLFsettings=None, - **kwargs - ): - - return self._generic_run_prep( - subFolder, - code_executor=tglf_executor, - code_executor_full=tglf_executor_full, - code_settings=TGLFsettings, - addControlFunction=self.run_specifications['control_function'], - controls_file=self.run_specifications['controls_file'], - **kwargs - ) - - def _run( - self, - tglf_executor, - tglf_executor_full={}, - **kwargs_TGLFrun - ): - """ - extraOptions and multipliers are not being grabbed from kwargs_TGLFrun, but from tglf_executor for WF - """ - - print("\n> Run TGLF") - - self._generic_run( - tglf_executor, - self.run_specifications, - **kwargs_TGLFrun - ) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Waveform if requested - # Cannot be in parallel to the previous run, because it needs the results of unstable ky - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if "runWaveForms" in kwargs_TGLFrun and kwargs_TGLFrun["runWaveForms"] is not None and len(kwargs_TGLFrun["runWaveForms"]) > 0: - self._run_wf(kwargs_TGLFrun["runWaveForms"], tglf_executor_full, **kwargs_TGLFrun) - - def _run_wf(self, kys, tglf_executor, **kwargs_TGLFrun): - """ - extraOptions and multipliers are not being grabbed from kwargs_TGLFrun, but from tglf_executor - """ - - if kwargs_TGLFrun.get("only_minimal_files", False): - raise Exception('[MITIM] Option to run WF with only minimal files is not available yet') - - if "runWaveForms" in kwargs_TGLFrun: - del kwargs_TGLFrun["runWaveForms"] - - # Grab these from tglf_executor - if "extraOptions" in kwargs_TGLFrun: - del kwargs_TGLFrun["extraOptions"] - if "multipliers" in kwargs_TGLFrun: - del kwargs_TGLFrun["multipliers"] - - self.ky_single = kys - ResultsFiles = copy.deepcopy(self.ResultsFiles) - self.ResultsFiles = copy.deepcopy(self.ResultsFiles_WF) - - self.FoldersTGLF_WF = {} - if self.ky_single is not None: - - tglf_executorWF = {} - for ky_single0 in self.ky_single: - print(f"> Running TGLF waveform analysis, ky~{ky_single0}") - - self.FoldersTGLF_WF[f"ky{ky_single0}"] = {} - for subFolderTGLF in tglf_executor: - - ky_single_orig = copy.deepcopy(ky_single0) - - FolderTGLF_old = tglf_executor[subFolderTGLF][list(tglf_executor[subFolderTGLF].keys())[0]]["folder"] - - self.ky_single = None - self.read(label=f"ky{ky_single0}", folder=FolderTGLF_old, cold_startWF = False) - self.ky_single = kys - - self.FoldersTGLF_WF[f"ky{ky_single0}"][ - FolderTGLF_old - ] = FolderTGLF_old / f"ky{ky_single0}" - - ky_singles = [] - for i, ir in enumerate(self.rhos): - # -------- Get the closest unstable mode to the one requested - if kwargs_TGLFrun.get("forceClosestUnstableWF", True): - - # Only unstable ones - kys_n = [] - for j in range(len(self.results[f"ky{ky_single0}"]["TGLFout"][i].ky)): - if self.results[f"ky{ky_single0}"]["TGLFout"][i].g[0, j] > 0.0: - kys_n.append(self.results[f"ky{ky_single0}"]["TGLFout"][i].ky[j]) - kys_n = np.array(kys_n) - # ---- - - closest_ky = kys_n[np.argmin(np.abs(kys_n - ky_single_orig))] - print(f"\t- rho = {ir:.3f}, requested ky={ky_single_orig:.3f}, & closest unstable ky based on previous run: ky={closest_ky:.3f}",typeMsg="i",) - ky_single = closest_ky - else: - ky_single = ky_single0 - - ky_singles.append(ky_single) - # ------------------------------------------------------------ - - kwargs_TGLFrun0 = copy.deepcopy(kwargs_TGLFrun) - if "extraOptions" in kwargs_TGLFrun: - extraOptions_WF = copy.deepcopy(kwargs_TGLFrun["extraOptions"]) - del kwargs_TGLFrun0["extraOptions"] - - else: - extraOptions_WF = {} - - extraOptions_WF = copy.deepcopy(tglf_executor[subFolderTGLF][list(tglf_executor[subFolderTGLF].keys())[0]]["extraOptions"]) - multipliers_WF = copy.deepcopy(tglf_executor[subFolderTGLF][list(tglf_executor[subFolderTGLF].keys())[0]]["multipliers"]) - - 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_executorWF, _, _ = self.prep_run( - (FolderTGLF_old / f"ky{ky_single0}").relative_to(FolderTGLF_old.parent), - tglf_executor=tglf_executorWF, - extraOptions=extraOptions_WF, - multipliers=multipliers_WF, - **kwargs_TGLFrun0, - ) - - # Run them all - self._run( - tglf_executorWF, - runWaveForms=[], - **kwargs_TGLFrun0, - ) - - # Recover previous stuff - self.ResultsFiles_WF = copy.deepcopy(self.ResultsFiles) - self.ResultsFiles = ResultsFiles - # ----------- def read( self, @@ -683,7 +579,7 @@ def read( # If no specified folder, check the last one if folder is None: - folder = self.FolderTGLFlast + folder = self.FolderSimLast # ----------------------------------------- @@ -2037,149 +1933,48 @@ def plot( legYN=contLab == 0, ) - # ~~~~~~~~~~~~~~ Scan options - - def runScan( + def _helper_wf( self, - subFolderTGLF, # 'scan1', - multipliers={}, - minimum_delta_abs={}, - variable="RLTS_1", - varUpDown=[0.5, 1.0, 1.5], - variables_scanTogether=[], - relativeChanges=True, - **kwargs_TGLFrun, + code_executor_full, + **kwargs ): + + runWaveForms = kwargs['runWaveForms'] if 'runWaveForms' in kwargs else None + forceClosestUnstableWF = kwargs['forceClosestUnstableWF'] if 'forceClosestUnstableWF' in kwargs else True - # ------------------------------------- - # Add baseline - # ------------------------------------- - if (1.0 not in varUpDown) and relativeChanges: - print("\n* Since variations vector did not include base case, I am adding it",typeMsg="i",) - varUpDown_new = [] - added = False - for i in varUpDown: - if i > 1.0 and not added: - varUpDown_new.append(1.0) - added = True - varUpDown_new.append(i) - else: - varUpDown_new = varUpDown - + if 'runWaveForms' in kwargs: + del kwargs['runWaveForms'] + if 'forceClosestUnstableWF' in kwargs: + del kwargs['forceClosestUnstableWF'] - tglf_executor, tglf_executor_full, folders, varUpDown_new = self._prepare_scan( - subFolderTGLF, - multipliers=multipliers, - minimum_delta_abs=minimum_delta_abs, - variable=variable, - varUpDown=varUpDown_new, - variables_scanTogether=variables_scanTogether, - relativeChanges=relativeChanges, - **kwargs_TGLFrun, - ) + # ********************************************************************************************* + # Waveform if requested (not in parallel because it needs the results of unstable ky) + # ********************************************************************************************* - # Run them all - self._run( - tglf_executor, - tglf_executor_full=tglf_executor_full, - **kwargs_TGLFrun, - ) + if runWaveForms is not None and len(runWaveForms) > 0: - # Read results - for cont_mult, mult in enumerate(varUpDown_new): - name = f"{variable}_{mult}" - self.read( - label=f"{self.subFolderTGLF_scan}_{name}", - folder=folders[cont_mult], - cold_startWF = False, - require_all_files=not kwargs_TGLFrun.get("only_minimal_files",False), - ) - - def _prepare_scan( + # Keep the same folder as before + self.keep_folder = copy.deepcopy(self.FolderSimLast) + + # Run WF + self._run_wf(runWaveForms, code_executor_full, forceClosestUnstableWF=forceClosestUnstableWF, **kwargs) + + # Get back to it + self.FolderSimLast = self.keep_folder + + def readScan( self, - subFolderTGLF, # 'scan1', - multipliers={}, - minimum_delta_abs={}, + label="scan1", + subFolderTGLF=None, variable="RLTS_1", - varUpDown=[0.5, 1.0, 1.5], - variables_scanTogether=[], - relativeChanges=True, - **kwargs_TGLFrun, - ): - """ - Multipliers will be modified by adding the scaning variables, but I don't want to modify the original - multipliers, as they may be passed to the next scan - - Set relativeChanges=False if varUpDown contains the exact values to change, not multipleiers - """ - multipliers_mod = copy.deepcopy(multipliers) - - self.subFolderTGLF_scan = subFolderTGLF - - if relativeChanges: - for i in range(len(varUpDown)): - varUpDown[i] = round(varUpDown[i], 6) - - print(f"\n- Proceeding to scan {variable}{' together with '+', '.join(variables_scanTogether) if len(variables_scanTogether)>0 else ''}:") - - tglf_executor = {} - tglf_executor_full = {} - folders = [] - for cont_mult, mult in enumerate(varUpDown): - mult = round(mult, 6) - - if relativeChanges: - print(f"\n + Multiplier: {mult} -----------------------------------------------------------------------------------------------------------") - else: - print(f"\n + Value: {mult} ----------------------------------------------------------------------------------------------------------------") - - multipliers_mod[variable] = mult - - for variable_scanTogether in variables_scanTogether: - multipliers_mod[variable_scanTogether] = mult - - name = f"{variable}_{mult}" - - species = self.inputs_files[self.rhos[0]] # Any rho will do - - multipliers_mod = completeVariation(multipliers_mod, species) - - if not relativeChanges: - for ikey in multipliers_mod: - 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"]: - kwargs_TGLFrun["Quasineutral"] = True - - # Only ask the cold_start in the first round - kwargs_TGLFrun["forceIfcold_start"] = cont_mult > 0 or ( - "forceIfcold_start" in kwargs_TGLFrun and kwargs_TGLFrun["forceIfcold_start"] - ) - - tglf_executor, tglf_executor_full, folderlast = self.prep_run( - f"{self.subFolderTGLF_scan}_{name}", - tglf_executor=tglf_executor, - tglf_executor_full=tglf_executor_full, - multipliers=multipliers_mod, - minimum_delta_abs=minimum_delta_abs, - **kwargs_TGLFrun, - ) - - folders.append(copy.deepcopy(folderlast)) - - return tglf_executor, tglf_executor_full, folders, varUpDown - - def readScan( - self, label="scan1", subFolderTGLF=None, variable="RLTS_1", positionIon=2 + positionIon=2 ): ''' positionIon is the index in the input.tglf file... so if you want for ion RLNS_5, positionIon=5 ''' if subFolderTGLF is None: - subFolderTGLF = self.subFolderTGLF_scan + subFolderTGLF = self.subFolder_scan self.scans[label] = {} self.scans[label]["variable"] = variable @@ -2320,6 +2115,7 @@ def plotScan( forceXposition=None, plotTGLFs=True, ): + unnormalization_successful = True for label in labels: unnormalization_successful = ( @@ -2898,14 +2694,14 @@ def runScanTurbulenceDrives( # Prepare all scans # ------------------------------------------ - tglf_executor, tglf_executor_full, folders = {}, {}, [] + code_executor, code_executor_full, folders = {}, {}, [] for cont, variable in enumerate(self.variablesDrives): # Only ask the cold_start in the first round kwargs_TGLFrun["forceIfcold_start"] = cont > 0 or ("forceIfcold_start" in kwargs_TGLFrun and kwargs_TGLFrun["forceIfcold_start"]) scan_name = f"{subFolderTGLF}_{variable}" # e.g. turbDrives_RLTS_1 - tglf_executor0, tglf_executor_full0, folders0, _ = self._prepare_scan( + code_executor0, code_executor_full0, folders0, _ = self._prepare_scan( scan_name, variable=variable, varUpDown=varUpDown_dict[variable], @@ -2913,8 +2709,8 @@ def runScanTurbulenceDrives( **kwargs_TGLFrun, ) - tglf_executor = tglf_executor | tglf_executor0 - tglf_executor_full = tglf_executor_full | tglf_executor_full0 + code_executor = code_executor | code_executor0 + code_executor_full = code_executor_full | code_executor_full0 folders += folders0 # ------------------------------------------ @@ -2922,8 +2718,8 @@ def runScanTurbulenceDrives( # ------------------------------------------ self._run( - tglf_executor, - tglf_executor_full=tglf_executor_full, + code_executor, + code_executor_full=code_executor_full, **kwargs_TGLFrun, ) @@ -2936,7 +2732,7 @@ def runScanTurbulenceDrives( for mult in varUpDown_dict[variable]: name = f"{variable}_{mult}" self.read( - label=f"{self.subFolderTGLF_scan}_{name}", + label=f"{self.subFolder_scan}_{name}", folder=folders[cont], cold_startWF = False, require_all_files=not kwargs_TGLFrun.get("only_minimal_files", False), @@ -3486,7 +3282,7 @@ def updateConvolution(self): ) = GACODEdefaults.convolution_CECE(self.d_perp_dict, dRdx=self.DRMAJDX_LOC) -def completeVariation(setVariations, species): +def completeVariation_TGLF(setVariations, species): ions_info = species.ions_info setVariations_new = copy.deepcopy(setVariations) diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 247a7835..86e9dd6a 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -369,7 +369,7 @@ def run( inputclass_TGLF = TGLFtools.TGLFinput() inputclass_TGLF = GACODErun.modifyInputs( inputclass_TGLF, - Settings=TGLFsettings, + code_settings=TGLFsettings, extraOptions=extraOptionsTGLF, addControlFunction=GACODEdefaults.addTGLFcontrol, NS=self.loc_n_ion + 1, @@ -634,7 +634,7 @@ def grab_tglf_objects(self, subfolder="tglf_runs", fromlabel="tgyro1", rhos=None inputsTGLF[rho] = inputclass tglf = TGLFtools.TGLF(rhos=rhos) - tglf.prep( + tglf.prep_using_tgyro( self.FolderGACODE / subfolder, specificInputs=inputsTGLF, inputgacode=self.FolderTGYRO / "input.gacode", diff --git a/src/mitim_tools/gacode_tools/scripts/read_tglf.py b/src/mitim_tools/gacode_tools/scripts/read_tglf.py index 09086fc4..54cb497d 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_tglf.py +++ b/src/mitim_tools/gacode_tools/scripts/read_tglf.py @@ -31,7 +31,7 @@ def main(): suffixes[i] = "" tglf = TGLFtools.TGLF() - tglf.prep_from_tglf( + tglf.prep_from_file( folders[0], folders[0] / f"input.tglf{suffixes[0]}", input_gacode=input_gacode ) for i, folder in enumerate(folders): diff --git a/src/mitim_tools/gacode_tools/scripts/run_tglf.py b/src/mitim_tools/gacode_tools/scripts/run_tglf.py index b146cef5..49bc6c3f 100644 --- a/src/mitim_tools/gacode_tools/scripts/run_tglf.py +++ b/src/mitim_tools/gacode_tools/scripts/run_tglf.py @@ -45,7 +45,7 @@ def main(): # ------------------------------------------------------------------------------ tglf = TGLFtools.TGLF() - tglf.prep_from_tglf(folder, input_tglf, input_gacode=input_gacode) + tglf.prep_from_file(folder, input_tglf, input_gacode=input_gacode) # ------------------------------------------------------------------------------ # Workflow diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index 972ac156..c99a8663 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -119,7 +119,7 @@ def TGLFinTRANSP(TGLFsettings, NS=3): return TGLFoptions -def addCGYROcontrol(Settings, rmin): +def addCGYROcontrol(code_settings, rmin): CGYROoptions = IOtools.generateMITIMNamelist( __mitimroot__ / "templates" / "input.cgyro.controls", caseInsensitive=False @@ -137,14 +137,14 @@ def addCGYROcontrol(Settings, rmin): ) as f: settings = json.load(f) - if str(Settings) in settings: - sett = settings[str(Settings)] + if str(code_settings) in settings: + sett = settings[str(code_settings)] label = sett["label"] for ikey in sett["controls"]: CGYROoptions[ikey] = sett["controls"][ikey] else: print( - "\t- Settings not found in input.cgyro.models.json, using defaults", + "\t- code_settings not found in input.cgyro.models.json, using defaults", typeMsg="w", ) label = "unspecified" diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 89993db6..ec2def7f 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -14,8 +14,6 @@ from mitim_tools.misc_tools.PLASMAtools import md_u - - class gacode_simulation: ''' Main class for running GACODE simulations. @@ -32,20 +30,29 @@ def __init__( self.nameRunid = "0" self.results, self.scans = {}, {} + + self.run_specifications = None - def _prep_direct( + def prep( self, mitim_state, # A MITIM state class FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) cold_start=False, # If True, do not use what it potentially inside the folder, run again forceIfcold_start=False, # Extra flag - state_converter='to_tglf', - input_class=None, - input_file='input.tglf' ): + ''' + This method prepares the GACODE run from a MITIM state class by setting up the necessary input files and directories. + ''' print("> Preparation run from input.gacode (direct conversion)") + if self.run_specifications is None: + raise Exception("[MITIM] Simulation child class did not define run specifications") + + state_converter = self.run_specifications['state_converter'] # e.g. to_tglf + input_class = self.run_specifications['input_class'] # e.g. TGLFinput + input_file = self.run_specifications['input_file'] # e.g. input.tglf + self.FolderGACODE = IOtools.expandPath(FolderGACODE) if cold_start or not self.FolderGACODE.exists(): @@ -88,32 +95,118 @@ def _prep_direct( self.NormalizationSets, cdf = NORMtools.normalizations(self.profiles) return cdf + + def run( + self, + subfolder, # 'neo1/', + code_settings=None, + extraOptions={}, + multipliers={}, + minimum_delta_abs={}, + ApplyCorrections=True, # Removing ions with too low density and that are fast species + Quasineutral=False, # Ensures quasineutrality. By default is False because I may want to run the file directly + launchSlurm=True, + cold_start=False, + forceIfcold_start=False, + extra_name="exe", + slurm_setup={"cores": 1,"minutes": 1}, # Cores per call (so, when running nR radii -> nR*4) + attempts_execution=1, + only_minimal_files=False, + ): - def _generic_run_prep( + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare inputs + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + code_executor, code_executor_full = self._run_prepare( + # + subfolder, + code_executor={}, + code_executor_full={}, + # + code_settings=code_settings, + extraOptions=extraOptions, + multipliers=multipliers, + # + cold_start=cold_start, + forceIfcold_start=forceIfcold_start, + only_minimal_files=only_minimal_files, + # + launchSlurm=launchSlurm, + slurm_setup=slurm_setup, + # + ApplyCorrections=ApplyCorrections, + minimum_delta_abs=minimum_delta_abs, + Quasineutral=Quasineutral, + ) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Run NEO + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self._run( + code_executor, + code_executor_full=code_executor_full, + code_settings=code_settings, + ApplyCorrections=ApplyCorrections, + Quasineutral=Quasineutral, + launchSlurm=launchSlurm, + cold_start=cold_start, + forceIfcold_start=forceIfcold_start, + extra_name=extra_name, + slurm_setup=slurm_setup, + only_minimal_files=only_minimal_files, + attempts_execution=attempts_execution, + ) + + return code_executor_full + + def _run_prepare( self, + # ******************************** + # Required options + # ******************************** subfolder_simulation, - rhos=None, code_executor=None, code_executor_full=None, + # ******************************** + # Run settings + # ******************************** code_settings=None, extraOptions={}, multipliers={}, + # ******************************** + # IO settings + # ******************************** cold_start=False, forceIfcold_start=False, only_minimal_files=False, - minimum_delta_abs={}, - ApplyCorrections=True, - Quasineutral=False, + # ******************************** + # Slurm settings (for warnings) + # ******************************** launchSlurm=True, - slurm_setup={ - "cores": 4, - "minutes": 5, - }, # Cores per call (so, when running nR radii -> nR*4) - addControlFunction=None, - controls_file='input.tglf.controls', - **kwargs + slurm_setup=None, + # ******************************** + # Additional settings to correct/modify inputs + # ******************************** + **kwargs_control ): + if slurm_setup is None: + slurm_setup = {"cores": self.run_specifications['default_cores'], "minutes": 5} + + if self.run_specifications is None: + raise Exception("[MITIM] Simulation child class did not define run specifications") + + # Because of historical relevance, I allow both TGLFsettings and code_settings #TODO #TOREMOVE + if "TGLFsettings" in kwargs_control: + if code_settings is not None: + raise Exception('[MITIM] Cannot use both TGLFsettings and code_settings') + else: + code_settings = kwargs_control["TGLFsettings"] + del kwargs_control["TGLFsettings"] + # ------------------------------------------------------------------------------------ + if code_executor is None: code_executor = {} if code_executor_full is None: @@ -123,8 +216,7 @@ def _generic_run_prep( # Prepare for run # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if rhos is None: - rhos = self.rhos + rhos = self.rhos inputs = copy.deepcopy(self.inputs_files) Folder_sim = self.FolderGACODE / subfolder_simulation @@ -156,21 +248,16 @@ def _generic_run_prep( # Change this specific run # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - ( - latest_inputsFile, - latest_inputsFileDict, - ) = change_and_write_code( + latest_inputsFile, latest_inputsFileDict = change_and_write_code( rhos, inputs, Folder_sim, code_settings=code_settings, extraOptions=extraOptions, multipliers=multipliers, - minimum_delta_abs=minimum_delta_abs, - ApplyCorrections=ApplyCorrections, - Quasineutral=Quasineutral, - addControlFunction=addControlFunction, - controls_file=controls_file, + addControlFunction=self.run_specifications['control_function'], + controls_file=self.run_specifications['controls_file'], + **kwargs_control ) code_executor_full[subfolder_simulation] = {} @@ -184,30 +271,35 @@ def _generic_run_prep( "multipliers": multipliers, } if irho in rhosEvaluate: - code_executor[subfolder_simulation][irho] = code_executor_full[subfolder_simulation][ - irho - ] + code_executor[subfolder_simulation][irho] = code_executor_full[subfolder_simulation][irho] # Check input file problems for irho in latest_inputsFileDict: latest_inputsFileDict[irho].anticipate_problems() # Check cores problem - expected_allocated_cores = int(len(rhosEvaluate) * slurm_setup["cores"]) - warning = 32 * 2 if launchSlurm: - print(f'\t- Slurm job will be submitted with {expected_allocated_cores} cores ({len(rhosEvaluate)} radii x {slurm_setup["cores"]} cores/radius)', - typeMsg="" if expected_allocated_cores < warning else "q",) + self._check_cores(rhosEvaluate, slurm_setup) + + self.FolderSimLast = Folder_sim - return code_executor, code_executor_full, Folder_sim + return code_executor, code_executor_full - def _generic_run( + def _check_cores(self, rhosEvaluate, slurm_setup, warning = 32 * 2): + expected_allocated_cores = int(len(rhosEvaluate) * slurm_setup["cores"]) + + print(f'\t- Slurm job will be submitted with {expected_allocated_cores} cores ({len(rhosEvaluate)} radii x {slurm_setup["cores"]} cores/radius)', + typeMsg="" if expected_allocated_cores < warning else "q",) + + def _run( self, code_executor, - run_specifications, **kwargs_run ): - + """ + extraOptions and multipliers are not being grabbed from kwargs_NEOrun, but from code_executor for WF + """ + if kwargs_run.get("only_minimal_files", False): filesToRetrieve = self.ResultsFiles_minimal else: @@ -221,16 +313,152 @@ def _generic_run( run_gacode_simulation( self.FolderGACODE, code_executor, - run_specifications=run_specifications, + run_specifications=self.run_specifications, filesToRetrieve=filesToRetrieve, minutes=kwargs_run.get("slurm_setup", {}).get("minutes", 5), - cores_simulation=kwargs_run.get("slurm_setup", {}).get("cores", 4), - name=f"{run_specifications['code']}_{self.nameRunid}{kwargs_run.get('extra_name', '')}", + cores_simulation=kwargs_run.get("slurm_setup", {}).get("cores", self.run_specifications['default_cores']), + name=f"{self.run_specifications['code']}_{self.nameRunid}{kwargs_run.get('extra_name', '')}", launchSlurm=kwargs_run.get("launchSlurm", True), attempts_execution=kwargs_run.get("attempts_execution", 1), ) else: - print(f"\t- {run_specifications['code'].upper()} not run because all results files found (please ensure consistency!)",typeMsg="i") + print(f"\t- {self.run_specifications['code'].upper()} not run because all results files found (please ensure consistency!)",typeMsg="i") + + def run_scan( + self, + subfolder, # 'scan1', + multipliers={}, + minimum_delta_abs={}, + variable="RLTS_1", + varUpDown=[0.5, 1.0, 1.5], + variables_scanTogether=[], + relativeChanges=True, + **kwargs_run, + ): + + # ------------------------------------- + # Add baseline + # ------------------------------------- + if (1.0 not in varUpDown) and relativeChanges: + print("\n* Since variations vector did not include base case, I am adding it",typeMsg="i",) + varUpDown_new = [] + added = False + for i in varUpDown: + if i > 1.0 and not added: + varUpDown_new.append(1.0) + added = True + varUpDown_new.append(i) + else: + varUpDown_new = varUpDown + + + code_executor, code_executor_full, folders, varUpDown_new = self._prepare_scan( + subfolder, + multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, + variable=variable, + varUpDown=varUpDown_new, + variables_scanTogether=variables_scanTogether, + relativeChanges=relativeChanges, + **kwargs_run, + ) + + # Run them all + self._run( + code_executor, + code_executor_full=code_executor_full, + **kwargs_run, + ) + + # Read results + for cont_mult, mult in enumerate(varUpDown_new): + name = f"{variable}_{mult}" + self.read( + label=f"{self.subFolder_scan}_{name}", + folder=folders[cont_mult], + cold_startWF = False, + require_all_files=not kwargs_run.get("only_minimal_files",False), + ) + + return code_executor_full + + def _prepare_scan( + self, + subfolder, # 'scan1', + multipliers={}, + minimum_delta_abs={}, + variable="RLTS_1", + varUpDown=[0.5, 1.0, 1.5], + variables_scanTogether=[], + relativeChanges=True, + **kwargs_run, + ): + """ + Multipliers will be modified by adding the scaning variables, but I don't want to modify the original + multipliers, as they may be passed to the next scan + + Set relativeChanges=False if varUpDown contains the exact values to change, not multipleiers + """ + + completeVariation = self.run_specifications['complete_variation'] + + multipliers_mod = copy.deepcopy(multipliers) + + self.subFolder_scan = subfolder + + if relativeChanges: + for i in range(len(varUpDown)): + varUpDown[i] = round(varUpDown[i], 6) + + print(f"\n- Proceeding to scan {variable}{' together with '+', '.join(variables_scanTogether) if len(variables_scanTogether)>0 else ''}:") + + code_executor = {} + code_executor_full = {} + folders = [] + for cont_mult, mult in enumerate(varUpDown): + mult = round(mult, 6) + + if relativeChanges: + print(f"\n + Multiplier: {mult} -----------------------------------------------------------------------------------------------------------") + else: + print(f"\n + Value: {mult} ----------------------------------------------------------------------------------------------------------------") + + multipliers_mod[variable] = mult + + for variable_scanTogether in variables_scanTogether: + multipliers_mod[variable_scanTogether] = mult + + name = f"{variable}_{mult}" + + species = self.inputs_files[self.rhos[0]] # Any rho will do + + if completeVariation is not None: + multipliers_mod = completeVariation(multipliers_mod, species) + + if not relativeChanges: + for ikey in multipliers_mod: + kwargs_run["extraOptions"][ikey] = multipliers_mod[ikey] + multipliers_mod = {} + + # Force ensure quasineutrality if the + if variable in ["AS_3", "AS_4", "AS_5", "AS_6"]: + kwargs_run["Quasineutral"] = True + + # Only ask the cold_start in the first round + kwargs_run["forceIfcold_start"] = cont_mult > 0 or ("forceIfcold_start" in kwargs_run and kwargs_run["forceIfcold_start"]) + + code_executor, code_executor_full = self._run_prepare( + f"{self.subFolder_scan}_{name}", + code_executor=code_executor, + code_executor_full=code_executor_full, + multipliers=multipliers_mod, + minimum_delta_abs=minimum_delta_abs, + **kwargs_run, + ) + + folders.append(copy.deepcopy(self.FolderSimLast)) + + return code_executor, code_executor_full, folders, varUpDown def change_and_write_code( @@ -245,6 +473,7 @@ def change_and_write_code( Quasineutral=False, addControlFunction=None, controls_file='input.tglf.controls', + **kwargs ): """ Received inputs classes and gives text. @@ -259,7 +488,7 @@ def change_and_write_code( print(f"\t- Changing input file for rho={rho:.4f}") input_sim_rho = modifyInputs( inputs[rho], - Settings=code_settings, + code_settings=code_settings, extraOptions=extraOptions, multipliers=multipliers, minimum_delta_abs=minimum_delta_abs, @@ -369,6 +598,7 @@ def run_gacode_simulation( attempts_execution = 1, max_jobs_at_once = None, ): + """ launchSlurm = True -> Launch as a batch job in the machine chosen launchSlurm = False -> Launch locally as a bash script @@ -687,7 +917,7 @@ def runTGYRO( def modifyInputs( input_class, - Settings=None, + code_settings=None, extraOptions={}, multipliers={}, minimum_delta_abs={}, @@ -702,15 +932,15 @@ def modifyInputs( GACODEdefaults.review_controls(multipliers, control = controls_file) # ------------------------------------------- - if Settings is not None: - CodeOptions = addControlFunction(Settings, **kwargs_to_function) + if code_settings is not None: + CodeOptions = addControlFunction(code_settings, **kwargs_to_function) # ~~~~~~~~~~ Change with presets - print(f" \t- Using presets Settings = {Settings}", typeMsg="i") + print(f" \t- Using presets code_settings = {code_settings}", typeMsg="i") input_class.controls = CodeOptions else: - print("\t- Input file was not modified by Settings, using what was there before",typeMsg="i") + print("\t- Input file was not modified by code_settings, using what was there before",typeMsg="i") # Make all upper case extraOptions = {ikey.upper(): value for ikey, value in extraOptions.items()} diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 807d40a5..38aa7e74 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -13,18 +13,20 @@ if cold_start and folder.exists(): os.system(f"rm -r {folder.resolve()}") -neo = NEOtools.NEO(rhos=np.linspace(0.1,0.95,20)) -neo.prep_direct(input_gacode, folder) +neo = NEOtools.NEO(rhos=np.linspace(0.1,0.95,10)) +neo.prep(input_gacode, folder) -neo.run('neo1/') +neo.run('neo1/', cold_start=cold_start) neo.read('NEO default') -neo.run('neo2/', extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}) +neo.run('neo2/', cold_start=cold_start, extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}) neo.read('NEO low res') -neo.run('neo3/', extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}, multipliers={'DLNTDR_1': 1.5}) +neo.run('neo3/', cold_start=cold_start, extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}, multipliers={'DLNTDR_1': 1.5}) neo.read('NEO low res + 50% aLTe') +# neo.run_scan('scan1', cold_start=cold_start, variable='DLNTDR_1', varUpDown=np.linspace(0.5, 1.5, 4)) + neo.plot(labels=['NEO default', 'NEO low res', 'NEO low res + 50% aLTe']) neo.fn.show() neo.fn.close() \ No newline at end of file diff --git a/tests/TGLF_workflow.py b/tests/TGLF_workflow.py index 34cdd186..09118dcf 100644 --- a/tests/TGLF_workflow.py +++ b/tests/TGLF_workflow.py @@ -13,11 +13,11 @@ os.system(f"rm -r {folder.resolve()}") tglf = TGLFtools.TGLF() -tglf.prep_from_tglf(folder, input_tglf) +tglf.prep_from_file(folder, input_tglf) tglf.run( - subFolderTGLF="run1/", - TGLFsettings=None, + "run1/", + code_settings=None, cold_start=cold_start, runWaveForms = [0.67, 10.0], forceIfcold_start=True, @@ -28,8 +28,8 @@ tglf.read(label="ES") tglf.run( - subFolderTGLF="run2/", - TGLFsettings=None, + "run2/", + code_settings=None, cold_start=cold_start, forceIfcold_start=True, extraOptions={"USE_BPER": True, "USE_BPAR": True}, diff --git a/tests/TGLFfull_workflow.py b/tests/TGLFfull_workflow.py index 88789a4f..031d077f 100644 --- a/tests/TGLFfull_workflow.py +++ b/tests/TGLFfull_workflow.py @@ -18,7 +18,7 @@ os.system(f"rm -r {folder}") tglf = TGLFtools.TGLF(cdf=cdf_file, time=2.5, avTime=0.02, rhos=np.array([0.6, 0.8])) -_ = tglf.prep(folder, cold_start=cold_start) +_ = tglf.prep_using_tgyro(folder, cold_start=cold_start) tglf.run( subFolderTGLF="runSAT2", diff --git a/tests/TGLFscan_workflow.py b/tests/TGLFscan_workflow.py index 5f5a7a70..47e1c99d 100644 --- a/tests/TGLFscan_workflow.py +++ b/tests/TGLFscan_workflow.py @@ -3,7 +3,7 @@ from mitim_tools.gacode_tools import TGLFtools, PROFILEStools from mitim_tools import __mitimroot__ -cold_start = True +cold_start = False (__mitimroot__ / 'tests' / 'scratch').mkdir(parents=True, exist_ok=True) @@ -14,7 +14,7 @@ os.system(f"rm -r {folder.resolve()}") tglf = TGLFtools.TGLF(rhos=[0.5, 0.7]) -tglf.prep_direct(input_gacode,folder, cold_start=cold_start) +tglf.prep(input_gacode,folder, cold_start=cold_start) tglf.runScan( subFolderTGLF = 'scan1', TGLFsettings = None, diff --git a/tests/VITALS_workflow.py b/tests/VITALS_workflow.py index df0ebf1f..f1c591dc 100644 --- a/tests/VITALS_workflow.py +++ b/tests/VITALS_workflow.py @@ -30,7 +30,7 @@ # ******************************************************************************** tglf = TGLFtools.TGLF(rhos=[rho]) -cdf = tglf.prep(folderWork, cold_start=cold_start, inputgacode=inputgacode) +cdf = tglf.prep_using_tgyro(folderWork, cold_start=cold_start, inputgacode=inputgacode) tglf.run(subFolderTGLF="run_base/", TGLFsettings=TGLFsettings, cold_start=cold_start) # ******************************************************************************** diff --git a/tests/data/FolderTRANSP/12345X01TR.DAT b/tests/data/FolderTRANSP/12345X01TR.DAT index 7c9d45a5..b4e80213 100644 --- a/tests/data/FolderTRANSP/12345X01TR.DAT +++ b/tests/data/FolderTRANSP/12345X01TR.DAT @@ -33,7 +33,7 @@ tgrid2 = 1e-3 ! Control of time resolution of 2D input data dtmaxg = 0.001 ! Max time step for MHD -!----- MPI Settings +!----- MPI code_settings nbi_pserve =0 ntoric_pserve =1 diff --git a/tutorials/TGLF_tutorial.py b/tutorials/TGLF_tutorial.py index c9927795..41adba28 100644 --- a/tutorials/TGLF_tutorial.py +++ b/tutorials/TGLF_tutorial.py @@ -10,7 +10,7 @@ tglf = TGLFtools.TGLF(rhos=[0.5, 0.7]) # Prepare the TGLF class -cdf = tglf.prep(folder, inputgacode=inputgacode_file, cold_start=False) +cdf = tglf.prep_using_tgyro(folder, inputgacode=inputgacode_file, cold_start=False) ''' *************************************************************************** diff --git a/tutorials/run_slurm_array_tutorial/test_launcher.py b/tutorials/run_slurm_array_tutorial/test_launcher.py index 229afbf4..28c299dd 100644 --- a/tutorials/run_slurm_array_tutorial/test_launcher.py +++ b/tutorials/run_slurm_array_tutorial/test_launcher.py @@ -10,7 +10,7 @@ print(f"Using partition: {partition}") -# Settings for slurm job +# code_settings for slurm job cpus = 2 hours = 1 memory = '100GB' From e9eb771810058a0a83946b19f91fc6903058e931 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 21 Aug 2025 16:08:04 -0400 Subject: [PATCH 179/385] misc --- src/mitim_tools/misc_tools/FARMINGtools.py | 28 ++++++---------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 9a23c753..005f279b 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -1059,12 +1059,8 @@ def create_slurm_execution_files( # ******* Basics commandSBATCH.append("#!/usr/bin/env bash") commandSBATCH.append(f"#SBATCH --job-name {nameJob}") - commandSBATCH.append( - f"#SBATCH --output {folderExecution}/slurm_output{label_log_files}.dat" - ) - commandSBATCH.append( - f"#SBATCH --error {folderExecution}/slurm_error{label_log_files}.dat" - ) + commandSBATCH.append(f"#SBATCH --output {folderExecution}/slurm_output{label_log_files}.dat") + commandSBATCH.append(f"#SBATCH --error {folderExecution}/slurm_error{label_log_files}.dat") if email is not None: commandSBATCH.append("#SBATCH --mail-user=" + email) @@ -1103,21 +1099,11 @@ def create_slurm_execution_files( # ******* Commands commandSBATCH.append("") - commandSBATCH.append( - 'echo "MITIM: Submitting SLURM job $SLURM_JOBID in $HOSTNAME (host: $SLURM_SUBMIT_HOST)"' - ) - commandSBATCH.append( - 'echo "MITIM: Nodes have $SLURM_CPUS_ON_NODE cores and $SLURM_JOB_NUM_NODES node(s) were allocated for this job"' - ) - commandSBATCH.append( - 'echo "MITIM: Each of the $SLURM_NTASKS tasks allocated will run with $SLURM_CPUS_PER_TASK cores, allocating $SRUN_CPUS_PER_TASK CPUs per srun"' - ) - commandSBATCH.append( - 'echo "***********************************************************************************************"' - ) - commandSBATCH.append( - 'echo ""' - ) + commandSBATCH.append('echo "MITIM: Submitting SLURM job $SLURM_JOBID in $HOSTNAME (host: $SLURM_SUBMIT_HOST)"') + commandSBATCH.append('echo "MITIM: Nodes have $SLURM_CPUS_ON_NODE cores and $SLURM_JOB_NUM_NODES node(s) were allocated for this job"') + commandSBATCH.append('echo "MITIM: Each of the $SLURM_NTASKS tasks allocated will run with $SLURM_CPUS_PER_TASK cores, allocating $SRUN_CPUS_PER_TASK CPUs per srun"') + commandSBATCH.append('echo "***********************************************************************************************"') + commandSBATCH.append('echo ""') commandSBATCH.append("") full_command = [modules_remote] if (modules_remote is not None) else [] From 1c5b176d7ca6c6f74aee3b5bea28204cedfa40c7 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 22 Aug 2025 10:51:13 -0400 Subject: [PATCH 180/385] Generalization of TGLF scans and application to NEO as well --- src/mitim_tools/gacode_tools/NEOtools.py | 104 +++++- src/mitim_tools/gacode_tools/TGLFtools.py | 308 +++++++----------- .../gacode_tools/utils/GACODErun.py | 91 +++++- tests/NEO_workflow.py | 9 +- tests/TGLFscan_workflow.py | 10 +- 5 files changed, 317 insertions(+), 205 deletions(-) diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index b6664ee3..15afebbe 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -48,7 +48,11 @@ def read( if folder is None: folder = self.FolderSimLast - self.results[label] = {'NEOout':[]} + self.results[label] = { + 'NEOout':[], + 'parsed': [], + "x": np.array(self.rhos), + } for rho in self.rhos: NEOout = NEOoutput( @@ -64,6 +68,7 @@ def read( self.results[label]['NEOout'].append(NEOout) + self.results[label]['parsed'].append(GACODErun.buildDictFromInput(NEOout.inputFile)) def plot( self, @@ -111,6 +116,93 @@ def plot( axQi.set_ylabel("$Q_i$ ($MW/m^2$)"); axQi.set_yscale('log') axGe.set_ylabel("$\\Gamma_e$ ($1E20/s/m^2$)"); #axGe.set_yscale('log') + + def read_scan( + self, + label="scan1", + subfolder=None, + variable="RLTS_1", + positionIon=2 + ): + + output_object = "NEOout" + + variable_mapping = { + 'scanned_variable': ["parsed", variable, None], + 'Qe_gb': [output_object, 'Qe', None], + 'Qi_gb': [output_object, 'Qi', None], + 'Ge_gb': [output_object, 'Ge', None], + 'Gi_gb': [output_object, 'GiAll', positionIon - 2], + 'Mt_gb': [output_object, 'Mt', None], + } + + variable_mapping_unn = { + 'Qe': [output_object, 'Qe_unn', None], + 'Qi': [output_object, 'Qi_unn', None], + 'Ge': [output_object, 'Ge_unn', None], + 'Gi': [output_object, 'GiAll_unn', positionIon - 2], + 'Mt': [output_object, 'Mt_unn', None], + } + + super().read_scan( + label=label, + subfolder=subfolder, + variable=variable, + positionIon=positionIon, + variable_mapping=variable_mapping, + variable_mapping_unn=variable_mapping_unn + ) + + def plot_scan( + self, + fn=None, + labels=["neo1"], + extratitle="", + fn_color=None, + colors=None, + ): + + if fn is None: + self.fn = GUItools.FigureNotebook("NEO Scan Notebook", geometry="1700x900", vertical=True) + else: + self.fn = fn + + fig1 = self.fn.add_figure(label=f"{extratitle}Summary", tab_color=fn_color) + + grid = plt.GridSpec(1, 3, hspace=0.7, wspace=0.2) + + if colors is None: + colors = GRAPHICStools.listColors() + + axQe = fig1.add_subplot(grid[0, 0]) + axQi = fig1.add_subplot(grid[0, 1]) + axGe = fig1.add_subplot(grid[0, 2]) + + cont = 0 + for label in labels: + for irho in range(len(self.rhos)): + + x = self.scans[label]['scanned_variable'][irho] + + axQe.plot(x, self.scans[label]['Qe'][irho], label=f'{label}, {self.rhos[irho]}', color=colors[cont], marker='o', linestyle='-') + axQi.plot(x, self.scans[label]['Qi'][irho], label=f'{label}, {self.rhos[irho]}', color=colors[cont], marker='o', linestyle='-') + axGe.plot(x, self.scans[label]['Ge'][irho], label=f'{label}, {self.rhos[irho]}', color=colors[cont], marker='o', linestyle='-') + + cont += 1 + + for ax in [axQe, axQi, axGe]: + ax.set_xlabel("Scanned variable") + GRAPHICStools.addDenseAxis(ax) + ax.legend(loc="best") + + axQe.set_ylabel("$Q_e$ ($MW/m^2$)"); + axQi.set_ylabel("$Q_i$ ($MW/m^2$)"); + axGe.set_ylabel("$\\Gamma_e$ ($1E20/s/m^2$)") + + plt.tight_layout() + + + # def prep(self, inputgacode, folder): # self.inputgacode = inputgacode # self.folder = IOtools.expandPath(folder) @@ -291,6 +383,16 @@ def read(self): self.roa = float(lines[0].split()[-1]) + + # ------------------------------------------------------------------------ + # Input file + # ------------------------------------------------------------------------ + + with open(self.FolderGACODE / ("input.neo" + self.suffix), "r") as fi: + lines = fi.readlines() + self.inputFile = "".join(lines) + + def unnormalize(self, normalization, rho=None): if normalization is not None: diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 10ed3462..54392f47 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -75,7 +75,7 @@ def __init__( # Run standalone TGLF (this will find the input.tglf in the previous folder, # and then copy to this specify TGLF run, and run it there) - tglf.run(subFolderTGLF='tglf1/',TGLFsettings=1,extraOptions={'NS':3}) + tglf.run(subfolder='tglf1/',TGLFsettings=1,extraOptions={'NS':3}) # Read results tglf.read(label='run1',folder='~/testTGLF/tglf1/') @@ -94,13 +94,13 @@ def __init__( cdf = tglf.prep_using_tgyro('~/testTGLF/') # Run - tglf.runScan('scan1/',TGLFsettings=1,varUpDown=np.linspace(0.5,2.0,20),variable='RLTS_2') + tglf.run_scan('scan1/',TGLFsettings=1,varUpDown=np.linspace(0.5,2.0,20),variable='RLTS_2') # Read scan - tglf.readScan(label='scan1',variable='RLTS_2') + tglf.read_scan(label='scan1',variable='RLTS_2') # Plot - plt.ion(); tglf.plotScan(labels=['scan1'],variableLabel='RLTS_2') + plt.ion(); tglf.plot_scan(labels=['scan1'],variableLabel='RLTS_2') **************************** ***** Special analysis ***** @@ -194,7 +194,7 @@ def __init__( def run( self, - subFolderTGLF, + subfolder, runWaveForms=None, # e.g. runWaveForms = [0.3,1.0] forceClosestUnstableWF=True, # Look at the growth rate spectrum and run exactly the ky of the closest unstable **kwargs_generic_run @@ -203,7 +203,7 @@ def run( I need to redefine the run method for the TGLF class because it has the option of producing WaveForms ''' - code_executor_full = super().run(subFolderTGLF, **kwargs_generic_run) + code_executor_full = super().run(subfolder, **kwargs_generic_run) kwargs_generic_run['runWaveForms'] = runWaveForms kwargs_generic_run['forceClosestUnstableWF'] = forceClosestUnstableWF @@ -211,23 +211,14 @@ def run( def run_scan( self, - subFolderTGLF, + subfolder, **kwargs, ): - code_executor_full = super().run_scan(subFolderTGLF,**kwargs) + code_executor_full = super().run_scan(subfolder,**kwargs) self._helper_wf(code_executor_full, **kwargs) - # TOREMOVE #TODO - def runScan( - self, - subFolderTGLF, - **kwargs, - ): - - self.run_scan(subFolderTGLF, **kwargs) - def _run_wf(self, kys, code_executor, forceClosestUnstableWF=True, **kwargs_TGLFrun): """ extraOptions and multipliers are not being grabbed from kwargs_TGLFrun, but from code_executor @@ -257,11 +248,11 @@ def _run_wf(self, kys, code_executor, forceClosestUnstableWF=True, **kwargs_TGLF print(f"> Running TGLF waveform analysis, ky~{ky_single0}") self.FoldersTGLF_WF[f"ky{ky_single0}"] = {} - for subFolderTGLF in code_executor: + for subfolder in code_executor: ky_single_orig = copy.deepcopy(ky_single0) - FolderTGLF_old = code_executor[subFolderTGLF][list(code_executor[subFolderTGLF].keys())[0]]["folder"] + FolderTGLF_old = code_executor[subfolder][list(code_executor[subfolder].keys())[0]]["folder"] self.ky_single = None self.read(label=f"ky{ky_single0}", folder=FolderTGLF_old, cold_startWF = False) @@ -301,8 +292,8 @@ def _run_wf(self, kys, code_executor, forceClosestUnstableWF=True, **kwargs_TGLF else: extraOptions_WF = {} - extraOptions_WF = copy.deepcopy(code_executor[subFolderTGLF][list(code_executor[subFolderTGLF].keys())[0]]["extraOptions"]) - multipliers_WF = copy.deepcopy(code_executor[subFolderTGLF][list(code_executor[subFolderTGLF].keys())[0]]["multipliers"]) + extraOptions_WF = copy.deepcopy(code_executor[subfolder][list(code_executor[subfolder].keys())[0]]["extraOptions"]) + multipliers_WF = copy.deepcopy(code_executor[subfolder][list(code_executor[subfolder].keys())[0]]["multipliers"]) extraOptions_WF["USE_TRANSPORT_MODEL"] = "F" extraOptions_WF["WRITE_WAVEFUNCTION_FLAG"] = 1 @@ -1961,150 +1952,76 @@ def _helper_wf( # Get back to it self.FolderSimLast = self.keep_folder + + #TODO #TOREMOVE + def runScan(self,subfolder,**kwargs): + self.run_scan(subfolder, **kwargs) + def readScan(self, **kwargs): + self.read_scan(**kwargs) + def plotScan(self,**kwargs): + self.plot_scan(**kwargs) - def readScan( + def read_scan( self, label="scan1", - subFolderTGLF=None, + subfolder=None, variable="RLTS_1", positionIon=2 ): - ''' - positionIon is the index in the input.tglf file... so if you want for ion RLNS_5, positionIon=5 - ''' - - if subFolderTGLF is None: - subFolderTGLF = self.subFolder_scan - - self.scans[label] = {} - self.scans[label]["variable"] = variable - self.scans[label]["positionBase"] = None - self.scans[label]["unnormalization_successful"] = True - self.scans[label]["results_tags"] = [] - - self.positionIon_scan = positionIon - - # ---- - x = [] - Qe, Qi, Ge, Gi, Mt, S = [],[],[],[],[],[] - Qe_gb, Qi_gb, Ge_gb, Gi_gb, Mt_gb, S_gb = [],[],[],[],[],[] - ky, g, f, eta1, eta2, itg, tem, etg = [],[],[],[],[],[],[],[] - etalow_g, etalow_f, etalow_k = [], [], [] - Qifast, Qifast_gb = [],[] - cont = 0 - for ikey in self.results: - isThisTheRightReadResults = (subFolderTGLF in ikey) and (variable== "_".join(ikey.split("_")[:-1]).split(subFolderTGLF + "_")[-1]) - - if isThisTheRightReadResults: - - self.scans[label]["results_tags"].append(ikey) - x0 = [] - Qe0, Qi0, Ge0, Gi0, Mt0, S0 = [],[],[],[],[],[] - Qe_gb0, Qi_gb0, Ge_gb0, Gi_gb0, Mt_gb0, S_gb0 = [],[],[],[],[],[] - ky0, g0, f0, eta10, eta20, itg0, tem0, etg0 = [],[],[],[],[],[],[],[] - etalow_g0, etalow_f0, etalow_k0 = [], [], [] - Qifast0, Qifast_gb0 = [],[] - for irho_cont in range(len(self.rhos)): - irho = np.where(self.results[ikey]["x"] == self.rhos[irho_cont])[0][0] - - # Unnormalized - x0.append(self.results[ikey]["parsed"][irho][variable]) - Qe_gb0.append(self.results[ikey]["TGLFout"][irho].Qe) - Qi_gb0.append(self.results[ikey]["TGLFout"][irho].Qi) - Qifast_gb0.append(self.results[ikey]["TGLFout"][irho].Qifast) - Ge_gb0.append(self.results[ikey]["TGLFout"][irho].Ge) - Gi_gb0.append(self.results[ikey]["TGLFout"][irho].GiAll[self.positionIon_scan - 2]) - Mt_gb0.append(self.results[ikey]["TGLFout"][irho].Mt) - S_gb0.append(self.results[ikey]["TGLFout"][irho].Se) - ky0.append(self.results[ikey]["TGLFout"][irho].ky) - g0.append(self.results[ikey]["TGLFout"][irho].g) - f0.append(self.results[ikey]["TGLFout"][irho].f) - eta10.append(self.results[ikey]["TGLFout"][irho].etas["metrics"]["eta_ITGTEM"]) - eta20.append(self.results[ikey]["TGLFout"][irho].etas["metrics"]["eta_ITGETG"]) - etalow_g0.append(self.results[ikey]["TGLFout"][irho].etas["metrics"]["g_lowk_max"]) - etalow_k0.append(self.results[ikey]["TGLFout"][irho].etas["metrics"]["k_lowk_max"]) - etalow_f0.append(self.results[ikey]["TGLFout"][irho].etas["metrics"]["f_lowk_max"]) - itg0.append(self.results[ikey]["TGLFout"][irho].etas["ITG"]["g_max"]) - tem0.append(self.results[ikey]["TGLFout"][irho].etas["TEM"]["g_max"]) - etg0.append(self.results[ikey]["TGLFout"][irho].etas["ETG"]["g_max"]) - - if self.results[ikey]["TGLFout"][irho].unnormalization_successful: - Qe0.append(self.results[ikey]["TGLFout"][irho].Qe_unn) - Qi0.append(self.results[ikey]["TGLFout"][irho].Qi_unn) - Qifast0.append(self.results[ikey]["TGLFout"][irho].Qifast_unn) - Ge0.append(self.results[ikey]["TGLFout"][irho].Ge_unn) - Gi0.append(self.results[ikey]["TGLFout"][irho].GiAll_unn[self.positionIon_scan - 2]) - Mt0.append(self.results[ikey]["TGLFout"][irho].Mt_unn) - S0.append(self.results[ikey]["TGLFout"][irho].Se_unn) - else: - self.scans[label]["unnormalization_successful"] = False - - x.append(x0) - Qe.append(Qe0) - Qi.append(Qi0) - Qifast.append(Qifast0) - Ge.append(Ge0) - Qe_gb.append(Qe_gb0) - Qi_gb.append(Qi_gb0) - Qifast_gb.append(Qifast_gb0) - Ge_gb.append(Ge_gb0) - Gi_gb.append(Gi_gb0) - Gi.append(Gi0) - Mt.append(Mt0) - S.append(S0) - ky.append(ky0) - g.append(g0) - f.append(f0) - eta1.append(eta10) - eta2.append(eta20) - etalow_g.append(etalow_g0) - etalow_f.append(etalow_f0) - etalow_k.append(etalow_k0) - itg.append(itg0) - tem.append(tem0) - etg.append(etg0) - - if float(ikey.split('_')[-1]) == 1.0: - self.scans[label]["positionBase"] = cont - cont += 1 + + output_object = "TGLFout" + + variable_mapping = { + 'scanned_variable': ["parsed", variable, None], + 'Qe_gb': [output_object, 'Qe', None], + 'Qi_gb': [output_object, 'Qi', None], + 'Ge_gb': [output_object, 'Ge', None], + 'Gi_gb': [output_object, 'GiAll', positionIon - 2], + 'Mt_gb': [output_object, 'Mt', None], + 'S_gb': [output_object, 'Se', None], + 'ky': [output_object, 'ky', None], + 'g': [output_object, 'g', None], + 'f': [output_object, 'f', None], + 'Qifast_gb': [output_object, 'Qifast', None], + 'eta_ITGETG': [output_object, 'eta_ITGETG', None], + 'eta_ITGTEM': [output_object, 'eta_ITGTEM', None], + 'g_lowk_max': [output_object, 'g_lowk_max', None], + 'f_lowk_max': [output_object, 'f_lowk_max', None], + 'k_lowk_max': [output_object, 'k_lowk_max', None], + 'g_ITG_max': [output_object, 'g_ITG_max', None], + 'g_ETG_max': [output_object, 'g_ETG_max', None], + 'g_TEM_max': [output_object, 'g_TEM_max', None], + } + + variable_mapping_unn = { + 'Qe': [output_object, 'Qe_unn', None], + 'Qi': [output_object, 'Qi_unn', None], + 'Ge': [output_object, 'Ge_unn', None], + 'Gi': [output_object, 'GiAll_unn', positionIon - 2], + 'Mt': [output_object, 'Mt_unn', None], + 'S': [output_object, 'Se_unn', None], + 'Qifast': [output_object, 'Qifast_unn', None], + } + + super().read_scan( + label=label, + subfolder=subfolder, + variable=variable, + positionIon=positionIon, + variable_mapping=variable_mapping, + variable_mapping_unn=variable_mapping_unn + ) - self.scans[label]["x"] = np.array(self.rhos) - self.scans[label]["xV"] = np.atleast_2d(np.transpose(x)) - self.scans[label]["Qe_gb"] = np.atleast_2d(np.transpose(Qe_gb)) - self.scans[label]["Qi_gb"] = np.atleast_2d(np.transpose(Qi_gb)) - self.scans[label]["Qifast_gb"] = np.atleast_2d(np.transpose(Qifast_gb)) - self.scans[label]["Ge_gb"] = np.atleast_2d(np.transpose(Ge_gb)) - self.scans[label]["Gi_gb"] = np.atleast_2d(np.transpose(Gi_gb)) - self.scans[label]["Mt_gb"] = np.atleast_2d(np.transpose(Mt_gb)) - self.scans[label]["S_gb"] = np.atleast_2d(np.transpose(S_gb)) - self.scans[label]["Qe"] = np.atleast_2d(np.transpose(Qe)) - self.scans[label]["Qi"] = np.atleast_2d(np.transpose(Qi)) - self.scans[label]["Qifast"] = np.atleast_2d(np.transpose(Qifast)) - self.scans[label]["Ge"] = np.atleast_2d(np.transpose(Ge)) - self.scans[label]["Gi"] = np.atleast_2d(np.transpose(Gi)) - self.scans[label]["Mt"] = np.atleast_2d(np.transpose(Mt)) - self.scans[label]["S"] = np.atleast_2d(np.transpose(S)) - self.scans[label]["eta1"] = np.atleast_2d(np.transpose(eta1)) - self.scans[label]["eta2"] = np.atleast_2d(np.transpose(eta2)) - self.scans[label]["itg"] = np.atleast_2d(np.transpose(itg)) - self.scans[label]["tem"] = np.atleast_2d(np.transpose(tem)) - self.scans[label]["etg"] = np.atleast_2d(np.transpose(etg)) - self.scans[label]["g_lowk_max"] = np.atleast_2d(np.transpose(etalow_g)) - self.scans[label]["f_lowk_max"] = np.atleast_2d(np.transpose(etalow_f)) - self.scans[label]["k_lowk_max"] = np.atleast_2d(np.transpose(etalow_k)) - self.scans[label]["ky"] = np.array(ky) - self.scans[label]["g"] = np.array(g) - self.scans[label]["f"] = np.array(f) - if len(self.scans[label]["ky"].shape) == 2: - self.scans[label]["ky"] = self.scans[label]["ky"].reshape((1, self.scans[label]["ky"].shape[0], self.scans[label]["ky"].shape[1])) - self.scans[label]["g"] = self.scans[label]["g"].reshape((1, self.scans[label]["g"].shape[0], self.scans[label]["g"].shape[1])) - self.scans[label]["f"] = self.scans[label]["f"].reshape((1, self.scans[label]["f"].shape[0], self.scans[label]["f"].shape[1])) - else: - self.scans[label]["ky"] = np.transpose(self.scans[label]["ky"], axes=[1, 0, 2]) - self.scans[label]["g"] = np.transpose(self.scans[label]["g"], axes=[1, 0, 2, 3]) - self.scans[label]["f"] = np.transpose(self.scans[label]["f"], axes=[1, 0, 2, 3]) + varS = ['ky', 'g', 'f'] + for var in varS: + if len(self.scans[label][var].shape) == 3: + axes_swap = [1, 2, 0] # [rho, scan, ky] + elif len(self.scans[label][var].shape) == 4: + axes_swap = [2, 3, 1, 0] # [rho, scan, nmode, ky] + + self.scans[label][var] = np.transpose(self.scans[label][var], axes=axes_swap) - def plotScan( + def plot_scan( self, labels=["scan1"], figs=None, @@ -2118,15 +2035,10 @@ def plotScan( unnormalization_successful = True for label in labels: - unnormalization_successful = ( - unnormalization_successful - and self.scans[label]["unnormalization_successful"] - ) + unnormalization_successful = unnormalization_successful and self.scans[label]["unnormalization_successful"] if figs is None: - self.fn = GUItools.FigureNotebook( - "TGLF Scan MITIM Notebook", geometry="1500x900", vertical=True - ) + self.fn = GUItools.FigureNotebook("TGLF Scan MITIM Notebook", geometry="1500x900", vertical=True) if unnormalization_successful: fig1 = self.fn.add_figure(label="Fluxes") fig1e = self.fn.add_figure(label="Fluxes (GB)") @@ -2188,7 +2100,7 @@ def plotScan( positionBase = self.scans[label]["positionBase"] - x = self.scans[label]["xV"] + x = self.scans[label]["scanned_variable"] if relativeX: xbase = x[:, positionBase : positionBase + 1] x = (x - xbase) / xbase * 100.0 @@ -2205,11 +2117,11 @@ def plotScan( self.scans[label]["Ge_gb"], self.scans[label]["Gi_gb"], ) - eta1, eta2 = self.scans[label]["eta1"], self.scans[label]["eta2"] + eta1, eta2 = self.scans[label]["eta_ITGETG"], self.scans[label]["eta_ITGTEM"] itg, tem, etg = ( - self.scans[label]["itg"], - self.scans[label]["tem"], - self.scans[label]["etg"], + self.scans[label]["g_ITG_max"], + self.scans[label]["g_TEM_max"], + self.scans[label]["g_ETG_max"], ) ky, g, f = ( self.scans[label]["ky"], @@ -2663,7 +2575,7 @@ def plotScan( def runScanTurbulenceDrives( self, - subFolderTGLF="drives1", + subfolder="drives1", varUpDown = None, # This setting supercedes the resolutionPoints and variation resolutionPoints=5, variation=0.5, @@ -2699,7 +2611,7 @@ def runScanTurbulenceDrives( # Only ask the cold_start in the first round kwargs_TGLFrun["forceIfcold_start"] = cont > 0 or ("forceIfcold_start" in kwargs_TGLFrun and kwargs_TGLFrun["forceIfcold_start"]) - scan_name = f"{subFolderTGLF}_{variable}" # e.g. turbDrives_RLTS_1 + scan_name = f"{subfolder}_{variable}" # e.g. turbDrives_RLTS_1 code_executor0, code_executor_full0, folders0, _ = self._prepare_scan( scan_name, @@ -2732,16 +2644,16 @@ def runScanTurbulenceDrives( for mult in varUpDown_dict[variable]: name = f"{variable}_{mult}" self.read( - label=f"{self.subFolder_scan}_{name}", + label=f"{self.subfolder_scan}_{name}", folder=folders[cont], cold_startWF = False, require_all_files=not kwargs_TGLFrun.get("only_minimal_files", False), ) cont += 1 - scan_name = f"{subFolderTGLF}_{variable}" # e.g. turbDrives_RLTS_1 + scan_name = f"{subfolder}_{variable}" # e.g. turbDrives_RLTS_1 - self.readScan(label=scan_name, variable=variable,positionIon=positionIon) + self.read_scan(label=scan_name, variable=variable,positionIon=positionIon) def plotScanTurbulenceDrives( self, label="drives1", figs=None, **kwargs_TGLFscanPlot @@ -2773,7 +2685,7 @@ def plotScanTurbulenceDrives( kwargs_TGLFscanPlot.pop("figs", None) - self.plotScan( + self.plot_scan( labels=labels, figs=figs1, variableLabel="X", @@ -2782,13 +2694,13 @@ def plotScanTurbulenceDrives( ) kwargs_TGLFscanPlot["plotTGLFs"] = False - self.plotScan( + self.plot_scan( labels=labels, figs=figs2, variableLabel="X", **kwargs_TGLFscanPlot ) def runAnalysis( self, - subFolderTGLF="analysis1", + subfolder="analysis1", label="analysis1", analysisType="chi_e", trace=[50.0, 174.0], @@ -2825,14 +2737,14 @@ def runAnalysis( np.linspace(1, 1 + variation / 2, 6)[1:], ) - self.runScan( - subFolderTGLF, + self.run_scan( + subfolder, varUpDown=varUpDown, variable=self.variable, **kwargs_TGLFrun, ) - self.readScan(label=label, variable=self.variable) + self.read_scan(label=label, variable=self.variable) if analysisType == "chi_e": Te_prof = self.NormalizationSets["SELECTED"]["Te_keV"] @@ -2849,7 +2761,7 @@ def runAnalysis( rho = self.NormalizationSets["SELECTED"]["rho"] a = self.NormalizationSets["SELECTED"]["rmin"][-1] - x = self.scans[label]["xV"] + x = self.scans[label]["scanned_variable"] yV = self.scans[label][self.variable_y] self.scans[label]["chi_inc"] = [] @@ -2905,16 +2817,16 @@ def runAnalysis( self.variable = f"RLNS_{position}" - self.runScan( - subFolderTGLF, + self.run_scan( + subfolder, varUpDown=varUpDown, variable=self.variable, **kwargs_TGLFrun, ) - self.readScan(label=label, variable=self.variable, positionIon=position) + self.read_scan(label=label, variable=self.variable, positionIon=position) - x = self.scans[label]["xV"] + x = self.scans[label]["scanned_variable"] yV = self.scans[label]["Gi"] self.variable_y = "Gi" @@ -2964,7 +2876,7 @@ def plotAnalysis(self, labels=["analysis1"], analysisType="chi_e", figs=None): variableLabel = "RLTS_2" elif analysisType == "Z": variableLabel = self.variable - self.plotScan( + self.plot_scan( labels=labels, figs=[fig2, fig2e, fig3], variableLabel=variableLabel ) @@ -3006,11 +2918,11 @@ def plotAnalysis(self, labels=["analysis1"], analysisType="chi_e", figs=None): ) ) - xV = self.scans[label]["xV"][irho] + xV = self.scans[label]["scanned_variable"][irho] Qe = self.scans[label][self.scans[label]["var_y"]][irho] xgrid = self.scans[label]["x_grid"][irho] ygrid = np.array(self.scans[label]["y_grid"][irho]) - xba = self.scans[label]["xV"][irho][ + xba = self.scans[label]["scanned_variable"][irho][ self.scans[label]["positionBase"] ] yba = np.interp(xba, xV, Qe) @@ -3158,7 +3070,7 @@ def plotAnalysis(self, labels=["analysis1"], analysisType="chi_e", figs=None): ax = ax00 ax.plot( - self.scans[label]["xV"][irho], + self.scans[label]["scanned_variable"][irho], np.array(self.scans[label]["Gi"][irho]), "o-", c=col, @@ -3171,7 +3083,7 @@ def plotAnalysis(self, labels=["analysis1"], analysisType="chi_e", figs=None): lw=0.5, c=col, ) - # ax.axvline(x=self.scans[label]['xV'][irho][self.scans[label]['positionBase']],ls='--',c=col,lw=1.) + # ax.axvline(x=self.scans[label]['scanned_variable'][irho][self.scans[label]['positionBase']],ls='--',c=col,lw=1.) cont += 1 @@ -3965,7 +3877,7 @@ def readTGLFresults( TGLFstd_TGLFout.append(TGLFout) inputclasses.append(TGLFout.inputclass) - parse = GACODErun.buildDictFromInput(TGLFout.inputFileTGLF) + parse = GACODErun.buildDictFromInput(TGLFout.inputFile) parsed.append(parse) results = { @@ -3997,7 +3909,7 @@ def __init__(self, FolderGACODE, suffix="",require_all_files=True): def postprocess(self): coeff, klow = 0.0, 0.8 - self.etas = processGrowthRates( + etas = processGrowthRates( self.ky, self.g[0, :], self.f[0, :], @@ -4007,6 +3919,16 @@ def postprocess(self): coeff=coeff, ) + self.eta_ITGETG = etas["metrics"]["eta_ITGETG"] + self.eta_ITGTEM = etas["metrics"]["eta_ITGTEM"] + self.g_lowk_max = etas["metrics"]["g_lowk_max"] + self.f_lowk_max = etas["metrics"]["f_lowk_max"] + self.k_lowk_max = etas["metrics"]["k_lowk_max"] + + self.g_ITG_max = etas["ITG"]["g_max"] + self.g_TEM_max = etas["TEM"]["g_max"] + self.g_ETG_max = etas["ETG"]["g_max"] + self.QeES = np.sum(self.SumFlux_Qe_phi) self.QeEM = np.sum(self.SumFlux_Qe_a) self.QiES = np.sum(self.SumFlux_Qi_phi) @@ -4512,7 +4434,7 @@ def read(self,require_all_files=True): with open(self.FolderGACODE / ("input.tglf" + self.suffix), "r") as fi: lines = fi.readlines() - self.inputFileTGLF = "".join(lines) + self.inputFile = "".join(lines) def unnormalize(self, normalization, rho=None, convolution_fun_fluct=None, factorTot_to_Perp=1.0): if normalization is not None: diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index ec2def7f..ae1aaea6 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -374,7 +374,7 @@ def run_scan( for cont_mult, mult in enumerate(varUpDown_new): name = f"{variable}_{mult}" self.read( - label=f"{self.subFolder_scan}_{name}", + label=f"{self.subfolder_scan}_{name}", folder=folders[cont_mult], cold_startWF = False, require_all_files=not kwargs_run.get("only_minimal_files",False), @@ -404,7 +404,7 @@ def _prepare_scan( multipliers_mod = copy.deepcopy(multipliers) - self.subFolder_scan = subfolder + self.subfolder_scan = subfolder if relativeChanges: for i in range(len(varUpDown)): @@ -448,7 +448,7 @@ def _prepare_scan( kwargs_run["forceIfcold_start"] = cont_mult > 0 or ("forceIfcold_start" in kwargs_run and kwargs_run["forceIfcold_start"]) code_executor, code_executor_full = self._run_prepare( - f"{self.subFolder_scan}_{name}", + f"{self.subfolder_scan}_{name}", code_executor=code_executor, code_executor_full=code_executor_full, multipliers=multipliers_mod, @@ -460,6 +460,91 @@ def _prepare_scan( return code_executor, code_executor_full, folders, varUpDown + def read_scan( + self, + label="scan1", + subfolder=None, + variable="RLTS_1", + positionIon=2, + variable_mapping=None, + variable_mapping_unn=None + ): + ''' + positionIon is the index in the input.tglf file... so if you want for ion RLNS_5, positionIon=5 + ''' + + if subfolder is None: + subfolder = self.subfolder_scan + + self.scans[label] = {} + self.scans[label]["variable"] = variable + self.scans[label]["positionBase"] = None + self.scans[label]["unnormalization_successful"] = True + self.scans[label]["results_tags"] = [] + + self.positionIon_scan = positionIon + + # ---- + + scan = {} + for ikey in variable_mapping | variable_mapping_unn: + scan[ikey] = [] + + cont = 0 + for ikey in self.results: + isThisTheRightReadResults = (subfolder in ikey) and (variable== "_".join(ikey.split("_")[:-1]).split(subfolder + "_")[-1]) + + if isThisTheRightReadResults: + + self.scans[label]["results_tags"].append(ikey) + + # Initialize lists + scan0 = {} + for ikey2 in variable_mapping | variable_mapping_unn: + scan0[ikey2] = [] + + # Loop over radii + for irho_cont in range(len(self.rhos)): + irho = np.where(self.results[ikey]["x"] == self.rhos[irho_cont])[0][0] + + for ikey2 in variable_mapping: + + obj = self.results[ikey][variable_mapping[ikey2][0]][irho] + if not hasattr(obj, '__dict__'): + obj_dict = obj + else: + obj_dict = obj.__dict__ + var0 = obj_dict[variable_mapping[ikey2][1]] + scan0[ikey2].append(var0 if variable_mapping[ikey2][2] is None else var0[variable_mapping[ikey2][2]]) + + # Unnormalized + self.scans[label]["unnormalization_successful"] = True + for ikey2 in variable_mapping_unn: + obj = self.results[ikey][variable_mapping_unn[ikey2][0]][irho] + if not hasattr(obj, '__dict__'): + obj_dict = obj + else: + obj_dict = obj.__dict__ + + if variable_mapping_unn[ikey2][1] not in obj_dict: + self.scans[label]["unnormalization_successful"] = False + break + var0 = obj_dict[variable_mapping_unn[ikey2][1]] + scan0[ikey2].append(var0 if variable_mapping_unn[ikey2][2] is None else var0[variable_mapping_unn[ikey2][2]]) + + for ikey2 in variable_mapping | variable_mapping_unn: + scan[ikey2].append(scan0[ikey2]) + + if float(ikey.split('_')[-1]) == 1.0: + self.scans[label]["positionBase"] = cont + cont += 1 + + self.scans[label]["x"] = np.array(self.rhos) + + for ikey2 in variable_mapping | variable_mapping_unn: + self.scans[label][ikey2] = np.atleast_2d(np.transpose(scan[ikey2])) + + def change_and_write_code( rhos, diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 38aa7e74..8f59565c 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -25,8 +25,11 @@ neo.run('neo3/', cold_start=cold_start, extraOptions={'N_ENERGY':5,'N_XI': 11, 'N_THETA': 11}, multipliers={'DLNTDR_1': 1.5}) neo.read('NEO low res + 50% aLTe') -# neo.run_scan('scan1', cold_start=cold_start, variable='DLNTDR_1', varUpDown=np.linspace(0.5, 1.5, 4)) - neo.plot(labels=['NEO default', 'NEO low res', 'NEO low res + 50% aLTe']) + +neo.run_scan('scan1', cold_start=cold_start, variable='DLNTDR_1', varUpDown=np.linspace(0.5, 1.5, 4)) +neo.read_scan(label='scan1',variable = 'DLNTDR_1') +neo.plot_scan(labels=['scan1'], fn = neo.fn) + neo.fn.show() -neo.fn.close() \ No newline at end of file +neo.fn.close() diff --git a/tests/TGLFscan_workflow.py b/tests/TGLFscan_workflow.py index 47e1c99d..0e93fd90 100644 --- a/tests/TGLFscan_workflow.py +++ b/tests/TGLFscan_workflow.py @@ -3,7 +3,7 @@ from mitim_tools.gacode_tools import TGLFtools, PROFILEStools from mitim_tools import __mitimroot__ -cold_start = False +cold_start = True (__mitimroot__ / 'tests' / 'scratch').mkdir(parents=True, exist_ok=True) @@ -16,20 +16,20 @@ tglf = TGLFtools.TGLF(rhos=[0.5, 0.7]) tglf.prep(input_gacode,folder, cold_start=cold_start) -tglf.runScan( subFolderTGLF = 'scan1', +tglf.run_scan( subfolder = 'scan1', TGLFsettings = None, cold_start = cold_start, runWaveForms = [0.67, 10.0], variable = 'RLTS_1', varUpDown = np.linspace(0.5,1.5,4)) -tglf.readScan(label='scan1',variable = 'RLTS_1') +tglf.read_scan(label='scan1',variable = 'RLTS_1') -tglf.plotScan(labels=['scan1']) +tglf.plot_scan(labels=['scan1']) tglf.fn.show() tglf.fn.close() tglf.runScanTurbulenceDrives( - subFolderTGLF = 'turb_drives', + subfolder = 'turb_drives', TGLFsettings = None, resolutionPoints=3, cold_start = cold_start) From e0c5d6e1c6cde98dd45437a57bfcb1483eb44cbb Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 22 Aug 2025 11:04:08 -0400 Subject: [PATCH 181/385] subFolderTGLF -> subfolder --- docs/capabilities/tglf_capabilities.rst | 4 ++-- docs/capabilities/vitals_capabilities.rst | 2 +- src/mitim_modules/portals/scripts/runTGLF.py | 4 ++-- .../portals/scripts/runTGLFdrivesfromPORTALS.py | 2 +- src/mitim_modules/portals/utils/PORTALSanalysis.py | 2 +- .../powertorch/physics_models/transport_tglfneo.py | 2 +- .../powertorch/physics_models/transport_tgyro.py | 4 ++-- src/mitim_modules/vitals/VITALSmain.py | 2 +- src/mitim_tools/gacode_tools/TGYROtools.py | 4 ++-- src/mitim_tools/gacode_tools/scripts/run_tglf.py | 6 +++--- src/mitim_tools/transp_tools/CDFtools.py | 10 +++++----- tests/TGLFfull_workflow.py | 6 +++--- tests/VITALS_workflow.py | 2 +- tutorials/TGLF_tutorial.py | 14 +++++++------- 14 files changed, 32 insertions(+), 32 deletions(-) diff --git a/docs/capabilities/tglf_capabilities.rst b/docs/capabilities/tglf_capabilities.rst index cfc34614..4cdb2f76 100644 --- a/docs/capabilities/tglf_capabilities.rst +++ b/docs/capabilities/tglf_capabilities.rst @@ -83,14 +83,14 @@ For example, the following two commands will run TGLF with saturation rule numbe .. code-block:: python - tglf.run( subFolderTGLF = 'yes_em_folder', + tglf.run( subfolder = 'yes_em_folder', TGLFsettings = 5, extraOptions = {}, cold_start = False ) tglf.read( label = 'yes_em' ) - tglf.run( subFolderTGLF = 'no_em_folder', + tglf.run( subfolder = 'no_em_folder', TGLFsettings = 5, extraOptions = {'USE_BPER':False}, cold_start = False ) diff --git a/docs/capabilities/vitals_capabilities.rst b/docs/capabilities/vitals_capabilities.rst index 5c21189f..7bf157d5 100644 --- a/docs/capabilities/vitals_capabilities.rst +++ b/docs/capabilities/vitals_capabilities.rst @@ -39,7 +39,7 @@ As a starting point of VITALS, you need to prepare and run TGLF for the base cas tglf = TGLFtools.TGLF( rhos = [ rho ] ) cdf = tglf.prep( folder, inputgacode = inputgacode_file) - tglf.run( subFolderTGLF = 'run_base', TGLFsettings = 5) + tglf.run( subfolder = 'run_base', TGLFsettings = 5) tglf.read( label = 'run_base' ) diff --git a/src/mitim_modules/portals/scripts/runTGLF.py b/src/mitim_modules/portals/scripts/runTGLF.py index 1cf8ea20..f92a7421 100644 --- a/src/mitim_modules/portals/scripts/runTGLF.py +++ b/src/mitim_modules/portals/scripts/runTGLF.py @@ -51,7 +51,7 @@ labels = [] for param in params: tglf.runScan( - subFolderTGLF="scan", + subfolder="scan", variable=param, varUpDown=varUpDown, TGLFsettings=TGLFsettings, @@ -69,7 +69,7 @@ else: tglf.runScanTurbulenceDrives( - subFolderTGLF="turb", + subfolder="turb", resolutionPoints=5, variation=var, variablesDrives=["RLTS_1", "RLTS_2", "RLNS_1", "XNUE", "TAUS_2", "BETAE"], diff --git a/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py b/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py index dffd6e9b..a0812ed8 100644 --- a/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py +++ b/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py @@ -41,7 +41,7 @@ tglf, TGLFsettings, extraOptions = portals.extractTGLF(positions=pos, evaluation=ev, modified_profiles=True, cold_start=cold_start) tglf.runScanTurbulenceDrives( - subFolderTGLF="turb", + subfolder="turb", resolutionPoints=num, variation=var, variablesDrives=["RLTS_1", "RLTS_2", "RLNS_1", "XNUE", "TAUS_2", "BETAE"], diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 64eb1efd..4039805d 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -639,7 +639,7 @@ def runTGLFfull( if "extraOptions" not in kwargsTGLF_this: kwargsTGLF_this["extraOptions"] = extraOptions - tglf.run(subFolderTGLF=f"tglf_{label}", cold_start=cold_start, **kwargsTGLF_this) + tglf.run(subfolder=f"tglf_{label}", cold_start=cold_start, **kwargsTGLF_this) # Read all previously run cases into a single class if tglf_object is None: diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 7dabee02..998b63bc 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -348,7 +348,7 @@ def _run_tglf_uncertainty_model( minutes = max(2, minutes) tglf.runScanTurbulenceDrives( - subFolderTGLF = name, + subfolder = name, variablesDrives = variables_to_scan, varUpDown = relative_scan, minimum_delta_abs = minimum_delta_abs, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 655be8b9..d5b4f038 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -121,7 +121,7 @@ def _evaluate_tglf_neo(self): # cold_start=cold_start) # tglf.run( - # subFolderTGLF="tglf_neo_original", + # subfolder="tglf_neo_original", # TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], # cold_start=cold_start, # forceIfcold_start=True, @@ -250,7 +250,7 @@ def tglf_scan_trick( minutes = max(2, minutes) tglf.runScanTurbulenceDrives( - subFolderTGLF = name, + subfolder = name, variablesDrives = variables_to_scan, varUpDown = relative_scan, minimum_delta_abs = minimum_delta_abs, diff --git a/src/mitim_modules/vitals/VITALSmain.py b/src/mitim_modules/vitals/VITALSmain.py index 998547ec..11fc6fe1 100644 --- a/src/mitim_modules/vitals/VITALSmain.py +++ b/src/mitim_modules/vitals/VITALSmain.py @@ -302,7 +302,7 @@ def runTGLF( folder_label = label tglf.run( - subFolderTGLF=f"{folder_label}", + subfolder=f"{folder_label}", cold_start=cold_start, TGLFsettings=self.TGLFparameters["TGLFsettings"], forceIfcold_start=True, diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 86e9dd6a..48a756b2 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -660,7 +660,7 @@ def runTGLF(self, fromlabel="tgyro1", rhos=None, cold_start=False): label = f"{self.nameRuns_default}_tglf1" self.tglf[fromlabel].run( - subFolderTGLF=f"{label}", + subfolder=f"{label}", TGLFsettings=None, ApplyCorrections=False, cold_start=cold_start, @@ -689,7 +689,7 @@ def runTGLFsensitivities(self, fromlabel="tgyro1", rho=0.5, cold_start=False): ) self.tglf[fromlabel].runScanTurbulenceDrives( - subFolderTGLF=f"{self.nameRuns_default}_tglf", + subfolder=f"{self.nameRuns_default}_tglf", TGLFsettings=None, ApplyCorrections=False, cold_start=cold_start, diff --git a/src/mitim_tools/gacode_tools/scripts/run_tglf.py b/src/mitim_tools/gacode_tools/scripts/run_tglf.py index 49bc6c3f..a2dc3869 100644 --- a/src/mitim_tools/gacode_tools/scripts/run_tglf.py +++ b/src/mitim_tools/gacode_tools/scripts/run_tglf.py @@ -52,12 +52,12 @@ def main(): # ------------------------------------------------------------------------------ if drives: - tglf.runScanTurbulenceDrives(subFolderTGLF="scan_turb", TGLFsettings=None) + tglf.runScanTurbulenceDrives(subfolder="scan_turb", TGLFsettings=None) tglf.plotScanTurbulenceDrives(label="scan_turb") elif scan is not None: tglf.runScan( - subFolderTGLF="scan1", + subfolder="scan1", variable=scan, varUpDown=np.linspace(0.2, 2.0, 5), TGLFsettings=None, @@ -67,7 +67,7 @@ def main(): tglf.plotScan(labels=["scan1"], variableLabel=scan) else: - tglf.run(subFolderTGLF="run1", TGLFsettings=None, cold_start=cold_start) + tglf.run(subfolder="run1", TGLFsettings=None, cold_start=cold_start) tglf.read(label="run1") tglf.plot(labels=["run1"]) diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index af3f97a0..02910ec4 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -14086,7 +14086,7 @@ def runTGLFstandalone( labelTGLF = kwargs_TGLFrun.get("label", "tglf1") self.TGLFstd[nameF].run( - subFolderTGLF=labelTGLF, + subfolder=labelTGLF, forceIfcold_start=forceIfcold_start, **kwargs_TGLFrun, ) @@ -14144,7 +14144,7 @@ def transportAnalysis( if typeAnalysis == "CHIPERT": self.TGLFstd[int(time * 1000)].runAnalysis( - subFolderTGLF="chi_per", + subfolder="chi_per", label="chi_pert", analysisType="e", TGLFsettings=TGLFsettings, @@ -14157,7 +14157,7 @@ def transportAnalysis( addTrace = [40, 173] self.TGLFstd[int(time * 1000)].runAnalysis( - subFolderTGLF="impurity", + subfolder="impurity", label="impurity", analysisType="Z", TGLFsettings=TGLFsettings, @@ -14173,7 +14173,7 @@ def transportAnalysis( if "FLUC" in typeAnalysis: self.TGLFstd[int(time * 1000)].run( - subFolderTGLF="fluctuations", + subfolder="fluctuations", TGLFsettings=TGLFsettings, forceIfcold_start=True, ) @@ -14697,7 +14697,7 @@ def compareChiPert( self.ChiPert_tglf = TGLFtools.TGLF(cdf=self.LocationCDF, time=time, rhos=rhos) self.ChiPert_tglf.prep(self.FolderCDF / "chi_per_calc", cold_start=cold_start) self.ChiPert_tglf.runAnalysis( - subFolderTGLF="chi_per", + subfolder="chi_per", label="chi_pert", analysisType="e", TGLFsettings=TGLFsettings, diff --git a/tests/TGLFfull_workflow.py b/tests/TGLFfull_workflow.py index 031d077f..3a90abd3 100644 --- a/tests/TGLFfull_workflow.py +++ b/tests/TGLFfull_workflow.py @@ -21,7 +21,7 @@ _ = tglf.prep_using_tgyro(folder, cold_start=cold_start) tglf.run( - subFolderTGLF="runSAT2", + subfolder="runSAT2", TGLFsettings=5, runWaveForms=[0.1,0.3], cold_start=cold_start, @@ -30,7 +30,7 @@ tglf.read(label="runSAT2", d_perp_cm={0.6: 0.5, 0.8: 0.5}) tglf.run( - subFolderTGLF="runSAT0", + subfolder="runSAT0", TGLFsettings=2, runWaveForms=[0.5], cold_start=cold_start, @@ -39,7 +39,7 @@ tglf.read(label="runSAT0", d_perp_cm={0.6: 0.5, 0.8: 0.5}) tglf.run( - subFolderTGLF="runSAT3", + subfolder="runSAT3", TGLFsettings=6, runWaveForms=[0.5], cold_start=cold_start, diff --git a/tests/VITALS_workflow.py b/tests/VITALS_workflow.py index f1c591dc..cd6f3a29 100644 --- a/tests/VITALS_workflow.py +++ b/tests/VITALS_workflow.py @@ -31,7 +31,7 @@ tglf = TGLFtools.TGLF(rhos=[rho]) cdf = tglf.prep_using_tgyro(folderWork, cold_start=cold_start, inputgacode=inputgacode) -tglf.run(subFolderTGLF="run_base/", TGLFsettings=TGLFsettings, cold_start=cold_start) +tglf.run(subfolder="run_base/", TGLFsettings=TGLFsettings, cold_start=cold_start) # ******************************************************************************** # Then, add experimental data of fluctuation information and error bars diff --git a/tutorials/TGLF_tutorial.py b/tutorials/TGLF_tutorial.py index 41adba28..12601fb6 100644 --- a/tutorials/TGLF_tutorial.py +++ b/tutorials/TGLF_tutorial.py @@ -20,7 +20,7 @@ # Run TGLF in subfolder tglf.run( - subFolderTGLF="yes_em_folder", + subfolder="yes_em_folder", TGLFsettings=5, extraOptions={}, cold_start=False @@ -31,7 +31,7 @@ # Run TGLF in a different subfolder with different settings tglf.run( - subFolderTGLF="no_em_folder", + subfolder="no_em_folder", TGLFsettings=5, extraOptions={"USE_BPER": False}, cold_start=False, @@ -49,7 +49,7 @@ *************************************************************************** ''' -tglf.runScan( subFolderTGLF = 'scan1', +tglf.runScan( subfolder = 'scan1', TGLFsettings = 5, cold_start = False, variable = 'RLTS_1', @@ -57,7 +57,7 @@ tglf.readScan(label='scan1',variable = 'RLTS_1') -tglf.runScan( subFolderTGLF = 'scan2', +tglf.runScan( subfolder = 'scan2', TGLFsettings = 5, cold_start = False, variable = 'RLTS_2', @@ -74,7 +74,7 @@ ''' tglf.runScanTurbulenceDrives( - subFolderTGLF = 'turb_drives', + subfolder = 'turb_drives', TGLFsettings = 5, cold_start = False) @@ -87,7 +87,7 @@ ''' tglf.runAnalysis( - subFolderTGLF = 'chi_e', + subfolder = 'chi_e', analysisType = 'chi_e', TGLFsettings = 5, cold_start = False, @@ -103,7 +103,7 @@ for i in[1,2,3,4,5,6]: tglf.run( - subFolderTGLF = f'settings{i}', + subfolder = f'settings{i}', runWaveForms = [0.67], TGLFsettings = i, cold_start = False) From 6ed6944c2c30068157e840da243b1a9b02995b13 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 22 Aug 2025 13:12:43 -0400 Subject: [PATCH 182/385] Changed default to use standalone TGLF and NEO in PORTALS --- src/mitim_modules/portals/PORTALSmain.py | 18 ++++++++++-------- .../physics_models/transport_tgyro.py | 14 +++++++------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index d8acafb3..3be47738 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -13,7 +13,6 @@ PORTALSoptimization, PORTALSanalysis, ) -from mitim_modules.powertorch.physics_models import targets_analytic, transport_tgyro, transport_cgyro from mitim_tools.opt_tools import STRATEGYtools from mitim_tools.opt_tools.utils import BOgraphics from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -193,11 +192,12 @@ def __init__( # Selection of model if CGYROrun: - transport_evaluator = transport_cgyro.cgyro_model + from mitim_modules.powertorch.physics_models.transport_cgyro import cgyro_model as transport_evaluator else: - transport_evaluator = transport_tgyro.tgyro_model - - target_evaluator = targets_analytic.analytical_model + # from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model as transport_evaluator + from mitim_modules.powertorch.physics_models.transport_tglfneo import tglfneo_model as transport_evaluator + + from mitim_modules.powertorch.physics_models.targets_analytic import analytical_model as target_evaluator self.PORTALSparameters = { "percentError": [5,10,1], # (%) Error (std, in percent) of model evaluation [TGLF (treated as minimum if scan trick), NEO, TARGET] @@ -467,7 +467,8 @@ def check_flags(self): print("\t- Requested fineTargetsResolution, so running powerstate target calculations",typeMsg="w") self.PORTALSparameters["target_evaluator_method"] = "powerstate" - if not issubclass(self.PORTALSparameters["transport_evaluator"], transport_tgyro.tgyro_model) and (self.PORTALSparameters["target_evaluator_method"] == "tgyro"): + from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model + if not issubclass(self.PORTALSparameters["transport_evaluator"], tgyro_model) and (self.PORTALSparameters["target_evaluator_method"] == "tgyro"): print("\t- Requested TGYRO targets, but transport evaluator is not tgyro, so changing to powerstate",typeMsg="w") self.PORTALSparameters["target_evaluator_method"] = "powerstate" @@ -482,7 +483,7 @@ def check_flags(self): if self.PORTALSparameters["target_evaluator_method"] == "tgyro" and self.PORTALSparameters['profiles_postprocessing_fun'] is not None: print("\t- Requested custom modification of postprocessing function but targets from TGYRO... are you sure?",typeMsg="q") - if self.PORTALSparameters["target_evaluator_method"] == "tgyro" and self.PORTALSparameters['transport_evaluator'] != transport_tgyro.tgyro_model: + if self.PORTALSparameters["target_evaluator_method"] == "tgyro" and self.PORTALSparameters['transport_evaluator'] != tgyro_model: print("\t- Requested TGYRO targets but transport evaluator is not TGYRO... are you sure?",typeMsg="q") key_rhos = "RoaLocations" if self.MODELparameters["RoaLocations"] is not None else "RhoLocations" @@ -554,7 +555,8 @@ def reuseTrainingTabular( self_copy.powerstate.transport_options["transport_evaluator"] = None self_copy.powerstate.target_options["target_evaluator_options"]["TypeTarget"] = "powerstate" else: - self_copy.powerstate.transport_options["transport_evaluator"] = transport_tgyro.tgyro_model + from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model + self_copy.powerstate.transport_options["transport_evaluator"] = tgyro_model _, dictOFs = runModelEvaluator( self_copy, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index d5b4f038..87eb8b60 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -999,15 +999,15 @@ def tgyro_to_powerstate(TGYROresults, # ********************************** # Store raw fluxes for better plotting later - powerstate.plasma["CZ_raw_tr_turb"] = torch.Tensor(TGYROresults.Gi_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) - powerstate.plasma["CZ_raw_tr_neoc"] = torch.Tensor(TGYROresults.Gi_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) + powerstate.plasma["GZ1E20m2_tr_turb"] = torch.Tensor(TGYROresults.Gi_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) + powerstate.plasma["GZ1E20m2_tr_neoc"] = torch.Tensor(TGYROresults.Gi_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) - powerstate.plasma["CZ_raw_tr_turb_stds"] = torch.Tensor(TGYROresults.Gi_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_raw_tr_neoc_stds"] = torch.Tensor(TGYROresults.Gi_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["GZ1E20m2_tr_turb_stds"] = torch.Tensor(TGYROresults.Gi_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["GZ1E20m2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Gi_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None if provideTargets: - powerstate.plasma["CZ_raw"] = torch.Tensor(TGYROresults.Gi_tar[impurityPosition, :, :nr]).to(powerstate.dfT) - powerstate.plasma["CZ_raw_stds"] = torch.Tensor(TGYROresults.Gi_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + powerstate.plasma["GZ1E20m2"] = torch.Tensor(TGYROresults.Gi_tar[impurityPosition, :, :nr]).to(powerstate.dfT) + powerstate.plasma["GZ1E20m2_stds"] = torch.Tensor(TGYROresults.Gi_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None powerstate.plasma["CZ_tr_turb"] = torch.Tensor(TGYROresults.Ci_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp powerstate.plasma["CZ_tr_neoc"] = torch.Tensor(TGYROresults.Ci_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp @@ -1042,7 +1042,7 @@ def tgyro_to_powerstate(TGYROresults, # Sum here turbulence and neoclassical, after modifications # ------------------------------------------------------------------------------------------------------------------------ - quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2', 'Ge1E20m2', 'CZ_raw'] + quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2', 'Ge1E20m2', 'GZ1E20m2'] for ikey in quantities: powerstate.plasma[ikey+"_tr"] = powerstate.plasma[ikey+"_tr_turb"] + powerstate.plasma[ikey+"_tr_neoc"] From c4af0041f7ae7194e8a8acb6646f7976e192be2e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 22 Aug 2025 14:33:04 -0400 Subject: [PATCH 183/385] Towards a tigher PORTALS-CGYRO coupling --- src/mitim_modules/portals/PORTALSmain.py | 1 + .../physics_models/transport_cgyro.py | 5 +- .../physics_models/transport_cgyroneo.py | 44 +++++++ .../physics_models/transport_tglfneo.py | 117 +++++------------- .../physics_models/transport_tgyro.py | 5 +- .../powertorch/utils/TRANSPORTtools.py | 117 ++++++++++++++++-- src/mitim_tools/gacode_tools/CGYROtools.py | 5 +- src/mitim_tools/gacode_tools/TGLFtools.py | 6 +- 8 files changed, 198 insertions(+), 102 deletions(-) create mode 100644 src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 3be47738..600ce99b 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -192,6 +192,7 @@ def __init__( # Selection of model if CGYROrun: + # from mitim_modules.powertorch.physics_models.transport_cgyroneo import cgyroneo_model as transport_evaluator from mitim_modules.powertorch.physics_models.transport_cgyro import cgyro_model as transport_evaluator else: # from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model as transport_evaluator diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index d8cb7fb2..c8f7ccfc 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -4,6 +4,7 @@ import numpy as np from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools, TGYROtools +from mitim_tools.plasmastate_tools import MITIMstate from mitim_modules.powertorch.physics_models import transport_tgyro from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -21,7 +22,7 @@ def evaluate(self): powerstate_orig = self._trick_cgyro(tgyro) # Process results - self._postprocess_results(tgyro, "cgyro_neo") + self._postprocess(tgyro, "cgyro_neo") # Some checks print("\t- Checking model modifications:") @@ -170,7 +171,7 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod Qi_criterion_stable = PORTALSparameters["Qi_criterion_stable"] try: - impurityPosition = PROFILEStools.impurity_location(PROFILEStools.gacode_state(unmodified_profiles), PORTALSparameters["ImpurityOfInterest"]) + impurityPosition = MITIMstate.impurity_location(PROFILEStools.gacode_state(unmodified_profiles), PORTALSparameters["ImpurityOfInterest"]) except ValueError: if 'nZ' in ProfilesPredicted: raise ValueError(f"Impurity {PORTALSparameters['ImpurityOfInterest']} not found in the profiles and needed for CGYRO evaluation") diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py new file mode 100644 index 00000000..6b99119c --- /dev/null +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -0,0 +1,44 @@ +from mitim_tools.gacode_tools import CGYROtools +from mitim_modules.powertorch.physics_models import transport_tglfneo +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +class cgyroneo_model(transport_tglfneo.tglfneo_model): + def __init__(self, powerstate, **kwargs): + super().__init__(powerstate, **kwargs) + + def evaluate_turbulence(self): + + pass + # # ------------------------------------------------------------------------------------------------------------------------ + # # Prepare CGYRO object + # # ------------------------------------------------------------------------------------------------------------------------ + + # rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] + + # tglf = TGLFtools.CGYRO(rhos=rho_locations) + + # _ = tglf.prep( + # self.powerstate.profiles_transport, + # self.folder, + # cold_start = cold_start, + # ) + + + # cgyro = CGYROtools.CGYRO() + + # cgyro.prep( + # self.folder, + # self.powerstate.profiles_transport.files[0], + # ) + + # cgyro.run( + # 'base_cgyro', + # roa = 0.55, + # CGYROsettings=0, + # submit_run=False + # ) + + # embed() + # #cgyro.read(label='base') + \ No newline at end of file diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 998b63bc..28ab1067 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -1,8 +1,7 @@ import shutil -import torch import numpy as np -from mitim_tools.misc_tools import IOtools, PLASMAtools +from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import TGLFtools, NEOtools from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -14,19 +13,12 @@ def __init__(self, powerstate, **kwargs): def produce_profiles(self): self._produce_profiles() - - def evaluate(self): - - tglf = self._evaluate_tglf() - neo = self._evaluate_neo() - - self._postprocess() # ************************************************************************************ # Private functions for the evaluation # ************************************************************************************ - def _evaluate_tglf(self): + def evaluate_turbulence(self): # ------------------------------------------------------------------------------------------------------------------------ # Grab options from powerstate @@ -68,11 +60,13 @@ def _evaluate_tglf(self): if use_tglf_scan_trick is None: + # ******************************************************************* # Just run TGLF once and apply an ad-hoc percent error to the results + # ******************************************************************* tglf.run( 'base_tglf', - TGLFsettings=TGLFsettings, + Settings=TGLFsettings, extraOptions=extraOptions, ApplyCorrections=False, launchSlurm= launchMODELviaSlurm, @@ -109,7 +103,9 @@ def _evaluate_tglf(self): else: + # ******************************************************************* # Run TGLF with scans to estimate the uncertainty + # ******************************************************************* Flux_base, Flux_mean, Flux_std = _run_tglf_uncertainty_model( tglf, @@ -125,26 +121,9 @@ def _evaluate_tglf(self): launchMODELviaSlurm=launchMODELviaSlurm, Qi_includes_fast=Qi_includes_fast, ) - - for i in range(len(tglf.profiles.Species)): - gacode_type = tglf.profiles.Species[i]['S'] - for rho in rho_locations: - tglf_type = tglf.inputs_files[0.25].ions_info[i+2]['type'] - - if gacode_type[:5] != tglf_type[:5]: - print(f"\t- For location {rho=:.2f}, ion specie #{i+1} ({tglf.profiles.Species[i]['N']}) is considered '{gacode_type}' by gacode but '{tglf_type}' by TGLF. Make sure this is consistent with your use case", typeMsg="w") - - if tglf_type == 'fast': - - if Qi_includes_fast: - print(f"\t\t\t* The fast ion considered by TGLF was summed into the Qi", typeMsg="i") - else: - print(f"\t\t\t* The fast ion considered by TGLF was NOT summed into the Qi", typeMsg="i") - - else: - - print(f"\t\t\t* The thermal ion considered by TGLF was summed into the Qi", typeMsg="i") - + + self._raise_warnings(tglf, rho_locations, Qi_includes_fast) + # ------------------------------------------------------------------------------------------------------------------------ # Pass the information to POWERSTATE # ------------------------------------------------------------------------------------------------------------------------ @@ -173,7 +152,7 @@ def _evaluate_tglf(self): return tglf - def _evaluate_neo(self): + def evaluate_neoclassical(self): # Options @@ -195,9 +174,7 @@ def _evaluate_neo(self): cold_start = cold_start, ) - neo.run( - 'base_neo', - ) + neo.run('base_neo') neo.read(label='base') @@ -223,56 +200,6 @@ def _evaluate_neo(self): self.powerstate.plasma["QieMWm3_tr_neoc_stds"] = Qe * 0.0 return neo - - def _postprocess(self): - - OriginalFimp = self.powerstate.transport_options["transport_evaluator_options"].get("OriginalFimp", 1.0) - - # ------------------------------------------------------------------------------------------------------------------------ - # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) - # ------------------------------------------------------------------------------------------------------------------------ - - variables = ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3'] - - for variable in variables: - for suffix in ['_tr_turb', '_tr_turb_stds', '_tr_neoc', '_tr_neoc_stds']: - - # Make them tensors and add a batch dimension - self.powerstate.plasma[f"{variable}{suffix}"] = torch.Tensor(self.powerstate.plasma[f"{variable}{suffix}"]).to(self.powerstate.dfT).unsqueeze(0) - - # Pad with zeros at rho=0.0 - self.powerstate.plasma[f"{variable}{suffix}"] = torch.cat(( - torch.zeros((1, 1)), - self.powerstate.plasma[f"{variable}{suffix}"], - ), dim=1) - - # ----------------------------------------------------------- - # Sum the turbulent and neoclassical contributions - # ----------------------------------------------------------- - - variables = ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2'] - - for variable in variables: - self.powerstate.plasma[f"{variable}_tr"] = self.powerstate.plasma[f"{variable}_tr_turb"] + self.powerstate.plasma[f"{variable}_tr_neoc"] - - # ----------------------------------------------------------- - # Convective fluxes (& Re-scale the GZ flux by the original impurity concentration) - # ----------------------------------------------------------- - - mapper_convective = { - 'Ce': 'Ge1E20m2', - 'CZ': 'GZ1E20m2', - } - - for key in mapper_convective.keys(): - for tt in ['','_turb', '_turb_stds', '_neoc', '_neoc_stds']: - - mult = 1.0 if key == 'Ce' else 1/OriginalFimp - - self.powerstate.plasma[f"{key}_tr{tt}"] = PLASMAtools.convective_flux( - self.powerstate.plasma["te"], - self.powerstate.plasma[f"{mapper_convective[key]}_tr{tt}"] - ) * mult def _profiles_to_store(self): @@ -288,6 +215,26 @@ def _profiles_to_store(self): else: print("\t- Could not move files", typeMsg="w") + def _raise_warnings(self, tglf, rho_locations, Qi_includes_fast): + + for i in range(len(tglf.profiles.Species)): + gacode_type = tglf.profiles.Species[i]['S'] + for rho in rho_locations: + tglf_type = tglf.inputs_files[0.25].ions_info[i+2]['type'] + + if gacode_type[:5] != tglf_type[:5]: + print(f"\t- For location {rho=:.2f}, ion specie #{i+1} ({tglf.profiles.Species[i]['N']}) is considered '{gacode_type}' by gacode but '{tglf_type}' by TGLF. Make sure this is consistent with your use case", typeMsg="w") + + if tglf_type == 'fast': + + if Qi_includes_fast: + print(f"\t\t\t* The fast ion considered by TGLF was summed into the Qi", typeMsg="i") + else: + print(f"\t\t\t* The fast ion considered by TGLF was NOT summed into the Qi", typeMsg="i") + + else: + print(f"\t\t\t* The thermal ion considered by TGLF was summed into the Qi", typeMsg="i") + def _run_tglf_uncertainty_model( tglf, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 87eb8b60..805f42b7 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -15,11 +15,12 @@ def __init__(self, powerstate, **kwargs): def produce_profiles(self): self._produce_profiles() + # TGYRO model is historical (#TODO #TOREMOVE) and therefore I'm not using the same evaluate as the rest, just keep it separate def evaluate(self): tgyro = self._evaluate_tglf_neo() - self._postprocess_results(tgyro, "tglf_neo") + self._postprocess(tgyro, "tglf_neo") # ************************************************************************************ # Private functions for TGLF and NEO evaluations @@ -137,7 +138,7 @@ def _evaluate_tglf_neo(self): return tgyro - def _postprocess_results(self, tgyro, label): + def _postprocess(self, tgyro, label): transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 270aeb7c..d6993872 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -1,6 +1,8 @@ +import torch +import numpy as np import copy import shutil -from mitim_tools.misc_tools import IOtools +from mitim_tools.misc_tools import IOtools, PLASMAtools from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -115,17 +117,116 @@ def _modify_profiles(self): print(f"\t- Impurity position has changed from {self.powerstate.impurityPosition} to {impurityPosition_new}",typeMsg="i") self.powerstate.impurityPosition_transport = p_new.Species.index(impurity_of_interest) + def evaluate(self): + + neoclassical = self.evaluate_neoclassical() + turbulence = self.evaluate_turbulence() + + self._postprocess() + + def _postprocess(self): + + OriginalFimp = self.powerstate.transport_options["transport_evaluator_options"].get("OriginalFimp", 1.0) + + # ------------------------------------------------------------------------------------------------------------------------ + # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) + # ------------------------------------------------------------------------------------------------------------------------ + + variables = ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3'] + + for variable in variables: + for suffix in ['_tr_turb', '_tr_turb_stds', '_tr_neoc', '_tr_neoc_stds']: + + # Make them tensors and add a batch dimension + self.powerstate.plasma[f"{variable}{suffix}"] = torch.Tensor(self.powerstate.plasma[f"{variable}{suffix}"]).to(self.powerstate.dfT).unsqueeze(0) + + # Pad with zeros at rho=0.0 + self.powerstate.plasma[f"{variable}{suffix}"] = torch.cat(( + torch.zeros((1, 1)), + self.powerstate.plasma[f"{variable}{suffix}"], + ), dim=1) + + # ----------------------------------------------------------- + # Sum the turbulent and neoclassical contributions + # ----------------------------------------------------------- + + variables = ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2'] + + for variable in variables: + self.powerstate.plasma[f"{variable}_tr"] = self.powerstate.plasma[f"{variable}_tr_turb"] + self.powerstate.plasma[f"{variable}_tr_neoc"] + + # ----------------------------------------------------------- + # Convective fluxes (& Re-scale the GZ flux by the original impurity concentration) + # ----------------------------------------------------------- + + mapper_convective = { + 'Ce': 'Ge1E20m2', + 'CZ': 'GZ1E20m2', + } + + for key in mapper_convective.keys(): + for tt in ['','_turb', '_turb_stds', '_neoc', '_neoc_stds']: + + mult = 1.0 if key == 'Ce' else 1/OriginalFimp + + self.powerstate.plasma[f"{key}_tr{tt}"] = PLASMAtools.convective_flux( + self.powerstate.plasma["te"], + self.powerstate.plasma[f"{mapper_convective[key]}_tr{tt}"] + ) * mult + # ---------------------------------------------------------------------------------------------------- # EVALUATE (custom part) # ---------------------------------------------------------------------------------------------------- - def evaluate(self): + def evaluate_turbulence(self): ''' - This needs to populate the following in self.powerstate.plasma - - QeMWm2, QeMWm2_tr, QeMWm2_tr_turb, QeMWm2_tr_neoc - Same for QiMWm2, Ce, CZ, MtJm2 - and their respective standard deviations + This needs to populate the following np.arrays in self.powerstate.plasma, with dimensions of rho: + - QeMWm2_tr_turb + - QiMWm2_tr_turb + - Ge1E20m2_tr_turb + - GZ1E20m2_tr_turb + - MtJm2_tr_turb + - QieMWm3_tr_turb (turbulence exchange) + and their respective standard deviations, e.g. QeMWm2_tr_turb_stds ''' - print(">> No transport fluxes to evaluate", typeMsg="w") - pass + print(">> No turbulent fluxes to evaluate", typeMsg="w") + + dim = self.powerstate.plasma['rho'].shape[-1]-1 + + for var in [ + 'QeMWm2', + 'QiMWm2', + 'Ge1E20m2', + 'GZ1E20m2', + 'MtJm2', + 'QieMWm3' + ]: + + self.powerstate.plasma[f"{var}_tr_turb"] = np.zeros(dim) + self.powerstate.plasma[f"{var}_tr_turb_stds"] = np.zeros(dim) + + def evaluate_neoclassical(self): + ''' + This needs to populate the following np.arrays in self.powerstate.plasma: + - QeMWm2_tr_neoc + - QiMWm2_tr_neoc + - Ge1E20m2_tr_neoc + - GZ1E20m2_tr_neoc + - MtJm2_tr_neoc + and their respective standard deviations, e.g. QeMWm2_tr_neoc_stds + ''' + print(">> No neoclassical fluxes to evaluate", typeMsg="w") + + dim = self.powerstate.plasma['rho'].shape[-1]-1 + + for var in [ + 'QeMWm2', + 'QiMWm2', + 'Ge1E20m2', + 'GZ1E20m2', + 'MtJm2', + ]: + + self.powerstate.plasma[f"{var}_tr_neoc"] = np.zeros(dim) + self.powerstate.plasma[f"{var}_tr_neoc_stds"] = np.zeros(dim) \ No newline at end of file diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 95877f4a..aea4ca36 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -71,7 +71,8 @@ def prep(self, folder, inputgacode_file): self.folder.mkdir(parents=True, exist_ok=True) self.inputgacode_file = self.folder / "input.gacode" - shutil.copy2(IOtools.expandPath(inputgacode_file), self.inputgacode_file) + if IOtools.expandPath(inputgacode_file) != self.inputgacode_file: + shutil.copy2(IOtools.expandPath(inputgacode_file), self.inputgacode_file) def _prerun( self, @@ -105,7 +106,7 @@ def _prerun( multipliers=multipliers, addControlFunction=GACODEdefaults.addCGYROcontrol, rmin=roa, - control_file = 'input.cgyro.controls' + controls_file = 'input.cgyro.controls' ) inputCGYRO.writeCurrentStatus() diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 54392f47..f70149af 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -54,7 +54,7 @@ def __init__( dictionaries. Plotting then can happen with more than one label of the same category. *Note* - The 'run' command uses input.tglf from the specified folder, but one can change the TGLFsettings presets, + The 'run' command uses input.tglf from the specified folder, but one can change the Settings presets, extraOptions and multipliers. The modified inputs is not rewritten in the actual folder, it is only written in the tmp folder on which the simulation takes place. @@ -75,7 +75,7 @@ def __init__( # Run standalone TGLF (this will find the input.tglf in the previous folder, # and then copy to this specify TGLF run, and run it there) - tglf.run(subfolder='tglf1/',TGLFsettings=1,extraOptions={'NS':3}) + tglf.run(subfolder='tglf1/',Settings=1,extraOptions={'NS':3}) # Read results tglf.read(label='run1',folder='~/testTGLF/tglf1/') @@ -94,7 +94,7 @@ def __init__( cdf = tglf.prep_using_tgyro('~/testTGLF/') # Run - tglf.run_scan('scan1/',TGLFsettings=1,varUpDown=np.linspace(0.5,2.0,20),variable='RLTS_2') + tglf.run_scan('scan1/',Settings=1,varUpDown=np.linspace(0.5,2.0,20),variable='RLTS_2') # Read scan tglf.read_scan(label='scan1',variable='RLTS_2') From ccab4e7842948eebc67e621716661679a516fe5e Mon Sep 17 00:00:00 2001 From: pabloprf Date: Fri, 22 Aug 2025 16:51:20 -0400 Subject: [PATCH 184/385] bug fix, profiles_postprocessing_fun must return the object --- src/mitim_modules/maestro/scripts/run_maestro.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index ab53fff4..eb37f077 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -145,6 +145,7 @@ def profiles_postprocessing_fun(file_profs): if enforce_same_density_gradients: p.enforce_same_density_gradients() p.write_state(file=file_profs) + return p beat_namelist['PORTALSparameters']['profiles_postprocessing_fun'] = profiles_postprocessing_fun else: From 64e427dc76f88e376bfb8f3066c5c88c3c6b8d03 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 22 Aug 2025 17:10:33 -0400 Subject: [PATCH 185/385] Added secondary criterion for transp finished --- src/mitim_tools/transp_tools/src/TRANSPsingularity.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py index 6ebbabc5..9fcdc89e 100644 --- a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py +++ b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py @@ -472,7 +472,7 @@ def interpretRun(infoSLURM, log_file): Case is not running (finished or failed) """ - if "TERMINATE THE RUN (NORMAL EXIT)" in "\n".join(log_file): + if "TERMINATE THE RUN (NORMAL EXIT)" in "\n".join(log_file) or "Finished TRANSP run app." in "\n".join(log_file): status = 1 info["info"]["status"] = "finished" elif ("Error termination" in "\n".join(log_file)) or ( @@ -489,10 +489,7 @@ def interpretRun(infoSLURM, log_file): status = -1 info["info"]["status"] = "stopped" else: - print( - "\t- No error nor termination found, assuming it is still running", - typeMsg="w", - ) + print("\t- No error nor termination found, assuming it is still running",typeMsg="w",) pringLogTail(log_file, typeMsg="i") status = 0 info["info"]["status"] = "running" From 2b5e88477259fb05338c96204886e17df5194b49 Mon Sep 17 00:00:00 2001 From: pabloprf Date: Fri, 22 Aug 2025 17:35:31 -0400 Subject: [PATCH 186/385] Bug fix catch type tglf ion --- .../powertorch/physics_models/transport_tglfneo.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 998b63bc..b41235d2 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -129,7 +129,7 @@ def _evaluate_tglf(self): for i in range(len(tglf.profiles.Species)): gacode_type = tglf.profiles.Species[i]['S'] for rho in rho_locations: - tglf_type = tglf.inputs_files[0.25].ions_info[i+2]['type'] + tglf_type = tglf.inputs_files[rho].ions_info[i+2]['type'] if gacode_type[:5] != tglf_type[:5]: print(f"\t- For location {rho=:.2f}, ion specie #{i+1} ({tglf.profiles.Species[i]['N']}) is considered '{gacode_type}' by gacode but '{tglf_type}' by TGLF. Make sure this is consistent with your use case", typeMsg="w") @@ -137,13 +137,13 @@ def _evaluate_tglf(self): if tglf_type == 'fast': if Qi_includes_fast: - print(f"\t\t\t* The fast ion considered by TGLF was summed into the Qi", typeMsg="i") + print("\t\t\t* The fast ion considered by TGLF was summed into the Qi", typeMsg="i") else: - print(f"\t\t\t* The fast ion considered by TGLF was NOT summed into the Qi", typeMsg="i") + print("\t\t\t* The fast ion considered by TGLF was NOT summed into the Qi", typeMsg="i") else: - print(f"\t\t\t* The thermal ion considered by TGLF was summed into the Qi", typeMsg="i") + print("\t\t\t* The thermal ion considered by TGLF was summed into the Qi", typeMsg="i") # ------------------------------------------------------------------------------------------------------------------------ # Pass the information to POWERSTATE From c862af16fc5c85eae1bf7e005d385954ff214409 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 22 Aug 2025 18:07:57 -0400 Subject: [PATCH 187/385] Generalization of code call --- src/mitim_tools/gacode_tools/NEOtools.py | 28 +++-------- src/mitim_tools/gacode_tools/TGLFtools.py | 42 ++++------------- .../gacode_tools/utils/GACODErun.py | 47 +++++++++++++++++-- tests/NEO_workflow.py | 2 +- 4 files changed, 60 insertions(+), 59 deletions(-) diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 15afebbe..45ec7d3b 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -1,5 +1,4 @@ import numpy as np -from pathlib import Path import matplotlib.pyplot as plt from mitim_tools import __version__ as mitim_version from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools @@ -7,8 +6,6 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -from mitim_tools.misc_tools.PLASMAtools import md_u - class NEO(GACODErun.gacode_simulation): def __init__( self, @@ -17,10 +14,13 @@ def __init__( super().__init__(rhos=rhos) + def code_call(folder, n, p, additional_command="", **kwargs): + return f" neo -e {folder} -n {n} -p {p} {additional_command} &\n" + self.run_specifications = { 'code': 'neo', 'input_file': 'input.neo', - 'code_call': 'neo -e', + 'code_call': code_call, 'control_function': GACODEdefaults.addNEOcontrol, 'controls_file': 'input.neo.controls', 'state_converter': 'to_neo', @@ -273,21 +273,10 @@ def check_if_files_exist(folder, list_files): return True - -class NEOinput: +class NEOinput(GACODErun.GACODEinput): def __init__(self, file=None): - self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None - - if self.file is not None and self.file.exists(): - with open(self.file, "r") as f: - lines = f.readlines() - file_txt = "".join(lines) - else: - file_txt = "" - input_dict = GACODErun.buildDictFromInput(file_txt) - - self.process(input_dict) - + super().__init__(file=file) + @classmethod def initialize_in_memory(cls, input_dict): instance = cls() @@ -299,9 +288,6 @@ def process(self, input_dict): self.controls = input_dict self.num_recorded = 100 - def anticipate_problems(self): - pass - def write_state(self, file=None): if file is None: diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index f70149af..4fb33bad 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -124,10 +124,13 @@ def __init__( super().__init__(rhos=rhos) + def code_call(folder, p, n = 1, additional_command="", **kwargs): + return f" tglf -e {folder} -n {n} -p {p} {additional_command} &\n" + self.run_specifications = { 'code': 'tglf', 'input_file': 'input.tglf', - 'code_call': 'tglf -e', + 'code_call': code_call, 'control_function': GACODEdefaults.addTGLFcontrol, 'controls_file': 'input.tglf.controls', 'state_converter': 'to_tglf', @@ -192,6 +195,7 @@ def __init__( "SELECTED": None, } + # This is redefined (from parent) because it has the option of producing WaveForms (very TGLF specific) def run( self, subfolder, @@ -199,9 +203,6 @@ def run( forceClosestUnstableWF=True, # Look at the growth rate spectrum and run exactly the ky of the closest unstable **kwargs_generic_run ): - ''' - I need to redefine the run method for the TGLF class because it has the option of producing WaveForms - ''' code_executor_full = super().run(subfolder, **kwargs_generic_run) @@ -209,6 +210,7 @@ def run( kwargs_generic_run['forceClosestUnstableWF'] = forceClosestUnstableWF self._helper_wf(code_executor_full, **kwargs_generic_run) + # This is redefined (from parent) because it has the option of producing WaveForms (very TGLF specific) def run_scan( self, subfolder, @@ -3258,25 +3260,9 @@ def reduceToControls(dict_all): return controls, plasma, geom -class TGLFinput: +class TGLFinput(GACODErun.GACODEinput): def __init__(self, file=None): - self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None - - if self.file is not None and self.file.exists(): - with open(self.file, "r") as f: - lines = f.readlines() - file_txt = "".join(lines) - else: - file_txt = "" - input_dict = GACODErun.buildDictFromInput(file_txt) - - self.process(input_dict) - - @classmethod - def initialize_in_memory(cls, input_dict): - instance = cls() - instance.process(input_dict) - return instance + super().__init__(file=file) def process(self, input_dict): @@ -3443,19 +3429,11 @@ def ensureQuasineutrality(self): diff = self.calcualteQuasineutralityError() print(f"\t- Oiriginal quasineutrality error: {diff:.1e}", typeMsg="i") - print( - f"\t- Modifying species {speciesMod} to ensure quasineutrality", - typeMsg="i", - ) + print(f"\t- Modifying species {speciesMod} to ensure quasineutrality", typeMsg="i") for i in speciesMod: self.species[i]["AS"] -= diff / self.species[i]["ZS"] / len(speciesMod) self.processSpecies() - print( - "\t- New quasineutrality error: {0:.1e}".format( - self.calcualteQuasineutralityError() - ), - typeMsg="i", - ) + print("\t- New quasineutrality error: {0:.1e}".format(self.calcualteQuasineutralityError()), typeMsg="i") def calcualteQuasineutralityError(self): fiZi = 0 diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index ae1aaea6..386771f6 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -68,6 +68,9 @@ def prep( else: self.profiles = mitim_state + # Keep a copy of the file + self.profiles.write_state(file=self.FolderGACODE / "input.gacode") + self.profiles.derive_quantities(mi_ref=md_u) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -691,7 +694,7 @@ def run_gacode_simulation( code = run_specifications.get('code', 'tglf') input_file = run_specifications.get('input_file', 'input.tglf') - code_call = run_specifications.get('code_call', 'tglf -e') + code_call = run_specifications.get('code_call', None) tmpFolder = FolderGACODE / f"tmp_{code}" IOtools.askNewFolder(tmpFolder, force=True) @@ -780,7 +783,7 @@ def run_gacode_simulation( # Loop over each folder and launch code, waiting if we've reached max_parallel_execution GACODEcommand += "for folder in \"${folders[@]}\"; do\n" - GACODEcommand += f" {code_call} \"$folder\" -n {cores_simulation} -p {gacode_job.folderExecution} &\n" + GACODEcommand += code_call(folder = '\"$folder\"', n = cores_simulation, p = gacode_job.folderExecution) GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" GACODEcommand += "done\n\n" GACODEcommand += "wait\n" @@ -803,7 +806,7 @@ def run_gacode_simulation( # Code launches GACODEcommand = "" for folder in folders_red: - GACODEcommand += f"{code_call} {folder} -n {cores_simulation} -p {gacode_job.folderExecution} &\n" + GACODEcommand += code_call(folder = folder, n = cores_simulation, p = gacode_job.folderExecution) GACODEcommand += "\nwait" # This is needed so that the script doesn't end before each job # Slurm setup @@ -832,7 +835,11 @@ def run_gacode_simulation( # Code launches indexed_folder = 'run"$SLURM_ARRAY_TASK_ID"' - GACODEcommand = f'{code_call} {indexed_folder} -n {cores_simulation} -p {gacode_job.folderExecution} 1> {gacode_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {gacode_job.folderExecution}/{indexed_folder}/slurm_error.dat\n' + GACODEcommand = code_call( + folder = indexed_folder, + n = cores_simulation, + p = gacode_job.folderExecution, + additional_command = f'1> {gacode_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {gacode_job.folderExecution}/{indexed_folder}/slurm_error.dat\n') # Slurm setup array_list = ",".join(array_list) @@ -1068,7 +1075,7 @@ def modifyInputs( input_class.plasma[ikey] = var_new else: # If the variable in extraOptions wasn't in there, consider it a control param - print("\t\t- Variable to change did not exist previously, creating now",typeMsg="i") + print(f"\t\t- Variable {ikey} to change did not exist previously, creating now",typeMsg="i") var_orig = None var_new = value_to_change_to input_class.controls[ikey] = var_new @@ -1809,3 +1816,33 @@ def defineNewGrid( plt.show() return x[imin:imax], y[imin:imax] + +class GACODEinput: + def __init__(self, file=None): + self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None + + if self.file is not None and self.file.exists(): + with open(self.file, "r") as f: + lines = f.readlines() + file_txt = "".join(lines) + else: + file_txt = "" + input_dict = buildDictFromInput(file_txt) + + self.process(input_dict) + + @classmethod + def initialize_in_memory(cls, input_dict): + instance = cls() + instance.process(input_dict) + return instance + + def anticipate_problems(self): + pass + + def remove_fast(self): + pass + + def removeLowDensitySpecie(self, *args): + pass + \ No newline at end of file diff --git a/tests/NEO_workflow.py b/tests/NEO_workflow.py index 8f59565c..35fef115 100644 --- a/tests/NEO_workflow.py +++ b/tests/NEO_workflow.py @@ -13,7 +13,7 @@ if cold_start and folder.exists(): os.system(f"rm -r {folder.resolve()}") -neo = NEOtools.NEO(rhos=np.linspace(0.1,0.95,10)) +neo = NEOtools.NEO(rhos=np.linspace(0.1,0.95,5)) neo.prep(input_gacode, folder) neo.run('neo1/', cold_start=cold_start) From b0117f6d7d983e2d1b61de7585fbb74d810ba2ad Mon Sep 17 00:00:00 2001 From: pabloprf Date: Fri, 22 Aug 2025 18:10:30 -0400 Subject: [PATCH 188/385] misc --- .../powertorch/physics_models/transport_tglfneo.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index f1ea4bdb..1f90258c 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -174,7 +174,11 @@ def evaluate_neoclassical(self): cold_start = cold_start, ) - neo.run('base_neo') + neo.run( + 'base_neo', + cold_start=cold_start, + forceIfcold_start=True, + ) neo.read(label='base') From 47c5fb51dfaa96de6f7b491fa6948da81e62b836 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 23 Aug 2025 15:12:03 -0400 Subject: [PATCH 189/385] Generalization of read commands and CGYRO as another TGLF/NEO --- src/mitim_tools/gacode_tools/CGYROtools.py | 235 ++++++++++++------ src/mitim_tools/gacode_tools/NEOtools.py | 49 +--- src/mitim_tools/gacode_tools/TGLFtools.py | 20 +- .../gacode_tools/utils/CGYROutils.py | 116 ++++++--- .../gacode_tools/utils/GACODEdefaults.py | 155 ++++-------- .../gacode_tools/utils/GACODErun.py | 81 +++++- .../plasmastate_tools/MITIMstate.py | 177 ++++++++++++- templates/input.cgyro.controls | 16 +- tests/CGYRO_workflow.py | 22 +- 9 files changed, 573 insertions(+), 298 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index aea4ca36..9707fa75 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -1,31 +1,48 @@ -import os import shutil import datetime import time +import copy from pathlib import Path import numpy as np import matplotlib.pyplot as plt +from mitim_tools import __version__ as mitim_version +from mitim_tools import __mitimroot__ from mitim_tools.gacode_tools.utils import GACODEdefaults, GACODErun, CGYROutils from mitim_tools.misc_tools import IOtools, GRAPHICStools, FARMINGtools from mitim_tools.gacode_tools.utils import GACODEplotting from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -class CGYRO: - def __init__(self): - - self.output_files_test = [ - "out.cgyro.equilibrium", - "out.cgyro.info", - "out.cgyro.mpi", - "input.cgyro.gen", - "out.cgyro.egrid", - "out.cgyro.grids", - "out.cgyro.memory", - "out.cgyro.rotation", - ] - - self.output_files = [ +class CGYRO(GACODErun.gacode_simulation): + def __init__( + self, + rhos=[0.4, 0.6], # rho locations of interest + ): + + super().__init__(rhos=rhos) + + def code_call(folder, p, n = 1, nomp = 1, additional_command="", **kwargs): + return f" cgyro -e {folder} -n {n} -nomp {nomp} {additional_command} -p {p} &\n" + + self.run_specifications = { + 'code': 'cgyro', + 'input_file': 'input.cgyro', + 'code_call': code_call, + 'control_function': GACODEdefaults.addCGYROcontrol, + 'controls_file': 'input.cgyro.controls', + 'state_converter': 'to_cgyro', + 'input_class': CGYROinput, + 'complete_variation': None, + 'default_cores': 16, # Default cores to use in the simulation + 'output_class': CGYROutils.CGYROout, + 'output_store': 'CGYROout' + } + + print("\n-----------------------------------------------------------------------------------------") + print("\t\t\t CGYRO class module") + print("-----------------------------------------------------------------------------------------\n") + + self.ResultsFiles = self.ResultsFiles_minimal = [ "bin.cgyro.geo", "bin.cgyro.kxky_e", "bin.cgyro.kxky_n", @@ -36,14 +53,12 @@ def __init__(self): "bin.cgyro.ky_cflux", "bin.cgyro.ky_flux", "bin.cgyro.phib", + "bin.cgyro.aparb", + "bin.cgyro.bparb", "bin.cgyro.restart", - "bin.cgyro.restart.old", "input.cgyro", "input.cgyro.gen", - "input.gacode", "mitim.out", - "mitim_bash.src", - "mitim_shell_executor.sh", "out.cgyro.egrid", "out.cgyro.equilibrium", "out.cgyro.freq", @@ -61,43 +76,52 @@ def __init__(self): "out.cgyro.version", ] - self.results = {} + self.output_files_test = [ + "out.cgyro.equilibrium", + "out.cgyro.info", + "out.cgyro.mpi", + "input.cgyro.gen", + "out.cgyro.egrid", + "out.cgyro.grids", + "out.cgyro.memory", + "out.cgyro.rotation", + ] - def prep(self, folder, inputgacode_file): + # def prep(self, folder, inputgacode_file): - # Prepare main folder with input.gacode - self.folder = IOtools.expandPath(folder) + # # Prepare main folder with input.gacode + # self.folder = IOtools.expandPath(folder) - self.folder.mkdir(parents=True, exist_ok=True) + # self.folder.mkdir(parents=True, exist_ok=True) - self.inputgacode_file = self.folder / "input.gacode" - if IOtools.expandPath(inputgacode_file) != self.inputgacode_file: - shutil.copy2(IOtools.expandPath(inputgacode_file), self.inputgacode_file) + # self.inputgacode_file = self.folder / "input.gacode" + # if IOtools.expandPath(inputgacode_file) != self.inputgacode_file: + # shutil.copy2(IOtools.expandPath(inputgacode_file), self.inputgacode_file) def _prerun( self, - subFolderCGYRO, + subfolder, roa=0.55, CGYROsettings=None, extraOptions={}, multipliers={}, ): - self.folderCGYRO = self.folder / Path(subFolderCGYRO) + self.folder = self.FolderGACODE / Path(subfolder) - self.folderCGYRO.mkdir(parents=True, exist_ok=True) + self.folder.mkdir(parents=True, exist_ok=True) - input_cgyro_file = self.folderCGYRO / "input.cgyro" + input_cgyro_file = self.folder / "input.cgyro" inputCGYRO = CGYROinput(file=input_cgyro_file) - inputgacode_file_this = self.folderCGYRO / "input.gacode" - shutil.copy2(self.inputgacode_file, inputgacode_file_this) + inputgacode_file_this = self.folder / "input.gacode" + self.profiles.write_state(inputgacode_file_this) ResultsFiles_new = [] - for i in self.output_files: + for i in self.ResultsFiles: if "mitim.out" not in i: ResultsFiles_new.append(i) - self.output_files = ResultsFiles_new + self.ResultsFiles = ResultsFiles_new inputCGYRO = GACODErun.modifyInputs( inputCGYRO, @@ -109,13 +133,14 @@ def _prerun( controls_file = 'input.cgyro.controls' ) - inputCGYRO.writeCurrentStatus() + inputCGYRO.write_state() return input_cgyro_file, inputgacode_file_this + def run_test( self, - subFolderCGYRO, + subfolder, roa=0.55, CGYROsettings=None, extraOptions={}, @@ -127,16 +152,16 @@ def run_test( print("\t- Cannot run CGYRO tests with scan_param, running just the base",typeMsg="i") input_cgyro_file, inputgacode_file_this = self._prerun( - subFolderCGYRO, + subfolder, roa=roa, CGYROsettings=CGYROsettings, extraOptions=extraOptions, multipliers=multipliers, ) - self.cgyro_job = FARMINGtools.mitim_job(self.folderCGYRO) + self.cgyro_job = FARMINGtools.mitim_job(self.folder) - name = f'mitim_cgyro_{subFolderCGYRO}_{roa:.6f}_test' + name = f'mitim_cgyro_{subfolder}_{roa:.6f}_test' self.cgyro_job.define_machine( "cgyro", @@ -159,16 +184,16 @@ def run_test( self.cgyro_job.run() - def run(self,subFolderCGYRO,test_run=False,**kwargs): + def run1(self,subfolder,test_run=False,**kwargs): if test_run: - self.run_test(subFolderCGYRO,**kwargs) + self.run_test(subfolder,**kwargs) else: - self.run_full(subFolderCGYRO,**kwargs) + self.run_full(subfolder,**kwargs) def run_full( self, - subFolderCGYRO, + subfolder, roa=0.55, CGYROsettings=None, extraOptions={}, @@ -187,16 +212,16 @@ def run_full( ): input_cgyro_file, inputgacode_file_this = self._prerun( - subFolderCGYRO, + subfolder, roa=roa, CGYROsettings=CGYROsettings, extraOptions=extraOptions, multipliers=multipliers, ) - self.cgyro_job = FARMINGtools.mitim_job(self.folderCGYRO) + self.cgyro_job = FARMINGtools.mitim_job(self.folder) - name = f'mitim_cgyro_{subFolderCGYRO}_{roa:.6f}' + name = f'mitim_cgyro_{subfolder}_{roa:.6f}' if scan_param is not None and submit_via_qsub: raise Exception(" Cannot use scan_param with submit_via_qsub=True, because it requires a different job for each value of the scan parameter.") @@ -223,7 +248,7 @@ def run_full( self.slurm_output = "batch.out" # --- - folder_run = self.folderCGYRO / subfolder + folder_run = self.folder / subfolder folder_run.mkdir(parents=True, exist_ok=True) # Copy the input.cgyro in the subfolder @@ -290,7 +315,7 @@ def run_full( output_folders = [] for i,value in enumerate(scan_param['values']): subfolder = f"scan{i}" - folder_run = self.folderCGYRO / subfolder + folder_run = self.folder / subfolder folder_run.mkdir(parents=True, exist_ok=True) # Copy the input.cgyro in the subfolder @@ -319,7 +344,7 @@ def run_full( control_file = 'input.cgyro.controls' ) - input_cgyro_file_this.writeCurrentStatus() + input_cgyro_file_this.write_state() # Copy the input.gacode file in the subfolder inputgacode_file_this = folder_run / "input.gacode" @@ -398,13 +423,32 @@ def delete(self): self.cgyro_job.run() - # --------------------------------------------------------------------------------------------------------- - # Reading and plotting - # --------------------------------------------------------------------------------------------------------- + # Re-defined to make specific arguments explicit + def read( + self, + tmin = 0.0, + minimal = False, + last_tmin_for_linear = True, + **kwargs + ): + + super().read( + tmin = tmin, + minimal = minimal, + last_tmin_for_linear = last_tmin_for_linear, + **kwargs) + + results = copy.deepcopy(self.results) + + #TODO accept more than one radii + self.results = {} + for label in results: + for i,rho in enumerate(self.rhos): + self.results[label+f'_{rho}'] = results[label]['CGYROout'][i] - def read(self, label="cgyro1", folder=None, tmin = 0.0, minimal = False, last_tmin_for_linear = True): + def read1(self, label="cgyro1", folder=None, tmin = 0.0, minimal = False, last_tmin_for_linear = True): - folder = IOtools.expandPath(folder) if folder is not None else self.folderCGYRO + folder = IOtools.expandPath(folder) if folder is not None else self.folder folders = sorted(list((folder).glob("scan*")), key=lambda f: int(''.join(filter(str.isdigit, f.name)))) @@ -1626,7 +1670,10 @@ def plot_2D(self, label="cgyro1", axs=None, times = None): number_times = len(axs)//3 if axs is not None else 4 - times = [self.results[label].t[-1-i*10] for i in range(number_times)] + try: + times = [self.results[label].t[-1-i*10] for i in range(number_times)] + except IndexError: + times = [self.results[label].t[-1-i*1] for i in range(number_times)] if axs is None: @@ -1865,42 +1912,66 @@ def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): ax.set_xlim(left=0) - -class CGYROinput: +class CGYROinput(GACODErun.GACODEinput): def __init__(self, file=None): - self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None - - if self.file is not None and self.file.exists(): - with open(self.file, "r") as f: - lines = f.readlines() - self.file_txt = "".join(lines) - else: - self.file_txt = "" - - self.controls = GACODErun.buildDictFromInput(self.file_txt) - - def writeCurrentStatus(self, file=None): - print("\t- Writting CGYRO input file") + super().__init__(file=file, controls_file= __mitimroot__ / "templates" / "input.cgyro.controls") + + def process(self, input_dict): + # Use standard processing + self._process(input_dict) + + # Get number of recorded species + self.num_recorded = 0 + if "N_SPECIES" in input_dict: + self.num_recorded = int(input_dict["N_SPECIES"]) + + def write_state(self, file=None): + if file is None: file = self.file + # Local formatter: floats -> 6 significant figures in exponential (uppercase), + # ints stay as ints, bools as 0/1, sequences space-separated with same rule. + def _fmt_num(x): + import numpy as _np + if isinstance(x, (bool, _np.bool_)): + return "True" if x else "False" + if isinstance(x, (_np.floating, float)): + # 6 significant figures in exponential => 5 digits after decimal + return f"{float(x):.5E}" + if isinstance(x, (_np.integer, int)): + return f"{int(x)}" + return str(x) + + def _fmt_value(val): + import numpy as _np + if isinstance(val, (list, tuple, _np.ndarray)): + # Flatten numpy arrays but keep ordering; join with spaces + if isinstance(val, _np.ndarray): + flat = val.flatten().tolist() + else: + flat = list(val) + return " ".join(_fmt_num(v) for v in flat) + return _fmt_num(val) + with open(file, "w") as f: - f.write( - "#-------------------------------------------------------------------------\n" - ) - f.write( - "# CGYRO input file modified by MITIM framework (Rodriguez-Fernandez, 2020)\n" - ) - f.write( - "#-------------------------------------------------------------------------" - ) + f.write("#-------------------------------------------------------------------------\n") + f.write(f"# CGYRO input file modified by MITIM {mitim_version}\n") + f.write("#-------------------------------------------------------------------------\n") f.write("\n\n# Control parameters\n") f.write("# ------------------\n\n") for ikey in self.controls: var = self.controls[ikey] - f.write(f"{ikey.ljust(23)} = {var}\n") + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") + + f.write("\n\n# Plasma/Geometry parameters\n") + f.write("# ------------------\n\n") + params = self.plasma | self.geom + for ikey in params: + var = params[ikey] + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") def _2D_mosaic(n_times): diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 45ec7d3b..d08d6ecb 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -27,6 +27,8 @@ def code_call(folder, n, p, additional_command="", **kwargs): 'input_class': NEOinput, 'complete_variation': None, 'default_cores': 1, # Default cores to use in the simulation + 'output_class': NEOoutput, + 'output_store': 'NEOout' } print("\n-----------------------------------------------------------------------------------------") @@ -34,41 +36,6 @@ def code_call(folder, n, p, additional_command="", **kwargs): print("-----------------------------------------------------------------------------------------\n") self.ResultsFiles = self.ResultsFiles_minimal = ['out.neo.transport_flux'] - - def read( - self, - label="neo1", - folder=None, # If None, search in the previously run folder - suffix=None, # If None, search with my standard _0.55 suffixes corresponding to rho of this TGLF class - **kwargs - ): - print("> Reading NEO results") - - # If no specified folder, check the last one - if folder is None: - folder = self.FolderSimLast - - self.results[label] = { - 'NEOout':[], - 'parsed': [], - "x": np.array(self.rhos), - } - for rho in self.rhos: - - NEOout = NEOoutput( - folder, - suffix=f"_{rho:.4f}" if suffix is None else suffix, - ) - - # Unnormalize - NEOout.unnormalize( - self.NormalizationSets["SELECTED"], - rho=rho, - ) - - self.results[label]['NEOout'].append(NEOout) - - self.results[label]['parsed'].append(GACODErun.buildDictFromInput(NEOout.inputFile)) def plot( self, @@ -285,7 +252,9 @@ def initialize_in_memory(cls, input_dict): def process(self, input_dict): #TODO - self.controls = input_dict + self.plasma = input_dict + self.controls = {} + self.geom = {} self.num_recorded = 100 def write_state(self, file=None): @@ -322,13 +291,15 @@ def _fmt_value(val): f.write(f"# NEO input file modified by MITIM {mitim_version}\n") f.write("#-------------------------------------------------------------------------\n") - for ikey in self.controls: - var = self.controls[ikey] + for ikey in self.plasma: + var = self.plasma[ikey] f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") -class NEOoutput: +class NEOoutput(GACODErun.GACODEoutput): def __init__(self, FolderGACODE, suffix=""): + super().__init__() + self.FolderGACODE, self.suffix = FolderGACODE, suffix if suffix == "": diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 4fb33bad..228ed6a5 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -137,6 +137,8 @@ def code_call(folder, p, n = 1, additional_command="", **kwargs): 'input_class': TGLFinput, 'complete_variation': completeVariation_TGLF, 'default_cores': 4, # Default cores to use in the simulation + 'output_class': TGLFoutput, + 'output_store': 'TGLFout' } print("\n-----------------------------------------------------------------------------------------") @@ -3541,12 +3543,7 @@ def _fmt_value(val): f.write("# -------\n") for ikey in self.species: if ikey > maxSpeciesTGLF: - print( - "\t- Maximum number of species in TGLF reached, not considering after {0} species".format( - maxSpeciesTGLF - ), - typeMsg="w", - ) + print(f"\t- Maximum number of species in TGLF reached, not considering after {maxSpeciesTGLF} species",typeMsg="w",) break if ikey == 1: extralab = " (electrons)" @@ -3868,14 +3865,16 @@ def readTGLFresults( return results -class TGLFoutput: - def __init__(self, FolderGACODE, suffix="",require_all_files=True): +class TGLFoutput(GACODErun.GACODEoutput): + def __init__(self, FolderGACODE, suffix="", require_all_files=True): + super().__init__() + self.FolderGACODE, self.suffix = FolderGACODE, suffix - if suffix == "": + if self.suffix == "": print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} without suffix") else: - print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} with suffix {suffix}") + print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} with suffix {self.suffix}") self.inputclass = TGLFinput(file=self.FolderGACODE / f"input.tglf{self.suffix}") self.roa = self.inputclass.geom["RMIN_LOC"] @@ -3914,6 +3913,7 @@ def postprocess(self): self.GeES = np.sum(self.SumFlux_Ge_phi) self.GeEM = np.sum(self.SumFlux_Ge_a) + # Redefined because of very specific TGLF stuff def read(self,require_all_files=True): # -------------------------------------------------------------------------------- # Ions to include? e.g. IncludeExtraIonsInQi = [2,3,4] -> This will sum to ion 1 diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index d10657df..4b0d25e3 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -4,6 +4,7 @@ import statsmodels.api as sm import matplotlib.pyplot as plt from mitim_tools.misc_tools import IOtools +from mitim_tools.gacode_tools.utils import GACODErun from mitim_tools.misc_tools.LOGtools import printMsg as print from pygacode.cgyro.data_plot import cgyrodata_plot from pygacode import gacodefuncs @@ -48,27 +49,15 @@ def __init__(self, labels, cgyro_data): self.Qe_mean = np.array(self.Qe_mean) self.Qi_mean = np.array(self.Qi_mean) -class CGYROout: - def __init__(self, folder, tmin=0.0, minimal=False, last_tmin_for_linear=True): - - original_dir = os.getcwd() +class CGYROout(GACODErun.GACODEoutput): + def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for_linear=True, **kwargs): + + super().__init__() self.folder = folder self.tmin = tmin - try: - print(f"\t- Reading CGYRO data from {self.folder.resolve()}") - self.cgyrodata = cgyrodata_plot(f"{self.folder.resolve()}{os.sep}") - except FileNotFoundError: - raise Exception(f"[MITIM] Could not find CGYRO data in {self.folder.resolve()}. Please check the folder path or run CGYRO first.") - except Exception as e: - print(f"\t- Error reading CGYRO data: {e}") - if print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): - os.chdir(self.folder) - os.system("cgyro -t") - self.cgyrodata = cgyrodata_plot(f"{self.folder.resolve()}{os.sep}") - - os.chdir(original_dir) + self.cgyrodata = self.read_using_cgyroplot(self.folder, suffix) # -------------------------------------------------------------- # Read inputs @@ -139,7 +128,10 @@ def __init__(self, folder, tmin=0.0, minimal=False, last_tmin_for_linear=True): self.cgyrodata.getbigfield() if 'kxky_phi' in self.cgyrodata.__dict__: - self._process_fluctuations() + try: + self._process_fluctuations() + except ValueError as e: + print(f'\t- Error processing fluctuations: {e}', typeMsg='w') else: print(f'\t- No fluctuations found in CGYRO data ({IOtools.clipstr(self.folder)}), skipping fluctuation processing and will not be able to plot default Notebook', typeMsg='w') else: @@ -147,6 +139,61 @@ def __init__(self, folder, tmin=0.0, minimal=False, last_tmin_for_linear=True): self._process_fluxes() self._saturate_signals() + + self.remove_symlinks() + + def read_using_cgyroplot(self, folder, suffix): + + original_dir = os.getcwd() + + # Handle files with suffix by creating temporary symbolic links + self.temp_links = [] + if suffix: + import glob + + # Find all files with the suffix pattern + pattern = f"{folder.resolve()}{os.sep}*{suffix}" + suffixed_files = glob.glob(pattern) + + for suffixed_file in suffixed_files: + # Create expected filename without suffix + original_name = suffixed_file.replace(suffix, '') + + # Only create symlink if the original doesn't exist and the suffixed file does + if not os.path.exists(original_name) and os.path.exists(suffixed_file): + try: + os.symlink(suffixed_file, original_name) + self.temp_links.append(original_name) + print(f"\t- Created temporary link: {os.path.basename(original_name)} -> {os.path.basename(suffixed_file)}") + except (OSError, FileExistsError) as e: + print(f"\t- Warning: Could not create symlink for {os.path.basename(suffixed_file)}: {e}", typeMsg='w') + + try: + print(f"\t- Reading CGYRO data from {folder.resolve()}") + cgyrodata = cgyrodata_plot(f"{folder.resolve()}{os.sep}") + except FileNotFoundError: + raise Exception(f"[MITIM] Could not find CGYRO data in {folder.resolve()}. Please check the folder path or run CGYRO first.") + except Exception as e: + print(f"\t- Error reading CGYRO data: {e}") + if print('- Could not read data, do you want me to try do "cgyro -t" in the folder?',typeMsg='q'): + os.chdir(folder) + os.system("cgyro -t") + cgyrodata = cgyrodata_plot(f"{folder.resolve()}{os.sep}") + finally: + + os.chdir(original_dir) + + return cgyrodata + + def remove_symlinks(self): + # Remove temporary symbolic links + for temp_link in self.temp_links: + try: + if os.path.islink(temp_link): + os.unlink(temp_link) + print(f"\t- Removed temporary link: {os.path.basename(temp_link)}") + except OSError as e: + print(f"\t- Warning: Could not remove temporary link {os.path.basename(temp_link)}: {e}", typeMsg='w') def _process_linear(self): @@ -220,20 +267,33 @@ def _process_fluctuations(self): self.__dict__[var] = self.__dict__[var] * self.artificial_rhos_factor # ******************************************************************** + # Case with dimensions: (nradial,ntoroidal,time) + if len(self.__dict__[var].shape) == 3: + axis_radial = 0 + axis_toroidal = 1 + var_ntor0 = self.__dict__[var][:,0,:] + var_ntorn = self.__dict__[var][:,1:,:] + # Case with dimensions: (nradial,ntoroidal,nions,time) + elif len(self.__dict__[var].shape) == 4: + axis_radial = 0 + axis_toroidal = 2 + var_ntor0 = self.__dict__[var][:,:,0,:] + var_ntorn = self.__dict__[var][:,:,1:,:] + # Sum over radial modes - self.__dict__[var+'_rms_sumnr'] = (abs(self.__dict__[var][:,:,:])**2).sum(axis=(0))**0.5 # (ntoroidal, time) - + self.__dict__[var+'_rms_sumnr'] = (abs(self.__dict__[var][:,:,:])**2).sum(axis=(axis_radial))**0.5 # (ntoroidal, time) or (nions, ntoroidal, time) + # Sum over radial modes AND separate n=0 and n>0 (sum) modes - self.__dict__[var+'_rms_sumnr_n0'] = (abs(self.__dict__[var][:,0,:])**2).sum(axis=0)**0.5 # (time) - self.__dict__[var+'_rms_sumnr_sumn1'] = (abs(self.__dict__[var][:,1:,:])**2).sum(axis=(0,1))**0.5 # (time) - + self.__dict__[var+'_rms_sumnr_n0'] = (abs(self.__dict__[var][:,0,:])**2).sum(axis=axis_radial)**0.5 # (time) or (nions, time) + self.__dict__[var+'_rms_sumnr_sumn1'] = (abs(self.__dict__[var][:,1:,:])**2).sum(axis=(axis_radial,axis_toroidal))**0.5 # (time) or (nions, time) + # Sum over radial modes and toroidal modes - self.__dict__[var+'_rms_sumnr_sumn'] = (abs(self.__dict__[var][:,:,:])**2).sum(axis=(0,1))**0.5 # (time) - + self.__dict__[var+'_rms_sumnr_sumn'] = (abs(self.__dict__[var][:,:,:])**2).sum(axis=(axis_radial,axis_toroidal))**0.5 # (time) or (nions, time) + # Separate n=0, n>0 (sum) modes, and all n (sum) modes - self.__dict__[var+'_rms_n0'] = (abs(self.__dict__[var][:,0,:])**2)**0.5 # (nradial,time) - self.__dict__[var+'_rms_sumn1'] = (abs(self.__dict__[var][:,1:,:])**2).sum(axis=(1))**0.5 # (nradial,time) - self.__dict__[var+'_rms_sumn'] = (abs(self.__dict__[var][:,:,:])**2).sum(axis=(1))**0.5 # (nradial,time) + self.__dict__[var+'_rms_n0'] = (abs(var_ntor0)**2)**0.5 # (nradial,time) + self.__dict__[var+'_rms_sumn1'] = (abs(var_ntorn)**2).sum(axis=(axis_toroidal))**0.5 # (nradial,time) + self.__dict__[var+'_rms_sumn'] = (abs(self.__dict__[var])**2).sum(axis=(axis_toroidal))**0.5 # (nradial,time) # Cross-phases self.neTe = _cross_phase(self.t, self.ne, self.Te) * 180/ np.pi # (nradial, ntoroidal, time) diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index c99a8663..8cd15494 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -59,6 +59,30 @@ def addNEOcontrol(*args, **kwargs): return NEOoptions +def addCGYROcontrol(code_settings, rmin=None, **kwargs): + + CGYROoptions = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.cgyro.controls", caseInsensitive=False) + + """ + ******************************************************************************** + Standard sets of TGLF control parameters + (rest of parameters are as defaults) + ******************************************************************************** + """ + + with open(__mitimroot__ / "templates" / "input.cgyro.models.json", "r") as f: + settings = json.load(f) + + if str(code_settings) in settings: + sett = settings[str(code_settings)] + for ikey in sett["controls"]: + CGYROoptions[ikey] = sett["controls"][ikey] + else: + print("\t- code_settings not found in input.cgyro.models.json, using defaults",typeMsg="w") + + return CGYROoptions + + def TGLFinTRANSP(TGLFsettings, NS=3): TGLFoptions = addTGLFcontrol(TGLFsettings, NS=NS) @@ -94,75 +118,22 @@ def TGLFinTRANSP(TGLFsettings, NS=3): # **** New ones that are TRANSP-specific TGLFoptions["TGLFMOD"] = 1 - TGLFoptions["NSPEC"] = ( - NS # Number of species used in tglf model (maximum 10 species allowed) - ) + TGLFoptions["NSPEC"] = NS # Number of species used in tglf model (maximum 10 species allowed) TGLFoptions["NLGRAD"] = False # Output flux gradients TGLFoptions["ALPHA_N"] = 0.0 # Scaling factor for vn shear TGLFoptions["ALPHA_T"] = 0.0 # Scaling factor for vt shear # TGLFoptions['ALPHA_DIA'] = 0.0 # Scaling factor for diamagnetic terms to exb shear TGLFoptions["CBETAE"] = 1.0 # Betae multiplier (needed for e-m calcs) TGLFoptions["CXNU"] = 1.0 # Collisionality multiplier - TGLFoptions["EM_STAB"] = ( - 0.0 # EM factor for the ion temperature gradient --- Is this the right default? - ) - TGLFoptions["PEVOLVING"] = ( - 0 # Evolving temperature and its gradients --- Is this the right default? - ) - TGLFoptions["kinetic_fast_ion"] = ( - 0 # Fast ion species model in TGLF --- Is this the right default? - ) + TGLFoptions["EM_STAB"] = 0.0 # EM factor for the ion temperature gradient --- Is this the right default? + TGLFoptions["PEVOLVING"] = 0 # Evolving temperature and its gradients --- Is this the right default? + TGLFoptions["kinetic_fast_ion"] = 0 # Fast ion species model in TGLF --- Is this the right default? # **** Other modifications TGLFoptions["UNITS"] = f"'{TGLFoptions['UNITS']}'" return TGLFoptions - -def addCGYROcontrol(code_settings, rmin): - - CGYROoptions = IOtools.generateMITIMNamelist( - __mitimroot__ / "templates" / "input.cgyro.controls", caseInsensitive=False - ) - - """ - ******************************************************************************** - Standard sets of TGLF control parameters - (rest of parameters are as defaults) - ******************************************************************************** - """ - - with open( - __mitimroot__ / "templates" / "input.cgyro.models.json", "r" - ) as f: - settings = json.load(f) - - if str(code_settings) in settings: - sett = settings[str(code_settings)] - label = sett["label"] - for ikey in sett["controls"]: - CGYROoptions[ikey] = sett["controls"][ikey] - else: - print( - "\t- code_settings not found in input.cgyro.models.json, using defaults", - typeMsg="w", - ) - label = "unspecified" - - CGYROoptions["RMIN"] = rmin - - # -------------------------------- - # From dictionary to text - # -------------------------------- - - CGYROinput = [""] - for ikey in CGYROoptions: - CGYROinput.append(f"{ikey} = {CGYROoptions[ikey]}") - CGYROinput.append("") - - return CGYROinput, CGYROoptions, label - - def addTGYROcontrol( num_it=0, howmany=8, @@ -228,74 +199,42 @@ def addTGYROcontrol( f"{int(cold_start)}" # 0: Start from beginning, 1: Continue from last iteration ) TGYROoptions["TGYRO_RELAX_ITERATIONS"] = f"{num_it}" # Number of iterations - TGYROoptions["TGYRO_WRITE_PROFILES_FLAG"] = ( - "1" # 1: Create new input.profiles at end, 0: Nothing, -1: At all iterations - ) + TGYROoptions["TGYRO_WRITE_PROFILES_FLAG"] = "1" # 1: Create new input.profiles at end, 0: Nothing, -1: At all iterations # ----------- Optimization - TGYROoptions["LOC_RESIDUAL_METHOD"] = ( - f"{solver_options['res_method']}" # 2: |F|, 3: |F|^2 - ) - TGYROoptions["TGYRO_ITERATION_METHOD"] = ( - f"{solver_options['tgyro_method']}" # 1: Standard local residual, 2 3 4 5 6 - ) - TGYROoptions["LOC_DX"] = ( - f"{solver_options['step_jac']}" # Step length for Jacobian calculation (df: 0.1), units of a/Lx - ) - TGYROoptions["LOC_DX_MAX"] = ( - f"{solver_options['step_max']}" # Max length for any Newton step (df: 1.0) - ) - TGYROoptions["LOC_RELAX"] = ( - f"{solver_options['relax_param']}" # Parameter 𝐶𝜂 controlling shrinkage of relaxation parameter - ) + TGYROoptions["LOC_RESIDUAL_METHOD"] = f"{solver_options['res_method']}" # 2: |F|, 3: |F|^2 + TGYROoptions["TGYRO_ITERATION_METHOD"] = f"{solver_options['tgyro_method']}" # 1: Standard local residual, 2 3 4 5 6 + TGYROoptions["LOC_DX"] = f"{solver_options['step_jac']}" # Step length for Jacobian calculation (df: 0.1), units of a/Lx + TGYROoptions["LOC_DX_MAX"] = f"{solver_options['step_max']}" # Max length for any Newton step (df: 1.0) + TGYROoptions["LOC_RELAX"] = f"{solver_options['relax_param']}" # Parameter 𝐶𝜂 controlling shrinkage of relaxation parameter # ----------- Prediction Options - TGYROoptions["LOC_SCENARIO"] = ( - f"{physics_options['TypeTarget']}" # 1: Static targets, 2: dynamic exchange, 3: alpha, rad, exchange change - ) + TGYROoptions["LOC_SCENARIO"] = f"{physics_options['TypeTarget']}" # 1: Static targets, 2: dynamic exchange, 3: alpha, rad, exchange change TGYROoptions["LOC_TI_FEEDBACK_FLAG"] = f"{Tipred}" # Evolve Ti? TGYROoptions["LOC_TE_FEEDBACK_FLAG"] = f"{Tepred}" # Evolve Te? TGYROoptions["LOC_ER_FEEDBACK_FLAG"] = f"{Erpred}" # Evolve Er? TGYROoptions["TGYRO_DEN_METHOD0"] = f"{nepred}" # Evolve ne? - TGYROoptions["LOC_PFLUX_METHOD"] = ( - f"{physics_options['ParticleFlux']}" # Particle flux method. 1 = zero target flux, 2 = beam, 3 = beam+wall - ) + TGYROoptions["LOC_PFLUX_METHOD"] = f"{physics_options['ParticleFlux']}" # Particle flux method. 1 = zero target flux, 2 = beam, 3 = beam+wall TGYROoptions["TGYRO_RMIN"] = f"{fromRho}" TGYROoptions["TGYRO_RMAX"] = f"{ToRho}" - TGYROoptions["TGYRO_USE_RHO"] = ( - f"{solver_options['UseRho']}" # 1: Grid provided in input.tgyro is for rho values - ) + TGYROoptions["TGYRO_USE_RHO"] = f"{solver_options['UseRho']}" # 1: Grid provided in input.tgyro is for rho values # ----------- Physics TGYROoptions["TGYRO_ROTATION_FLAG"] = "1" # Trigger rotation physics? - TGYROoptions["TGYRO_NEO_METHOD"] = ( - f"{physics_options['neoclassical']}" # 0: None, 1: H&H, 2: NEO - ) - TGYROoptions["TGYRO_TGLF_REVISION"] = ( - "0" # 0: Use input.tglf in folders, instead of GA defaults. - ) - TGYROoptions["TGYRO_EXPWD_FLAG"] = ( - f"{physics_options['TurbulentExchange']}" # Add turbulent exchange to exchange powers in targets? - ) + TGYROoptions["TGYRO_NEO_METHOD"] = f"{physics_options['neoclassical']}" # 0: None, 1: H&H, 2: NEO + TGYROoptions["TGYRO_TGLF_REVISION"] = "0" # 0: Use input.tglf in folders, instead of GA defaults. + TGYROoptions["TGYRO_EXPWD_FLAG"] = f"{physics_options['TurbulentExchange']}" # Add turbulent exchange to exchange powers in targets? # TGYROoptions['TGYRO_ZEFF_FLAG'] = '1' # 1: Use Zeff from input.gacode # ----------- Assumptions for i in physics_options["quasineutrality"]: - TGYROoptions[f"TGYRO_DEN_METHOD{i}"] = ( - "-1" # Species used to ensure quasineutrality - ) + TGYROoptions[f"TGYRO_DEN_METHOD{i}"] = "-1" # Species used to ensure quasineutrality # TGYROoptions['LOC_NUM_EQUIL_FLAG'] = f"{physics_options['usingINPUTgeo']}" # 0: Use Miller, 1: Use numerical equilibrium (not valid for TGLF_scans) #DEPRECATED IN LATEST VERSIONS - TGYROoptions["LOC_LOCK_PROFILE_FLAG"] = ( - f"{physics_options['InputType']}" # 0: Re-compute profiles from coarse gradients grid, 1: Use exact profiles (only valid at first iteration) - ) - TGYROoptions["TGYRO_CONSISTENT_FLAG"] = ( - f"{physics_options['GradientsType']}" # 0: Finite-difference gradients used from input.gacode, 1: Gradients from coarse profiles? - ) + TGYROoptions["LOC_LOCK_PROFILE_FLAG"] = f"{physics_options['InputType']}" # 0: Re-compute profiles from coarse gradients grid, 1: Use exact profiles (only valid at first iteration) + TGYROoptions["TGYRO_CONSISTENT_FLAG"] = f"{physics_options['GradientsType']}" # 0: Finite-difference gradients used from input.gacode, 1: Gradients from coarse profiles? TGYROoptions["LOC_EVOLVE_GRAD_ONLY_FLAG"] = "0" # 1: Do not change absolute values - TGYROoptions["TGYRO_PTOT_FLAG"] = ( - f"{physics_options['PtotType']}" # 0: Compute pressure from profiles, 1: correct from input.gacode PTOT profile - ) + TGYROoptions["TGYRO_PTOT_FLAG"] = f"{physics_options['PtotType']}" # 0: Compute pressure from profiles, 1: correct from input.gacode PTOT profile # ----------- Radii @@ -407,12 +346,12 @@ def convolution_CECE(d_perp_dict, dRdx=1.0): def review_controls(TGLFoptions, control = "input.tglf.controls"): - TGLFoptions_check = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / control, caseInsensitive=False) + options_check = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / control, caseInsensitive=False) # Add plasma too potential_flags = ['NS', 'SIGN_BT', 'SIGN_IT', 'VEXB', 'VEXB_SHEAR', 'BETAE', 'XNUE', 'ZEFF', 'DEBYE'] for flag in potential_flags: - TGLFoptions_check[flag] = None + options_check[flag] = None for option in TGLFoptions: @@ -421,5 +360,5 @@ def review_controls(TGLFoptions, control = "input.tglf.controls"): # Do not fail with e.g. P_PRIME_LOC isGeometry = option.split('_')[-1] in ['LOC'] - if (not isSpecie) and (not isGeometry) and (option not in TGLFoptions_check): + if (not isSpecie) and (not isGeometry) and (option not in options_check): print(f"\t- Option {option} not in {control}, prone to errors", typeMsg="q") diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 386771f6..b8a2fe1b 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -112,10 +112,13 @@ def run( cold_start=False, forceIfcold_start=False, extra_name="exe", - slurm_setup={"cores": 1,"minutes": 1}, # Cores per call (so, when running nR radii -> nR*4) + slurm_setup=None, # Cores per call (so, when running nR radii -> nR*4) attempts_execution=1, only_minimal_files=False, ): + + if slurm_setup is None: + slurm_setup = {"cores": self.run_specifications['default_cores'], "minutes": 5} # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Prepare inputs @@ -463,6 +466,44 @@ def _prepare_scan( return code_executor, code_executor_full, folders, varUpDown + def read( + self, + label="run1", + folder=None, # If None, search in the previously run folder + suffix=None, # If None, search with my standard _0.55 suffixes corresponding to rho of this TGLF class + **kwargs_to_class_output + ): + print("> Reading simulation results") + + class_output = [self.run_specifications['output_class'], self.run_specifications['output_store']] + + # If no specified folder, check the last one + if folder is None: + folder = self.FolderSimLast + + self.results[label] = { + class_output[1]:[], + 'parsed': [], + "x": np.array(self.rhos), + } + for rho in self.rhos: + + SIMout = class_output[0]( + folder, + suffix=f"_{rho:.4f}" if suffix is None else suffix, + **kwargs_to_class_output + ) + + # Unnormalize + SIMout.unnormalize( + self.NormalizationSets["SELECTED"], + rho=rho, + ) + + self.results[label][class_output[1]].append(SIMout) + + self.results[label]['parsed'].append(buildDictFromInput(SIMout.inputFile) if SIMout.inputFile else None) + def read_scan( self, label="scan1", @@ -619,14 +660,14 @@ def change_and_write_code( return inputFile, mod_input_file -def inputToVariable(finalFolder, rhos, file='input.tglf'): +def inputToVariable(folder, rhos, file='input.tglf'): """ Entire text file to variable """ inputFilesTGLF = {} - for cont, rho in enumerate(rhos): - fileN = finalFolder / f"{file}_{rho:.4f}" + for rho in rhos: + fileN = folder / f"{file}_{rho:.4f}" with open(fileN, "r") as f: lines = f.readlines() @@ -766,6 +807,10 @@ def run_gacode_simulation( # Simply bash, no slurm if not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): + if cores_simulation > max_cores_per_node: + print(f"\t - Detected {cores_simulation} cores required, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node = cores_simulation + max_parallel_execution = max_cores_per_node // cores_simulation # Make sure we don't overload the machine when running locally (assuming no farming trans-node) print(f"\t- {code.upper()} will be executed as bash script (total cores: {total_cores_required}, cores per simulation: {cores_simulation}). MITIM will launch {total_simulation_executions // max_parallel_execution+1} sequential executions",typeMsg="i") @@ -1817,9 +1862,18 @@ def defineNewGrid( return x[imin:imax], y[imin:imax] +class GACODEoutput: + def __init__(self, *args, **kwargs): + self.inputFile = None + + def unnormalize(self, *args, **kwargs): + print("No unnormalization implemented.") + class GACODEinput: - def __init__(self, file=None): + def __init__(self, file=None, controls_file=None): self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None + + self.controls_file = controls_file if self.file is not None and self.file.exists(): with open(self.file, "r") as f: @@ -1830,6 +1884,7 @@ def __init__(self, file=None): input_dict = buildDictFromInput(file_txt) self.process(input_dict) + @classmethod def initialize_in_memory(cls, input_dict): @@ -1837,6 +1892,22 @@ def initialize_in_memory(cls, input_dict): instance.process(input_dict) return instance + def _process(self, input_dict): + + if self.controls_file is not None: + options_check = [key for key in IOtools.generateMITIMNamelist(self.controls_file, caseInsensitive=False).keys()] + else: + options_check = [] + + self.controls, self.plasma, self.geom = {}, {}, {} + for key in input_dict.keys(): + if key in options_check: + self.controls[key] = input_dict[key] + else: + self.plasma[key] = input_dict[key] + + self.num_recorded = 100 + def anticipate_problems(self): pass diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 48e6f05b..e65b9ac3 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2308,7 +2308,7 @@ def interpolator(y): 'ZS': self.Species[i]['Z'], 'MASS': self.Species[i]['A']/mass_ref, 'RLNS': interpolator(self.derived['aLni'][:,i]), - 'RLTS': interpolator(self.derived['aLTi'][:,0] if self.Species[i]['S'] == 'therm' else self.derived["aLTi"][:,i]), + 'RLTS': interpolator(self.derived["aLTi"][:,i]), 'TAUS': interpolator(self.derived["tite_all"][:,i]), 'AS': interpolator(self.derived['fi'][:,i]), 'VPAR': interpolator(vpar), @@ -2416,8 +2416,18 @@ def to_neo(self, r=[0.5], r_is_rho = True): s_delta = self.derived["r"] * self._deriv_gacode(self.profiles["delta(-)"]) s_zeta = self.derived["r"] * self._deriv_gacode(self.profiles["zeta(-)"]) - omega_rot = self.profiles["w0(rad/s)"] / self.derived['c_s'] * self.derived['a'] - omega_rot_deriv = self._deriv_gacode(self.profiles["w0(rad/s)"])/ self.derived['c_s'] * self.derived['a']**2 + # Rotations + rmaj = self.derived['Rmajoa'] + cs = self.derived['c_s'] + a = self.derived['a'] + mach = self.profiles["w0(rad/s)"] * (self.derived['Rmajoa']*a) + gamma_p = self._deriv_gacode(self.profiles["w0(rad/s)"]) * (self.derived['Rmajoa']*a) + + # NEO definition: 'OMEGA_ROT=',mach_loc/rmaj_loc/cs_loc + omega_rot = mach / rmaj / cs # Equivalent to: self.profiles["w0(rad/s)"] / self.derived['c_s'] * a + + # NEO definition: 'OMEGA_ROT_DERIV=',-gamma_p_loc*a/cs_loc/rmaj_loc + omega_rot_deriv = gamma_p * a / cs / rmaj # Equivalent to: self._deriv_gacode(self.profiles["w0(rad/s)"])/ self.derived['c_s'] * self.derived['a']**2 inputsNEO = {} for rho in r: @@ -2445,7 +2455,7 @@ def interpolator(y): 'Z': self.Species[i]['Z'], 'MASS': self.Species[i]['A']/mass_ref, 'DLNNDR': interpolator(self.derived['aLni'][:,i]), - 'DLNTDR': interpolator(self.derived['aLTi'][:,0] if self.Species[i]['S'] == 'therm' else self.derived["aLTi"][:,i]), + 'DLNTDR': interpolator(self.derived["aLTi"][:,i]), 'TEMP': interpolator(self.derived["tite_all"][:,i]), 'DENS': interpolator(self.derived['fi'][:,i]), } @@ -2531,6 +2541,165 @@ def interpolator(y): return inputsNEO + def to_cgyro(self, r=[0.5], r_is_rho = True): + + # <> Function to interpolate a curve <> + from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function + + # Determine if the input radius is rho toroidal or r/a + if r_is_rho: + r_interpolation = self.profiles['rho(-)'] + else: + r_interpolation = self.derived['roa'] + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Prepare the inputs + # --------------------------------------------------------------------------------------------------------------------------------------- + + # Determine the mass reference + mass_ref = 2.0 + + sign_it = int(-np.sign(self.profiles["current(MA)"][-1])) + sign_bt = int(-np.sign(self.profiles["bcentr(T)"][-1])) + + s_kappa = self.derived["r"] / self.profiles["kappa(-)"] * self._deriv_gacode(self.profiles["kappa(-)"]) + s_delta = self.derived["r"] * self._deriv_gacode(self.profiles["delta(-)"]) + s_zeta = self.derived["r"] * self._deriv_gacode(self.profiles["zeta(-)"]) + + # Rotations + cs = self.derived['c_s'] + a = self.derived['a'] + mach = self.profiles["w0(rad/s)"] * (self.derived['Rmajoa']*a) + gamma_p = -self._deriv_gacode(self.profiles["w0(rad/s)"]) * (self.derived['Rmajoa']*a) + gamma_e = -self._deriv_gacode(self.profiles["w0(rad/s)"]) * (self.profiles['rmin(m)'] / self.profiles['q(-)']) + + # CGYRO definition: 'MACH=',mach_loc/cs_loc + mach = mach / cs + + # CGYRO definition: 'GAMMA_P=',gamma_p_loc*a/cs_loc + gamma_p = gamma_p * a / cs + + # CGYRO definition: 'GAMMA_E=',gamma_e_loc*a/cs_loc + gamma_e = gamma_e * a / cs + + + # Because in MITIMstate I keep Bunit always positive, but CGYRO routines may need it negative? #TODO + sign_Bunit = np.sign(self.profiles['torfluxa(Wb/radian)'][0]) + + + inputsCGYRO = {} + for rho in r: + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Define interpolator at this rho + # --------------------------------------------------------------------------------------------------------------------------------------- + + def interpolator(y): + return interpolation_function(rho, r_interpolation,y).item() + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Controls come from options + # --------------------------------------------------------------------------------------------------------------------------------------- + + controls = GACODEdefaults.addCGYROcontrol(0) + controls['PROFILE_MODEL'] = 1 + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Species come from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + species = {} + for i in range(len(self.Species)): + species[i+1] = { + 'Z': self.Species[i]['Z'], + 'MASS': self.Species[i]['A']/mass_ref, + 'DLNNDR': interpolator(self.derived['aLni'][:,i]), + 'DLNTDR': interpolator(self.derived["aLTi"][:,i]), + 'TEMP': interpolator(self.derived["tite_all"][:,i]), + 'DENS': interpolator(self.derived['fi'][:,i]), + } + + ie = i+2 + + species[ie] = { + 'Z': -1.0, + 'MASS': 0.000272445, + 'DLNNDR': interpolator(self.derived['aLne']), + 'DLNTDR': interpolator(self.derived['aLTe']), + 'TEMP': 1.0, + 'DENS': 1.0, + } + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Plasma comes from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + plasma = { + 'N_SPECIES': len(species), + 'IPCCW': sign_bt, + 'BTCCW': sign_it, + 'MACH': interpolator(mach), + 'GAMMA_E': interpolator(gamma_e), + 'GAMMA_P': interpolator(gamma_p), + 'NU_EE': interpolator(self.derived['xnue']), + 'BETAE_UNIT': interpolator(self.derived['betae']), + 'LAMBDA_STAR': interpolator(self.derived['debye']) * sign_Bunit, + } + + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Geometry comes from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + parameters = { + 'RMIN': self.derived['roa'], + 'RMAJ': self.derived['Rmajoa'], + 'SHIFT': self._deriv_gacode(self.profiles["rmaj(m)"]), + 'ZMAG': self.derived["Zmagoa"], + 'DZMAG': self._deriv_gacode(self.profiles["zmag(m)"]), + 'Q': np.abs(self.profiles["q(-)"]), + 'S': self.derived["s_hat"], + 'KAPPA': self.profiles["kappa(-)"], + 'S_KAPPA': s_kappa, + 'DELTA': self.profiles["delta(-)"], + 'S_DELTA': s_delta, + 'ZETA': self.profiles["zeta(-)"], + 'S_ZETA': s_zeta, + } + + # Add MXH and derivatives + for ikey in self.profiles: + if 'shape_cos' in ikey or 'shape_sin' in ikey: + + # TGLF only accepts 6, as of July 2025 + if int(ikey[-4]) > 6: + continue + + key_mod = ikey.upper().split('(')[0] + + parameters[key_mod] = self.profiles[ikey] + parameters[f"{key_mod.split('_')[0]}_S_{key_mod.split('_')[-1]}"] = self.derived["r"] * self._deriv_gacode(self.profiles[ikey]) + + geom = {} + for k in parameters: + par = torch.nan_to_num(torch.from_numpy(parameters[k]) if type(parameters[k]) is np.ndarray else parameters[k], nan=0.0, posinf=1E10, neginf=-1E10) + geom[k] = interpolator(par) + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Merging + # --------------------------------------------------------------------------------------------------------------------------------------- + + input_dict = {**controls, **plasma, **geom} + + for i in range(len(species)): + for k in species[i+1]: + input_dict[f'{k}_{i+1}'] = species[i+1][k] + + inputsCGYRO[rho] = input_dict + + return inputsCGYRO + + def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', times = [0.0,1.0], Vsurf = 0.0): print("\t- Converting to TRANSP") diff --git a/templates/input.cgyro.controls b/templates/input.cgyro.controls index e2d85fa3..a1f0bf0a 100644 --- a/templates/input.cgyro.controls +++ b/templates/input.cgyro.controls @@ -17,10 +17,9 @@ NONLINEAR_FLAG=1 #Radius used for simulation -RMIN=0.55 -#Use Experimental or Specified Inputs -PROFILE_MODEL=2 +#Data source +PROFILE_MODEL=1 # 1: Use inputs in this file; 2: Useinput.gacode (indicate RMIN=) #Geometry type (1=s-alpha,2=MXH) EQUILIBRIUM_MODEL=2 @@ -128,18 +127,11 @@ BETAE_UNIT_SCALE=1.0 BETA_STAR_SCALE=1.0 LAMBDA_STAR_SCALE=1.0 - -#============================== -#Species Specification -#============================== - -#Number of gyrokinetic species -N_SPECIES=2 - EXCH_FLAG=1 +# Gradient scalings DLNTDR_SCALE_1 = 1.0 DLNTDR_SCALE_2 = 1.0 DLNTDR_SCALE_3 = 1.0 DLNTDR_SCALE_4 = 1.0 -DLNTDR_SCALE_5 = 1.0 \ No newline at end of file +DLNTDR_SCALE_5 = 1.0 diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index 15ff7a8e..1b0ebc0b 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -12,26 +12,28 @@ folder.mkdir(parents=True, exist_ok=True) -cgyro = CGYROtools.CGYRO() +cgyro = CGYROtools.CGYRO(rhos = [0.5, 0.7]) -cgyro.prep(folder,gacode_file) +cgyro.prep( + gacode_file, + folder) cgyro.run( 'linear', - roa = 0.55, - CGYROsettings=0, + code_settings=0, extraOptions={ 'KY':0.3, - 'MAX_TIME': 1E1, # Short, I just want to test the run + 'MAX_TIME': 10.0, # Short, I just want to test the run }, - submit_via_qsub=False # NERSC: True #TODO change this + slurm_setup={'cores':4} + #submit_via_qsub=False # NERSC: True #TODO change this ) -cgyro.check(every_n_minutes=1) -cgyro.fetch() -cgyro.delete() +# cgyro.check(every_n_minutes=1) +# cgyro.fetch() +# cgyro.delete() cgyro.read(label="cgyro1") -cgyro.plot(labels=["cgyro1"]) +cgyro.plot(labels=["cgyro1_0.5","cgyro1_0.7"]) cgyro.fn.show() From e4d5b32ec2f65b6c4768be8363eb3d598f121140 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 23 Aug 2025 16:12:48 -0400 Subject: [PATCH 190/385] Full CGYRO workflow working, inheriting from gacode simulation --- src/mitim_tools/gacode_tools/CGYROtools.py | 2485 +++++++++-------- .../gacode_tools/utils/CGYROutils.py | 6 +- .../gacode_tools/utils/GACODErun.py | 11 +- tests/CGYRO_workflow.py | 42 +- 4 files changed, 1315 insertions(+), 1229 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 9707fa75..5b6d5359 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -87,400 +87,34 @@ def code_call(folder, p, n = 1, nomp = 1, additional_command="", **kwargs): "out.cgyro.rotation", ] - # def prep(self, folder, inputgacode_file): - - # # Prepare main folder with input.gacode - # self.folder = IOtools.expandPath(folder) - - # self.folder.mkdir(parents=True, exist_ok=True) - - # self.inputgacode_file = self.folder / "input.gacode" - # if IOtools.expandPath(inputgacode_file) != self.inputgacode_file: - # shutil.copy2(IOtools.expandPath(inputgacode_file), self.inputgacode_file) - - def _prerun( - self, - subfolder, - roa=0.55, - CGYROsettings=None, - extraOptions={}, - multipliers={}, - ): - - self.folder = self.FolderGACODE / Path(subfolder) - - self.folder.mkdir(parents=True, exist_ok=True) - - input_cgyro_file = self.folder / "input.cgyro" - inputCGYRO = CGYROinput(file=input_cgyro_file) - - inputgacode_file_this = self.folder / "input.gacode" - self.profiles.write_state(inputgacode_file_this) - - ResultsFiles_new = [] - for i in self.ResultsFiles: - if "mitim.out" not in i: - ResultsFiles_new.append(i) - self.ResultsFiles = ResultsFiles_new - - inputCGYRO = GACODErun.modifyInputs( - inputCGYRO, - code_settings=CGYROsettings, - extraOptions=extraOptions, - multipliers=multipliers, - addControlFunction=GACODEdefaults.addCGYROcontrol, - rmin=roa, - controls_file = 'input.cgyro.controls' - ) - - inputCGYRO.write_state() - - return input_cgyro_file, inputgacode_file_this - - - def run_test( - self, - subfolder, - roa=0.55, - CGYROsettings=None, - extraOptions={}, - multipliers={}, - **kwargs - ): - - if 'scan_param' in kwargs: - print("\t- Cannot run CGYRO tests with scan_param, running just the base",typeMsg="i") - - input_cgyro_file, inputgacode_file_this = self._prerun( - subfolder, - roa=roa, - CGYROsettings=CGYROsettings, - extraOptions=extraOptions, - multipliers=multipliers, - ) - - self.cgyro_job = FARMINGtools.mitim_job(self.folder) - - name = f'mitim_cgyro_{subfolder}_{roa:.6f}_test' - - self.cgyro_job.define_machine( - "cgyro", - name, - slurm_settings={ - "name": name, - "minutes": 5, - "cpuspertask": 1, - "ntasks": 1, - }, - ) - - CGYROcommand = "cgyro -t ." - - self.cgyro_job.prep( - CGYROcommand, - input_files=[input_cgyro_file, inputgacode_file_this], - output_files=self.output_files_test, - ) - - self.cgyro_job.run() - - def run1(self,subfolder,test_run=False,**kwargs): - - if test_run: - self.run_test(subfolder,**kwargs) - else: - self.run_full(subfolder,**kwargs) - - def run_full( - self, - subfolder, - roa=0.55, - CGYROsettings=None, - extraOptions={}, - multipliers={}, - scan_param = None, # {'variable': 'KY', 'values': [0.2,0.3,0.4]} - enforce_equality = None, # e.g. {'DLNTDR_SCALE_2': 'DLNTDR_SCALE_1', 'DLNTDR_SCALE_3': 'DLNTDR_SCALE_1'} - minutes = 5, - n = 16, - nomp = 1, - cpuspertask=None, # if None, will default to 1 - queue=None, #if blank will default to the one in settings - mem=None, # in MB - submit_via_qsub=True, #TODO fix this, works only at NERSC? no scans? - clean_folder_going_in=True, # Make sure the scratch folder is removed before running (unless I want a restart!) - submit_run=True, # False if I just want to check and fetch the job that was already submitted (e.g. via qsub or slurm) - ): - - input_cgyro_file, inputgacode_file_this = self._prerun( - subfolder, - roa=roa, - CGYROsettings=CGYROsettings, - extraOptions=extraOptions, - multipliers=multipliers, - ) - - self.cgyro_job = FARMINGtools.mitim_job(self.folder) - - name = f'mitim_cgyro_{subfolder}_{roa:.6f}' - - if scan_param is not None and submit_via_qsub: - raise Exception(" Cannot use scan_param with submit_via_qsub=True, because it requires a different job for each value of the scan parameter.") - - if submit_via_qsub: - - self.cgyro_job.define_machine( - "cgyro", - name, - launchSlurm=False, - ) - - subfolder = "scan0" - queue = "-queue " + self.cgyro_job.machineSettings['slurm']['partition'] if "partition" in self.cgyro_job.machineSettings['slurm'] else "" - - if mem is not None: - CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} -mem {mem} {queue} -w 0:{minutes}:00 -s' - else: - CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} {queue} -w 0:{minutes}:00 -s' - - if "account" in self.cgyro_job.machineSettings["slurm"] and self.cgyro_job.machineSettings["slurm"]["account"] is not None: - CGYROcommand += f" -repo {self.cgyro_job.machineSettings['slurm']['account']}" - - self.slurm_output = "batch.out" - - # --- - folder_run = self.folder / subfolder - folder_run.mkdir(parents=True, exist_ok=True) - - # Copy the input.cgyro in the subfolder - input_cgyro_file_this = folder_run / "input.cgyro" - shutil.copy2(input_cgyro_file, input_cgyro_file_this) - - # Copy the input.gacode file in the subfolder - inputgacode_file_this = folder_run / "input.gacode" - shutil.copy2(self.inputgacode_file, inputgacode_file_this) - - # Prepare the input and output folders - input_folders = [folder_run] - output_folders = [subfolder] - - else: - - if scan_param is None: - job_array = None - folder = 'scan0' - scan_param = {'variable': None, 'values': [0]} # Dummy scan parameter to avoid issues with the code below - else: - # Array - job_array = '' - for i,value in enumerate(scan_param['values']): - if job_array != '': - job_array += ',' - job_array += str(i) - - folder = 'scan"$SLURM_ARRAY_TASK_ID"' - - # Machine - self.cgyro_job.define_machine( - "cgyro", - name, - slurm_settings={ - "name": name, - "minutes": minutes, - "ntasks": n, - "job_array": job_array, - # Validate n and nomp before assigning cpuspertask - }, - ) - - if cpuspertask is not None: - if not isinstance(cpuspertask, int): - raise TypeError(" cpuspertask must be an integer") - self.cgyro_job.slurm_settings["cpuspertask"] = cpuspertask - - - if queue is not None: - self.cgyro_job.machineSettings['slurm']['partition'] = queue - - if mem is not None: - self.cgyro_job.slurm_settings['mem'] = mem - - # if not self.cgyro_job.launchSlurm: - # raise Exception(" Cannot run CGYRO scans without slurm") - - # Command to run cgyro - CGYROcommand = f'cgyro -e {folder} -n {n} -nomp {nomp} -p {self.cgyro_job.folderExecution}' - - # Scans - input_folders = [] - output_folders = [] - for i,value in enumerate(scan_param['values']): - subfolder = f"scan{i}" - folder_run = self.folder / subfolder - folder_run.mkdir(parents=True, exist_ok=True) - - # Copy the input.cgyro in the subfolder - input_cgyro_file_this = folder_run / "input.cgyro" - shutil.copy2(input_cgyro_file, input_cgyro_file_this) - - # Modify the input.cgyro file with the scan parameter - extraOptions_this = extraOptions.copy() - if scan_param['variable'] is not None: - extraOptions_this[scan_param['variable']] = value - - - # If there is an enforce_equality, apply it - if enforce_equality is not None: - for key in enforce_equality: - extraOptions_this[key] = extraOptions_this[enforce_equality[key]] - - inputCGYRO = CGYROinput(file=input_cgyro_file_this) - input_cgyro_file_this = GACODErun.modifyInputs( - inputCGYRO, - code_settings=CGYROsettings, - extraOptions=extraOptions_this, - multipliers=multipliers, - addControlFunction=GACODEdefaults.addCGYROcontrol, - rmin=roa, - control_file = 'input.cgyro.controls' - ) - - input_cgyro_file_this.write_state() - - # Copy the input.gacode file in the subfolder - inputgacode_file_this = folder_run / "input.gacode" - shutil.copy2(self.inputgacode_file, inputgacode_file_this) - - # Prepare the input and output folders - input_folders.append(folder_run) - output_folders.append(subfolder) - - self.slurm_output = "slurm_output.dat" - - # First submit the job with gacode_qsub, which will submit the cgyro job via slurm, with name - self.cgyro_job.prep( - CGYROcommand, - input_folders = input_folders, - output_folders=output_folders, - ) - - if submit_run: - self.cgyro_job.run( - waitYN=False, - check_if_files_received=False, - removeScratchFolders=False, - removeScratchFolders_goingIn=clean_folder_going_in, - ) - - # Prepare how to search for the job without waiting for it - name_default_submission_qsub = Path(self.cgyro_job.folderExecution).name - - self.cgyro_job.launchSlurm = True - self.cgyro_job.slurm_settings['name'] = name_default_submission_qsub - - - def check(self, every_n_minutes=5): - - if self.cgyro_job.launchSlurm: - print("- Checker job status") - - while True: - self.cgyro_job.check(file_output = self.slurm_output) - print(f'\t- Current status (as of {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}): {self.cgyro_job.status} ({self.cgyro_job.infoSLURM["STATE"]})') - if self.cgyro_job.status == 2: - break - else: - print(f"\t- Waiting {every_n_minutes} minutes") - time.sleep(every_n_minutes * 60) - else: - print("- Not checking status because this was run command line (not slurm)") - - print("\n\t* Job considered finished",typeMsg="i") - - def fetch(self): - """ - For a job that has been submitted but not waited for, once it is done, get the results - """ - - print("\n\n\t- Fetching results") - - if self.cgyro_job.launchSlurm: - self.cgyro_job.connect() - self.cgyro_job.retrieve() - self.cgyro_job.close() - else: - print("- Not retrieving results because this was run command line (not slurm)") - - def delete(self): - - print("\n\n\t- Deleting job") - - self.cgyro_job.launchSlurm = False - - self.cgyro_job.prep( - f"scancel -n {self.cgyro_job.slurm_settings['name']}", - label_log_files="_finish", - ) - - self.cgyro_job.run() - # Re-defined to make specific arguments explicit def read( self, tmin = 0.0, minimal = False, last_tmin_for_linear = True, - **kwargs - ): - - super().read( - tmin = tmin, - minimal = minimal, - last_tmin_for_linear = last_tmin_for_linear, - **kwargs) - - results = copy.deepcopy(self.results) - - #TODO accept more than one radii - self.results = {} - for label in results: - for i,rho in enumerate(self.rhos): - self.results[label+f'_{rho}'] = results[label]['CGYROout'][i] - - def read1(self, label="cgyro1", folder=None, tmin = 0.0, minimal = False, last_tmin_for_linear = True): - - folder = IOtools.expandPath(folder) if folder is not None else self.folder - - folders = sorted(list((folder).glob("scan*")), key=lambda f: int(''.join(filter(str.isdigit, f.name)))) - - if len(folders) == 0: - folders = [folder] - attach_name = False - else: - print(f"\t- Found {len(folders)} scan folders in {folder.resolve()}:") - for f in folders: - print(f"\t\t- {f.name}") - attach_name = True - - data = {} - labels = [] - for folder in folders: - - if attach_name: - label1 = f"{label}_{folder.name}" - else: - label1 = label - - data[label1] = CGYROutils.CGYROout(folder, tmin=tmin, minimal=minimal, last_tmin_for_linear=last_tmin_for_linear) - labels.append(label1) + **kwargs + ): + + super().read( + tmin = tmin, + minimal = minimal, + last_tmin_for_linear = last_tmin_for_linear, + **kwargs) - self.results.update(data) - - if attach_name: - self.results[label] = CGYROutils.CGYROlinear_scan(labels, data) + def _labelize(self, labels): - def plot(self, labels=[""], include_2D=True, common_colorbar=True): + # If it has radii, we need to correct the labels + self.results_all = copy.deepcopy(self.results) + self.results = {} + labels_with_rho = [] + for label in labels: + for i,rho in enumerate(self.rhos): + labels_with_rho.append(f"{label}_{rho}") + self.results[f'{label}_{rho}'] = self.results_all[label]['CGYROout'][i] + labels = labels_with_rho + # ------------------------------------------------ - # If it has scans, we need to correct the labels labels_corrected = [] for i in range(len(labels)): @@ -492,9 +126,22 @@ def plot(self, labels=[""], include_2D=True, common_colorbar=True): labels = labels_corrected # ------------------------------------------------ + return labels + + def plot( + self, + labels=[""], + fn=None, + include_2D=True, + common_colorbar=True): - from mitim_tools.misc_tools.GUItools import FigureNotebook - self.fn = FigureNotebook("CGYRO Notebook", geometry="1600x1000") + labels = self._labelize(labels) + + if fn is None: + from mitim_tools.misc_tools.GUItools import FigureNotebook + self.fn = FigureNotebook("CGYRO Notebook", geometry="1600x1000") + else: + self.fn = fn fig = self.fn.add_figure(label="Fluxes (time)") axsFluxes_t = fig.subplot_mosaic( @@ -686,6 +333,8 @@ def plot(self, labels=[""], include_2D=True, common_colorbar=True): cb.update_ticks() #cb.set_label(f"{var} (common range)") + self.results = self.results_all + def _plot_trace(self, ax, label, variable, c="b", lw=1, ls="-", label_plot='', meanstd=True, var_meanstd= None): t = self.results[label].t @@ -808,273 +457,637 @@ def plot_fluxes(self, axs=None, label="", c="b", lw=1, plotLegend=True): if plotLegend: ax.legend(loc='best', prop={'size': 8},) - # Ion energy fluxes - ax = axs["C"] - self._plot_trace(ax,label,"Qi",c=c,lw=lw,ls=ls[0],label_plot=f"{label}, Total") - self._plot_trace(ax,label,"Qi_EM",c=c,lw=lw,ls=ls[1],label_plot=f"{label}, EM ($A_\\parallel$+$A_\\perp$)", meanstd=False) - + # Ion energy fluxes + ax = axs["C"] + self._plot_trace(ax,label,"Qi",c=c,lw=lw,ls=ls[0],label_plot=f"{label}, Total") + self._plot_trace(ax,label,"Qi_EM",c=c,lw=lw,ls=ls[1],label_plot=f"{label}, EM ($A_\\parallel$+$A_\\perp$)", meanstd=False) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$Q_i$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion energy fluxes') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + + # Ion species energy fluxes + ax = axs["D"] + for j, i in enumerate(self.results[label].ions_flags): + self._plot_trace(ax,label,self.results[label].Qi_all[j],c=c,lw=lw,ls=ls[j],label_plot=f"{label}, {self.results[label].all_names[i]}", meanstd=False) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$Q_i$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion energy fluxes (separate species)') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_fluxes_ky(self, axs=None, label="", c="b", lw=1, plotLegend=True): + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + AC + BD + """ + ) + + ls = GRAPHICStools.listLS() + + # Electron energy flux + ax = axs["A"] + ax.plot(self.results[label].ky, self.results[label].Qe_ky_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Qe_ky_mean-self.results[label].Qe_ky_std, self.results[label].Qe_ky_mean+self.results[label].Qe_ky_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$Q_e$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron energy flux vs. $k_\\theta\\rho_s$') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + # Electron particle flux + ax = axs["B"] + ax.plot(self.results[label].ky, self.results[label].Ge_ky_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Ge_ky_mean-self.results[label].Ge_ky_std, self.results[label].Ge_ky_mean+self.results[label].Ge_ky_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\Gamma_e$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron particle flux vs. $k_\\theta\\rho_s$') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + # Ion energy flux + ax = axs["C"] + ax.plot(self.results[label].ky, self.results[label].Qi_ky_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Qi_ky_mean-self.results[label].Qi_ky_std, self.results[label].Qi_ky_mean+self.results[label].Qi_ky_std, color=c, alpha=0.2) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$Q_i$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion energy fluxes vs. $k_\\theta\\rho_s$') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + # Ion species energy fluxes + ax = axs["D"] + for j, i in enumerate(self.results[label].ions_flags): + ax.plot(self.results[label].ky, self.results[label].Qi_all_ky_mean[j],ls[j]+'o', markersize=5, color=c, label=f"{label}, {self.results[label].all_names[i]}") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$Q_i$ (GB)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion energy fluxes vs. $k_\\theta\\rho_s$(separate species)') + if plotLegend: + ax.legend(loc='best', prop={'size': 8},) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_intensities(self, axs = None, label= "cgyro1", c="b", addText=True): + + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + ACEG + BDFH + """ + ) + + ls = GRAPHICStools.listLS() + + ax = axs["A"] + ax.plot(self.results[label].t, self.results[label].phi_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].phi_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].phi_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta \\phi/\\phi_0$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Potential intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta\phi/\phi_0|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + ax = axs["B"] + if 'apar' in self.results[label].__dict__: + ax.plot(self.results[label].t, self.results[label].apar_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}, $A_\\parallel$") + ax.plot(self.results[label].t, self.results[label].bpar_rms_sumnr_sumn*100.0, '--', c=c, lw=2, label=f"{label}, $B_\\parallel$") + ax.legend(loc='best', prop={'size': 8},) + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta F_\\parallel/F_{\\parallel,0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('EM potential intensity fluctuations') + + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta F_\parallel/F_{\parallel,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + + ax = axs["C"] + ax.plot(self.results[label].t, self.results[label].ne_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].ne_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].ne_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta n_e/n_{e,0}/n_{e0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron Density intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta n_e/n_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + + ax = axs["D"] + ax.plot(self.results[label].t, self.results[label].Te_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].Te_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].Te_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta T_e/T_{e,0}/T_{e0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Electron Temperature intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta T_e/T_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + + + ax = axs["E"] + ax.plot(self.results[label].t, self.results[label].ni_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].ni_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].ni_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta n_i/n_{i,0}/n_{i0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion Density intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta n_i/n_{i,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + + ax = axs["F"] + ax.plot(self.results[label].t, self.results[label].Ti_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") + ax.plot(self.results[label].t, self.results[label].Ti_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") + ax.plot(self.results[label].t, self.results[label].Ti_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) + ax.set_ylabel("$\\delta T_i/T_{i,0}/T_{i0}$ (%)") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion Temperature intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta T_i/T_{i,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + ax = axs["G"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].t, self.results[label].ni_all_rms_sumnr_sumn[ion]*100.0, ls[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]}") + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$Q_i$ (GB)") + ax.set_ylabel("$\\delta n_i/n_{i,0}/n_{i0}$ (%)") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ion energy fluxes') - if plotLegend: - ax.legend(loc='best', prop={'size': 8},) + ax.set_title('Ions (all) Density intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) - # Ion species energy fluxes - ax = axs["D"] - for j, i in enumerate(self.results[label].ions_flags): - self._plot_trace(ax,label,self.results[label].Qi_all[j],c=c,lw=lw,ls=ls[j],label_plot=f"{label}, {self.results[label].all_names[i]}", meanstd=False) - + + ax = axs["H"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].t, self.results[label].Ti_all_rms_sumnr_sumn[ion]*100.0, ls[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]}") + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$Q_i$ (GB)") + ax.set_ylabel("$\\delta T_i/T_{i,0}/n_{i0}$ (%)") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ion energy fluxes (separate species)') - if plotLegend: - ax.legend(loc='best', prop={'size': 8},) + ax.set_title('Ions (all) Temperature intensity fluctuations') + ax.legend(loc='best', prop={'size': 8},) + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) - def plot_fluxes_ky(self, axs=None, label="", c="b", lw=1, plotLegend=True): + def plot_intensities_ky(self, axs=None, label="", c="b", addText=True): if axs is None: plt.ion() fig = plt.figure(figsize=(18, 9)) axs = fig.subplot_mosaic( """ - AC - BD + ACEG + BDFH """ ) ls = GRAPHICStools.listLS() - # Electron energy flux + # Potential intensity ax = axs["A"] - ax.plot(self.results[label].ky, self.results[label].Qe_ky_mean, '-o', markersize=5, color=c, label=label+' (mean)') - ax.fill_between(self.results[label].ky, self.results[label].Qe_ky_mean-self.results[label].Qe_ky_std, self.results[label].Qe_ky_mean+self.results[label].Qe_ky_std, color=c, alpha=0.2) + ax.plot(self.results[label].ky, self.results[label].phi_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].phi_rms_sumnr_mean-self.results[label].phi_rms_sumnr_std, self.results[label].phi_rms_sumnr_mean+self.results[label].phi_rms_sumnr_std, color=c, alpha=0.2) ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$Q_e$ (GB)") + ax.set_ylabel(r"$\delta\phi/\phi_0$") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Electron energy flux vs. $k_\\theta\\rho_s$') - if plotLegend: - ax.legend(loc='best', prop={'size': 8},) + ax.set_title('Potential intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) ax.axhline(0.0, color='k', ls='--', lw=1) - # Electron particle flux + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta\phi/\phi_0|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + # EM potential intensity ax = axs["B"] - ax.plot(self.results[label].ky, self.results[label].Ge_ky_mean, '-o', markersize=5, color=c, label=label+' (mean)') - ax.fill_between(self.results[label].ky, self.results[label].Ge_ky_mean-self.results[label].Ge_ky_std, self.results[label].Ge_ky_mean+self.results[label].Ge_ky_std, color=c, alpha=0.2) - + if 'apar' in self.results[label].__dict__: + ax.plot(self.results[label].ky, self.results[label].apar_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+', $A_\\parallel$ (mean)') + ax.fill_between(self.results[label].ky, self.results[label].apar_rms_sumnr_mean-self.results[label].apar_rms_sumnr_std, self.results[label].apar_rms_sumnr_mean+self.results[label].apar_rms_sumnr_std, color=c, alpha=0.2) + ax.plot(self.results[label].ky, self.results[label].bpar_rms_sumnr_mean, '--', markersize=5, color=c, label=label+', $B_\\parallel$ (mean)') + ax.fill_between(self.results[label].ky, self.results[label].bpar_rms_sumnr_mean-self.results[label].bpar_rms_sumnr_std, self.results[label].bpar_rms_sumnr_mean+self.results[label].bpar_rms_sumnr_std, color=c, alpha=0.2) + + ax.legend(loc='best', prop={'size': 8},) + ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\Gamma_e$ (GB)") + ax.set_ylabel(r"$\delta F_\parallel/F_{\parallel,0}$") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Electron particle flux vs. $k_\\theta\\rho_s$') - if plotLegend: - ax.legend(loc='best', prop={'size': 8},) + ax.set_title('EM potential intensity vs. $k_\\theta\\rho_s$') + ax.axhline(0.0, color='k', ls='--', lw=1) - # Ion energy flux + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta F_\parallel/F_{\parallel,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + # Electron particle intensity ax = axs["C"] - ax.plot(self.results[label].ky, self.results[label].Qi_ky_mean, '-o', markersize=5, color=c, label=label+' (mean)') - ax.fill_between(self.results[label].ky, self.results[label].Qi_ky_mean-self.results[label].Qi_ky_std, self.results[label].Qi_ky_mean+self.results[label].Qi_ky_std, color=c, alpha=0.2) + ax.plot(self.results[label].ky, self.results[label].ne_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].ne_rms_sumnr_mean-self.results[label].ne_rms_sumnr_std, self.results[label].ne_rms_sumnr_mean+self.results[label].ne_rms_sumnr_std, color=c, alpha=0.2) ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$Q_i$ (GB)") + ax.set_ylabel("$\\delta n_e/n_{e,0}$") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ion energy fluxes vs. $k_\\theta\\rho_s$') - if plotLegend: - ax.legend(loc='best', prop={'size': 8},) + ax.set_title('Electron particle intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) ax.axhline(0.0, color='k', ls='--', lw=1) - GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta n_e/n_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) - # Ion species energy fluxes + # Electron temperature intensity ax = axs["D"] - for j, i in enumerate(self.results[label].ions_flags): - ax.plot(self.results[label].ky, self.results[label].Qi_all_ky_mean[j],ls[j]+'o', markersize=5, color=c, label=f"{label}, {self.results[label].all_names[i]}") + ax.plot(self.results[label].ky, self.results[label].Te_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Te_rms_sumnr_mean-self.results[label].Te_rms_sumnr_std, self.results[label].Te_rms_sumnr_mean+self.results[label].Te_rms_sumnr_std, color=c, alpha=0.2) - ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$Q_i$ (GB)") + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta T_e/T_{e,0}$") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ion energy fluxes vs. $k_\\theta\\rho_s$(separate species)') - if plotLegend: - ax.legend(loc='best', prop={'size': 8},) + ax.set_title('Electron temperature intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta T_e/T_{e,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + # Ion particle intensity + ax = axs["E"] + ax.plot(self.results[label].ky, self.results[label].ni_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].ni_rms_sumnr_mean-self.results[label].ni_rms_sumnr_std, self.results[label].ni_rms_sumnr_mean+self.results[label].ni_rms_sumnr_std, color=c, alpha=0.2) - GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta n_i/n_{i,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion particle intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + # Add mathematical definitions text + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta n_i/n_{i,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + # Ion temperature intensity + ax = axs["F"] + ax.plot(self.results[label].ky, self.results[label].Ti_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') + ax.fill_between(self.results[label].ky, self.results[label].Ti_rms_sumnr_mean-self.results[label].Ti_rms_sumnr_std, self.results[label].Ti_rms_sumnr_mean+self.results[label].Ti_rms_sumnr_std, color=c, alpha=0.2) - def plot_intensities(self, axs = None, label= "cgyro1", c="b", addText=True): + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta T_i/T_{i,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ion temperature intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + if addText: + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n_r}|\delta T_i/T_{i,0}|^2\rangle}$', + transform=ax.transAxes, + fontsize=12, + verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + + # Ion particle intensity + ax = axs["G"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].ky, self.results[label].ni_all_rms_sumnr_mean[ion], ls[ion]+'o', markersize=5, color=c, label=f"{label}, {self.results[label].all_names[ion]} (mean)") + + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta n_i/n_{i,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ions (all) particle intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + + + # Ion temperature intensity + ax = axs["H"] + for ion in self.results[label].ions_flags: + ax.plot(self.results[label].ky, self.results[label].Ti_all_rms_sumnr_mean[ion], ls[ion]+'o', markersize=5, color=c, label=f"{label}, {self.results[label].all_names[ion]} (mean)") + + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\delta T_i/T_{i,0}$") + GRAPHICStools.addDenseAxis(ax) + ax.set_title('Ions (all) temperature intensity vs. $k_\\theta\\rho_s$') + ax.legend(loc='best', prop={'size': 8},) + ax.axhline(0.0, color='k', ls='--', lw=1) + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_intensities_kx(self, axs=None, label="", c="b", addText=True): if axs is None: plt.ion() fig = plt.figure(figsize=(18, 9)) axs = fig.subplot_mosaic( """ - ACEG - BDFH + AC + BD """ ) - - ls = GRAPHICStools.listLS() - - ax = axs["A"] - ax.plot(self.results[label].t, self.results[label].phi_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") - ax.plot(self.results[label].t, self.results[label].phi_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") - ax.plot(self.results[label].t, self.results[label].phi_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") - - ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\delta \\phi/\\phi_0$ (%)") - GRAPHICStools.addDenseAxis(ax) - ax.set_title('Potential intensity fluctuations') - ax.legend(loc='best', prop={'size': 8},) - - # Add mathematical definitions text - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta\phi/\phi_0|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) - ax = axs["B"] - if 'apar' in self.results[label].__dict__: - ax.plot(self.results[label].t, self.results[label].apar_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}, $A_\\parallel$") - ax.plot(self.results[label].t, self.results[label].bpar_rms_sumnr_sumn*100.0, '--', c=c, lw=2, label=f"{label}, $B_\\parallel$") - ax.legend(loc='best', prop={'size': 8},) + # Potential intensity + ax = axs["A"] + ax.plot(self.results[label].kx, self.results[label].phi_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+' (mean)') + ax.plot(self.results[label].kx, self.results[label].phi_rms_n0_mean, '-.', markersize=0.5, lw=0.5, color=c, label=label+', $n=0$ (mean)') + ax.plot(self.results[label].kx, self.results[label].phi_rms_sumn1_mean, '--', markersize=0.5, lw=0.5, color=c, label=label+', $n>0$ (mean)') - ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\delta F_\\parallel/F_{\\parallel,0}$ (%)") + ax.set_xlabel("$k_{x}$") + ax.set_ylabel("$\\delta \\phi/\\phi_0$") GRAPHICStools.addDenseAxis(ax) - ax.set_title('EM potential intensity fluctuations') + ax.set_title('Potential intensity vs kx') + ax.legend(loc='best', prop={'size': 8},) + ax.set_yscale('log') - # Add mathematical definitions text if addText: ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta F_\parallel/F_{\parallel,0}|^2\rangle}$', + r'$\sqrt{\langle\sum_{n}|\delta\phi/\phi_0|^2\rangle}$', transform=ax.transAxes, fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + # EM potential intensity + ax = axs["C"] + if 'apar' in self.results[label].__dict__: + ax.plot(self.results[label].kx, self.results[label].apar_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+', $A_\\parallel$ (mean)') + ax.plot(self.results[label].kx, self.results[label].bpar_rms_sumn_mean, '--', markersize=1.0, lw=1.0, color=c, label=label+', $B_\\parallel$ (mean)') + ax.legend(loc='best', prop={'size': 8},) - ax = axs["C"] - ax.plot(self.results[label].t, self.results[label].ne_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") - ax.plot(self.results[label].t, self.results[label].ne_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") - ax.plot(self.results[label].t, self.results[label].ne_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") - - ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\delta n_e/n_{e,0}/n_{e0}$ (%)") + + ax.set_xlabel("$k_{x}$") + ax.set_ylabel("$\\delta F_\\parallel/F_{\\parallel,0}$") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Electron Density intensity fluctuations') - ax.legend(loc='best', prop={'size': 8},) + ax.set_title('EM potential intensity vs kx') + ax.set_yscale('log') # Add mathematical definitions text if addText: ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta n_e/n_{e,0}|^2\rangle}$', + r'$\sqrt{\langle\sum_{n}|\delta F_\parallel/F_{\parallel,0}|^2\rangle}$', transform=ax.transAxes, fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + # Electron particle intensity + ax = axs["B"] + ax.plot(self.results[label].kx, self.results[label].ne_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+' (mean)') + ax.plot(self.results[label].kx, self.results[label].ne_rms_n0_mean, '-.', markersize=0.5, lw=0.5, color=c, label=label+', $n=0$ (mean)') + ax.plot(self.results[label].kx, self.results[label].ne_rms_sumn1_mean, '--', markersize=0.5, lw=0.5, color=c, label=label+', $n>0$ (mean)') - ax = axs["D"] - ax.plot(self.results[label].t, self.results[label].Te_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") - ax.plot(self.results[label].t, self.results[label].Te_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") - ax.plot(self.results[label].t, self.results[label].Te_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") - - ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\delta T_e/T_{e,0}/T_{e0}$ (%)") + ax.set_xlabel("$k_{x}$") + ax.set_ylabel("$\\delta n_e/n_{e,0}$") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Electron Temperature intensity fluctuations') + ax.set_title('Electron particle intensity vs kx') ax.legend(loc='best', prop={'size': 8},) + ax.set_yscale('log') # Add mathematical definitions text if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta T_e/T_{e,0}|^2\rangle}$', + ax.text(0.02, 0.95, + r'$\sqrt{\langle\sum_{n}|\delta n_e/n_{e,0}|^2\rangle}$', transform=ax.transAxes, fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + # Electron temperature intensity + ax = axs["D"] + ax.plot(self.results[label].kx, self.results[label].Te_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+' (mean)') + ax.plot(self.results[label].kx, self.results[label].Te_rms_n0_mean, '-.', markersize=0.5, lw=0.5, color=c, label=label+', $n=0$ (mean)') + ax.plot(self.results[label].kx, self.results[label].Te_rms_sumn1_mean, '--', markersize=0.5, lw=0.5, color=c, label=label+', $n>0$ (mean)') - - - - ax = axs["E"] - ax.plot(self.results[label].t, self.results[label].ni_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") - ax.plot(self.results[label].t, self.results[label].ni_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") - ax.plot(self.results[label].t, self.results[label].ni_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") - - ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\delta n_i/n_{i,0}/n_{i0}$ (%)") + ax.set_xlabel("$k_{x}$") + ax.set_ylabel("$\\delta T_e/T_{e,0}$") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ion Density intensity fluctuations') + ax.set_title('Electron temperature intensity vs kx') ax.legend(loc='best', prop={'size': 8},) - - # Add mathematical definitions text + ax.set_yscale('log') + if addText: ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta n_i/n_{i,0}|^2\rangle}$', + r'$\sqrt{\langle\sum_{n}|\delta T_e/T_{e,0}|^2\rangle}$', transform=ax.transAxes, fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + def plot_turbulence(self, axs = None, label= "cgyro1", c="b", kys = None): + + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + axs = fig.subplot_mosaic( + """ + AC + BD + """ + ) - ax = axs["F"] - ax.plot(self.results[label].t, self.results[label].Ti_rms_sumnr_sumn*100.0, '-', c=c, lw=2, label=f"{label}") - ax.plot(self.results[label].t, self.results[label].Ti_rms_sumnr_n0*100.0, '-.', c=c, lw=0.5, label=f"{label}, $n=0$") - ax.plot(self.results[label].t, self.results[label].Ti_rms_sumnr_sumn1*100.0, '--', c=c, lw=0.5, label=f"{label}, $n>0$") + # Is no kys provided, select just 3: first, last and middle + if kys is None: + ikys = [0] + if len(self.results[label].ky) > 1: + ikys.append(-1) + if len(self.results[label].ky) > 2: + ikys.append(len(self.results[label].ky) // 2) + + ikys = np.unique(ikys) + else: + ikys = [self.results[label].ky.index(ky) for ky in kys if ky in self.results[label].ky] + # Growth rate as function of time + ax = axs["A"] + for i,ky in enumerate(ikys): + self._plot_trace( + ax, + label, + self.results[label].g[ky, :], + c=c, + ls = GRAPHICStools.listLS()[i], + lw=1, + label_plot=f"$k_{{\\theta}}\\rho_s={np.abs(self.results[label].ky[ky]):.2f}$", + var_meanstd = [self.results[label].g_mean[ky], self.results[label].g_std[ky]], + ) + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\delta T_i/T_{i,0}/T_{i0}$ (%)") + ax.set_ylabel("$\\gamma$ (norm.)") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ion Temperature intensity fluctuations') + ax.set_title('Growth rate vs time') ax.legend(loc='best', prop={'size': 8},) - # Add mathematical definitions text - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}\sum_{n_r}|\delta T_i/T_{i,0}|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) - - - ax = axs["G"] - for ion in self.results[label].ions_flags: - ax.plot(self.results[label].t, self.results[label].ni_all_rms_sumnr_sumn[ion]*100.0, ls[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]}") - + # Frequency as function of time + ax = axs["B"] + for i,ky in enumerate(ikys): + self._plot_trace( + ax, + label, + self.results[label].f[ky, :], + c=c, + ls = GRAPHICStools.listLS()[i], + lw=1, + label_plot=f"$k_{{\\theta}}\\rho_s={np.abs(self.results[label].ky[ky]):.2f}$", + var_meanstd = [self.results[label].f_mean[ky], self.results[label].f_std[ky]], + ) + ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\delta n_i/n_{i,0}/n_{i0}$ (%)") + ax.set_ylabel("$\\omega$ (norm.)") GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ions (all) Density intensity fluctuations') + ax.set_title('Real Frequency vs time') ax.legend(loc='best', prop={'size': 8},) - - ax = axs["H"] - for ion in self.results[label].ions_flags: - ax.plot(self.results[label].t, self.results[label].Ti_all_rms_sumnr_sumn[ion]*100.0, ls[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]}") - - ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\delta T_i/T_{i,0}/n_{i0}$ (%)") + # Mean+Std Growth rate as function of ky + ax = axs["C"] + ax.errorbar(self.results[label].ky, self.results[label].g_mean, yerr=self.results[label].g_std, fmt='-o', markersize=5, color=c, label=label+' (mean+std)') + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\gamma$ (norm.)") + ax.set_title('Saturated Growth Rate') GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ions (all) Temperature intensity fluctuations') ax.legend(loc='best', prop={'size': 8},) - - + + # Mean+Std Frequency as function of ky + ax = axs["D"] + ax.errorbar(self.results[label].ky, self.results[label].f_mean, yerr=self.results[label].f_std, fmt='-o', markersize=5, color=c, label=label+' (mean+std)') + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\omega$ (norm.)") + ax.set_title('Saturated Real Frequency') + GRAPHICStools.addDenseAxis(ax) + ax.legend(loc='best', prop={'size': 8},) + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) - def plot_intensities_ky(self, axs=None, label="", c="b", addText=True): + def plot_cross_phases(self, axs = None, label= "cgyro1", c="b"): + if axs is None: plt.ion() fig = plt.figure(figsize=(18, 9)) @@ -1087,829 +1100,875 @@ def plot_intensities_ky(self, axs=None, label="", c="b", addText=True): ) ls = GRAPHICStools.listLS() - - # Potential intensity + m = GRAPHICStools.listmarkers() + ax = axs["A"] - ax.plot(self.results[label].ky, self.results[label].phi_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') - ax.fill_between(self.results[label].ky, self.results[label].phi_rms_sumnr_mean-self.results[label].phi_rms_sumnr_std, self.results[label].phi_rms_sumnr_mean+self.results[label].phi_rms_sumnr_std, color=c, alpha=0.2) + ax.plot(self.results[label].ky, self.results[label].neTe_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].neTe_kx0_mean-self.results[label].neTe_kx0_std, self.results[label].neTe_kx0_mean+self.results[label].neTe_kx0_std, color=c, alpha=0.2) ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel(r"$\delta\phi/\phi_0$") + ax.set_ylabel("$n_e-T_e$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) GRAPHICStools.addDenseAxis(ax) - ax.set_title('Potential intensity vs. $k_\\theta\\rho_s$') - ax.legend(loc='best', prop={'size': 8},) ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$n_e-T_e$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) - # Add mathematical definitions text - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n_r}|\delta\phi/\phi_0|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) - # EM potential intensity ax = axs["B"] - if 'apar' in self.results[label].__dict__: - ax.plot(self.results[label].ky, self.results[label].apar_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+', $A_\\parallel$ (mean)') - ax.fill_between(self.results[label].ky, self.results[label].apar_rms_sumnr_mean-self.results[label].apar_rms_sumnr_std, self.results[label].apar_rms_sumnr_mean+self.results[label].apar_rms_sumnr_std, color=c, alpha=0.2) - ax.plot(self.results[label].ky, self.results[label].bpar_rms_sumnr_mean, '--', markersize=5, color=c, label=label+', $B_\\parallel$ (mean)') - ax.fill_between(self.results[label].ky, self.results[label].bpar_rms_sumnr_mean-self.results[label].bpar_rms_sumnr_std, self.results[label].bpar_rms_sumnr_mean+self.results[label].bpar_rms_sumnr_std, color=c, alpha=0.2) - - ax.legend(loc='best', prop={'size': 8},) + ax.plot(self.results[label].ky, self.results[label].niTi_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].niTi_kx0_mean-self.results[label].niTi_kx0_std, self.results[label].niTi_kx0_mean+self.results[label].niTi_kx0_std, color=c, alpha=0.2) ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel(r"$\delta F_\parallel/F_{\parallel,0}$") + ax.set_ylabel("$n_i-T_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) GRAPHICStools.addDenseAxis(ax) - ax.set_title('EM potential intensity vs. $k_\\theta\\rho_s$') - ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$n_i-T_i$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) - # Add mathematical definitions text - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n_r}|\delta F_\parallel/F_{\parallel,0}|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) - - - # Electron particle intensity ax = axs["C"] - ax.plot(self.results[label].ky, self.results[label].ne_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') - ax.fill_between(self.results[label].ky, self.results[label].ne_rms_sumnr_mean-self.results[label].ne_rms_sumnr_std, self.results[label].ne_rms_sumnr_mean+self.results[label].ne_rms_sumnr_std, color=c, alpha=0.2) + ax.plot(self.results[label].ky, self.results[label].phine_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].phine_kx0_mean-self.results[label].phine_kx0_std, self.results[label].phine_kx0_mean+self.results[label].phine_kx0_std, color=c, alpha=0.2) ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\delta n_e/n_{e,0}$") + ax.set_ylabel("$\\phi-n_e$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) GRAPHICStools.addDenseAxis(ax) - ax.set_title('Electron particle intensity vs. $k_\\theta\\rho_s$') - ax.legend(loc='best', prop={'size': 8},) ax.axhline(0.0, color='k', ls='--', lw=1) - - # Add mathematical definitions text - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n_r}|\delta n_e/n_{e,0}|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + ax.set_title('$\\phi-n_e$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) - # Electron temperature intensity ax = axs["D"] - ax.plot(self.results[label].ky, self.results[label].Te_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') - ax.fill_between(self.results[label].ky, self.results[label].Te_rms_sumnr_mean-self.results[label].Te_rms_sumnr_std, self.results[label].Te_rms_sumnr_mean+self.results[label].Te_rms_sumnr_std, color=c, alpha=0.2) - + ax.plot(self.results[label].ky, self.results[label].phini_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].phini_kx0_mean-self.results[label].phini_kx0_std, self.results[label].phini_kx0_mean+self.results[label].phini_kx0_std, color=c, alpha=0.2) ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\delta T_e/T_{e,0}$") + ax.set_ylabel("$\\phi-n_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) GRAPHICStools.addDenseAxis(ax) - ax.set_title('Electron temperature intensity vs. $k_\\theta\\rho_s$') - ax.legend(loc='best', prop={'size': 8},) ax.axhline(0.0, color='k', ls='--', lw=1) - - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n_r}|\delta T_e/T_{e,0}|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) - - - # Ion particle intensity + ax.set_title('$\\phi-n_i$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + ax = axs["E"] - ax.plot(self.results[label].ky, self.results[label].ni_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') - ax.fill_between(self.results[label].ky, self.results[label].ni_rms_sumnr_mean-self.results[label].ni_rms_sumnr_std, self.results[label].ni_rms_sumnr_mean+self.results[label].ni_rms_sumnr_std, color=c, alpha=0.2) + ax.plot(self.results[label].ky, self.results[label].phiTe_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].phiTe_kx0_mean-self.results[label].phiTe_kx0_std, self.results[label].phiTe_kx0_mean+self.results[label].phiTe_kx0_std, color=c, alpha=0.2) ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\delta n_i/n_{i,0}$") + ax.set_ylabel("$\\phi-T_e$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ion particle intensity vs. $k_\\theta\\rho_s$') - ax.legend(loc='best', prop={'size': 8},) ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$\\phi-T_e$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) - # Add mathematical definitions text - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n_r}|\delta n_i/n_{i,0}|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) - # Ion temperature intensity ax = axs["F"] - ax.plot(self.results[label].ky, self.results[label].Ti_rms_sumnr_mean, '-o', markersize=5, color=c, label=label+' (mean)') - ax.fill_between(self.results[label].ky, self.results[label].Ti_rms_sumnr_mean-self.results[label].Ti_rms_sumnr_std, self.results[label].Ti_rms_sumnr_mean+self.results[label].Ti_rms_sumnr_std, color=c, alpha=0.2) - + ax.plot(self.results[label].ky, self.results[label].phiTi_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") + ax.fill_between(self.results[label].ky, self.results[label].phiTi_kx0_mean-self.results[label].phiTi_kx0_std, self.results[label].phiTi_kx0_mean+self.results[label].phiTi_kx0_std, color=c, alpha=0.2) ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\delta T_i/T_{i,0}$") + ax.set_ylabel("$\\phi-T_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ion temperature intensity vs. $k_\\theta\\rho_s$') - ax.legend(loc='best', prop={'size': 8},) ax.axhline(0.0, color='k', ls='--', lw=1) - - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n_r}|\delta T_i/T_{i,0}|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + ax.set_title('$\\phi-T_i$ cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) - # Ion particle intensity ax = axs["G"] for ion in self.results[label].ions_flags: - ax.plot(self.results[label].ky, self.results[label].ni_all_rms_sumnr_mean[ion], ls[ion]+'o', markersize=5, color=c, label=f"{label}, {self.results[label].all_names[ion]} (mean)") - + ax.plot(self.results[label].ky, self.results[label].phiTi_all_kx0_mean[ion], ls[ion]+m[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]} (mean)", markersize=4) ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\delta n_i/n_{i,0}$") + ax.set_ylabel("$\\phi-T_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ions (all) particle intensity vs. $k_\\theta\\rho_s$') - ax.legend(loc='best', prop={'size': 8},) ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$\\phi-T_i$ (all) cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) - - # Ion temperature intensity + ax = axs["H"] for ion in self.results[label].ions_flags: - ax.plot(self.results[label].ky, self.results[label].Ti_all_rms_sumnr_mean[ion], ls[ion]+'o', markersize=5, color=c, label=f"{label}, {self.results[label].all_names[ion]} (mean)") + ax.plot(self.results[label].ky, self.results[label].phini_all_kx0_mean[ion], ls[ion]+m[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]} (mean)", markersize=4) + + ax.set_xlabel("$k_{\\theta} \\rho_s$") + ax.set_ylabel("$\\phi-n_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) + GRAPHICStools.addDenseAxis(ax) + ax.axhline(0.0, color='k', ls='--', lw=1) + ax.set_title('$\\phi-n_i$ (all) cross-phase ($k_x=0$)') + ax.legend(loc='best', prop={'size': 8},) + + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_ballooning(self, time = None, label="cgyro1", c="b", axs=None): + + if axs is None: + plt.ion() + fig = plt.figure(figsize=(18, 9)) + + axs = fig.subplot_mosaic( + """ + 135 + 246 + """ + ) + + if time is None: + time = np.min([self.results[label].tmin, self.results[label].tmax_fluct]) + + it = np.argmin(np.abs(self.results[label].t - time)) + + colorsC, _ = GRAPHICStools.colorTableFade( + len(self.results[label].ky), + startcolor=c, + endcolor=c, + alphalims=[1.0, 0.4], + ) + + ax = axs['1'] + for ky in range(len(self.results[label].ky)): + for var, axsT in zip( + ["phi_ballooning", "apar_ballooning", "bpar_ballooning"], + [[axs['1'], axs['2']], [axs['3'], axs['4']], [axs['5'], axs['6']]], + ): + + f = self.results[label].__dict__[var][:, it] + y1 = np.real(f) + y2 = np.imag(f) + x = self.results[label].theta_ballooning / np.pi + + # Normalize + y1_max = np.max(np.abs(y1)) + y2_max = np.max(np.abs(y2)) + y1 /= y1_max + y2 /= y2_max + + ax = axsT[0] + ax.plot( + x, + y1, + color=colorsC[ky], + ls="-", + label=f"$k_{{\\theta}}\\rho_s={np.abs( self.results[label].ky[ky]):.2f}$ (max {y1_max:.2e})", + ) + ax = axsT[1] + ax.plot( + x, + y2, + color=colorsC[ky], + ls="-", + label=f"$k_{{\\theta}}\\rho_s={np.abs( self.results[label].ky[ky]):.2f}$ (max {y2_max:.2e})", + ) + + + ax = axs['1'] + ax.set_xlabel("$\\theta/\\pi$ (normalized to maximum)") + ax.set_ylabel("Re($\\delta\\phi$)") + ax.set_title("$\\delta\\phi$") + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + ax.set_xlim([-2 * np.pi, 2 * np.pi]) + + ax = axs['3'] + ax.set_xlabel("$\\theta/\\pi$ (normalized to maximum)") + ax.set_ylabel("Re($\\delta A\\parallel$)") + ax.set_title("$\\delta A\\parallel$") + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + ax = axs['5'] + ax.set_xlabel("$\\theta/\\pi$ (normalized to maximum)") + ax.set_ylabel("Re($\\delta B\\parallel$)") + ax.set_title("$\\delta B\\parallel$") + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + ax = axs['2'] + ax.set_xlabel("$\\theta/\\pi$") + ax.set_ylabel("Im($\\delta\\phi$)") + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + ax = axs['4'] + ax.set_xlabel("$\\theta/\\pi$") + ax.set_ylabel("Im($\\delta A\\parallel$)") + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + ax = axs['6'] + ax.set_xlabel("$\\theta/\\pi$") + ax.set_ylabel("Im($\\delta B\\parallel$)") + ax.legend(loc="best", prop={"size": 8}) + GRAPHICStools.addDenseAxis(ax) + + + for ax in [axs['1'], axs['3'], axs['5'], axs['2'], axs['4'], axs['6']]: + ax.axvline(x=0, lw=0.5, ls="--", c="k") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + + + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + + def plot_2D(self, label="cgyro1", axs=None, times = None): + + if times is None: + times = [] + + number_times = len(axs)//3 if axs is not None else 4 + + try: + times = [self.results[label].t[-1-i*10] for i in range(number_times)] + except IndexError: + times = [self.results[label].t[-1-i*1] for i in range(number_times)] + + if axs is None: + + mosaic = _2D_mosaic(len(times)) + + plt.ion() + fig = plt.figure(figsize=(18, 9)) + axs = fig.subplot_mosaic(mosaic) + + # Pre-calculate global min/max for each field type across all times + phi_values = [] + n_values = [] + e_values = [] + + for time in times: + it = np.argmin(np.abs(self.results[label].t - time)) + + # Get phi values + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_phi', it = it) + phi_values.append(fp) + + # Get n values + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_n',species = self.results[label].electron_flag, it = it) + n_values.append(fp) + + # Get e values + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_e',species = self.results[label].electron_flag, it = it) + e_values.append(fp) + + # Calculate global ranges + phi_max = np.max([np.max(np.abs(fp)) for fp in phi_values]) + phi_min, phi_max = -phi_max, +phi_max + + n_max = np.max([np.max(np.abs(fp)) for fp in n_values]) + n_min, n_max = -n_max, +n_max + + e_max = np.max([np.max(np.abs(fp)) for fp in e_values]) + e_min, e_max = -e_max, +e_max + colorbars = [] # Store colorbar references + # Now plot with consistent colorbar ranges + for time_i, time in enumerate(times): + + print(f"\t- Plotting 2D turbulence for {label} at time {time}") + + it = np.argmin(np.abs(self.results[label].t - time)) + + cfig = axs[str(time_i+1)].get_figure() + + # Phi plot + ax = axs[str(time_i+1)] + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_phi', it = it) - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\delta T_i/T_{i,0}$") - GRAPHICStools.addDenseAxis(ax) - ax.set_title('Ions (all) temperature intensity vs. $k_\\theta\\rho_s$') - ax.legend(loc='best', prop={'size': 8},) - ax.axhline(0.0, color='k', ls='--', lw=1) - - GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + cs1 = ax.contourf(xp,yp,np.transpose(fp),levels=np.arange(phi_min,phi_max,(phi_max-phi_min)/256),cmap=plt.get_cmap('jet')) + cphi = cfig.colorbar(cs1, ax=ax) - def plot_intensities_kx(self, axs=None, label="", c="b", addText=True): - if axs is None: - plt.ion() - fig = plt.figure(figsize=(18, 9)) + ax.set_xlabel("$x/\\rho_s$") + ax.set_ylabel("$y/\\rho_s$") + ax.set_title(f"$\\delta\\phi/\\phi_0$ (t={self.results[label].t[it]} $a/c_s$)") + ax.set_aspect('equal') - axs = fig.subplot_mosaic( - """ - AC - BD - """ - ) + # N plot + ax = axs[str(time_i+1+len(times))] + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_n',species = self.results[label].electron_flag, it = it) - # Potential intensity - ax = axs["A"] - ax.plot(self.results[label].kx, self.results[label].phi_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+' (mean)') - ax.plot(self.results[label].kx, self.results[label].phi_rms_n0_mean, '-.', markersize=0.5, lw=0.5, color=c, label=label+', $n=0$ (mean)') - ax.plot(self.results[label].kx, self.results[label].phi_rms_sumn1_mean, '--', markersize=0.5, lw=0.5, color=c, label=label+', $n>0$ (mean)') + cs2 = ax.contourf(xp,yp,np.transpose(fp),levels=np.arange(n_min,n_max,(n_max-n_min)/256),cmap=plt.get_cmap('jet')) + cn = cfig.colorbar(cs2, ax=ax) - ax.set_xlabel("$k_{x}$") - ax.set_ylabel("$\\delta \\phi/\\phi_0$") - GRAPHICStools.addDenseAxis(ax) - ax.set_title('Potential intensity vs kx') - ax.legend(loc='best', prop={'size': 8},) - ax.set_yscale('log') - - # Add mathematical definitions text - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}|\delta\phi/\phi_0|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + ax.set_xlabel("$x/\\rho_s$") + ax.set_ylabel("$y/\\rho_s$") + ax.set_title(f"$\\delta n_e/n_{{e,0}}$ (t={self.results[label].t[it]} $a/c_s$)") + ax.set_aspect('equal') - # EM potential intensity - ax = axs["C"] - if 'apar' in self.results[label].__dict__: - ax.plot(self.results[label].kx, self.results[label].apar_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+', $A_\\parallel$ (mean)') - ax.plot(self.results[label].kx, self.results[label].bpar_rms_sumn_mean, '--', markersize=1.0, lw=1.0, color=c, label=label+', $B_\\parallel$ (mean)') + # E plot + ax = axs[str(time_i+1+len(times)*2)] + xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_e',species = self.results[label].electron_flag, it = it) - ax.legend(loc='best', prop={'size': 8},) + cs3 = ax.contourf(xp,yp,np.transpose(fp),levels=np.arange(e_min,e_max,(e_max-e_min)/256),cmap=plt.get_cmap('jet')) + ce = cfig.colorbar(cs3, ax=ax) + ax.set_xlabel("$x/\\rho_s$") + ax.set_ylabel("$y/\\rho_s$") + ax.set_title(f"$\\delta E_e/E_{{e,0}}$ (t={self.results[label].t[it]} $a/c_s$)") + ax.set_aspect('equal') + + # Store the colorbar objects with their associated contour plots + colorbars.append({ + 'phi': cphi, + 'n': cn, + 'e': ce + }) - ax.set_xlabel("$k_{x}$") - ax.set_ylabel("$\\delta F_\\parallel/F_{\\parallel,0}$") - GRAPHICStools.addDenseAxis(ax) - ax.set_title('EM potential intensity vs kx') - ax.set_yscale('log') + GRAPHICStools.adjust_subplots(axs=axs, vertical=0.4, horizontal=0.3) - # Add mathematical definitions text - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}|\delta F_\parallel/F_{\parallel,0}|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + return colorbars + + def _to_real_space(self, variable = 'kxky_phi', species = None, label="cgyro1", theta_plot = 0, it = -1): + + # from pygacode + def maptoreal_fft(nr,nn,nx,ny,c): + d = np.zeros([nx,nn],dtype=complex) + for i in range(nr): + p = i-nr//2 + if -p < 0: + k = -p+nx + else: + k = -p + d[k,0:nn] = np.conj(c[i,0:nn]) + f = np.fft.irfft2(d,s=[nx,ny],norm='forward')*0.5 - # Electron particle intensity - ax = axs["B"] - ax.plot(self.results[label].kx, self.results[label].ne_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+' (mean)') - ax.plot(self.results[label].kx, self.results[label].ne_rms_n0_mean, '-.', markersize=0.5, lw=0.5, color=c, label=label+', $n=0$ (mean)') - ax.plot(self.results[label].kx, self.results[label].ne_rms_sumn1_mean, '--', markersize=0.5, lw=0.5, color=c, label=label+', $n>0$ (mean)') + # Correct for half-sum + f = 2*f - ax.set_xlabel("$k_{x}$") - ax.set_ylabel("$\\delta n_e/n_{e,0}$") - GRAPHICStools.addDenseAxis(ax) - ax.set_title('Electron particle intensity vs kx') - ax.legend(loc='best', prop={'size': 8},) - ax.set_yscale('log') + return f - # Add mathematical definitions text - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}|\delta n_e/n_{e,0}|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) - - # Electron temperature intensity - ax = axs["D"] - ax.plot(self.results[label].kx, self.results[label].Te_rms_sumn_mean, '-o', markersize=1.0, lw=1.0, color=c, label=label+' (mean)') - ax.plot(self.results[label].kx, self.results[label].Te_rms_n0_mean, '-.', markersize=0.5, lw=0.5, color=c, label=label+', $n=0$ (mean)') - ax.plot(self.results[label].kx, self.results[label].Te_rms_sumn1_mean, '--', markersize=0.5, lw=0.5, color=c, label=label+', $n>0$ (mean)') + # Real space + nr = self.results[label].cgyrodata.n_radial + nn = self.results[label].cgyrodata.n_n + craw = self.results[label].cgyrodata.__dict__[variable] + + itheta = np.argmin(np.abs(self.results[label].theta_stored-theta_plot)) + if species is None: + c = craw[:,itheta,:,it] + else: + c = craw[:,itheta,species,:,it] - ax.set_xlabel("$k_{x}$") - ax.set_ylabel("$\\delta T_e/T_{e,0}$") - GRAPHICStools.addDenseAxis(ax) - ax.set_title('Electron temperature intensity vs kx') - ax.legend(loc='best', prop={'size': 8},) - ax.set_yscale('log') + nx = self.results[label].cgyrodata.__dict__[variable].shape[0] + ny = nx - if addText: - ax.text(0.02, 0.95, - r'$\sqrt{\langle\sum_{n}|\delta T_e/T_{e,0}|^2\rangle}$', - transform=ax.transAxes, - fontsize=12, - verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + # Arrays + x = np.arange(nx)*2*np.pi/nx + y = np.arange(ny)*2*np.pi/ny + f = maptoreal_fft(nr,nn,nx,ny,c) - GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + # Physical maxima + ky1 = self.results[label].cgyrodata.ky[1] if len(self.results[label].cgyrodata.ky) > 1 else self.results[label].cgyrodata.ky[0] + xmax = self.results[label].cgyrodata.length + ymax = (2*np.pi)/np.abs(ky1) + xp = x/(2*np.pi)*xmax + yp = y/(2*np.pi)*ymax - def plot_turbulence(self, axs = None, label= "cgyro1", c="b", kys = None): + # Periodic extensions + xp = np.append(xp,xmax) + yp = np.append(yp,ymax) + fp = np.zeros([nx+1,ny+1]) + fp[0:nx,0:ny] = f[:,:] + fp[-1,:] = fp[0,:] + fp[:,-1] = fp[:,0] - if axs is None: - plt.ion() - fig = plt.figure(figsize=(18, 9)) - - axs = fig.subplot_mosaic( - """ - AC - BD - """ - ) + return xp, yp, fp + + def plot_quick_linear(self, labels=["cgyro1"], fig=None): + + colors = GRAPHICStools.listColors() - # Is no kys provided, select just 3: first, last and middle - if kys is None: - ikys = [0] - if len(self.results[label].ky) > 1: - ikys.append(-1) - if len(self.results[label].ky) > 2: - ikys.append(len(self.results[label].ky) // 2) - - ikys = np.unique(ikys) - else: - ikys = [self.results[label].ky.index(ky) for ky in kys if ky in self.results[label].ky] + if fig is None: + fig = plt.figure(figsize=(15,9)) - # Growth rate as function of time - ax = axs["A"] - for i,ky in enumerate(ikys): - self._plot_trace( - ax, - label, - self.results[label].g[ky, :], - c=c, - ls = GRAPHICStools.listLS()[i], - lw=1, - label_plot=f"$k_{{\\theta}}\\rho_s={np.abs(self.results[label].ky[ky]):.2f}$", - var_meanstd = [self.results[label].g_mean[ky], self.results[label].g_std[ky]], - ) + axs = fig.subplot_mosaic( + """ + 12 + 34 + """ + ) - ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\gamma$ (norm.)") - GRAPHICStools.addDenseAxis(ax) - ax.set_title('Growth rate vs time') - ax.legend(loc='best', prop={'size': 8},) + def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): + + for cont, label in enumerate(labels): + c = self.results[label] + baseColor = colors[cont+start_cont+1] + colorsC, _ = GRAPHICStools.colorTableFade( + len(c.ky), + startcolor=baseColor, + endcolor=baseColor, + alphalims=[1.0, 0.4], + ) + + ax = axs['1'] + for ky in range(len(c.ky)): + ax.plot( + c.t, + c.g[ky,:], + color=colorsC[ky], + label=f"$k_{{\\theta}}\\rho_s={np.abs(c.ky[ky]):.2f}$", + ) + + ax = axs['2'] + for ky in range(len(c.ky)): + ax.plot( + c.t, + c.f[ky,:], + color=colorsC[ky], + label=f"$k_{{\\theta}}\\rho_s={np.abs(c.ky[ky]):.2f}$", + ) - # Frequency as function of time - ax = axs["B"] - for i,ky in enumerate(ikys): - self._plot_trace( - ax, - label, - self.results[label].f[ky, :], - c=c, - ls = GRAPHICStools.listLS()[i], + GACODEplotting.plotTGLFspectrum( + [axs['3'], axs['4']], + self.results[label_base].ky, + self.results[label_base].g_mean, + freq=self.results[label_base].f_mean, + coeff=0.0, + c=col_lin, + ls="-", lw=1, - label_plot=f"$k_{{\\theta}}\\rho_s={np.abs(self.results[label].ky[ky]):.2f}$", - var_meanstd = [self.results[label].f_mean[ky], self.results[label].f_std[ky]], + label="", + facecolors=colors, + markersize=50, + alpha=1.0, + titles=["Growth Rate", "Real Frequency"], + removeLow=1e-4, + ylabel=True, ) - ax.set_xlabel("$t$ ($a/c_s$)"); #ax.set_xlim(left=0.0) - ax.set_ylabel("$\\omega$ (norm.)") - GRAPHICStools.addDenseAxis(ax) - ax.set_title('Real Frequency vs time') - ax.legend(loc='best', prop={'size': 8},) + return cont + - # Mean+Std Growth rate as function of ky - ax = axs["C"] - ax.errorbar(self.results[label].ky, self.results[label].g_mean, yerr=self.results[label].g_std, fmt='-o', markersize=5, color=c, label=label+' (mean+std)') - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\gamma$ (norm.)") - ax.set_title('Saturated Growth Rate') - GRAPHICStools.addDenseAxis(ax) - ax.legend(loc='best', prop={'size': 8},) - # Mean+Std Frequency as function of ky - ax = axs["D"] - ax.errorbar(self.results[label].ky, self.results[label].f_mean, yerr=self.results[label].f_std, fmt='-o', markersize=5, color=c, label=label+' (mean+std)') - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\omega$ (norm.)") - ax.set_title('Saturated Real Frequency') - GRAPHICStools.addDenseAxis(ax) - ax.legend(loc='best', prop={'size': 8},) + labels = self._labelize(labels) - GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) - - def plot_cross_phases(self, axs = None, label= "cgyro1", c="b"): + # Make it linear object for nice plotting + labels = self._kyfy(labels) + + co = -1 + for i,label0 in enumerate(labels): + if isinstance(self.results[label0], CGYROutils.CGYROlinear_scan): + co = _plot_linear_stability(axs, self.results[label0].labels, label0, start_cont=co, col_lin=colors[i]) + else: + co = _plot_linear_stability(axs, [label0], label0, start_cont=co, col_lin=colors[i]) - if axs is None: - plt.ion() - fig = plt.figure(figsize=(18, 9)) + ax = axs['1'] + ax.set_xlabel("Time $(a/c_s)$") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + ax.set_ylabel("$\\gamma$ $(c_s/a)$") + ax.set_title("Growth Rate") + ax.set_xlim(left=0) + ax.legend() + ax = axs['2'] + ax.set_xlabel("Time $(a/c_s)$") + ax.set_ylabel("$\\omega$ $(c_s/a)$") + ax.set_title("Real Frequency") + ax.axhline(y=0, lw=0.5, ls="--", c="k") + ax.set_xlim(left=0) - axs = fig.subplot_mosaic( - """ - ACEG - BDFH - """ + def _kyfy(self,labels_original): + ''' + This function transforms the original labels into the linear scan + e.g. from labels: + ['scan1_KY_0.3_0.5', + 'scan1_KY_0.3_0.7', + 'scan1_KY_0.4_0.5', + 'scan1_KY_0.4_0.7'] + to: + ['scan1_0.5', + 'scan1_0.7'] + where these are the CGYROlinear_scan object + ''' + + labelsD = {} + for label in labels_original: + parts = label.split('_') + if len(parts) >= 4 and parts[1] == "KY": + # Extract the base name (scan1), middle value (0.3/0.4), and last value (0.5/0.7) + base_name = parts[0] + middle_value = float(parts[2]) + last_value = parts[3] + + # Create the new key format: base_name + "_" + last_value + new_key = f"{base_name}_{last_value}" + + # Add the middle value to the list for this key + if new_key not in labelsD: + labelsD[new_key] = [] + labelsD[new_key].append(label) + + labels = [] + for label in labelsD: + self.results[label] = CGYROutils.CGYROlinear_scan(labelsD[label], self.results ) - - ls = GRAPHICStools.listLS() - m = GRAPHICStools.listmarkers() - - ax = axs["A"] - ax.plot(self.results[label].ky, self.results[label].neTe_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") - ax.fill_between(self.results[label].ky, self.results[label].neTe_kx0_mean-self.results[label].neTe_kx0_std, self.results[label].neTe_kx0_mean+self.results[label].neTe_kx0_std, color=c, alpha=0.2) + labels.append(label) - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$n_e-T_e$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) - GRAPHICStools.addDenseAxis(ax) - ax.axhline(0.0, color='k', ls='--', lw=1) - ax.set_title('$n_e-T_e$ cross-phase ($k_x=0$)') - ax.legend(loc='best', prop={'size': 8},) + return labels - ax = axs["B"] - ax.plot(self.results[label].ky, self.results[label].niTi_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") - ax.fill_between(self.results[label].ky, self.results[label].niTi_kx0_mean-self.results[label].niTi_kx0_std, self.results[label].niTi_kx0_mean+self.results[label].niTi_kx0_std, color=c, alpha=0.2) + # def prep(self, folder, inputgacode_file): - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$n_i-T_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) - GRAPHICStools.addDenseAxis(ax) - ax.axhline(0.0, color='k', ls='--', lw=1) - ax.set_title('$n_i-T_i$ cross-phase ($k_x=0$)') - ax.legend(loc='best', prop={'size': 8},) + # # Prepare main folder with input.gacode + # self.folder = IOtools.expandPath(folder) - ax = axs["C"] - ax.plot(self.results[label].ky, self.results[label].phine_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") - ax.fill_between(self.results[label].ky, self.results[label].phine_kx0_mean-self.results[label].phine_kx0_std, self.results[label].phine_kx0_mean+self.results[label].phine_kx0_std, color=c, alpha=0.2) + # self.folder.mkdir(parents=True, exist_ok=True) - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\phi-n_e$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) - GRAPHICStools.addDenseAxis(ax) - ax.axhline(0.0, color='k', ls='--', lw=1) - ax.set_title('$\\phi-n_e$ cross-phase ($k_x=0$)') - ax.legend(loc='best', prop={'size': 8},) + # self.inputgacode_file = self.folder / "input.gacode" + # if IOtools.expandPath(inputgacode_file) != self.inputgacode_file: + # shutil.copy2(IOtools.expandPath(inputgacode_file), self.inputgacode_file) - ax = axs["D"] - ax.plot(self.results[label].ky, self.results[label].phini_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") - ax.fill_between(self.results[label].ky, self.results[label].phini_kx0_mean-self.results[label].phini_kx0_std, self.results[label].phini_kx0_mean+self.results[label].phini_kx0_std, color=c, alpha=0.2) - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\phi-n_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) - GRAPHICStools.addDenseAxis(ax) - ax.axhline(0.0, color='k', ls='--', lw=1) - ax.set_title('$\\phi-n_i$ cross-phase ($k_x=0$)') - ax.legend(loc='best', prop={'size': 8},) + def _prerun( + self, + subfolder, + roa=0.55, + CGYROsettings=None, + extraOptions={}, + multipliers={}, + ): + self.folder = self.FolderGACODE / Path(subfolder) - ax = axs["E"] - ax.plot(self.results[label].ky, self.results[label].phiTe_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") - ax.fill_between(self.results[label].ky, self.results[label].phiTe_kx0_mean-self.results[label].phiTe_kx0_std, self.results[label].phiTe_kx0_mean+self.results[label].phiTe_kx0_std, color=c, alpha=0.2) + self.folder.mkdir(parents=True, exist_ok=True) - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\phi-T_e$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) - GRAPHICStools.addDenseAxis(ax) - ax.axhline(0.0, color='k', ls='--', lw=1) - ax.set_title('$\\phi-T_e$ cross-phase ($k_x=0$)') - ax.legend(loc='best', prop={'size': 8},) - + input_cgyro_file = self.folder / "input.cgyro" + inputCGYRO = CGYROinput(file=input_cgyro_file) - ax = axs["F"] - ax.plot(self.results[label].ky, self.results[label].phiTi_kx0_mean, '-o', c=c, lw=2, label=f"{label} (mean)") - ax.fill_between(self.results[label].ky, self.results[label].phiTi_kx0_mean-self.results[label].phiTi_kx0_std, self.results[label].phiTi_kx0_mean+self.results[label].phiTi_kx0_std, color=c, alpha=0.2) - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\phi-T_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) - GRAPHICStools.addDenseAxis(ax) - ax.axhline(0.0, color='k', ls='--', lw=1) - ax.set_title('$\\phi-T_i$ cross-phase ($k_x=0$)') - ax.legend(loc='best', prop={'size': 8},) - - - ax = axs["G"] - for ion in self.results[label].ions_flags: - ax.plot(self.results[label].ky, self.results[label].phiTi_all_kx0_mean[ion], ls[ion]+m[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]} (mean)", markersize=4) + inputgacode_file_this = self.folder / "input.gacode" + self.profiles.write_state(inputgacode_file_this) - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\phi-T_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) - GRAPHICStools.addDenseAxis(ax) - ax.axhline(0.0, color='k', ls='--', lw=1) - ax.set_title('$\\phi-T_i$ (all) cross-phase ($k_x=0$)') - ax.legend(loc='best', prop={'size': 8},) - + ResultsFiles_new = [] + for i in self.ResultsFiles: + if "mitim.out" not in i: + ResultsFiles_new.append(i) + self.ResultsFiles = ResultsFiles_new - ax = axs["H"] - for ion in self.results[label].ions_flags: - ax.plot(self.results[label].ky, self.results[label].phini_all_kx0_mean[ion], ls[ion]+m[ion], c=c, lw=1, label=f"{label}, {self.results[label].all_names[ion]} (mean)", markersize=4) - - ax.set_xlabel("$k_{\\theta} \\rho_s$") - ax.set_ylabel("$\\phi-n_i$ cross-phase (degrees)"); ax.set_ylim([-180, 180]) - GRAPHICStools.addDenseAxis(ax) - ax.axhline(0.0, color='k', ls='--', lw=1) - ax.set_title('$\\phi-n_i$ (all) cross-phase ($k_x=0$)') - ax.legend(loc='best', prop={'size': 8},) + inputCGYRO = GACODErun.modifyInputs( + inputCGYRO, + code_settings=CGYROsettings, + extraOptions=extraOptions, + multipliers=multipliers, + addControlFunction=GACODEdefaults.addCGYROcontrol, + rmin=roa, + controls_file = 'input.cgyro.controls' + ) + + inputCGYRO.write_state() + + return input_cgyro_file, inputgacode_file_this + + + def run_test( + self, + subfolder, + roa=0.55, + CGYROsettings=None, + extraOptions={}, + multipliers={}, + **kwargs + ): + if 'scan_param' in kwargs: + print("\t- Cannot run CGYRO tests with scan_param, running just the base",typeMsg="i") - GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + input_cgyro_file, inputgacode_file_this = self._prerun( + subfolder, + roa=roa, + CGYROsettings=CGYROsettings, + extraOptions=extraOptions, + multipliers=multipliers, + ) + + self.cgyro_job = FARMINGtools.mitim_job(self.folder) + + name = f'mitim_cgyro_{subfolder}_{roa:.6f}_test' + + self.cgyro_job.define_machine( + "cgyro", + name, + slurm_settings={ + "name": name, + "minutes": 5, + "cpuspertask": 1, + "ntasks": 1, + }, + ) + + CGYROcommand = "cgyro -t ." + + self.cgyro_job.prep( + CGYROcommand, + input_files=[input_cgyro_file, inputgacode_file_this], + output_files=self.output_files_test, + ) + + self.cgyro_job.run() + + def run1(self,subfolder,test_run=False,**kwargs): + if test_run: + self.run_test(subfolder,**kwargs) + else: + self.run_full(subfolder,**kwargs) + def run_full( + self, + subfolder, + roa=0.55, + CGYROsettings=None, + extraOptions={}, + multipliers={}, + scan_param = None, # {'variable': 'KY', 'values': [0.2,0.3,0.4]} + enforce_equality = None, # e.g. {'DLNTDR_SCALE_2': 'DLNTDR_SCALE_1', 'DLNTDR_SCALE_3': 'DLNTDR_SCALE_1'} + minutes = 5, + n = 16, + nomp = 1, + cpuspertask=None, # if None, will default to 1 + queue=None, #if blank will default to the one in settings + mem=None, # in MB + submit_via_qsub=True, #TODO fix this, works only at NERSC? no scans? + clean_folder_going_in=True, # Make sure the scratch folder is removed before running (unless I want a restart!) + submit_run=True, # False if I just want to check and fetch the job that was already submitted (e.g. via qsub or slurm) + ): + + input_cgyro_file, inputgacode_file_this = self._prerun( + subfolder, + roa=roa, + CGYROsettings=CGYROsettings, + extraOptions=extraOptions, + multipliers=multipliers, + ) + self.cgyro_job = FARMINGtools.mitim_job(self.folder) + name = f'mitim_cgyro_{subfolder}_{roa:.6f}' + if scan_param is not None and submit_via_qsub: + raise Exception(" Cannot use scan_param with submit_via_qsub=True, because it requires a different job for each value of the scan parameter.") - def plot_ballooning(self, time = None, label="cgyro1", c="b", axs=None): - - if axs is None: - plt.ion() - fig = plt.figure(figsize=(18, 9)) + if submit_via_qsub: - axs = fig.subplot_mosaic( - """ - 135 - 246 - """ + self.cgyro_job.define_machine( + "cgyro", + name, + launchSlurm=False, ) - if time is None: - time = np.min([self.results[label].tmin, self.results[label].tmax_fluct]) - - it = np.argmin(np.abs(self.results[label].t - time)) + subfolder = "scan0" + queue = "-queue " + self.cgyro_job.machineSettings['slurm']['partition'] if "partition" in self.cgyro_job.machineSettings['slurm'] else "" - colorsC, _ = GRAPHICStools.colorTableFade( - len(self.results[label].ky), - startcolor=c, - endcolor=c, - alphalims=[1.0, 0.4], - ) + if mem is not None: + CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} -mem {mem} {queue} -w 0:{minutes}:00 -s' + else: + CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} {queue} -w 0:{minutes}:00 -s' - ax = axs['1'] - for ky in range(len(self.results[label].ky)): - for var, axsT in zip( - ["phi_ballooning", "apar_ballooning", "bpar_ballooning"], - [[axs['1'], axs['2']], [axs['3'], axs['4']], [axs['5'], axs['6']]], - ): + if "account" in self.cgyro_job.machineSettings["slurm"] and self.cgyro_job.machineSettings["slurm"]["account"] is not None: + CGYROcommand += f" -repo {self.cgyro_job.machineSettings['slurm']['account']}" - f = self.results[label].__dict__[var][:, it] - y1 = np.real(f) - y2 = np.imag(f) - x = self.results[label].theta_ballooning / np.pi + self.slurm_output = "batch.out" - # Normalize - y1_max = np.max(np.abs(y1)) - y2_max = np.max(np.abs(y2)) - y1 /= y1_max - y2 /= y2_max + # --- + folder_run = self.folder / subfolder + folder_run.mkdir(parents=True, exist_ok=True) - ax = axsT[0] - ax.plot( - x, - y1, - color=colorsC[ky], - ls="-", - label=f"$k_{{\\theta}}\\rho_s={np.abs( self.results[label].ky[ky]):.2f}$ (max {y1_max:.2e})", - ) - ax = axsT[1] - ax.plot( - x, - y2, - color=colorsC[ky], - ls="-", - label=f"$k_{{\\theta}}\\rho_s={np.abs( self.results[label].ky[ky]):.2f}$ (max {y2_max:.2e})", - ) + # Copy the input.cgyro in the subfolder + input_cgyro_file_this = folder_run / "input.cgyro" + shutil.copy2(input_cgyro_file, input_cgyro_file_this) + # Copy the input.gacode file in the subfolder + inputgacode_file_this = folder_run / "input.gacode" + shutil.copy2(self.inputgacode_file, inputgacode_file_this) - ax = axs['1'] - ax.set_xlabel("$\\theta/\\pi$ (normalized to maximum)") - ax.set_ylabel("Re($\\delta\\phi$)") - ax.set_title("$\\delta\\phi$") - ax.legend(loc="best", prop={"size": 8}) - GRAPHICStools.addDenseAxis(ax) + # Prepare the input and output folders + input_folders = [folder_run] + output_folders = [subfolder] - ax.set_xlim([-2 * np.pi, 2 * np.pi]) + else: - ax = axs['3'] - ax.set_xlabel("$\\theta/\\pi$ (normalized to maximum)") - ax.set_ylabel("Re($\\delta A\\parallel$)") - ax.set_title("$\\delta A\\parallel$") - ax.legend(loc="best", prop={"size": 8}) - GRAPHICStools.addDenseAxis(ax) + if scan_param is None: + job_array = None + folder = 'scan0' + scan_param = {'variable': None, 'values': [0]} # Dummy scan parameter to avoid issues with the code below + else: + # Array + job_array = '' + for i,value in enumerate(scan_param['values']): + if job_array != '': + job_array += ',' + job_array += str(i) - ax = axs['5'] - ax.set_xlabel("$\\theta/\\pi$ (normalized to maximum)") - ax.set_ylabel("Re($\\delta B\\parallel$)") - ax.set_title("$\\delta B\\parallel$") - ax.legend(loc="best", prop={"size": 8}) - GRAPHICStools.addDenseAxis(ax) + folder = 'scan"$SLURM_ARRAY_TASK_ID"' - ax = axs['2'] - ax.set_xlabel("$\\theta/\\pi$") - ax.set_ylabel("Im($\\delta\\phi$)") - ax.legend(loc="best", prop={"size": 8}) - GRAPHICStools.addDenseAxis(ax) + # Machine + self.cgyro_job.define_machine( + "cgyro", + name, + slurm_settings={ + "name": name, + "minutes": minutes, + "ntasks": n, + "job_array": job_array, + # Validate n and nomp before assigning cpuspertask + }, + ) - ax = axs['4'] - ax.set_xlabel("$\\theta/\\pi$") - ax.set_ylabel("Im($\\delta A\\parallel$)") - ax.legend(loc="best", prop={"size": 8}) - GRAPHICStools.addDenseAxis(ax) + if cpuspertask is not None: + if not isinstance(cpuspertask, int): + raise TypeError(" cpuspertask must be an integer") + self.cgyro_job.slurm_settings["cpuspertask"] = cpuspertask + - ax = axs['6'] - ax.set_xlabel("$\\theta/\\pi$") - ax.set_ylabel("Im($\\delta B\\parallel$)") - ax.legend(loc="best", prop={"size": 8}) - GRAPHICStools.addDenseAxis(ax) + if queue is not None: + self.cgyro_job.machineSettings['slurm']['partition'] = queue + if mem is not None: + self.cgyro_job.slurm_settings['mem'] = mem - for ax in [axs['1'], axs['3'], axs['5'], axs['2'], axs['4'], axs['6']]: - ax.axvline(x=0, lw=0.5, ls="--", c="k") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - - - GRAPHICStools.adjust_subplots(axs=axs, vertical=0.3, horizontal=0.3) + # if not self.cgyro_job.launchSlurm: + # raise Exception(" Cannot run CGYRO scans without slurm") - def plot_2D(self, label="cgyro1", axs=None, times = None): - - if times is None: - times = [] - - number_times = len(axs)//3 if axs is not None else 4 + # Command to run cgyro + CGYROcommand = f'cgyro -e {folder} -n {n} -nomp {nomp} -p {self.cgyro_job.folderExecution}' - try: - times = [self.results[label].t[-1-i*10] for i in range(number_times)] - except IndexError: - times = [self.results[label].t[-1-i*1] for i in range(number_times)] + # Scans + input_folders = [] + output_folders = [] + for i,value in enumerate(scan_param['values']): + subfolder = f"scan{i}" + folder_run = self.folder / subfolder + folder_run.mkdir(parents=True, exist_ok=True) - if axs is None: + # Copy the input.cgyro in the subfolder + input_cgyro_file_this = folder_run / "input.cgyro" + shutil.copy2(input_cgyro_file, input_cgyro_file_this) - mosaic = _2D_mosaic(len(times)) + # Modify the input.cgyro file with the scan parameter + extraOptions_this = extraOptions.copy() + if scan_param['variable'] is not None: + extraOptions_this[scan_param['variable']] = value - plt.ion() - fig = plt.figure(figsize=(18, 9)) - axs = fig.subplot_mosaic(mosaic) - # Pre-calculate global min/max for each field type across all times - phi_values = [] - n_values = [] - e_values = [] - - for time in times: - it = np.argmin(np.abs(self.results[label].t - time)) - - # Get phi values - xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_phi', it = it) - phi_values.append(fp) - - # Get n values - xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_n',species = self.results[label].electron_flag, it = it) - n_values.append(fp) - - # Get e values - xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_e',species = self.results[label].electron_flag, it = it) - e_values.append(fp) - - # Calculate global ranges - phi_max = np.max([np.max(np.abs(fp)) for fp in phi_values]) - phi_min, phi_max = -phi_max, +phi_max - - n_max = np.max([np.max(np.abs(fp)) for fp in n_values]) - n_min, n_max = -n_max, +n_max - - e_max = np.max([np.max(np.abs(fp)) for fp in e_values]) - e_min, e_max = -e_max, +e_max + # If there is an enforce_equality, apply it + if enforce_equality is not None: + for key in enforce_equality: + extraOptions_this[key] = extraOptions_this[enforce_equality[key]] + + inputCGYRO = CGYROinput(file=input_cgyro_file_this) + input_cgyro_file_this = GACODErun.modifyInputs( + inputCGYRO, + code_settings=CGYROsettings, + extraOptions=extraOptions_this, + multipliers=multipliers, + addControlFunction=GACODEdefaults.addCGYROcontrol, + rmin=roa, + control_file = 'input.cgyro.controls' + ) - colorbars = [] # Store colorbar references - # Now plot with consistent colorbar ranges - for time_i, time in enumerate(times): - - print(f"\t- Plotting 2D turbulence for {label} at time {time}") - - it = np.argmin(np.abs(self.results[label].t - time)) - - cfig = axs[str(time_i+1)].get_figure() - - # Phi plot - ax = axs[str(time_i+1)] - xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_phi', it = it) + input_cgyro_file_this.write_state() - cs1 = ax.contourf(xp,yp,np.transpose(fp),levels=np.arange(phi_min,phi_max,(phi_max-phi_min)/256),cmap=plt.get_cmap('jet')) - cphi = cfig.colorbar(cs1, ax=ax) + # Copy the input.gacode file in the subfolder + inputgacode_file_this = folder_run / "input.gacode" + shutil.copy2(self.inputgacode_file, inputgacode_file_this) - ax.set_xlabel("$x/\\rho_s$") - ax.set_ylabel("$y/\\rho_s$") - ax.set_title(f"$\\delta\\phi/\\phi_0$ (t={self.results[label].t[it]} $a/c_s$)") - ax.set_aspect('equal') + # Prepare the input and output folders + input_folders.append(folder_run) + output_folders.append(subfolder) - # N plot - ax = axs[str(time_i+1+len(times))] - xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_n',species = self.results[label].electron_flag, it = it) + self.slurm_output = "slurm_output.dat" - cs2 = ax.contourf(xp,yp,np.transpose(fp),levels=np.arange(n_min,n_max,(n_max-n_min)/256),cmap=plt.get_cmap('jet')) - cn = cfig.colorbar(cs2, ax=ax) + # First submit the job with gacode_qsub, which will submit the cgyro job via slurm, with name + self.cgyro_job.prep( + CGYROcommand, + input_folders = input_folders, + output_folders=output_folders, + ) - ax.set_xlabel("$x/\\rho_s$") - ax.set_ylabel("$y/\\rho_s$") - ax.set_title(f"$\\delta n_e/n_{{e,0}}$ (t={self.results[label].t[it]} $a/c_s$)") - ax.set_aspect('equal') + if submit_run: + self.cgyro_job.run( + waitYN=False, + check_if_files_received=False, + removeScratchFolders=False, + removeScratchFolders_goingIn=clean_folder_going_in, + ) - # E plot - ax = axs[str(time_i+1+len(times)*2)] - xp, yp, fp = self._to_real_space(label=label, variable = 'kxky_e',species = self.results[label].electron_flag, it = it) + # Prepare how to search for the job without waiting for it + name_default_submission_qsub = Path(self.cgyro_job.folderExecution).name - cs3 = ax.contourf(xp,yp,np.transpose(fp),levels=np.arange(e_min,e_max,(e_max-e_min)/256),cmap=plt.get_cmap('jet')) - ce = cfig.colorbar(cs3, ax=ax) + self.cgyro_job.launchSlurm = True + self.cgyro_job.slurm_settings['name'] = name_default_submission_qsub - ax.set_xlabel("$x/\\rho_s$") - ax.set_ylabel("$y/\\rho_s$") - ax.set_title(f"$\\delta E_e/E_{{e,0}}$ (t={self.results[label].t[it]} $a/c_s$)") - ax.set_aspect('equal') - - # Store the colorbar objects with their associated contour plots - colorbars.append({ - 'phi': cphi, - 'n': cn, - 'e': ce - }) - GRAPHICStools.adjust_subplots(axs=axs, vertical=0.4, horizontal=0.3) + def check(self, every_n_minutes=5): - return colorbars - - - def _to_real_space(self, variable = 'kxky_phi', species = None, label="cgyro1", theta_plot = 0, it = -1): - - # from pygacode - def maptoreal_fft(nr,nn,nx,ny,c): + if self.cgyro_job.launchSlurm: + print("- Checker job status") - d = np.zeros([nx,nn],dtype=complex) - for i in range(nr): - p = i-nr//2 - if -p < 0: - k = -p+nx + while True: + self.cgyro_job.check(file_output = self.slurm_output) + print(f'\t- Current status (as of {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}): {self.cgyro_job.status} ({self.cgyro_job.infoSLURM["STATE"]})') + if self.cgyro_job.status == 2: + break else: - k = -p - d[k,0:nn] = np.conj(c[i,0:nn]) - f = np.fft.irfft2(d,s=[nx,ny],norm='forward')*0.5 + print(f"\t- Waiting {every_n_minutes} minutes") + time.sleep(every_n_minutes * 60) + else: + print("- Not checking status because this was run command line (not slurm)") - # Correct for half-sum - f = 2*f + print("\n\t* Job considered finished",typeMsg="i") - return f + def fetch(self): + """ + For a job that has been submitted but not waited for, once it is done, get the results + """ - # Real space - nr = self.results[label].cgyrodata.n_radial - nn = self.results[label].cgyrodata.n_n - craw = self.results[label].cgyrodata.__dict__[variable] - - itheta = np.argmin(np.abs(self.results[label].theta_stored-theta_plot)) - if species is None: - c = craw[:,itheta,:,it] + print("\n\n\t- Fetching results") + + if self.cgyro_job.launchSlurm: + self.cgyro_job.connect() + self.cgyro_job.retrieve() + self.cgyro_job.close() else: - c = craw[:,itheta,species,:,it] + print("- Not retrieving results because this was run command line (not slurm)") - nx = self.results[label].cgyrodata.__dict__[variable].shape[0] - ny = nx - - # Arrays - x = np.arange(nx)*2*np.pi/nx - y = np.arange(ny)*2*np.pi/ny - f = maptoreal_fft(nr,nn,nx,ny,c) - - # Physical maxima - ky1 = self.results[label].cgyrodata.ky[1] if len(self.results[label].cgyrodata.ky) > 1 else self.results[label].cgyrodata.ky[0] - xmax = self.results[label].cgyrodata.length - ymax = (2*np.pi)/np.abs(ky1) - xp = x/(2*np.pi)*xmax - yp = y/(2*np.pi)*ymax + def delete(self): - # Periodic extensions - xp = np.append(xp,xmax) - yp = np.append(yp,ymax) - fp = np.zeros([nx+1,ny+1]) - fp[0:nx,0:ny] = f[:,:] - fp[-1,:] = fp[0,:] - fp[:,-1] = fp[:,0] - - return xp, yp, fp - - def plot_quick_linear(self, labels=["cgyro1"], fig=None): - colors = GRAPHICStools.listColors() + print("\n\n\t- Deleting job") - if fig is None: - fig = plt.figure(figsize=(15,9)) + self.cgyro_job.launchSlurm = False - axs = fig.subplot_mosaic( - """ - 12 - 34 - """ + self.cgyro_job.prep( + f"scancel -n {self.cgyro_job.slurm_settings['name']}", + label_log_files="_finish", ) - - def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): - for cont, label in enumerate(labels): - c = self.results[label] - baseColor = colors[cont+start_cont+1] - colorsC, _ = GRAPHICStools.colorTableFade( - len(c.ky), - startcolor=baseColor, - endcolor=baseColor, - alphalims=[1.0, 0.4], - ) + self.cgyro_job.run() - ax = axs['1'] - for ky in range(len(c.ky)): - ax.plot( - c.t, - c.g[ky,:], - color=colorsC[ky], - label=f"$k_{{\\theta}}\\rho_s={np.abs(c.ky[ky]):.2f}$", - ) + def read1(self, label="cgyro1", folder=None, tmin = 0.0, minimal = False, last_tmin_for_linear = True): - ax = axs['2'] - for ky in range(len(c.ky)): - ax.plot( - c.t, - c.f[ky,:], - color=colorsC[ky], - label=f"$k_{{\\theta}}\\rho_s={np.abs(c.ky[ky]):.2f}$", - ) + folder = IOtools.expandPath(folder) if folder is not None else self.folder - GACODEplotting.plotTGLFspectrum( - [axs['3'], axs['4']], - self.results[label_base].ky, - self.results[label_base].g_mean, - freq=self.results[label_base].f_mean, - coeff=0.0, - c=col_lin, - ls="-", - lw=1, - label="", - facecolors=colors, - markersize=50, - alpha=1.0, - titles=["Growth Rate", "Real Frequency"], - removeLow=1e-4, - ylabel=True, - ) - - return cont + folders = sorted(list((folder).glob("scan*")), key=lambda f: int(''.join(filter(str.isdigit, f.name)))) + + if len(folders) == 0: + folders = [folder] + attach_name = False + else: + print(f"\t- Found {len(folders)} scan folders in {folder.resolve()}:") + for f in folders: + print(f"\t\t- {f.name}") + attach_name = True - co = -1 - for i,label0 in enumerate(labels): - if isinstance(self.results[label0], CGYROutils.CGYROlinear_scan): - co = _plot_linear_stability(axs, self.results[label0].labels, label0, start_cont=co, col_lin=colors[i]) + data = {} + labels = [] + for folder in folders: + + if attach_name: + label1 = f"{label}_{folder.name}" else: - co = _plot_linear_stability(axs, [label0], label0, start_cont=co, col_lin=colors[i]) + label1 = label + + data[label1] = CGYROutils.CGYROout(folder, tmin=tmin, minimal=minimal, last_tmin_for_linear=last_tmin_for_linear) + labels.append(label1) + + self.results.update(data) + + if attach_name: + self.results[label] = CGYROutils.CGYROlinear_scan(labels, data) - ax = axs['1'] - ax.set_xlabel("Time $(a/c_s)$") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - ax.set_ylabel("$\\gamma$ $(c_s/a)$") - ax.set_title("Growth Rate") - ax.set_xlim(left=0) - ax.legend() - ax = axs['2'] - ax.set_xlabel("Time $(a/c_s)$") - ax.set_ylabel("$\\omega$ $(c_s/a)$") - ax.set_title("Real Frequency") - ax.axhline(y=0, lw=0.5, ls="--", c="k") - ax.set_xlim(left=0) class CGYROinput(GACODErun.GACODEinput): diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 4b0d25e3..211d6efd 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -1,6 +1,7 @@ import os import scipy import numpy as np +from pathlib import Path import statsmodels.api as sm import matplotlib.pyplot as plt from mitim_tools.misc_tools import IOtools @@ -56,6 +57,9 @@ def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for self.folder = folder self.tmin = tmin + + if isinstance(self.folder, str): + self.folder = Path(self.folder) self.cgyrodata = self.read_using_cgyroplot(self.folder, suffix) @@ -143,7 +147,7 @@ def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for self.remove_symlinks() def read_using_cgyroplot(self, folder, suffix): - + original_dir = os.getcwd() # Handle files with suffix by creating temporary symbolic links diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index b8a2fe1b..6d949a9a 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -495,10 +495,13 @@ def read( ) # Unnormalize - SIMout.unnormalize( - self.NormalizationSets["SELECTED"], - rho=rho, - ) + if 'NormalizationSets' in self.__dict__: + SIMout.unnormalize( + self.NormalizationSets["SELECTED"], + rho=rho, + ) + else: + print("No normalization sets found.") self.results[label][class_output[1]].append(SIMout) diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index 1b0ebc0b..0265e66a 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -14,26 +14,46 @@ cgyro = CGYROtools.CGYRO(rhos = [0.5, 0.7]) -cgyro.prep( - gacode_file, - folder) +cgyro.prep(gacode_file,folder) + +# --------------- +# Standalone run +# --------------- cgyro.run( 'linear', code_settings=0, extraOptions={ - 'KY':0.3, + 'KY':0.5, 'MAX_TIME': 10.0, # Short, I just want to test the run }, - slurm_setup={'cores':4} - #submit_via_qsub=False # NERSC: True #TODO change this + slurm_setup={ + 'cores':8 + } ) +cgyro.read(label="cgyro1") +cgyro.plot(labels=["cgyro1"]) -# cgyro.check(every_n_minutes=1) -# cgyro.fetch() -# cgyro.delete() +# --------------- +# Scan of KY +# --------------- -cgyro.read(label="cgyro1") +cgyro.run_scan( + 'scan1', + cold_start=cold_start, + extraOptions={ + 'MAX_TIME': 10.0, + }, + variable='KY', + varUpDown=[0.3,0.4], + slurm_setup={ + 'cores':4 + } + ) + +cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) + +fig = cgyro.fn.add_figure(label="Quick linear") +cgyro.plot_quick_linear(labels=["scan1_KY_0.3","scan1_KY_0.4"], fig = fig) -cgyro.plot(labels=["cgyro1_0.5","cgyro1_0.7"]) cgyro.fn.show() From 77ad435bdbc1be321fee203c3220e00fce94dcb1 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 23 Aug 2025 16:39:52 -0400 Subject: [PATCH 191/385] Towards a PORTALS-CGYRO direct interaction --- .../physics_models/transport_cgyroneo.py | 58 +++++++++-------- .../gacode_tools/utils/GACODErun.py | 63 ++++++++++--------- templates/input.cgyro.models.json | 6 ++ tests/CGYRO_workflow.py | 8 ++- 4 files changed, 80 insertions(+), 55 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index 6b99119c..f7021545 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -9,36 +9,44 @@ def __init__(self, powerstate, **kwargs): def evaluate_turbulence(self): - pass - # # ------------------------------------------------------------------------------------------------------------------------ - # # Prepare CGYRO object - # # ------------------------------------------------------------------------------------------------------------------------ + rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] - # rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] + transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] - # tglf = TGLFtools.CGYRO(rhos=rho_locations) + cold_start = transport_evaluator_options.get("cold_start", False) - # _ = tglf.prep( - # self.powerstate.profiles_transport, - # self.folder, - # cold_start = cold_start, - # ) + # ------------------------------------------------------------------------------------------------------------------------ + # Prepare CGYRO object + # ------------------------------------------------------------------------------------------------------------------------ + + rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] + + cgyro = CGYROtools.CGYRO(rhos=rho_locations) + _ = cgyro.prep( + self.powerstate.profiles_transport, + self.folder, + ) - # cgyro = CGYROtools.CGYRO() + cgyro = CGYROtools.CGYRO( + rhos = rho_locations + ) - # cgyro.prep( - # self.folder, - # self.powerstate.profiles_transport.files[0], - # ) + cgyro.prep( + self.powerstate.profiles_transport.files[0], + self.folder, + ) - # cgyro.run( - # 'base_cgyro', - # roa = 0.55, - # CGYROsettings=0, - # submit_run=False - # ) + _ = cgyro.run( + 'base_cgyro', + full_submission=False, + code_settings=1, + cold_start=cold_start, + forceIfcold_start=True, + ) - # embed() - # #cgyro.read(label='base') - \ No newline at end of file + # cgyro.read( + # label='base_cgyro' + # ) + + embed() \ No newline at end of file diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 6d949a9a..bca9900d 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -115,6 +115,7 @@ def run( slurm_setup=None, # Cores per call (so, when running nR radii -> nR*4) attempts_execution=1, only_minimal_files=False, + full_submission=True # Flag to submit the job or just prepare everything but do not submit via MITIM ): if slurm_setup is None: @@ -163,6 +164,7 @@ def run( slurm_setup=slurm_setup, only_minimal_files=only_minimal_files, attempts_execution=attempts_execution, + full_submission=full_submission ) return code_executor_full @@ -300,6 +302,7 @@ def _check_cores(self, rhosEvaluate, slurm_setup, warning = 32 * 2): def _run( self, code_executor, + full_submission=True, **kwargs_run ): """ @@ -326,6 +329,7 @@ def _run( name=f"{self.run_specifications['code']}_{self.nameRunid}{kwargs_run.get('extra_name', '')}", launchSlurm=kwargs_run.get("launchSlurm", True), attempts_execution=kwargs_run.get("attempts_execution", 1), + full_submission=full_submission, ) else: print(f"\t- {self.run_specifications['code'].upper()} not run because all results files found (please ensure consistency!)",typeMsg="i") @@ -729,6 +733,7 @@ def run_gacode_simulation( launchSlurm = True, attempts_execution = 1, max_jobs_at_once = None, + full_submission = True, ): """ @@ -929,43 +934,45 @@ def run_gacode_simulation( shellPostCommands=shellPostCommands, ) - gacode_job.run( - removeScratchFolders=True, - attempts_execution=attempts_execution - ) + if full_submission: + + gacode_job.run( + removeScratchFolders=True, + attempts_execution=attempts_execution + ) - # --------------------------------------------- - # Organize - # --------------------------------------------- + # --------------------------------------------- + # Organize + # --------------------------------------------- - print("\t- Retrieving files and changing names for storing") - fineall = True - for subfolder_sim in code_executor: + print("\t- Retrieving files and changing names for storing") + fineall = True + for subfolder_sim in code_executor: - for i, rho in enumerate(code_executor[subfolder_sim].keys()): - for file in filesToRetrieve: - original_file = f"{file}_{rho:.4f}{extraFlag}" - final_destination = ( - code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" - ) - final_destination.unlink(missing_ok=True) + for i, rho in enumerate(code_executor[subfolder_sim].keys()): + for file in filesToRetrieve: + original_file = f"{file}_{rho:.4f}{extraFlag}" + final_destination = ( + code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" + ) + final_destination.unlink(missing_ok=True) - temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" - temp_file.replace(final_destination) + temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" + temp_file.replace(final_destination) - fineall = fineall and final_destination.exists() + fineall = fineall and final_destination.exists() - if not final_destination.exists(): - print(f"\t!! file {file} ({original_file}) could not be retrived",typeMsg="w",) + if not final_destination.exists(): + print(f"\t!! file {file} ({original_file}) could not be retrived",typeMsg="w",) - if fineall: - print("\t\t- All files were successfully retrieved") + if fineall: + print("\t\t- All files were successfully retrieved") - # Remove temporary folder - shutil.rmtree(tmpFolder) + # Remove temporary folder + shutil.rmtree(tmpFolder) - else: - print("\t\t- Some files were not retrieved", typeMsg="w") + else: + print("\t\t- Some files were not retrieved", typeMsg="w") diff --git a/templates/input.cgyro.models.json b/templates/input.cgyro.models.json index 4b45ea68..421b0ebb 100644 --- a/templates/input.cgyro.models.json +++ b/templates/input.cgyro.models.json @@ -4,5 +4,11 @@ "controls": { "NONLINEAR_FLAG": 0 } + }, + "1": { + "label": "Nonlinear", + "controls": { + "NONLINEAR_FLAG": 1 + } } } \ No newline at end of file diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index 0265e66a..c2c50969 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -29,7 +29,9 @@ }, slurm_setup={ 'cores':8 - } + }, + cold_start=cold_start, + forceIfcold_start=True, ) cgyro.read(label="cgyro1") cgyro.plot(labels=["cgyro1"]) @@ -48,7 +50,9 @@ varUpDown=[0.3,0.4], slurm_setup={ 'cores':4 - } + }, + cold_start=cold_start, + forceIfcold_start=True, ) cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) From 6717a328d62de4f183faae36aef00b94338f686f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 24 Aug 2025 14:08:06 -0400 Subject: [PATCH 192/385] Further generalization of gacodeinput --- src/mitim_tools/gacode_tools/CGYROtools.py | 61 +-------------- src/mitim_tools/gacode_tools/NEOtools.py | 57 +------------- src/mitim_tools/gacode_tools/TGLFtools.py | 58 ++++++-------- .../gacode_tools/utils/GACODErun.py | 77 +++++++++++++++---- 4 files changed, 89 insertions(+), 164 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 5b6d5359..8894e92e 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -5,7 +5,6 @@ from pathlib import Path import numpy as np import matplotlib.pyplot as plt -from mitim_tools import __version__ as mitim_version from mitim_tools import __mitimroot__ from mitim_tools.gacode_tools.utils import GACODEdefaults, GACODErun, CGYROutils from mitim_tools.misc_tools import IOtools, GRAPHICStools, FARMINGtools @@ -1969,68 +1968,12 @@ def read1(self, label="cgyro1", folder=None, tmin = 0.0, minimal = False, last_t if attach_name: self.results[label] = CGYROutils.CGYROlinear_scan(labels, data) - - class CGYROinput(GACODErun.GACODEinput): def __init__(self, file=None): super().__init__(file=file, controls_file= __mitimroot__ / "templates" / "input.cgyro.controls") - def process(self, input_dict): - - # Use standard processing - self._process(input_dict) - - # Get number of recorded species - self.num_recorded = 0 - if "N_SPECIES" in input_dict: - self.num_recorded = int(input_dict["N_SPECIES"]) - - def write_state(self, file=None): - - if file is None: - file = self.file - - # Local formatter: floats -> 6 significant figures in exponential (uppercase), - # ints stay as ints, bools as 0/1, sequences space-separated with same rule. - def _fmt_num(x): - import numpy as _np - if isinstance(x, (bool, _np.bool_)): - return "True" if x else "False" - if isinstance(x, (_np.floating, float)): - # 6 significant figures in exponential => 5 digits after decimal - return f"{float(x):.5E}" - if isinstance(x, (_np.integer, int)): - return f"{int(x)}" - return str(x) - - def _fmt_value(val): - import numpy as _np - if isinstance(val, (list, tuple, _np.ndarray)): - # Flatten numpy arrays but keep ordering; join with spaces - if isinstance(val, _np.ndarray): - flat = val.flatten().tolist() - else: - flat = list(val) - return " ".join(_fmt_num(v) for v in flat) - return _fmt_num(val) - - with open(file, "w") as f: - f.write("#-------------------------------------------------------------------------\n") - f.write(f"# CGYRO input file modified by MITIM {mitim_version}\n") - f.write("#-------------------------------------------------------------------------\n") - - f.write("\n\n# Control parameters\n") - f.write("# ------------------\n\n") - for ikey in self.controls: - var = self.controls[ikey] - f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") - - f.write("\n\n# Plasma/Geometry parameters\n") - f.write("# ------------------\n\n") - params = self.plasma | self.geom - for ikey in params: - var = params[ikey] - f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") + self.code = "CGYRO" + self.n_species = 'N_SPECIES' def _2D_mosaic(n_times): diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index d08d6ecb..cf45da5f 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -1,6 +1,5 @@ import numpy as np import matplotlib.pyplot as plt -from mitim_tools import __version__ as mitim_version from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools from mitim_tools.gacode_tools.utils import GACODErun, GACODEdefaults from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -244,60 +243,12 @@ class NEOinput(GACODErun.GACODEinput): def __init__(self, file=None): super().__init__(file=file) - @classmethod - def initialize_in_memory(cls, input_dict): - instance = cls() - instance.process(input_dict) - return instance - - def process(self, input_dict): - #TODO - self.plasma = input_dict - self.controls = {} - self.geom = {} - self.num_recorded = 100 - - def write_state(self, file=None): - - if file is None: - file = self.file - - # Local formatter: floats -> 6 significant figures in exponential (uppercase), - # ints stay as ints, bools as 0/1, sequences space-separated with same rule. - def _fmt_num(x): - import numpy as _np - if isinstance(x, (bool, _np.bool_)): - return "True" if x else "False" - if isinstance(x, (_np.floating, float)): - # 6 significant figures in exponential => 5 digits after decimal - return f"{float(x):.5E}" - if isinstance(x, (_np.integer, int)): - return f"{int(x)}" - return str(x) - - def _fmt_value(val): - import numpy as _np - if isinstance(val, (list, tuple, _np.ndarray)): - # Flatten numpy arrays but keep ordering; join with spaces - if isinstance(val, _np.ndarray): - flat = val.flatten().tolist() - else: - flat = list(val) - return " ".join(_fmt_num(v) for v in flat) - return _fmt_num(val) - - with open(file, "w") as f: - f.write("#-------------------------------------------------------------------------\n") - f.write(f"# NEO input file modified by MITIM {mitim_version}\n") - f.write("#-------------------------------------------------------------------------\n") - - for ikey in self.plasma: - var = self.plasma[ikey] - f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") - + self.code = 'NEO' + self.n_species = 'N_SPECIES' + class NEOoutput(GACODErun.GACODEoutput): - def __init__(self, FolderGACODE, suffix=""): + def __init__(self, FolderGACODE, suffix="", **kwargs): super().__init__() self.FolderGACODE, self.suffix = FolderGACODE, suffix diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 228ed6a5..9076b27a 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -514,7 +514,7 @@ def prep_from_file( # input_tglf_file inputclass = TGLFinput(file=input_tglf_file) - roa = inputclass.geom["RMIN_LOC"] + roa = inputclass.plasma["RMIN_LOC"] print(f"\t- This file correspond to r/a={roa} according to RMIN_LOC") if self.NormalizationSets["input_gacode"] is not None: @@ -3179,10 +3179,8 @@ def updateConvolution(self): ) else: for i in self.latest_inputsFileTGLFDict: - if "DRMAJDX_LOC" in self.latest_inputsFileTGLFDict[i].geom: - self.DRMAJDX_LOC[i] = self.latest_inputsFileTGLFDict[i].geom[ - "DRMAJDX_LOC" - ] + if "DRMAJDX_LOC" in self.latest_inputsFileTGLFDict[i].plasma: + self.DRMAJDX_LOC[i] = self.latest_inputsFileTGLFDict[i].plasma["DRMAJDX_LOC"] else: self.DRMAJDX_LOC[i] = 0.0 print( @@ -3238,8 +3236,10 @@ def completeVariation_TGLF(setVariations, species): return setVariations_new def reduceToControls(dict_all): - controls, plasma, geom = {}, {}, {} + controls, plasma = {}, {} for ikey in dict_all: + + # Plasma if ikey in [ "VEXB", "VEXB_SHEAR", @@ -3253,25 +3253,29 @@ def reduceToControls(dict_all): ]: plasma[ikey] = dict_all[ikey] + # Geometry elif (len(ikey.split("_")) > 1) and ( (ikey.split("_")[-1] in ["SA", "LOC"]) or (ikey.split("_")[0] in ["SHAPE"]) ): - geom[ikey] = dict_all[ikey] + plasma[ikey] = dict_all[ikey] + # Controls else: controls[ikey] = dict_all[ikey] - return controls, plasma, geom + return controls, plasma class TGLFinput(GACODErun.GACODEinput): def __init__(self, file=None): super().__init__(file=file) + self.code = 'TGLF' + self.n_species = "NS" + def process(self, input_dict): # Get number of recorded species - self.num_recorded = 0 - if "NS" in input_dict: - self.num_recorded = int(input_dict["NS"]) + if self.n_species in input_dict: + self.num_recorded = int(input_dict[self.n_species]) # Species ----------- self.species = {} @@ -3300,7 +3304,7 @@ def process(self, input_dict): specie[var] = 0.0 self.species[i + 1] = specie - self.controls, self.plasma, self.geom = reduceToControls(controls_all) + self.controls, self.plasma = reduceToControls(controls_all) self.processSpecies() def processSpecies(self, MinMultiplierToBeFast=2.0): @@ -3480,7 +3484,6 @@ def addTraceSpecie(self, ZS, MASS, AS=1e-6, position=None, increaseNS=True, posi return position def write_state(self, file=None): - print("\t- Writting TGLF input file") maxSpeciesTGLF = 6 # TGLF cannot handle more than 6 species @@ -3513,7 +3516,7 @@ def _fmt_value(val): with open(file, "w") as f: f.write("#-------------------------------------------------------------------------\n") - f.write(f"# TGLF input file modified by MITIM {mitim_version}\n") + f.write(f"# {self.code} input file modified by MITIM {mitim_version}\n") f.write("#-------------------------------------------------------------------------") f.write("\n\n# Control parameters\n") @@ -3522,16 +3525,10 @@ def _fmt_value(val): var = self.controls[ikey] f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") - f.write("\n\n# Geometry parameters\n") - f.write("# ------------------\n\n") - for ikey in self.geom: - var = self.geom[ikey] - f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") - - f.write("\n\n# Plasma parameters\n") + f.write("\n\n# Plasma/Geometry parameters\n") f.write("# ------------------\n\n") for ikey in self.plasma: - if ikey == "NS": + if ikey == self.n_species: var = np.min([self.plasma[ikey], maxSpeciesTGLF]) if var < self.plasma[ikey]: print(f"\t- Maximum number of species in TGLF reached, not considering after {maxSpeciesTGLF} species",typeMsg="w",) @@ -3739,19 +3736,8 @@ def plotPlasma(self, axs=None, color="b", legends=True): ax.plot(x, y, "-o", lw=1, color=color) ax.set_xticks(x) ax.set_xticklabels(x, rotation=90, fontsize=6) - ax.set_ylabel("PLASMA") + ax.set_ylabel("PLASMA & GEOMETRY") - ax = axs[1] - - x, y = [], [] - for i in self.geom: - x.append(i) - y.append(self.geom[i]) - x, y = np.array(x), np.array(y) - ax.plot(x, y, "-o", lw=1, color=color) - ax.set_xticks(x) - ax.set_xticklabels(x, rotation=90, fontsize=6) - ax.set_ylabel("GEOMETRY") def plotControls(self, axs=None, color="b", markersize=5): if axs is None: @@ -3766,7 +3752,7 @@ def plotControls(self, axs=None, color="b", markersize=5): cont = 0 x, y = x1, y1 - dicts = [self.controls] # ,self.geom,self.plasma] + dicts = [self.controls] for dictT in dicts: for i in dictT: @@ -3877,7 +3863,7 @@ def __init__(self, FolderGACODE, suffix="", require_all_files=True): print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} with suffix {self.suffix}") self.inputclass = TGLFinput(file=self.FolderGACODE / f"input.tglf{self.suffix}") - self.roa = self.inputclass.geom["RMIN_LOC"] + self.roa = self.inputclass.plasma["RMIN_LOC"] self.read(require_all_files=require_all_files) self.postprocess() diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index bca9900d..b016bce8 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -5,6 +5,7 @@ from pathlib import Path import matplotlib.pyplot as plt from scipy.interpolate import interp1d +from mitim_tools import __version__ as mitim_version from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.gacode_tools.utils import GACODEdefaults, NORMtools from mitim_tools.transp_tools.utils import NTCCtools @@ -1102,29 +1103,26 @@ def modifyInputs( 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 input_class.species[1] except: isspecie = False + # is a species parameter? if isspecie: specie = int(ikey.split("_")[-1]) varK = "_".join(ikey.split("_")[:-1]) var_orig = input_class.species[specie][varK] var_new = value_to_change_to input_class.species[specie][varK] = var_new + # is a another parameter? else: if ikey in input_class.controls: var_orig = input_class.controls[ikey] var_new = value_to_change_to input_class.controls[ikey] = var_new - elif 'geom' in input_class.__dict__ and ikey in input_class.geom: - var_orig = input_class.geom[ikey] - var_new = value_to_change_to - input_class.geom[ikey] = var_new - elif 'plasma' in input_class.__dict__ and ikey in input_class.plasma: + elif ikey in input_class.plasma: var_orig = input_class.plasma[ikey] var_new = value_to_change_to input_class.plasma[ikey] = var_new @@ -1156,11 +1154,6 @@ def modifyInputs( var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.controls[ikey] = var_new - elif ikey in input_class.geom: - var_orig = input_class.geom[ikey] - var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) - input_class.geom[ikey] = var_new - elif ikey in input_class.plasma: var_orig = input_class.plasma[ikey] var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) @@ -1894,29 +1887,81 @@ def __init__(self, file=None, controls_file=None): input_dict = buildDictFromInput(file_txt) self.process(input_dict) + + self.code = '' + self.n_species = None - @classmethod def initialize_in_memory(cls, input_dict): instance = cls() instance.process(input_dict) return instance - def _process(self, input_dict): + def process(self, input_dict): if self.controls_file is not None: options_check = [key for key in IOtools.generateMITIMNamelist(self.controls_file, caseInsensitive=False).keys()] else: options_check = [] - self.controls, self.plasma, self.geom = {}, {}, {} + self.controls, self.plasma = {}, {} for key in input_dict.keys(): if key in options_check: self.controls[key] = input_dict[key] else: self.plasma[key] = input_dict[key] - self.num_recorded = 100 + # Get number of recorded species + if self.n_species is not None and self.n_species in input_dict: + self.num_recorded = int(input_dict[self.n_species]) + else: + self.num_recorded = 100 + + def write_state(self, file=None): + + if file is None: + file = self.file + + # Local formatter: floats -> 6 significant figures in exponential (uppercase), + # ints stay as ints, bools as 0/1, sequences space-separated with same rule. + def _fmt_num(x): + import numpy as _np + if isinstance(x, (bool, _np.bool_)): + return "True" if x else "False" + if isinstance(x, (_np.floating, float)): + # 6 significant figures in exponential => 5 digits after decimal + return f"{float(x):.5E}" + if isinstance(x, (_np.integer, int)): + return f"{int(x)}" + return str(x) + + def _fmt_value(val): + import numpy as _np + if isinstance(val, (list, tuple, _np.ndarray)): + # Flatten numpy arrays but keep ordering; join with spaces + if isinstance(val, _np.ndarray): + flat = val.flatten().tolist() + else: + flat = list(val) + return " ".join(_fmt_num(v) for v in flat) + return _fmt_num(val) + + with open(file, "w") as f: + f.write("#-------------------------------------------------------------------------\n") + f.write(f"# {self.code} input file modified by MITIM {mitim_version}\n") + f.write("#-------------------------------------------------------------------------\n") + + f.write("\n\n# Control parameters\n") + f.write("# ------------------\n\n") + for ikey in self.controls: + var = self.controls[ikey] + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") + + f.write("\n\n# Plasma/Geometry parameters\n") + f.write("# ------------------\n\n") + for ikey in self.plasma: + var = self.plasma[ikey] + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") def anticipate_problems(self): pass From e94f2f2cd96ca208a4f209cb576ed25f4d720969 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 24 Aug 2025 16:50:13 -0400 Subject: [PATCH 193/385] GX fully working; missing parallelization and checks about variables --- src/mitim_tools/gacode_tools/CGYROtools.py | 11 +- src/mitim_tools/gacode_tools/NEOtools.py | 12 +- src/mitim_tools/gacode_tools/TGLFtools.py | 13 +- .../gacode_tools/utils/GACODEdefaults.py | 10 +- .../gacode_tools/utils/GACODErun.py | 11 +- src/mitim_tools/gyrokinetics_tools/GXtools.py | 279 ++++++++++++++++++ src/mitim_tools/misc_tools/FARMINGtools.py | 4 +- .../plasmastate_tools/MITIMstate.py | 162 ++++++++-- templates/config_user_example.json | 1 + templates/input.gx.controls | 58 ++++ tests/GX_workflow.py | 36 +++ 11 files changed, 548 insertions(+), 49 deletions(-) create mode 100644 src/mitim_tools/gyrokinetics_tools/GXtools.py create mode 100644 templates/input.gx.controls create mode 100644 tests/GX_workflow.py diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 8894e92e..81b4fc9a 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -1970,11 +1970,12 @@ def read1(self, label="cgyro1", folder=None, tmin = 0.0, minimal = False, last_t class CGYROinput(GACODErun.GACODEinput): def __init__(self, file=None): - super().__init__(file=file, controls_file= __mitimroot__ / "templates" / "input.cgyro.controls") - - self.code = "CGYRO" - self.n_species = 'N_SPECIES' - + super().__init__( + file=file, + controls_file= __mitimroot__ / "templates" / "input.cgyro.controls", + code="CGYRO", + n_species='N_SPECIES', + ) def _2D_mosaic(n_times): diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index cf45da5f..4fe0a31d 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -3,6 +3,7 @@ from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools from mitim_tools.gacode_tools.utils import GACODErun, GACODEdefaults from mitim_tools.misc_tools.LOGtools import printMsg as print +from mitim_tools import __mitimroot__ from IPython import embed class NEO(GACODErun.gacode_simulation): @@ -241,11 +242,12 @@ def check_if_files_exist(folder, list_files): class NEOinput(GACODErun.GACODEinput): def __init__(self, file=None): - super().__init__(file=file) - - self.code = 'NEO' - self.n_species = 'N_SPECIES' - + super().__init__( + file=file, + controls_file= __mitimroot__ / "templates" / "input.neo.controls", + code='NEO', + n_species='N_SPECIES' + ) class NEOoutput(GACODErun.GACODEoutput): def __init__(self, FolderGACODE, suffix="", **kwargs): diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 9076b27a..a44460cd 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -5,7 +5,7 @@ import xarray as xr import matplotlib.pyplot as plt from mitim_tools import __version__ as mitim_version - +from mitim_tools import __mitimroot__ from mitim_tools.gacode_tools import TGYROtools from mitim_tools.misc_tools import ( IOtools, @@ -3263,13 +3263,14 @@ def reduceToControls(dict_all): return controls, plasma - class TGLFinput(GACODErun.GACODEinput): def __init__(self, file=None): - super().__init__(file=file) - - self.code = 'TGLF' - self.n_species = "NS" + super().__init__( + file=file, + controls_file= __mitimroot__ / "templates" / "input.tglf.controls", + code = 'TGLF', + n_species = 'NS' + ) def process(self, input_dict): diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index 8cd15494..10d4983c 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -55,9 +55,15 @@ def addTGLFcontrol(TGLFsettings, NS=2, minimal=False): def addNEOcontrol(*args, **kwargs): - NEOoptions = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.neo.controls", caseInsensitive=False) + options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.neo.controls", caseInsensitive=False) - return NEOoptions + return options + +def addGXcontrol(*args, **kwargs): + + options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.gx.controls", caseInsensitive=False) + + return options def addCGYROcontrol(code_settings, rmin=None, **kwargs): diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index b016bce8..8fe6d921 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -1091,7 +1091,7 @@ def modifyInputs( print("\t- Input file was not modified by code_settings, using what was there before",typeMsg="i") # Make all upper case - extraOptions = {ikey.upper(): value for ikey, value in extraOptions.items()} + #extraOptions = {ikey.upper(): value for ikey, value in extraOptions.items()} # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Change with external options -> Input directly, not as multiplier @@ -1873,10 +1873,12 @@ def unnormalize(self, *args, **kwargs): print("No unnormalization implemented.") class GACODEinput: - def __init__(self, file=None, controls_file=None): + def __init__(self, file=None, controls_file=None, code='', n_species=None): self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None self.controls_file = controls_file + self.code = code + self.n_species = n_species if self.file is not None and self.file.exists(): with open(self.file, "r") as f: @@ -1887,10 +1889,7 @@ def __init__(self, file=None, controls_file=None): input_dict = buildDictFromInput(file_txt) self.process(input_dict) - - self.code = '' - self.n_species = None - + @classmethod def initialize_in_memory(cls, input_dict): instance = cls() diff --git a/src/mitim_tools/gyrokinetics_tools/GXtools.py b/src/mitim_tools/gyrokinetics_tools/GXtools.py new file mode 100644 index 00000000..86c94d3d --- /dev/null +++ b/src/mitim_tools/gyrokinetics_tools/GXtools.py @@ -0,0 +1,279 @@ +import numpy as np +import netCDF4 +import matplotlib.pyplot as plt +from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools +from mitim_tools.gacode_tools.utils import GACODErun, GACODEdefaults +from mitim_tools.misc_tools.LOGtools import printMsg as print +from mitim_tools import __mitimroot__ +from mitim_tools import __version__ as mitim_version +from IPython import embed + +class GX(GACODErun.gacode_simulation): + def __init__( + self, + rhos=[0.4, 0.6], # rho locations of interest + ): + + super().__init__(rhos=rhos) + + def code_call(folder, n, p, additional_command="", **kwargs): + return f" gx {folder}/gxplasma.in &\n" + + self.run_specifications = { + 'code': 'gx', + 'input_file': 'gxplasma.in', + 'code_call': code_call, + 'control_function': GACODEdefaults.addGXcontrol, + 'controls_file': 'input.gx.controls', + 'state_converter': 'to_gx', + 'input_class': GXinput, + 'complete_variation': None, + 'default_cores': 4, # Default gpus to use in the simulation + 'output_class': GXoutput, + 'output_store': 'GXout' + } + + print("\n-----------------------------------------------------------------------------------------") + print("\t\t\t GX class module") + print("-----------------------------------------------------------------------------------------\n") + + self.ResultsFiles = self.ResultsFiles_minimal = [ + 'gxplasma.eik.out', + 'gxplasma.eiknc.nc', + 'gxplasma.gx_geo.log', + 'gxplasma.restart.nc', + 'gxplasma.big.nc', + 'gxplasma.out.nc' + ] + + def plot( + self, + fn=None, + labels=["gx1"], + extratitle="", + fn_color=None, + colors=None, + ): + + if fn is None: + self.fn = GUItools.FigureNotebook("GX MITIM Notebook", geometry="1700x900", vertical=True) + else: + self.fn = fn + + fig1 = self.fn.add_figure(label=f"{extratitle}Summary", tab_color=fn_color) + + grid = plt.GridSpec(1, 2, hspace=0.7, wspace=0.2) + + if colors is None: + colors = GRAPHICStools.listColors() + + ax1 = fig1.add_subplot(grid[0, 0]) + ax2 = fig1.add_subplot(grid[0, 1]) + + i = 0 + for label in labels: + for irho in range(len(self.rhos)): + c = self.results[label]['GXout'][irho] + for iky in range(len(c.ky)): + ax1.plot(c.t, c.w[:, iky], label=f"{label} rho={self.rhos[irho]} ky={c.ky[iky]}", color=colors[i]) + ax2.plot(c.t, c.g[:, iky], label=f"{label} rho={self.rhos[irho]} ky={c.ky[iky]}", color=colors[i]) + i += 1 + + for ax in [ax1, ax2]: + ax.set_xlabel("Time") + ax.set_xlim(left=0) + GRAPHICStools.addDenseAxis(ax) + ax1.set_ylabel("Real frequency") + ax1.legend(loc='best', prop={'size': 4}) + + ax2.set_ylabel("Growth rate") + +class GXinput(GACODErun.GACODEinput): + def __init__(self, file=None): + super().__init__( + file=file, + controls_file= __mitimroot__ / "templates" / "input.gx.controls", + code='GX', + n_species='nspecies' + ) + + # GX has a very particular way to write its state + def write_state(self, file=None): + + if file is None: + file = self.file + + + with open(file, "w") as f: + f.write("#-------------------------------------------------------------------------\n") + f.write(f"# {self.code} input file modified by MITIM {mitim_version}\n") + f.write("#-------------------------------------------------------------------------\n") + + # title: [controls], [plasma] + blocks = { + 'Dimensions': + [ ['ntheta', 'nperiod', 'ny', 'nx', 'nhermite', 'nlaguerre', 'nspecies'], [] ], + 'Domain': + [ ['y0', 'boundary'], [] ], + 'Physics': + [ ['nonlinear_mode', 'ei_colls'], ['beta'] ], + 'Time': + [ ['t_max', 'scheme'], [] ], + 'Initialization': + [ ['ikpar_init', 'init_field', 'init_amp', 'gaussian_init'], [] ], + 'Geometry': + [ + ['geo_option'], + ['rhoc', 'Rmaj', 'R_geo', 'shift', 'qinp', 'shat', 'akappa', 'akappri', 'tri', 'tripri', 'betaprim'] + ], + 'Dissipation': + [ ['closure_model', 'hypercollisions', 'nu_hyper_m', 'p_hyper_m', 'nu_hyper_l', 'p_hyper_l', 'hyper', 'D_hyper', 'p_hyper'], [] ], + 'Restart': + [ ['restart', 'save_for_restart'], [] ], + 'Diagnostics': + [ ['nwrite', 'omega', 'fluxes', 'fields'], [] ] + } + + param_written = [] + for block_name, params in blocks.items(): + param_written = self._write_block(f, f"[{block_name}]", params, param_written) + + param_written = self._write_block_species(f, param_written) + + def _write_block_species(self, f, param_written): + + # Local formatter: floats -> 6 significant figures in exponential (uppercase), + # ints stay as ints, bools as 0/1, sequences space-separated with same rule. + def _fmt_num(x): + import numpy as _np + if isinstance(x, (bool, _np.bool_)): + return "true" if x else "false" + if isinstance(x, (_np.floating, float)): + # 6 significant figures in exponential => 5 digits after decimal + return f"{float(x):.5E}" + if isinstance(x, (_np.integer, int)): + return f"{int(x)}" + return str(x) + + def _fmt_value(val): + import numpy as _np + if isinstance(val, (list, tuple, _np.ndarray)): + # Flatten numpy arrays but keep ordering; join with spaces + if isinstance(val, _np.ndarray): + flat = val.flatten().tolist() + else: + flat = list(val) + return " ".join(_fmt_num(v) for v in flat) + return _fmt_num(val) + + self.num_recorded = 0 + for i in range(100): + if f"z_{i+1}" in self.plasma: + self.num_recorded += 1 + else: + break + + z, dens, temp, mass, fprim, tprim, vnewk, typeS = '[ ', '[ ', '[ ', '[ ', '[ ', '[ ', '[ ', '[ ' + for i in range(self.num_recorded): + typeS += f'"{_fmt_value(self.plasma[f"type_{i+1}"])}", ' + z += f'{_fmt_value(self.plasma[f"z_{i+1}"])}, ' + mass += f'{_fmt_value(self.plasma[f"mass_{i+1}"])}, ' + dens += f'{_fmt_value(self.plasma[f"dens_{i+1}"])}, ' + temp += f'{_fmt_value(self.plasma[f"temp_{i+1}"])}, ' + fprim += f'{_fmt_value(self.plasma[f"fprim_{i+1}"])}, ' + tprim += f'{_fmt_value(self.plasma[f"tprim_{i+1}"])}, ' + vnewk += f'{_fmt_value(self.plasma[f"vnewk_{i+1}"])}, ' + + param_written.append(f"type_{i+1}") + param_written.append(f"z_{i+1}") + param_written.append(f"dens_{i+1}") + param_written.append(f"temp_{i+1}") + param_written.append(f"mass_{i+1}") + param_written.append(f"fprim_{i+1}") + param_written.append(f"tprim_{i+1}") + param_written.append(f"vnewk_{i+1}") + + f.write("[species]\n") + f.write(f" z = {z[:-2]} ]\n") + f.write(f" dens = {dens[:-2]} ]\n") + f.write(f" temp = {temp[:-2]} ]\n") + f.write(f" mass = {mass[:-2]} ]\n") + f.write(f" fprim = {fprim[:-2]} ]\n") + f.write(f" tprim = {tprim[:-2]} ]\n") + f.write(f" vnewk = {vnewk[:-2]} ]\n") + f.write(f" type = {typeS[:-2]} ]\n") + f.write("\n") + + return param_written + + def _write_block(self,f,name, param, param_written): + + # Local formatter: floats -> 6 significant figures in exponential (uppercase), + # ints stay as ints, bools as 0/1, sequences space-separated with same rule. + def _fmt_num(x): + import numpy as _np + if isinstance(x, (bool, _np.bool_)): + return "true" if x else "false" + if isinstance(x, (_np.floating, float)): + # 6 significant figures in exponential => 5 digits after decimal + return f"{float(x):.5E}" + if isinstance(x, (_np.integer, int)): + return f"{int(x)}" + return str(x) + + def _fmt_value(val): + import numpy as _np + if isinstance(val, (list, tuple, _np.ndarray)): + # Flatten numpy arrays but keep ordering; join with spaces + if isinstance(val, _np.ndarray): + flat = val.flatten().tolist() + else: + flat = list(val) + return " ".join(_fmt_num(v) for v in flat) + return _fmt_num(val) + + f.write(f'{name}\n') + for p in param[0]: + f.write(f" {p.ljust(23)} = {_fmt_value(self.controls[p])}\n") + param_written.append(p) + for p in param[1]: + f.write(f" {p.ljust(23)} = {_fmt_value(self.plasma[p])}\n") + param_written.append(p) + f.write(f'\n') + + return param_written + + +class GXoutput(GACODErun.GACODEoutput): + def __init__(self, FolderGACODE, suffix="", **kwargs): + super().__init__() + + self.FolderGACODE, self.suffix = FolderGACODE, suffix + + if suffix == "": + print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} without suffix") + else: + print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} with suffix {suffix}") + + self.inputclass = GXinput(file=self.FolderGACODE / f"gxplasma.in{self.suffix}") + + self.read() + + def read(self): + + data = netCDF4.Dataset(self.FolderGACODE / f"gxplasma.out.nc{self.suffix}") + + self.t = data.groups['Grids'].variables['time'][:] # (time) + + # Growth rates + ikx = 0 + self.ky = data.groups['Grids'].variables['ky'][1:] # (ky) + self.w = data.groups['Diagnostics'].variables['omega_kxkyt'][:,1:,ikx,0] # (time, ky) + self.g = data.groups['Diagnostics'].variables['omega_kxkyt'][:,1:,ikx,1] # (time, ky) + + # get fluxes + # Qi = data.groups['Fluxes'].variables['qflux'][:,1] + # Qe = data.groups['Fluxes'].variables['qflux'][:,0] + # Ge = data.groups['Fluxes'].variables['pflux'][:,0] + + diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 005f279b..e2c6bc10 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -1106,7 +1106,9 @@ def create_slurm_execution_files( commandSBATCH.append('echo ""') commandSBATCH.append("") - full_command = [modules_remote] if (modules_remote is not None) else [] + # If modules, add them, but also make sure I expand the potential aliases that they may have! + full_command = ["shopt -s expand_aliases",modules_remote] if (modules_remote is not None) else [] + full_command.extend(command) for c in full_command: commandSBATCH.append(c) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index e65b9ac3..ef7569d0 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2267,7 +2267,7 @@ def to_tglf(self, r=[0.5], TGLFsettings=1, r_is_rho = True): # Prepare the inputs for TGLF # --------------------------------------------------------------------------------------------------------------------------------------- - inputsTGLF = {} + input_parameters = {} for rho in r: # --------------------------------------------------------------------------------------------------------------------------------------- @@ -2369,27 +2369,26 @@ def interpolator(y): parameters[key_mod] = self.profiles[ikey] parameters[f"{key_mod.split('_')[0]}_S_{key_mod.split('_')[-1]}"] = self.derived["r"] * self._deriv_gacode(self.profiles[ikey]) - geom = {} for k in parameters: par = torch.nan_to_num(torch.from_numpy(parameters[k]) if type(parameters[k]) is np.ndarray else parameters[k], nan=0.0, posinf=1E10, neginf=-1E10) - geom[k] = interpolator(par) + plasma[k] = interpolator(par) - geom['BETA_LOC'] = 0.0 - geom['KX0_LOC'] = 0.0 + plasma['BETA_LOC'] = 0.0 + plasma['KX0_LOC'] = 0.0 # --------------------------------------------------------------------------------------------------------------------------------------- # Merging # --------------------------------------------------------------------------------------------------------------------------------------- - input_dict = {**controls, **plasma, **geom} + input_dict = controls | plasma for i in range(len(species)): for k in species[i+1]: input_dict[f'{k}_{i+1}'] = species[i+1][k] - inputsTGLF[rho] = input_dict + input_parameters[rho] = input_dict - return inputsTGLF + return input_parameters def to_neo(self, r=[0.5], r_is_rho = True): @@ -2429,7 +2428,7 @@ def to_neo(self, r=[0.5], r_is_rho = True): # NEO definition: 'OMEGA_ROT_DERIV=',-gamma_p_loc*a/cs_loc/rmaj_loc omega_rot_deriv = gamma_p * a / cs / rmaj # Equivalent to: self._deriv_gacode(self.profiles["w0(rad/s)"])/ self.derived['c_s'] * self.derived['a']**2 - inputsNEO = {} + input_parameters = {} for rho in r: # --------------------------------------------------------------------------------------------------------------------------------------- @@ -2461,7 +2460,6 @@ def interpolator(y): } ie = i+2 - species[ie] = { 'Z': -1.0, 'MASS': 0.000272445, @@ -2522,24 +2520,23 @@ def interpolator(y): parameters[key_mod] = self.profiles[ikey] parameters[f"{key_mod.split('_')[0]}_S_{key_mod.split('_')[-1]}"] = self.derived["r"] * self._deriv_gacode(self.profiles[ikey]) - geom = {} for k in parameters: par = torch.nan_to_num(torch.from_numpy(parameters[k]) if type(parameters[k]) is np.ndarray else parameters[k], nan=0.0, posinf=1E10, neginf=-1E10) - geom[k] = interpolator(par) + plasma[k] = interpolator(par) # --------------------------------------------------------------------------------------------------------------------------------------- # Merging # --------------------------------------------------------------------------------------------------------------------------------------- - input_dict = {**controls, **plasma, **geom} + input_dict = controls | plasma for i in range(len(species)): for k in species[i+1]: input_dict[f'{k}_{i+1}'] = species[i+1][k] - inputsNEO[rho] = input_dict + input_parameters[rho] = input_dict - return inputsNEO + return input_parameters def to_cgyro(self, r=[0.5], r_is_rho = True): @@ -2582,12 +2579,10 @@ def to_cgyro(self, r=[0.5], r_is_rho = True): # CGYRO definition: 'GAMMA_E=',gamma_e_loc*a/cs_loc gamma_e = gamma_e * a / cs - # Because in MITIMstate I keep Bunit always positive, but CGYRO routines may need it negative? #TODO sign_Bunit = np.sign(self.profiles['torfluxa(Wb/radian)'][0]) - - inputsCGYRO = {} + input_parameters = {} for rho in r: # --------------------------------------------------------------------------------------------------------------------------------------- @@ -2620,7 +2615,6 @@ def interpolator(y): } ie = i+2 - species[ie] = { 'Z': -1.0, 'MASS': 0.000272445, @@ -2680,25 +2674,145 @@ def interpolator(y): parameters[key_mod] = self.profiles[ikey] parameters[f"{key_mod.split('_')[0]}_S_{key_mod.split('_')[-1]}"] = self.derived["r"] * self._deriv_gacode(self.profiles[ikey]) - geom = {} for k in parameters: par = torch.nan_to_num(torch.from_numpy(parameters[k]) if type(parameters[k]) is np.ndarray else parameters[k], nan=0.0, posinf=1E10, neginf=-1E10) - geom[k] = interpolator(par) + plasma[k] = interpolator(par) # --------------------------------------------------------------------------------------------------------------------------------------- # Merging # --------------------------------------------------------------------------------------------------------------------------------------- - input_dict = {**controls, **plasma, **geom} + input_dict = controls | plasma for i in range(len(species)): for k in species[i+1]: input_dict[f'{k}_{i+1}'] = species[i+1][k] - inputsCGYRO[rho] = input_dict + input_parameters[rho] = input_dict + + return input_parameters + + def to_gx(self, r=[0.5], r_is_rho = True): + + # <> Function to interpolate a curve <> + from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function + + # Determine if the input radius is rho toroidal or r/a + if r_is_rho: + r_interpolation = self.profiles['rho(-)'] + else: + r_interpolation = self.derived['roa'] + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Prepare the inputs + # --------------------------------------------------------------------------------------------------------------------------------------- + + # Determine the mass reference + mass_ref = 2.0 + + + import scipy.constants as const + p_normalized = self.profiles['ptot(Pa)'] / (8*np.pi) * (2*const.mu_0) + betaprim = -(8*np.pi / self.derived['B_unit']**2) * np.gradient(p_normalized, self.derived['roa']) + + #TODO #to check + s_kappa = self.derived["r"] / self.profiles["kappa(-)"] * self._deriv_gacode(self.profiles["kappa(-)"]) + s_delta = self.derived["r"] * self._deriv_gacode(self.profiles["delta(-)"]) + + + + input_parameters = {} + for rho in r: + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Define interpolator at this rho + # --------------------------------------------------------------------------------------------------------------------------------------- - return inputsCGYRO + def interpolator(y): + return interpolation_function(rho, r_interpolation,y).item() + # --------------------------------------------------------------------------------------------------------------------------------------- + # Controls come from options + # --------------------------------------------------------------------------------------------------------------------------------------- + + controls = GACODEdefaults.addGXcontrol() + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Species come from profiles + # --------------------------------------------------------------------------------------------------------------------------------------- + + species = {} + + # Ions + for i in range(len(self.Species)): + + nu_ii = self.derived['xnue'] * (self.Species[i]['Z']/self.profiles['ze'][0])**4 * (self.profiles['ni(10^19/m^3)'][:,0]/self.profiles['ne(10^19/m^3)']) * (self.profiles['mass'][i]/self.profiles['masse'][0])**-0.5 * (self.profiles['ti(keV)'][:,0]/self.profiles['te(keV)'])**-1.5 + + species[i+1] = { + 'z': self.Species[i]['Z'], + 'mass': self.Species[i]['A']/mass_ref, + 'temp': interpolator(self.derived["tite_all"][:,i]), + 'dens': interpolator(self.derived['fi'][:,i]), + 'fprim': interpolator(self.derived['aLni'][:,i]), + 'tprim': interpolator(self.derived["aLTi"][:,i]), + 'vnewk': interpolator(nu_ii), + 'type': 'ion', + } + + # Electrons + ie = i+2 + species[ie] = { + 'z': -1.0, + 'mass': 0.000272445, + 'temp': 1.0, + 'dens': 1.0, + 'fprim': interpolator(self.derived['aLne']), + 'tprim': interpolator(self.derived['aLTe']), + 'vnewk': interpolator(self.derived['xnue']), + 'type': 'electron' + } + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Plasma and geometry + # --------------------------------------------------------------------------------------------------------------------------------------- + + plasma = { + 'nspecies': len(species)-1 # do not count electrons} + } + + parameters = { + 'beta': self.derived['betae'], + 'rhoc': self.derived['roa'], + 'Rmaj': self.derived['Rmajoa'], + 'R_geo': self.derived['Rmajoa'] / abs(self.derived['B_unit'] / self.derived['B0']), + 'shift': self._deriv_gacode(self.profiles["rmaj(m)"]), + 'qinp': np.abs(self.profiles["q(-)"]), + 'shat': self.derived["s_hat"], + 'akappa': self.profiles["kappa(-)"], + 'akappri': s_kappa, + 'tri': self.profiles["delta(-)"], + 'tripri': s_delta, + 'betaprim': betaprim, + } + + for k in parameters: + par = torch.nan_to_num(torch.from_numpy(parameters[k]) if type(parameters[k]) is np.ndarray else parameters[k], nan=0.0, posinf=1E10, neginf=-1E10) + plasma[k] = interpolator(par) + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Merging + # --------------------------------------------------------------------------------------------------------------------------------------- + + input_dict = controls | plasma + + for i in range(len(species)): + for k in species[i+1]: + input_dict[f'{k}_{i+1}'] = species[i+1][k] + + input_parameters[rho] = input_dict + + return input_parameters + def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', times = [0.0,1.0], Vsurf = 0.0): diff --git a/templates/config_user_example.json b/templates/config_user_example.json index df1223fe..a1644094 100644 --- a/templates/config_user_example.json +++ b/templates/config_user_example.json @@ -8,6 +8,7 @@ "tglf": "engaging", "neo": "engaging", "cgyro": "engaging", + "gx": "engaging", "astra": "engaging", "eq": "mfews", "scruncher": "mfews", diff --git a/templates/input.gx.controls b/templates/input.gx.controls new file mode 100644 index 00000000..e4bf3841 --- /dev/null +++ b/templates/input.gx.controls @@ -0,0 +1,58 @@ +#============================================================= +# GX example +#============================================================= + +[Dimensions] + ntheta = 24 # number of points along field line (theta) per 2pi segment + nperiod = 1 # number of 2pi segments along field line is 2*nperiod-1 + ny = 106 # number of real-space grid-points in y, nky = 1 + (ny-1)/3 --> 36 ky modes + nx = 382 # number of real-space grid-points in x, nkx = 1 + 2*(nx-1)/3 --> 255 kx modes + + nhermite = 8 # number of hermite moments (v_parallel resolution) + nlaguerre = 4 # number of laguerre moments (mu B resolution) + nspecies = 1 # number of evolved kinetic species (adiabatic electrons don't count towards nspecies) + +[Domain] + y0 = 14.0 # controls box length in y (in units of rho_ref) and minimum ky, so that ky_min*rho_ref = 1/y0 --> ky_min = 0.0714, ky_max = 2.57, L_y = 88 + boundary = "linked" # use twist-shift boundary conditions along field line + +[Physics] + nonlinear_mode = false # this is a linear calculation + ei_colls = false # turn off electron-ion collisions + +[Time] + t_max = 1000.0 # end time (in units of L_ref/vt_ref) + scheme = "rk4" # use RK3 timestepping scheme (with adaptive timestepping) + +[Initialization] + gaussian_init = true # initial perturbation is a gaussian in theta + ikpar_init = 0 # parallel wavenumber of initial perturbation + init_field = "density" # initial condition set in density + init_amp = 1.0e-3 # amplitude of initial condition + +[Geometry] + geo_option = "miller" # use Miller geometry + +[Dissipation] + closure_model = "none" # no closure assumptions (just truncation) + hypercollisions = true # use hypercollision model + nu_hyper_m = 0.1 # coefficient of hermite hypercollisions + p_hyper_m = 6 # exponent of hermite hypercollisions + nu_hyper_l = 0.1 # coefficient of laguerre hypercollisions + p_hyper_l = 6 # exponent of laguerre hypercollisions + + hyper = true # use hyperdiffusion + D_hyper = 0.05 # coefficient of hyperdiffusion + p_hyper = 2 # exponent of hyperdiffusion is 2*p_hyper = 4 + +[Restart] + restart = false + save_for_restart = true + +[Diagnostics] + nwrite = 1000 # write diagnostics every nwrite timesteps + omega = true # compute and write linear growth rate and frequency + fluxes = true # compute and write heat and particle fluxes + fields = true # compute and write electric and magnetic fields + +[species] diff --git a/tests/GX_workflow.py b/tests/GX_workflow.py new file mode 100644 index 00000000..d4932b11 --- /dev/null +++ b/tests/GX_workflow.py @@ -0,0 +1,36 @@ +import os +from mitim_tools.gacode_tools.PROFILEStools import gacode_state +from mitim_tools.gyrokinetics_tools import GXtools +from mitim_tools import __mitimroot__ + +cold_start = True + +(__mitimroot__ / 'tests' / 'scratch').mkdir(parents=True, exist_ok=True) + +folder = __mitimroot__ / "tests" / "scratch" / "gx_test" +input_gacode = __mitimroot__ / "tests" / "data" / "input.gacode" + +# Only 1 ion species +p = gacode_state(input_gacode) +p.lumpSpecies(ions_list=[1,2,3,4]) +# + +if cold_start and folder.exists(): + os.system(f"rm -r {folder.resolve()}") + +gx = GXtools.GX(rhos=[0.5,0.7]) +gx.prep(p, folder) + +gx.run( + 'gx1/', + cold_start=cold_start, + extraOptions={ + 't_max':5.0, # Short, I just want to test the run + }, + ) +gx.read('gx1') + +gx.plot(labels=['gx1']) + +gx.fn.show() +gx.fn.close() From f6a81c55e2950f5095ede4181b76eeef4baf076a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sun, 24 Aug 2025 19:19:11 -0400 Subject: [PATCH 194/385] Nicer GX --- .../gacode_tools/utils/GACODEdefaults.py | 49 ++-- src/mitim_tools/gyrokinetics_tools/GXtools.py | 212 ++++++++++++------ .../plasmastate_tools/MITIMstate.py | 8 +- templates/input.gx.controls | 10 +- templates/input.gx.models.json | 16 ++ tests/GX_workflow.py | 8 +- tests/TGLF_workflow.py | 21 +- 7 files changed, 225 insertions(+), 99 deletions(-) create mode 100644 templates/input.gx.models.json diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index 10d4983c..56a6d6a1 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -7,7 +7,7 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print -def addTGLFcontrol(TGLFsettings, NS=2, minimal=False): +def addTGLFcontrol(code_settings, NS=2, minimal=False): """ ******************************************************************************** Define dictionary to start with @@ -15,8 +15,8 @@ def addTGLFcontrol(TGLFsettings, NS=2, minimal=False): """ # Minimum working set - if minimal or TGLFsettings == 0: - TGLFoptions = { + if minimal or code_settings == 0: + options = { "USE_MHD_RULE": True, "USE_BPER": False, "USE_BPAR": False, @@ -31,8 +31,8 @@ def addTGLFcontrol(TGLFsettings, NS=2, minimal=False): # Define every flag else: - TGLFoptions = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.tglf.controls", caseInsensitive=False) - TGLFoptions["NMODES"] = NS + 2 + options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.tglf.controls", caseInsensitive=False) + options["NMODES"] = NS + 2 """ ******************************************************************************** @@ -44,34 +44,51 @@ def addTGLFcontrol(TGLFsettings, NS=2, minimal=False): with open(__mitimroot__ / "templates" / "input.tglf.models.json", "r") as f: settings = json.load(f) - if str(TGLFsettings) in settings: - sett = settings[str(TGLFsettings)] + if str(code_settings) in settings: + sett = settings[str(code_settings)] for ikey in sett["controls"]: - TGLFoptions[ikey] = sett["controls"][ikey] + options[ikey] = sett["controls"][ikey] else: - print("\t- TGLFsettings not found in input.tglf.models.json, using defaults",typeMsg="w",) + print(f"\t- {code_settings = } not found in input.tglf.models.json, using defaults",typeMsg="w",) - return TGLFoptions + return options -def addNEOcontrol(*args, **kwargs): +def addNEOcontrol(code_settings,*args, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.neo.controls", caseInsensitive=False) return options -def addGXcontrol(*args, **kwargs): +def addGXcontrol(code_settings,*args, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.gx.controls", caseInsensitive=False) + """ + ******************************************************************************** + Standard sets of control parameters + (rest of parameters are as defaults) + ******************************************************************************** + """ + + with open(__mitimroot__ / "templates" / "input.gx.models.json", "r") as f: + settings = json.load(f) + + if str(code_settings) in settings: + sett = settings[str(code_settings)] + for ikey in sett["controls"]: + options[ikey] = sett["controls"][ikey] + else: + print("\t- code_settings not found in input.cgyro.models.json, using defaults",typeMsg="w") + return options def addCGYROcontrol(code_settings, rmin=None, **kwargs): - CGYROoptions = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.cgyro.controls", caseInsensitive=False) + options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.cgyro.controls", caseInsensitive=False) """ ******************************************************************************** - Standard sets of TGLF control parameters + Standard sets of control parameters (rest of parameters are as defaults) ******************************************************************************** """ @@ -82,11 +99,11 @@ def addCGYROcontrol(code_settings, rmin=None, **kwargs): if str(code_settings) in settings: sett = settings[str(code_settings)] for ikey in sett["controls"]: - CGYROoptions[ikey] = sett["controls"][ikey] + options[ikey] = sett["controls"][ikey] else: print("\t- code_settings not found in input.cgyro.models.json, using defaults",typeMsg="w") - return CGYROoptions + return options def TGLFinTRANSP(TGLFsettings, NS=3): diff --git a/src/mitim_tools/gyrokinetics_tools/GXtools.py b/src/mitim_tools/gyrokinetics_tools/GXtools.py index 86c94d3d..d8c8c752 100644 --- a/src/mitim_tools/gyrokinetics_tools/GXtools.py +++ b/src/mitim_tools/gyrokinetics_tools/GXtools.py @@ -1,4 +1,3 @@ -import numpy as np import netCDF4 import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools @@ -17,7 +16,7 @@ def __init__( super().__init__(rhos=rhos) def code_call(folder, n, p, additional_command="", **kwargs): - return f" gx {folder}/gxplasma.in &\n" + return f" gx {folder}/gxplasma.in > {folder}/gxplasma.mitim.log &\n" self.run_specifications = { 'code': 'gx', @@ -43,7 +42,8 @@ def code_call(folder, n, p, additional_command="", **kwargs): 'gxplasma.gx_geo.log', 'gxplasma.restart.nc', 'gxplasma.big.nc', - 'gxplasma.out.nc' + 'gxplasma.out.nc', + 'gxplasma.mitim.log' ] def plot( @@ -59,16 +59,60 @@ def plot( self.fn = GUItools.FigureNotebook("GX MITIM Notebook", geometry="1700x900", vertical=True) else: self.fn = fn - - fig1 = self.fn.add_figure(label=f"{extratitle}Summary", tab_color=fn_color) - - grid = plt.GridSpec(1, 2, hspace=0.7, wspace=0.2) if colors is None: colors = GRAPHICStools.listColors() - ax1 = fig1.add_subplot(grid[0, 0]) - ax2 = fig1.add_subplot(grid[0, 1]) + # Fluxes + fig = self.fn.add_figure(label=f"{extratitle}Fluxes", tab_color=fn_color) + + grid = plt.GridSpec(1, 3, hspace=0.7, wspace=0.2) + + ax1 = fig.add_subplot(grid[0, 0]) + ax2 = fig.add_subplot(grid[0, 1]) + ax3 = fig.add_subplot(grid[0, 2]) + + i = 0 + for label in labels: + for irho in range(len(self.rhos)): + c = self.results[label]['GXout'][irho] + + typeLs = '-' if c.t.shape[0]>20 else '-s' + + ax1.plot(c.t, c.Qe, typeLs, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) + ax2.plot(c.t, c.Qi, typeLs, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) + ax3.plot(c.t, c.Ge, typeLs, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) + + i += 1 + + for ax in [ax1, ax2, ax3]: + ax.set_xlabel("Time ($L_{ref}/c_s$)") + ax.set_xlim(left=0) + GRAPHICStools.addDenseAxis(ax) + + ax1.set_title('Electron heat flux') + ax1.set_ylabel("Electron heat flux ($Q_e/Q_{GB}$)") + ax1.legend(loc='best', prop={'size': 8}) + + ax2.set_title('Ion heat flux') + ax2.set_ylabel("Ion heat flux ($Q_i/Q_{GB}$)") + ax2.legend(loc='best', prop={'size': 8}) + + ax3.set_title('Electron particle flux') + ax3.set_ylabel("Electron particle flux ($\\Gamma_e/\\Gamma_{GB}$)") + ax3.legend(loc='best', prop={'size': 8}) + + plt.tight_layout() + + + # Linear stability + fig = self.fn.add_figure(label=f"{extratitle}Linear stability", tab_color=fn_color) + + grid = plt.GridSpec(2, 2, hspace=0.7, wspace=0.2) + + + ax1 = fig.add_subplot(grid[0, 0]) + ax2 = fig.add_subplot(grid[1, 0]) i = 0 for label in labels: @@ -80,14 +124,38 @@ def plot( i += 1 for ax in [ax1, ax2]: - ax.set_xlabel("Time") + ax.set_xlabel("Time ($L_{ref}/c_s$)") ax.set_xlim(left=0) GRAPHICStools.addDenseAxis(ax) ax1.set_ylabel("Real frequency") ax1.legend(loc='best', prop={'size': 4}) - ax2.set_ylabel("Growth rate") + + + ax3 = fig.add_subplot(grid[0, 1]) + ax4 = fig.add_subplot(grid[1, 1]) + + i = 0 + for label in labels: + for irho in range(len(self.rhos)): + c = self.results[label]['GXout'][irho] + ax3.plot(c.ky, c.w[-1, :], label=f"{label} rho={self.rhos[irho]}", color=colors[i]) + ax4.plot(c.ky, c.g[-1, :], label=f"{label} rho={self.rhos[irho]}", color=colors[i]) + i += 1 + + for ax in [ax3, ax4]: + ax.set_xlabel("$k_\\theta\\rho_s$") + ax.set_xlim(left=0) + GRAPHICStools.addDenseAxis(ax) + + ax3.set_ylabel("Real frequency") + ax3.legend(loc='best', prop={'size': 4}) + ax4.set_ylabel("Growth rate") + + plt.tight_layout() + + class GXinput(GACODErun.GACODEinput): def __init__(self, file=None): super().__init__( @@ -111,35 +179,81 @@ def write_state(self, file=None): # title: [controls], [plasma] blocks = { - 'Dimensions': + '': + [ ['debug'], [] ], + '[Dimensions]': [ ['ntheta', 'nperiod', 'ny', 'nx', 'nhermite', 'nlaguerre', 'nspecies'], [] ], - 'Domain': + '[Domain]': [ ['y0', 'boundary'], [] ], - 'Physics': + '[Physics]': [ ['nonlinear_mode', 'ei_colls'], ['beta'] ], - 'Time': + '[Time]': [ ['t_max', 'scheme'], [] ], - 'Initialization': + '[Initialization]': [ ['ikpar_init', 'init_field', 'init_amp', 'gaussian_init'], [] ], - 'Geometry': + '[Geometry]': [ ['geo_option'], ['rhoc', 'Rmaj', 'R_geo', 'shift', 'qinp', 'shat', 'akappa', 'akappri', 'tri', 'tripri', 'betaprim'] ], - 'Dissipation': + '[Dissipation]': [ ['closure_model', 'hypercollisions', 'nu_hyper_m', 'p_hyper_m', 'nu_hyper_l', 'p_hyper_l', 'hyper', 'D_hyper', 'p_hyper'], [] ], - 'Restart': - [ ['restart', 'save_for_restart'], [] ], - 'Diagnostics': + '[Restart]': + [ ['restart', 'save_for_restart', 'nsave'], [] ], + '[Diagnostics]': [ ['nwrite', 'omega', 'fluxes', 'fields'], [] ] } param_written = [] for block_name, params in blocks.items(): - param_written = self._write_block(f, f"[{block_name}]", params, param_written) + param_written = self._write_block(f, f"{block_name}", params, param_written) param_written = self._write_block_species(f, param_written) + # Check that parameters were all considerd in the blocks + for param in self.controls | self.plasma: + if param not in param_written: + print(f"Warning: {param} not written to file", typeMsg="q") + + def _write_block(self,f,name, param, param_written): + + # Local formatter: floats -> 6 significant figures in exponential (uppercase), + # ints stay as ints, bools as 0/1, sequences space-separated with same rule. + def _fmt_num(x): + import numpy as _np + if isinstance(x, (bool, _np.bool_)): + return "true" if x else "false" + if isinstance(x, (_np.floating, float)): + # 6 significant figures in exponential => 5 digits after decimal + return f"{float(x):.5E}" + if isinstance(x, (_np.integer, int)): + return f"{int(x)}" + return str(x) + + def _fmt_value(val): + import numpy as _np + if isinstance(val, (list, tuple, _np.ndarray)): + # Flatten numpy arrays but keep ordering; join with spaces + if isinstance(val, _np.ndarray): + flat = val.flatten().tolist() + else: + flat = list(val) + return " ".join(_fmt_num(v) for v in flat) + return _fmt_num(val) + + f.write(f'{name}\n') + for p in param[0]: + if p in self.controls: + f.write(f" {p.ljust(23)} = {_fmt_value(self.controls[p])}\n") + param_written.append(p) + for p in param[1]: + if p in self.plasma: + f.write(f" {p.ljust(23)} = {_fmt_value(self.plasma[p])}\n") + param_written.append(p) + f.write(f'\n') + + return param_written + def _write_block_species(self, f, param_written): # Local formatter: floats -> 6 significant figures in exponential (uppercase), @@ -206,44 +320,6 @@ def _fmt_value(val): return param_written - def _write_block(self,f,name, param, param_written): - - # Local formatter: floats -> 6 significant figures in exponential (uppercase), - # ints stay as ints, bools as 0/1, sequences space-separated with same rule. - def _fmt_num(x): - import numpy as _np - if isinstance(x, (bool, _np.bool_)): - return "true" if x else "false" - if isinstance(x, (_np.floating, float)): - # 6 significant figures in exponential => 5 digits after decimal - return f"{float(x):.5E}" - if isinstance(x, (_np.integer, int)): - return f"{int(x)}" - return str(x) - - def _fmt_value(val): - import numpy as _np - if isinstance(val, (list, tuple, _np.ndarray)): - # Flatten numpy arrays but keep ordering; join with spaces - if isinstance(val, _np.ndarray): - flat = val.flatten().tolist() - else: - flat = list(val) - return " ".join(_fmt_num(v) for v in flat) - return _fmt_num(val) - - f.write(f'{name}\n') - for p in param[0]: - f.write(f" {p.ljust(23)} = {_fmt_value(self.controls[p])}\n") - param_written.append(p) - for p in param[1]: - f.write(f" {p.ljust(23)} = {_fmt_value(self.plasma[p])}\n") - param_written.append(p) - f.write(f'\n') - - return param_written - - class GXoutput(GACODErun.GACODEoutput): def __init__(self, FolderGACODE, suffix="", **kwargs): super().__init__() @@ -271,9 +347,15 @@ def read(self): self.w = data.groups['Diagnostics'].variables['omega_kxkyt'][:,1:,ikx,0] # (time, ky) self.g = data.groups['Diagnostics'].variables['omega_kxkyt'][:,1:,ikx,1] # (time, ky) - # get fluxes - # Qi = data.groups['Fluxes'].variables['qflux'][:,1] - # Qe = data.groups['Fluxes'].variables['qflux'][:,0] - # Ge = data.groups['Fluxes'].variables['pflux'][:,0] - + # Fluxes + Q = data.groups['Diagnostics'].variables['HeatFlux_st'] # (time, species) + G = data.groups['Diagnostics'].variables['ParticleFlux_st'] # (time, species) + + # Assume electrons are always last + self.Qe = Q[:,-1] + self.QiAll = Q[:,:-1] + self.Qi = self.QiAll.sum(axis=1) + self.Ge = G[:,-1] + self.GiAll = G[:,:-1] + self.Gi = self.GiAll.sum(axis=1) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index ef7569d0..fa27ec91 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2276,14 +2276,12 @@ def to_tglf(self, r=[0.5], TGLFsettings=1, r_is_rho = True): def interpolator(y): return interpolation_function(rho, r_interpolation,y).item() - - TGLFoptions = GACODEdefaults.addTGLFcontrol(TGLFsettings) # --------------------------------------------------------------------------------------------------------------------------------------- # Controls come from options # --------------------------------------------------------------------------------------------------------------------------------------- - controls = TGLFoptions + controls = GACODEdefaults.addTGLFcontrol(TGLFsettings) # --------------------------------------------------------------------------------------------------------------------------------------- # Species come from profiles @@ -2442,7 +2440,7 @@ def interpolator(y): # Controls come from options # --------------------------------------------------------------------------------------------------------------------------------------- - controls = GACODEdefaults.addNEOcontrol() + controls = GACODEdefaults.addNEOcontrol(0) # --------------------------------------------------------------------------------------------------------------------------------------- # Species come from profiles @@ -2735,7 +2733,7 @@ def interpolator(y): # Controls come from options # --------------------------------------------------------------------------------------------------------------------------------------- - controls = GACODEdefaults.addGXcontrol() + controls = GACODEdefaults.addGXcontrol(0) # --------------------------------------------------------------------------------------------------------------------------------------- # Species come from profiles diff --git a/templates/input.gx.controls b/templates/input.gx.controls index e4bf3841..5bf028a1 100644 --- a/templates/input.gx.controls +++ b/templates/input.gx.controls @@ -2,6 +2,8 @@ # GX example #============================================================= + debug = false + [Dimensions] ntheta = 24 # number of points along field line (theta) per 2pi segment nperiod = 1 # number of 2pi segments along field line is 2*nperiod-1 @@ -17,7 +19,7 @@ boundary = "linked" # use twist-shift boundary conditions along field line [Physics] - nonlinear_mode = false # this is a linear calculation + nonlinear_mode = true # this is a linear calculation ei_colls = false # turn off electron-ion collisions [Time] @@ -48,11 +50,11 @@ [Restart] restart = false save_for_restart = true - + nsave = 100 + [Diagnostics] nwrite = 1000 # write diagnostics every nwrite timesteps + nsave = 1000 # save restart file every nsave timesteps omega = true # compute and write linear growth rate and frequency fluxes = true # compute and write heat and particle fluxes fields = true # compute and write electric and magnetic fields - -[species] diff --git a/templates/input.gx.models.json b/templates/input.gx.models.json new file mode 100644 index 00000000..b14b6f86 --- /dev/null +++ b/templates/input.gx.models.json @@ -0,0 +1,16 @@ +{ + "0": { + "label": "Linear with 10 ky modes", + "controls": { + "nonlinear_mode": false, + "nx": 1, + "ny": 28 + } + }, + "1": { + "label": "Nonlinear", + "controls": { + "nonlinear_mode": true + } + } +} \ No newline at end of file diff --git a/tests/GX_workflow.py b/tests/GX_workflow.py index d4932b11..aecdccd1 100644 --- a/tests/GX_workflow.py +++ b/tests/GX_workflow.py @@ -10,10 +10,9 @@ folder = __mitimroot__ / "tests" / "scratch" / "gx_test" input_gacode = __mitimroot__ / "tests" / "data" / "input.gacode" -# Only 1 ion species +# Reduce the ion species p = gacode_state(input_gacode) -p.lumpSpecies(ions_list=[1,2,3,4]) -# +p.lumpImpurities() if cold_start and folder.exists(): os.system(f"rm -r {folder.resolve()}") @@ -24,8 +23,9 @@ gx.run( 'gx1/', cold_start=cold_start, + code_settings=0, # Linear extraOptions={ - 't_max':5.0, # Short, I just want to test the run + 't_max':50.0, # Short, I just want to test the run }, ) gx.read('gx1') diff --git a/tests/TGLF_workflow.py b/tests/TGLF_workflow.py index 09118dcf..4c154920 100644 --- a/tests/TGLF_workflow.py +++ b/tests/TGLF_workflow.py @@ -17,7 +17,7 @@ tglf.run( "run1/", - code_settings=None, + code_settings=0, cold_start=cold_start, runWaveForms = [0.67, 10.0], forceIfcold_start=True, @@ -25,20 +25,31 @@ slurm_setup={"cores": 4, "minutes": 1}, ) -tglf.read(label="ES") +tglf.read(label="ES (0)") tglf.run( "run2/", - code_settings=None, + code_settings=0, cold_start=cold_start, forceIfcold_start=True, extraOptions={"USE_BPER": True, "USE_BPAR": True}, slurm_setup={"cores": 4, "minutes": 1}, ) -tglf.read(label="EM") +tglf.read(label="EM (0)") -tglf.plot(labels=["EM","ES"]) +tglf.run( + "run3/", + code_settings=6, + cold_start=cold_start, + forceIfcold_start=True, + extraOptions={"USE_BPER": True, "USE_BPAR": True}, + slurm_setup={"cores": 4, "minutes": 1}, +) + +tglf.read(label="EM (6)") + +tglf.plot(labels=["EM (0)","ES (0)", "EM (6)"]) # Required if running in non-interactive mode tglf.fn.show() \ No newline at end of file From 5dc30691e6fe856534b02d80cdfce082257e6f88 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 25 Aug 2025 11:51:11 -0400 Subject: [PATCH 195/385] Improved GX --- src/mitim_tools/gyrokinetics_tools/GXtools.py | 69 ++++++++++--------- templates/input.gx.controls | 18 +++-- tests/GX_workflow.py | 10 ++- 3 files changed, 53 insertions(+), 44 deletions(-) diff --git a/src/mitim_tools/gyrokinetics_tools/GXtools.py b/src/mitim_tools/gyrokinetics_tools/GXtools.py index d8c8c752..d556f3bb 100644 --- a/src/mitim_tools/gyrokinetics_tools/GXtools.py +++ b/src/mitim_tools/gyrokinetics_tools/GXtools.py @@ -92,15 +92,15 @@ def plot( ax1.set_title('Electron heat flux') ax1.set_ylabel("Electron heat flux ($Q_e/Q_{GB}$)") - ax1.legend(loc='best', prop={'size': 8}) + ax1.legend(loc='best', prop={'size': 12}) ax2.set_title('Ion heat flux') ax2.set_ylabel("Ion heat flux ($Q_i/Q_{GB}$)") - ax2.legend(loc='best', prop={'size': 8}) + ax2.legend(loc='best', prop={'size': 12}) ax3.set_title('Electron particle flux') ax3.set_ylabel("Electron particle flux ($\\Gamma_e/\\Gamma_{GB}$)") - ax3.legend(loc='best', prop={'size': 8}) + ax3.legend(loc='best', prop={'size': 12}) plt.tight_layout() @@ -118,9 +118,12 @@ def plot( for label in labels: for irho in range(len(self.rhos)): c = self.results[label]['GXout'][irho] + + typeLs = '-' if c.t.shape[0]>20 else '-s' + for iky in range(len(c.ky)): - ax1.plot(c.t, c.w[:, iky], label=f"{label} rho={self.rhos[irho]} ky={c.ky[iky]}", color=colors[i]) - ax2.plot(c.t, c.g[:, iky], label=f"{label} rho={self.rhos[irho]} ky={c.ky[iky]}", color=colors[i]) + ax1.plot(c.t, c.w[:, iky], typeLs, label=f"{label} rho={self.rhos[irho]} ky={c.ky[iky]}", color=colors[i]) + ax2.plot(c.t, c.g[:, iky], typeLs, label=f"{label} rho={self.rhos[irho]} ky={c.ky[iky]}", color=colors[i]) i += 1 for ax in [ax1, ax2]: @@ -140,8 +143,8 @@ def plot( for label in labels: for irho in range(len(self.rhos)): c = self.results[label]['GXout'][irho] - ax3.plot(c.ky, c.w[-1, :], label=f"{label} rho={self.rhos[irho]}", color=colors[i]) - ax4.plot(c.ky, c.g[-1, :], label=f"{label} rho={self.rhos[irho]}", color=colors[i]) + ax3.plot(c.ky, c.w[-1, :], '-s', markersize = 5, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) + ax4.plot(c.ky, c.g[-1, :], '-s', markersize = 5, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) i += 1 for ax in [ax3, ax4]: @@ -150,8 +153,10 @@ def plot( GRAPHICStools.addDenseAxis(ax) ax3.set_ylabel("Real frequency") - ax3.legend(loc='best', prop={'size': 4}) + ax3.legend(loc='best', prop={'size': 12}) + ax3.axhline(y=0, color='k', linestyle='--', linewidth=1) ax4.set_ylabel("Growth rate") + ax4.set_ylim(bottom=0) plt.tight_layout() @@ -182,13 +187,13 @@ def write_state(self, file=None): '': [ ['debug'], [] ], '[Dimensions]': - [ ['ntheta', 'nperiod', 'ny', 'nx', 'nhermite', 'nlaguerre', 'nspecies'], [] ], + [ ['ntheta', 'nperiod', 'ny', 'nx', 'nhermite', 'nlaguerre'], ['nspecies'] ], '[Domain]': [ ['y0', 'boundary'], [] ], '[Physics]': [ ['nonlinear_mode', 'ei_colls'], ['beta'] ], '[Time]': - [ ['t_max', 'scheme'], [] ], + [ ['t_max', 'scheme', 'dt', 'nstep'], [] ], '[Initialization]': [ ['ikpar_init', 'init_field', 'init_amp', 'gaussian_init'], [] ], '[Geometry]': @@ -244,11 +249,13 @@ def _fmt_value(val): f.write(f'{name}\n') for p in param[0]: if p in self.controls: - f.write(f" {p.ljust(23)} = {_fmt_value(self.controls[p])}\n") + if self.controls[p] is not None: + f.write(f" {p.ljust(23)} = {_fmt_value(self.controls[p])}\n") param_written.append(p) for p in param[1]: if p in self.plasma: - f.write(f" {p.ljust(23)} = {_fmt_value(self.plasma[p])}\n") + if self.plasma[p] is not None: + f.write(f" {p.ljust(23)} = {_fmt_value(self.plasma[p])}\n") param_written.append(p) f.write(f'\n') @@ -281,41 +288,41 @@ def _fmt_value(val): return _fmt_num(val) self.num_recorded = 0 - for i in range(100): + for i in range(1000): if f"z_{i+1}" in self.plasma: self.num_recorded += 1 else: break - z, dens, temp, mass, fprim, tprim, vnewk, typeS = '[ ', '[ ', '[ ', '[ ', '[ ', '[ ', '[ ', '[ ' + z, dens, temp, mass, fprim, tprim, vnewk, typeS = '[ ', '[ ', '[ ', '[ ', '[ ', '[ ', '[ ', '[ ' for i in range(self.num_recorded): - typeS += f'"{_fmt_value(self.plasma[f"type_{i+1}"])}", ' - z += f'{_fmt_value(self.plasma[f"z_{i+1}"])}, ' - mass += f'{_fmt_value(self.plasma[f"mass_{i+1}"])}, ' - dens += f'{_fmt_value(self.plasma[f"dens_{i+1}"])}, ' - temp += f'{_fmt_value(self.plasma[f"temp_{i+1}"])}, ' - fprim += f'{_fmt_value(self.plasma[f"fprim_{i+1}"])}, ' - tprim += f'{_fmt_value(self.plasma[f"tprim_{i+1}"])}, ' - vnewk += f'{_fmt_value(self.plasma[f"vnewk_{i+1}"])}, ' + typeS += f'"{_fmt_value(self.plasma[f"type_{i+1}"])}", ' + z += f'{_fmt_value(self.plasma[f"z_{i+1}"])}, ' + mass += f'{_fmt_value(self.plasma[f"mass_{i+1}"])}, ' + dens += f'{_fmt_value(self.plasma[f"dens_{i+1}"])}, ' + temp += f'{_fmt_value(self.plasma[f"temp_{i+1}"])}, ' + fprim += f'{_fmt_value(self.plasma[f"fprim_{i+1}"])}, ' + tprim += f'{_fmt_value(self.plasma[f"tprim_{i+1}"])}, ' + vnewk += f'{_fmt_value(self.plasma[f"vnewk_{i+1}"])}, ' param_written.append(f"type_{i+1}") param_written.append(f"z_{i+1}") + param_written.append(f"mass_{i+1}") param_written.append(f"dens_{i+1}") param_written.append(f"temp_{i+1}") - param_written.append(f"mass_{i+1}") param_written.append(f"fprim_{i+1}") param_written.append(f"tprim_{i+1}") param_written.append(f"vnewk_{i+1}") f.write("[species]\n") - f.write(f" z = {z[:-2]} ]\n") - f.write(f" dens = {dens[:-2]} ]\n") - f.write(f" temp = {temp[:-2]} ]\n") - f.write(f" mass = {mass[:-2]} ]\n") - f.write(f" fprim = {fprim[:-2]} ]\n") - f.write(f" tprim = {tprim[:-2]} ]\n") - f.write(f" vnewk = {vnewk[:-2]} ]\n") - f.write(f" type = {typeS[:-2]} ]\n") + f.write(f" z = {z[:-4]} ]\n") + f.write(f" dens = {dens[:-4]} ]\n") + f.write(f" temp = {temp[:-4]} ]\n") + f.write(f" mass = {mass[:-4]} ]\n") + f.write(f" fprim = {fprim[:-4]} ]\n") + f.write(f" tprim = {tprim[:-4]} ]\n") + f.write(f" vnewk = {vnewk[:-4]} ]\n") + f.write(f" type = {typeS[:-4]} ]\n") f.write("\n") return param_written diff --git a/templates/input.gx.controls b/templates/input.gx.controls index 5bf028a1..4fa41b76 100644 --- a/templates/input.gx.controls +++ b/templates/input.gx.controls @@ -5,17 +5,16 @@ debug = false [Dimensions] - ntheta = 24 # number of points along field line (theta) per 2pi segment - nperiod = 1 # number of 2pi segments along field line is 2*nperiod-1 + ntheta = 32 # number of points along field line (theta) per 2pi segment + nperiod = 2 # number of 2pi segments along field line is 2*nperiod-1 ny = 106 # number of real-space grid-points in y, nky = 1 + (ny-1)/3 --> 36 ky modes nx = 382 # number of real-space grid-points in x, nkx = 1 + 2*(nx-1)/3 --> 255 kx modes - nhermite = 8 # number of hermite moments (v_parallel resolution) - nlaguerre = 4 # number of laguerre moments (mu B resolution) - nspecies = 1 # number of evolved kinetic species (adiabatic electrons don't count towards nspecies) + nhermite = 48 # number of hermite moments (v_parallel resolution) + nlaguerre = 16 # number of laguerre moments (mu B resolution) [Domain] - y0 = 14.0 # controls box length in y (in units of rho_ref) and minimum ky, so that ky_min*rho_ref = 1/y0 --> ky_min = 0.0714, ky_max = 2.57, L_y = 88 + y0 = 20.0 # controls box length in y (in units of rho_ref) and minimum ky, so that ky_min*rho_ref = 1/y0 --> ky_min = 0.0714, ky_max = 2.57, L_y = 88 boundary = "linked" # use twist-shift boundary conditions along field line [Physics] @@ -33,7 +32,7 @@ init_amp = 1.0e-3 # amplitude of initial condition [Geometry] - geo_option = "miller" # use Miller geometry + geo_option = "miller" # use Miller geometry (values provided below) [Dissipation] closure_model = "none" # no closure assumptions (just truncation) @@ -50,11 +49,10 @@ [Restart] restart = false save_for_restart = true - nsave = 100 - + nsave = 100 # save restart file every nsave timesteps + [Diagnostics] nwrite = 1000 # write diagnostics every nwrite timesteps - nsave = 1000 # save restart file every nsave timesteps omega = true # compute and write linear growth rate and frequency fluxes = true # compute and write heat and particle fluxes fields = true # compute and write electric and magnetic fields diff --git a/tests/GX_workflow.py b/tests/GX_workflow.py index aecdccd1..0d4ddf08 100644 --- a/tests/GX_workflow.py +++ b/tests/GX_workflow.py @@ -10,9 +10,10 @@ folder = __mitimroot__ / "tests" / "scratch" / "gx_test" input_gacode = __mitimroot__ / "tests" / "data" / "input.gacode" -# Reduce the ion species +# Reduce the ion species to just 1 p = gacode_state(input_gacode) -p.lumpImpurities() +p.lumpSpecies(ions_list=[1,2,3,4]) +# if cold_start and folder.exists(): os.system(f"rm -r {folder.resolve()}") @@ -25,7 +26,10 @@ cold_start=cold_start, code_settings=0, # Linear extraOptions={ - 't_max':50.0, # Short, I just want to test the run + 't_max':500.0, + # 'dt': 1.0, # timestep size (in units of L_ref/vt_ref) + # 'nstep': 10, + # 'nwrite': 1, }, ) gx.read('gx1') From 6d239f1cca46a1b31e28f17277e2cda502610ca8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 09:33:31 -0400 Subject: [PATCH 196/385] Added specification of machine's GPUs --- src/mitim_tools/misc_tools/CONFIGread.py | 1 + templates/config_user_example.json | 15 ++++++++++----- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/mitim_tools/misc_tools/CONFIGread.py b/src/mitim_tools/misc_tools/CONFIGread.py index b759c969..e59f64ba 100644 --- a/src/mitim_tools/misc_tools/CONFIGread.py +++ b/src/mitim_tools/misc_tools/CONFIGread.py @@ -125,6 +125,7 @@ def machineSettings( "folderWork": scratch, "slurm": {}, "cores_per_node": s[machine].get("cores_per_node", None), + "gpus_per_node": s[machine].get("gpus_per_node", 0), "isTunnelSameMachine": ( bool(s[machine]["isTunnelSameMachine"]) if "isTunnelSameMachine" in s[machine] diff --git a/templates/config_user_example.json b/templates/config_user_example.json index a1644094..42f9f150 100644 --- a/templates/config_user_example.json +++ b/templates/config_user_example.json @@ -23,7 +23,8 @@ "username": "exampleusername", "scratch": "/Users/exampleusername/scratch/", "modules": "", - "cores_per_node": 8 + "cores_per_node": 8, + "gpus_per_node": 0 }, "perlmutter": { "machine": "perlmutter.nersc.gov", @@ -38,7 +39,8 @@ "exclusive": false, "email": "optional@email" }, - "cores_per_node": 32 + "cores_per_node": 32, + "gpus_per_node": 0 }, "engaging": { "machine": "eofe7.mit.edu", @@ -49,7 +51,8 @@ }, "scratch": "/nobackup1/exampleusername/", "modules": "", - "cores_per_node": 32 + "cores_per_node": 32, + "gpus_per_node": 0 }, "mfews": { "machine": "mfews02.psfc.mit.edu", @@ -57,7 +60,8 @@ "port": 9224, "scratch": "/home/exampleusername/scratch/", "modules": "", - "cores_per_node": 8 + "cores_per_node": 8, + "gpus_per_node": 0 }, "globus": { "username": "exampleusername", @@ -78,6 +82,7 @@ "email": "optional@email", "exclude": "node584" }, - "cores_per_node": 8 + "cores_per_node": 8, + "gpus_per_node": 0 } } \ No newline at end of file From 04f62c51c3554ba0e6d62b3550aba8abfe91f275 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 13:31:48 -0400 Subject: [PATCH 197/385] misc GX --- templates/input.gx.controls | 23 ++++++++++++++--------- templates/input.gx.models.json | 3 +++ tests/CGYRO_workflow.py | 5 ++--- tests/GX_workflow.py | 13 ++++++------- 4 files changed, 25 insertions(+), 19 deletions(-) diff --git a/templates/input.gx.controls b/templates/input.gx.controls index 4fa41b76..f6e52998 100644 --- a/templates/input.gx.controls +++ b/templates/input.gx.controls @@ -6,7 +6,7 @@ [Dimensions] ntheta = 32 # number of points along field line (theta) per 2pi segment - nperiod = 2 # number of 2pi segments along field line is 2*nperiod-1 + nperiod = 1 # number of 2pi segments along field line is 2*nperiod-1 ny = 106 # number of real-space grid-points in y, nky = 1 + (ny-1)/3 --> 36 ky modes nx = 382 # number of real-space grid-points in x, nkx = 1 + 2*(nx-1)/3 --> 255 kx modes @@ -23,7 +23,7 @@ [Time] t_max = 1000.0 # end time (in units of L_ref/vt_ref) - scheme = "rk4" # use RK3 timestepping scheme (with adaptive timestepping) + scheme = "rk3" # use RK3 timestepping scheme (with adaptive timestepping) [Initialization] gaussian_init = true # initial perturbation is a gaussian in theta @@ -37,22 +37,27 @@ [Dissipation] closure_model = "none" # no closure assumptions (just truncation) hypercollisions = true # use hypercollision model - nu_hyper_m = 0.1 # coefficient of hermite hypercollisions + hyper = true # use hyperdiffusion + HB_hyper = false # use Hammett-Belli hyperdiffusivity model + + nu_hyper_m = 0.5 # coefficient of hermite hypercollisions + nu_hyper_l = 0.5 # coefficient of laguerre hypercollisions + D_hyper = 0.5 # coefficient of hyperdiffusion p_hyper_m = 6 # exponent of hermite hypercollisions - nu_hyper_l = 0.1 # coefficient of laguerre hypercollisions p_hyper_l = 6 # exponent of laguerre hypercollisions - - hyper = true # use hyperdiffusion - D_hyper = 0.05 # coefficient of hyperdiffusion p_hyper = 2 # exponent of hyperdiffusion is 2*p_hyper = 4 + D_H = 0.5 # coefficient of H-B hyperdiffusion + w_osc = 0.0 # frequency parameter in the H-B model + p_HB = 2 # exponent for the H-B model [Restart] restart = false save_for_restart = true - nsave = 100 # save restart file every nsave timesteps + nsave = 100 # save restart file every nsave timesteps [Diagnostics] - nwrite = 1000 # write diagnostics every nwrite timesteps + nwrite = 100 # write diagnostics every nwrite timesteps (this is NOT a/c_s) (=1 doesn't work) omega = true # compute and write linear growth rate and frequency fluxes = true # compute and write heat and particle fluxes fields = true # compute and write electric and magnetic fields + moments = true # write moments on the grid diff --git a/templates/input.gx.models.json b/templates/input.gx.models.json index b14b6f86..8bd4c885 100644 --- a/templates/input.gx.models.json +++ b/templates/input.gx.models.json @@ -3,6 +3,9 @@ "label": "Linear with 10 ky modes", "controls": { "nonlinear_mode": false, + "hyper": false, + "HB_hyper": false, + "nperiod": 2, "nx": 1, "ny": 28 } diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index c2c50969..d13e8e9b 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -25,7 +25,7 @@ code_settings=0, extraOptions={ 'KY':0.5, - 'MAX_TIME': 10.0, # Short, I just want to test the run + 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file }, slurm_setup={ 'cores':8 @@ -42,9 +42,8 @@ cgyro.run_scan( 'scan1', - cold_start=cold_start, extraOptions={ - 'MAX_TIME': 10.0, + 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file }, variable='KY', varUpDown=[0.3,0.4], diff --git a/tests/GX_workflow.py b/tests/GX_workflow.py index 0d4ddf08..0054592e 100644 --- a/tests/GX_workflow.py +++ b/tests/GX_workflow.py @@ -12,13 +12,13 @@ # Reduce the ion species to just 1 p = gacode_state(input_gacode) -p.lumpSpecies(ions_list=[1,2,3,4]) -# +p.lumpIons() +# -------------------------------- if cold_start and folder.exists(): os.system(f"rm -r {folder.resolve()}") -gx = GXtools.GX(rhos=[0.5,0.7]) +gx = GXtools.GX(rhos=[0.5, 0.6]) gx.prep(p, folder) gx.run( @@ -26,10 +26,9 @@ cold_start=cold_start, code_settings=0, # Linear extraOptions={ - 't_max':500.0, - # 'dt': 1.0, # timestep size (in units of L_ref/vt_ref) - # 'nstep': 10, - # 'nwrite': 1, + 't_max':10.0, # Run up to 50.0 a/c_s + 'y0' :5.0, # kymin = 1/y0 = 0.2 + 'ny': 34, # nky = 1 + (ny-1)/3 = 12 -> ky_range = 0.2 - 2.4 }, ) gx.read('gx1') From f2cff7993177686e460d4dbdc19dd5b3db8b0d6c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 13:32:07 -0400 Subject: [PATCH 198/385] Lump all ions --- src/mitim_tools/plasmastate_tools/MITIMstate.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index fa27ec91..cb8734af 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1417,6 +1417,10 @@ def lumpImpurities(self): self.lumpSpecies(ions_list=self.ion_list_impurities) + def lumpIons(self): + + self.lumpSpecies(ions_list=self.ion_list_main+self.ion_list_impurities) + def lumpDT(self): if self.DTplasmaBool: @@ -2775,7 +2779,7 @@ def interpolator(y): # --------------------------------------------------------------------------------------------------------------------------------------- plasma = { - 'nspecies': len(species)-1 # do not count electrons} + 'nspecies': len(species) } parameters = { From 1723391b0e9c601d01307508e7570795055e99bc Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 16:00:16 -0400 Subject: [PATCH 199/385] Implemented json idea for transport models in PORTALS --- .../physics_models/transport_analytic.py | 12 +- .../physics_models/transport_cgyroneo.py | 57 ++++++++-- .../physics_models/transport_tglfneo.py | 64 ++++++----- .../powertorch/utils/TRANSPORTtools.py | 105 ++++++++++++++++-- 4 files changed, 184 insertions(+), 54 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py index 57ac1922..2558b3bc 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -43,14 +43,14 @@ def evaluate(self): self.powerstate.plasma["a"].unsqueeze(-1), ) - self.powerstate.plasma["QeMWm2_tr_turb"] = Pe_tr * 2 / 3 - self.powerstate.plasma["QiMWm2_tr_turb"] = Pi_tr * 2 / 3 + self.__dict__["QeMWm2_tr_turb"] = Pe_tr * 2 / 3 + self.__dict__["QiMWm2_tr_turb"] = Pi_tr * 2 / 3 - self.powerstate.plasma["QeMWm2_tr_neoc"] = Pe_tr * 1 / 3 - self.powerstate.plasma["QiMWm2_tr_neoc"] = Pi_tr * 1 / 3 + self.__dict__["QeMWm2_tr_neoc"] = Pe_tr * 1 / 3 + self.__dict__["QiMWm2_tr_neoc"] = Pi_tr * 1 / 3 - self.powerstate.plasma["QeMWm2_tr"] = self.powerstate.plasma["QeMWm2_tr_turb"] + self.powerstate.plasma["QeMWm2_tr_neoc"] - self.powerstate.plasma["QiMWm2_tr"] = self.powerstate.plasma["QiMWm2_tr_turb"] + self.powerstate.plasma["QiMWm2_tr_neoc"] + self.__dict__["QeMWm2_tr"] = self.__dict__["QeMWm2_tr_turb"] + self.__dict__["QeMWm2_tr_neoc"] + self.__dict__["QiMWm2_tr"] = self.__dict__["QiMWm2_tr_turb"] + self.__dict__["QiMWm2_tr_neoc"] # ------------------------------------------------------------------ # SURROGATE diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index f7021545..f5e421a6 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -1,8 +1,12 @@ +from mitim_tools.misc_tools import IOtools +from functools import partial from mitim_tools.gacode_tools import CGYROtools +from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_modules.powertorch.physics_models import transport_tglfneo from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +# Inherit from transport_tglfneo.tglfneo_model so that I have the NEO evaluator class cgyroneo_model(transport_tglfneo.tglfneo_model): def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) @@ -14,6 +18,8 @@ def evaluate_turbulence(self): transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] cold_start = transport_evaluator_options.get("cold_start", False) + + automatic_run = False # ------------------------------------------------------------------------------------------------------------------------ # Prepare CGYRO object @@ -37,16 +43,43 @@ def evaluate_turbulence(self): self.folder, ) - _ = cgyro.run( - 'base_cgyro', - full_submission=False, - code_settings=1, - cold_start=cold_start, - forceIfcold_start=True, - ) - - # cgyro.read( - # label='base_cgyro' - # ) + if automatic_run: + raise Exception("[MITIM] automatic_run not implemented yet") + + # cgyro.read( + # label='base_cgyro' + # ) + + # TRANSPORTtools.write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb') - embed() \ No newline at end of file + else: + + # Just prepare everything but do not submit (full_submission=False) + _ = cgyro.run( + 'base_cgyro', + full_submission=False, + code_settings=1, + cold_start=cold_start, + forceIfcold_start=True, + ) + + # Wait until the user has placed the json file in the right folder + + attempts = 0 + while (self.folder / 'fluxes_turb.json').exists() is False: + if attempts > 0: + print(f"\n\n !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!", typeMsg='i') + print(f"\tMITIM could not find the file", typeMsg='i') + print(f" !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! \n\n", typeMsg='i') + logic_to_wait(self.folder) + attempts += 1 + +def logic_to_wait(folder): + print(f" **** CGYRO prepared. Please, run CGYRO from the simulation setup in folder: ", typeMsg='i') + print(f"\t {folder}/base_cgyro\n", typeMsg='i') + print(f" **** When finished, the fluxes_turb.json file should be placed in:", typeMsg='i') + print(f"\t {folder}/fluxes_turb.json\n", typeMsg='i') + print(f" **** When you have done that, please write an 'exit' and enter (for continuing)\n", typeMsg='i') + + embed() + diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 1f90258c..b387b2ca 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -2,6 +2,7 @@ import shutil import numpy as np from mitim_tools.misc_tools import IOtools +from functools import partial from mitim_tools.gacode_tools import TGLFtools, NEOtools from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -17,7 +18,7 @@ def produce_profiles(self): # ************************************************************************************ # Private functions for the evaluation # ************************************************************************************ - + @IOtools.hook_method(after=partial(TRANSPORTtools.write_json, file_name = 'fluxes_turb.json', suffix= 'turb')) def evaluate_turbulence(self): # ------------------------------------------------------------------------------------------------------------------------ @@ -125,33 +126,34 @@ def evaluate_turbulence(self): self._raise_warnings(tglf, rho_locations, Qi_includes_fast) # ------------------------------------------------------------------------------------------------------------------------ - # Pass the information to POWERSTATE + # Pass the information # ------------------------------------------------------------------------------------------------------------------------ - self.powerstate.plasma["QeMWm2_tr_turb"] = Flux_mean[0] - self.powerstate.plasma["QeMWm2_tr_turb_stds"] = Flux_std[0] + self.QeMWm2_tr_turb = Flux_mean[0] + self.QeMWm2_tr_turb_stds = Flux_std[0] - self.powerstate.plasma["QiMWm2_tr_turb"] = Flux_mean[1] - self.powerstate.plasma["QiMWm2_tr_turb_stds"] = Flux_std[1] + self.QiMWm2_tr_turb = Flux_mean[1] + self.QiMWm2_tr_turb_stds = Flux_std[1] - self.powerstate.plasma["Ge1E20m2_tr_turb"] = Flux_mean[2] - self.powerstate.plasma["Ge1E20m2_tr_turb_stds"] = Flux_std[2] + self.Ge1E20m2_tr_turb = Flux_mean[2] + self.Ge1E20m2_tr_turb_stds = Flux_std[2] - self.powerstate.plasma["GZ1E20m2_tr_turb"] = Flux_mean[3] - self.powerstate.plasma["GZ1E20m2_tr_turb_stds"] = Flux_std[3] + self.GZ1E20m2_tr_turb = Flux_mean[3] + self.GZ1E20m2_tr_turb_stds = Flux_std[3] - self.powerstate.plasma["MtJm2_tr_turb"] = Flux_mean[4] - self.powerstate.plasma["MtJm2_tr_turb_stds"] = Flux_std[4] + self.MtJm2_tr_turb = Flux_mean[4] + self.MtJm2_tr_turb_stds = Flux_std[4] if provideTurbulentExchange: - self.powerstate.plasma["QieMWm3_tr_turb"] = Flux_mean[5] - self.powerstate.plasma["QieMWm3_tr_turb_stds"] = Flux_std[5] + self.QieMWm3_tr_turb = Flux_mean[5] + self.QieMWm3_tr_turb_stds = Flux_std[5] else: - self.powerstate.plasma["QieMWm3_tr_turb"] = Flux_mean[5] * 0.0 - self.powerstate.plasma["QieMWm3_tr_turb_stds"] = Flux_std[5] * 0.0 + self.QieMWm3_tr_turb = Flux_mean[5] * 0.0 + self.QieMWm3_tr_turb_stds = Flux_std[5] * 0.0 return tglf + @IOtools.hook_method(after=partial(TRANSPORTtools.write_json, file_name = 'fluxes_neoc.json', suffix= 'neoc')) def evaluate_neoclassical(self): # Options @@ -188,20 +190,24 @@ def evaluate_neoclassical(self): GZ = np.array([neo.results['base']['NEOout'][i].GiAll_unn[impurityPosition] for i in range(len(rho_locations))]) Mt = np.array([neo.results['base']['NEOout'][i].Mt_unn for i in range(len(rho_locations))]) - self.powerstate.plasma["QeMWm2_tr_neoc"] = Qe - self.powerstate.plasma["QiMWm2_tr_neoc"] = Qi - self.powerstate.plasma["Ge1E20m2_tr_neoc"] = Ge - self.powerstate.plasma["GZ1E20m2_tr_neoc"] = GZ - self.powerstate.plasma["MtJm2_tr_neoc"] = Mt + # ------------------------------------------------------------------------------------------------------------------------ + # Pass the information + # ------------------------------------------------------------------------------------------------------------------------ + + self.QeMWm2_tr_neoc = Qe + self.QiMWm2_tr_neoc = Qi + self.Ge1E20m2_tr_neoc = Ge + self.GZ1E20m2_tr_neoc = GZ + self.MtJm2_tr_neoc = Mt - self.powerstate.plasma["QeMWm2_tr_neoc_stds"] = abs(Qe) * percentError[1]/100.0 - self.powerstate.plasma["QiMWm2_tr_neoc_stds"] = abs(Qi) * percentError[1]/100.0 - self.powerstate.plasma["Ge1E20m2_tr_neoc_stds"] = abs(Ge) * percentError[1]/100.0 - self.powerstate.plasma["GZ1E20m2_tr_neoc_stds"] = abs(GZ) * percentError[1]/100.0 - self.powerstate.plasma["MtJm2_tr_neoc_stds"] = abs(Mt) * percentError[1]/100.0 - - self.powerstate.plasma["QieMWm3_tr_neoc"] = Qe * 0.0 - self.powerstate.plasma["QieMWm3_tr_neoc_stds"] = Qe * 0.0 + self.QeMWm2_tr_neoc_stds = abs(Qe) * percentError[1]/100.0 + self.QiMWm2_tr_neoc_stds = abs(Qi) * percentError[1]/100.0 + self.Ge1E20m2_tr_neoc_stds = abs(Ge) * percentError[1]/100.0 + self.GZ1E20m2_tr_neoc_stds = abs(GZ) * percentError[1]/100.0 + self.MtJm2_tr_neoc_stds = abs(Mt) * percentError[1]/100.0 + + self.QieMWm3_tr_neoc = Qe * 0.0 + self.QieMWm3_tr_neoc_stds = Qe * 0.0 return neo diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index d6993872..31eb2533 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -1,5 +1,8 @@ +import json +from matplotlib.pylab import f import torch import numpy as np +from functools import partial import copy import shutil from mitim_tools.misc_tools import IOtools, PLASMAtools @@ -7,6 +10,58 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): + ''' + For tracking and reproducibility (e.g. external runs), we want to write a json file + containing the simulation results. JSON should look like: + + { + 'r/a': ... + 'fluxes_mean': + { + 'QeMWm2': ... + 'QiMWm2': ... + 'Ge1E20m2': ... + 'GZ1E20m2': ... + 'MtJm2': ... + 'QieMWm3': ... + }, + 'fluxes_stds': + { + 'QeMWm2': ... + 'QiMWm2': ... + 'Ge1E20m2': ... + 'GZ1E20m2': ... + 'MtJm2': ... + 'QieMWm3': ... + }, + 'additional': ... + } + ''' + + with open(self.folder / file_name, 'w') as f: + + rho = self.powerstate.plasma["rho"][0, 1:].cpu().numpy() + + fluxes_mean = {} + fluxes_stds = {} + + for var in ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3']: + fluxes_mean[var] = self.__dict__[f"{var}_tr_{suffix}"].tolist() + fluxes_stds[var] = self.__dict__[f"{var}_tr_{suffix}_stds"].tolist() + + json_dict = { + 'r/a': rho.tolist(), + 'fluxes_mean': fluxes_mean, + 'fluxes_stds': fluxes_stds, + 'additional': { + } + } + + json.dump(json_dict, f, indent=4) + + print(f"\t* Written JSON with {suffix} information to {self.folder / file_name}") + class power_transport: ''' Default class for power transport models, change "evaluate" method to implement a new model and produce_profiles if the model requires written input.gacode written @@ -119,9 +174,28 @@ def _modify_profiles(self): def evaluate(self): + ''' + ****************************************************************************************************** + Evaluate neoclassical and turbulent transport. + These functions use a hook to write the .json files to communicate the results to powerstate.plasma + ****************************************************************************************************** + ''' neoclassical = self.evaluate_neoclassical() turbulence = self.evaluate_turbulence() + ''' + ****************************************************************************************************** + From the json to powerstate.plasma + ****************************************************************************************************** + ''' + self._populate_from_json(file_name = 'fluxes_turb.json', suffix= 'turb') + self._populate_from_json(file_name = 'fluxes_neoc.json', suffix= 'neoc') + + ''' + ****************************************************************************************************** + Post-process the data: add turb and neoc, tensorize and transformations + ****************************************************************************************************** + ''' self._postprocess() def _postprocess(self): @@ -174,12 +248,28 @@ def _postprocess(self): self.powerstate.plasma[f"{mapper_convective[key]}_tr{tt}"] ) * mult + def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): + + ''' + Populate the powerstate.plasma with the results from the json file + ''' + + with open(self.folder / file_name, 'r') as f: + json_dict = json.load(f) + + for var in ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3']: + self.powerstate.plasma[f"{var}_tr_{suffix}"] = np.array(json_dict['fluxes_mean'][var]) + self.powerstate.plasma[f"{var}_tr_{suffix}_stds"] = np.array(json_dict['fluxes_stds'][var]) + + print(f"\t* Populated powerstate.plasma with JSON data from {self.folder / file_name}") + # ---------------------------------------------------------------------------------------------------- # EVALUATE (custom part) # ---------------------------------------------------------------------------------------------------- + @IOtools.hook_method(after=partial(write_json, file_name = 'fluxes_turb.json', suffix= 'turb')) def evaluate_turbulence(self): ''' - This needs to populate the following np.arrays in self.powerstate.plasma, with dimensions of rho: + This needs to populate the following np.arrays in self., with dimensions of rho: - QeMWm2_tr_turb - QiMWm2_tr_turb - Ge1E20m2_tr_turb @@ -202,12 +292,13 @@ def evaluate_turbulence(self): 'QieMWm3' ]: - self.powerstate.plasma[f"{var}_tr_turb"] = np.zeros(dim) - self.powerstate.plasma[f"{var}_tr_turb_stds"] = np.zeros(dim) - + self.__dict__[f"{var}_tr_turb"] = np.zeros(dim) + self.__dict__[f"{var}_tr_turb_stds"] = np.zeros(dim) + + @IOtools.hook_method(after=partial(write_json, file_name = 'fluxes_neoc.json', suffix= 'neoc')) def evaluate_neoclassical(self): ''' - This needs to populate the following np.arrays in self.powerstate.plasma: + This needs to populate the following np.arrays in self.: - QeMWm2_tr_neoc - QiMWm2_tr_neoc - Ge1E20m2_tr_neoc @@ -228,5 +319,5 @@ def evaluate_neoclassical(self): 'MtJm2', ]: - self.powerstate.plasma[f"{var}_tr_neoc"] = np.zeros(dim) - self.powerstate.plasma[f"{var}_tr_neoc_stds"] = np.zeros(dim) \ No newline at end of file + self.__dict__[f"{var}_tr_neoc"] = np.zeros(dim) + self.__dict__[f"{var}_tr_neoc_stds"] = np.zeros(dim) \ No newline at end of file From 78cf60b99cdcf64b927457bb5f59b81db8e1d246 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 16:49:30 -0400 Subject: [PATCH 200/385] Better handling of slurm_settings when launching a gacode simulation --- .../physics_models/transport_analytic.py | 12 +- .../physics_models/transport_cgyroneo.py | 1 + src/mitim_tools/misc_tools/FARMINGtools.py | 136 ++++++++---------- tests/POWERTORCH_workflow.py | 4 +- 4 files changed, 72 insertions(+), 81 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py index 2558b3bc..f49d41d5 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -43,14 +43,14 @@ def evaluate(self): self.powerstate.plasma["a"].unsqueeze(-1), ) - self.__dict__["QeMWm2_tr_turb"] = Pe_tr * 2 / 3 - self.__dict__["QiMWm2_tr_turb"] = Pi_tr * 2 / 3 + self.QeMWm2_tr_turb = Pe_tr * 2 / 3 + self.QiMWm2_tr_turb = Pi_tr * 2 / 3 - self.__dict__["QeMWm2_tr_neoc"] = Pe_tr * 1 / 3 - self.__dict__["QiMWm2_tr_neoc"] = Pi_tr * 1 / 3 + self.QeMWm2_tr_neoc = Pe_tr * 1 / 3 + self.QiMWm2_tr_neoc = Pi_tr * 1 / 3 - self.__dict__["QeMWm2_tr"] = self.__dict__["QeMWm2_tr_turb"] + self.__dict__["QeMWm2_tr_neoc"] - self.__dict__["QiMWm2_tr"] = self.__dict__["QiMWm2_tr_turb"] + self.__dict__["QiMWm2_tr_neoc"] + self.QeMWm2_tr = self.QeMWm2_tr_turb + self.QeMWm2_tr_neoc + self.QiMWm2_tr = self.QiMWm2_tr_turb + self.QiMWm2_tr_neoc # ------------------------------------------------------------------ # SURROGATE diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index f5e421a6..a1468c69 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -11,6 +11,7 @@ class cgyroneo_model(transport_tglfneo.tglfneo_model): def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) + # Do not hook here def evaluate_turbulence(self): rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index e2c6bc10..c5392c0c 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -67,7 +67,7 @@ def define_machine( code, nameScratch, launchSlurm=True, - slurm_settings={}, + slurm_settings=None, ): # Separated in case I need to quickly grab the machine settings self.define_machine_quick(code, nameScratch, slurm_settings=slurm_settings) @@ -81,7 +81,7 @@ def define_machine( # Print Slurm info if self.launchSlurm: print("\t- Slurm Settings:") - print("\t\t- Job settings:") + print("\t\t- Job settings (different than MITIM default):") for key in self.slurm_settings: if self.slurm_settings[key] is not None: print(f"\t\t\t- {key}: {self.slurm_settings[key]}") @@ -91,20 +91,13 @@ def define_machine( 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 + def define_machine_quick(self, code, nameScratch, slurm_settings=None): - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Defaults for slurm + self.slurm_settings = slurm_settings if slurm_settings is not None else {} + + # In case there's no name, I need it 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.slurm_settings.setdefault("mem", None) - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - + self.machineSettings = CONFIGread.machineSettings( code=code, nameScratch=nameScratch, @@ -183,21 +176,14 @@ def run( command_str_mod, self.folderExecution, modules_remote=self.machineSettings["modules"], - job_array=self.slurm_settings["job_array"] if "job_array" in self.slurm_settings else None, - job_array_limit=self.slurm_settings["job_array_limit"] if "job_array_limit" in self.slurm_settings else None, folder_local=self.folder_local, shellPreCommands=self.shellPreCommands, shellPostCommands=self.shellPostCommands, - nameJob=self.slurm_settings["name"] if "name" in self.slurm_settings else "test", - minutes=self.slurm_settings["minutes"] if "minutes" in self.slurm_settings else 5, - nodes=self.slurm_settings["nodes"] if "nodes" in self.slurm_settings else None, - ntasks=self.slurm_settings["ntasks"] if "ntasks" in self.slurm_settings else 1, - cpuspertask=self.slurm_settings["cpuspertask"] if "cpuspertask" in self.slurm_settings else 4, - slurm=self.machineSettings["slurm"], - memory_req_by_job=self.slurm_settings["mem"] if "mem" in self.slurm_settings else None, - launchSlurm=self.launchSlurm, label_log_files=self.label_log_files, wait_until_sbatch=waitYN, + slurm=self.machineSettings["slurm"], + launchSlurm=self.launchSlurm, + slurm_settings=self.slurm_settings, ) # ****************************************************** @@ -988,40 +974,43 @@ def SerialProcedure(Function, Params, howmany): def create_slurm_execution_files( command, - folder_remote, + folderExecution, modules_remote=None, - slurm={}, folder_local=None, shellPreCommands=None, shellPostCommands=None, - launchSlurm=True, - nameJob="test", - minutes=5, - ntasks=1, - cpuspertask=4, - memory_req_by_job=None, - job_array=None, - job_array_limit=None, # If job_array is not None, this is the limit of the array size at once - nodes=None, label_log_files="", wait_until_sbatch=True, + slurm={}, + launchSlurm=True, + slurm_settings = None ): - if isinstance(command, str): - command = [command] - - if shellPostCommands is None: - shellPostCommands = [] - - if shellPreCommands is None: - shellPreCommands = [] - - folderExecution = folder_remote + fileSBATCH = folder_local / f"mitim_bash{label_log_files}.src" fileSHELL = folder_local / f"mitim_shell_executor{label_log_files}.sh" fileSBATCH_remote = f"{folderExecution}/mitim_bash{label_log_files}.src" - - minutes = int(minutes) - + + # --------------------------------------------------- + # slurm_settings indicate the job resource allocation + # --------------------------------------------------- + + nameJob = slurm_settings.setdefault("name", "mitim_job") + minutes = int(slurm_settings.setdefault("minutes", 10)) + memory_req_by_job = slurm_settings.setdefault("mem", None) + + nodes = slurm_settings.setdefault("nodes", None) + ntasks = slurm_settings.setdefault("ntasks", 1) + cpuspertask = slurm_settings.setdefault("cpuspertask", 1) + ntaskspernode = slurm_settings.setdefault("ntaskspernode", None) + gpuspertask = slurm_settings.setdefault("gpuspertask", None) + + job_array = slurm_settings.setdefault("job_array", None) + job_array_limit = slurm_settings.setdefault("job_array_limit", None) + + # --------------------------------------------------- + # slurm indicate the machine specifications as given by the config instead of individual job + # --------------------------------------------------- + partition = slurm.setdefault("partition", None) email = slurm.setdefault("email", None) exclude = slurm.setdefault("exclude", None) @@ -1029,7 +1018,7 @@ def create_slurm_execution_files( constraint = slurm.setdefault("constraint", None) memory_req_by_config = slurm.setdefault("mem", None) request_exclusive_node = slurm.setdefault("exclusive", False) - + if memory_req_by_job == 0 : print("\t\t- Entire node memory requested by job, overwriting memory requested by config file", typeMsg="i") memory_req = memory_req_by_job @@ -1040,6 +1029,13 @@ def create_slurm_execution_files( if memory_req_by_config is not None: print(f"\t\t- Memory requested by config file ({memory_req_by_config})", typeMsg="i") memory_req = memory_req_by_config + + if minutes >= 60: + hours = minutes // 60 + minutes = minutes - hours * 60 + time_com = f"{str(hours).zfill(2)}:{str(minutes).zfill(2)}:00" + else: + time_com = f"{str(minutes).zfill(2)}:00" """ ******************************************************************************************** @@ -1047,58 +1043,51 @@ def create_slurm_execution_files( ******************************************************************************************** """ - if minutes >= 60: - hours = minutes // 60 - minutes = minutes - hours * 60 - time_com = f"{str(hours).zfill(2)}:{str(minutes).zfill(2)}:00" - else: - time_com = f"{str(minutes).zfill(2)}:00" + command = [command] if isinstance(command, str) else command + shellPreCommands = [] if shellPreCommands is None else shellPreCommands + shellPostCommands = [] if shellPostCommands is None else shellPostCommands + # ~~~~ Construct SLURM header ~~~~~~~~~~~~~~~ commandSBATCH = [] - # ******* Basics commandSBATCH.append("#!/usr/bin/env bash") commandSBATCH.append(f"#SBATCH --job-name {nameJob}") commandSBATCH.append(f"#SBATCH --output {folderExecution}/slurm_output{label_log_files}.dat") commandSBATCH.append(f"#SBATCH --error {folderExecution}/slurm_error{label_log_files}.dat") + commandSBATCH.append(f"#SBATCH --time {time_com}") if email is not None: commandSBATCH.append("#SBATCH --mail-user=" + email) - - # ******* Partition / Billing if partition is not None: 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 not None: commandSBATCH.append(f"#SBATCH --array={job_array}{f'%{job_array_limit} ' if job_array_limit is not None else ''}") elif request_exclusive_node: commandSBATCH.append("#SBATCH --exclusive") - - # ******* CPU setup if nodes is not None: commandSBATCH.append(f"#SBATCH --nodes {nodes}") - commandSBATCH.append(f"#SBATCH --ntasks {ntasks}") - commandSBATCH.append(f"#SBATCH --cpus-per-task {cpuspertask}") - + if ntasks is not None: + commandSBATCH.append(f"#SBATCH --ntasks {ntasks}") + if ntaskspernode is not None: + commandSBATCH.append(f"#SBATCH --ntasks-per-node {ntaskspernode}") + if cpuspertask is not None: + commandSBATCH.append(f"#SBATCH --cpus-per-task {cpuspertask}") + if gpuspertask is not None: + commandSBATCH.append(f"#SBATCH --gpus-per-task {gpuspertask}") if exclude is not None: commandSBATCH.append(f"#SBATCH --exclude={exclude}") commandSBATCH.append("#SBATCH --profile=all") - + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - commandSBATCH.append("export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK") - - # ******* Commands + # ~~~~ Commands ~~~~~~~~~~~~~~~ commandSBATCH.append("") + commandSBATCH.append("export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK") commandSBATCH.append('echo "MITIM: Submitting SLURM job $SLURM_JOBID in $HOSTNAME (host: $SLURM_SUBMIT_HOST)"') commandSBATCH.append('echo "MITIM: Nodes have $SLURM_CPUS_ON_NODE cores and $SLURM_JOB_NUM_NODES node(s) were allocated for this job"') commandSBATCH.append('echo "MITIM: Each of the $SLURM_NTASKS tasks allocated will run with $SLURM_CPUS_PER_TASK cores, allocating $SRUN_CPUS_PER_TASK CPUs per srun"') @@ -1114,6 +1103,7 @@ def create_slurm_execution_files( commandSBATCH.append(c) commandSBATCH.append("") + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ wait_txt = " --wait" if wait_until_sbatch else "" if launchSlurm: @@ -1154,7 +1144,7 @@ def create_slurm_execution_files( ******************************************************************************************** """ - comm = f"cd {folder_remote} && chmod +x {fileSBATCH_remote} && chmod +x mitim_shell_executor{label_log_files}.sh && ./mitim_shell_executor{label_log_files}.sh > mitim.out" + comm = f"cd {folderExecution} && chmod +x {fileSBATCH_remote} && chmod +x mitim_shell_executor{label_log_files}.sh && ./mitim_shell_executor{label_log_files}.sh > mitim.out" return comm, fileSBATCH.resolve(), fileSHELL.resolve() diff --git a/tests/POWERTORCH_workflow.py b/tests/POWERTORCH_workflow.py index 3eab027f..2306a9a4 100644 --- a/tests/POWERTORCH_workflow.py +++ b/tests/POWERTORCH_workflow.py @@ -3,7 +3,7 @@ import numpy as np from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.powertorch import STATEtools -from mitim_modules.powertorch.physics import TRANSPORTtools +from mitim_modules.powertorch.physics_models import transport_analytic from mitim_tools import __mitimroot__ # Inputs @@ -14,7 +14,7 @@ evolution_options = { 'ProfilePredicted': ['te', 'ti'], 'rhoPredicted': rho }, - transport_options = { 'transport_evaluator': TRANSPORTtools.diffusion_model, + transport_options = { 'transport_evaluator': transport_analytic.diffusion_model, 'transport_evaluator_options': { 'chi_e': torch.ones(rho.shape[0]).to(rho)*0.8, 'chi_i': torch.ones(rho.shape[0]).to(rho)*1.2 From db4bb848e79f3cf2dc39d22940afc3762b731232 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 16:50:29 -0400 Subject: [PATCH 201/385] Each code handles its own slurm specifications now --- src/mitim_tools/gacode_tools/CGYROtools.py | 33 +- src/mitim_tools/gacode_tools/NEOtools.py | 22 + src/mitim_tools/gacode_tools/TGLFtools.py | 22 + .../gacode_tools/utils/GACODErun.py | 525 +++++++++--------- src/mitim_tools/gyrokinetics_tools/GXtools.py | 41 +- 5 files changed, 361 insertions(+), 282 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 81b4fc9a..f11fb625 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -7,7 +7,7 @@ import matplotlib.pyplot as plt from mitim_tools import __mitimroot__ from mitim_tools.gacode_tools.utils import GACODEdefaults, GACODErun, CGYROutils -from mitim_tools.misc_tools import IOtools, GRAPHICStools, FARMINGtools +from mitim_tools.misc_tools import IOtools, GRAPHICStools, FARMINGtools, CONFIGread from mitim_tools.gacode_tools.utils import GACODEplotting from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -23,10 +23,41 @@ def __init__( def code_call(folder, p, n = 1, nomp = 1, additional_command="", **kwargs): return f" cgyro -e {folder} -n {n} -nomp {nomp} {additional_command} -p {p} &\n" + def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): + + slurm_settings = { + "name": name, + "minutes": minutes, + } + + # Gather if this is a GPU enabled machine + machineSettings = CONFIGread.machineSettings(code='cgyro') + + if type_of_submission == "slurm_standard": + + slurm_settings['ntasks'] = total_cores_required + + if machineSettings['gpus_per_node'] > 0: + slurm_settings['gpuspertask'] = cores_per_code_call + else: + slurm_settings['cpuspertask'] = cores_per_code_call + + elif type_of_submission == "slurm_array": + + slurm_settings['ntasks'] = 1 + if machineSettings['gpus_per_node'] > 0: + slurm_settings['gpuspertask'] = cores_per_code_call + else: + slurm_settings['cpuspertask'] = cores_per_code_call + slurm_settings['job_array'] = ",".join(array_list) + + return slurm_settings + self.run_specifications = { 'code': 'cgyro', 'input_file': 'input.cgyro', 'code_call': code_call, + 'code_slurm_settings': code_slurm_settings, 'control_function': GACODEdefaults.addCGYROcontrol, 'controls_file': 'input.cgyro.controls', 'state_converter': 'to_cgyro', diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 4fe0a31d..cd4ff534 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -17,10 +17,32 @@ def __init__( def code_call(folder, n, p, additional_command="", **kwargs): return f" neo -e {folder} -n {n} -p {p} {additional_command} &\n" + def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): + + slurm_settings = { + "name": name, + "minutes": minutes, + 'job_array_limit': None, # Limit to this number at most running jobs at the same time? + } + + if type_of_submission == "slurm_standard": + + slurm_settings['ntasks'] = total_cores_required + slurm_settings['cpuspertask'] = cores_per_code_call + + elif type_of_submission == "slurm_array": + + slurm_settings['ntasks'] = 1 + slurm_settings['cpuspertask'] = cores_per_code_call + slurm_settings['job_array'] = ",".join(array_list) + + return slurm_settings + self.run_specifications = { 'code': 'neo', 'input_file': 'input.neo', 'code_call': code_call, + 'code_slurm_settings': code_slurm_settings, 'control_function': GACODEdefaults.addNEOcontrol, 'controls_file': 'input.neo.controls', 'state_converter': 'to_neo', diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index a44460cd..e83b4982 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -127,10 +127,32 @@ def __init__( def code_call(folder, p, n = 1, additional_command="", **kwargs): return f" tglf -e {folder} -n {n} -p {p} {additional_command} &\n" + def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): + + slurm_settings = { + "name": name, + "minutes": minutes, + 'job_array_limit': None, # Limit to this number at most running jobs at the same time? + } + + if type_of_submission == "slurm_standard": + + slurm_settings['ntasks'] = total_cores_required + slurm_settings['cpuspertask'] = cores_per_code_call + + elif type_of_submission == "slurm_array": + + slurm_settings['ntasks'] = 1 + slurm_settings['cpuspertask'] = cores_per_code_call + slurm_settings['job_array'] = ",".join(array_list) + + return slurm_settings + self.run_specifications = { 'code': 'tglf', 'input_file': 'input.tglf', 'code_call': code_call, + 'code_slurm_settings': code_slurm_settings, 'control_function': GACODEdefaults.addTGLFcontrol, 'controls_file': 'input.tglf.controls', 'state_converter': 'to_tglf', diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 8fe6d921..802f9f77 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -120,7 +120,7 @@ def run( ): if slurm_setup is None: - slurm_setup = {"cores": self.run_specifications['default_cores'], "minutes": 5} + slurm_setup = {"cores": self.run_specifications['default_cores'], "minutes": 10} # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Prepare inputs @@ -287,8 +287,8 @@ def _run_prepare( latest_inputsFileDict[irho].anticipate_problems() # Check cores problem - if launchSlurm: - self._check_cores(rhosEvaluate, slurm_setup) + # if launchSlurm: + # self._check_cores(rhosEvaluate, slurm_setup) self.FolderSimLast = Folder_sim @@ -319,21 +319,252 @@ def _run( for subfolder_simulation in code_executor: c += len(code_executor[subfolder_simulation]) - if c > 0: - run_gacode_simulation( - self.FolderGACODE, - code_executor, - run_specifications=self.run_specifications, - filesToRetrieve=filesToRetrieve, - minutes=kwargs_run.get("slurm_setup", {}).get("minutes", 5), - cores_simulation=kwargs_run.get("slurm_setup", {}).get("cores", self.run_specifications['default_cores']), - name=f"{self.run_specifications['code']}_{self.nameRunid}{kwargs_run.get('extra_name', '')}", - launchSlurm=kwargs_run.get("launchSlurm", True), - attempts_execution=kwargs_run.get("attempts_execution", 1), - full_submission=full_submission, - ) - else: + if c == 0: + print(f"\t- {self.run_specifications['code'].upper()} not run because all results files found (please ensure consistency!)",typeMsg="i") + + else: + + # ---------------------------------------------------------------------------------------------------------------- + # Run simulation + # ---------------------------------------------------------------------------------------------------------------- + """ + launchSlurm = True -> Launch as a batch job in the machine chosen + launchSlurm = False -> Launch locally as a bash script + """ + + # Get code info + code = self.run_specifications.get('code', 'tglf') + input_file = self.run_specifications.get('input_file', 'input.tglf') + code_call = self.run_specifications.get('code_call', None) + code_slurm_settings = self.run_specifications.get('code_slurm_settings', None) + + # Get execution info + minutes = kwargs_run.get("slurm_setup", {}).get("minutes", 5) + cores_per_code_call = kwargs_run.get("slurm_setup", {}).get("cores", self.run_specifications['default_cores']) + launchSlurm = kwargs_run.get("launchSlurm", True) + + extraFlag = kwargs_run.get('extra_name', '') + name = f"{self.run_specifications['code']}_{self.nameRunid}{extraFlag}" + + attempts_execution = kwargs_run.get("attempts_execution", 1) + + tmpFolder = self.FolderGACODE / f"tmp_{code}" + IOtools.askNewFolder(tmpFolder, force=True) + + gacode_job = FARMINGtools.mitim_job(tmpFolder) + + gacode_job.define_machine_quick(code,f"mitim_{name}") + + folders, folders_red = [], [] + for subfolder_sim in code_executor: + + rhos = list(code_executor[subfolder_sim].keys()) + + # --------------------------------------------- + # Prepare files and folders + # --------------------------------------------- + + for i, rho in enumerate(rhos): + print(f"\t- Preparing {code.upper()} execution ({subfolder_sim}) at rho={rho:.4f}") + + folder_sim_this = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" + folders.append(folder_sim_this) + + folder_sim_this_rel = folder_sim_this.relative_to(tmpFolder) + folders_red.append(folder_sim_this_rel.as_posix() if gacode_job.machineSettings['machine'] != 'local' else str(folder_sim_this_rel)) + + folder_sim_this.mkdir(parents=True, exist_ok=True) + + input_file_sim = folder_sim_this / input_file + with open(input_file_sim, "w") as f: + f.write(code_executor[subfolder_sim][rho]["inputs"]) + + # --------------------------------------------- + # Prepare command + # --------------------------------------------- + + # Grab machine local limits ------------------------------------------------- + max_cores_per_node = FARMINGtools.mitim_job.grab_machine_settings(code)["cores_per_node"] + + # If the run is local and not slurm, let's check the number of cores + if (FARMINGtools.mitim_job.grab_machine_settings(code)["machine"] == "local") and not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): + + cores_in_machine = int(os.cpu_count()) + cores_allocated = int(os.environ.get('SLURM_CPUS_PER_TASK')) if os.environ.get('SLURM_CPUS_PER_TASK') is not None else None + + if cores_allocated is not None: + if max_cores_per_node is None or (cores_allocated < max_cores_per_node): + print(f"\t - Detected {cores_allocated} cores allocated by SLURM, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node = cores_allocated + elif cores_in_machine is not None: + if max_cores_per_node is None or (cores_in_machine < max_cores_per_node): + print(f"\t - Detected {cores_in_machine} cores in machine, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node = cores_in_machine + else: + # Default to just 16 just in case + if max_cores_per_node is None: + max_cores_per_node = 16 + else: + # For remote execution, default to just 16 just in case + if max_cores_per_node is None: + max_cores_per_node = 16 + # --------------------------------------------------------------------------- + + # Grab the total number of cores of this job -------------------------------- + total_simulation_executions = len(rhos) * len(code_executor) + total_cores_required = int(cores_per_code_call) * total_simulation_executions + # --------------------------------------------------------------------------- + + if not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): + type_of_submission = "bash" + elif total_cores_required < max_cores_per_node: + type_of_submission = "slurm_standard" + elif total_cores_required >= max_cores_per_node: + type_of_submission = "slurm_array" + + shellPreCommands = None + shellPostCommands = None + + # Simply bash, no slurm + if type_of_submission == "bash": + + if cores_per_code_call > max_cores_per_node: + print(f"\t - Detected {cores_per_code_call} cores required, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node = cores_per_code_call + + max_parallel_execution = max_cores_per_node // cores_per_code_call # Make sure we don't overload the machine when running locally (assuming no farming trans-node) + + print(f"\t- {code.upper()} will be executed as bash script (total cores: {total_cores_required}, cores per simulation: {cores_per_code_call}). MITIM will launch {total_simulation_executions // max_parallel_execution+1} sequential executions",typeMsg="i") + + # Build the bash script with job control enabled and a loop to limit parallel jobs + GACODEcommand = "#!/usr/bin/env bash\n" + GACODEcommand += "set -m\n" # Enable job control even in non-interactive mode + GACODEcommand += f"max_parallel_execution={max_parallel_execution}\n\n" # Set the maximum number of parallel processes + + # Create a bash array of folders + GACODEcommand += "folders=(\n" + for folder in folders_red: + GACODEcommand += f' "{folder}"\n' + GACODEcommand += ")\n\n" + + # Loop over each folder and launch code, waiting if we've reached max_parallel_execution + GACODEcommand += "for folder in \"${folders[@]}\"; do\n" + GACODEcommand += code_call(folder = '\"$folder\"', n = cores_per_code_call, p = gacode_job.folderExecution) + GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" + GACODEcommand += "done\n\n" + GACODEcommand += "wait\n" + + # Standard job + elif type_of_submission == "slurm_standard": + + print(f"\t- {code.upper()} will be executed in SLURM as standard job (cpus: {total_cores_required})",typeMsg="i") + + # Code launches + GACODEcommand = "" + for folder in folders_red: + GACODEcommand += code_call(folder = folder, n = cores_per_code_call, p = gacode_job.folderExecution) + GACODEcommand += "\nwait" # This is needed so that the script doesn't end before each job + + # Job array + elif type_of_submission == "slurm_array": + + print(f"\t- {code.upper()} will be executed in SLURM as job array due to its size (cpus: {total_cores_required})",typeMsg="i") + + # As a pre-command, organize all folders in a simpler way + shellPreCommands = [] + shellPostCommands = [] + array_list = [] + for i, folder in enumerate(folders_red): + array_list.append(f"{i}") + folder_temp_array = f"run{i}" + folder_actual = folder + shellPreCommands.append(f"mkdir {gacode_job.folderExecution}/{folder_temp_array}; cp {gacode_job.folderExecution}/{folder_actual}/* {gacode_job.folderExecution}/{folder_temp_array}/.") + shellPostCommands.append(f"cp {gacode_job.folderExecution}/{folder_temp_array}/* {gacode_job.folderExecution}/{folder_actual}/.; rm -r {gacode_job.folderExecution}/{folder_temp_array}") + + # Code launches + indexed_folder = 'run"$SLURM_ARRAY_TASK_ID"' + GACODEcommand = code_call( + folder = indexed_folder, + n = cores_per_code_call, + p = gacode_job.folderExecution, + additional_command = f'1> {gacode_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {gacode_job.folderExecution}/{indexed_folder}/slurm_error.dat\n') + + # --------------------------------------------- + # Execute + # --------------------------------------------- + + slurm_settings = code_slurm_settings( + name=code, + minutes=minutes, + total_cores_required=total_cores_required, + cores_per_code_call=cores_per_code_call, + type_of_submission=type_of_submission, + array_list=array_list if type_of_submission == "slurm_array" else None + ) + + gacode_job.define_machine( + code, + f"mitim_{name}", + launchSlurm=launchSlurm, + slurm_settings=slurm_settings, + ) + + # I would like the mitim_job to check if the retrieved folders were complete + check_files_in_folder = {} + for folder in folders_red: + check_files_in_folder[folder] = filesToRetrieve + # --------------------------------------------- + + gacode_job.prep( + GACODEcommand, + input_folders=folders, + output_folders=folders_red, + check_files_in_folder=check_files_in_folder, + shellPreCommands=shellPreCommands, + shellPostCommands=shellPostCommands, + ) + + if full_submission: + + gacode_job.run( + removeScratchFolders=True, + attempts_execution=attempts_execution + ) + + # --------------------------------------------- + # Organize + # --------------------------------------------- + + print("\t- Retrieving files and changing names for storing") + fineall = True + for subfolder_sim in code_executor: + + for i, rho in enumerate(code_executor[subfolder_sim].keys()): + for file in filesToRetrieve: + original_file = f"{file}_{rho:.4f}" + final_destination = ( + code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" + ) + final_destination.unlink(missing_ok=True) + + temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" + temp_file.replace(final_destination) + + fineall = fineall and final_destination.exists() + + if not final_destination.exists(): + print(f"\t!! file {file} ({original_file}) could not be retrived",typeMsg="w",) + + if fineall: + print("\t\t- All files were successfully retrieved") + + # Remove temporary folder + shutil.rmtree(tmpFolder) + + else: + print("\t\t- Some files were not retrieved", typeMsg="w") + def run_scan( self, @@ -722,262 +953,6 @@ def cold_start_checker( return rhosEvaluate -def run_gacode_simulation( - FolderGACODE, - code_executor, - run_specifications = None, - minutes = 5, - cores_simulation = 4, - extraFlag = "", - filesToRetrieve = None, - name = "", - launchSlurm = True, - attempts_execution = 1, - max_jobs_at_once = None, - full_submission = True, -): - - """ - launchSlurm = True -> Launch as a batch job in the machine chosen - launchSlurm = False -> Launch locally as a bash script - """ - - code = run_specifications.get('code', 'tglf') - input_file = run_specifications.get('input_file', 'input.tglf') - code_call = run_specifications.get('code_call', None) - - tmpFolder = FolderGACODE / f"tmp_{code}" - IOtools.askNewFolder(tmpFolder, force=True) - - gacode_job = FARMINGtools.mitim_job(tmpFolder) - - gacode_job.define_machine_quick(code,f"mitim_{name}") - - folders, folders_red = [], [] - for subfolder_sim in code_executor: - - rhos = list(code_executor[subfolder_sim].keys()) - - # --------------------------------------------- - # Prepare files and folders - # --------------------------------------------- - - for i, rho in enumerate(rhos): - print(f"\t- Preparing {code.upper()} execution ({subfolder_sim}) at rho={rho:.4f}") - - folder_sim_this = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" - folders.append(folder_sim_this) - - folder_sim_this_rel = folder_sim_this.relative_to(tmpFolder) - folders_red.append(folder_sim_this_rel.as_posix() if gacode_job.machineSettings['machine'] != 'local' else str(folder_sim_this_rel)) - - folder_sim_this.mkdir(parents=True, exist_ok=True) - - input_file_sim = folder_sim_this / input_file - with open(input_file_sim, "w") as f: - f.write(code_executor[subfolder_sim][rho]["inputs"]) - - # --------------------------------------------- - # Prepare command - # --------------------------------------------- - - # Grab machine local limits ------------------------------------------------- - max_cores_per_node = FARMINGtools.mitim_job.grab_machine_settings(code)["cores_per_node"] - - # If the run is local and not slurm, let's check the number of cores - if (FARMINGtools.mitim_job.grab_machine_settings(code)["machine"] == "local") and not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): - - cores_in_machine = int(os.cpu_count()) - cores_allocated = int(os.environ.get('SLURM_CPUS_PER_TASK')) if os.environ.get('SLURM_CPUS_PER_TASK') is not None else None - - if cores_allocated is not None: - if max_cores_per_node is None or (cores_allocated < max_cores_per_node): - print(f"\t - Detected {cores_allocated} cores allocated by SLURM, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") - max_cores_per_node = cores_allocated - elif cores_in_machine is not None: - if max_cores_per_node is None or (cores_in_machine < max_cores_per_node): - print(f"\t - Detected {cores_in_machine} cores in machine, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") - max_cores_per_node = cores_in_machine - else: - # Default to just 16 just in case - if max_cores_per_node is None: - max_cores_per_node = 16 - else: - # For remote execution, default to just 16 just in case - if max_cores_per_node is None: - max_cores_per_node = 16 - # --------------------------------------------------------------------------- - - # Grab the total number of cores of this job -------------------------------- - total_simulation_executions = len(rhos) * len(code_executor) - total_cores_required = int(cores_simulation) * total_simulation_executions - # --------------------------------------------------------------------------- - - # Simply bash, no slurm - if not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): - - if cores_simulation > max_cores_per_node: - print(f"\t - Detected {cores_simulation} cores required, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") - max_cores_per_node = cores_simulation - - max_parallel_execution = max_cores_per_node // cores_simulation # Make sure we don't overload the machine when running locally (assuming no farming trans-node) - - print(f"\t- {code.upper()} will be executed as bash script (total cores: {total_cores_required}, cores per simulation: {cores_simulation}). MITIM will launch {total_simulation_executions // max_parallel_execution+1} sequential executions",typeMsg="i") - - # Build the bash script with job control enabled and a loop to limit parallel jobs - GACODEcommand = "#!/usr/bin/env bash\n" - GACODEcommand += "set -m\n" # Enable job control even in non-interactive mode - GACODEcommand += f"max_parallel_execution={max_parallel_execution}\n\n" # Set the maximum number of parallel processes - - # Create a bash array of folders - GACODEcommand += "folders=(\n" - for folder in folders_red: - GACODEcommand += f' "{folder}"\n' - GACODEcommand += ")\n\n" - - # Loop over each folder and launch code, waiting if we've reached max_parallel_execution - GACODEcommand += "for folder in \"${folders[@]}\"; do\n" - GACODEcommand += code_call(folder = '\"$folder\"', n = cores_simulation, p = gacode_job.folderExecution) - GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" - GACODEcommand += "done\n\n" - GACODEcommand += "wait\n" - - # Slurm setup - array_list = None - job_array_limit = None - shellPreCommands = None - shellPostCommands = None - ntasks = total_cores_required - cpuspertask = cores_simulation - - else: - - # Standard job - if total_cores_required < max_cores_per_node: - - print(f"\t- {code.upper()} will be executed in SLURM as standard job (cpus: {total_cores_required})",typeMsg="i") - - # Code launches - GACODEcommand = "" - for folder in folders_red: - GACODEcommand += code_call(folder = folder, n = cores_simulation, p = gacode_job.folderExecution) - GACODEcommand += "\nwait" # This is needed so that the script doesn't end before each job - - # Slurm setup - array_list = None - job_array_limit = None - shellPreCommands = None - shellPostCommands = None - ntasks = total_simulation_executions - cpuspertask = cores_simulation - - # Job array - else: - - print(f"\t- {code.upper()} will be executed in SLURM as job array due to its size (cpus: {total_cores_required})",typeMsg="i") - - # As a pre-command, organize all folders in a simpler way - shellPreCommands = [] - shellPostCommands = [] - array_list = [] - for i, folder in enumerate(folders_red): - array_list.append(f"{i}") - folder_temp_array = f"run{i}" - folder_actual = folder - shellPreCommands.append(f"mkdir {gacode_job.folderExecution}/{folder_temp_array}; cp {gacode_job.folderExecution}/{folder_actual}/* {gacode_job.folderExecution}/{folder_temp_array}/.") - shellPostCommands.append(f"cp {gacode_job.folderExecution}/{folder_temp_array}/* {gacode_job.folderExecution}/{folder_actual}/.; rm -r {gacode_job.folderExecution}/{folder_temp_array}") - - # Code launches - indexed_folder = 'run"$SLURM_ARRAY_TASK_ID"' - GACODEcommand = code_call( - folder = indexed_folder, - n = cores_simulation, - p = gacode_job.folderExecution, - additional_command = f'1> {gacode_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {gacode_job.folderExecution}/{indexed_folder}/slurm_error.dat\n') - - # Slurm setup - array_list = ",".join(array_list) - ntasks = 1 - cpuspertask = cores_simulation - job_array_limit = max_jobs_at_once # Limit to this number at most running jobs at the same time - - # --------------------------------------------- - # Execute - # --------------------------------------------- - - gacode_job.define_machine( - code, - f"mitim_{name}", - launchSlurm=launchSlurm, - slurm_settings={ - "minutes": minutes, - "ntasks": ntasks, - "name": name, - "cpuspertask": cpuspertask, - "job_array": array_list, - "job_array_limit": job_array_limit, - #"nodes": 1, - }, - ) - - # I would like the mitim_job to check if the retrieved folders were complete - check_files_in_folder = {} - for folder in folders_red: - check_files_in_folder[folder] = filesToRetrieve - # --------------------------------------------- - - gacode_job.prep( - GACODEcommand, - input_folders=folders, - output_folders=folders_red, - check_files_in_folder=check_files_in_folder, - shellPreCommands=shellPreCommands, - shellPostCommands=shellPostCommands, - ) - - if full_submission: - - gacode_job.run( - removeScratchFolders=True, - attempts_execution=attempts_execution - ) - - # --------------------------------------------- - # Organize - # --------------------------------------------- - - print("\t- Retrieving files and changing names for storing") - fineall = True - for subfolder_sim in code_executor: - - for i, rho in enumerate(code_executor[subfolder_sim].keys()): - for file in filesToRetrieve: - original_file = f"{file}_{rho:.4f}{extraFlag}" - final_destination = ( - code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" - ) - final_destination.unlink(missing_ok=True) - - temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" - temp_file.replace(final_destination) - - fineall = fineall and final_destination.exists() - - if not final_destination.exists(): - print(f"\t!! file {file} ({original_file}) could not be retrived",typeMsg="w",) - - if fineall: - print("\t\t- All files were successfully retrieved") - - # Remove temporary folder - shutil.rmtree(tmpFolder) - - else: - print("\t\t- Some files were not retrieved", typeMsg="w") - - - - def runTGYRO( folderWork, outputFiles=None, @@ -1879,6 +1854,8 @@ def __init__(self, file=None, controls_file=None, code='', n_species=None): self.controls_file = controls_file self.code = code self.n_species = n_species + + self.num_recorded = 100 if self.file is not None and self.file.exists(): with open(self.file, "r") as f: @@ -1913,8 +1890,6 @@ def process(self, input_dict): # Get number of recorded species if self.n_species is not None and self.n_species in input_dict: self.num_recorded = int(input_dict[self.n_species]) - else: - self.num_recorded = 100 def write_state(self, file=None): diff --git a/src/mitim_tools/gyrokinetics_tools/GXtools.py b/src/mitim_tools/gyrokinetics_tools/GXtools.py index d556f3bb..4cf7fcf6 100644 --- a/src/mitim_tools/gyrokinetics_tools/GXtools.py +++ b/src/mitim_tools/gyrokinetics_tools/GXtools.py @@ -1,6 +1,6 @@ import netCDF4 import matplotlib.pyplot as plt -from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools +from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools, CONFIGread from mitim_tools.gacode_tools.utils import GACODErun, GACODEdefaults from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __mitimroot__ @@ -16,12 +16,40 @@ def __init__( super().__init__(rhos=rhos) def code_call(folder, n, p, additional_command="", **kwargs): - return f" gx {folder}/gxplasma.in > {folder}/gxplasma.mitim.log &\n" + return f" gx -n {n} {folder}/gxplasma.in > {folder}/gxplasma.mitim.log &\n" + + def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): + + slurm_settings = { + "name": name, + "minutes": minutes, + } + + # Gather if this is a GPU enabled machine + machineSettings = CONFIGread.machineSettings(code='gx') + + if machineSettings['gpus_per_node'] == 0: + raise Exception("[MITIM] GX needs GPUs to run, but the selected machine does not have any GPU configured. Please select another machine in the config file with gpus_per_node>0.") + + if type_of_submission == "slurm_standard": + + slurm_settings['ntasks'] = total_cores_required + slurm_settings['gpuspertask'] = cores_per_code_call + slurm_settings['job_array'] = None + + elif type_of_submission == "slurm_array": + + slurm_settings['ntasks'] = 1 + slurm_settings['gpuspertask'] = cores_per_code_call + slurm_settings['job_array'] = ",".join(array_list) + + return slurm_settings self.run_specifications = { 'code': 'gx', 'input_file': 'gxplasma.in', 'code_call': code_call, + 'code_slurm_settings': code_slurm_settings, 'control_function': GACODEdefaults.addGXcontrol, 'controls_file': 'input.gx.controls', 'state_converter': 'to_gx', @@ -202,11 +230,11 @@ def write_state(self, file=None): ['rhoc', 'Rmaj', 'R_geo', 'shift', 'qinp', 'shat', 'akappa', 'akappri', 'tri', 'tripri', 'betaprim'] ], '[Dissipation]': - [ ['closure_model', 'hypercollisions', 'nu_hyper_m', 'p_hyper_m', 'nu_hyper_l', 'p_hyper_l', 'hyper', 'D_hyper', 'p_hyper'], [] ], + [ ['closure_model', 'hypercollisions', 'nu_hyper_m', 'p_hyper_m', 'nu_hyper_l', 'p_hyper_l', 'hyper', 'D_hyper', 'p_hyper', 'D_H', 'w_osc', 'p_HB', 'HB_hyper'], [] ], '[Restart]': [ ['restart', 'save_for_restart', 'nsave'], [] ], '[Diagnostics]': - [ ['nwrite', 'omega', 'fluxes', 'fields'], [] ] + [ ['nwrite', 'omega', 'fluxes', 'fields', 'moments'], [] ] } param_written = [] @@ -315,14 +343,15 @@ def _fmt_value(val): param_written.append(f"vnewk_{i+1}") f.write("[species]\n") + f.write(f" type = {typeS[:-4]} ]\n") f.write(f" z = {z[:-4]} ]\n") + f.write(f" mass = {mass[:-4]} ]\n") f.write(f" dens = {dens[:-4]} ]\n") f.write(f" temp = {temp[:-4]} ]\n") - f.write(f" mass = {mass[:-4]} ]\n") f.write(f" fprim = {fprim[:-4]} ]\n") f.write(f" tprim = {tprim[:-4]} ]\n") f.write(f" vnewk = {vnewk[:-4]} ]\n") - f.write(f" type = {typeS[:-4]} ]\n") + f.write("\n") return param_written From b93b497e4b2d1bc504cde5d4a5ca5d04c53cd335 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 16:54:24 -0400 Subject: [PATCH 202/385] Adjustments to TGYRO model to follow same process --- pyproject.toml | 2 +- src/mitim_modules/portals/PORTALSmain.py | 1 - .../physics_models/transport_tgyro.py | 20 ++++++++++++------- src/mitim_tools/__init__.py | 2 +- 4 files changed, 15 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 22a4ec74..865df759 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "MITIM" -version = "3.0.0" +version = "4.0.0" description = "MIT Integrated Modeling Suite for Fusion Applications" readme = "README.md" requires-python = ">=3.10, <3.13" # Notes: 3.9 has issues with the latest BOTORCH, 3.13 has issues with tensorflow (nn) and omfit_classesv (omfit_new) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 600ce99b..493c6038 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -195,7 +195,6 @@ def __init__( # from mitim_modules.powertorch.physics_models.transport_cgyroneo import cgyroneo_model as transport_evaluator from mitim_modules.powertorch.physics_models.transport_cgyro import cgyro_model as transport_evaluator else: - # from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model as transport_evaluator from mitim_modules.powertorch.physics_models.transport_tglfneo import tglfneo_model as transport_evaluator from mitim_modules.powertorch.physics_models.targets_analytic import analytical_model as target_evaluator diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 805f42b7..b8dbe1bb 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -15,17 +15,23 @@ def __init__(self, powerstate, **kwargs): def produce_profiles(self): self._produce_profiles() - # TGYRO model is historical (#TODO #TOREMOVE) and therefore I'm not using the same evaluate as the rest, just keep it separate - def evaluate(self): - - tgyro = self._evaluate_tglf_neo() - - self._postprocess(tgyro, "tglf_neo") - # ************************************************************************************ # Private functions for TGLF and NEO evaluations # ************************************************************************************ + # Do nothing here + def evaluate_neoclassical(self): + pass + + # Evaluate both + def evaluate_turbulence(self): + + tgyro = self._evaluate_tglf_neo() + self._postprocess(tgyro, "tglf_neo") + + TRANSPORTtools.write_json(self, file_name = 'fluxes_neoc.json', suffix= 'neoc') + TRANSPORTtools.write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb') + def _evaluate_tglf_neo(self): transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] diff --git a/src/mitim_tools/__init__.py b/src/mitim_tools/__init__.py index b4ef771a..1bdd8728 100644 --- a/src/mitim_tools/__init__.py +++ b/src/mitim_tools/__init__.py @@ -1,4 +1,4 @@ from mitim_tools.misc_tools.CONFIGread import config_manager from pathlib import Path -__version__ = "3.0.0" +__version__ = "4.0.0" __mitimroot__ = Path(__file__).resolve().parents[2] From 8daa0796ef7ba507a46d8a711896f5d88cedac63 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 17:18:49 -0400 Subject: [PATCH 203/385] Recover TGYRO way --- .../physics_models/transport_tgyro.py | 156 +++++++----------- 1 file changed, 62 insertions(+), 94 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index b8dbe1bb..0b15e484 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -27,7 +27,7 @@ def evaluate_neoclassical(self): def evaluate_turbulence(self): tgyro = self._evaluate_tglf_neo() - self._postprocess(tgyro, "tglf_neo") + self._postprocess_tgyro(tgyro, "tglf_neo") TRANSPORTtools.write_json(self, file_name = 'fluxes_neoc.json', suffix= 'neoc') TRANSPORTtools.write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb') @@ -144,7 +144,7 @@ def _evaluate_tglf_neo(self): return tgyro - def _postprocess(self, tgyro, label): + def _postprocess_tgyro(self, tgyro, label): transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] @@ -158,9 +158,9 @@ def _postprocess(self, tgyro, label): impurityPosition = self.powerstate.impurityPosition_transport # Produce right quantities (TGYRO -> powerstate.plasma) - self.powerstate = tgyro_to_powerstate( + tgyro_to_powerstate( + self, tgyro.results[label], - self.powerstate, Qi_includes_fast=Qi_includes_fast, impurityPosition=impurityPosition, UseFineGridTargets=UseFineGridTargets, @@ -899,8 +899,9 @@ def defineReferenceFluxes( # This is where the definitions for the summation variables happen for mitim and PORTALSplot # ------------------------------------------------------------------------------------------------------------------------------------------------------ -def tgyro_to_powerstate(TGYROresults, - powerstate, +def tgyro_to_powerstate( + self, + TGYROresults, forceZeroParticleFlux=False, Qi_includes_fast=False, impurityPosition=1, @@ -919,23 +920,19 @@ def tgyro_to_powerstate(TGYROresults, if UseFineGridTargets: TGYROresults.useFineGridTargets(impurityPosition=impurityPosition) - nr = powerstate.plasma['rho'].shape[-1] - if powerstate.plasma['rho'].shape[-1] != TGYROresults.rho.shape[-1]: + nr = self.powerstate.plasma['rho'].shape[-1] + if self.powerstate.plasma['rho'].shape[-1] != TGYROresults.rho.shape[-1]: print('\t- TGYRO was run with an extra point in the grid, treating it carefully now') # ********************************** # *********** Electron Energy Fluxes # ********************************** - powerstate.plasma["QeMWm2_tr_turb"] = torch.Tensor(TGYROresults.Qe_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QeMWm2_tr_neoc"] = torch.Tensor(TGYROresults.Qe_sim_neo[:, :nr]).to(powerstate.dfT) + self.QeMWm2_tr_turb = TGYROresults.Qe_sim_turb[0, 1:nr] + self.QeMWm2_tr_neoc = TGYROresults.Qe_sim_neo[0, 1:nr] - powerstate.plasma["QeMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Qe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["QeMWm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Qe_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["QeMWm2"] = torch.Tensor(TGYROresults.Qe_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QeMWm2_stds"] = torch.Tensor(TGYROresults.Qe_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + self.QeMWm2_tr_turb_stds = TGYROresults.Qe_sim_turb_stds[0, 1:nr] + self.QeMWm2_tr_neoc_stds = TGYROresults.Qe_sim_neo_stds[0, 1:nr] # ********************************** # *********** Ion Energy Fluxes @@ -943,116 +940,87 @@ def tgyro_to_powerstate(TGYROresults, if Qi_includes_fast: - powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QiMWm2_tr_neoc"] = torch.Tensor(TGYROresults.QiIons_sim_neo[:, :nr]).to(powerstate.dfT) + self.QiMWm2_tr_turb = TGYROresults.QiIons_sim_turb[0, 1:nr] + self.QiMWm2_tr_neoc = TGYROresults.QiIons_sim_neo[0, 1:nr] - powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["QiMWm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + self.QiMWm2_tr_turb_stds = TGYROresults.QiIons_sim_turb_stds[0, 1:nr] + self.QiMWm2_tr_neoc_stds = TGYROresults.QiIons_sim_neo_stds[0, 1:nr] else: - powerstate.plasma["QiMWm2_tr_turb"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QiMWm2_tr_neoc"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr[:, :nr]).to(powerstate.dfT) + self.QiMWm2_tr_turb = TGYROresults.QiIons_sim_turb_thr[0, 1:nr] + self.QiMWm2_tr_neoc = TGYROresults.QiIons_sim_neo_thr[0, 1:nr] - powerstate.plasma["QiMWm2_tr_turb_stds"] = torch.Tensor(TGYROresults.QiIons_sim_turb_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["QiMWm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.QiIons_sim_neo_thr_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["QiMWm2"] = torch.Tensor(TGYROresults.Qi_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QiMWm2_stds"] = torch.Tensor(TGYROresults.Qi_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + self.QiMWm2_tr_turb_stds = TGYROresults.QiIons_sim_turb_thr_stds[0, 1:nr] + self.QiMWm2_tr_neoc_stds = TGYROresults.QiIons_sim_neo_thr_stds[0, 1:nr] # ********************************** # *********** Momentum Fluxes # ********************************** - powerstate.plasma["MtJm2_tr_turb"] = torch.Tensor(TGYROresults.Mt_sim_turb[:, :nr]).to(powerstate.dfT) # So far, let's include fast in momentum - powerstate.plasma["MtJm2_tr_neoc"] = torch.Tensor(TGYROresults.Mt_sim_neo[:, :nr]).to(powerstate.dfT) + self.MtJm2_tr_turb = TGYROresults.Mt_sim_turb[0, 1:nr] # So far, let's include fast in momentum + self.MtJm2_tr_neoc = TGYROresults.Mt_sim_neo[0, 1:nr] - powerstate.plasma["MtJm2_tr_turb_stds"] = torch.Tensor(TGYROresults.Mt_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["MtJm2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Mt_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + self.MtJm2_tr_turb_stds = TGYROresults.Mt_sim_turb_stds[0, 1:nr] + self.MtJm2_tr_neoc_stds = TGYROresults.Mt_sim_neo_stds[0, 1:nr] - if provideTargets: - powerstate.plasma["MtJm2"] = torch.Tensor(TGYROresults.Mt_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["MtJm2_stds"] = torch.Tensor(TGYROresults.Mt_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - # ********************************** # *********** Particle Fluxes # ********************************** # Store raw fluxes for better plotting later - powerstate.plasma["Ge1E20m2_tr_turb"] = torch.Tensor(TGYROresults.Ge_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ge1E20m2_tr_neoc"] = torch.Tensor(TGYROresults.Ge_sim_neo[:, :nr]).to(powerstate.dfT) - - powerstate.plasma["Ge1E20m2_tr_turb_stds"] = torch.Tensor(TGYROresults.Ge_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Ge1E20m2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ge_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["Ge1E20m2"] = torch.Tensor(TGYROresults.Ge_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ge1E20m2_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - - powerstate.plasma["Ce_tr_turb"] = torch.Tensor(TGYROresults.Ce_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_tr_neoc"] = torch.Tensor(TGYROresults.Ce_sim_neo[:, :nr]).to(powerstate.dfT) - - powerstate.plasma["Ce_tr_turb_stds"] = torch.Tensor(TGYROresults.Ce_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["Ce_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ce_sim_neo_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["Ce"] = torch.Tensor(TGYROresults.Ce_tar[:, :nr]).to(powerstate.dfT) - powerstate.plasma["Ce_stds"] = torch.Tensor(TGYROresults.Ce_tar_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + self.Ge1E20m2_tr_turb = TGYROresults.Ge_sim_turb[0, 1:nr] + self.Ge1E20m2_tr_neoc = TGYROresults.Ge_sim_neo[0, 1:nr] + self.Ge1E20m2_tr_turb_stds = TGYROresults.Ge_sim_turb_stds[0, 1:nr] + self.Ge1E20m2_tr_neoc_stds = TGYROresults.Ge_sim_neo_stds[0, 1:nr] + # ********************************** # *********** Impurity Fluxes # ********************************** # Store raw fluxes for better plotting later - powerstate.plasma["GZ1E20m2_tr_turb"] = torch.Tensor(TGYROresults.Gi_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) - powerstate.plasma["GZ1E20m2_tr_neoc"] = torch.Tensor(TGYROresults.Gi_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) + self.GZ1E20m2_tr_turb = TGYROresults.Gi_sim_turb[impurityPosition, 0, 1:nr] + self.GZ1E20m2_tr_neoc = TGYROresults.Gi_sim_neo[impurityPosition, 0, 1:nr] - powerstate.plasma["GZ1E20m2_tr_turb_stds"] = torch.Tensor(TGYROresults.Gi_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - powerstate.plasma["GZ1E20m2_tr_neoc_stds"] = torch.Tensor(TGYROresults.Gi_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["GZ1E20m2"] = torch.Tensor(TGYROresults.Gi_tar[impurityPosition, :, :nr]).to(powerstate.dfT) - powerstate.plasma["GZ1E20m2_stds"] = torch.Tensor(TGYROresults.Gi_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None - - powerstate.plasma["CZ_tr_turb"] = torch.Tensor(TGYROresults.Ci_sim_turb[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - powerstate.plasma["CZ_tr_neoc"] = torch.Tensor(TGYROresults.Ci_sim_neo[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - - powerstate.plasma["CZ_tr_turb_stds"] = torch.Tensor(TGYROresults.Ci_sim_turb_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None - powerstate.plasma["CZ_tr_neoc_stds"] = torch.Tensor(TGYROresults.Ci_sim_neo_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None - - if provideTargets: - powerstate.plasma["CZ"] = torch.Tensor(TGYROresults.Ci_tar[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp - powerstate.plasma["CZ_stds"] = torch.Tensor(TGYROresults.Ci_tar_stds[impurityPosition, :, :nr]).to(powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None + self.GZ1E20m2_tr_turb_stds = TGYROresults.Gi_sim_turb_stds[impurityPosition, 0, 1:nr] + self.GZ1E20m2_tr_neoc_stds = TGYROresults.Gi_sim_neo_stds[impurityPosition, 0, 1:nr] # ********************************** # *********** Energy Exchange # ********************************** if provideTurbulentExchange: - powerstate.plasma["QieMWm3_tr_turb"] = torch.Tensor(TGYROresults.EXe_sim_turb[:, :nr]).to(powerstate.dfT) - powerstate.plasma["QieMWm3_tr_turb_stds"] = torch.Tensor(TGYROresults.EXe_sim_turb_stds[:, :nr]).to(powerstate.dfT) if TGYROresults.tgyro_stds else None + self.QieMWm3_tr_turb = TGYROresults.EXe_sim_turb[0, 1:nr] + self.QieMWm3_tr_turb_stds = TGYROresults.EXe_sim_turb_stds[0, 1:nr] else: - powerstate.plasma["QieMWm3_tr_turb"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 - powerstate.plasma["QieMWm3_tr_turb_stds"] = powerstate.plasma["QeMWm2_tr_turb"] * 0.0 + self.QieMWm3_tr_turb = self.QeMWm2_tr_turb * 0.0 + self.QieMWm3_tr_turb_stds = self.QeMWm2_tr_turb * 0.0 - # ********************************** - # *********** Traget extra - # ********************************** - - if forceZeroParticleFlux and provideTargets: - powerstate.plasma["Ce"] = powerstate.plasma["Ce"] * 0.0 - powerstate.plasma["Ce_stds"] = powerstate.plasma["Ce_stds"] * 0.0 - - # ------------------------------------------------------------------------------------------------------------------------ - # Sum here turbulence and neoclassical, after modifications - # ------------------------------------------------------------------------------------------------------------------------ + self.QieMWm3_tr_neoc = self.QeMWm2_tr_turb * 0.0 + self.QieMWm3_tr_neoc_stds = self.QeMWm2_tr_turb_stds * 0.0 - quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2', 'Ge1E20m2', 'GZ1E20m2'] - for ikey in quantities: - powerstate.plasma[ikey+"_tr"] = powerstate.plasma[ikey+"_tr_turb"] + powerstate.plasma[ikey+"_tr_neoc"] - - return powerstate + # ********************************** + # *********** Targets + # *********************************** + if provideTargets: + self.powerstate.plasma["QeMWm2"] = torch.Tensor(TGYROresults.Qe_tar[0, 1:nr]).to(self.powerstate.dfT) + self.powerstate.plasma["QeMWm2_stds"] = torch.Tensor(TGYROresults.Qe_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None + self.powerstate.plasma["QiMWm2"] = torch.Tensor(TGYROresults.Qi_tar[0, 1:nr]).to(self.powerstate.dfT) + self.powerstate.plasma["QiMWm2_stds"] = torch.Tensor(TGYROresults.Qi_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None + self.powerstate.plasma["MtJm2"] = torch.Tensor(TGYROresults.Mt_tar[0, 1:nr]).to(self.powerstate.dfT) + self.powerstate.plasma["MtJm2_stds"] = torch.Tensor(TGYROresults.Mt_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None + self.powerstate.plasma["Ge1E20m2"] = torch.Tensor(TGYROresults.Ge_tar[0, 1:nr]).to(self.powerstate.dfT) + self.powerstate.plasma["Ge1E20m2_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None + self.powerstate.plasma["Ce"] = torch.Tensor(TGYROresults.Ce_tar[0, 1:nr]).to(self.powerstate.dfT) + self.powerstate.plasma["Ce_stds"] = torch.Tensor(TGYROresults.Ce_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None + self.powerstate.plasma["GZ1E20m2"] = torch.Tensor(TGYROresults.Gi_tar[impurityPosition, 0, 1:nr]).to(self.powerstate.dfT) + self.powerstate.plasma["GZ1E20m2_stds"] = torch.Tensor(TGYROresults.Gi_tar_stds[impurityPosition, 0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None + self.powerstate.plasma["CZ"] = torch.Tensor(TGYROresults.Ci_tar[impurityPosition, 0, 1:nr]).to(self.powerstate.dfT) / OriginalFimp + self.powerstate.plasma["CZ_stds"] = torch.Tensor(TGYROresults.Ci_tar_stds[impurityPosition, 0, 1:nr]).to(self.powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None + + if forceZeroParticleFlux: + self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce"] * 0.0 + self.powerstate.plasma["Ce_stds"] = self.powerstate.plasma["Ce_stds"] * 0.0 From ee0d75c8b1c38c3b9ba88efe43a511fd21b2e3a8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 17:25:25 -0400 Subject: [PATCH 204/385] Recovered old cgyro way --- src/mitim_modules/powertorch/physics_models/transport_cgyro.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index c8f7ccfc..5756dbe0 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -22,7 +22,7 @@ def evaluate(self): powerstate_orig = self._trick_cgyro(tgyro) # Process results - self._postprocess(tgyro, "cgyro_neo") + self._postprocess_tgyro(tgyro, "cgyro_neo") # Some checks print("\t- Checking model modifications:") From 1f67bc827de6d34582bce88cb7f31b10dac68372 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 18:16:17 -0400 Subject: [PATCH 205/385] Development of new CGYRO model for PORTALS --- src/mitim_modules/portals/PORTALSmain.py | 3 +- .../portals/utils/PORTALSplot.py | 4 +- .../physics_models/transport_cgyro.py | 4 + .../physics_models/transport_cgyroneo.py | 128 +++++++++++++++--- .../powertorch/utils/POWERplot.py | 4 +- .../powertorch/utils/TARGETStools.py | 6 +- .../powertorch/utils/TRANSPORTtools.py | 13 +- tests/GX_workflow.py | 2 +- 8 files changed, 130 insertions(+), 34 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 493c6038..f0e56c76 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -192,8 +192,7 @@ def __init__( # Selection of model if CGYROrun: - # from mitim_modules.powertorch.physics_models.transport_cgyroneo import cgyroneo_model as transport_evaluator - from mitim_modules.powertorch.physics_models.transport_cgyro import cgyro_model as transport_evaluator + from mitim_modules.powertorch.physics_models.transport_cgyroneo import cgyroneo_model as transport_evaluator else: from mitim_modules.powertorch.physics_models.transport_tglfneo import tglfneo_model as transport_evaluator diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index d9f0ebad..a1e9a068 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -256,7 +256,7 @@ def PORTALSanalyzer_plotMetrics( if axnZ_f is not None: axnZ_f.plot(rho, power.plasma['GZ1E20m2_tr_turb'].cpu().numpy()+power.plasma['GZ1E20m2_tr_neoc'].cpu().numpy(), "-", c=col, lw=lwt, alpha=alph) - axnZ_f.plot(rho, power.plasma['CZ_raw'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) + axnZ_f.plot(rho, power.plasma['GZ1E20m2'].cpu().numpy(), "--", c=col, lw=lwt, alpha=alph) if axw0_f is not None: axw0_f.plot( @@ -3013,7 +3013,7 @@ def plotFluxComparison( Qe_tar = power.plasma['QeMWm2'].cpu().numpy()[0][ixF:] Qi_tar = power.plasma['QiMWm2'].cpu().numpy()[0][ixF:] Ge_tar = power.plasma['Ge1E20m2'].cpu().numpy()[0][ixF:] * (1-int(forceZeroParticleFlux)) - GZ_tar = power.plasma['CZ_raw'].cpu().numpy()[0][ixF:] + GZ_tar = power.plasma['GZ1E20m2'].cpu().numpy()[0][ixF:] Mt_tar = power.plasma['MtJm2'].cpu().numpy()[0][ixF:] # Plot ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 5756dbe0..8cc16c1a 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -1,3 +1,7 @@ +''' +TO BE COMPLETELY REMOVED +''' + import copy import shutil import torch diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index a1468c69..edf98c1e 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -1,7 +1,6 @@ -from mitim_tools.misc_tools import IOtools -from functools import partial +import json +import numpy as np from mitim_tools.gacode_tools import CGYROtools -from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_modules.powertorch.physics_models import transport_tglfneo from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -44,6 +43,17 @@ def evaluate_turbulence(self): self.folder, ) + # Just prepare everything but do not submit (full_submission=False) + _ = cgyro.run( + 'base_cgyro', + full_submission=automatic_run, + code_settings=1, + cold_start=cold_start, + forceIfcold_start=True, + ) + + + # Full run here if automatic_run: raise Exception("[MITIM] automatic_run not implemented yet") @@ -53,21 +63,16 @@ def evaluate_turbulence(self): # TRANSPORTtools.write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb') + # Wait until the user has placed the json file in the right folder else: - # Just prepare everything but do not submit (full_submission=False) - _ = cgyro.run( - 'base_cgyro', - full_submission=False, - code_settings=1, - cold_start=cold_start, - forceIfcold_start=True, - ) - - # Wait until the user has placed the json file in the right folder - + pre_checks(self) + + file_path = self.folder / 'fluxes_turb.json' + attempts = 0 - while (self.folder / 'fluxes_turb.json').exists() is False: + all_good = False + while (file_path.exists() is False) or (not all_good): if attempts > 0: print(f"\n\n !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!", typeMsg='i') print(f"\tMITIM could not find the file", typeMsg='i') @@ -75,12 +80,99 @@ def evaluate_turbulence(self): logic_to_wait(self.folder) attempts += 1 + if file_path.exists(): + all_good = post_checks(self) + +def pre_checks(self): + + plasma = self.powerstate.plasma + + txt = "\nFluxes to be matched by CGYRO ( TARGETS - NEO ):" + + # Print gradients + for var, varn in zip( + ["r/a ", "rho ", "a/LTe", "a/LTi", "a/Lne", "a/LnZ", "a/Lw0"], + ["roa", "rho", "aLte", "aLti", "aLne", "aLnZ", "aLw0"], + ): + txt += f"\n{var} = " + for j in range(plasma["rho"].shape[1] - 1): + txt += f"{plasma[varn][0,j+1]:.6f} " + + # Print target fluxes + for var, varn in zip( + ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ge (1E20m2/s)", "GZ (1E20m2/s)", "Mt (J/m^2) "], + ["QeMWm2", "QiMWm2", "Ge1E20m2", "GZ1E20m2", "MtJm2"], + ): + txt += f"\n{var} = " + for j in range(plasma["rho"].shape[1] - 1): + txt += f"{plasma[varn][0,j+1]-self.__dict__[f'{varn}_tr_neoc'][j]:.4e} " + + print(txt) + def logic_to_wait(folder): - print(f" **** CGYRO prepared. Please, run CGYRO from the simulation setup in folder: ", typeMsg='i') + print(f"\n**** CGYRO prepared. Please, run CGYRO from the simulation setup in folder: ", typeMsg='i') print(f"\t {folder}/base_cgyro\n", typeMsg='i') print(f" **** When finished, the fluxes_turb.json file should be placed in:", typeMsg='i') print(f"\t {folder}/fluxes_turb.json\n", typeMsg='i') - print(f" **** When you have done that, please write an 'exit' and enter (for continuing)\n", typeMsg='i') + print(f" **** When you have done that, please write 'exit' and click enter (for continuing and reading that file)\n", typeMsg='i') + + +def post_checks(self, rtol = 1e-2): + + with open(self.folder / 'fluxes_turb.json', 'r') as f: + json_dict = json.load(f) + + additional_info_from_json = json_dict.get('additional_info', {}) + + all_good = True + + if len(additional_info_from_json) == 0: + print(f"\t- No additional info found in fluxes_turb.json to be compared with", typeMsg='i') + + else: + print(f"\t- Additional info found in fluxes_turb.json:", typeMsg='i') + for k, v in additional_info_from_json.items(): + vP = self.powerstate.plasma[k].cpu().numpy()[0,1:] + print(f"\t {k} from JSON : {[round(i,4) for i in v]}", typeMsg='i') + print(f"\t {k} from POWERSTATE: {[round(i,4) for i in vP]}", typeMsg='i') + + if not np.allclose(v, vP, rtol=rtol): + all_good = print(f"{k} does not match with a relative tolerance of {rtol}:", typeMsg='q') + + return all_good + +def write_json_CGYRO(roa, fluxes_mean, fluxes_stds, additional_info, file = 'fluxes_turb.json'): + ''' + Helper to write JSON + roa must be an array: [0.25, 0.35, ...] + fluxes_mean must be a dictionary with the fields and arrays: + 'QeMWm2': [0.1, 0.2, ...], + 'QiMWm2': ..., + 'Ge1E20m2': ..., + 'GZ1E20m2': ..., + 'MtJm2': ..., + 'QieMWm3': ... + same for fluxes_stds + additional_info must be a dictionary with any additional information to include in the JSON and compare to powerstate, + for example: + 'aLte': [0.2, 0.5, ...], + 'aLti': [0.3, 0.6, ...], + 'aLne': [0.3, 0.6, ...], + 'Qgb': [0.4, 0.7, ...] + ''' - embed() + with open(file, 'w') as f: + + fluxes_mean = {} + fluxes_stds = {} + additional_info_extended = additional_info | {'roa': roa.tolist()} + + json_dict = { + 'fluxes_mean': fluxes_mean, + 'fluxes_stds': fluxes_stds, + 'additional_info': additional_info_extended + } + + json.dump(json_dict, f, indent=4) + diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index 458ac6e0..bf29adaa 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -51,9 +51,9 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co if "nZ" in self.ProfilesPredicted: # If this model provides the raw particle flux, go for it - if 'CZ_raw_tr' in self.plasma: + if 'GZ1E20m2_tr' in self.plasma: set_plots.append( - [ 'nZ', 'aLnZ', 'CZ_raw_tr', 'CZ_raw', + [ 'nZ', 'aLnZ', 'GZ1E20m2_tr', 'GZ1E20m2', 'Impurity Density','$n_Z$ ($10^{20}m^{-3}$)','$a/Ln_Z$','$\\Gamma_Z$ (GB)','$\\Gamma_Z$ ($10^{20}m^{-3}/s$)', 1E-1,"Ggb"]) else: diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index e53fd6ea..0322beec 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -119,7 +119,7 @@ def postprocessing(self, forceZeroParticleFlux=False, relative_error_assumed=1.0 self.powerstate.plasma["QeMWm2"] = self.powerstate.plasma["QeMWm2_fixedtargets"] + self.P[: self.P.shape[0]//2, :] # MW/m^2 self.powerstate.plasma["QiMWm2"] = self.powerstate.plasma["QiMWm2_fixedtargets"] + self.P[self.P.shape[0]//2 :, :] # MW/m^2 self.powerstate.plasma["Ge1E20m2"] = self.powerstate.plasma["Ge_fixedtargets"] # 1E20/s/m^2 - self.powerstate.plasma["CZ_raw"] = self.powerstate.plasma["GZ_fixedtargets"] # 1E20/s/m^2 + self.powerstate.plasma["GZ1E20m2"] = self.powerstate.plasma["GZ_fixedtargets"] # 1E20/s/m^2 self.powerstate.plasma["MtJm2"] = self.powerstate.plasma["MtJm2_fixedtargets"] # J/m^2 if forceZeroParticleFlux: @@ -127,13 +127,13 @@ def postprocessing(self, forceZeroParticleFlux=False, relative_error_assumed=1.0 # Convective fluxes self.powerstate.plasma["Ce"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["Ge1E20m2"]) # MW/m^2 - self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["CZ_raw"]) # MW/m^2 + self.powerstate.plasma["CZ"] = PLASMAtools.convective_flux(self.powerstate.plasma["te"], self.powerstate.plasma["GZ1E20m2"]) # MW/m^2 # ************************************************************************************************** # Error # ************************************************************************************************** - variables_to_error = ["QeMWm2", "QiMWm2", "Ce", "CZ", "MtJm2", "Ge1E20m2", "CZ_raw"] + variables_to_error = ["QeMWm2", "QiMWm2", "Ce", "CZ", "MtJm2", "Ge1E20m2", "GZ1E20m2"] for i in variables_to_error: self.powerstate.plasma[i + "_stds"] = abs(self.powerstate.plasma[i]) * relative_error_assumed / 100 diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 31eb2533..f8faae26 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -16,7 +16,7 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): containing the simulation results. JSON should look like: { - 'r/a': ... + 'fluxes_mean': { 'QeMWm2': ... @@ -35,14 +35,14 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): 'MtJm2': ... 'QieMWm3': ... }, - 'additional': ... + 'additional_info': { + 'rho': rho.tolist(), + } } ''' with open(self.folder / file_name, 'w') as f: - rho = self.powerstate.plasma["rho"][0, 1:].cpu().numpy() - fluxes_mean = {} fluxes_stds = {} @@ -51,10 +51,11 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): fluxes_stds[var] = self.__dict__[f"{var}_tr_{suffix}_stds"].tolist() json_dict = { - 'r/a': rho.tolist(), 'fluxes_mean': fluxes_mean, 'fluxes_stds': fluxes_stds, - 'additional': { + 'additional_info': { + 'rho': self.powerstate.plasma["rho"][0, 1:].cpu().numpy().tolist(), + 'roa': self.powerstate.plasma["roa"][0, 1:].cpu().numpy().tolist(), } } diff --git a/tests/GX_workflow.py b/tests/GX_workflow.py index 0054592e..52a3a437 100644 --- a/tests/GX_workflow.py +++ b/tests/GX_workflow.py @@ -26,7 +26,7 @@ cold_start=cold_start, code_settings=0, # Linear extraOptions={ - 't_max':10.0, # Run up to 50.0 a/c_s + 't_max':10.0, # Run up to 1 a/c_s 'y0' :5.0, # kymin = 1/y0 = 0.2 'ny': 34, # nky = 1 + (ny-1)/3 = 12 -> ky_range = 0.2 - 2.4 }, From b4d624483244dd5bac12ce4a985117adff0769e6 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 19:12:40 -0400 Subject: [PATCH 206/385] misc corrections --- .../powertorch/physics_models/transport_tglfneo.py | 4 ++-- tests/GX_workflow.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index b387b2ca..13070c55 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -67,7 +67,7 @@ def evaluate_turbulence(self): tglf.run( 'base_tglf', - Settings=TGLFsettings, + code_settings=TGLFsettings, extraOptions=extraOptions, ApplyCorrections=False, launchSlurm= launchMODELviaSlurm, @@ -187,7 +187,7 @@ def evaluate_neoclassical(self): Qe = np.array([neo.results['base']['NEOout'][i].Qe_unn for i in range(len(rho_locations))]) Qi = np.array([neo.results['base']['NEOout'][i].Qi_unn for i in range(len(rho_locations))]) Ge = np.array([neo.results['base']['NEOout'][i].Ge_unn for i in range(len(rho_locations))]) - GZ = np.array([neo.results['base']['NEOout'][i].GiAll_unn[impurityPosition] for i in range(len(rho_locations))]) + GZ = np.array([neo.results['base']['NEOout'][i].GiAll_unn[impurityPosition-1] for i in range(len(rho_locations))]) Mt = np.array([neo.results['base']['NEOout'][i].Mt_unn for i in range(len(rho_locations))]) # ------------------------------------------------------------------------------------------------------------------------ diff --git a/tests/GX_workflow.py b/tests/GX_workflow.py index 52a3a437..433dcbc2 100644 --- a/tests/GX_workflow.py +++ b/tests/GX_workflow.py @@ -26,7 +26,7 @@ cold_start=cold_start, code_settings=0, # Linear extraOptions={ - 't_max':10.0, # Run up to 1 a/c_s + 't_max':1.0, # Run up to 1 a/c_s 'y0' :5.0, # kymin = 1/y0 = 0.2 'ny': 34, # nky = 1 + (ny-1)/3 = 12 -> ky_range = 0.2 - 2.4 }, From 47d7119a065d3901b3a30b47ff355779d956cabc Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 22:49:56 -0400 Subject: [PATCH 207/385] misc fixes --- src/mitim_tools/gacode_tools/CGYROtools.py | 2 +- src/mitim_tools/gacode_tools/NEOtools.py | 2 +- src/mitim_tools/gacode_tools/TGLFtools.py | 2 +- .../gacode_tools/utils/GACODErun.py | 26 ++++++++++++------- src/mitim_tools/gyrokinetics_tools/GXtools.py | 8 +++--- src/mitim_tools/misc_tools/FARMINGtools.py | 4 +-- tests/GX_workflow.py | 10 ++++--- 7 files changed, 33 insertions(+), 21 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index f11fb625..e3cab9eb 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -21,7 +21,7 @@ def __init__( super().__init__(rhos=rhos) def code_call(folder, p, n = 1, nomp = 1, additional_command="", **kwargs): - return f" cgyro -e {folder} -n {n} -nomp {nomp} {additional_command} -p {p} &\n" + return f"cgyro -e {folder} -n {n} -nomp {nomp} {additional_command} -p {p}" def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index cd4ff534..28596862 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -15,7 +15,7 @@ def __init__( super().__init__(rhos=rhos) def code_call(folder, n, p, additional_command="", **kwargs): - return f" neo -e {folder} -n {n} -p {p} {additional_command} &\n" + return f"neo -e {folder} -n {n} -p {p} {additional_command}" def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index e83b4982..6f547415 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -125,7 +125,7 @@ def __init__( super().__init__(rhos=rhos) def code_call(folder, p, n = 1, additional_command="", **kwargs): - return f" tglf -e {folder} -n {n} -p {p} {additional_command} &\n" + return f"tglf -e {folder} -n {n} -p {p} {additional_command}" def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index 802f9f77..ae3aadc7 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -385,10 +385,12 @@ def _run( # --------------------------------------------- # Grab machine local limits ------------------------------------------------- - max_cores_per_node = FARMINGtools.mitim_job.grab_machine_settings(code)["cores_per_node"] + machineSettings = FARMINGtools.mitim_job.grab_machine_settings(code) + max_cores_per_node = machineSettings["cores_per_node"] # If the run is local and not slurm, let's check the number of cores - if (FARMINGtools.mitim_job.grab_machine_settings(code)["machine"] == "local") and not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): + if (machineSettings["machine"] == "local") and \ + not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): cores_in_machine = int(os.cpu_count()) cores_allocated = int(os.environ.get('SLURM_CPUS_PER_TASK')) if os.environ.get('SLURM_CPUS_PER_TASK') is not None else None @@ -416,21 +418,27 @@ def _run( total_cores_required = int(cores_per_code_call) * total_simulation_executions # --------------------------------------------------------------------------- + # If it's GPUS enable machine, do the comparison based on it + if machineSettings['gpus_per_node'] == 0: + max_cores_per_node_compare = max_cores_per_node + else: + print(f"\t - Detected {machineSettings['gpus_per_node']} GPUs in machine, using this value as maximum for non-arra execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node_compare = machineSettings['gpus_per_node'] + if not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): type_of_submission = "bash" - elif total_cores_required < max_cores_per_node: + elif total_cores_required < max_cores_per_node_compare: type_of_submission = "slurm_standard" - elif total_cores_required >= max_cores_per_node: + elif total_cores_required >= max_cores_per_node_compare: type_of_submission = "slurm_array" - shellPreCommands = None - shellPostCommands = None + shellPreCommands, shellPostCommands = None, None # Simply bash, no slurm if type_of_submission == "bash": if cores_per_code_call > max_cores_per_node: - print(f"\t - Detected {cores_per_code_call} cores required, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + print(f"\t- Detected {cores_per_code_call} cores required, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") max_cores_per_node = cores_per_code_call max_parallel_execution = max_cores_per_node // cores_per_code_call # Make sure we don't overload the machine when running locally (assuming no farming trans-node) @@ -450,7 +458,7 @@ def _run( # Loop over each folder and launch code, waiting if we've reached max_parallel_execution GACODEcommand += "for folder in \"${folders[@]}\"; do\n" - GACODEcommand += code_call(folder = '\"$folder\"', n = cores_per_code_call, p = gacode_job.folderExecution) + GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = gacode_job.folderExecution)}\n' GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" GACODEcommand += "done\n\n" GACODEcommand += "wait\n" @@ -463,7 +471,7 @@ def _run( # Code launches GACODEcommand = "" for folder in folders_red: - GACODEcommand += code_call(folder = folder, n = cores_per_code_call, p = gacode_job.folderExecution) + GACODEcommand += f' {code_call(folder = folder, n = cores_per_code_call, p = gacode_job.folderExecution)} &\n' GACODEcommand += "\nwait" # This is needed so that the script doesn't end before each job # Job array diff --git a/src/mitim_tools/gyrokinetics_tools/GXtools.py b/src/mitim_tools/gyrokinetics_tools/GXtools.py index 4cf7fcf6..009b1473 100644 --- a/src/mitim_tools/gyrokinetics_tools/GXtools.py +++ b/src/mitim_tools/gyrokinetics_tools/GXtools.py @@ -16,7 +16,7 @@ def __init__( super().__init__(rhos=rhos) def code_call(folder, n, p, additional_command="", **kwargs): - return f" gx -n {n} {folder}/gxplasma.in > {folder}/gxplasma.mitim.log &\n" + return f"gx -n {n} {folder}/gxplasma.in > {folder}/gxplasma.mitim.log" def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): @@ -34,13 +34,13 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call if type_of_submission == "slurm_standard": slurm_settings['ntasks'] = total_cores_required - slurm_settings['gpuspertask'] = cores_per_code_call + slurm_settings['gpuspertask'] = 1 # Because of MPI, each task needs a GPU, and I'm passing cores_per_code_call per task slurm_settings['job_array'] = None elif type_of_submission == "slurm_array": - slurm_settings['ntasks'] = 1 - slurm_settings['gpuspertask'] = cores_per_code_call + slurm_settings['ntasks'] = cores_per_code_call + slurm_settings['gpuspertask'] = 1 slurm_settings['job_array'] = ",".join(array_list) return slurm_settings diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index c5392c0c..fd6920f1 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -999,8 +999,8 @@ def create_slurm_execution_files( memory_req_by_job = slurm_settings.setdefault("mem", None) nodes = slurm_settings.setdefault("nodes", None) - ntasks = slurm_settings.setdefault("ntasks", 1) - cpuspertask = slurm_settings.setdefault("cpuspertask", 1) + ntasks = slurm_settings.setdefault("ntasks", None) + cpuspertask = slurm_settings.setdefault("cpuspertask", None) ntaskspernode = slurm_settings.setdefault("ntaskspernode", None) gpuspertask = slurm_settings.setdefault("gpuspertask", None) diff --git a/tests/GX_workflow.py b/tests/GX_workflow.py index 433dcbc2..76e66904 100644 --- a/tests/GX_workflow.py +++ b/tests/GX_workflow.py @@ -24,12 +24,16 @@ gx.run( 'gx1/', cold_start=cold_start, - code_settings=0, # Linear + code_settings=0, # Linear extraOptions={ - 't_max':1.0, # Run up to 1 a/c_s - 'y0' :5.0, # kymin = 1/y0 = 0.2 + 't_max':5.0, # Run up to 5 a/c_s (should take ~2min using 8 A100s) + 'y0' :5.0, # kymin = 1/y0 = 0.2 'ny': 34, # nky = 1 + (ny-1)/3 = 12 -> ky_range = 0.2 - 2.4 }, + slurm_setup = { + "cores": 4, # Each of the two radius with 4 GPUs each + "minutes": 10 + } ) gx.read('gx1') From 2a330eebc4043c36241f72506386b303e570eb6c Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 23:34:05 -0400 Subject: [PATCH 208/385] Recovered capability to check and fetch, now for all simulations --- .../physics_models/transport_cgyroneo.py | 36 +- src/mitim_tools/gacode_tools/CGYROtools.py | 369 ------------------ .../gacode_tools/utils/GACODErun.py | 218 ++++++----- tests/CGYRO_workflow.py | 44 ++- 4 files changed, 165 insertions(+), 502 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index edf98c1e..a63dd388 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -19,7 +19,7 @@ def evaluate_turbulence(self): cold_start = transport_evaluator_options.get("cold_start", False) - automatic_run = False + run_type = 'prep' # ------------------------------------------------------------------------------------------------------------------------ # Prepare CGYRO object @@ -42,29 +42,27 @@ def evaluate_turbulence(self): self.powerstate.profiles_transport.files[0], self.folder, ) - - # Just prepare everything but do not submit (full_submission=False) - _ = cgyro.run( - 'base_cgyro', - full_submission=automatic_run, - code_settings=1, - cold_start=cold_start, - forceIfcold_start=True, - ) - - - # Full run here - if automatic_run: - raise Exception("[MITIM] automatic_run not implemented yet") - + + if run_type in ['normal', 'submit']: + raise Exception("[MITIM] Automatic submission or full run not implemented") + # cgyro.read( # label='base_cgyro' # ) # TRANSPORTtools.write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb') - - # Wait until the user has placed the json file in the right folder - else: + + elif run_type == 'prep': + + _ = cgyro.run( + 'base_cgyro', + run_type = run_type, + code_settings=1, + cold_start=cold_start, + forceIfcold_start=True, + ) + + # Wait until the user has placed the json file in the right folder pre_checks(self) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index e3cab9eb..0e8bee1c 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -1630,375 +1630,6 @@ def _kyfy(self,labels_original): return labels - - # def prep(self, folder, inputgacode_file): - - # # Prepare main folder with input.gacode - # self.folder = IOtools.expandPath(folder) - - # self.folder.mkdir(parents=True, exist_ok=True) - - # self.inputgacode_file = self.folder / "input.gacode" - # if IOtools.expandPath(inputgacode_file) != self.inputgacode_file: - # shutil.copy2(IOtools.expandPath(inputgacode_file), self.inputgacode_file) - - def _prerun( - self, - subfolder, - roa=0.55, - CGYROsettings=None, - extraOptions={}, - multipliers={}, - ): - - self.folder = self.FolderGACODE / Path(subfolder) - - self.folder.mkdir(parents=True, exist_ok=True) - - input_cgyro_file = self.folder / "input.cgyro" - inputCGYRO = CGYROinput(file=input_cgyro_file) - - inputgacode_file_this = self.folder / "input.gacode" - self.profiles.write_state(inputgacode_file_this) - - ResultsFiles_new = [] - for i in self.ResultsFiles: - if "mitim.out" not in i: - ResultsFiles_new.append(i) - self.ResultsFiles = ResultsFiles_new - - inputCGYRO = GACODErun.modifyInputs( - inputCGYRO, - code_settings=CGYROsettings, - extraOptions=extraOptions, - multipliers=multipliers, - addControlFunction=GACODEdefaults.addCGYROcontrol, - rmin=roa, - controls_file = 'input.cgyro.controls' - ) - - inputCGYRO.write_state() - - return input_cgyro_file, inputgacode_file_this - - - def run_test( - self, - subfolder, - roa=0.55, - CGYROsettings=None, - extraOptions={}, - multipliers={}, - **kwargs - ): - - if 'scan_param' in kwargs: - print("\t- Cannot run CGYRO tests with scan_param, running just the base",typeMsg="i") - - input_cgyro_file, inputgacode_file_this = self._prerun( - subfolder, - roa=roa, - CGYROsettings=CGYROsettings, - extraOptions=extraOptions, - multipliers=multipliers, - ) - - self.cgyro_job = FARMINGtools.mitim_job(self.folder) - - name = f'mitim_cgyro_{subfolder}_{roa:.6f}_test' - - self.cgyro_job.define_machine( - "cgyro", - name, - slurm_settings={ - "name": name, - "minutes": 5, - "cpuspertask": 1, - "ntasks": 1, - }, - ) - - CGYROcommand = "cgyro -t ." - - self.cgyro_job.prep( - CGYROcommand, - input_files=[input_cgyro_file, inputgacode_file_this], - output_files=self.output_files_test, - ) - - self.cgyro_job.run() - - def run1(self,subfolder,test_run=False,**kwargs): - - if test_run: - self.run_test(subfolder,**kwargs) - else: - self.run_full(subfolder,**kwargs) - - def run_full( - self, - subfolder, - roa=0.55, - CGYROsettings=None, - extraOptions={}, - multipliers={}, - scan_param = None, # {'variable': 'KY', 'values': [0.2,0.3,0.4]} - enforce_equality = None, # e.g. {'DLNTDR_SCALE_2': 'DLNTDR_SCALE_1', 'DLNTDR_SCALE_3': 'DLNTDR_SCALE_1'} - minutes = 5, - n = 16, - nomp = 1, - cpuspertask=None, # if None, will default to 1 - queue=None, #if blank will default to the one in settings - mem=None, # in MB - submit_via_qsub=True, #TODO fix this, works only at NERSC? no scans? - clean_folder_going_in=True, # Make sure the scratch folder is removed before running (unless I want a restart!) - submit_run=True, # False if I just want to check and fetch the job that was already submitted (e.g. via qsub or slurm) - ): - - input_cgyro_file, inputgacode_file_this = self._prerun( - subfolder, - roa=roa, - CGYROsettings=CGYROsettings, - extraOptions=extraOptions, - multipliers=multipliers, - ) - - self.cgyro_job = FARMINGtools.mitim_job(self.folder) - - name = f'mitim_cgyro_{subfolder}_{roa:.6f}' - - if scan_param is not None and submit_via_qsub: - raise Exception(" Cannot use scan_param with submit_via_qsub=True, because it requires a different job for each value of the scan parameter.") - - if submit_via_qsub: - - self.cgyro_job.define_machine( - "cgyro", - name, - launchSlurm=False, - ) - - subfolder = "scan0" - queue = "-queue " + self.cgyro_job.machineSettings['slurm']['partition'] if "partition" in self.cgyro_job.machineSettings['slurm'] else "" - - if mem is not None: - CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} -mem {mem} {queue} -w 0:{minutes}:00 -s' - else: - CGYROcommand = f'gacode_qsub -e {subfolder} -n {n} -nomp {nomp} {queue} -w 0:{minutes}:00 -s' - - if "account" in self.cgyro_job.machineSettings["slurm"] and self.cgyro_job.machineSettings["slurm"]["account"] is not None: - CGYROcommand += f" -repo {self.cgyro_job.machineSettings['slurm']['account']}" - - self.slurm_output = "batch.out" - - # --- - folder_run = self.folder / subfolder - folder_run.mkdir(parents=True, exist_ok=True) - - # Copy the input.cgyro in the subfolder - input_cgyro_file_this = folder_run / "input.cgyro" - shutil.copy2(input_cgyro_file, input_cgyro_file_this) - - # Copy the input.gacode file in the subfolder - inputgacode_file_this = folder_run / "input.gacode" - shutil.copy2(self.inputgacode_file, inputgacode_file_this) - - # Prepare the input and output folders - input_folders = [folder_run] - output_folders = [subfolder] - - else: - - if scan_param is None: - job_array = None - folder = 'scan0' - scan_param = {'variable': None, 'values': [0]} # Dummy scan parameter to avoid issues with the code below - else: - # Array - job_array = '' - for i,value in enumerate(scan_param['values']): - if job_array != '': - job_array += ',' - job_array += str(i) - - folder = 'scan"$SLURM_ARRAY_TASK_ID"' - - # Machine - self.cgyro_job.define_machine( - "cgyro", - name, - slurm_settings={ - "name": name, - "minutes": minutes, - "ntasks": n, - "job_array": job_array, - # Validate n and nomp before assigning cpuspertask - }, - ) - - if cpuspertask is not None: - if not isinstance(cpuspertask, int): - raise TypeError(" cpuspertask must be an integer") - self.cgyro_job.slurm_settings["cpuspertask"] = cpuspertask - - - if queue is not None: - self.cgyro_job.machineSettings['slurm']['partition'] = queue - - if mem is not None: - self.cgyro_job.slurm_settings['mem'] = mem - - # if not self.cgyro_job.launchSlurm: - # raise Exception(" Cannot run CGYRO scans without slurm") - - # Command to run cgyro - CGYROcommand = f'cgyro -e {folder} -n {n} -nomp {nomp} -p {self.cgyro_job.folderExecution}' - - # Scans - input_folders = [] - output_folders = [] - for i,value in enumerate(scan_param['values']): - subfolder = f"scan{i}" - folder_run = self.folder / subfolder - folder_run.mkdir(parents=True, exist_ok=True) - - # Copy the input.cgyro in the subfolder - input_cgyro_file_this = folder_run / "input.cgyro" - shutil.copy2(input_cgyro_file, input_cgyro_file_this) - - # Modify the input.cgyro file with the scan parameter - extraOptions_this = extraOptions.copy() - if scan_param['variable'] is not None: - extraOptions_this[scan_param['variable']] = value - - - # If there is an enforce_equality, apply it - if enforce_equality is not None: - for key in enforce_equality: - extraOptions_this[key] = extraOptions_this[enforce_equality[key]] - - inputCGYRO = CGYROinput(file=input_cgyro_file_this) - input_cgyro_file_this = GACODErun.modifyInputs( - inputCGYRO, - code_settings=CGYROsettings, - extraOptions=extraOptions_this, - multipliers=multipliers, - addControlFunction=GACODEdefaults.addCGYROcontrol, - rmin=roa, - control_file = 'input.cgyro.controls' - ) - - input_cgyro_file_this.write_state() - - # Copy the input.gacode file in the subfolder - inputgacode_file_this = folder_run / "input.gacode" - shutil.copy2(self.inputgacode_file, inputgacode_file_this) - - # Prepare the input and output folders - input_folders.append(folder_run) - output_folders.append(subfolder) - - self.slurm_output = "slurm_output.dat" - - # First submit the job with gacode_qsub, which will submit the cgyro job via slurm, with name - self.cgyro_job.prep( - CGYROcommand, - input_folders = input_folders, - output_folders=output_folders, - ) - - if submit_run: - self.cgyro_job.run( - waitYN=False, - check_if_files_received=False, - removeScratchFolders=False, - removeScratchFolders_goingIn=clean_folder_going_in, - ) - - # Prepare how to search for the job without waiting for it - name_default_submission_qsub = Path(self.cgyro_job.folderExecution).name - - self.cgyro_job.launchSlurm = True - self.cgyro_job.slurm_settings['name'] = name_default_submission_qsub - - - def check(self, every_n_minutes=5): - - if self.cgyro_job.launchSlurm: - print("- Checker job status") - - while True: - self.cgyro_job.check(file_output = self.slurm_output) - print(f'\t- Current status (as of {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}): {self.cgyro_job.status} ({self.cgyro_job.infoSLURM["STATE"]})') - if self.cgyro_job.status == 2: - break - else: - print(f"\t- Waiting {every_n_minutes} minutes") - time.sleep(every_n_minutes * 60) - else: - print("- Not checking status because this was run command line (not slurm)") - - print("\n\t* Job considered finished",typeMsg="i") - - def fetch(self): - """ - For a job that has been submitted but not waited for, once it is done, get the results - """ - - print("\n\n\t- Fetching results") - - if self.cgyro_job.launchSlurm: - self.cgyro_job.connect() - self.cgyro_job.retrieve() - self.cgyro_job.close() - else: - print("- Not retrieving results because this was run command line (not slurm)") - - def delete(self): - - print("\n\n\t- Deleting job") - - self.cgyro_job.launchSlurm = False - - self.cgyro_job.prep( - f"scancel -n {self.cgyro_job.slurm_settings['name']}", - label_log_files="_finish", - ) - - self.cgyro_job.run() - - def read1(self, label="cgyro1", folder=None, tmin = 0.0, minimal = False, last_tmin_for_linear = True): - - folder = IOtools.expandPath(folder) if folder is not None else self.folder - - folders = sorted(list((folder).glob("scan*")), key=lambda f: int(''.join(filter(str.isdigit, f.name)))) - - if len(folders) == 0: - folders = [folder] - attach_name = False - else: - print(f"\t- Found {len(folders)} scan folders in {folder.resolve()}:") - for f in folders: - print(f"\t\t- {f.name}") - attach_name = True - - data = {} - labels = [] - for folder in folders: - - if attach_name: - label1 = f"{label}_{folder.name}" - else: - label1 = label - - data[label1] = CGYROutils.CGYROout(folder, tmin=tmin, minimal=minimal, last_tmin_for_linear=last_tmin_for_linear) - labels.append(label1) - - self.results.update(data) - - if attach_name: - self.results[label] = CGYROutils.CGYROlinear_scan(labels, data) - class CGYROinput(GACODErun.GACODEinput): def __init__(self, file=None): super().__init__( diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index ae3aadc7..c8328900 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -1,4 +1,6 @@ import shutil +import datetime +import time import os import copy import numpy as np @@ -116,7 +118,7 @@ def run( slurm_setup=None, # Cores per call (so, when running nR radii -> nR*4) attempts_execution=1, only_minimal_files=False, - full_submission=True # Flag to submit the job or just prepare everything but do not submit via MITIM + run_type = 'normal', # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit ): if slurm_setup is None: @@ -165,7 +167,7 @@ def run( slurm_setup=slurm_setup, only_minimal_files=only_minimal_files, attempts_execution=attempts_execution, - full_submission=full_submission + run_type=run_type ) return code_executor_full @@ -303,7 +305,7 @@ def _check_cores(self, rhosEvaluate, slurm_setup, warning = 32 * 2): def _run( self, code_executor, - full_submission=True, + run_type = 'normal', # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit **kwargs_run ): """ @@ -352,9 +354,9 @@ def _run( tmpFolder = self.FolderGACODE / f"tmp_{code}" IOtools.askNewFolder(tmpFolder, force=True) - gacode_job = FARMINGtools.mitim_job(tmpFolder) + self.simulation_job = FARMINGtools.mitim_job(tmpFolder) - gacode_job.define_machine_quick(code,f"mitim_{name}") + self.simulation_job.define_machine_quick(code,f"mitim_{name}") folders, folders_red = [], [] for subfolder_sim in code_executor: @@ -372,7 +374,7 @@ def _run( folders.append(folder_sim_this) folder_sim_this_rel = folder_sim_this.relative_to(tmpFolder) - folders_red.append(folder_sim_this_rel.as_posix() if gacode_job.machineSettings['machine'] != 'local' else str(folder_sim_this_rel)) + folders_red.append(folder_sim_this_rel.as_posix() if self.simulation_job.machineSettings['machine'] != 'local' else str(folder_sim_this_rel)) folder_sim_this.mkdir(parents=True, exist_ok=True) @@ -390,7 +392,7 @@ def _run( # If the run is local and not slurm, let's check the number of cores if (machineSettings["machine"] == "local") and \ - not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): + not (launchSlurm and ("partition" in self.simulation_job.machineSettings["slurm"])): cores_in_machine = int(os.cpu_count()) cores_allocated = int(os.environ.get('SLURM_CPUS_PER_TASK')) if os.environ.get('SLURM_CPUS_PER_TASK') is not None else None @@ -425,7 +427,7 @@ def _run( print(f"\t - Detected {machineSettings['gpus_per_node']} GPUs in machine, using this value as maximum for non-arra execution (vs {max_cores_per_node} specified)",typeMsg="i") max_cores_per_node_compare = machineSettings['gpus_per_node'] - if not (launchSlurm and ("partition" in gacode_job.machineSettings["slurm"])): + if not (launchSlurm and ("partition" in self.simulation_job.machineSettings["slurm"])): type_of_submission = "bash" elif total_cores_required < max_cores_per_node_compare: type_of_submission = "slurm_standard" @@ -458,7 +460,7 @@ def _run( # Loop over each folder and launch code, waiting if we've reached max_parallel_execution GACODEcommand += "for folder in \"${folders[@]}\"; do\n" - GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = gacode_job.folderExecution)}\n' + GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = self.simulation_job.folderExecution)}\n' GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" GACODEcommand += "done\n\n" GACODEcommand += "wait\n" @@ -471,7 +473,7 @@ def _run( # Code launches GACODEcommand = "" for folder in folders_red: - GACODEcommand += f' {code_call(folder = folder, n = cores_per_code_call, p = gacode_job.folderExecution)} &\n' + GACODEcommand += f' {code_call(folder = folder, n = cores_per_code_call, p = self.simulation_job.folderExecution)} &\n' GACODEcommand += "\nwait" # This is needed so that the script doesn't end before each job # Job array @@ -487,16 +489,16 @@ def _run( array_list.append(f"{i}") folder_temp_array = f"run{i}" folder_actual = folder - shellPreCommands.append(f"mkdir {gacode_job.folderExecution}/{folder_temp_array}; cp {gacode_job.folderExecution}/{folder_actual}/* {gacode_job.folderExecution}/{folder_temp_array}/.") - shellPostCommands.append(f"cp {gacode_job.folderExecution}/{folder_temp_array}/* {gacode_job.folderExecution}/{folder_actual}/.; rm -r {gacode_job.folderExecution}/{folder_temp_array}") + shellPreCommands.append(f"mkdir {self.simulation_job.folderExecution}/{folder_temp_array}; cp {self.simulation_job.folderExecution}/{folder_actual}/* {self.simulation_job.folderExecution}/{folder_temp_array}/.") + shellPostCommands.append(f"cp {self.simulation_job.folderExecution}/{folder_temp_array}/* {self.simulation_job.folderExecution}/{folder_actual}/.; rm -r {self.simulation_job.folderExecution}/{folder_temp_array}") # Code launches indexed_folder = 'run"$SLURM_ARRAY_TASK_ID"' GACODEcommand = code_call( folder = indexed_folder, n = cores_per_code_call, - p = gacode_job.folderExecution, - additional_command = f'1> {gacode_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {gacode_job.folderExecution}/{indexed_folder}/slurm_error.dat\n') + p = self.simulation_job.folderExecution, + additional_command = f'1> {self.simulation_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {self.simulation_job.folderExecution}/{indexed_folder}/slurm_error.dat\n') # --------------------------------------------- # Execute @@ -511,7 +513,7 @@ def _run( array_list=array_list if type_of_submission == "slurm_array" else None ) - gacode_job.define_machine( + self.simulation_job.define_machine( code, f"mitim_{name}", launchSlurm=launchSlurm, @@ -524,7 +526,7 @@ def _run( check_files_in_folder[folder] = filesToRetrieve # --------------------------------------------- - gacode_job.prep( + self.simulation_job.prep( GACODEcommand, input_folders=folders, output_folders=folders_red, @@ -533,45 +535,121 @@ def _run( shellPostCommands=shellPostCommands, ) - if full_submission: - - gacode_job.run( + # Submit run and wait + if run_type == 'normal': + + self.simulation_job.run( removeScratchFolders=True, attempts_execution=attempts_execution ) + + self._organize_results(code_executor, tmpFolder, filesToRetrieve) + + # Submit run but do not wait; the user should do checks and fetch results + elif run_type == 'submit': + self.simulation_job.run( + waitYN=False, + check_if_files_received=False, + removeScratchFolders=False, + removeScratchFolders_goingIn=kwargs_run.get("cold_start", False), + ) - # --------------------------------------------- - # Organize - # --------------------------------------------- + self.kwargs_organize = { + "code_executor": code_executor, + "tmpFolder": tmpFolder, + "filesToRetrieve": filesToRetrieve + } - print("\t- Retrieving files and changing names for storing") - fineall = True - for subfolder_sim in code_executor: + self.slurm_output = "slurm_output.dat" - for i, rho in enumerate(code_executor[subfolder_sim].keys()): - for file in filesToRetrieve: - original_file = f"{file}_{rho:.4f}" - final_destination = ( - code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" - ) - final_destination.unlink(missing_ok=True) + # Prepare how to search for the job without waiting for it + self.simulation_job.launchSlurm = True + self.simulation_job.slurm_settings['name'] = Path(self.simulation_job.folderExecution).name - temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" - temp_file.replace(final_destination) + def check(self, every_n_minutes=None): - fineall = fineall and final_destination.exists() + if self.simulation_job.launchSlurm: + print("- Checker job status") - if not final_destination.exists(): - print(f"\t!! file {file} ({original_file}) could not be retrived",typeMsg="w",) + while True: + self.simulation_job.check(file_output = self.slurm_output) + print(f'\t- Current status (as of {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}): {self.simulation_job.status} ({self.simulation_job.infoSLURM["STATE"]})') + if self.simulation_job.status == 2: + print("\n\t* Job considered finished (please do .fetch() to retrieve results)",typeMsg="i") + break + elif every_n_minutes is None: + print("\n\t* Job not finished yet") + break + else: + print(f"\n\t* Waiting {every_n_minutes} minutes") + time.sleep(every_n_minutes * 60) + else: + print("- Not checking status because this was run command line (not slurm)") - if fineall: - print("\t\t- All files were successfully retrieved") + def fetch(self): + """ + For a job that has been submitted but not waited for, once it is done, get the results + """ - # Remove temporary folder - shutil.rmtree(tmpFolder) + print("\n\n\t- Fetching results") - else: - print("\t\t- Some files were not retrieved", typeMsg="w") + if self.simulation_job.launchSlurm: + self.simulation_job.connect() + self.simulation_job.retrieve() + self.simulation_job.close() + + self._organize_results(**self.kwargs_organize) + + else: + print("- Not retrieving results because this was run command line (not slurm)") + + def delete(self): + + print("\n\n\t- Deleting job") + + self.simulation_job.launchSlurm = False + + self.simulation_job.prep( + f"scancel -n {self.simulation_job.slurm_settings['name']}", + label_log_files="_finish", + ) + + self.simulation_job.run() + + def _organize_results(self, code_executor, tmpFolder, filesToRetrieve): + + # --------------------------------------------- + # Organize + # --------------------------------------------- + + print("\t- Retrieving files and changing names for storing") + fineall = True + for subfolder_sim in code_executor: + + for i, rho in enumerate(code_executor[subfolder_sim].keys()): + for file in filesToRetrieve: + original_file = f"{file}_{rho:.4f}" + final_destination = ( + code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" + ) + final_destination.unlink(missing_ok=True) + + temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" + temp_file.replace(final_destination) + + fineall = fineall and final_destination.exists() + + if not final_destination.exists(): + print(f"\t!! file {file} ({original_file}) could not be retrived",typeMsg="w",) + + if fineall: + print("\t\t- All files were successfully retrieved") + + # Remove temporary folder + shutil.rmtree(tmpFolder) + + else: + print("\t\t- Some files were not retrieved", typeMsg="w") def run_scan( @@ -1280,60 +1358,6 @@ def CDFtoTRXPLoutput( grids=grids, ) - -def executeCGYRO( - FolderCGYRO, - linesCGYRO, - fileProfiles, - outputFiles=["out.cgyro.run"], - name="", - numcores=32, -): - FolderCGYRO.mkdir(parents=True, exist_ok=True) - - 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 - # --------------- - - fileCGYRO = FolderCGYRO / f"input.cgyro" - with open(fileCGYRO, "w") as f: - f.write("\n".join(linesCGYRO)) - - # --------------- - # Execution command - # --------------- - - folderExecution = cgyro_job.machineSettings["folderWork"] - CGYROcommand = f"cgyro -e . -n {numcores} -p {folderExecution}" - - shellPreCommands = [] - - # --------------- - # Execute - # --------------- - - cgyro_job.prep( - CGYROcommand, - input_files=[fileProfiles, fileCGYRO], - output_files=outputFiles, - shellPreCommands=shellPreCommands, - ) - - cgyro_job.run() - - def runTRXPL( FolderTRXPL, timeRun, diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index d13e8e9b..c96be6d0 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -20,6 +20,8 @@ # Standalone run # --------------- +run_type = 'submit' # 'normal': submit and wait; 'submit': Just prepare and submit, do not wait [requies cgyro.check() and cgyro.fetch()] + cgyro.run( 'linear', code_settings=0, @@ -28,11 +30,18 @@ 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file }, slurm_setup={ - 'cores':8 + 'cores':8, + 'minutes': 10, }, cold_start=cold_start, forceIfcold_start=True, + run_type=run_type, ) + +if run_type == 'submit': + cgyro.check(every_n_minutes=1) + cgyro.fetch() + cgyro.read(label="cgyro1") cgyro.plot(labels=["cgyro1"]) @@ -40,23 +49,24 @@ # Scan of KY # --------------- -cgyro.run_scan( - 'scan1', - extraOptions={ - 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file - }, - variable='KY', - varUpDown=[0.3,0.4], - slurm_setup={ - 'cores':4 - }, - cold_start=cold_start, - forceIfcold_start=True, - ) +# cgyro.run_scan( +# 'scan1', +# extraOptions={ +# 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file +# }, +# variable='KY', +# varUpDown=[0.3,0.4], +# slurm_setup={ +# 'cores':4 +# }, +# cold_start=cold_start, +# forceIfcold_start=True, +# run_type='normal' +# ) -cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) +# cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) -fig = cgyro.fn.add_figure(label="Quick linear") -cgyro.plot_quick_linear(labels=["scan1_KY_0.3","scan1_KY_0.4"], fig = fig) +# fig = cgyro.fn.add_figure(label="Quick linear") +# cgyro.plot_quick_linear(labels=["scan1_KY_0.3","scan1_KY_0.4"], fig = fig) cgyro.fn.show() From 714d5e91ad78f20539315334ede8b6dbd71ab60f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 23:49:55 -0400 Subject: [PATCH 209/385] Naming changing... moving to simulation_tools --- src/mitim_tools/gacode_tools/CGYROtools.py | 13 +- src/mitim_tools/gacode_tools/NEOtools.py | 7 +- src/mitim_tools/gacode_tools/TGLFtools.py | 9 +- .../gacode_tools/utils/CGYROutils.py | 4 +- .../gacode_tools/utils/GACODErun.py | 1293 ----------------- src/mitim_tools/misc_tools/FARMINGtools.py | 2 +- src/mitim_tools/simulation_tools/SIMtools.py | 1285 ++++++++++++++++ .../physics}/GXtools.py | 7 +- tests/CGYRO_workflow.py | 40 +- tests/GX_workflow.py | 2 +- 10 files changed, 1328 insertions(+), 1334 deletions(-) create mode 100644 src/mitim_tools/simulation_tools/SIMtools.py rename src/mitim_tools/{gyrokinetics_tools => simulation_tools/physics}/GXtools.py (98%) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 0e8bee1c..998b72df 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -1,18 +1,15 @@ -import shutil -import datetime -import time import copy -from pathlib import Path import numpy as np import matplotlib.pyplot as plt from mitim_tools import __mitimroot__ -from mitim_tools.gacode_tools.utils import GACODEdefaults, GACODErun, CGYROutils -from mitim_tools.misc_tools import IOtools, GRAPHICStools, FARMINGtools, CONFIGread +from mitim_tools.gacode_tools.utils import GACODEdefaults, CGYROutils +from mitim_tools.simulation_tools import SIMtools +from mitim_tools.misc_tools import GRAPHICStools, CONFIGread from mitim_tools.gacode_tools.utils import GACODEplotting from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -class CGYRO(GACODErun.gacode_simulation): +class CGYRO(SIMtools.mitim_simulation): def __init__( self, rhos=[0.4, 0.6], # rho locations of interest @@ -1630,7 +1627,7 @@ def _kyfy(self,labels_original): return labels -class CGYROinput(GACODErun.GACODEinput): +class CGYROinput(SIMtools.GACODEinput): def __init__(self, file=None): super().__init__( file=file, diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 28596862..c32b56b2 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -2,11 +2,12 @@ import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools from mitim_tools.gacode_tools.utils import GACODErun, GACODEdefaults +from mitim_tools.simulation_tools import SIMtools from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __mitimroot__ from IPython import embed -class NEO(GACODErun.gacode_simulation): +class NEO(SIMtools.mitim_simulation): def __init__( self, rhos=[0.4, 0.6], # rho locations of interest @@ -262,7 +263,7 @@ def check_if_files_exist(folder, list_files): return True -class NEOinput(GACODErun.GACODEinput): +class NEOinput(SIMtools.GACODEinput): def __init__(self, file=None): super().__init__( file=file, @@ -271,7 +272,7 @@ def __init__(self, file=None): n_species='N_SPECIES' ) -class NEOoutput(GACODErun.GACODEoutput): +class NEOoutput(SIMtools.GACODEoutput): def __init__(self, FolderGACODE, suffix="", **kwargs): super().__init__() diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 6f547415..4d72c604 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -21,6 +21,7 @@ GACODEplotting, GACODErun, ) +from mitim_tools.simulation_tools import SIMtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -28,7 +29,7 @@ MAX_TGLF_SPECIES = 6 -class TGLF(GACODErun.gacode_simulation): +class TGLF(SIMtools.mitim_simulation): def __init__( self, rhos=[0.4, 0.6], # rho locations of interest @@ -3285,7 +3286,7 @@ def reduceToControls(dict_all): return controls, plasma -class TGLFinput(GACODErun.GACODEinput): +class TGLFinput(SIMtools.GACODEinput): def __init__(self, file=None): super().__init__( file=file, @@ -3861,7 +3862,7 @@ def readTGLFresults( TGLFstd_TGLFout.append(TGLFout) inputclasses.append(TGLFout.inputclass) - parse = GACODErun.buildDictFromInput(TGLFout.inputFile) + parse = SIMtools.buildDictFromInput(TGLFout.inputFile) parsed.append(parse) results = { @@ -3874,7 +3875,7 @@ def readTGLFresults( return results -class TGLFoutput(GACODErun.GACODEoutput): +class TGLFoutput(SIMtools.GACODEoutput): def __init__(self, FolderGACODE, suffix="", require_all_files=True): super().__init__() diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 211d6efd..ca480e39 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -5,7 +5,7 @@ import statsmodels.api as sm import matplotlib.pyplot as plt from mitim_tools.misc_tools import IOtools -from mitim_tools.gacode_tools.utils import GACODErun +from mitim_tools.simulation_tools import SIMtools from mitim_tools.misc_tools.LOGtools import printMsg as print from pygacode.cgyro.data_plot import cgyrodata_plot from pygacode import gacodefuncs @@ -50,7 +50,7 @@ def __init__(self, labels, cgyro_data): self.Qe_mean = np.array(self.Qe_mean) self.Qi_mean = np.array(self.Qi_mean) -class CGYROout(GACODErun.GACODEoutput): +class CGYROout(SIMtools.GACODEoutput): def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for_linear=True, **kwargs): super().__init__() diff --git a/src/mitim_tools/gacode_tools/utils/GACODErun.py b/src/mitim_tools/gacode_tools/utils/GACODErun.py index c8328900..ccce9612 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODErun.py +++ b/src/mitim_tools/gacode_tools/utils/GACODErun.py @@ -1,15 +1,7 @@ import shutil -import datetime -import time -import os -import copy import numpy as np -from pathlib import Path import matplotlib.pyplot as plt from scipy.interpolate import interp1d -from mitim_tools import __version__ as mitim_version -from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.gacode_tools.utils import GACODEdefaults, NORMtools from mitim_tools.transp_tools.utils import NTCCtools from mitim_tools.misc_tools import FARMINGtools, IOtools, MATHtools, GRAPHICStools from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -17,1028 +9,6 @@ from mitim_tools.misc_tools.PLASMAtools import md_u -class gacode_simulation: - ''' - Main class for running GACODE simulations. - ''' - def __init__( - self, - rhos=[0.4, 0.6], # rho locations of interest - ): - self.rhos = np.array(rhos) if rhos is not None else None - - self.ResultsFiles = [] - self.ResultsFiles_minimal = [] - - self.nameRunid = "0" - - self.results, self.scans = {}, {} - - self.run_specifications = None - - def prep( - self, - mitim_state, # A MITIM state class - FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) - cold_start=False, # If True, do not use what it potentially inside the folder, run again - forceIfcold_start=False, # Extra flag - ): - ''' - This method prepares the GACODE run from a MITIM state class by setting up the necessary input files and directories. - ''' - - print("> Preparation run from input.gacode (direct conversion)") - - if self.run_specifications is None: - raise Exception("[MITIM] Simulation child class did not define run specifications") - - state_converter = self.run_specifications['state_converter'] # e.g. to_tglf - input_class = self.run_specifications['input_class'] # e.g. TGLFinput - input_file = self.run_specifications['input_file'] # e.g. input.tglf - - self.FolderGACODE = IOtools.expandPath(FolderGACODE) - - if cold_start or not self.FolderGACODE.exists(): - IOtools.askNewFolder(self.FolderGACODE, force=forceIfcold_start) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare state - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if isinstance(mitim_state, str) or isinstance(mitim_state, Path): - # If a string, assume it's a path to input.gacode - self.profiles = PROFILEStools.gacode_state(mitim_state) - else: - self.profiles = mitim_state - - # Keep a copy of the file - self.profiles.write_state(file=self.FolderGACODE / "input.gacode") - - self.profiles.derive_quantities(mi_ref=md_u) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize from state - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - # Call the method dynamically based on state_converter - conversion_method = getattr(self.profiles, state_converter) - self.inputs_files = conversion_method(r=self.rhos, r_is_rho=True) - - for rho in self.inputs_files: - - # Initialize class - self.inputs_files[rho] = input_class.initialize_in_memory(self.inputs_files[rho]) - - # Write input.tglf file - self.inputs_files[rho].file = self.FolderGACODE / f'{input_file}_{rho:.4f}' - self.inputs_files[rho].write_state() - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Definining normalizations - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - print("> Setting up normalizations") - self.NormalizationSets, cdf = NORMtools.normalizations(self.profiles) - - return cdf - - def run( - self, - subfolder, # 'neo1/', - code_settings=None, - extraOptions={}, - multipliers={}, - minimum_delta_abs={}, - ApplyCorrections=True, # Removing ions with too low density and that are fast species - Quasineutral=False, # Ensures quasineutrality. By default is False because I may want to run the file directly - launchSlurm=True, - cold_start=False, - forceIfcold_start=False, - extra_name="exe", - slurm_setup=None, # Cores per call (so, when running nR radii -> nR*4) - attempts_execution=1, - only_minimal_files=False, - run_type = 'normal', # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit - ): - - if slurm_setup is None: - slurm_setup = {"cores": self.run_specifications['default_cores'], "minutes": 10} - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare inputs - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - code_executor, code_executor_full = self._run_prepare( - # - subfolder, - code_executor={}, - code_executor_full={}, - # - code_settings=code_settings, - extraOptions=extraOptions, - multipliers=multipliers, - # - cold_start=cold_start, - forceIfcold_start=forceIfcold_start, - only_minimal_files=only_minimal_files, - # - launchSlurm=launchSlurm, - slurm_setup=slurm_setup, - # - ApplyCorrections=ApplyCorrections, - minimum_delta_abs=minimum_delta_abs, - Quasineutral=Quasineutral, - ) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Run NEO - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - self._run( - code_executor, - code_executor_full=code_executor_full, - code_settings=code_settings, - ApplyCorrections=ApplyCorrections, - Quasineutral=Quasineutral, - launchSlurm=launchSlurm, - cold_start=cold_start, - forceIfcold_start=forceIfcold_start, - extra_name=extra_name, - slurm_setup=slurm_setup, - only_minimal_files=only_minimal_files, - attempts_execution=attempts_execution, - run_type=run_type - ) - - return code_executor_full - - def _run_prepare( - self, - # ******************************** - # Required options - # ******************************** - subfolder_simulation, - code_executor=None, - code_executor_full=None, - # ******************************** - # Run settings - # ******************************** - code_settings=None, - extraOptions={}, - multipliers={}, - # ******************************** - # IO settings - # ******************************** - cold_start=False, - forceIfcold_start=False, - only_minimal_files=False, - # ******************************** - # Slurm settings (for warnings) - # ******************************** - launchSlurm=True, - slurm_setup=None, - # ******************************** - # Additional settings to correct/modify inputs - # ******************************** - **kwargs_control - ): - - if slurm_setup is None: - slurm_setup = {"cores": self.run_specifications['default_cores'], "minutes": 5} - - if self.run_specifications is None: - raise Exception("[MITIM] Simulation child class did not define run specifications") - - # Because of historical relevance, I allow both TGLFsettings and code_settings #TODO #TOREMOVE - if "TGLFsettings" in kwargs_control: - if code_settings is not None: - raise Exception('[MITIM] Cannot use both TGLFsettings and code_settings') - else: - code_settings = kwargs_control["TGLFsettings"] - del kwargs_control["TGLFsettings"] - # ------------------------------------------------------------------------------------ - - if code_executor is None: - code_executor = {} - if code_executor_full is None: - code_executor_full = {} - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare for run - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - rhos = self.rhos - - inputs = copy.deepcopy(self.inputs_files) - Folder_sim = self.FolderGACODE / subfolder_simulation - - ResultsFiles_new = [] - for i in self.ResultsFiles: - if "mitim.out" not in i: - ResultsFiles_new.append(i) - self.ResultsFiles = ResultsFiles_new - - if only_minimal_files: - filesToRetrieve = self.ResultsFiles_minimal - else: - filesToRetrieve = self.ResultsFiles - - # Do I need to run all radii? - rhosEvaluate = cold_start_checker( - rhos, - filesToRetrieve, - Folder_sim, - cold_start=cold_start, - ) - - if len(rhosEvaluate) == len(rhos): - # All radii need to be evaluated - IOtools.askNewFolder(Folder_sim, force=forceIfcold_start) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Change this specific run - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - latest_inputsFile, latest_inputsFileDict = change_and_write_code( - rhos, - inputs, - Folder_sim, - code_settings=code_settings, - extraOptions=extraOptions, - multipliers=multipliers, - addControlFunction=self.run_specifications['control_function'], - controls_file=self.run_specifications['controls_file'], - **kwargs_control - ) - - code_executor_full[subfolder_simulation] = {} - code_executor[subfolder_simulation] = {} - for irho in self.rhos: - code_executor_full[subfolder_simulation][irho] = { - "folder": Folder_sim, - "dictionary": latest_inputsFileDict[irho], - "inputs": latest_inputsFile[irho], - "extraOptions": extraOptions, - "multipliers": multipliers, - } - if irho in rhosEvaluate: - code_executor[subfolder_simulation][irho] = code_executor_full[subfolder_simulation][irho] - - # Check input file problems - for irho in latest_inputsFileDict: - latest_inputsFileDict[irho].anticipate_problems() - - # Check cores problem - # if launchSlurm: - # self._check_cores(rhosEvaluate, slurm_setup) - - self.FolderSimLast = Folder_sim - - return code_executor, code_executor_full - - def _check_cores(self, rhosEvaluate, slurm_setup, warning = 32 * 2): - expected_allocated_cores = int(len(rhosEvaluate) * slurm_setup["cores"]) - - print(f'\t- Slurm job will be submitted with {expected_allocated_cores} cores ({len(rhosEvaluate)} radii x {slurm_setup["cores"]} cores/radius)', - typeMsg="" if expected_allocated_cores < warning else "q",) - - def _run( - self, - code_executor, - run_type = 'normal', # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit - **kwargs_run - ): - """ - extraOptions and multipliers are not being grabbed from kwargs_NEOrun, but from code_executor for WF - """ - - if kwargs_run.get("only_minimal_files", False): - filesToRetrieve = self.ResultsFiles_minimal - else: - filesToRetrieve = self.ResultsFiles - - c = 0 - for subfolder_simulation in code_executor: - c += len(code_executor[subfolder_simulation]) - - if c == 0: - - print(f"\t- {self.run_specifications['code'].upper()} not run because all results files found (please ensure consistency!)",typeMsg="i") - - else: - - # ---------------------------------------------------------------------------------------------------------------- - # Run simulation - # ---------------------------------------------------------------------------------------------------------------- - """ - launchSlurm = True -> Launch as a batch job in the machine chosen - launchSlurm = False -> Launch locally as a bash script - """ - - # Get code info - code = self.run_specifications.get('code', 'tglf') - input_file = self.run_specifications.get('input_file', 'input.tglf') - code_call = self.run_specifications.get('code_call', None) - code_slurm_settings = self.run_specifications.get('code_slurm_settings', None) - - # Get execution info - minutes = kwargs_run.get("slurm_setup", {}).get("minutes", 5) - cores_per_code_call = kwargs_run.get("slurm_setup", {}).get("cores", self.run_specifications['default_cores']) - launchSlurm = kwargs_run.get("launchSlurm", True) - - extraFlag = kwargs_run.get('extra_name', '') - name = f"{self.run_specifications['code']}_{self.nameRunid}{extraFlag}" - - attempts_execution = kwargs_run.get("attempts_execution", 1) - - tmpFolder = self.FolderGACODE / f"tmp_{code}" - IOtools.askNewFolder(tmpFolder, force=True) - - self.simulation_job = FARMINGtools.mitim_job(tmpFolder) - - self.simulation_job.define_machine_quick(code,f"mitim_{name}") - - folders, folders_red = [], [] - for subfolder_sim in code_executor: - - rhos = list(code_executor[subfolder_sim].keys()) - - # --------------------------------------------- - # Prepare files and folders - # --------------------------------------------- - - for i, rho in enumerate(rhos): - print(f"\t- Preparing {code.upper()} execution ({subfolder_sim}) at rho={rho:.4f}") - - folder_sim_this = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" - folders.append(folder_sim_this) - - folder_sim_this_rel = folder_sim_this.relative_to(tmpFolder) - folders_red.append(folder_sim_this_rel.as_posix() if self.simulation_job.machineSettings['machine'] != 'local' else str(folder_sim_this_rel)) - - folder_sim_this.mkdir(parents=True, exist_ok=True) - - input_file_sim = folder_sim_this / input_file - with open(input_file_sim, "w") as f: - f.write(code_executor[subfolder_sim][rho]["inputs"]) - - # --------------------------------------------- - # Prepare command - # --------------------------------------------- - - # Grab machine local limits ------------------------------------------------- - machineSettings = FARMINGtools.mitim_job.grab_machine_settings(code) - max_cores_per_node = machineSettings["cores_per_node"] - - # If the run is local and not slurm, let's check the number of cores - if (machineSettings["machine"] == "local") and \ - not (launchSlurm and ("partition" in self.simulation_job.machineSettings["slurm"])): - - cores_in_machine = int(os.cpu_count()) - cores_allocated = int(os.environ.get('SLURM_CPUS_PER_TASK')) if os.environ.get('SLURM_CPUS_PER_TASK') is not None else None - - if cores_allocated is not None: - if max_cores_per_node is None or (cores_allocated < max_cores_per_node): - print(f"\t - Detected {cores_allocated} cores allocated by SLURM, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") - max_cores_per_node = cores_allocated - elif cores_in_machine is not None: - if max_cores_per_node is None or (cores_in_machine < max_cores_per_node): - print(f"\t - Detected {cores_in_machine} cores in machine, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") - max_cores_per_node = cores_in_machine - else: - # Default to just 16 just in case - if max_cores_per_node is None: - max_cores_per_node = 16 - else: - # For remote execution, default to just 16 just in case - if max_cores_per_node is None: - max_cores_per_node = 16 - # --------------------------------------------------------------------------- - - # Grab the total number of cores of this job -------------------------------- - total_simulation_executions = len(rhos) * len(code_executor) - total_cores_required = int(cores_per_code_call) * total_simulation_executions - # --------------------------------------------------------------------------- - - # If it's GPUS enable machine, do the comparison based on it - if machineSettings['gpus_per_node'] == 0: - max_cores_per_node_compare = max_cores_per_node - else: - print(f"\t - Detected {machineSettings['gpus_per_node']} GPUs in machine, using this value as maximum for non-arra execution (vs {max_cores_per_node} specified)",typeMsg="i") - max_cores_per_node_compare = machineSettings['gpus_per_node'] - - if not (launchSlurm and ("partition" in self.simulation_job.machineSettings["slurm"])): - type_of_submission = "bash" - elif total_cores_required < max_cores_per_node_compare: - type_of_submission = "slurm_standard" - elif total_cores_required >= max_cores_per_node_compare: - type_of_submission = "slurm_array" - - shellPreCommands, shellPostCommands = None, None - - # Simply bash, no slurm - if type_of_submission == "bash": - - if cores_per_code_call > max_cores_per_node: - print(f"\t- Detected {cores_per_code_call} cores required, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") - max_cores_per_node = cores_per_code_call - - max_parallel_execution = max_cores_per_node // cores_per_code_call # Make sure we don't overload the machine when running locally (assuming no farming trans-node) - - print(f"\t- {code.upper()} will be executed as bash script (total cores: {total_cores_required}, cores per simulation: {cores_per_code_call}). MITIM will launch {total_simulation_executions // max_parallel_execution+1} sequential executions",typeMsg="i") - - # Build the bash script with job control enabled and a loop to limit parallel jobs - GACODEcommand = "#!/usr/bin/env bash\n" - GACODEcommand += "set -m\n" # Enable job control even in non-interactive mode - GACODEcommand += f"max_parallel_execution={max_parallel_execution}\n\n" # Set the maximum number of parallel processes - - # Create a bash array of folders - GACODEcommand += "folders=(\n" - for folder in folders_red: - GACODEcommand += f' "{folder}"\n' - GACODEcommand += ")\n\n" - - # Loop over each folder and launch code, waiting if we've reached max_parallel_execution - GACODEcommand += "for folder in \"${folders[@]}\"; do\n" - GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = self.simulation_job.folderExecution)}\n' - GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" - GACODEcommand += "done\n\n" - GACODEcommand += "wait\n" - - # Standard job - elif type_of_submission == "slurm_standard": - - print(f"\t- {code.upper()} will be executed in SLURM as standard job (cpus: {total_cores_required})",typeMsg="i") - - # Code launches - GACODEcommand = "" - for folder in folders_red: - GACODEcommand += f' {code_call(folder = folder, n = cores_per_code_call, p = self.simulation_job.folderExecution)} &\n' - GACODEcommand += "\nwait" # This is needed so that the script doesn't end before each job - - # Job array - elif type_of_submission == "slurm_array": - - print(f"\t- {code.upper()} will be executed in SLURM as job array due to its size (cpus: {total_cores_required})",typeMsg="i") - - # As a pre-command, organize all folders in a simpler way - shellPreCommands = [] - shellPostCommands = [] - array_list = [] - for i, folder in enumerate(folders_red): - array_list.append(f"{i}") - folder_temp_array = f"run{i}" - folder_actual = folder - shellPreCommands.append(f"mkdir {self.simulation_job.folderExecution}/{folder_temp_array}; cp {self.simulation_job.folderExecution}/{folder_actual}/* {self.simulation_job.folderExecution}/{folder_temp_array}/.") - shellPostCommands.append(f"cp {self.simulation_job.folderExecution}/{folder_temp_array}/* {self.simulation_job.folderExecution}/{folder_actual}/.; rm -r {self.simulation_job.folderExecution}/{folder_temp_array}") - - # Code launches - indexed_folder = 'run"$SLURM_ARRAY_TASK_ID"' - GACODEcommand = code_call( - folder = indexed_folder, - n = cores_per_code_call, - p = self.simulation_job.folderExecution, - additional_command = f'1> {self.simulation_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {self.simulation_job.folderExecution}/{indexed_folder}/slurm_error.dat\n') - - # --------------------------------------------- - # Execute - # --------------------------------------------- - - slurm_settings = code_slurm_settings( - name=code, - minutes=minutes, - total_cores_required=total_cores_required, - cores_per_code_call=cores_per_code_call, - type_of_submission=type_of_submission, - array_list=array_list if type_of_submission == "slurm_array" else None - ) - - self.simulation_job.define_machine( - code, - f"mitim_{name}", - launchSlurm=launchSlurm, - slurm_settings=slurm_settings, - ) - - # I would like the mitim_job to check if the retrieved folders were complete - check_files_in_folder = {} - for folder in folders_red: - check_files_in_folder[folder] = filesToRetrieve - # --------------------------------------------- - - self.simulation_job.prep( - GACODEcommand, - input_folders=folders, - output_folders=folders_red, - check_files_in_folder=check_files_in_folder, - shellPreCommands=shellPreCommands, - shellPostCommands=shellPostCommands, - ) - - # Submit run and wait - if run_type == 'normal': - - self.simulation_job.run( - removeScratchFolders=True, - attempts_execution=attempts_execution - ) - - self._organize_results(code_executor, tmpFolder, filesToRetrieve) - - # Submit run but do not wait; the user should do checks and fetch results - elif run_type == 'submit': - self.simulation_job.run( - waitYN=False, - check_if_files_received=False, - removeScratchFolders=False, - removeScratchFolders_goingIn=kwargs_run.get("cold_start", False), - ) - - self.kwargs_organize = { - "code_executor": code_executor, - "tmpFolder": tmpFolder, - "filesToRetrieve": filesToRetrieve - } - - self.slurm_output = "slurm_output.dat" - - # Prepare how to search for the job without waiting for it - self.simulation_job.launchSlurm = True - self.simulation_job.slurm_settings['name'] = Path(self.simulation_job.folderExecution).name - - def check(self, every_n_minutes=None): - - if self.simulation_job.launchSlurm: - print("- Checker job status") - - while True: - self.simulation_job.check(file_output = self.slurm_output) - print(f'\t- Current status (as of {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}): {self.simulation_job.status} ({self.simulation_job.infoSLURM["STATE"]})') - if self.simulation_job.status == 2: - print("\n\t* Job considered finished (please do .fetch() to retrieve results)",typeMsg="i") - break - elif every_n_minutes is None: - print("\n\t* Job not finished yet") - break - else: - print(f"\n\t* Waiting {every_n_minutes} minutes") - time.sleep(every_n_minutes * 60) - else: - print("- Not checking status because this was run command line (not slurm)") - - def fetch(self): - """ - For a job that has been submitted but not waited for, once it is done, get the results - """ - - print("\n\n\t- Fetching results") - - if self.simulation_job.launchSlurm: - self.simulation_job.connect() - self.simulation_job.retrieve() - self.simulation_job.close() - - self._organize_results(**self.kwargs_organize) - - else: - print("- Not retrieving results because this was run command line (not slurm)") - - def delete(self): - - print("\n\n\t- Deleting job") - - self.simulation_job.launchSlurm = False - - self.simulation_job.prep( - f"scancel -n {self.simulation_job.slurm_settings['name']}", - label_log_files="_finish", - ) - - self.simulation_job.run() - - def _organize_results(self, code_executor, tmpFolder, filesToRetrieve): - - # --------------------------------------------- - # Organize - # --------------------------------------------- - - print("\t- Retrieving files and changing names for storing") - fineall = True - for subfolder_sim in code_executor: - - for i, rho in enumerate(code_executor[subfolder_sim].keys()): - for file in filesToRetrieve: - original_file = f"{file}_{rho:.4f}" - final_destination = ( - code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" - ) - final_destination.unlink(missing_ok=True) - - temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" - temp_file.replace(final_destination) - - fineall = fineall and final_destination.exists() - - if not final_destination.exists(): - print(f"\t!! file {file} ({original_file}) could not be retrived",typeMsg="w",) - - if fineall: - print("\t\t- All files were successfully retrieved") - - # Remove temporary folder - shutil.rmtree(tmpFolder) - - else: - print("\t\t- Some files were not retrieved", typeMsg="w") - - - def run_scan( - self, - subfolder, # 'scan1', - multipliers={}, - minimum_delta_abs={}, - variable="RLTS_1", - varUpDown=[0.5, 1.0, 1.5], - variables_scanTogether=[], - relativeChanges=True, - **kwargs_run, - ): - - # ------------------------------------- - # Add baseline - # ------------------------------------- - if (1.0 not in varUpDown) and relativeChanges: - print("\n* Since variations vector did not include base case, I am adding it",typeMsg="i",) - varUpDown_new = [] - added = False - for i in varUpDown: - if i > 1.0 and not added: - varUpDown_new.append(1.0) - added = True - varUpDown_new.append(i) - else: - varUpDown_new = varUpDown - - - code_executor, code_executor_full, folders, varUpDown_new = self._prepare_scan( - subfolder, - multipliers=multipliers, - minimum_delta_abs=minimum_delta_abs, - variable=variable, - varUpDown=varUpDown_new, - variables_scanTogether=variables_scanTogether, - relativeChanges=relativeChanges, - **kwargs_run, - ) - - # Run them all - self._run( - code_executor, - code_executor_full=code_executor_full, - **kwargs_run, - ) - - # Read results - for cont_mult, mult in enumerate(varUpDown_new): - name = f"{variable}_{mult}" - self.read( - label=f"{self.subfolder_scan}_{name}", - folder=folders[cont_mult], - cold_startWF = False, - require_all_files=not kwargs_run.get("only_minimal_files",False), - ) - - return code_executor_full - - def _prepare_scan( - self, - subfolder, # 'scan1', - multipliers={}, - minimum_delta_abs={}, - variable="RLTS_1", - varUpDown=[0.5, 1.0, 1.5], - variables_scanTogether=[], - relativeChanges=True, - **kwargs_run, - ): - """ - Multipliers will be modified by adding the scaning variables, but I don't want to modify the original - multipliers, as they may be passed to the next scan - - Set relativeChanges=False if varUpDown contains the exact values to change, not multipleiers - """ - - completeVariation = self.run_specifications['complete_variation'] - - multipliers_mod = copy.deepcopy(multipliers) - - self.subfolder_scan = subfolder - - if relativeChanges: - for i in range(len(varUpDown)): - varUpDown[i] = round(varUpDown[i], 6) - - print(f"\n- Proceeding to scan {variable}{' together with '+', '.join(variables_scanTogether) if len(variables_scanTogether)>0 else ''}:") - - code_executor = {} - code_executor_full = {} - folders = [] - for cont_mult, mult in enumerate(varUpDown): - mult = round(mult, 6) - - if relativeChanges: - print(f"\n + Multiplier: {mult} -----------------------------------------------------------------------------------------------------------") - else: - print(f"\n + Value: {mult} ----------------------------------------------------------------------------------------------------------------") - - multipliers_mod[variable] = mult - - for variable_scanTogether in variables_scanTogether: - multipliers_mod[variable_scanTogether] = mult - - name = f"{variable}_{mult}" - - species = self.inputs_files[self.rhos[0]] # Any rho will do - - if completeVariation is not None: - multipliers_mod = completeVariation(multipliers_mod, species) - - if not relativeChanges: - for ikey in multipliers_mod: - kwargs_run["extraOptions"][ikey] = multipliers_mod[ikey] - multipliers_mod = {} - - # Force ensure quasineutrality if the - if variable in ["AS_3", "AS_4", "AS_5", "AS_6"]: - kwargs_run["Quasineutral"] = True - - # Only ask the cold_start in the first round - kwargs_run["forceIfcold_start"] = cont_mult > 0 or ("forceIfcold_start" in kwargs_run and kwargs_run["forceIfcold_start"]) - - code_executor, code_executor_full = self._run_prepare( - f"{self.subfolder_scan}_{name}", - code_executor=code_executor, - code_executor_full=code_executor_full, - multipliers=multipliers_mod, - minimum_delta_abs=minimum_delta_abs, - **kwargs_run, - ) - - folders.append(copy.deepcopy(self.FolderSimLast)) - - return code_executor, code_executor_full, folders, varUpDown - - def read( - self, - label="run1", - folder=None, # If None, search in the previously run folder - suffix=None, # If None, search with my standard _0.55 suffixes corresponding to rho of this TGLF class - **kwargs_to_class_output - ): - print("> Reading simulation results") - - class_output = [self.run_specifications['output_class'], self.run_specifications['output_store']] - - # If no specified folder, check the last one - if folder is None: - folder = self.FolderSimLast - - self.results[label] = { - class_output[1]:[], - 'parsed': [], - "x": np.array(self.rhos), - } - for rho in self.rhos: - - SIMout = class_output[0]( - folder, - suffix=f"_{rho:.4f}" if suffix is None else suffix, - **kwargs_to_class_output - ) - - # Unnormalize - if 'NormalizationSets' in self.__dict__: - SIMout.unnormalize( - self.NormalizationSets["SELECTED"], - rho=rho, - ) - else: - print("No normalization sets found.") - - self.results[label][class_output[1]].append(SIMout) - - self.results[label]['parsed'].append(buildDictFromInput(SIMout.inputFile) if SIMout.inputFile else None) - - def read_scan( - self, - label="scan1", - subfolder=None, - variable="RLTS_1", - positionIon=2, - variable_mapping=None, - variable_mapping_unn=None - ): - ''' - positionIon is the index in the input.tglf file... so if you want for ion RLNS_5, positionIon=5 - ''' - - if subfolder is None: - subfolder = self.subfolder_scan - - self.scans[label] = {} - self.scans[label]["variable"] = variable - self.scans[label]["positionBase"] = None - self.scans[label]["unnormalization_successful"] = True - self.scans[label]["results_tags"] = [] - - self.positionIon_scan = positionIon - - # ---- - - scan = {} - for ikey in variable_mapping | variable_mapping_unn: - scan[ikey] = [] - - cont = 0 - for ikey in self.results: - isThisTheRightReadResults = (subfolder in ikey) and (variable== "_".join(ikey.split("_")[:-1]).split(subfolder + "_")[-1]) - - if isThisTheRightReadResults: - - self.scans[label]["results_tags"].append(ikey) - - # Initialize lists - scan0 = {} - for ikey2 in variable_mapping | variable_mapping_unn: - scan0[ikey2] = [] - - # Loop over radii - for irho_cont in range(len(self.rhos)): - irho = np.where(self.results[ikey]["x"] == self.rhos[irho_cont])[0][0] - - for ikey2 in variable_mapping: - - obj = self.results[ikey][variable_mapping[ikey2][0]][irho] - if not hasattr(obj, '__dict__'): - obj_dict = obj - else: - obj_dict = obj.__dict__ - var0 = obj_dict[variable_mapping[ikey2][1]] - scan0[ikey2].append(var0 if variable_mapping[ikey2][2] is None else var0[variable_mapping[ikey2][2]]) - - # Unnormalized - self.scans[label]["unnormalization_successful"] = True - for ikey2 in variable_mapping_unn: - obj = self.results[ikey][variable_mapping_unn[ikey2][0]][irho] - if not hasattr(obj, '__dict__'): - obj_dict = obj - else: - obj_dict = obj.__dict__ - - if variable_mapping_unn[ikey2][1] not in obj_dict: - self.scans[label]["unnormalization_successful"] = False - break - var0 = obj_dict[variable_mapping_unn[ikey2][1]] - scan0[ikey2].append(var0 if variable_mapping_unn[ikey2][2] is None else var0[variable_mapping_unn[ikey2][2]]) - - for ikey2 in variable_mapping | variable_mapping_unn: - scan[ikey2].append(scan0[ikey2]) - - if float(ikey.split('_')[-1]) == 1.0: - self.scans[label]["positionBase"] = cont - cont += 1 - - self.scans[label]["x"] = np.array(self.rhos) - - for ikey2 in variable_mapping | variable_mapping_unn: - self.scans[label][ikey2] = np.atleast_2d(np.transpose(scan[ikey2])) - - - -def change_and_write_code( - rhos, - inputs0, - Folder_sim, - code_settings=None, - extraOptions={}, - multipliers={}, - minimum_delta_abs={}, - ApplyCorrections=True, - Quasineutral=False, - addControlFunction=None, - controls_file='input.tglf.controls', - **kwargs -): - """ - Received inputs classes and gives text. - ApplyCorrections refer to removing ions with too low density and that are fast species - """ - - inputs = copy.deepcopy(inputs0) - - mod_input_file = {} - ns_max = [] - for i, rho in enumerate(rhos): - print(f"\t- Changing input file for rho={rho:.4f}") - input_sim_rho = modifyInputs( - inputs[rho], - code_settings=code_settings, - extraOptions=extraOptions, - multipliers=multipliers, - minimum_delta_abs=minimum_delta_abs, - position_change=i, - addControlFunction=addControlFunction, - controls_file=controls_file, - NS=inputs[rho].num_recorded, - ) - - input_file = input_sim_rho.file.name.split('_')[0] - - newfile = Folder_sim / f"{input_file}_{rho:.4f}" - - if code_settings is not None: - # Apply corrections - if ApplyCorrections: - print("\t- Applying corrections") - input_sim_rho.removeLowDensitySpecie() - input_sim_rho.remove_fast() - - # Ensure that plasma to run is quasineutral - if Quasineutral: - input_sim_rho.ensureQuasineutrality() - else: - print('\t- Not applying corrections because settings is None') - - input_sim_rho.write_state(file=newfile) - - mod_input_file[rho] = input_sim_rho - - ns_max.append(inputs[rho].num_recorded) - - # Convert back to a string because that's how the run operates - inputFile = inputToVariable(Folder_sim, rhos, file=input_file) - - if (np.diff(ns_max) > 0).any(): - print("> Each radial location has its own number of species... probably because of removal of fast or low density...",typeMsg="w") - print("\t * Reading of simulation results will fail... consider doing something before launching run",typeMsg="q") - - return inputFile, mod_input_file - - -def inputToVariable(folder, rhos, file='input.tglf'): - """ - Entire text file to variable - """ - - inputFilesTGLF = {} - for rho in rhos: - fileN = folder / f"{file}_{rho:.4f}" - - with open(fileN, "r") as f: - lines = f.readlines() - inputFilesTGLF[rho] = "".join(lines) - - return inputFilesTGLF - - -def cold_start_checker( - rhos, - ResultsFiles, - Folder_sim, - cold_start=False, - print_each_time=False, -): - """ - This function checks if the TGLF inputs are already in the folder. If they are, it returns True - """ - cont_each = 0 - if cold_start: - rhosEvaluate = rhos - else: - rhosEvaluate = [] - for ir in rhos: - existsRho = True - for j in ResultsFiles: - ffi = Folder_sim / f"{j}_{ir:.4f}" - existsThis = ffi.exists() - existsRho = existsRho and existsThis - if not existsThis: - if print_each_time: - print(f"\t* {ffi} does not exist") - else: - cont_each += 1 - if not existsRho: - rhosEvaluate.append(ir) - - if not print_each_time and cont_each > 0: - print(f'\t* {cont_each} files from expected set are missing') - - if len(rhosEvaluate) < len(rhos) and len(rhosEvaluate) > 0: - print("~ Not all radii are found, but not removing folder and running only those that are needed",typeMsg="i",) - - return rhosEvaluate - - def runTGYRO( folderWork, outputFiles=None, @@ -1123,121 +93,6 @@ def runTGYRO( tgyro_job.run() - -def modifyInputs( - input_class, - code_settings=None, - extraOptions={}, - multipliers={}, - minimum_delta_abs={}, - position_change=0, - addControlFunction=None, - controls_file = 'input.tglf.controls', - **kwargs_to_function, -): - - # Check that those are valid flags - GACODEdefaults.review_controls(extraOptions, control = controls_file) - GACODEdefaults.review_controls(multipliers, control = controls_file) - # ------------------------------------------- - - if code_settings is not None: - CodeOptions = addControlFunction(code_settings, **kwargs_to_function) - - # ~~~~~~~~~~ Change with presets - print(f" \t- Using presets code_settings = {code_settings}", typeMsg="i") - input_class.controls = CodeOptions - - else: - print("\t- Input file was not modified by code_settings, using what was there before",typeMsg="i") - - # Make all upper case - #extraOptions = {ikey.upper(): value for ikey, value in extraOptions.items()} - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Change with external options -> Input directly, not as multiplier - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - 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] - - try: - isspecie = ikey.split("_")[0] in input_class.species[1] - except: - isspecie = False - - # is a species parameter? - if isspecie: - specie = int(ikey.split("_")[-1]) - varK = "_".join(ikey.split("_")[:-1]) - var_orig = input_class.species[specie][varK] - var_new = value_to_change_to - input_class.species[specie][varK] = var_new - # is a another parameter? - else: - if ikey in input_class.controls: - var_orig = input_class.controls[ikey] - var_new = value_to_change_to - input_class.controls[ikey] = var_new - elif ikey in input_class.plasma: - var_orig = input_class.plasma[ikey] - var_new = value_to_change_to - input_class.plasma[ikey] = var_new - else: - # If the variable in extraOptions wasn't in there, consider it a control param - print(f"\t\t- Variable {ikey} to change did not exist previously, creating now",typeMsg="i") - var_orig = None - var_new = value_to_change_to - input_class.controls[ikey] = var_new - - print(f"\t\t- Changing {ikey} from {var_orig} to {var_new}",typeMsg="i",) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Change with multipliers -> Input directly, not as multiplier - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if len(multipliers) > 0: - print("\t\t- Variables change:") - for ikey in multipliers: - # is a specie one? - if "species" in input_class.__dict__.keys() and ikey.split("_")[0] in input_class.species[1]: - specie = int(ikey.split("_")[-1]) - varK = "_".join(ikey.split("_")[:-1]) - var_orig = input_class.species[specie][varK] - var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) - input_class.species[specie][varK] = var_new - else: - if ikey in input_class.controls: - var_orig = input_class.controls[ikey] - var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) - input_class.controls[ikey] = var_new - - elif ikey in input_class.plasma: - var_orig = input_class.plasma[ikey] - var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) - input_class.plasma[ikey] = var_new - - else: - print("\t- Variable to scan did not exist in original file, add it as extraOptions first",typeMsg="w",) - - print(f"\t\t\t- Changing {ikey} from {var_orig} to {var_new} (x{multipliers[ikey]})") - - return input_class - -def multiplier_input(var_orig, multiplier, minimum_delta_abs = None): - - delta = var_orig * (multiplier - 1.0) - - if minimum_delta_abs is not None: - if (multiplier != 1.0) and abs(delta) < minimum_delta_abs: - print(f"\t\t\t- delta = {delta} is smaller than minimum_delta_abs = {minimum_delta_abs}, enforcing",typeMsg="i") - delta = np.sign(delta) * minimum_delta_abs - - return var_orig + delta - def findNamelist(LocationCDF, folderWork=None, nameRunid="10000", ForceFirst=True): # ----------------------------------------------------------- # Find namelist @@ -1264,7 +119,6 @@ def findNamelist(LocationCDF, folderWork=None, nameRunid="10000", ForceFirst=Tru return LocationNML, dummy - def prepareTGYRO( LocationCDF, LocationNML, @@ -1304,7 +158,6 @@ def prepareTGYRO( includeGEQ=includeGEQ, ) - def CDFtoTRXPLoutput( LocationCDF, LocationNML, @@ -1419,7 +272,6 @@ def runTRXPL( ) trxpl_job.run() - def runPROFILES_GEN( FolderTGLF, nameFiles="10001", @@ -1486,7 +338,6 @@ def runPROFILES_GEN( print(f"\t\t- Proceeding to run PROFILES_GEN with: {txt}") pgen_job.run() - def runVGEN( workingFolder, numcores=32, @@ -1563,45 +414,10 @@ def runVGEN( return file_new - -def buildDictFromInput(inputFile): - parsed = {} - - lines = inputFile.split("\n") - for line in lines: - if "=" in line: - splits = [i.split()[0] for i in line.split("=")] - if ("." in splits[1]) and (splits[1][0].split()[0] != "."): - parsed[splits[0].split()[0]] = float(splits[1].split()[0]) - else: - try: - parsed[splits[0].split()[0]] = int(splits[1].split()[0]) - except: - parsed[splits[0].split()[0]] = splits[1].split()[0] - - for i in parsed: - if isinstance(parsed[i], str): - if ( - parsed[i].lower() == "t" - or parsed[i].lower() == "true" - or parsed[i].lower() == ".true." - ): - parsed[i] = True - elif ( - parsed[i].lower() == "f" - or parsed[i].lower() == "false" - or parsed[i].lower() == ".false." - ): - parsed[i] = False - - return parsed - - # ---------------------------------------------------------------------- # Reading/Writing routines # ---------------------------------------------------------------------- - def obtainFluctuationLevel( ky, Amplitude, @@ -1643,7 +459,6 @@ def obtainFluctuationLevel( return fluctSim * 100.0 * factorTot_to_Perp - def obtainNTphase( ky, nTphase, @@ -1664,7 +479,6 @@ def obtainNTphase( return neTe - def integrateSpectrum( xOriginal, yOriginal, @@ -1806,7 +620,6 @@ def integrateSpectrum( return integ - def defineNewGrid( xOriginal1, yOriginal1, @@ -1872,109 +685,3 @@ def defineNewGrid( return x[imin:imax], y[imin:imax] -class GACODEoutput: - def __init__(self, *args, **kwargs): - self.inputFile = None - - def unnormalize(self, *args, **kwargs): - print("No unnormalization implemented.") - -class GACODEinput: - def __init__(self, file=None, controls_file=None, code='', n_species=None): - self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None - - self.controls_file = controls_file - self.code = code - self.n_species = n_species - - self.num_recorded = 100 - - if self.file is not None and self.file.exists(): - with open(self.file, "r") as f: - lines = f.readlines() - file_txt = "".join(lines) - else: - file_txt = "" - input_dict = buildDictFromInput(file_txt) - - self.process(input_dict) - - @classmethod - def initialize_in_memory(cls, input_dict): - instance = cls() - instance.process(input_dict) - return instance - - def process(self, input_dict): - - if self.controls_file is not None: - options_check = [key for key in IOtools.generateMITIMNamelist(self.controls_file, caseInsensitive=False).keys()] - else: - options_check = [] - - self.controls, self.plasma = {}, {} - for key in input_dict.keys(): - if key in options_check: - self.controls[key] = input_dict[key] - else: - self.plasma[key] = input_dict[key] - - # Get number of recorded species - if self.n_species is not None and self.n_species in input_dict: - self.num_recorded = int(input_dict[self.n_species]) - - def write_state(self, file=None): - - if file is None: - file = self.file - - # Local formatter: floats -> 6 significant figures in exponential (uppercase), - # ints stay as ints, bools as 0/1, sequences space-separated with same rule. - def _fmt_num(x): - import numpy as _np - if isinstance(x, (bool, _np.bool_)): - return "True" if x else "False" - if isinstance(x, (_np.floating, float)): - # 6 significant figures in exponential => 5 digits after decimal - return f"{float(x):.5E}" - if isinstance(x, (_np.integer, int)): - return f"{int(x)}" - return str(x) - - def _fmt_value(val): - import numpy as _np - if isinstance(val, (list, tuple, _np.ndarray)): - # Flatten numpy arrays but keep ordering; join with spaces - if isinstance(val, _np.ndarray): - flat = val.flatten().tolist() - else: - flat = list(val) - return " ".join(_fmt_num(v) for v in flat) - return _fmt_num(val) - - with open(file, "w") as f: - f.write("#-------------------------------------------------------------------------\n") - f.write(f"# {self.code} input file modified by MITIM {mitim_version}\n") - f.write("#-------------------------------------------------------------------------\n") - - f.write("\n\n# Control parameters\n") - f.write("# ------------------\n\n") - for ikey in self.controls: - var = self.controls[ikey] - f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") - - f.write("\n\n# Plasma/Geometry parameters\n") - f.write("# ------------------\n\n") - for ikey in self.plasma: - var = self.plasma[ikey] - f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") - - def anticipate_problems(self): - pass - - def remove_fast(self): - pass - - def removeLowDensitySpecie(self, *args): - pass - \ No newline at end of file diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index fd6920f1..156c76f9 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -620,7 +620,7 @@ def retrieve(self, check_if_files_received=True, check_files_in_folder={}): if received: print("\t\t- All correct", typeMsg="i") else: - print("\t* Not all received, trying once again", typeMsg="w") + print("\t* Not all received, trying once again", typeMsg="i") time.sleep(10) _ = self.retrieve(check_if_files_received=False) received = self.check_all_received(check_files_in_folder=check_files_in_folder) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py new file mode 100644 index 00000000..8790c695 --- /dev/null +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -0,0 +1,1285 @@ +import shutil +import datetime +import time +import os +import copy +import numpy as np +from pathlib import Path +from mitim_tools import __version__ as mitim_version +from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.gacode_tools.utils import GACODEdefaults, NORMtools +from mitim_tools.misc_tools import FARMINGtools, IOtools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +from mitim_tools.misc_tools.PLASMAtools import md_u + +class mitim_simulation: + ''' + Main class for running GACODE simulations. + ''' + def __init__( + self, + rhos=[0.4, 0.6], # rho locations of interest + ): + self.rhos = np.array(rhos) if rhos is not None else None + + self.ResultsFiles = [] + self.ResultsFiles_minimal = [] + + self.nameRunid = "0" + + self.results, self.scans = {}, {} + + self.run_specifications = None + + def prep( + self, + mitim_state, # A MITIM state class + FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) + cold_start=False, # If True, do not use what it potentially inside the folder, run again + forceIfcold_start=False, # Extra flag + ): + ''' + This method prepares the GACODE run from a MITIM state class by setting up the necessary input files and directories. + ''' + + print("> Preparation run from input.gacode (direct conversion)") + + if self.run_specifications is None: + raise Exception("[MITIM] Simulation child class did not define run specifications") + + state_converter = self.run_specifications['state_converter'] # e.g. to_tglf + input_class = self.run_specifications['input_class'] # e.g. TGLFinput + input_file = self.run_specifications['input_file'] # e.g. input.tglf + + self.FolderGACODE = IOtools.expandPath(FolderGACODE) + + if cold_start or not self.FolderGACODE.exists(): + IOtools.askNewFolder(self.FolderGACODE, force=forceIfcold_start) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare state + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + if isinstance(mitim_state, str) or isinstance(mitim_state, Path): + # If a string, assume it's a path to input.gacode + self.profiles = PROFILEStools.gacode_state(mitim_state) + else: + self.profiles = mitim_state + + # Keep a copy of the file + self.profiles.write_state(file=self.FolderGACODE / "input.gacode") + + self.profiles.derive_quantities(mi_ref=md_u) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize from state + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + # Call the method dynamically based on state_converter + conversion_method = getattr(self.profiles, state_converter) + self.inputs_files = conversion_method(r=self.rhos, r_is_rho=True) + + for rho in self.inputs_files: + + # Initialize class + self.inputs_files[rho] = input_class.initialize_in_memory(self.inputs_files[rho]) + + # Write input.tglf file + self.inputs_files[rho].file = self.FolderGACODE / f'{input_file}_{rho:.4f}' + self.inputs_files[rho].write_state() + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Definining normalizations + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + print("> Setting up normalizations") + self.NormalizationSets, cdf = NORMtools.normalizations(self.profiles) + + return cdf + + def run( + self, + subfolder, # 'neo1/', + code_settings=None, + extraOptions={}, + multipliers={}, + minimum_delta_abs={}, + ApplyCorrections=True, # Removing ions with too low density and that are fast species + Quasineutral=False, # Ensures quasineutrality. By default is False because I may want to run the file directly + launchSlurm=True, + cold_start=False, + forceIfcold_start=False, + extra_name="exe", + slurm_setup=None, # Cores per call (so, when running nR radii -> nR*4) + attempts_execution=1, + only_minimal_files=False, + run_type = 'normal', # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit + ): + + if slurm_setup is None: + slurm_setup = {"cores": self.run_specifications['default_cores'], "minutes": 10} + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare inputs + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + code_executor, code_executor_full = self._run_prepare( + # + subfolder, + code_executor={}, + code_executor_full={}, + # + code_settings=code_settings, + extraOptions=extraOptions, + multipliers=multipliers, + # + cold_start=cold_start, + forceIfcold_start=forceIfcold_start, + only_minimal_files=only_minimal_files, + # + launchSlurm=launchSlurm, + slurm_setup=slurm_setup, + # + ApplyCorrections=ApplyCorrections, + minimum_delta_abs=minimum_delta_abs, + Quasineutral=Quasineutral, + ) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Run NEO + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self._run( + code_executor, + code_executor_full=code_executor_full, + code_settings=code_settings, + ApplyCorrections=ApplyCorrections, + Quasineutral=Quasineutral, + launchSlurm=launchSlurm, + cold_start=cold_start, + forceIfcold_start=forceIfcold_start, + extra_name=extra_name, + slurm_setup=slurm_setup, + only_minimal_files=only_minimal_files, + attempts_execution=attempts_execution, + run_type=run_type + ) + + return code_executor_full + + def _run_prepare( + self, + # ******************************** + # Required options + # ******************************** + subfolder_simulation, + code_executor=None, + code_executor_full=None, + # ******************************** + # Run settings + # ******************************** + code_settings=None, + extraOptions={}, + multipliers={}, + # ******************************** + # IO settings + # ******************************** + cold_start=False, + forceIfcold_start=False, + only_minimal_files=False, + # ******************************** + # Slurm settings (for warnings) + # ******************************** + launchSlurm=True, + slurm_setup=None, + # ******************************** + # Additional settings to correct/modify inputs + # ******************************** + **kwargs_control + ): + + if slurm_setup is None: + slurm_setup = {"cores": self.run_specifications['default_cores'], "minutes": 5} + + if self.run_specifications is None: + raise Exception("[MITIM] Simulation child class did not define run specifications") + + # Because of historical relevance, I allow both TGLFsettings and code_settings #TODO #TOREMOVE + if "TGLFsettings" in kwargs_control: + if code_settings is not None: + raise Exception('[MITIM] Cannot use both TGLFsettings and code_settings') + else: + code_settings = kwargs_control["TGLFsettings"] + del kwargs_control["TGLFsettings"] + # ------------------------------------------------------------------------------------ + + if code_executor is None: + code_executor = {} + if code_executor_full is None: + code_executor_full = {} + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare for run + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + rhos = self.rhos + + inputs = copy.deepcopy(self.inputs_files) + Folder_sim = self.FolderGACODE / subfolder_simulation + + ResultsFiles_new = [] + for i in self.ResultsFiles: + if "mitim.out" not in i: + ResultsFiles_new.append(i) + self.ResultsFiles = ResultsFiles_new + + if only_minimal_files: + filesToRetrieve = self.ResultsFiles_minimal + else: + filesToRetrieve = self.ResultsFiles + + # Do I need to run all radii? + rhosEvaluate = cold_start_checker( + rhos, + filesToRetrieve, + Folder_sim, + cold_start=cold_start, + ) + + if len(rhosEvaluate) == len(rhos): + # All radii need to be evaluated + IOtools.askNewFolder(Folder_sim, force=forceIfcold_start) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Change this specific run + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + latest_inputsFile, latest_inputsFileDict = change_and_write_code( + rhos, + inputs, + Folder_sim, + code_settings=code_settings, + extraOptions=extraOptions, + multipliers=multipliers, + addControlFunction=self.run_specifications['control_function'], + controls_file=self.run_specifications['controls_file'], + **kwargs_control + ) + + code_executor_full[subfolder_simulation] = {} + code_executor[subfolder_simulation] = {} + for irho in self.rhos: + code_executor_full[subfolder_simulation][irho] = { + "folder": Folder_sim, + "dictionary": latest_inputsFileDict[irho], + "inputs": latest_inputsFile[irho], + "extraOptions": extraOptions, + "multipliers": multipliers, + } + if irho in rhosEvaluate: + code_executor[subfolder_simulation][irho] = code_executor_full[subfolder_simulation][irho] + + # Check input file problems + for irho in latest_inputsFileDict: + latest_inputsFileDict[irho].anticipate_problems() + + # Check cores problem + # if launchSlurm: + # self._check_cores(rhosEvaluate, slurm_setup) + + self.FolderSimLast = Folder_sim + + return code_executor, code_executor_full + + def _check_cores(self, rhosEvaluate, slurm_setup, warning = 32 * 2): + expected_allocated_cores = int(len(rhosEvaluate) * slurm_setup["cores"]) + + print(f'\t- Slurm job will be submitted with {expected_allocated_cores} cores ({len(rhosEvaluate)} radii x {slurm_setup["cores"]} cores/radius)', + typeMsg="" if expected_allocated_cores < warning else "q",) + + def _run( + self, + code_executor, + run_type = 'normal', # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit + **kwargs_run + ): + """ + extraOptions and multipliers are not being grabbed from kwargs_NEOrun, but from code_executor for WF + """ + + if kwargs_run.get("only_minimal_files", False): + filesToRetrieve = self.ResultsFiles_minimal + else: + filesToRetrieve = self.ResultsFiles + + c = 0 + for subfolder_simulation in code_executor: + c += len(code_executor[subfolder_simulation]) + + if c == 0: + + print(f"\t- {self.run_specifications['code'].upper()} not run because all results files found (please ensure consistency!)",typeMsg="i") + + else: + + # ---------------------------------------------------------------------------------------------------------------- + # Run simulation + # ---------------------------------------------------------------------------------------------------------------- + """ + launchSlurm = True -> Launch as a batch job in the machine chosen + launchSlurm = False -> Launch locally as a bash script + """ + + # Get code info + code = self.run_specifications.get('code', 'tglf') + input_file = self.run_specifications.get('input_file', 'input.tglf') + code_call = self.run_specifications.get('code_call', None) + code_slurm_settings = self.run_specifications.get('code_slurm_settings', None) + + # Get execution info + minutes = kwargs_run.get("slurm_setup", {}).get("minutes", 5) + cores_per_code_call = kwargs_run.get("slurm_setup", {}).get("cores", self.run_specifications['default_cores']) + launchSlurm = kwargs_run.get("launchSlurm", True) + + extraFlag = kwargs_run.get('extra_name', '') + name = f"{self.run_specifications['code']}_{self.nameRunid}{extraFlag}" + + attempts_execution = kwargs_run.get("attempts_execution", 1) + + tmpFolder = self.FolderGACODE / f"tmp_{code}" + IOtools.askNewFolder(tmpFolder, force=True) + + self.simulation_job = FARMINGtools.mitim_job(tmpFolder) + + self.simulation_job.define_machine_quick(code,f"mitim_{name}") + + folders, folders_red = [], [] + for subfolder_sim in code_executor: + + rhos = list(code_executor[subfolder_sim].keys()) + + # --------------------------------------------- + # Prepare files and folders + # --------------------------------------------- + + for i, rho in enumerate(rhos): + print(f"\t- Preparing {code.upper()} execution ({subfolder_sim}) at rho={rho:.4f}") + + folder_sim_this = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" + folders.append(folder_sim_this) + + folder_sim_this_rel = folder_sim_this.relative_to(tmpFolder) + folders_red.append(folder_sim_this_rel.as_posix() if self.simulation_job.machineSettings['machine'] != 'local' else str(folder_sim_this_rel)) + + folder_sim_this.mkdir(parents=True, exist_ok=True) + + input_file_sim = folder_sim_this / input_file + with open(input_file_sim, "w") as f: + f.write(code_executor[subfolder_sim][rho]["inputs"]) + + # --------------------------------------------- + # Prepare command + # --------------------------------------------- + + # Grab machine local limits ------------------------------------------------- + machineSettings = FARMINGtools.mitim_job.grab_machine_settings(code) + max_cores_per_node = machineSettings["cores_per_node"] + + # If the run is local and not slurm, let's check the number of cores + if (machineSettings["machine"] == "local") and \ + not (launchSlurm and ("partition" in self.simulation_job.machineSettings["slurm"])): + + cores_in_machine = int(os.cpu_count()) + cores_allocated = int(os.environ.get('SLURM_CPUS_PER_TASK')) if os.environ.get('SLURM_CPUS_PER_TASK') is not None else None + + if cores_allocated is not None: + if max_cores_per_node is None or (cores_allocated < max_cores_per_node): + print(f"\t - Detected {cores_allocated} cores allocated by SLURM, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node = cores_allocated + elif cores_in_machine is not None: + if max_cores_per_node is None or (cores_in_machine < max_cores_per_node): + print(f"\t - Detected {cores_in_machine} cores in machine, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node = cores_in_machine + else: + # Default to just 16 just in case + if max_cores_per_node is None: + max_cores_per_node = 16 + else: + # For remote execution, default to just 16 just in case + if max_cores_per_node is None: + max_cores_per_node = 16 + # --------------------------------------------------------------------------- + + # Grab the total number of cores of this job -------------------------------- + total_simulation_executions = len(rhos) * len(code_executor) + total_cores_required = int(cores_per_code_call) * total_simulation_executions + # --------------------------------------------------------------------------- + + # If it's GPUS enable machine, do the comparison based on it + if machineSettings['gpus_per_node'] == 0: + max_cores_per_node_compare = max_cores_per_node + else: + print(f"\t - Detected {machineSettings['gpus_per_node']} GPUs in machine, using this value as maximum for non-arra execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node_compare = machineSettings['gpus_per_node'] + + if not (launchSlurm and ("partition" in self.simulation_job.machineSettings["slurm"])): + type_of_submission = "bash" + elif total_cores_required < max_cores_per_node_compare: + type_of_submission = "slurm_standard" + elif total_cores_required >= max_cores_per_node_compare: + type_of_submission = "slurm_array" + + shellPreCommands, shellPostCommands = None, None + + # Simply bash, no slurm + if type_of_submission == "bash": + + if cores_per_code_call > max_cores_per_node: + print(f"\t- Detected {cores_per_code_call} cores required, using this value as maximum for local execution (vs {max_cores_per_node} specified)",typeMsg="i") + max_cores_per_node = cores_per_code_call + + max_parallel_execution = max_cores_per_node // cores_per_code_call # Make sure we don't overload the machine when running locally (assuming no farming trans-node) + + print(f"\t- {code.upper()} will be executed as bash script (total cores: {total_cores_required}, cores per simulation: {cores_per_code_call}). MITIM will launch {total_simulation_executions // max_parallel_execution+1} sequential executions",typeMsg="i") + + # Build the bash script with job control enabled and a loop to limit parallel jobs + GACODEcommand = "#!/usr/bin/env bash\n" + GACODEcommand += "set -m\n" # Enable job control even in non-interactive mode + GACODEcommand += f"max_parallel_execution={max_parallel_execution}\n\n" # Set the maximum number of parallel processes + + # Create a bash array of folders + GACODEcommand += "folders=(\n" + for folder in folders_red: + GACODEcommand += f' "{folder}"\n' + GACODEcommand += ")\n\n" + + # Loop over each folder and launch code, waiting if we've reached max_parallel_execution + GACODEcommand += "for folder in \"${folders[@]}\"; do\n" + GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = self.simulation_job.folderExecution)}\n' + GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" + GACODEcommand += "done\n\n" + GACODEcommand += "wait\n" + + # Standard job + elif type_of_submission == "slurm_standard": + + print(f"\t- {code.upper()} will be executed in SLURM as standard job (cpus: {total_cores_required})",typeMsg="i") + + # Code launches + GACODEcommand = "" + for folder in folders_red: + GACODEcommand += f' {code_call(folder = folder, n = cores_per_code_call, p = self.simulation_job.folderExecution)} &\n' + GACODEcommand += "\nwait" # This is needed so that the script doesn't end before each job + + # Job array + elif type_of_submission == "slurm_array": + + print(f"\t- {code.upper()} will be executed in SLURM as job array due to its size (cpus: {total_cores_required})",typeMsg="i") + + # As a pre-command, organize all folders in a simpler way + shellPreCommands = [] + shellPostCommands = [] + array_list = [] + for i, folder in enumerate(folders_red): + array_list.append(f"{i}") + folder_temp_array = f"run{i}" + folder_actual = folder + shellPreCommands.append(f"mkdir {self.simulation_job.folderExecution}/{folder_temp_array}; cp {self.simulation_job.folderExecution}/{folder_actual}/* {self.simulation_job.folderExecution}/{folder_temp_array}/.") + shellPostCommands.append(f"cp {self.simulation_job.folderExecution}/{folder_temp_array}/* {self.simulation_job.folderExecution}/{folder_actual}/.; rm -r {self.simulation_job.folderExecution}/{folder_temp_array}") + + # Code launches + indexed_folder = 'run"$SLURM_ARRAY_TASK_ID"' + GACODEcommand = code_call( + folder = indexed_folder, + n = cores_per_code_call, + p = self.simulation_job.folderExecution, + additional_command = f'1> {self.simulation_job.folderExecution}/{indexed_folder}/slurm_output.dat 2> {self.simulation_job.folderExecution}/{indexed_folder}/slurm_error.dat\n') + + # --------------------------------------------- + # Execute + # --------------------------------------------- + + slurm_settings = code_slurm_settings( + name=code, + minutes=minutes, + total_cores_required=total_cores_required, + cores_per_code_call=cores_per_code_call, + type_of_submission=type_of_submission, + array_list=array_list if type_of_submission == "slurm_array" else None + ) + + self.simulation_job.define_machine( + code, + f"mitim_{name}", + launchSlurm=launchSlurm, + slurm_settings=slurm_settings, + ) + + # I would like the mitim_job to check if the retrieved folders were complete + check_files_in_folder = {} + for folder in folders_red: + check_files_in_folder[folder] = filesToRetrieve + # --------------------------------------------- + + self.simulation_job.prep( + GACODEcommand, + input_folders=folders, + output_folders=folders_red, + check_files_in_folder=check_files_in_folder, + shellPreCommands=shellPreCommands, + shellPostCommands=shellPostCommands, + ) + + # Submit run and wait + if run_type == 'normal': + + self.simulation_job.run( + removeScratchFolders=True, + attempts_execution=attempts_execution + ) + + self._organize_results(code_executor, tmpFolder, filesToRetrieve) + + # Submit run but do not wait; the user should do checks and fetch results + elif run_type == 'submit': + self.simulation_job.run( + waitYN=False, + check_if_files_received=False, + removeScratchFolders=False, + removeScratchFolders_goingIn=kwargs_run.get("cold_start", False), + ) + + self.kwargs_organize = { + "code_executor": code_executor, + "tmpFolder": tmpFolder, + "filesToRetrieve": filesToRetrieve + } + + self.slurm_output = "slurm_output.dat" + + # Prepare how to search for the job without waiting for it + self.simulation_job.launchSlurm = True + self.simulation_job.slurm_settings['name'] = Path(self.simulation_job.folderExecution).name + + def check(self, every_n_minutes=None): + + if self.simulation_job.launchSlurm: + print("- Checker job status") + + while True: + self.simulation_job.check(file_output = self.slurm_output) + print(f'\t- Current status (as of {datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}): {self.simulation_job.status} ({self.simulation_job.infoSLURM["STATE"]})') + if self.simulation_job.status == 2: + print("\n\t* Job considered finished (please do .fetch() to retrieve results)",typeMsg="i") + break + elif every_n_minutes is None: + print("\n\t* Job not finished yet") + break + else: + print(f"\n\t* Waiting {every_n_minutes} minutes") + time.sleep(every_n_minutes * 60) + else: + print("- Not checking status because this was run command line (not slurm)") + + def fetch(self): + """ + For a job that has been submitted but not waited for, once it is done, get the results + """ + + print("\n\n\t- Fetching results") + + if self.simulation_job.launchSlurm: + self.simulation_job.connect() + self.simulation_job.retrieve() + self.simulation_job.close() + + self._organize_results(**self.kwargs_organize) + + else: + print("- Not retrieving results because this was run command line (not slurm)") + + def delete(self): + + print("\n\n\t- Deleting job") + + self.simulation_job.launchSlurm = False + + self.simulation_job.prep( + f"scancel -n {self.simulation_job.slurm_settings['name']}", + label_log_files="_finish", + ) + + self.simulation_job.run() + + def _organize_results(self, code_executor, tmpFolder, filesToRetrieve): + + # --------------------------------------------- + # Organize + # --------------------------------------------- + + print("\t- Retrieving files and changing names for storing") + fineall = True + for subfolder_sim in code_executor: + + for i, rho in enumerate(code_executor[subfolder_sim].keys()): + for file in filesToRetrieve: + original_file = f"{file}_{rho:.4f}" + final_destination = ( + code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" + ) + final_destination.unlink(missing_ok=True) + + temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" + temp_file.replace(final_destination) + + fineall = fineall and final_destination.exists() + + if not final_destination.exists(): + print(f"\t!! file {file} ({original_file}) could not be retrived",typeMsg="w",) + + if fineall: + print("\t\t- All files were successfully retrieved") + + # Remove temporary folder + shutil.rmtree(tmpFolder) + + else: + print("\t\t- Some files were not retrieved", typeMsg="w") + + + def run_scan( + self, + subfolder, # 'scan1', + multipliers={}, + minimum_delta_abs={}, + variable="RLTS_1", + varUpDown=[0.5, 1.0, 1.5], + variables_scanTogether=[], + relativeChanges=True, + **kwargs_run, + ): + + # ------------------------------------- + # Add baseline + # ------------------------------------- + if (1.0 not in varUpDown) and relativeChanges: + print("\n* Since variations vector did not include base case, I am adding it",typeMsg="i",) + varUpDown_new = [] + added = False + for i in varUpDown: + if i > 1.0 and not added: + varUpDown_new.append(1.0) + added = True + varUpDown_new.append(i) + else: + varUpDown_new = varUpDown + + + code_executor, code_executor_full, folders, varUpDown_new = self._prepare_scan( + subfolder, + multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, + variable=variable, + varUpDown=varUpDown_new, + variables_scanTogether=variables_scanTogether, + relativeChanges=relativeChanges, + **kwargs_run, + ) + + # Run them all + self._run( + code_executor, + code_executor_full=code_executor_full, + **kwargs_run, + ) + + # Read results + for cont_mult, mult in enumerate(varUpDown_new): + name = f"{variable}_{mult}" + self.read( + label=f"{self.subfolder_scan}_{name}", + folder=folders[cont_mult], + cold_startWF = False, + require_all_files=not kwargs_run.get("only_minimal_files",False), + ) + + return code_executor_full + + def _prepare_scan( + self, + subfolder, # 'scan1', + multipliers={}, + minimum_delta_abs={}, + variable="RLTS_1", + varUpDown=[0.5, 1.0, 1.5], + variables_scanTogether=[], + relativeChanges=True, + **kwargs_run, + ): + """ + Multipliers will be modified by adding the scaning variables, but I don't want to modify the original + multipliers, as they may be passed to the next scan + + Set relativeChanges=False if varUpDown contains the exact values to change, not multipleiers + """ + + completeVariation = self.run_specifications['complete_variation'] + + multipliers_mod = copy.deepcopy(multipliers) + + self.subfolder_scan = subfolder + + if relativeChanges: + for i in range(len(varUpDown)): + varUpDown[i] = round(varUpDown[i], 6) + + print(f"\n- Proceeding to scan {variable}{' together with '+', '.join(variables_scanTogether) if len(variables_scanTogether)>0 else ''}:") + + code_executor = {} + code_executor_full = {} + folders = [] + for cont_mult, mult in enumerate(varUpDown): + mult = round(mult, 6) + + if relativeChanges: + print(f"\n + Multiplier: {mult} -----------------------------------------------------------------------------------------------------------") + else: + print(f"\n + Value: {mult} ----------------------------------------------------------------------------------------------------------------") + + multipliers_mod[variable] = mult + + for variable_scanTogether in variables_scanTogether: + multipliers_mod[variable_scanTogether] = mult + + name = f"{variable}_{mult}" + + species = self.inputs_files[self.rhos[0]] # Any rho will do + + if completeVariation is not None: + multipliers_mod = completeVariation(multipliers_mod, species) + + if not relativeChanges: + for ikey in multipliers_mod: + kwargs_run["extraOptions"][ikey] = multipliers_mod[ikey] + multipliers_mod = {} + + # Force ensure quasineutrality if the + if variable in ["AS_3", "AS_4", "AS_5", "AS_6"]: + kwargs_run["Quasineutral"] = True + + # Only ask the cold_start in the first round + kwargs_run["forceIfcold_start"] = cont_mult > 0 or ("forceIfcold_start" in kwargs_run and kwargs_run["forceIfcold_start"]) + + code_executor, code_executor_full = self._run_prepare( + f"{self.subfolder_scan}_{name}", + code_executor=code_executor, + code_executor_full=code_executor_full, + multipliers=multipliers_mod, + minimum_delta_abs=minimum_delta_abs, + **kwargs_run, + ) + + folders.append(copy.deepcopy(self.FolderSimLast)) + + return code_executor, code_executor_full, folders, varUpDown + + def read( + self, + label="run1", + folder=None, # If None, search in the previously run folder + suffix=None, # If None, search with my standard _0.55 suffixes corresponding to rho of this TGLF class + **kwargs_to_class_output + ): + print("> Reading simulation results") + + class_output = [self.run_specifications['output_class'], self.run_specifications['output_store']] + + # If no specified folder, check the last one + if folder is None: + folder = self.FolderSimLast + + self.results[label] = { + class_output[1]:[], + 'parsed': [], + "x": np.array(self.rhos), + } + for rho in self.rhos: + + SIMout = class_output[0]( + folder, + suffix=f"_{rho:.4f}" if suffix is None else suffix, + **kwargs_to_class_output + ) + + # Unnormalize + if 'NormalizationSets' in self.__dict__: + SIMout.unnormalize( + self.NormalizationSets["SELECTED"], + rho=rho, + ) + else: + print("No normalization sets found.") + + self.results[label][class_output[1]].append(SIMout) + + self.results[label]['parsed'].append(buildDictFromInput(SIMout.inputFile) if SIMout.inputFile else None) + + def read_scan( + self, + label="scan1", + subfolder=None, + variable="RLTS_1", + positionIon=2, + variable_mapping=None, + variable_mapping_unn=None + ): + ''' + positionIon is the index in the input.tglf file... so if you want for ion RLNS_5, positionIon=5 + ''' + + if subfolder is None: + subfolder = self.subfolder_scan + + self.scans[label] = {} + self.scans[label]["variable"] = variable + self.scans[label]["positionBase"] = None + self.scans[label]["unnormalization_successful"] = True + self.scans[label]["results_tags"] = [] + + self.positionIon_scan = positionIon + + # ---- + + scan = {} + for ikey in variable_mapping | variable_mapping_unn: + scan[ikey] = [] + + cont = 0 + for ikey in self.results: + isThisTheRightReadResults = (subfolder in ikey) and (variable== "_".join(ikey.split("_")[:-1]).split(subfolder + "_")[-1]) + + if isThisTheRightReadResults: + + self.scans[label]["results_tags"].append(ikey) + + # Initialize lists + scan0 = {} + for ikey2 in variable_mapping | variable_mapping_unn: + scan0[ikey2] = [] + + # Loop over radii + for irho_cont in range(len(self.rhos)): + irho = np.where(self.results[ikey]["x"] == self.rhos[irho_cont])[0][0] + + for ikey2 in variable_mapping: + + obj = self.results[ikey][variable_mapping[ikey2][0]][irho] + if not hasattr(obj, '__dict__'): + obj_dict = obj + else: + obj_dict = obj.__dict__ + var0 = obj_dict[variable_mapping[ikey2][1]] + scan0[ikey2].append(var0 if variable_mapping[ikey2][2] is None else var0[variable_mapping[ikey2][2]]) + + # Unnormalized + self.scans[label]["unnormalization_successful"] = True + for ikey2 in variable_mapping_unn: + obj = self.results[ikey][variable_mapping_unn[ikey2][0]][irho] + if not hasattr(obj, '__dict__'): + obj_dict = obj + else: + obj_dict = obj.__dict__ + + if variable_mapping_unn[ikey2][1] not in obj_dict: + self.scans[label]["unnormalization_successful"] = False + break + var0 = obj_dict[variable_mapping_unn[ikey2][1]] + scan0[ikey2].append(var0 if variable_mapping_unn[ikey2][2] is None else var0[variable_mapping_unn[ikey2][2]]) + + for ikey2 in variable_mapping | variable_mapping_unn: + scan[ikey2].append(scan0[ikey2]) + + if float(ikey.split('_')[-1]) == 1.0: + self.scans[label]["positionBase"] = cont + cont += 1 + + self.scans[label]["x"] = np.array(self.rhos) + + for ikey2 in variable_mapping | variable_mapping_unn: + self.scans[label][ikey2] = np.atleast_2d(np.transpose(scan[ikey2])) + +def change_and_write_code( + rhos, + inputs0, + Folder_sim, + code_settings=None, + extraOptions={}, + multipliers={}, + minimum_delta_abs={}, + ApplyCorrections=True, + Quasineutral=False, + addControlFunction=None, + controls_file='input.tglf.controls', + **kwargs +): + """ + Received inputs classes and gives text. + ApplyCorrections refer to removing ions with too low density and that are fast species + """ + + inputs = copy.deepcopy(inputs0) + + mod_input_file = {} + ns_max = [] + for i, rho in enumerate(rhos): + print(f"\t- Changing input file for rho={rho:.4f}") + input_sim_rho = modifyInputs( + inputs[rho], + code_settings=code_settings, + extraOptions=extraOptions, + multipliers=multipliers, + minimum_delta_abs=minimum_delta_abs, + position_change=i, + addControlFunction=addControlFunction, + controls_file=controls_file, + NS=inputs[rho].num_recorded, + ) + + input_file = input_sim_rho.file.name.split('_')[0] + + newfile = Folder_sim / f"{input_file}_{rho:.4f}" + + if code_settings is not None: + # Apply corrections + if ApplyCorrections: + print("\t- Applying corrections") + input_sim_rho.removeLowDensitySpecie() + input_sim_rho.remove_fast() + + # Ensure that plasma to run is quasineutral + if Quasineutral: + input_sim_rho.ensureQuasineutrality() + else: + print('\t- Not applying corrections because settings is None') + + input_sim_rho.write_state(file=newfile) + + mod_input_file[rho] = input_sim_rho + + ns_max.append(inputs[rho].num_recorded) + + # Convert back to a string because that's how the run operates + inputFile = inputToVariable(Folder_sim, rhos, file=input_file) + + if (np.diff(ns_max) > 0).any(): + print("> Each radial location has its own number of species... probably because of removal of fast or low density...",typeMsg="w") + print("\t * Reading of simulation results will fail... consider doing something before launching run",typeMsg="q") + + return inputFile, mod_input_file + +def inputToVariable(folder, rhos, file='input.tglf'): + """ + Entire text file to variable + """ + + inputFilesTGLF = {} + for rho in rhos: + fileN = folder / f"{file}_{rho:.4f}" + + with open(fileN, "r") as f: + lines = f.readlines() + inputFilesTGLF[rho] = "".join(lines) + + return inputFilesTGLF + +def cold_start_checker( + rhos, + ResultsFiles, + Folder_sim, + cold_start=False, + print_each_time=False, +): + """ + This function checks if the TGLF inputs are already in the folder. If they are, it returns True + """ + cont_each = 0 + if cold_start: + rhosEvaluate = rhos + else: + rhosEvaluate = [] + for ir in rhos: + existsRho = True + for j in ResultsFiles: + ffi = Folder_sim / f"{j}_{ir:.4f}" + existsThis = ffi.exists() + existsRho = existsRho and existsThis + if not existsThis: + if print_each_time: + print(f"\t* {ffi} does not exist") + else: + cont_each += 1 + if not existsRho: + rhosEvaluate.append(ir) + + if not print_each_time and cont_each > 0: + print(f'\t* {cont_each} files from expected set are missing') + + if len(rhosEvaluate) < len(rhos) and len(rhosEvaluate) > 0: + print("~ Not all radii are found, but not removing folder and running only those that are needed",typeMsg="i",) + + return rhosEvaluate + +def modifyInputs( + input_class, + code_settings=None, + extraOptions={}, + multipliers={}, + minimum_delta_abs={}, + position_change=0, + addControlFunction=None, + controls_file = 'input.tglf.controls', + **kwargs_to_function, +): + + # Check that those are valid flags + GACODEdefaults.review_controls(extraOptions, control = controls_file) + GACODEdefaults.review_controls(multipliers, control = controls_file) + # ------------------------------------------- + + if code_settings is not None: + CodeOptions = addControlFunction(code_settings, **kwargs_to_function) + + # ~~~~~~~~~~ Change with presets + print(f" \t- Using presets code_settings = {code_settings}", typeMsg="i") + input_class.controls = CodeOptions + + else: + print("\t- Input file was not modified by code_settings, using what was there before",typeMsg="i") + + # Make all upper case + #extraOptions = {ikey.upper(): value for ikey, value in extraOptions.items()} + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Change with external options -> Input directly, not as multiplier + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + 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] + + try: + isspecie = ikey.split("_")[0] in input_class.species[1] + except: + isspecie = False + + # is a species parameter? + if isspecie: + specie = int(ikey.split("_")[-1]) + varK = "_".join(ikey.split("_")[:-1]) + var_orig = input_class.species[specie][varK] + var_new = value_to_change_to + input_class.species[specie][varK] = var_new + # is a another parameter? + else: + if ikey in input_class.controls: + var_orig = input_class.controls[ikey] + var_new = value_to_change_to + input_class.controls[ikey] = var_new + elif ikey in input_class.plasma: + var_orig = input_class.plasma[ikey] + var_new = value_to_change_to + input_class.plasma[ikey] = var_new + else: + # If the variable in extraOptions wasn't in there, consider it a control param + print(f"\t\t- Variable {ikey} to change did not exist previously, creating now",typeMsg="i") + var_orig = None + var_new = value_to_change_to + input_class.controls[ikey] = var_new + + print(f"\t\t- Changing {ikey} from {var_orig} to {var_new}",typeMsg="i",) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Change with multipliers -> Input directly, not as multiplier + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if len(multipliers) > 0: + print("\t\t- Variables change:") + for ikey in multipliers: + # is a specie one? + if "species" in input_class.__dict__.keys() and ikey.split("_")[0] in input_class.species[1]: + specie = int(ikey.split("_")[-1]) + varK = "_".join(ikey.split("_")[:-1]) + var_orig = input_class.species[specie][varK] + var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + input_class.species[specie][varK] = var_new + else: + if ikey in input_class.controls: + var_orig = input_class.controls[ikey] + var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + input_class.controls[ikey] = var_new + + elif ikey in input_class.plasma: + var_orig = input_class.plasma[ikey] + var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + input_class.plasma[ikey] = var_new + + else: + print("\t- Variable to scan did not exist in original file, add it as extraOptions first",typeMsg="w",) + + print(f"\t\t\t- Changing {ikey} from {var_orig} to {var_new} (x{multipliers[ikey]})") + + return input_class + +def multiplier_input(var_orig, multiplier, minimum_delta_abs = None): + + delta = var_orig * (multiplier - 1.0) + + if minimum_delta_abs is not None: + if (multiplier != 1.0) and abs(delta) < minimum_delta_abs: + print(f"\t\t\t- delta = {delta} is smaller than minimum_delta_abs = {minimum_delta_abs}, enforcing",typeMsg="i") + delta = np.sign(delta) * minimum_delta_abs + + return var_orig + delta + +def buildDictFromInput(inputFile): + parsed = {} + + lines = inputFile.split("\n") + for line in lines: + if "=" in line: + splits = [i.split()[0] for i in line.split("=")] + if ("." in splits[1]) and (splits[1][0].split()[0] != "."): + parsed[splits[0].split()[0]] = float(splits[1].split()[0]) + else: + try: + parsed[splits[0].split()[0]] = int(splits[1].split()[0]) + except: + parsed[splits[0].split()[0]] = splits[1].split()[0] + + for i in parsed: + if isinstance(parsed[i], str): + if ( + parsed[i].lower() == "t" + or parsed[i].lower() == "true" + or parsed[i].lower() == ".true." + ): + parsed[i] = True + elif ( + parsed[i].lower() == "f" + or parsed[i].lower() == "false" + or parsed[i].lower() == ".false." + ): + parsed[i] = False + + return parsed + +class GACODEoutput: + def __init__(self, *args, **kwargs): + self.inputFile = None + + def unnormalize(self, *args, **kwargs): + print("No unnormalization implemented.") + +class GACODEinput: + def __init__(self, file=None, controls_file=None, code='', n_species=None): + self.file = IOtools.expandPath(file) if isinstance(file, (str, Path)) else None + + self.controls_file = controls_file + self.code = code + self.n_species = n_species + + self.num_recorded = 100 + + if self.file is not None and self.file.exists(): + with open(self.file, "r") as f: + lines = f.readlines() + file_txt = "".join(lines) + else: + file_txt = "" + input_dict = buildDictFromInput(file_txt) + + self.process(input_dict) + + @classmethod + def initialize_in_memory(cls, input_dict): + instance = cls() + instance.process(input_dict) + return instance + + def process(self, input_dict): + + if self.controls_file is not None: + options_check = [key for key in IOtools.generateMITIMNamelist(self.controls_file, caseInsensitive=False).keys()] + else: + options_check = [] + + self.controls, self.plasma = {}, {} + for key in input_dict.keys(): + if key in options_check: + self.controls[key] = input_dict[key] + else: + self.plasma[key] = input_dict[key] + + # Get number of recorded species + if self.n_species is not None and self.n_species in input_dict: + self.num_recorded = int(input_dict[self.n_species]) + + def write_state(self, file=None): + + if file is None: + file = self.file + + # Local formatter: floats -> 6 significant figures in exponential (uppercase), + # ints stay as ints, bools as 0/1, sequences space-separated with same rule. + def _fmt_num(x): + import numpy as _np + if isinstance(x, (bool, _np.bool_)): + return "True" if x else "False" + if isinstance(x, (_np.floating, float)): + # 6 significant figures in exponential => 5 digits after decimal + return f"{float(x):.5E}" + if isinstance(x, (_np.integer, int)): + return f"{int(x)}" + return str(x) + + def _fmt_value(val): + import numpy as _np + if isinstance(val, (list, tuple, _np.ndarray)): + # Flatten numpy arrays but keep ordering; join with spaces + if isinstance(val, _np.ndarray): + flat = val.flatten().tolist() + else: + flat = list(val) + return " ".join(_fmt_num(v) for v in flat) + return _fmt_num(val) + + with open(file, "w") as f: + f.write("#-------------------------------------------------------------------------\n") + f.write(f"# {self.code} input file modified by MITIM {mitim_version}\n") + f.write("#-------------------------------------------------------------------------\n") + + f.write("\n\n# Control parameters\n") + f.write("# ------------------\n\n") + for ikey in self.controls: + var = self.controls[ikey] + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") + + f.write("\n\n# Plasma/Geometry parameters\n") + f.write("# ------------------\n\n") + for ikey in self.plasma: + var = self.plasma[ikey] + f.write(f"{ikey.ljust(23)} = {_fmt_value(var)}\n") + + def anticipate_problems(self): + pass + + def remove_fast(self): + pass + + def removeLowDensitySpecie(self, *args): + pass + \ No newline at end of file diff --git a/src/mitim_tools/gyrokinetics_tools/GXtools.py b/src/mitim_tools/simulation_tools/physics/GXtools.py similarity index 98% rename from src/mitim_tools/gyrokinetics_tools/GXtools.py rename to src/mitim_tools/simulation_tools/physics/GXtools.py index 009b1473..6c6bc225 100644 --- a/src/mitim_tools/gyrokinetics_tools/GXtools.py +++ b/src/mitim_tools/simulation_tools/physics/GXtools.py @@ -2,12 +2,13 @@ import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools, CONFIGread from mitim_tools.gacode_tools.utils import GACODErun, GACODEdefaults +from mitim_tools.simulation_tools import SIMtools from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __mitimroot__ from mitim_tools import __version__ as mitim_version from IPython import embed -class GX(GACODErun.gacode_simulation): +class GX(SIMtools.mitim_simulation): def __init__( self, rhos=[0.4, 0.6], # rho locations of interest @@ -189,7 +190,7 @@ def plot( plt.tight_layout() -class GXinput(GACODErun.GACODEinput): +class GXinput(SIMtools.GACODEinput): def __init__(self, file=None): super().__init__( file=file, @@ -356,7 +357,7 @@ def _fmt_value(val): return param_written -class GXoutput(GACODErun.GACODEoutput): +class GXoutput(SIMtools.GACODEoutput): def __init__(self, FolderGACODE, suffix="", **kwargs): super().__init__() diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index c96be6d0..8438bf42 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -49,24 +49,26 @@ # Scan of KY # --------------- -# cgyro.run_scan( -# 'scan1', -# extraOptions={ -# 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file -# }, -# variable='KY', -# varUpDown=[0.3,0.4], -# slurm_setup={ -# 'cores':4 -# }, -# cold_start=cold_start, -# forceIfcold_start=True, -# run_type='normal' -# ) - -# cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) - -# fig = cgyro.fn.add_figure(label="Quick linear") -# cgyro.plot_quick_linear(labels=["scan1_KY_0.3","scan1_KY_0.4"], fig = fig) +run_type = 'normal' + +cgyro.run_scan( + 'scan1', + extraOptions={ + 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file + }, + variable='KY', + varUpDown=[0.3,0.4], + slurm_setup={ + 'cores':4 + }, + cold_start=cold_start, + forceIfcold_start=True, + run_type=run_type + ) + +cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) + +fig = cgyro.fn.add_figure(label="Quick linear") +cgyro.plot_quick_linear(labels=["scan1_KY_0.3","scan1_KY_0.4"], fig = fig) cgyro.fn.show() diff --git a/tests/GX_workflow.py b/tests/GX_workflow.py index 76e66904..9c3f4c41 100644 --- a/tests/GX_workflow.py +++ b/tests/GX_workflow.py @@ -1,6 +1,6 @@ import os from mitim_tools.gacode_tools.PROFILEStools import gacode_state -from mitim_tools.gyrokinetics_tools import GXtools +from mitim_tools.simulation_tools.physics import GXtools from mitim_tools import __mitimroot__ cold_start = True From ab14f8b49f437e20f70c0db076c19244cd57366b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 26 Aug 2025 23:58:50 -0400 Subject: [PATCH 210/385] misc --- src/mitim_tools/simulation_tools/SIMtools.py | 10 ++++++++++ tests/CGYRO_workflow.py | 4 ++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 8790c695..f006b6f6 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -322,6 +322,8 @@ def _run( print(f"\t- {self.run_specifications['code'].upper()} not run because all results files found (please ensure consistency!)",typeMsg="i") + self.simulation_job = None + else: # ---------------------------------------------------------------------------------------------------------------- @@ -565,6 +567,10 @@ def _run( def check(self, every_n_minutes=None): + if self.simulation_job is None: + print("- Not checking status because simulation job is not defined (not run)", typeMsg="i") + return + if self.simulation_job.launchSlurm: print("- Checker job status") @@ -588,6 +594,10 @@ def fetch(self): For a job that has been submitted but not waited for, once it is done, get the results """ + if self.simulation_job is None: + print("- Not fetching because simulation job is not defined (not run)", typeMsg="i") + return + print("\n\n\t- Fetching results") if self.simulation_job.launchSlurm: diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index 8438bf42..72c649d2 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -30,7 +30,7 @@ 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file }, slurm_setup={ - 'cores':8, + 'cores':16, # Each CGYRO instance (each radius will have this number of cores or gpus) 'minutes': 10, }, cold_start=cold_start, @@ -59,7 +59,7 @@ variable='KY', varUpDown=[0.3,0.4], slurm_setup={ - 'cores':4 + 'cores':16 }, cold_start=cold_start, forceIfcold_start=True, From 6e94942caeb79294a0d0bf23c2a813dee63355ee Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 11:41:37 -0400 Subject: [PATCH 211/385] Transitioned to GB json files --- .../physics_models/transport_cgyroneo.py | 86 +++++++----- .../physics_models/transport_tglfneo.py | 90 ++++++------- .../physics_models/transport_tgyro.py | 60 ++++----- .../powertorch/utils/TARGETStools.py | 2 + .../powertorch/utils/TRANSPORTtools.py | 126 +++++++++++++----- 5 files changed, 220 insertions(+), 144 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index a63dd388..1e9555c6 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -69,12 +69,14 @@ def evaluate_turbulence(self): file_path = self.folder / 'fluxes_turb.json' attempts = 0 - all_good = False + all_good = post_checks(self) if file_path.exists() else False while (file_path.exists() is False) or (not all_good): if attempts > 0: - print(f"\n\n !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!", typeMsg='i') - print(f"\tMITIM could not find the file", typeMsg='i') - print(f" !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! \n\n", typeMsg='i') + print(f"\n !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!", typeMsg='i') + print(f"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!", typeMsg='i') + print(f" MITIM could not find the file... looping back", typeMsg='i') + print(f"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!", typeMsg='i') + print(f"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!", typeMsg='i') logic_to_wait(self.folder) attempts += 1 @@ -85,37 +87,38 @@ def pre_checks(self): plasma = self.powerstate.plasma - txt = "\nFluxes to be matched by CGYRO ( TARGETS - NEO ):" + txt = "\nFluxes to be matched by CGYRO ( Target - Neoclassical ):" # Print gradients for var, varn in zip( ["r/a ", "rho ", "a/LTe", "a/LTi", "a/Lne", "a/LnZ", "a/Lw0"], ["roa", "rho", "aLte", "aLti", "aLne", "aLnZ", "aLw0"], ): - txt += f"\n{var} = " + txt += f"\n{var} = " for j in range(plasma["rho"].shape[1] - 1): txt += f"{plasma[varn][0,j+1]:.6f} " # Print target fluxes for var, varn in zip( - ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ge (1E20m2/s)", "GZ (1E20m2/s)", "Mt (J/m^2) "], - ["QeMWm2", "QiMWm2", "Ge1E20m2", "GZ1E20m2", "MtJm2"], + ["Qe (GB)", "Qi (GB)", "Ge (GB)", "GZ (GB)", "Mt (GB)"], + ["QeGB", "QiGB", "GeGB", "GZGB", "MtGB"], ): txt += f"\n{var} = " for j in range(plasma["rho"].shape[1] - 1): - txt += f"{plasma[varn][0,j+1]-self.__dict__[f'{varn}_tr_neoc'][j]:.4e} " + txt += f"{plasma[varn][0,j+1]-self.__dict__[f'{varn}_neoc'][j]:.4e} " print(txt) def logic_to_wait(folder): - print(f"\n**** CGYRO prepared. Please, run CGYRO from the simulation setup in folder: ", typeMsg='i') + print(f"\n**** CGYRO prepared. Please, run CGYRO from the simulation setup in folder:\n", typeMsg='i') print(f"\t {folder}/base_cgyro\n", typeMsg='i') - print(f" **** When finished, the fluxes_turb.json file should be placed in:", typeMsg='i') + print(f"**** When finished, the fluxes_turb.json file should be placed in:\n", typeMsg='i') print(f"\t {folder}/fluxes_turb.json\n", typeMsg='i') - print(f" **** When you have done that, please write 'exit' and click enter (for continuing and reading that file)\n", typeMsg='i') + while not print(f"**** When you have done that, please say yes", typeMsg='q'): + pass -def post_checks(self, rtol = 1e-2): +def post_checks(self, rtol = 1e-3): with open(self.folder / 'fluxes_turb.json', 'r') as f: json_dict = json.load(f) @@ -135,35 +138,49 @@ def post_checks(self, rtol = 1e-2): print(f"\t {k} from POWERSTATE: {[round(i,4) for i in vP]}", typeMsg='i') if not np.allclose(v, vP, rtol=rtol): - all_good = print(f"{k} does not match with a relative tolerance of {rtol}:", typeMsg='q') + all_good = print(f"{k} does not match with a relative tolerance of {rtol*100.0:.2f}%:", typeMsg='q') return all_good -def write_json_CGYRO(roa, fluxes_mean, fluxes_stds, additional_info, file = 'fluxes_turb.json'): +def write_json_CGYRO(roa, fluxes_mean, fluxes_stds, additional_info = None, file = 'fluxes_turb.json'): ''' + ********************* Helper to write JSON - roa must be an array: [0.25, 0.35, ...] - fluxes_mean must be a dictionary with the fields and arrays: - 'QeMWm2': [0.1, 0.2, ...], - 'QiMWm2': ..., - 'Ge1E20m2': ..., - 'GZ1E20m2': ..., - 'MtJm2': ..., - 'QieMWm3': ... - same for fluxes_stds - additional_info must be a dictionary with any additional information to include in the JSON and compare to powerstate, - for example: - 'aLte': [0.2, 0.5, ...], - 'aLti': [0.3, 0.6, ...], - 'aLne': [0.3, 0.6, ...], - 'Qgb': [0.4, 0.7, ...] + ********************* + roa + Must be an array: [0.25, 0.35, ...] + fluxes_mean + Must be a dictionary with the fields and arrays: + 'QeMWm2': [0.1, 0.2, ...], + 'QiMWm2': ..., + 'Ge1E20m2': ..., + 'GZ1E20m2': ..., + 'MtJm2': ..., + 'QieMWm3': .. + or, alternatively (or complementary), in GB units: + 'QeGB': [0.1, 0.2, ...], + 'QiGB': ..., + 'GeGB': ..., + 'GZGB': ..., + 'MtGB': ..., + 'QieGB': .. + fluxes_stds + Exact same structure as fluxes_mean + additional_info + A dictionary with any additional information to include in the JSON and compare to powerstate, + for example (and recommended): + 'aLte': [0.2, 0.5, ...], + 'aLti': [0.3, 0.6, ...], + 'aLne': [0.3, 0.6, ...], + 'Qgb': [0.4, 0.7, ...], + 'rho': [0.2, 0.5, ...], ''' - - with open(file, 'w') as f: + if additional_info is None: + additional_info = {} - fluxes_mean = {} - fluxes_stds = {} + with open(file, 'w') as f: + additional_info_extended = additional_info | {'roa': roa.tolist()} json_dict = { @@ -173,4 +190,3 @@ def write_json_CGYRO(roa, fluxes_mean, fluxes_stds, additional_info, file = 'flu } json.dump(json_dict, f, indent=4) - diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 13070c55..15fd4289 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -84,16 +84,16 @@ def evaluate_turbulence(self): tglf.read(label='base',require_all_files=False) - Qe = np.array([tglf.results['base']['TGLFout'][i].Qe_unn for i in range(len(rho_locations))]) - Qi = np.array([tglf.results['base']['TGLFout'][i].Qi_unn for i in range(len(rho_locations))]) - Ge = np.array([tglf.results['base']['TGLFout'][i].Ge_unn for i in range(len(rho_locations))]) - GZ = np.array([tglf.results['base']['TGLFout'][i].GiAll_unn[impurityPosition] for i in range(len(rho_locations))]) - Mt = np.array([tglf.results['base']['TGLFout'][i].Mt_unn for i in range(len(rho_locations))]) - S = np.array([tglf.results['base']['TGLFout'][i].Se_unn for i in range(len(rho_locations))]) + Qe = np.array([tglf.results['base']['TGLFout'][i].Qe for i in range(len(rho_locations))]) + Qi = np.array([tglf.results['base']['TGLFout'][i].Qi for i in range(len(rho_locations))]) + Ge = np.array([tglf.results['base']['TGLFout'][i].Ge for i in range(len(rho_locations))]) + GZ = np.array([tglf.results['base']['TGLFout'][i].GiAll[impurityPosition] for i in range(len(rho_locations))]) + Mt = np.array([tglf.results['base']['TGLFout'][i].Mt for i in range(len(rho_locations))]) + S = np.array([tglf.results['base']['TGLFout'][i].Se for i in range(len(rho_locations))]) if Qi_includes_fast: - Qifast = [tglf.results['base']['TGLFout'][i].Qifast_unn for i in range(len(rho_locations))] + Qifast = [tglf.results['base']['TGLFout'][i].Qifast for i in range(len(rho_locations))] if Qifast.sum() != 0.0: print(f"\t- Qi includes fast ions, adding their contribution") @@ -129,27 +129,27 @@ def evaluate_turbulence(self): # Pass the information # ------------------------------------------------------------------------------------------------------------------------ - self.QeMWm2_tr_turb = Flux_mean[0] - self.QeMWm2_tr_turb_stds = Flux_std[0] + self.QeGB_turb = Flux_mean[0] + self.QeGB_turb_stds = Flux_std[0] - self.QiMWm2_tr_turb = Flux_mean[1] - self.QiMWm2_tr_turb_stds = Flux_std[1] + self.QiGB_turb = Flux_mean[1] + self.QiGB_turb_stds = Flux_std[1] - self.Ge1E20m2_tr_turb = Flux_mean[2] - self.Ge1E20m2_tr_turb_stds = Flux_std[2] + self.GeGB_turb = Flux_mean[2] + self.GeGB_turb_stds = Flux_std[2] - self.GZ1E20m2_tr_turb = Flux_mean[3] - self.GZ1E20m2_tr_turb_stds = Flux_std[3] + self.GZGB_turb = Flux_mean[3] + self.GZGB_turb_stds = Flux_std[3] - self.MtJm2_tr_turb = Flux_mean[4] - self.MtJm2_tr_turb_stds = Flux_std[4] + self.MtGB_turb = Flux_mean[4] + self.MtGB_turb_stds = Flux_std[4] if provideTurbulentExchange: - self.QieMWm3_tr_turb = Flux_mean[5] - self.QieMWm3_tr_turb_stds = Flux_std[5] + self.QieGB_turb = Flux_mean[5] + self.QieGB_turb_stds = Flux_std[5] else: - self.QieMWm3_tr_turb = Flux_mean[5] * 0.0 - self.QieMWm3_tr_turb_stds = Flux_std[5] * 0.0 + self.QieGB_turb = Flux_mean[5] * 0.0 + self.QieGB_turb_stds = Flux_std[5] * 0.0 return tglf @@ -184,30 +184,30 @@ def evaluate_neoclassical(self): neo.read(label='base') - Qe = np.array([neo.results['base']['NEOout'][i].Qe_unn for i in range(len(rho_locations))]) - Qi = np.array([neo.results['base']['NEOout'][i].Qi_unn for i in range(len(rho_locations))]) - Ge = np.array([neo.results['base']['NEOout'][i].Ge_unn for i in range(len(rho_locations))]) - GZ = np.array([neo.results['base']['NEOout'][i].GiAll_unn[impurityPosition-1] for i in range(len(rho_locations))]) - Mt = np.array([neo.results['base']['NEOout'][i].Mt_unn for i in range(len(rho_locations))]) + Qe = np.array([neo.results['base']['NEOout'][i].Qe for i in range(len(rho_locations))]) + Qi = np.array([neo.results['base']['NEOout'][i].Qi for i in range(len(rho_locations))]) + Ge = np.array([neo.results['base']['NEOout'][i].Ge for i in range(len(rho_locations))]) + GZ = np.array([neo.results['base']['NEOout'][i].GiAll[impurityPosition-1] for i in range(len(rho_locations))]) + Mt = np.array([neo.results['base']['NEOout'][i].Mt for i in range(len(rho_locations))]) # ------------------------------------------------------------------------------------------------------------------------ # Pass the information # ------------------------------------------------------------------------------------------------------------------------ - self.QeMWm2_tr_neoc = Qe - self.QiMWm2_tr_neoc = Qi - self.Ge1E20m2_tr_neoc = Ge - self.GZ1E20m2_tr_neoc = GZ - self.MtJm2_tr_neoc = Mt + self.QeGB_neoc = Qe + self.QiGB_neoc = Qi + self.GeGB_neoc = Ge + self.GZGB_neoc = GZ + self.MtGB_neoc = Mt - self.QeMWm2_tr_neoc_stds = abs(Qe) * percentError[1]/100.0 - self.QiMWm2_tr_neoc_stds = abs(Qi) * percentError[1]/100.0 - self.Ge1E20m2_tr_neoc_stds = abs(Ge) * percentError[1]/100.0 - self.GZ1E20m2_tr_neoc_stds = abs(GZ) * percentError[1]/100.0 - self.MtJm2_tr_neoc_stds = abs(Mt) * percentError[1]/100.0 + self.QeGB_neoc_stds = abs(Qe) * percentError[1]/100.0 + self.QiGB_neoc_stds = abs(Qi) * percentError[1]/100.0 + self.GeGB_neoc_stds = abs(Ge) * percentError[1]/100.0 + self.GZGB_neoc_stds = abs(GZ) * percentError[1]/100.0 + self.MtGB_neoc_stds = abs(Mt) * percentError[1]/100.0 - self.QieMWm3_tr_neoc = Qe * 0.0 - self.QieMWm3_tr_neoc_stds = Qe * 0.0 + self.QieGB_neoc = Qe * 0.0 + self.QieGB_neoc_stds = Qe * 0.0 return neo @@ -342,13 +342,13 @@ def _run_tglf_uncertainty_model( for vari in variables_to_scan: jump = tglf.scans[f'{name}_{vari}']['Qe'].shape[-1] - Qe[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qe'] - Qi[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qi'] - Qifast[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qifast'] - Ge[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Ge'] - GZ[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Gi'] - Mt[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Mt'] - S[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['S'] + Qe[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qe_gb'] + Qi[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qi_gb'] + Qifast[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qifast_gb'] + Ge[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Ge_gb'] + GZ[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Gi_gb'] + Mt[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Mt_gb'] + S[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['S_gb'] cont += jump if Qi_includes_fast: diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 0b15e484..bb705f4b 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -928,11 +928,11 @@ def tgyro_to_powerstate( # *********** Electron Energy Fluxes # ********************************** - self.QeMWm2_tr_turb = TGYROresults.Qe_sim_turb[0, 1:nr] - self.QeMWm2_tr_neoc = TGYROresults.Qe_sim_neo[0, 1:nr] + self.QeGB_turb = TGYROresults.QeGB_sim_turb[0, 1:nr] + self.QeGB_neoc = TGYROresults.QeGB_sim_neo[0, 1:nr] - self.QeMWm2_tr_turb_stds = TGYROresults.Qe_sim_turb_stds[0, 1:nr] - self.QeMWm2_tr_neoc_stds = TGYROresults.Qe_sim_neo_stds[0, 1:nr] + self.QeGB_turb_stds = TGYROresults.QeGB_sim_turb_stds[0, 1:nr] + self.QeGB_neoc_stds = TGYROresults.QeGB_sim_neo_stds[0, 1:nr] # ********************************** # *********** Ion Energy Fluxes @@ -940,65 +940,65 @@ def tgyro_to_powerstate( if Qi_includes_fast: - self.QiMWm2_tr_turb = TGYROresults.QiIons_sim_turb[0, 1:nr] - self.QiMWm2_tr_neoc = TGYROresults.QiIons_sim_neo[0, 1:nr] + self.QiGB_turb = TGYROresults.QiIons_sim_turb[0, 1:nr] + self.QiGB_neoc = TGYROresults.QiIons_sim_neo[0, 1:nr] - self.QiMWm2_tr_turb_stds = TGYROresults.QiIons_sim_turb_stds[0, 1:nr] - self.QiMWm2_tr_neoc_stds = TGYROresults.QiIons_sim_neo_stds[0, 1:nr] + self.QiGB_turb_stds = TGYROresults.QiGBIons_sim_turb_stds[0, 1:nr] + self.QiGB_neoc_stds = TGYROresults.QiGBIons_sim_neo_stds[0, 1:nr] else: - self.QiMWm2_tr_turb = TGYROresults.QiIons_sim_turb_thr[0, 1:nr] - self.QiMWm2_tr_neoc = TGYROresults.QiIons_sim_neo_thr[0, 1:nr] + self.QiGB_turb = TGYROresults.QiGBIons_sim_turb_thr[0, 1:nr] + self.QiGB_neoc = TGYROresults.QiGBIons_sim_neo_thr[0, 1:nr] - self.QiMWm2_tr_turb_stds = TGYROresults.QiIons_sim_turb_thr_stds[0, 1:nr] - self.QiMWm2_tr_neoc_stds = TGYROresults.QiIons_sim_neo_thr_stds[0, 1:nr] + self.QiGB_turb_stds = TGYROresults.QiGBIons_sim_turb_thr_stds[0, 1:nr] + self.QiGB_neoc_stds = TGYROresults.QiGBIons_sim_neo_thr_stds[0, 1:nr] # ********************************** # *********** Momentum Fluxes # ********************************** - self.MtJm2_tr_turb = TGYROresults.Mt_sim_turb[0, 1:nr] # So far, let's include fast in momentum - self.MtJm2_tr_neoc = TGYROresults.Mt_sim_neo[0, 1:nr] + self.MtGB_turb = TGYROresults.MtGB_sim_turb[0, 1:nr] # So far, let's include fast in momentum + self.MtGB_neoc = TGYROresults.MtGB_sim_neo[0, 1:nr] - self.MtJm2_tr_turb_stds = TGYROresults.Mt_sim_turb_stds[0, 1:nr] - self.MtJm2_tr_neoc_stds = TGYROresults.Mt_sim_neo_stds[0, 1:nr] + self.MtGB_turb_stds = TGYROresults.MtGB_sim_turb_stds[0, 1:nr] + self.MtGB_neoc_stds = TGYROresults.MtGB_sim_neo_stds[0, 1:nr] # ********************************** # *********** Particle Fluxes # ********************************** # Store raw fluxes for better plotting later - self.Ge1E20m2_tr_turb = TGYROresults.Ge_sim_turb[0, 1:nr] - self.Ge1E20m2_tr_neoc = TGYROresults.Ge_sim_neo[0, 1:nr] + self.GeGB_turb = TGYROresults.GeGB_sim_turb[0, 1:nr] + self.GeGB_neoc = TGYROresults.GeGB_sim_neo[0, 1:nr] - self.Ge1E20m2_tr_turb_stds = TGYROresults.Ge_sim_turb_stds[0, 1:nr] - self.Ge1E20m2_tr_neoc_stds = TGYROresults.Ge_sim_neo_stds[0, 1:nr] + self.GeGB_turb_stds = TGYROresults.GeGB_sim_turb_stds[0, 1:nr] + self.GeGB_neoc_stds = TGYROresults.GeGB_sim_neo_stds[0, 1:nr] # ********************************** # *********** Impurity Fluxes # ********************************** # Store raw fluxes for better plotting later - self.GZ1E20m2_tr_turb = TGYROresults.Gi_sim_turb[impurityPosition, 0, 1:nr] - self.GZ1E20m2_tr_neoc = TGYROresults.Gi_sim_neo[impurityPosition, 0, 1:nr] + self.GZGB_turb = TGYROresults.GiGB_sim_turb[impurityPosition, 0, 1:nr] + self.GZGB_neoc = TGYROresults.GiGB_sim_neo[impurityPosition, 0, 1:nr] - self.GZ1E20m2_tr_turb_stds = TGYROresults.Gi_sim_turb_stds[impurityPosition, 0, 1:nr] - self.GZ1E20m2_tr_neoc_stds = TGYROresults.Gi_sim_neo_stds[impurityPosition, 0, 1:nr] + self.GZGB_turb_stds = TGYROresults.GiGB_sim_turb_stds[impurityPosition, 0, 1:nr] + self.GZGB_neoc_stds = TGYROresults.GiGB_sim_neo_stds[impurityPosition, 0, 1:nr] # ********************************** # *********** Energy Exchange # ********************************** if provideTurbulentExchange: - self.QieMWm3_tr_turb = TGYROresults.EXe_sim_turb[0, 1:nr] - self.QieMWm3_tr_turb_stds = TGYROresults.EXe_sim_turb_stds[0, 1:nr] + self.QieGB_turb = TGYROresults.EXeGB_sim_turb[0, 1:nr] + self.QieGB_turb_stds = TGYROresults.EXeGB_sim_turb_stds[0, 1:nr] else: - self.QieMWm3_tr_turb = self.QeMWm2_tr_turb * 0.0 - self.QieMWm3_tr_turb_stds = self.QeMWm2_tr_turb * 0.0 + self.QieGB_turb = self.QeGB_turb * 0.0 + self.QieGB_turb_stds = self.QeGB_turb * 0.0 - self.QieMWm3_tr_neoc = self.QeMWm2_tr_turb * 0.0 - self.QieMWm3_tr_neoc_stds = self.QeMWm2_tr_turb_stds * 0.0 + self.QieGB_neoc = self.QeGB_turb * 0.0 + self.QieGB_neoc_stds = self.QeGB_turb_stds * 0.0 # ********************************** # *********** Targets diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 0322beec..133448d2 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -144,6 +144,8 @@ def postprocessing(self, forceZeroParticleFlux=False, relative_error_assumed=1.0 self.powerstate.plasma["QeGB"] = self.powerstate.plasma["QeMWm2"] / self.powerstate.plasma["Qgb"] self.powerstate.plasma["QiGB"] = self.powerstate.plasma["QiMWm2"] / self.powerstate.plasma["Qgb"] + self.powerstate.plasma["GeGB"] = self.powerstate.plasma["Ge1E20m2"] / self.powerstate.plasma["Ggb"] + self.powerstate.plasma["GZGB"] = self.powerstate.plasma["GZ1E20m2"] / self.powerstate.plasma["Ggb"] self.powerstate.plasma["CeGB"] = self.powerstate.plasma["Ce"] / self.powerstate.plasma["Qgb"] self.powerstate.plasma["CZGB"] = self.powerstate.plasma["CZ"] / self.powerstate.plasma["Qgb"] self.powerstate.plasma["MtGB"] = self.powerstate.plasma["MtJm2"] / self.powerstate.plasma["Pgb"] diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index f8faae26..82d8b52c 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -19,21 +19,21 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): 'fluxes_mean': { - 'QeMWm2': ... - 'QiMWm2': ... - 'Ge1E20m2': ... - 'GZ1E20m2': ... - 'MtJm2': ... - 'QieMWm3': ... + 'QeGB': ... + 'QiGB': ... + 'GeGB': ... + 'GZGB': ... + 'MtGB': ... + 'QieGB': ... }, 'fluxes_stds': { - 'QeMWm2': ... - 'QiMWm2': ... - 'Ge1E20m2': ... - 'GZ1E20m2': ... - 'MtJm2': ... - 'QieMWm3': ... + 'QeGB': ... + 'QiGB': ... + 'GeGB': ... + 'GZGB': ... + 'MtGB': ... + 'QieGB': ... }, 'additional_info': { 'rho': rho.tolist(), @@ -46,9 +46,9 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): fluxes_mean = {} fluxes_stds = {} - for var in ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3']: - fluxes_mean[var] = self.__dict__[f"{var}_tr_{suffix}"].tolist() - fluxes_stds[var] = self.__dict__[f"{var}_tr_{suffix}_stds"].tolist() + for var in ['QeGB', 'QiGB', 'GeGB', 'GZGB', 'MtGB', 'QieGB']: + fluxes_mean[var] = self.__dict__[f"{var}_{suffix}"].tolist() + fluxes_stds[var] = self.__dict__[f"{var}_{suffix}_stds"].tolist() json_dict = { 'fluxes_mean': fluxes_mean, @@ -56,6 +56,10 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): 'additional_info': { 'rho': self.powerstate.plasma["rho"][0, 1:].cpu().numpy().tolist(), 'roa': self.powerstate.plasma["roa"][0, 1:].cpu().numpy().tolist(), + 'Qgb': self.powerstate.plasma["Qgb"][0, 1:].cpu().numpy().tolist(), + 'aLte': self.powerstate.plasma["aLte"][0, 1:].cpu().numpy().tolist(), + 'aLti': self.powerstate.plasma["aLti"][0, 1:].cpu().numpy().tolist(), + 'aLne': self.powerstate.plasma["aLne"][0, 1:].cpu().numpy().tolist(), } } @@ -254,16 +258,69 @@ def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): ''' Populate the powerstate.plasma with the results from the json file ''' + + print(f"\t* Populating powerstate.plasma with JSON data from {self.folder / file_name}") with open(self.folder / file_name, 'r') as f: json_dict = json.load(f) + + # See if the file has GB or real units + units_GB, units_real = False, False + if 'QeGB' in json_dict['fluxes_mean']: + units_GB = True + if 'QeMWm2' in json_dict['fluxes_mean']: + units_real = True + + units = 'both' if (units_GB and units_real) else 'GB' if units_GB else 'real' if units_real else 'none' + + if units == 'real': + + print("\t\t- File has fluxes in real units... populating powerstate directly") + + for var in ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3']: + self.powerstate.plasma[f"{var}_tr_{suffix}"] = np.array(json_dict['fluxes_mean'][var]) + self.powerstate.plasma[f"{var}_tr_{suffix}_stds"] = np.array(json_dict['fluxes_stds'][var]) + + elif units == 'GB' or units == 'both': + + mapper = { + 'QeGB': ['Qgb', 'QeMWm2'], + 'QiGB': ['Qgb', 'QiMWm2'], + 'GeGB': ['Ggb', 'Ge1E20m2'], + 'GZGB': ['Ggb', 'GZ1E20m2'], + 'MtGB': ['Pgb', 'MtJm2'], + 'QieGB': ['Sgb', 'QieMWm3'], + } - for var in ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3']: - self.powerstate.plasma[f"{var}_tr_{suffix}"] = np.array(json_dict['fluxes_mean'][var]) - self.powerstate.plasma[f"{var}_tr_{suffix}_stds"] = np.array(json_dict['fluxes_stds'][var]) + dum = {} + for var in mapper: + gb = self.powerstate.plasma[f"{mapper[var][0]}"][0,1:].cpu().numpy() + dum[f"{mapper[var][1]}_tr_{suffix}"] = np.array(json_dict['fluxes_mean'][var]) * gb + dum[f"{mapper[var][1]}_tr_{suffix}_stds"] = np.array(json_dict['fluxes_stds'][var]) * gb + + if units == 'GB': + + print("\t\t- File has fluxes in GB units... using GB units from powerstate to convert to real units") + + for var in mapper: + self.powerstate.plasma[f"{mapper[var][1]}_tr_{suffix}"] = dum[f"{mapper[var][1]}_tr_{suffix}"] + self.powerstate.plasma[f"{mapper[var][1]}_tr_{suffix}_stds"] = dum[f"{mapper[var][1]}_tr_{suffix}_stds"] + + elif units == 'both': + + print("\t\t- File has fluxes in both GB and real units... using real units and checking consistency") + + for var in mapper: + if not np.allclose(self.powerstate.plasma[f"{mapper[var][1]}_tr_{suffix}"], dum[f"{mapper[var][1]}_tr_{suffix}"]): + print(f"\t\t\t- Inconsistent values found for {mapper[var][1]}_tr_{suffix}") + + for var in ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3']: + self.powerstate.plasma[f"{var}_tr_{suffix}"] = np.array(json_dict['fluxes_mean'][var]) + self.powerstate.plasma[f"{var}_tr_{suffix}_stds"] = np.array(json_dict['fluxes_stds'][var]) + + else: + raise ValueError("[MITIM] Unknown units in JSON file") - print(f"\t* Populated powerstate.plasma with JSON data from {self.folder / file_name}") - # ---------------------------------------------------------------------------------------------------- # EVALUATE (custom part) # ---------------------------------------------------------------------------------------------------- @@ -271,13 +328,13 @@ def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): def evaluate_turbulence(self): ''' This needs to populate the following np.arrays in self., with dimensions of rho: - - QeMWm2_tr_turb - - QiMWm2_tr_turb - - Ge1E20m2_tr_turb - - GZ1E20m2_tr_turb - - MtJm2_tr_turb - - QieMWm3_tr_turb (turbulence exchange) - and their respective standard deviations, e.g. QeMWm2_tr_turb_stds + - QeGB_turb + - QiGB_turb + - GeGB_turb + - GZGB_turb + - MtGB_turb + - QieGB_turb (turbulence exchange) + and their respective standard deviations, e.g. QeGB_turb_stds ''' print(">> No turbulent fluxes to evaluate", typeMsg="w") @@ -300,12 +357,12 @@ def evaluate_turbulence(self): def evaluate_neoclassical(self): ''' This needs to populate the following np.arrays in self.: - - QeMWm2_tr_neoc - - QiMWm2_tr_neoc - - Ge1E20m2_tr_neoc - - GZ1E20m2_tr_neoc - - MtJm2_tr_neoc - and their respective standard deviations, e.g. QeMWm2_tr_neoc_stds + - QeGB_neoc + - QiGB_neoc + - GeGB_tr_neoc + - GZGB_tr_neoc + - MtGB_tr_neoc + and their respective standard deviations, e.g. QeGB_tr_neoc_stds ''' print(">> No neoclassical fluxes to evaluate", typeMsg="w") @@ -321,4 +378,5 @@ def evaluate_neoclassical(self): ]: self.__dict__[f"{var}_tr_neoc"] = np.zeros(dim) - self.__dict__[f"{var}_tr_neoc_stds"] = np.zeros(dim) \ No newline at end of file + self.__dict__[f"{var}_tr_neoc_stds"] = np.zeros(dim) + \ No newline at end of file From 4074d9b0e3e5c134f2b51614ab04747069c2d336 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 11:42:11 -0400 Subject: [PATCH 212/385] cpus per task allocations --- src/mitim_tools/gacode_tools/CGYROtools.py | 2 +- src/mitim_tools/gacode_tools/TGLFtools.py | 2 +- src/mitim_tools/gacode_tools/TGYROtools.py | 5 +++-- src/mitim_tools/simulation_tools/SIMtools.py | 10 +++++----- 4 files changed, 10 insertions(+), 9 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 998b72df..8ce216c1 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -32,7 +32,7 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call if type_of_submission == "slurm_standard": - slurm_settings['ntasks'] = total_cores_required + slurm_settings['ntasks'] = total_cores_required // cores_per_code_call if machineSettings['gpus_per_node'] > 0: slurm_settings['gpuspertask'] = cores_per_code_call diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 4d72c604..06af82fe 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -138,7 +138,7 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call if type_of_submission == "slurm_standard": - slurm_settings['ntasks'] = total_cores_required + slurm_settings['ntasks'] = total_cores_required // cores_per_code_call slurm_settings['cpuspertask'] = cores_per_code_call elif type_of_submission == "slurm_array": diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index 48a756b2..a53813d6 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -10,6 +10,7 @@ ) from mitim_tools.gacode_tools import TGLFtools from mitim_tools.gacode_tools.utils import GACODEinterpret, GACODEdefaults, GACODErun +from mitim_tools.simulation_tools import SIMtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed @@ -367,7 +368,7 @@ def run( f"\t\t- Creating only-controls input.tglf file in {IOtools.clipstr(str(self.FolderTGYRO_tmp.resolve()))}input.tglf" ) inputclass_TGLF = TGLFtools.TGLFinput() - inputclass_TGLF = GACODErun.modifyInputs( + inputclass_TGLF = SIMtools.modifyInputs( inputclass_TGLF, code_settings=TGLFsettings, extraOptions=extraOptionsTGLF, @@ -4547,7 +4548,7 @@ def __init__(self, input_profiles, file=None, onlyThermal=False, limitSpecies=10 else: self.file_txt = "" - self.input_dict = GACODErun.buildDictFromInput(self.file_txt) + self.input_dict = SIMtools.buildDictFromInput(self.file_txt) # Species self.species = input_profiles.Species diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index f006b6f6..2b71a030 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -35,16 +35,16 @@ def __init__( def prep( self, - mitim_state, # A MITIM state class - FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) - cold_start=False, # If True, do not use what it potentially inside the folder, run again - forceIfcold_start=False, # Extra flag + mitim_state, # A MITIM state class + FolderGACODE, # Main folder where all caculations happen (runs will be in subfolders) + cold_start=False, # If True, do not use what it potentially inside the folder, run again + forceIfcold_start=False, # Extra flag ): ''' This method prepares the GACODE run from a MITIM state class by setting up the necessary input files and directories. ''' - print("> Preparation run from input.gacode (direct conversion)") + print("> Preparation run from MITIM state class (direct conversion)") if self.run_specifications is None: raise Exception("[MITIM] Simulation child class did not define run specifications") From 5ef8243a46e592f92fc0ba9b35efe461fecc4bcb Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 12:00:07 -0400 Subject: [PATCH 213/385] Corrected cores --- .../powertorch/physics_models/transport_cgyroneo.py | 1 - src/mitim_tools/simulation_tools/SIMtools.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index 1e9555c6..90828c13 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -117,7 +117,6 @@ def logic_to_wait(folder): while not print(f"**** When you have done that, please say yes", typeMsg='q'): pass - def post_checks(self, rtol = 1e-3): with open(self.folder / 'fluxes_turb.json', 'r') as f: diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 2b71a030..0155449d 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -459,7 +459,7 @@ def _run( # Loop over each folder and launch code, waiting if we've reached max_parallel_execution GACODEcommand += "for folder in \"${folders[@]}\"; do\n" - GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = self.simulation_job.folderExecution)}\n' + GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = self.simulation_job.folderExecution)} &\n' GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" GACODEcommand += "done\n\n" GACODEcommand += "wait\n" From 14d4eb1a0b039fdae7e7bfe747f112674e6ef899 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 13:14:54 -0400 Subject: [PATCH 214/385] changed from model_complete to transport_simulation_folder --- src/mitim_modules/portals/PORTALSmain.py | 2 +- src/mitim_modules/portals/utils/PORTALSanalysis.py | 4 ++-- src/mitim_modules/portals/utils/PORTALSoptimization.py | 6 +++--- src/mitim_modules/powertorch/STATEtools.py | 4 ++-- .../powertorch/physics_models/transport_tgyro.py | 2 +- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index f0e56c76..a498c728 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -601,7 +601,7 @@ def runModelEvaluator( # Prep run # --------------------------------------------------------------------------------------------------- - folder_model = FolderEvaluation / "model_complete" + folder_model = FolderEvaluation / "transport_simulation_folder" folder_model.mkdir(parents=True, exist_ok=True) # --------------------------------------------------------------------------------------------------- diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 4039805d..a875ff6f 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -179,12 +179,12 @@ def prep_metrics(self, ilast=None): self.profiles_next = None x_train_num = self.step.train_X.shape[0] - file = self.opt_fun.folder / "Execution" / f"Evaluation.{x_train_num}" / "model_complete" / "input.gacode_unmodified" + file = self.opt_fun.folder / "Execution" / f"Evaluation.{x_train_num}" / "transport_simulation_folder" / "input.gacode_unmodified" if file.exists(): print("\t\t- Reading next profile to evaluate (from folder)") self.profiles_next = PROFILEStools.gacode_state(file, derive_quantities=False) - file = self.opt_fun.folder / "Execution" / f"Evaluation.{x_train_num}" / "model_complete" / "input.gacode.new" + file = self.opt_fun.folder / "Execution" / f"Evaluation.{x_train_num}" / "transport_simulation_folder" / "input.gacode.new" if file.exists(): self.profiles_next_new = PROFILEStools.gacode_state( file, derive_quantities=False diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index af6065f6..15384cac 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -78,10 +78,10 @@ def initialization_simple_relax(self): newname = f"{namingConvention}_{i}" # Delte destination first - if (ff / "model_complete").exists(): - IOtools.shutil_rmtree(ff / "model_complete") + if (ff / "transport_simulation_folder").exists(): + IOtools.shutil_rmtree(ff / "transport_simulation_folder") - shutil.copytree(MainFolder / newname / "model_complete", ff / "model_complete") #### delete first + shutil.copytree(MainFolder / newname / "transport_simulation_folder", ff / "transport_simulation_folder") #### delete first return Xopt.cpu().numpy() diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index e8dccaa9..0bf96325 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -367,13 +367,13 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): if folder_main is not None: folder = IOtools.expandPath(folder_main) / f"{namingConvention}_{cont}" if issubclass(self.transport_options["transport_evaluator"], TRANSPORTtools.power_transport): - (folder / "model_complete").mkdir(parents=True, exist_ok=True) + (folder / "transport_simulation_folder").mkdir(parents=True, exist_ok=True) # *************************************************************************************************************** # Calculate # *************************************************************************************************************** - folder_run = folder / "model_complete" if folder_main is not None else IOtools.expandPath('~/scratch/') + folder_run = folder / "transport_simulation_folder" if folder_main is not None else IOtools.expandPath('~/scratch/') QTransport, QTarget, _, _ = self.calculate(X, nameRun=nameRun, folder=folder_run, evaluation_number=cont) cont += 1 diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index bb705f4b..5ac4e2d4 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -554,7 +554,7 @@ def dummyCDF(GeneralFolder, FolderEvaluation): This routine creates path to a dummy CDF file in FolderEvaluation, with the name "simulation_evaluation.CDF" GeneralFolder, e.g. ~/runs_portals/run10/ - FolderEvaluation, e.g. ~/runs_portals/run10000/Execution/Evaluation.0/model_complete/ + FolderEvaluation, e.g. ~/runs_portals/run10000/Execution/Evaluation.0/transport_simulation_folder/ """ # ------- Name construction for scratch folders in parallel ---------------- From 5ca15b6de51e54eac5eb6094c4bd2582a9d5e56f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 15:20:10 -0400 Subject: [PATCH 215/385] SMARTS relevant: Always run base TGLF even if CGYRO selected --- .../physics_models/transport_cgyroneo.py | 14 +++++--------- .../powertorch/physics_models/transport_tglfneo.py | 7 ++++++- templates/input.cgyro.controls | 13 +------------ 3 files changed, 12 insertions(+), 22 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index 90828c13..5be29364 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -12,6 +12,11 @@ def __init__(self, powerstate, **kwargs): # Do not hook here def evaluate_turbulence(self): + + # Run base TGLF always, to keep track of discrepancies! -------------------------------------- + self.powerstate.transport_options["transport_evaluator_options"]["use_tglf_scan_trick"] = None + self._evaluate_tglf() + # -------------------------------------------------------------------------------------------- rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] @@ -34,15 +39,6 @@ def evaluate_turbulence(self): self.folder, ) - cgyro = CGYROtools.CGYRO( - rhos = rho_locations - ) - - cgyro.prep( - self.powerstate.profiles_transport.files[0], - self.folder, - ) - if run_type in ['normal', 'submit']: raise Exception("[MITIM] Automatic submission or full run not implemented") diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 15fd4289..92a3a6b3 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -18,9 +18,14 @@ def produce_profiles(self): # ************************************************************************************ # Private functions for the evaluation # ************************************************************************************ + @IOtools.hook_method(after=partial(TRANSPORTtools.write_json, file_name = 'fluxes_turb.json', suffix= 'turb')) - def evaluate_turbulence(self): + def evaluate_turbulence(self): + self._evaluate_tglf() + # Have it separate such that I can call it from the CGYRO class but not with the decorator + def _evaluate_tglf(self): + # ------------------------------------------------------------------------------------------------------------------------ # Grab options from powerstate # ------------------------------------------------------------------------------------------------------------------------ diff --git a/templates/input.cgyro.controls b/templates/input.cgyro.controls index a1f0bf0a..47306e1f 100644 --- a/templates/input.cgyro.controls +++ b/templates/input.cgyro.controls @@ -1,15 +1,4 @@ -#============================================================= -#CGYRO ion-scale input file -#============================================================= -# -#Simulation notes: ITER Baseline Scenario ~134 x 118 rho_s_D box ; -#0.053 < k_theta rhos_D <1.2. Focused on r/a=0.55 -#ExB shear on and new profiles set at 430 acs -#Startup phase starts with ExB shear on. Profiles adopted from -#Holland JPP 2023 and Mantica PPCF 2020 -#------------------------------------------------------------- - -#============================ +#======================== #Basic Simulation Parameters #============================ From 884931a1908bde33bd690ad6db1c98b472c03e3f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 15:25:04 -0400 Subject: [PATCH 216/385] SMARTS relevant: Added PORTALSparameters['only_minimal_files_TGLF'] --- src/mitim_modules/portals/PORTALSmain.py | 1 + src/mitim_modules/portals/utils/PORTALSinit.py | 1 + .../powertorch/physics_models/transport_tglfneo.py | 4 +++- 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index a498c728..2e8e86b9 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -222,6 +222,7 @@ def __init__( "additional_params_in_surrogate": additional_params_in_surrogate, "use_tglf_scan_trick": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta "keep_full_model_folder": True, # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) + "only_minimal_files_TGLF": True, # If True, only retrieve the minimal set of files from TGLF runs, to avoid large folders "cores_per_tglf_instance": 1, # Number of cores to use per TGLF instance } diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 4cf5a56c..8eae7c16 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -137,6 +137,7 @@ def initializeProblem( "forceZeroParticleFlux": portals_fun.PORTALSparameters["forceZeroParticleFlux"], "percentError": portals_fun.PORTALSparameters["percentError"], "use_tglf_scan_trick": portals_fun.PORTALSparameters["use_tglf_scan_trick"], + "only_minimal_files_TGLF": portals_fun.PORTALSparameters["only_minimal_files_TGLF"], "cold_start": False, "MODELparameters": portals_fun.MODELparameters, "impurityPosition": position_of_impurity, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 92a3a6b3..87b545c6 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -43,6 +43,8 @@ def _evaluate_tglf(self): use_tglf_scan_trick = transport_evaluator_options.get("use_tglf_scan_trick", None) cores_per_tglf_instance = transport_evaluator_options.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) + only_minimal_files_TGLF = transport_evaluator_options.get("only_minimal_files_TGLF", True) + # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) impurityPosition = self.powerstate.impurityPosition_transport @@ -84,7 +86,7 @@ def _evaluate_tglf(self): "minutes": 2, }, attempts_execution=2, - only_minimal_files=True, + only_minimal_files=only_minimal_files_TGLF, ) tglf.read(label='base',require_all_files=False) From 617fadbf285614560e566ca6db47121dacbaa3b5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 15:32:00 -0400 Subject: [PATCH 217/385] Instead of only_minimal_files_TGLF, use "keep_tglf_files": "base", # minimal: only retrieve minimal files always; base: retrieve all files when running TGLF base; none: always minimal files --- src/mitim_modules/portals/PORTALSmain.py | 2 +- src/mitim_modules/portals/utils/PORTALSinit.py | 2 +- .../powertorch/physics_models/transport_tglfneo.py | 8 +++++--- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 2e8e86b9..cb6bdf53 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -222,7 +222,7 @@ def __init__( "additional_params_in_surrogate": additional_params_in_surrogate, "use_tglf_scan_trick": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta "keep_full_model_folder": True, # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) - "only_minimal_files_TGLF": True, # If True, only retrieve the minimal set of files from TGLF runs, to avoid large folders + "keep_tglf_files": "base", # minimal: only retrieve minimal files always; base: retrieve all files when running TGLF base; none: always minimal files "cores_per_tglf_instance": 1, # Number of cores to use per TGLF instance } diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 8eae7c16..e1cc00d0 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -137,7 +137,7 @@ def initializeProblem( "forceZeroParticleFlux": portals_fun.PORTALSparameters["forceZeroParticleFlux"], "percentError": portals_fun.PORTALSparameters["percentError"], "use_tglf_scan_trick": portals_fun.PORTALSparameters["use_tglf_scan_trick"], - "only_minimal_files_TGLF": portals_fun.PORTALSparameters["only_minimal_files_TGLF"], + "keep_tglf_files": portals_fun.PORTALSparameters["keep_tglf_files"], "cold_start": False, "MODELparameters": portals_fun.MODELparameters, "impurityPosition": position_of_impurity, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 87b545c6..e43652bd 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -43,7 +43,7 @@ def _evaluate_tglf(self): use_tglf_scan_trick = transport_evaluator_options.get("use_tglf_scan_trick", None) cores_per_tglf_instance = transport_evaluator_options.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) - only_minimal_files_TGLF = transport_evaluator_options.get("only_minimal_files_TGLF", True) + keep_tglf_files = transport_evaluator_options.get("keep_tglf_files", "minimal") # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) impurityPosition = self.powerstate.impurityPosition_transport @@ -86,7 +86,7 @@ def _evaluate_tglf(self): "minutes": 2, }, attempts_execution=2, - only_minimal_files=only_minimal_files_TGLF, + only_minimal_files=keep_tglf_files in ['minimal'] ) tglf.read(label='base',require_all_files=False) @@ -128,6 +128,7 @@ def _evaluate_tglf(self): cores_per_tglf_instance=cores_per_tglf_instance, launchMODELviaSlurm=launchMODELviaSlurm, Qi_includes_fast=Qi_includes_fast, + only_minimal_files=keep_tglf_files in ['minimal', 'base'] ) self._raise_warnings(tglf, rho_locations, Qi_includes_fast) @@ -268,6 +269,7 @@ def _run_tglf_uncertainty_model( cores_per_tglf_instance = 4, # e.g. 4 core per radius, since this is going to launch ~ Nr=5 x (Nv=6 x Nd=2 + 1) = 65 TGLFs at once launchMODELviaSlurm=False, Qi_includes_fast=False, + only_minimal_files=True, # Since I only care about fluxes here, do not retrieve all the files ): print(f"\t- Running TGLF standalone scans ({delta = }) to determine relative errors") @@ -329,7 +331,7 @@ def _run_tglf_uncertainty_model( extra_name = f'{extra_name}_{name}', positionIon=impurityPosition+2, attempts_execution=2, - only_minimal_files=True, # Since I only care about fluxes here, do not retrieve all the files + only_minimal_files=only_minimal_files, launchSlurm=launchMODELviaSlurm, ) From 33f2daac19390418d1a72565d5eec719358f08cd Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 15:43:07 -0400 Subject: [PATCH 218/385] SMARTS relevant: TGLF base case separated from the drives scan --- .../physics_models/transport_tglfneo.py | 99 +++++++++++-------- 1 file changed, 56 insertions(+), 43 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index e43652bd..031c5871 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -66,48 +66,52 @@ def _evaluate_tglf(self): # Run TGLF # ------------------------------------------------------------------------------------------------------------------------ - if use_tglf_scan_trick is None: + # Run base + tglf.run( + 'base_tglf', + code_settings=TGLFsettings, + extraOptions=extraOptions, + ApplyCorrections=False, + launchSlurm= launchMODELviaSlurm, + cold_start= cold_start, + forceIfcold_start=True, + extra_name= self.name, + slurm_setup={ + "cores": cores_per_tglf_instance, + "minutes": 2, + }, + attempts_execution=2, + only_minimal_files=keep_tglf_files in ['minimal'] + ) + + tglf.read(label='base',require_all_files=False) + + # Grab values + Qe = np.array([tglf.results['base']['TGLFout'][i].Qe for i in range(len(rho_locations))]) + Qi = np.array([tglf.results['base']['TGLFout'][i].Qi for i in range(len(rho_locations))]) + Ge = np.array([tglf.results['base']['TGLFout'][i].Ge for i in range(len(rho_locations))]) + GZ = np.array([tglf.results['base']['TGLFout'][i].GiAll[impurityPosition] for i in range(len(rho_locations))]) + Mt = np.array([tglf.results['base']['TGLFout'][i].Mt for i in range(len(rho_locations))]) + S = np.array([tglf.results['base']['TGLFout'][i].Se for i in range(len(rho_locations))]) + + if Qi_includes_fast: - # ******************************************************************* - # Just run TGLF once and apply an ad-hoc percent error to the results - # ******************************************************************* + Qifast = [tglf.results['base']['TGLFout'][i].Qifast for i in range(len(rho_locations))] - tglf.run( - 'base_tglf', - code_settings=TGLFsettings, - extraOptions=extraOptions, - ApplyCorrections=False, - launchSlurm= launchMODELviaSlurm, - cold_start= cold_start, - forceIfcold_start=True, - extra_name= self.name, - slurm_setup={ - "cores": cores_per_tglf_instance, - "minutes": 2, - }, - attempts_execution=2, - only_minimal_files=keep_tglf_files in ['minimal'] - ) - - tglf.read(label='base',require_all_files=False) + if Qifast.sum() != 0.0: + print(f"\t- Qi includes fast ions, adding their contribution") + Qi += Qifast - Qe = np.array([tglf.results['base']['TGLFout'][i].Qe for i in range(len(rho_locations))]) - Qi = np.array([tglf.results['base']['TGLFout'][i].Qi for i in range(len(rho_locations))]) - Ge = np.array([tglf.results['base']['TGLFout'][i].Ge for i in range(len(rho_locations))]) - GZ = np.array([tglf.results['base']['TGLFout'][i].GiAll[impurityPosition] for i in range(len(rho_locations))]) - Mt = np.array([tglf.results['base']['TGLFout'][i].Mt for i in range(len(rho_locations))]) - S = np.array([tglf.results['base']['TGLFout'][i].Se for i in range(len(rho_locations))]) - - if Qi_includes_fast: - - Qifast = [tglf.results['base']['TGLFout'][i].Qifast for i in range(len(rho_locations))] - - if Qifast.sum() != 0.0: - print(f"\t- Qi includes fast ions, adding their contribution") - Qi += Qifast + Flux_base = np.array([Qe, Qi, Ge, GZ, Mt, S]) - Flux_mean = np.array([Qe, Qi, Ge, GZ, Mt, S]) - Flux_std = abs(Flux_mean)*percentError[0]/100.0 + if use_tglf_scan_trick is None: + + # ******************************************************************* + # Just apply an ad-hoc percent error to the results + # ******************************************************************* + + Flux_mean = Flux_base + Flux_std = abs(Flux_mean)*percentError[0]/100.0 else: @@ -115,10 +119,11 @@ def _evaluate_tglf(self): # Run TGLF with scans to estimate the uncertainty # ******************************************************************* - Flux_base, Flux_mean, Flux_std = _run_tglf_uncertainty_model( + Flux_mean, Flux_std = _run_tglf_uncertainty_model( tglf, rho_locations, self.powerstate.ProfilesPredicted, + Flux_base = Flux_base, TGLFsettings=TGLFsettings, extraOptionsTGLF=extraOptions, impurityPosition=impurityPosition, @@ -258,6 +263,7 @@ def _run_tglf_uncertainty_model( tglf, rho_locations, ProfilesPredicted, + Flux_base = None, TGLFsettings=None, extraOptionsTGLF=None, impurityPosition=1, @@ -321,7 +327,7 @@ def _run_tglf_uncertainty_model( TGLFsettings = TGLFsettings, extraOptions = extraOptionsTGLF, ApplyCorrections = False, - add_baseline_to = 'first', + add_baseline_to = 'none', cold_start=cold_start, forceIfcold_start=True, slurm_setup={ @@ -364,6 +370,15 @@ def _run_tglf_uncertainty_model( print(f"\t- Qi includes fast ions, adding their contribution") Qi += Qifast + # Add the base that was calculated earlier + if Flux_base is not None: + Qe = np.append(np.atleast_2d(Flux_base[0]).T, Qe, axis=1) + Qi = np.append(np.atleast_2d(Flux_base[1]).T, Qi, axis=1) + Ge = np.append(np.atleast_2d(Flux_base[2]).T, Ge, axis=1) + GZ = np.append(np.atleast_2d(Flux_base[3]).T, GZ, axis=1) + Mt = np.append(np.atleast_2d(Flux_base[4]).T, Mt, axis=1) + S = np.append(np.atleast_2d(Flux_base[5]).T, S, axis=1) + # Calculate the standard deviation of the scans, that's going to be the reported stds def calculate_mean_std(Q): @@ -388,9 +403,7 @@ def calculate_mean_std(Q): S_point, S_std = calculate_mean_std(S) #TODO: Careful with fast particles - - Flux_base = [Qe[:,0], Qi[:,0], Ge[:,0], GZ[:,0], Mt[:,0], S[:,0]] Flux_mean = [Qe_point, Qi_point, Ge_point, GZ_point, Mt_point, S_point] Flux_std = [Qe_std, Qi_std, Ge_std, GZ_std, Mt_std, S_std] - return Flux_base, Flux_mean, Flux_std + return Flux_mean, Flux_std From c981b0a30ff1a5f765ecb5ca1f0db14789a613ae Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 16:33:40 -0400 Subject: [PATCH 219/385] SMARTS relevant: Whenever a simulation is farmed out, a log file indicating how it was run is written to the folder --- src/mitim_tools/misc_tools/FARMINGtools.py | 95 ++++++++++++++++++-- src/mitim_tools/simulation_tools/SIMtools.py | 12 +-- 2 files changed, 94 insertions(+), 13 deletions(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 156c76f9..fcb54ef1 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -2,6 +2,7 @@ Set of tools to farm out simulations to run in either remote clusters or locally, serially or parallel """ +from math import log from tqdm import tqdm import os import shutil @@ -56,11 +57,17 @@ """ class mitim_job: - def __init__(self, folder_local): + def __init__( + self, + folder_local, + log_simulation_file = None # If not None, log information of how the simulation went to this file + ): + if not isinstance(folder_local, (str, Path)): raise TypeError('MITIM job folder must be a valid string or pathlib.Path object to a local directory') self.folder_local = IOtools.expandPath(folder_local) self.jobid = None + self.log_simulation_file = log_simulation_file def define_machine( self, @@ -266,6 +273,7 @@ def full_process( wait_for_all_commands=wait_for_all_commands, printYN=True, timeoutSecs=timeoutSecs if timeoutSecs < 1e6 else None, + log_file=self.log_simulation_file ) # ~~~~~~ Retrieve @@ -463,9 +471,7 @@ def create_scratch_folder(self): return output, error def send(self): - print( - f'\t* Sending files{" to remote server" if self.ssh is not None else ""}:' - ) + print(f'\t* Sending files{" to remote server" if self.ssh is not None else ""}:') # Create a tarball of the local directory print("\t\t- Tarballing (local side)") @@ -512,12 +518,87 @@ def send(self): print("\t\t- Removing tarball (remote side)") self.execute(f"rm {self.folderExecution}/mitim_send.tar.gz") - def execute(self, command_str, **kwargs): + def execute(self, command_str, log_file=None, **kwargs): if self.ssh is not None: - return self.execute_remote(command_str, **kwargs) + output, error = self.execute_remote(command_str, **kwargs) else: - return self.execute_local(command_str, **kwargs) + output, error = self.execute_local(command_str, **kwargs) + + # Write information file about where and how the run took place + if log_file is not None: + self.write_information_file(command_str, output, error, file=log_file) + + return output, error + + def write_information_file(self, command, output, error, file = 'mitim_simulation.log'): + """ + Write a log file with information about where the simulation happened (local/remote), + user, host, ssh settings if remote, and head/tail of output and error. + """ + import getpass + import platform + from datetime import datetime + + # Prepare context info + now = datetime.now().strftime('%Y-%m-%d %H:%M:%S') + is_remote = self.ssh is not None + lines = [] + lines.append("==================== MITIM Simulation Execution Log ====================\n") + lines.append(f"Date (finished): {now}") + lines.append(f"Execution Type: {'Remote' if is_remote else 'Local'}\n") + lines.append("--- Execution Details ---") + if is_remote: + lines.append(f"SSH User: {getattr(self, 'target_user', 'N/A')}") + lines.append(f"SSH Host: {getattr(self, 'target_host', 'N/A')}") + lines.append(f"Remote Folder: {self.folderExecution}") + else: + lines.append(f"User: {getpass.getuser()}") + lines.append(f"Host: {platform.node()}") + lines.append(f"Folder: {self.folderExecution}") + lines.append("") + + def get_head_tail(data, n=20): + try: + text = data.decode("utf-8", errors="replace") + except Exception: + text = str(data) + lines_ = text.splitlines() + head = lines_[:n] + tail = lines_[-n:] if len(lines_) > n else [] + return head, tail + + out_head, out_tail = get_head_tail(output) + err_head, err_tail = get_head_tail(error) + + lines.append("--- Output (Head) ---") + lines.extend(out_head if out_head else [""]) + lines.append("") + lines.append("--- Output (Tail) ---") + lines.extend(out_tail if out_tail else [""]) + lines.append("") + lines.append("--- Error (Head) ---") + lines.extend(err_head if err_head else [""]) + lines.append("") + lines.append("--- Error (Tail) ---") + lines.extend(err_tail if err_tail else [""]) + lines.append("") + lines.append(f"--- Command ---") + lines.append(f"{command}") + lines.append(f"\n--- Input Files ---") + lines.extend([str(file) for file in self.input_files]) + lines.append(f"\n--- Input Folders ---") + lines.extend([str(folder) for folder in self.input_folders]) + lines.append(f"\n--- Output Files ---") + lines.extend([str(file) for file in self.output_files]) + lines.append(f"\n--- Output Folders ---") + lines.extend([str(folder) for folder in self.output_folders]) + lines.append("\n=======================================================================\n") + + # Write to file (file can be Path or str) + file_path = file if isinstance(file, (str, Path)) else str(file) + with open(file_path, "w", encoding="utf-8") as f: + f.write("\n".join(lines)) def execute_remote( self, diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 0155449d..08fb8e31 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -353,7 +353,9 @@ def _run( tmpFolder = self.FolderGACODE / f"tmp_{code}" IOtools.askNewFolder(tmpFolder, force=True) - self.simulation_job = FARMINGtools.mitim_job(tmpFolder) + kkeys = [keys for keys in code_executor.keys()] + log_simulation_file=self.FolderGACODE / f"mitim_simulation_{kkeys[0]}.log" # Refer with the first folder + self.simulation_job = FARMINGtools.mitim_job(tmpFolder, log_simulation_file=log_simulation_file) self.simulation_job.define_machine_quick(code,f"mitim_{name}") @@ -633,12 +635,11 @@ def _organize_results(self, code_executor, tmpFolder, filesToRetrieve): fineall = True for subfolder_sim in code_executor: - for i, rho in enumerate(code_executor[subfolder_sim].keys()): + for rho in code_executor[subfolder_sim].keys(): for file in filesToRetrieve: original_file = f"{file}_{rho:.4f}" - final_destination = ( - code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" - ) + final_destination = code_executor[subfolder_sim][rho]['folder'] / f"{original_file}" + final_destination.unlink(missing_ok=True) temp_file = tmpFolder / subfolder_sim / f"rho_{rho:.4f}" / f"{file}" @@ -658,7 +659,6 @@ def _organize_results(self, code_executor, tmpFolder, filesToRetrieve): else: print("\t\t- Some files were not retrieved", typeMsg="w") - def run_scan( self, subfolder, # 'scan1', From 9bbb954306c305fd71c475df20dc8b215ab477d8 Mon Sep 17 00:00:00 2001 From: pabloprf Date: Wed, 27 Aug 2025 17:18:10 -0400 Subject: [PATCH 220/385] Bug fix slurm submitter --- src/mitim_tools/misc_tools/FARMINGtools.py | 3 +++ src/mitim_tools/opt_tools/scripts/slurm.py | 14 ++++++++------ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index fcb54ef1..50042300 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -1075,6 +1075,9 @@ def create_slurm_execution_files( # slurm_settings indicate the job resource allocation # --------------------------------------------------- + if slurm_settings is None: + slurm_settings = {} + nameJob = slurm_settings.setdefault("name", "mitim_job") minutes = int(slurm_settings.setdefault("minutes", 10)) memory_req_by_job = slurm_settings.setdefault("mem", None) diff --git a/src/mitim_tools/opt_tools/scripts/slurm.py b/src/mitim_tools/opt_tools/scripts/slurm.py index c24dc425..dda7006b 100644 --- a/src/mitim_tools/opt_tools/scripts/slurm.py +++ b/src/mitim_tools/opt_tools/scripts/slurm.py @@ -44,14 +44,16 @@ def run_slurm( _, fileSBATCH, _ = FARMINGtools.create_slurm_execution_files( command, - folder_remote=folder, + folder, folder_local=folder, - nameJob=nameJob, slurm={"partition": partition, 'exclude': exclude}, - minutes=int(60 * hours), - ntasks=1, - cpuspertask=n, - memory_req_by_job=mem, + slurm_settings = { + 'nameJob': nameJob, + 'minutes': int(60 * hours), + 'ntasks': 1, + 'cpuspertask': n, + 'memory_req_by_job': mem, + } ) command_execution = f"sbatch {fileSBATCH}" From 42cc2c1c72d5f800ce9bdf9b06bf4cfbb3554d30 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 17:52:18 -0400 Subject: [PATCH 221/385] bug fix --- src/mitim_tools/gacode_tools/CGYROtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 8ce216c1..bbe9e0a1 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -18,7 +18,7 @@ def __init__( super().__init__(rhos=rhos) def code_call(folder, p, n = 1, nomp = 1, additional_command="", **kwargs): - return f"cgyro -e {folder} -n {n} -nomp {nomp} {additional_command} -p {p}" + return f"cgyro -e {folder} -n {n} -nomp {nomp} -p {p} {additional_command}" def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): From c9f9bedff3877a5aca07c987544a8bca6c900e8a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 27 Aug 2025 18:21:21 -0400 Subject: [PATCH 222/385] First implementation of automatic PORTALS-CGYRO working --- src/mitim_modules/portals/PORTALSmain.py | 8 ++- .../physics_models/transport_cgyroneo.py | 63 ++++++++++++++----- .../physics_models/transport_tglfneo.py | 34 +++++----- src/mitim_tools/simulation_tools/SIMtools.py | 4 ++ tests/CGYRO_workflow.py | 54 ++++++++-------- 5 files changed, 105 insertions(+), 58 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index cb6bdf53..6acaff66 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -169,7 +169,13 @@ def __init__( "Tfast_ratio": False, # Keep the ratio of Tfast/Te constant throughout the Te evolution "ensureMachNumber": None, # Change w0 to match this Mach number when Ti varies }, - "transport_model": {"TGLFsettings": 6, "extraOptionsTGLF": {}} + "transport_model": { + "run_type": "normal", # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit + "TGLFsettings": 6, + "extraOptionsTGLF": {}, + "CGYROsettings": 1, + "extraOptionsCGYRO": {} + } } for key in self.MODELparameters.keys(): diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index 5be29364..f81a19e6 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -24,7 +24,11 @@ def evaluate_turbulence(self): cold_start = transport_evaluator_options.get("cold_start", False) - run_type = 'prep' + run_type = transport_evaluator_options["MODELparameters"]["transport_model"]["run_type"] + CGYROsettings = transport_evaluator_options["MODELparameters"]["transport_model"]["CGYROsettings"] + extraOptionsCGYRO = transport_evaluator_options["MODELparameters"]["transport_model"]["extraOptionsCGYRO"] + every_n_minutesCGYRO = transport_evaluator_options["MODELparameters"]["transport_model"]["every_n_minutesCGYRO"] + tminCGYRO = transport_evaluator_options["MODELparameters"]["transport_model"]["tminCGYRO"] # ------------------------------------------------------------------------------------------------------------------------ # Prepare CGYRO object @@ -39,25 +43,53 @@ def evaluate_turbulence(self): self.folder, ) + _ = cgyro.run( + 'base_cgyro', + run_type = run_type, + code_settings=CGYROsettings, + extraOptions=extraOptionsCGYRO, + cold_start=cold_start, + forceIfcold_start=True, + ) + if run_type in ['normal', 'submit']: - raise Exception("[MITIM] Automatic submission or full run not implemented") + + if run_type in ['submit']: + cgyro.check(every_n_minutes=every_n_minutesCGYRO) + cgyro.fetch() - # cgyro.read( - # label='base_cgyro' - # ) + cgyro.read( + label='base_cgyro', + tmin = tminCGYRO + ) + + # ------------------------------------------------------------------------------------------------------------------------ + # Pass the information to what power_transport expects + # ------------------------------------------------------------------------------------------------------------------------ - # TRANSPORTtools.write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb') + self.QeGB_turb = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Qe_mean for i in range(len(rho_locations))]) + self.QeGB_turb_stds = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Qe_std for i in range(len(rho_locations))]) + + self.QiGB_turb = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Qi_mean for i in range(len(rho_locations))]) + self.QiGB_turb_stds = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Qi_std for i in range(len(rho_locations))]) + + self.GeGB_turb = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Ge_mean for i in range(len(rho_locations))]) + self.GeGB_turb_stds = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Ge_std for i in range(len(rho_locations))]) - elif run_type == 'prep': + self.GZGB_turb = self.QeGB_turb*0.0 #TODO + self.GZGB_turb_stds = self.QeGB_turb*0.0 #TODO - _ = cgyro.run( - 'base_cgyro', - run_type = run_type, - code_settings=1, - cold_start=cold_start, - forceIfcold_start=True, - ) - + self.MtGB_turb = self.QeGB_turb*0.0 #TODO + self.MtGB_turb_stds = self.QeGB_turb*0.0 #TODO + + self.QieGB_turb = self.QeGB_turb*0.0 #TODO + self.QieGB_turb_stds = self.QeGB_turb*0.0 #TODO + + from mitim_modules.powertorch.utils.TRANSPORTtools import write_json + write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb') + + elif run_type == 'prep': + # Wait until the user has placed the json file in the right folder pre_checks(self) @@ -79,6 +111,7 @@ def evaluate_turbulence(self): if file_path.exists(): all_good = post_checks(self) + def pre_checks(self): plasma = self.powerstate.plasma diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 031c5871..1df7c115 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -23,7 +23,7 @@ def produce_profiles(self): def evaluate_turbulence(self): self._evaluate_tglf() - # Have it separate such that I can call it from the CGYRO class but not with the decorator + # Have it separate such that I can call it from the CGYRO class but without the decorator def _evaluate_tglf(self): # ------------------------------------------------------------------------------------------------------------------------ @@ -42,7 +42,6 @@ def _evaluate_tglf(self): percentError = transport_evaluator_options.get("percentError", [5, 1, 0.5]) use_tglf_scan_trick = transport_evaluator_options.get("use_tglf_scan_trick", None) cores_per_tglf_instance = transport_evaluator_options.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) - keep_tglf_files = transport_evaluator_options.get("keep_tglf_files", "minimal") # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) @@ -63,10 +62,9 @@ def _evaluate_tglf(self): ) # ------------------------------------------------------------------------------------------------------------------------ - # Run TGLF + # Run TGLF (base) # ------------------------------------------------------------------------------------------------------------------------ - - # Run base + tglf.run( 'base_tglf', code_settings=TGLFsettings, @@ -103,6 +101,10 @@ def _evaluate_tglf(self): Qi += Qifast Flux_base = np.array([Qe, Qi, Ge, GZ, Mt, S]) + + # ------------------------------------------------------------------------------------------------------------------------ + # Evaluate TGLF uncertainty + # ------------------------------------------------------------------------------------------------------------------------ if use_tglf_scan_trick is None: @@ -139,7 +141,7 @@ def _evaluate_tglf(self): self._raise_warnings(tglf, rho_locations, Qi_includes_fast) # ------------------------------------------------------------------------------------------------------------------------ - # Pass the information + # Pass the information to what power_transport expects # ------------------------------------------------------------------------------------------------------------------------ self.QeGB_turb = Flux_mean[0] @@ -157,27 +159,27 @@ def _evaluate_tglf(self): self.MtGB_turb = Flux_mean[4] self.MtGB_turb_stds = Flux_std[4] - if provideTurbulentExchange: - self.QieGB_turb = Flux_mean[5] - self.QieGB_turb_stds = Flux_std[5] - else: - self.QieGB_turb = Flux_mean[5] * 0.0 - self.QieGB_turb_stds = Flux_std[5] * 0.0 + self.QieGB_turb = Flux_mean[5] if provideTurbulentExchange else Flux_mean[5]*0.0 + self.QieGB_turb_stds = Flux_std[5] if provideTurbulentExchange else Flux_std[5]*0.0 return tglf @IOtools.hook_method(after=partial(TRANSPORTtools.write_json, file_name = 'fluxes_neoc.json', suffix= 'neoc')) def evaluate_neoclassical(self): - # Options - + # ------------------------------------------------------------------------------------------------------------------------ + # Grab options from powerstate + # ------------------------------------------------------------------------------------------------------------------------ + transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] cold_start = transport_evaluator_options.get("cold_start", False) percentError = transport_evaluator_options.get("percentError", [5, 1, 0.5]) impurityPosition = self.powerstate.impurityPosition_transport + # ------------------------------------------------------------------------------------------------------------------------ # Run + # ------------------------------------------------------------------------------------------------------------------------ rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] @@ -204,7 +206,7 @@ def evaluate_neoclassical(self): Mt = np.array([neo.results['base']['NEOout'][i].Mt for i in range(len(rho_locations))]) # ------------------------------------------------------------------------------------------------------------------------ - # Pass the information + # Pass the information to what power_transport expects # ------------------------------------------------------------------------------------------------------------------------ self.QeGB_neoc = Qe @@ -213,12 +215,14 @@ def evaluate_neoclassical(self): self.GZGB_neoc = GZ self.MtGB_neoc = Mt + # Uncertainties is just a percent of the value self.QeGB_neoc_stds = abs(Qe) * percentError[1]/100.0 self.QiGB_neoc_stds = abs(Qi) * percentError[1]/100.0 self.GeGB_neoc_stds = abs(Ge) * percentError[1]/100.0 self.GZGB_neoc_stds = abs(GZ) * percentError[1]/100.0 self.MtGB_neoc_stds = abs(Mt) * percentError[1]/100.0 + # No neoclassical exchange self.QieGB_neoc = Qe * 0.0 self.QieGB_neoc_stds = Qe * 0.0 diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 08fb8e31..c0b4fae2 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -548,6 +548,10 @@ def _run( # Submit run but do not wait; the user should do checks and fetch results elif run_type == 'submit': + + if type_of_submission == "slurm_array": + raise Exception("[MITIM] run_type='submit' with slurm_array not implemented yet because of issues about folders being moved around, TBD") + self.simulation_job.run( waitYN=False, check_if_files_received=False, diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index 72c649d2..4f27f67e 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -45,30 +45,30 @@ cgyro.read(label="cgyro1") cgyro.plot(labels=["cgyro1"]) -# --------------- -# Scan of KY -# --------------- - -run_type = 'normal' - -cgyro.run_scan( - 'scan1', - extraOptions={ - 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file - }, - variable='KY', - varUpDown=[0.3,0.4], - slurm_setup={ - 'cores':16 - }, - cold_start=cold_start, - forceIfcold_start=True, - run_type=run_type - ) - -cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) - -fig = cgyro.fn.add_figure(label="Quick linear") -cgyro.plot_quick_linear(labels=["scan1_KY_0.3","scan1_KY_0.4"], fig = fig) - -cgyro.fn.show() +# # --------------- +# # Scan of KY +# # --------------- + +# run_type = 'normal' + +# cgyro.run_scan( +# 'scan1', +# extraOptions={ +# 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file +# }, +# variable='KY', +# varUpDown=[0.3,0.4], +# slurm_setup={ +# 'cores':16 +# }, +# cold_start=cold_start, +# forceIfcold_start=True, +# run_type=run_type +# ) + +# cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) + +# fig = cgyro.fn.add_figure(label="Quick linear") +# cgyro.plot_quick_linear(labels=["scan1_KY_0.3","scan1_KY_0.4"], fig = fig) + +# cgyro.fn.show() From 366ffd7f903018086a88fef14301efd3fc812d6c Mon Sep 17 00:00:00 2001 From: pabloprf Date: Wed, 27 Aug 2025 21:54:57 -0400 Subject: [PATCH 223/385] bug fix --- src/mitim_tools/misc_tools/FARMINGtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 50042300..503fc002 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -1080,7 +1080,7 @@ def create_slurm_execution_files( nameJob = slurm_settings.setdefault("name", "mitim_job") minutes = int(slurm_settings.setdefault("minutes", 10)) - memory_req_by_job = slurm_settings.setdefault("mem", None) + memory_req_by_job = slurm_settings.setdefault("memory_req_by_job", None) nodes = slurm_settings.setdefault("nodes", None) ntasks = slurm_settings.setdefault("ntasks", None) From 704859d0ee70ac86d3aa67d23d054e3a5c391ed4 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 14:04:28 -0400 Subject: [PATCH 224/385] Substantial changes to naming convention for portals parameters --- regressions/portals_regressions.py | 26 +- .../maestro/scripts/run_maestro.py | 2 +- .../maestro/tmp_tests/maestro_test1.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 285 +++++++++--------- src/mitim_modules/portals/PORTALStools.py | 34 +-- .../portals/utils/PORTALSanalysis.py | 67 ++-- .../portals/utils/PORTALSinit.py | 163 +++++----- .../portals/utils/PORTALSoptimization.py | 4 +- .../portals/utils/PORTALSplot.py | 94 +++--- src/mitim_modules/powertorch/STATEtools.py | 36 +-- .../physics_models/transport_analytic.py | 4 +- .../physics_models/transport_cgyro.py | 33 +- .../physics_models/transport_cgyroneo.py | 24 +- .../physics_models/transport_tglfneo.py | 93 +++--- .../physics_models/transport_tgyro.py | 136 +++------ .../powertorch/scripts/calculateTargets.py | 44 +-- .../powertorch/utils/POWERplot.py | 14 +- .../powertorch/utils/TARGETStools.py | 2 +- .../powertorch/utils/TRANSFORMtools.py | 8 +- .../powertorch/utils/TRANSPORTtools.py | 4 +- templates/maestro_namelist.json | 4 +- tests/PORTALS_workflow.py | 19 +- tests/POWERTORCH_workflow.py | 2 +- tutorials/PORTALS_tutorial.py | 14 +- 24 files changed, 532 insertions(+), 582 deletions(-) diff --git a/regressions/portals_regressions.py b/regressions/portals_regressions.py index 36ce13d6..7a32ab14 100644 --- a/regressions/portals_regressions.py +++ b/regressions/portals_regressions.py @@ -55,12 +55,12 @@ def conditions_regressions(variables): portals_fun = PORTALSmain.portals(folderWork) portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - portals_fun.INITparameters["remove_fast"] = True + portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True - portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti"] + portals_fun.portals_parameters["model_parameters"]["predicted_channels"] = ["te", "ti"] portals_fun.optimization_options["acquisition_options"]["optimizers"] = ["botorch"] - portals_fun.PORTALSparameters["transport_evaluator"] = TRANSPORTtools.diffusion_model + portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator"] = TRANSPORTtools.diffusion_model transport_evaluator_options = {'chi_e': torch.ones(5)*0.5,'chi_i': torch.ones(5)*2.0} portals_fun.prep(inputgacode, folderWork, transport_evaluator_options=transport_evaluator_options) @@ -93,13 +93,13 @@ def conditions_regressions(variables): portals_fun = PORTALSmain.portals(folderWork) portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - portals_fun.MODELparameters["RhoLocations"] = [0.25, 0.45, 0.65, 0.85] - portals_fun.INITparameters["remove_fast"] = True - portals_fun.INITparameters["quasineutrality"] = True - portals_fun.INITparameters["enforce_same_aLn"] = True - portals_fun.MODELparameters["transport_model"]["TGLFsettings"] = 2 + portals_fun.portals_parameters["model_parameters"]["radii_rho"] = [0.25, 0.45, 0.65, 0.85] + portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True + portals_fun.portals_parameters["initialization_parameters"]["quasineutrality"] = True + portals_fun.portals_parameters["initialization_parameters"]["enforce_same_aLn"] = True + portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["TGLFsettings"] = 2 - portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti", "ne"] + portals_fun.portals_parameters["model_parameters"]["predicted_channels"] = ["te", "ti", "ne"] portals_fun.prep(inputgacode, folderWork) mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=False, askQuestions=False) @@ -134,12 +134,12 @@ def conditions_regressions(variables): # portals_fun = PORTALSmain.portals(folderWork) # portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 # portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - # portals_fun.INITparameters["remove_fast"] = True + # portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True - # portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti", "ne",'nZ','w0'] + # portals_fun.portals_parameters["model_parameters"]["predicted_channels"] = ["te", "ti", "ne",'nZ','w0'] - # portals_fun.PORTALSparameters["ImpurityOfInterest"] = 'W' - # portals_fun.PORTALSparameters["turbulent_exchange_as_surrogate"] = True + # portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"] = 'W' + # portals_fun.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"] = True # portals_fun.prep(inputgacode, folderWork) # mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=False, askQuestions=False) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index eb37f077..a85f8a1a 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -146,7 +146,7 @@ def profiles_postprocessing_fun(file_profs): p.enforce_same_density_gradients() p.write_state(file=file_profs) return p - beat_namelist['PORTALSparameters']['profiles_postprocessing_fun'] = profiles_postprocessing_fun + beat_namelist['PORTALSparameters']['main_parameters']['profiles_postprocessing_fun'] = profiles_postprocessing_fun else: raise ValueError(f"[MITIM] {beat_type} beat not found in the MAESTRO namelist") diff --git a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py b/src/mitim_modules/maestro/tmp_tests/maestro_test1.py index 32846cb6..04c021b9 100644 --- a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py +++ b/src/mitim_modules/maestro/tmp_tests/maestro_test1.py @@ -42,7 +42,7 @@ } # To see what values this namelist can take: mitim_modules/portals/PORTALSmain.py: __init__() -portals_namelist = { "PORTALSparameters": {"launchEvaluationsAsSlurmJobs": True,"forceZeroParticleFlux": True, 'use_tglf_scan_trick': 0.02}, +portals_namelist = { "main_parameters": {"launchEvaluationsAsSlurmJobs": True,"forceZeroParticleFlux": True, 'use_tglf_scan_trick': 0.02}, "MODELparameters": { "RoaLocations": [0.35,0.55,0.75,0.875,0.9], "ProfilesPredicted": ["te", "ti", "ne"], "Physics_options": {"TypeTarget": 3}, diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 6acaff66..7fa098b1 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -115,7 +115,34 @@ def __init__( # Default (please change to your desire after instancing the object) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.potential_flags = {'INITparameters': [], 'MODELparameters': [], 'PORTALSparameters': []} + self.portals_parameters = {} + + """ + Physics-informed parameters to fit surrogates + --------------------------------------------- + """ + + ( + portals_transformation_variables, + portals_transformation_variables_trace, + ) = PORTALStools.default_portals_transformation_variables(additional_params = additional_params_in_surrogate) + + """ + PORTALS-specific parameters (i.e. on the transport solver level, not model) + ---------------------------------------------------------------------------- + """ + + self.portals_parameters["main_parameters"] = { + "portals_transformation_variables": portals_transformation_variables, # Physics-informed parameters to fit surrogates + "portals_transformation_variables_trace": portals_transformation_variables_trace, # Physics-informed parameters to fit surrogates for trace impurities + "turbulent_exchange_as_surrogate": False, # Run turbulent exchange as surrogate? + "Pseudo_multipliers": [1.0]*5, # [Qe,Qi,Ge] multipliers to calculate pseudo + "applyImpurityGammaTrick": True, # If True, fit model to GZ/nZ, valid on the trace limit + "UseOriginalImpurityConcentrationAsWeight": 1.0, # If not None, using UseOriginalImpurityConcentrationAsWeight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis + "fImp_orig": 1.0, + "additional_params_in_surrogate": additional_params_in_surrogate, + "keep_full_model_folder": True, # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) + } """ Parameters to initialize files @@ -126,75 +153,23 @@ def __init__( to run these corrections. """ - self.INITparameters = { - "recalculate_ptot": True, # Recompute PTOT to match kinetic profiles (after removals) - "quasineutrality": False, # Make sure things are quasineutral by changing the *MAIN* ion (D,T or both) (after removals) - "removeIons": [], # Remove this ion from the input.gacode (if D,T,Z, eliminate T with [2]) - "remove_fast": False, # Automatically detect which are fast ions and remove them - "thermalize_fast": False, # Do not remove fast, keep their diluiton effect but make them thermal + self.portals_parameters["initialization_parameters"] = { + "recalculate_ptot": True, # Recompute PTOT to match kinetic profiles (after removals) + "quasineutrality": False, # Make sure things are quasineutral by changing the *MAIN* ion (D,T or both) (after removals) + "removeIons": [], # Remove this ion from the input.gacode (if D,T,Z, eliminate T with [2]) + "remove_fast": False, # Automatically detect which are fast ions and remove them + "thermalize_fast": False, # Do not remove fast, keep their diluiton effect but make them thermal "enforce_same_aLn": False, # Make all ion density gradients equal to electrons "groupQIONE": False, "ensure_positive_Gamma": False, "ensureMachNumber": None, } - for key in self.INITparameters.keys(): - self.potential_flags['INITparameters'].append(key) - - """ - Parameters to run the model - --------------------------- - The corrections are applied prior to each evaluation, so that things are consistent. - Here, do not include things that are not specific for a given iteration. Otherwise if they are general - changes to input.gacode, then that should go into INITparameters. - - if MODELparameters contains RoaLocations, use that instead of RhoLocations - """ - - self.MODELparameters = { - "RhoLocations": [0.3, 0.45, 0.6, 0.75, 0.9], - "RoaLocations": None, - "ProfilesPredicted": ["te", "ti", "ne"], # ['nZ','w0'] - "Physics_options": { - "TypeTarget": 3, - "TurbulentExchange": 0, # In PORTALS TGYRO evaluations, let's always calculate turbulent exchange, but NOT include it in targets! - "PtotType": 1, # In PORTALS TGYRO evaluations, let's always use the PTOT column (so control of that comes from the ap) - "GradientsType": 0, # In PORTALS TGYRO evaluations, we need to not recompute gradients - "InputType": 1, # In PORTALS TGYRO evaluations, we need to use exact profiles - }, - "applyCorrections": { - "Ti_thermals": True, # Keep all thermal ion temperatures equal to the main Ti - "ni_thermals": True, # Adjust for quasineutrality by modifying the thermal ion densities together with ne - "recalculate_ptot": True, # Recompute PTOT to insert in input file each time - "Tfast_ratio": False, # Keep the ratio of Tfast/Te constant throughout the Te evolution - "ensureMachNumber": None, # Change w0 to match this Mach number when Ti varies - }, - "transport_model": { - "run_type": "normal", # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit - "TGLFsettings": 6, - "extraOptionsTGLF": {}, - "CGYROsettings": 1, - "extraOptionsCGYRO": {} - } - } - - for key in self.MODELparameters.keys(): - self.potential_flags['MODELparameters'].append(key) - - """ - Physics-informed parameters to fit surrogates - --------------------------------------------- - """ - - ( - portals_transformation_variables, - portals_transformation_variables_trace, - ) = PORTALStools.default_portals_transformation_variables(additional_params = additional_params_in_surrogate) - - """ - Parameters to run PORTALS - ----------------------- - """ + ''' + model_parameters + ---------------- + These parameters are communicated to the powertorch object. + ''' # Selection of model if CGYROrun: @@ -204,36 +179,98 @@ def __init__( from mitim_modules.powertorch.physics_models.targets_analytic import analytical_model as target_evaluator - self.PORTALSparameters = { - "percentError": [5,10,1], # (%) Error (std, in percent) of model evaluation [TGLF (treated as minimum if scan trick), NEO, TARGET] - "transport_evaluator": transport_evaluator, - "target_evaluator": target_evaluator, - "target_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) - "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) - "Qi_includes_fast": False, # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) - "portals_transformation_variables": portals_transformation_variables, # Physics-informed parameters to fit surrogates - "portals_transformation_variables_trace": portals_transformation_variables_trace, # Physics-informed parameters to fit surrogates for trace impurities - "Qi_criterion_stable": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable - "percentError_stable": 5.0, # (%) For CGYRO runs, minimum error based on target if case is considered stable - "forceZeroParticleFlux": False, # If True, ignore particle flux profile and assume zero for all radii - "turbulent_exchange_as_surrogate": False, # Run turbulent exchange as surrogate? - "profiles_postprocessing_fun": None, # Function to post-process input.gacode only BEFORE passing to transport codes - "Pseudo_multipliers": [1.0]*5, # [Qe,Qi,Ge] multipliers to calculate pseudo + self.portals_parameters["model_parameters"] = { + + # Specification of radial locations (roa wins over rho, if provided) + "radii_rho": [0.3, 0.45, 0.6, 0.75, 0.9], + "radii_roa": None, + + # Channels to be predicted + "predicted_channels": ["te", "ti", "ne"], # ['nZ','w0'] + + # ------------------------- + # Transport model settings + # ------------------------- + "transport_parameters": { + + # Transport model class + + "transport_evaluator": transport_evaluator, + + # Simulation kwargs to be passed directly to run and read commands + + "transport_evaluator_options": { + "tglf": { + "run": { + "code_settings": 6, + "extraOptions": {}, + }, + "read": {}, + "use_scan_trick_for_stds": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta + "keep_files": "base", # minimal: only retrieve minimal files always; base: retrieve all files when running TGLF base; none: always minimal files + "cores_per_tglf_instance": 1, # Number of cores to use per TGLF instance + "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) + "percent_error": 5, # (%) Error (std, in percent) of model evaluation TGLF if not scan trick + "Qi_includes_fast": False, # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) + }, + "neo": { + "run": {}, + "read": {}, + "percent_error": 10 # (%) Error (std, in percent) of model evaluation + }, + "cgyro": { + "run": { + "code_settings": 1, + "extraOptions": {}, + "run_type": "normal", # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit + }, + "read": { + "tmin": 0.0 + }, + "Qi_criterion_stable": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable + "percentError_stable": 5.0, # (%) For CGYRO runs, minimum error based on target if case is considered stable + }, + }, + + # Corrections to be applied to each iteration input.gacode file + + "profiles_postprocessing_fun": None, # Function to post-process input.gacode only BEFORE passing to transport codes + + "applyCorrections": { + "Ti_thermals": True, # Keep all thermal ion temperatures equal to the main Ti + "ni_thermals": True, # Adjust for quasineutrality by modifying the thermal ion densities together with ne + "recalculate_ptot": True, # Recompute PTOT to insert in input file each time + "Tfast_ratio": False, # Keep the ratio of Tfast/Te constant throughout the Te evolution + "ensureMachNumber": None, # Change w0 to match this Mach number when Ti varies + }, + }, + + # ------------------------- + # Target model settings + # ------------------------- + "target_parameters": { + "target_evaluator": target_evaluator, + "target_evaluator_options": { + "TypeTarget": 3, + "target_evaluator": target_evaluator, + "target_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) + "forceZeroParticleFlux": False, # If True, ignore particle flux profile and assume zero for all radii + "fineTargetsResolution": 20, # If not None, calculate targets with this radial resolution (defaults target_evaluator_method to powerstate) + "percent_error": 1 # (%) Error (std, in percent) of model evaluation + } + }, "ImpurityOfInterest": None, # Impurity to do flux-matching for if nZ enabled (name of first impurity instance AFTER postprocessing), e.g. "W" - "applyImpurityGammaTrick": True, # If True, fit model to GZ/nZ, valid on the trace limit - "UseOriginalImpurityConcentrationAsWeight": 1.0, # If not None, using UseOriginalImpurityConcentrationAsWeight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis - "fImp_orig": 1.0, - "fineTargetsResolution": 20, # If not None, calculate targets with this radial resolution (defaults target_evaluator_method to powerstate) - "hardCodedCGYRO": None, # If not None, use this hard-coded CGYRO evaluation - "additional_params_in_surrogate": additional_params_in_surrogate, - "use_tglf_scan_trick": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta - "keep_full_model_folder": True, # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) - "keep_tglf_files": "base", # minimal: only retrieve minimal files always; base: retrieve all files when running TGLF base; none: always minimal files - "cores_per_tglf_instance": 1, # Number of cores to use per TGLF instance } - for key in self.PORTALSparameters.keys(): - self.potential_flags['PORTALSparameters'].append(key) + # Grab all the flags here in a way that, after changing the dictionary extenrally, I make sure it's the same flags as PORTALS expects + + def _grab_flags_dictionary(d): + keys = {} + for key in d.keys(): + keys[key] = _grab_flags_dictionary(d[key]) if isinstance(d[key], dict) else None + return keys + + self.potential_flags = _grab_flags_dictionary(self.portals_parameters) def prep( self, @@ -284,15 +321,15 @@ def prep( ymax_rel0 = copy.deepcopy(ymax_rel) ymax_rel = {} - for prof in self.MODELparameters["ProfilesPredicted"]: - ymax_rel[prof] = np.array( [ymax_rel0] * len(self.MODELparameters[key_rhos]) ) + for prof in self.portals_parameters["model_parameters"]["predicted_channels"]: + ymax_rel[prof] = np.array( [ymax_rel0] * len(self.portals_parameters["model_parameters"][key_rhos]) ) if IOtools.isfloat(ymin_rel): ymin_rel0 = copy.deepcopy(ymin_rel) ymin_rel = {} - for prof in self.MODELparameters["ProfilesPredicted"]: - ymin_rel[prof] = np.array( [ymin_rel0] * len(self.MODELparameters[key_rhos]) ) + for prof in self.portals_parameters["model_parameters"]["predicted_channels"]: + ymin_rel[prof] = np.array( [ymin_rel0] * len(self.portals_parameters["model_parameters"][key_rhos]) ) if enforceFiniteTemperatureGradients is not None: for prof in ['te', 'ti']: @@ -305,7 +342,7 @@ def prep( self, self.folder, fileGACODE, - self.INITparameters, + self.portals_parameters, ymax_rel, ymin_rel, start_from_folder=start_from_folder, @@ -317,7 +354,6 @@ def prep( tensor_options = self.tensor_options, seedInitial=seedInitial, checkForSpecies=askQuestions, - transport_evaluator_options=transport_evaluator_options, ) print(">> PORTALS initalization module (END)", typeMsg="i") @@ -382,7 +418,7 @@ def run(self, paramsfile, resultsfile): name, numPORTALS=numPORTALS, dictOFs=dictOFs, - remove_folder_upon_completion=not self.PORTALSparameters["keep_full_model_folder"], + remove_folder_upon_completion=not self.portals_parameters["main_parameters"]["keep_full_model_folder"], ) # Write results @@ -441,7 +477,7 @@ def scalarized_objective(self, Y): res must have shape (dim1...N) """ - of, cal, _, res = PORTALStools.calculate_residuals(self.powerstate, self.PORTALSparameters,specific_vars=var_dict) + of, cal, _, res = PORTALStools.calculate_residuals(self.powerstate, self.portals_parameters,specific_vars=var_dict) return of, cal, res @@ -459,40 +495,19 @@ def check_flags(self): print(">> PORTALS flags pre-check") # Check that I haven't added a deprecated variable that I expect some behavior from - for key in self.potential_flags.keys(): - for flag in self.__dict__[key]: - if flag not in self.potential_flags[key]: - print( - f"\t- {key}['{flag}'] is an unexpected variable, prone to errors or misinterpretation", - typeMsg="q", - ) - # ---------------------------------------------------------------------------------- - - if self.PORTALSparameters["fineTargetsResolution"] is not None: - if self.PORTALSparameters["target_evaluator_method"] != "powerstate": - print("\t- Requested fineTargetsResolution, so running powerstate target calculations",typeMsg="w") - self.PORTALSparameters["target_evaluator_method"] = "powerstate" - - from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model - if not issubclass(self.PORTALSparameters["transport_evaluator"], tgyro_model) and (self.PORTALSparameters["target_evaluator_method"] == "tgyro"): - print("\t- Requested TGYRO targets, but transport evaluator is not tgyro, so changing to powerstate",typeMsg="w") - self.PORTALSparameters["target_evaluator_method"] = "powerstate" - - if ("InputType" not in self.MODELparameters["Physics_options"]) or self.MODELparameters["Physics_options"]["InputType"] != 1: - print("\t- In PORTALS TGYRO evaluations, we need to use exact profiles (InputType=1)",typeMsg="i") - self.MODELparameters["Physics_options"]["InputType"] = 1 - - if ("GradientsType" not in self.MODELparameters["Physics_options"]) or self.MODELparameters["Physics_options"]["GradientsType"] != 0: - print("\t- In PORTALS TGYRO evaluations, we need to not recompute gradients (GradientsType=0)",typeMsg="i") - self.MODELparameters["Physics_options"]["GradientsType"] = 0 - - if self.PORTALSparameters["target_evaluator_method"] == "tgyro" and self.PORTALSparameters['profiles_postprocessing_fun'] is not None: - print("\t- Requested custom modification of postprocessing function but targets from TGYRO... are you sure?",typeMsg="q") - - if self.PORTALSparameters["target_evaluator_method"] == "tgyro" and self.PORTALSparameters['transport_evaluator'] != tgyro_model: - print("\t- Requested TGYRO targets but transport evaluator is not TGYRO... are you sure?",typeMsg="q") + + def _check_flags_dictionary(d, d_check): + for key in d.keys(): + if key not in d_check: + print(f"\t- {key} is an unexpected variable, prone to errors or misinterpretation",typeMsg="q") + elif not isinstance(d[key], dict): + continue + else: + _check_flags_dictionary(d[key], d_check[key]) + + _check_flags_dictionary(self.portals_parameters, self.potential_flags) - key_rhos = "RoaLocations" if self.MODELparameters["RoaLocations"] is not None else "RhoLocations" + key_rhos = "radii_roa" if self.portals_parameters["model_parameters"]["radii_roa"] is not None else "radii_rho" return key_rhos @@ -615,9 +630,9 @@ def runModelEvaluator( # Prepare evaluating vector X # --------------------------------------------------------------------------------------------------- - X = torch.zeros(len(powerstate.ProfilesPredicted) * (powerstate.plasma["rho"].shape[1] - 1)).to(powerstate.dfT) + X = torch.zeros(len(powerstate.predicted_channels) * (powerstate.plasma["rho"].shape[1] - 1)).to(powerstate.dfT) cont = 0 - for ikey in powerstate.ProfilesPredicted: + for ikey in powerstate.predicted_channels: for ix in range(powerstate.plasma["rho"].shape[1] - 1): X[cont] = dictDVs[f"aL{ikey}_{ix+1}"]["value"] cont += 1 @@ -654,7 +669,7 @@ def runModelEvaluator( def map_powerstate_to_portals(powerstate, dictOFs): - for var in powerstate.ProfilesPredicted: + for var in powerstate.predicted_channels: # Write in OFs for i in range(powerstate.plasma["rho"].shape[1] - 1): # Ignore position 0, which is rho=0 if var == "te": diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 81389be2..6538ece0 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -291,7 +291,7 @@ def constructEvaluationProfiles(X, surrogate_parameters, recalculateTargets=Fals # Obtain modified profiles CPs = torch.zeros((X.shape[0], num_x + 1)).to(X) - for iprof, var in enumerate(powerstate.ProfilesPredicted): + for iprof, var in enumerate(powerstate.predicted_channels): # Specific part of the input vector that deals with this profile and introduce to CP vector (that starts with 0,0) CPs[:, 1:] = X[:, (iprof * num_x) : (iprof * num_x) + num_x] @@ -352,7 +352,7 @@ def stopping_criteria_portals(mitim_bo, parameters = {}): -def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): +def calculate_residuals(powerstate, portals_parameters, specific_vars=None): """ Notes ----- @@ -399,7 +399,7 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added # ------------------------------------------------------------------------- - if PORTALSparameters["turbulent_exchange_as_surrogate"]: + if portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]: QieMWm2_tr_turb = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) else: QieMWm2_tr_turb = torch.zeros(dfT.shape).to(dfT) @@ -413,7 +413,7 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): torch.Tensor().to(dfT), torch.Tensor().to(dfT), ) - for prof in powerstate.ProfilesPredicted: + for prof in powerstate.predicted_channels: if prof == "te": var = "Qe" elif prof == "ti": @@ -452,28 +452,28 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): if var == "Qe": of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][0], - cal0 * PORTALSparameters["Pseudo_multipliers"][0], + of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][0], + cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][0], ) elif var == "Qi": of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][1], - cal0 * PORTALSparameters["Pseudo_multipliers"][1], + of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][1], + cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][1], ) elif var == "Ge": of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][2], - cal0 * PORTALSparameters["Pseudo_multipliers"][2], + of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][2], + cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][2], ) elif var == "GZ": of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][3], - cal0 * PORTALSparameters["Pseudo_multipliers"][3], + of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][3], + cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][3], ) elif var == "MtJm2": of0, cal0 = ( - of0 * PORTALSparameters["Pseudo_multipliers"][4], - cal0 * PORTALSparameters["Pseudo_multipliers"][4], + of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][4], + cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][4], ) of, cal = torch.cat((of, of0), dim=-1), torch.cat((cal, cal0), dim=-1) @@ -491,7 +491,7 @@ def calculate_residuals(powerstate, PORTALSparameters, specific_vars=None): return of, cal, source, res -def calculate_residuals_distributions(powerstate, PORTALSparameters): +def calculate_residuals_distributions(powerstate, portals_parameters): """ - Works with tensors - It should be independent on how many dimensions it has, except that the last dimension is the multi-ofs @@ -532,7 +532,7 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added # ------------------------------------------------------------------------- - if PORTALSparameters["turbulent_exchange_as_surrogate"]: + if portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]: QieMWm2_tr_turb = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) QieMWm2_tr_turb_stds = computeTurbExchangeIndividual(var_dict["Qie_tr_turb_stds"], powerstate) else: @@ -545,7 +545,7 @@ def calculate_residuals_distributions(powerstate, PORTALSparameters): of, cal = torch.Tensor().to(dfT), torch.Tensor().to(dfT) ofE, calE = torch.Tensor().to(dfT), torch.Tensor().to(dfT) - for prof in powerstate.ProfilesPredicted: + for prof in powerstate.predicted_channels: if prof == "te": var = "Qe" elif prof == "ti": diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index a875ff6f..8fe72fd4 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -150,17 +150,15 @@ def prep_metrics(self, ilast=None): self.rhos = self.mitim_runs[0]['powerstate'].plasma['rho'][0,1:].cpu().numpy() self.roa = self.mitim_runs[0]['powerstate'].plasma['roa'][0,1:].cpu().numpy() - self.PORTALSparameters = self.opt_fun.mitim_model.optimization_object.PORTALSparameters - self.MODELparameters = self.opt_fun.mitim_model.optimization_object.MODELparameters + self.portals_parameters = self.opt_fun.mitim_model.optimization_object.portals_parameters # Useful flags - self.ProfilesPredicted = self.MODELparameters["ProfilesPredicted"] + self.predicted_channels = self.portals_parameters["model_parameters"]["predicted_channels"] - self.runWithImpurity = self.powerstate.impurityPosition if "nZ" in self.ProfilesPredicted else None + self.runWithImpurity = self.powerstate.impurityPosition if "nZ" in self.predicted_channels else None - self.runWithRotation = "w0" in self.ProfilesPredicted - self.includeFast = self.PORTALSparameters["Qi_includes_fast"] - self.forceZeroParticleFlux = self.PORTALSparameters["forceZeroParticleFlux"] + self.runWithRotation = "w0" in self.predicted_channels + self.forceZeroParticleFlux = self.portals_parameters["model_parameters"]["target_parameters"]["target_evaluator_options"]["forceZeroParticleFlux"] # Profiles and tgyro results print("\t- Reading profiles and tgyros for each evaluation") @@ -170,7 +168,7 @@ def prep_metrics(self, ilast=None): self.powerstates.append(self.mitim_runs[i]["powerstate"]) # runWithImpurity_transport is stored after powerstate has run transport - self.runWithImpurity_transport = self.powerstates[0].impurityPosition_transport if "nZ" in self.ProfilesPredicted else None + self.runWithImpurity_transport = self.powerstates[0].impurityPosition_transport if "nZ" in self.predicted_channels else None if len(self.powerstates) <= self.ibest: print("\t- PORTALS was read after new residual was computed but before pickle was written!",typeMsg="w") @@ -229,7 +227,7 @@ def prep_metrics(self, ilast=None): _, _, source, res = PORTALStools.calculate_residuals( power, - self.PORTALSparameters, + self.portals_parameters, ) # Make sense of tensor "source" which are defining the entire predictive set in @@ -239,7 +237,7 @@ def prep_metrics(self, ilast=None): GZ_resR = np.zeros(self.rhos.shape[0]) Mt_resR = np.zeros(self.rhos.shape[0]) cont = 0 - for prof in self.MODELparameters["ProfilesPredicted"]: + for prof in self.portals_parameters["model_parameters"]["predicted_channels"]: for ix in range(self.rhos.shape[0]): if prof == "te": Qe_resR[ix] = source[0, cont].abs() @@ -273,7 +271,7 @@ def prep_metrics(self, ilast=None): y2_std, ) = PORTALStools.calculate_residuals_distributions( power, - self.PORTALSparameters, + self.portals_parameters, ) QR, chiR = PLASMAtools.RicciMetric( @@ -330,14 +328,14 @@ def prep_metrics(self, ilast=None): self.resCheck = ( self.resTeM + self.resTiM + self.resneM + self.resnZM + self.resw0M - ) / len(self.MODELparameters["ProfilesPredicted"]) + ) / len(self.portals_parameters["model_parameters"]["predicted_channels"]) # --------------------------------------------------------------------------------------------------------------------- # Jacobian # --------------------------------------------------------------------------------------------------------------------- DeltaQ1 = [] - for i in self.MODELparameters["ProfilesPredicted"]: + for i in self.portals_parameters["model_parameters"]["predicted_channels"]: if i == "te": DeltaQ1.append(-self.resTe) if i == "ti": @@ -505,8 +503,7 @@ def extractPORTALS(self, evaluation=None, folder=None, modified_profiles=False): portals_fun = portals(folder) # Transfer settings - portals_fun.PORTALSparameters = portals_fun_original.PORTALSparameters - portals_fun.MODELparameters = portals_fun_original.MODELparameters + portals_fun.portals_parameters = portals_fun_original.portals_parameters # PRINTING print( @@ -514,7 +511,7 @@ def extractPORTALS(self, evaluation=None, folder=None, modified_profiles=False): **************************************************************************************************** > MITIM has extracted PORTALS class to run in {IOtools.clipstr(folder)}, to proceed: 1. Modify any parameter as required - portals_fun.PORTALSparameters, portals_fun.MODELparameters, portals_fun.optimization_options + portals_fun.portals_parameters 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: @@ -547,12 +544,12 @@ def extractTGYRO(self, folder=None, cold_start=False, evaluation=0, modified_pro folder, profilesclass_custom=profiles, cold_start=cold_start, forceIfcold_start=True ) - TGLFsettings = self.MODELparameters["transport_model"]["TGLFsettings"] - extraOptionsTGLF = self.MODELparameters["transport_model"]["extraOptionsTGLF"] + TGLFsettings = self.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["TGLFsettings"] + extraOptionsTGLF = self.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["extraOptionsTGLF"] PredictionSet = [ - int("te" in self.MODELparameters["ProfilesPredicted"]), - int("ti" in self.MODELparameters["ProfilesPredicted"]), - int("ne" in self.MODELparameters["ProfilesPredicted"]), + int("te" in self.portals_parameters["model_parameters"]["predicted_channels"]), + int("ti" in self.portals_parameters["model_parameters"]["predicted_channels"]), + int("ne" in self.portals_parameters["model_parameters"]["predicted_channels"]), ] return tgyro, self.rhos, PredictionSet, TGLFsettings, extraOptionsTGLF @@ -566,12 +563,12 @@ def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=F """ NOTE on radial location extraction: Two possible options for the rho locations to use: - 1. self.MODELparameters["RhoLocations"] -> the ones PORTALS sent to TGYRO + 1. self.portals_parameters["model_parameters"]["radii_rho"] -> 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. #TODO: This should be fixed in the future, we should never send to TGYRO more than 7 decimal places of any variable """ - rhos_considered = self.MODELparameters["RhoLocations"] + rhos_considered = self.portals_parameters["model_parameters"]["radii_rho"] if positions is None: rhos = rhos_considered @@ -596,8 +593,8 @@ def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=F tglf = TGLFtools.TGLF(rhos=rhos) _ = tglf.prep_using_tgyro(folder, cold_start=cold_start, inputgacode=inputgacode) - TGLFsettings = self.MODELparameters["transport_model"]["TGLFsettings"] - extraOptions = self.MODELparameters["transport_model"]["extraOptionsTGLF"] + TGLFsettings = self.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["TGLFsettings"] + extraOptions = self.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["extraOptionsTGLF"] return tglf, TGLFsettings, extraOptions @@ -990,12 +987,12 @@ def plotMetrics(self, extra_lab="", **kwargs): figG = self.fn.add_figure(label=f"{extra_lab} - Sequence") # ----------------------------------------------------------------- - axs, axsM = STATEtools.add_axes_powerstate_plot(figMain, num_kp=np.max([3,len(self.powerstates[-1].ProfilesPredicted)])) + axs, axsM = STATEtools.add_axes_powerstate_plot(figMain, num_kp=np.max([3,len(self.powerstates[-1].predicted_channels)])) colors = GRAPHICStools.listColors() axsGrads_extra = [] cont = 0 - for i in range(np.max([3,len(self.powerstates[-1].ProfilesPredicted)])): + for i in range(np.max([3,len(self.powerstates[-1].predicted_channels)])): axsGrads_extra.append(axs[cont]) axsGrads_extra.append(axs[cont+1]) cont += 4 @@ -1012,8 +1009,8 @@ def plotMetrics(self, extra_lab="", **kwargs): self.powerstates[i].profiles.plot_gradients( axsGrads_extra, color=colors[i], - plotImpurity=self.powerstates[-1].impurityPosition if 'nZ' in self.powerstates[-1].ProfilesPredicted else None, - plotRotation='w0' in self.powerstates[0].ProfilesPredicted, + plotImpurity=self.powerstates[-1].impurityPosition if 'nZ' in self.powerstates[-1].predicted_channels else None, + plotRotation='w0' in self.powerstates[0].predicted_channels, ls='-', lw=0.5, lastRho=self.powerstates[0].plasma["rho"][-1, -1].item(), @@ -1027,8 +1024,8 @@ def plotMetrics(self, extra_lab="", **kwargs): self.profiles[-1].plot_gradients( axsGrads_extra, color=colors[i+1], - plotImpurity=self.powerstates[-1].impurityPosition_transport if 'nZ' in self.powerstates[-1].ProfilesPredicted else None, - plotRotation='w0' in self.powerstates[0].ProfilesPredicted, + plotImpurity=self.powerstates[-1].impurityPosition_transport if 'nZ' in self.powerstates[-1].predicted_channels else None, + plotRotation='w0' in self.powerstates[0].predicted_channels, ls='-', lw=1.0, lastRho=self.powerstates[0].plasma["rho"][-1, -1].item(), @@ -1055,8 +1052,8 @@ def plotMetrics(self, extra_lab="", **kwargs): p.profiles.plot_gradients( axsGrads, color=colors[i], - plotImpurity=p.impurityPosition if 'nZ' in p.ProfilesPredicted else None, - plotRotation='w0' in p.ProfilesPredicted, + plotImpurity=p.impurityPosition if 'nZ' in p.predicted_channels else None, + plotRotation='w0' in p.predicted_channels, lastRho=self.powerstates[0].plasma["rho"][-1, -1].item(), label=f"profile #{i}", ) @@ -1067,8 +1064,8 @@ def plotMetrics(self, extra_lab="", **kwargs): prof.plot_gradients( axsGrads, color=colors[i+1], - plotImpurity=p.impurityPosition_transport if 'nZ' in p.ProfilesPredicted else None, - plotRotation='w0' in p.ProfilesPredicted, + plotImpurity=p.impurityPosition_transport if 'nZ' in p.predicted_channels else None, + plotRotation='w0' in p.predicted_channels, lastRho=p.plasma["rho"][-1, -1].item(), label="next", ) diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index e1cc00d0..fb64b7a4 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -18,7 +18,7 @@ def initializeProblem( portals_fun, folderWork, fileStart, - INITparameters, + portals_parameters, RelVar_y_max, RelVar_y_min, limitsAreRelative=True, @@ -27,7 +27,6 @@ def initializeProblem( dvs_fixed=None, start_from_folder=None, define_ranges_from_profiles=None, - transport_evaluator_options=None, seedInitial=None, checkForSpecies=True, tensor_options = { @@ -39,7 +38,6 @@ def initializeProblem( Notes: - Specification of points occur in rho coordinate, although internally the work is r/a cold_start = True if cold_start from beginning - - I can give transport_evaluator_options directly (e.g. if I want chis or something) - define_ranges_from_profiles must be PROFILES class """ @@ -73,34 +71,34 @@ def initializeProblem( profiles = PROFILEStools.gacode_state(initialization_file) # About radial locations - if portals_fun.MODELparameters["RoaLocations"] is not None: - roa = portals_fun.MODELparameters["RoaLocations"] + if portals_fun.portals_parameters["model_parameters"]["radii_roa"] is not None: + roa = portals_fun.portals_parameters["model_parameters"]["radii_roa"] rho = np.interp(roa, profiles.derived["roa"], profiles.profiles["rho(-)"]) print("\t * r/a provided, transforming to rho:") print(f"\t\t r/a = {roa}") print(f"\t\t rho = {rho}") - portals_fun.MODELparameters["RhoLocations"] = rho + portals_fun.portals_parameters["model_parameters"]["radii_rho"] = rho if ( - len(INITparameters["removeIons"]) > 0 - or INITparameters["remove_fast"] - or INITparameters["quasineutrality"] - or INITparameters["enforce_same_aLn"] - or INITparameters["recalculate_ptot"] + len(portals_parameters["initialization_parameters"]["removeIons"]) > 0 + or portals_parameters["initialization_parameters"]["remove_fast"] + or portals_parameters["initialization_parameters"]["quasineutrality"] + or portals_parameters["initialization_parameters"]["enforce_same_aLn"] + or portals_parameters["initialization_parameters"]["recalculate_ptot"] ): - profiles.correct(options=INITparameters) + profiles.correct(options=portals_parameters["initialization_parameters"]) - if portals_fun.PORTALSparameters["ImpurityOfInterest"] is not None: - position_of_impurity = MITIMstate.impurity_location(profiles, portals_fun.PORTALSparameters["ImpurityOfInterest"]) + if portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"] is not None: + position_of_impurity = MITIMstate.impurity_location(profiles, portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"]) else: position_of_impurity = 1 - if portals_fun.PORTALSparameters["UseOriginalImpurityConcentrationAsWeight"] is not None and portals_fun.PORTALSparameters["ImpurityOfInterest"] is not None: + if portals_fun.portals_parameters["main_parameters"]["UseOriginalImpurityConcentrationAsWeight"] is not None and portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"] is not None: f0 = profiles.Species[position_of_impurity]["n0"] / profiles.profiles['ne(10^19/m^3)'][0] - portals_fun.PORTALSparameters["fImp_orig"] = f0/portals_fun.PORTALSparameters["UseOriginalImpurityConcentrationAsWeight"] - print(f'\t- Ion {portals_fun.PORTALSparameters["ImpurityOfInterest"]} has original central concentration of {f0:.2e}, using its inverse multiplied by {portals_fun.PORTALSparameters["UseOriginalImpurityConcentrationAsWeight"]} as scaling factor of GZ -> {portals_fun.PORTALSparameters["fImp_orig"]:.2e}',typeMsg="i") + portals_fun.portals_parameters["main_parameters"]["fImp_orig"] = f0/portals_fun.portals_parameters["main_parameters"]["UseOriginalImpurityConcentrationAsWeight"] + print(f'\t- Ion {portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"]} has original central concentration of {f0:.2e}, using its inverse multiplied by {portals_fun.portals_parameters["main_parameters"]["UseOriginalImpurityConcentrationAsWeight"]} as scaling factor of GZ -> {portals_fun.portals_parameters["main_parameters"]["fImp_orig"]:.2e}',typeMsg="i") else: - portals_fun.PORTALSparameters["fImp_orig"] = 1.0 + portals_fun.portals_parameters["main_parameters"]["fImp_orig"] = 1.0 # Check if I will be able to calculate radiation speciesNotFound = [] @@ -112,7 +110,7 @@ def initializeProblem( # Print warning or question to be careful! if len(speciesNotFound) > 0: - if portals_fun.MODELparameters["Physics_options"]["TypeTarget"] == 3: + if portals_fun.portals_parameters["model_parameters"]["target_parameters"]["TypeTarget"] == 3: answerYN = print(f"\t- Species {speciesNotFound} not found in radiation database, radiation will be zero in PORTALS... is this ok for your predictions?",typeMsg="q" if checkForSpecies else "w") if checkForSpecies and (not answerYN): @@ -124,31 +122,7 @@ def initializeProblem( # Prepare and defaults - xCPs = torch.from_numpy(np.array(portals_fun.MODELparameters["RhoLocations"])).to(dfT) - - if transport_evaluator_options is None: - transport_evaluator_options = { - "launchMODELviaSlurm": portals_fun.PORTALSparameters["launchEvaluationsAsSlurmJobs"], - "Qi_includes_fast": portals_fun.PORTALSparameters["Qi_includes_fast"], - "TurbulentExchange": portals_fun.PORTALSparameters["turbulent_exchange_as_surrogate"], - "profiles_postprocessing_fun": portals_fun.PORTALSparameters["profiles_postprocessing_fun"], - "UseFineGridTargets": portals_fun.PORTALSparameters["fineTargetsResolution"], - "OriginalFimp": portals_fun.PORTALSparameters["fImp_orig"], - "forceZeroParticleFlux": portals_fun.PORTALSparameters["forceZeroParticleFlux"], - "percentError": portals_fun.PORTALSparameters["percentError"], - "use_tglf_scan_trick": portals_fun.PORTALSparameters["use_tglf_scan_trick"], - "keep_tglf_files": portals_fun.PORTALSparameters["keep_tglf_files"], - "cold_start": False, - "MODELparameters": portals_fun.MODELparameters, - "impurityPosition": position_of_impurity, - } - - if "extra_params" not in transport_evaluator_options: - transport_evaluator_options["extra_params"] = { - "PORTALSparameters": portals_fun.PORTALSparameters, - "folder": portals_fun.folder, - } - + xCPs = torch.from_numpy(np.array(portals_fun.portals_parameters["model_parameters"]["radii_rho"])).to(dfT) """ *************************************************************************************************** @@ -156,25 +130,23 @@ def initializeProblem( *************************************************************************************************** """ + transport_parameters = portals_fun.portals_parameters["model_parameters"]["transport_parameters"] + + # Add folder and cold_start to the simulation options + transport_options = transport_parameters | {"folder": portals_fun.folder, "cold_start": False} + + target_options = portals_fun.portals_parameters["model_parameters"]["target_parameters"] + portals_fun.powerstate = STATEtools.powerstate( profiles, evolution_options={ - "ProfilePredicted": portals_fun.MODELparameters["ProfilesPredicted"], + "ProfilePredicted": portals_fun.portals_parameters["model_parameters"]["predicted_channels"], "rhoPredicted": xCPs, "impurityPosition": position_of_impurity, - "fineTargetsResolution": portals_fun.PORTALSparameters["fineTargetsResolution"], - }, - transport_options={ - "transport_evaluator": portals_fun.PORTALSparameters["transport_evaluator"], - "transport_evaluator_options": transport_evaluator_options, - }, - target_options={ - "target_evaluator": portals_fun.PORTALSparameters["target_evaluator"], - "target_evaluator_options": { - "TypeTarget": portals_fun.MODELparameters["Physics_options"]["TypeTarget"], - "target_evaluator_method": portals_fun.PORTALSparameters["target_evaluator_method"]}, }, - tensor_options = tensor_options + transport_options=transport_options, + target_options=target_options, + tensor_options=tensor_options ) # After resolution and corrections, store. @@ -185,13 +157,13 @@ def initializeProblem( # Store parameterization in dictCPs_base (to define later the relative variations) and modify profiles class with parameterized profiles dictCPs_base = {} - for name in portals_fun.MODELparameters["ProfilesPredicted"]: + for name in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]: dictCPs_base[name] = portals_fun.powerstate.update_var(name, var=None)[0, :] # Maybe it was provided from earlier run if start_from_folder is not None: dictCPs_base = grabPrevious(start_from_folder, dictCPs_base) - for name in portals_fun.MODELparameters["ProfilesPredicted"]: + for name in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]: _ = portals_fun.powerstate.update_var( name, var=dictCPs_base[name].unsqueeze(0) ) @@ -199,7 +171,7 @@ def initializeProblem( # Write this updated profiles class (with parameterized profiles) _ = portals_fun.powerstate.from_powerstate( write_input_gacode=FolderInitialization / "input.gacode", - postprocess_input_gacode=portals_fun.MODELparameters["applyCorrections"], + postprocess_input_gacode=portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["applyCorrections"], ) # Original complete targets @@ -214,22 +186,25 @@ def initializeProblem( powerstate_extra = STATEtools.powerstate( define_ranges_from_profiles, evolution_options={ - "ProfilePredicted": portals_fun.MODELparameters["ProfilesPredicted"], + "ProfilePredicted": portals_fun.portals_parameters["model_parameters"]["predicted_channels"], "rhoPredicted": xCPs, "impurityPosition": position_of_impurity, - "fineTargetsResolution": portals_fun.PORTALSparameters["fineTargetsResolution"], + "fineTargetsResolution": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["fineTargetsResolution"], }, target_options={ - "target_evaluator": portals_fun.PORTALSparameters["target_evaluator"], - "transport_evaluator_options": { - "TypeTarget": portals_fun.MODELparameters["Physics_options"]["TypeTarget"], - "target_evaluator_method": portals_fun.PORTALSparameters["target_evaluator_method"]}, + "target_evaluator": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["target_evaluator"], + "target_evaluator_options": { + "TypeTarget": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["TypeTarget"], + "target_evaluator_method": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["target_evaluator_method"], + "forceZeroParticleFlux": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["target_evaluator_options"]["forceZeroParticleFlux"], + "percent_error": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["percent_error"] + }, }, tensor_options = tensor_options ) dictCPs_base_extra = {} - for name in portals_fun.MODELparameters["ProfilesPredicted"]: + for name in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]: dictCPs_base_extra[name] = powerstate_extra.update_var(name, var=None)[0, :] dictCPs_base = dictCPs_base_extra @@ -285,7 +260,7 @@ def initializeProblem( elif ikey == "w0": var = "Mt" - for i in range(len(portals_fun.MODELparameters["RhoLocations"])): + for i in range(len(portals_fun.portals_parameters["model_parameters"]["radii_rho"])): ofs.append(f"{var}_tr_turb_{i+1}") ofs.append(f"{var}_tr_neoc_{i+1}") @@ -293,13 +268,13 @@ def initializeProblem( name_objectives.append(f"{var}Res_{i+1}") - if portals_fun.PORTALSparameters["turbulent_exchange_as_surrogate"]: - for i in range(len(portals_fun.MODELparameters["RhoLocations"])): + if portals_fun.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]: + for i in range(len(portals_fun.portals_parameters["model_parameters"]["radii_rho"])): ofs.append(f"Qie_tr_turb_{i+1}") name_transformed_ofs = [] for of in ofs: - if ("GZ" in of) and (portals_fun.PORTALSparameters["applyImpurityGammaTrick"]): + if ("GZ" in of) and (portals_fun.portals_parameters["main_parameters"]["applyImpurityGammaTrick"]): lab = f"{of} (GB MOD)" else: lab = f"{of} (GB)" @@ -328,14 +303,14 @@ def initializeProblem( # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Variables = {} - for ikey in portals_fun.PORTALSparameters["portals_transformation_variables"]: + for ikey in portals_fun.portals_parameters["main_parameters"]["portals_transformation_variables"]: Variables[ikey] = prepportals_transformation_variables(portals_fun, ikey) portals_fun.surrogate_parameters = { "transformationInputs": PORTALStools.input_transform_portals, "transformationOutputs": PORTALStools.output_transform_portals, "powerstate": portals_fun.powerstate, - "applyImpurityGammaTrick": portals_fun.PORTALSparameters["applyImpurityGammaTrick"], + "applyImpurityGammaTrick": portals_fun.portals_parameters["main_parameters"]["applyImpurityGammaTrick"], "surrogate_transformation_variables_alltimes": Variables, "surrogate_transformation_variables_lasttime": copy.deepcopy(Variables[list(Variables.keys())[-1]]), "parameters_combined": {}, @@ -343,9 +318,9 @@ def initializeProblem( def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValues=False): allOuts = portals_fun.optimization_options["problem_options"]["ofs"] - portals_transformation_variables = portals_fun.PORTALSparameters["portals_transformation_variables"][ikey] - portals_transformation_variables_trace = portals_fun.PORTALSparameters["portals_transformation_variables_trace"][ikey] - additional_params_in_surrogate = portals_fun.PORTALSparameters["additional_params_in_surrogate"] + portals_transformation_variables = portals_fun.portals_parameters["main_parameters"]["portals_transformation_variables"][ikey] + portals_transformation_variables_trace = portals_fun.portals_parameters["main_parameters"]["portals_transformation_variables_trace"][ikey] + additional_params_in_surrogate = portals_fun.portals_parameters["main_parameters"]["additional_params_in_surrogate"] Variables = {} for output in allOuts: @@ -378,17 +353,17 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue isAbsValFixed = False Variations = { - "aLte": "te" in portals_fun.MODELparameters["ProfilesPredicted"], - "aLti": "ti" in portals_fun.MODELparameters["ProfilesPredicted"], - "aLne": "ne" in portals_fun.MODELparameters["ProfilesPredicted"], - "aLw0": "w0" in portals_fun.MODELparameters["ProfilesPredicted"], - "te": ("te" in portals_fun.MODELparameters["ProfilesPredicted"]) + "aLte": "te" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "aLti": "ti" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "aLne": "ne" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "aLw0": "w0" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "te": ("te" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) and (not isAbsValFixed), - "ti": ("ti" in portals_fun.MODELparameters["ProfilesPredicted"]) + "ti": ("ti" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) and (not isAbsValFixed), - "ne": ("ne" in portals_fun.MODELparameters["ProfilesPredicted"]) + "ne": ("ne" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) and (not isAbsValFixed), - "w0": ("w0" in portals_fun.MODELparameters["ProfilesPredicted"]) + "w0": ("w0" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) and (not isAbsValFixed), } @@ -413,20 +388,20 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue isAbsValFixed = False Variations = { - "aLte": "te" in portals_fun.MODELparameters["ProfilesPredicted"], - "aLti": "ti" in portals_fun.MODELparameters["ProfilesPredicted"], - "aLne": "ne" in portals_fun.MODELparameters["ProfilesPredicted"], - "aLw0": "w0" in portals_fun.MODELparameters["ProfilesPredicted"], - "aLnZ": "nZ" in portals_fun.MODELparameters["ProfilesPredicted"], - "te": ("te" in portals_fun.MODELparameters["ProfilesPredicted"]) + "aLte": "te" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "aLti": "ti" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "aLne": "ne" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "aLw0": "w0" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "aLnZ": "nZ" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "te": ("te" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) and (not isAbsValFixed), - "ti": ("ti" in portals_fun.MODELparameters["ProfilesPredicted"]) + "ti": ("ti" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) and (not isAbsValFixed), - "ne": ("ne" in portals_fun.MODELparameters["ProfilesPredicted"]) + "ne": ("ne" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) and (not isAbsValFixed), - "w0": ("w0" in portals_fun.MODELparameters["ProfilesPredicted"]) + "w0": ("w0" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) and (not isAbsValFixed), - "nZ": ("nZ" in portals_fun.MODELparameters["ProfilesPredicted"]) + "nZ": ("nZ" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) and (not isAbsValFixed), } diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index 15384cac..05f47e40 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -140,13 +140,13 @@ def flux_match_surrogate( # Define transport calculation function as a surrogate model transport_options['transport_evaluator'] = transport_analytic.surrogate - transport_options['transport_evaluator_options'] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} + transport_options["transport_evaluator_options"] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} # Create powerstate with the same options as the original portals but with the new profiles powerstate = STATEtools.powerstate( profiles, evolution_options={ - "ProfilePredicted": step.surrogate_parameters["powerstate"].ProfilesPredicted, + "ProfilePredicted": step.surrogate_parameters["powerstate"].predicted_channels, "rhoPredicted": step.surrogate_parameters["powerstate"].plasma["rho"][0,1:], "impurityPosition": step.surrogate_parameters["powerstate"].impurityPosition, "fineTargetsResolution": step.surrogate_parameters["powerstate"].fineTargetsResolution, diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index a1e9a068..78f934e4 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -40,14 +40,14 @@ def PORTALSanalyzer_plotMetrics( plt.ion() fig = plt.figure(figsize=(15, 8)) - numprofs = len(self.ProfilesPredicted) + numprofs = len(self.predicted_channels) grid = plt.GridSpec(nrows=8, ncols=numprofs + 1, hspace=0.3, wspace=0.35) cont = 0 # Te - if "te" in self.ProfilesPredicted: + if "te" in self.predicted_channels: axTe = fig.add_subplot(grid[:4, cont]) axTe.set_title("Electron Temperature") axTe_g = fig.add_subplot(grid[4:6, cont]) @@ -56,7 +56,7 @@ def PORTALSanalyzer_plotMetrics( else: axTe = axTe_g = axTe_f = None - if "ti" in self.ProfilesPredicted: + if "ti" in self.predicted_channels: axTi = fig.add_subplot(grid[:4, cont]) axTi.set_title("Ion Temperature") axTi_g = fig.add_subplot(grid[4:6, cont]) @@ -66,7 +66,7 @@ def PORTALSanalyzer_plotMetrics( axTi = axTi_g = axTi_f = None - if "ne" in self.ProfilesPredicted: + if "ne" in self.predicted_channels: axne = fig.add_subplot(grid[:4, cont]) axne.set_title("Electron Density") axne_g = fig.add_subplot(grid[4:6, cont]) @@ -478,7 +478,7 @@ def PORTALSanalyzer_plotMetrics( ax.set_xticklabels([]) ax = axC - if "te" in self.ProfilesPredicted: + if "te" in self.predicted_channels: v = self.resTeM ax.plot( self.evaluations, @@ -489,7 +489,7 @@ def PORTALSanalyzer_plotMetrics( markersize=2, label=self.labelsFluxes["te"], ) - if "ti" in self.ProfilesPredicted: + if "ti" in self.predicted_channels: v = self.resTiM ax.plot( self.evaluations, @@ -500,7 +500,7 @@ def PORTALSanalyzer_plotMetrics( markersize=2, label=self.labelsFluxes["ti"], ) - if "ne" in self.ProfilesPredicted: + if "ne" in self.predicted_channels: v = self.resneM ax.plot( self.evaluations, @@ -511,7 +511,7 @@ def PORTALSanalyzer_plotMetrics( markersize=2, label=self.labelsFluxes["ne"], ) - if "nZ" in self.ProfilesPredicted: + if "nZ" in self.predicted_channels: v = self.resnZM ax.plot( self.evaluations, @@ -522,7 +522,7 @@ def PORTALSanalyzer_plotMetrics( markersize=2, label=self.labelsFluxes["nZ"], ) - if "w0" in self.ProfilesPredicted: + if "w0" in self.predicted_channels: v = self.resw0M ax.plot( self.evaluations, @@ -546,7 +546,7 @@ def PORTALSanalyzer_plotMetrics( ): if (indexUse is None) or (indexUse >= len(self.powerstates)): continue - if "te" in self.ProfilesPredicted: + if "te" in self.predicted_channels: v = self.resTeM ax.plot( [self.evaluations[indexUse]], @@ -555,7 +555,7 @@ def PORTALSanalyzer_plotMetrics( color=col, markersize=4, ) - if "ti" in self.ProfilesPredicted: + if "ti" in self.predicted_channels: v = self.resTiM ax.plot( [self.evaluations[indexUse]], @@ -564,7 +564,7 @@ def PORTALSanalyzer_plotMetrics( color=col, markersize=4, ) - if "ne" in self.ProfilesPredicted: + if "ne" in self.predicted_channels: v = self.resneM ax.plot( [self.evaluations[indexUse]], @@ -573,7 +573,7 @@ def PORTALSanalyzer_plotMetrics( color=col, markersize=4, ) - if "nZ" in self.ProfilesPredicted: + if "nZ" in self.predicted_channels: v = self.resnZM ax.plot( [self.evaluations[indexUse]], @@ -582,7 +582,7 @@ def PORTALSanalyzer_plotMetrics( color=col, markersize=4, ) - if "w0" in self.ProfilesPredicted: + if "w0" in self.predicted_channels: v = self.resw0M ax.plot( [self.evaluations[indexUse]], @@ -1128,7 +1128,7 @@ def PORTALSanalyzer_plotExpected( # ---- Plot - numprofs = len(self.ProfilesPredicted) + numprofs = len(self.predicted_channels) if numprofs <= 4: wspace = 0.3 @@ -1138,7 +1138,7 @@ def PORTALSanalyzer_plotExpected( grid = plt.GridSpec(nrows=4, ncols=numprofs, hspace=0.2, wspace=wspace) cont = 0 - if "te" in self.ProfilesPredicted: + if "te" in self.predicted_channels: axTe = fig.add_subplot(grid[0, cont]) axTe.set_title("Electron Temperature") axTe_g = fig.add_subplot(grid[1, cont], sharex=axTe) @@ -1147,7 +1147,7 @@ def PORTALSanalyzer_plotExpected( cont += 1 else: axTe = axTe_g = axTe_f = axTe_r = None - if "ti" in self.ProfilesPredicted: + if "ti" in self.predicted_channels: axTi = fig.add_subplot(grid[0, cont], sharex=axTe) axTi.set_title("Ion Temperature") axTi_g = fig.add_subplot(grid[1, cont], sharex=axTe) @@ -1156,7 +1156,7 @@ def PORTALSanalyzer_plotExpected( cont += 1 else: axTi = axTi_g = axTi_f = axTi_r = None - if "ne" in self.ProfilesPredicted: + if "ne" in self.predicted_channels: axne = fig.add_subplot(grid[0, cont], sharex=axTe) axne.set_title("Electron Density") axne_g = fig.add_subplot(grid[1, cont], sharex=axTe) @@ -1207,7 +1207,7 @@ def PORTALSanalyzer_plotExpected( rho = p.profiles["rho(-)"] roa = p.derived["roa"] - rhoVals = self.MODELparameters["RhoLocations"] + rhoVals = self.portals_parameters["model_parameters"]["radii_rho"] roaVals = np.interp(rhoVals, rho, roa) lastX = roaVals[-1] @@ -1395,7 +1395,7 @@ def PORTALSanalyzer_plotExpected( rho = self.profiles_next_new.profiles["rho(-)"] - rhoVals = self.MODELparameters["RhoLocations"] + rhoVals = self.portals_parameters["model_parameters"]["radii_rho"] roaVals = np.interp(rhoVals, rho, roa) p0 = self.powerstates[plotPoints[0]].profiles @@ -1852,7 +1852,7 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): fig4 = fn.add_figure(label="PROFILES Comparison", tab_color=fn_color) grid = plt.GridSpec( 2, - np.max([3, len(self.ProfilesPredicted)]), + np.max([3, len(self.predicted_channels)]), hspace=0.3, wspace=0.3, ) @@ -1892,10 +1892,10 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): axs4, color=colors[i], label=label, - lastRho=self.MODELparameters["RhoLocations"][-1], + lastRho=self.portals_parameters["model_parameters"]["radii_rho"][-1], alpha=alpha, useRoa=True, - RhoLocationsPlot=self.MODELparameters["RhoLocations"], + RhoLocationsPlot=self.portals_parameters["model_parameters"]["radii_rho"], plotImpurity=self.runWithImpurity, plotRotation=self.runWithRotation, autoscale=i == 3, @@ -1908,7 +1908,7 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): # ------------------------------------------------------- fig = fn.add_figure(label="Powerstate", tab_color=fn_color) - axs, axsM = STATEtools.add_axes_powerstate_plot(fig,num_kp=len(self.ProfilesPredicted)) + axs, axsM = STATEtools.add_axes_powerstate_plot(fig,num_kp=len(self.predicted_channels)) for indeces,c in zip(indecesPlot,["g","r","m"]): if indeces is not None: @@ -1926,7 +1926,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): fig = plt.figure() pps = np.max( - [3, len(self.ProfilesPredicted)] + [3, len(self.predicted_channels)] ) # Because plotGradients require at least Te, Ti, ne grid = plt.GridSpec(2, pps, hspace=0.3, wspace=0.3) axsR = [] @@ -1959,7 +1959,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): p.plot_gradients( axsR, color="b", - lastRho=self.MODELparameters["RhoLocations"][-1], + lastRho=self.portals_parameters["model_parameters"]["radii_rho"][-1], ms=ms, lw=1.0, label="Initial (#0)", @@ -1976,7 +1976,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): p.plot_gradients( axsR, color="r", - lastRho=self.MODELparameters["RhoLocations"][-1], + lastRho=self.portals_parameters["model_parameters"]["radii_rho"][-1], ms=ms, lw=0.3, ls="-o" if self.opt_fun.mitim_model.avoidPoints is not None else "-.o", @@ -1988,7 +1988,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): p.plot_gradients( axsR, color="g", - lastRho=self.MODELparameters["RhoLocations"][-1], + lastRho=self.portals_parameters["model_parameters"]["radii_rho"][-1], ms=ms, lw=1.0, label=f"Best (#{self.opt_fun.res.best_absolute_index})", @@ -2011,10 +2011,10 @@ def PORTALSanalyzer_plotModelComparison( if (fig is None) and (axs is None): plt.ion() - fig = plt.figure(figsize=(15, 6 if len(self.ProfilesPredicted)+int(self.PORTALSparameters["turbulent_exchange_as_surrogate"]) < 4 else 10)) + fig = plt.figure(figsize=(15, 6 if len(self.predicted_channels)+int(self.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]) < 4 else 10)) if axs is None: - if len(self.ProfilesPredicted)+int(self.PORTALSparameters["turbulent_exchange_as_surrogate"]) < 4: + if len(self.predicted_channels)+int(self.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]) < 4: axs = fig.subplots(ncols=3) else: axs = fig.subplots(ncols=3, nrows=2) @@ -2027,7 +2027,7 @@ def PORTALSanalyzer_plotModelComparison( metrics = {} # te - if 'te' in self.ProfilesPredicted: + if 'te' in self.predicted_channels: 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" @@ -2052,7 +2052,7 @@ def PORTALSanalyzer_plotModelComparison( cont += 1 # ti - if 'ti' in self.ProfilesPredicted: + if 'ti' in self.predicted_channels: 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" @@ -2077,7 +2077,7 @@ def PORTALSanalyzer_plotModelComparison( cont += 1 # ne - if 'ne' in self.ProfilesPredicted: + if 'ne' in self.predicted_channels: 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" @@ -2119,7 +2119,7 @@ def PORTALSanalyzer_plotModelComparison( cont += 1 - if "nZ" in self.ProfilesPredicted: + if "nZ" in self.predicted_channels: impurity_search = self.runWithImpurity_transport @@ -2165,7 +2165,7 @@ def PORTALSanalyzer_plotModelComparison( cont += 1 - if "w0" in self.ProfilesPredicted: + if "w0" in self.predicted_channels: if UseTGLFfull_x is not None: raise Exception("Momentum plot not implemented yet") # w0 @@ -2201,7 +2201,7 @@ def PORTALSanalyzer_plotModelComparison( cont += 1 - if self.PORTALSparameters["turbulent_exchange_as_surrogate"]: + if self.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]: if UseTGLFfull_x is not None: raise Exception("Turbulent exchange plot not implemented yet") # Sexch @@ -2419,8 +2419,8 @@ def varToReal(y, mitim_model): cont = 0 Qe, Qi, Ge, GZ, Mt = [], [], [], [], [] Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = [], [], [], [], [] - for prof in mitim_model.optimization_object.MODELparameters["ProfilesPredicted"]: - for rad in mitim_model.optimization_object.MODELparameters["RhoLocations"]: + for prof in mitim_model.optimization_object.portals_parameters["model_parameters"]["predicted_channels"]: + for rad in mitim_model.optimization_object.portals_parameters["model_parameters"]["radii_rho"]: if prof == "te": Qe.append(of[0, cont]) Qe_tar.append(cal[0, cont]) @@ -2493,7 +2493,7 @@ def plotVars( .plasma["roa"][0, 1:] .cpu() .cpu().numpy() - ) # mitim_model.optimization_object.MODELparameters['RhoLocations'] + ) # mitim_model.optimization_object.portals_parameters["model_parameters"]['radii_rho'] try: Qe, Qi, Ge, GZ, Mt, Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = varToReal( @@ -3254,7 +3254,7 @@ def plotFluxComparison( def produceInfoRanges( self_complete, bounds, axsR, label="", color="k", lw=0.2, alpha=0.05 ): - rhos = np.append([0], self_complete.MODELparameters["RhoLocations"]) + rhos = np.append([0], self_complete.portals_parameters["model_parameters"]["radii_rho"]) aLTe, aLTi, aLne, aLnZ, aLw0 = ( np.zeros((len(rhos), 2)), np.zeros((len(rhos), 2)), @@ -3275,19 +3275,19 @@ def produceInfoRanges( if f"aLw0_{i+1}" in bounds: aLw0[i + 1, :] = bounds[f"aLw0_{i+1}"] - X = torch.zeros(((len(rhos) - 1) * len(self_complete.MODELparameters["ProfilesPredicted"]), 2)) + X = torch.zeros(((len(rhos) - 1) * len(self_complete.portals_parameters["model_parameters"]["predicted_channels"]), 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.MODELparameters["ProfilesPredicted"]: + if "ne" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLne[1:, :]) cont += 1 - if "nZ" in self_complete.MODELparameters["ProfilesPredicted"]: + if "nZ" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLnZ[1:, :]) cont += 1 - if "w0" in self_complete.MODELparameters["ProfilesPredicted"]: + if "w0" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLw0[1:, :]) cont += 1 @@ -3338,7 +3338,7 @@ def produceInfoRanges( ) cont = 0 - if "ne" in self_complete.MODELparameters["ProfilesPredicted"]: + if "ne" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: GRAPHICStools.fillGraph( axsR[3 + cont + 1], powerstate.plasma["rho"][0], @@ -3361,7 +3361,7 @@ def produceInfoRanges( ) cont += 2 - if "nZ" in self_complete.MODELparameters["ProfilesPredicted"]: + if "nZ" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: GRAPHICStools.fillGraph( axsR[3 + cont + 1], powerstate.plasma["rho"][0], @@ -3384,7 +3384,7 @@ def produceInfoRanges( ) cont += 2 - if "w0" in self_complete.MODELparameters["ProfilesPredicted"]: + if "w0" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: GRAPHICStools.fillGraph( axsR[3 + cont + 1], powerstate.plasma["rho"][0], diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 0bf96325..df83ddb6 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -36,9 +36,8 @@ def __init__( - profiles_object: Object for gacode_state or others - evolution_options: - rhoPredicted: radial grid (MUST NOT CONTAIN ZERO, it will be added internally) - - ProfilesPredicted: list of profiles to predict + - predicted_channels: list of profiles to predict - impurityPosition: int = position of the impurity in the ions set - - fineTargetsResolution: int = resolution of the fine targets - transport_options: dictionary with transport_evaluator and transport_evaluator_options - target_options: dictionary with target_evaluator and target_evaluator_options ''' @@ -57,10 +56,11 @@ def __init__( "target_evaluator": targets_analytic.analytical_model, "target_evaluator_options": { "TypeTarget": 3, - "target_evaluator_method": "powerstate" + "target_evaluator_method": "powerstate", + "forceZeroParticleFlux": False }, } - + if tensor_options is None: tensor_options = { "dtype": torch.double, @@ -77,10 +77,10 @@ def __init__( self.target_options = target_options # Default options - self.ProfilesPredicted = evolution_options.get("ProfilePredicted", ["te", "ti", "ne"]) + self.predicted_channels = evolution_options.get("ProfilePredicted", ["te", "ti", "ne"]) self.impurityPosition = evolution_options.get("impurityPosition", 1) self.impurityPosition_transport = copy.deepcopy(self.impurityPosition) - self.fineTargetsResolution = evolution_options.get("fineTargetsResolution", None) + self.fineTargetsResolution = target_options.get("fineTargetsResolution", None) self.scaleIonDensities = evolution_options.get("scaleIonDensities", True) rho_vec = evolution_options.get("rhoPredicted", [0.2, 0.4, 0.6, 0.8]) @@ -96,7 +96,7 @@ def _ensure_ne_before_nz(lst): # Swap "ne" and "nZ" positions lst[ne_index], lst[nz_index] = lst[nz_index], lst[ne_index] return lst - self.ProfilesPredicted = _ensure_ne_before_nz(self.ProfilesPredicted) + self.predicted_channels = _ensure_ne_before_nz(self.predicted_channels) # Default type and device tensor self.dfT = torch.randn((2, 2), **tensor_options) @@ -105,7 +105,7 @@ def _ensure_ne_before_nz(lst): Potential profiles to evolve (aLX) and their corresponding flux matching ------------------------------------------------------------------------ The order in the P and P_tr (and therefore the source S) - tensors will be the same as in self.ProfilesPredicted + tensors will be the same as in self.predicted_channels ''' self.profile_map = { "te": ("QeMWm2", "QeMWm2_tr"), @@ -131,7 +131,7 @@ def _ensure_ne_before_nz(lst): ), torch.Tensor().to(self.dfT) self.labelsFM = [] - for profile in self.ProfilesPredicted: + for profile in self.predicted_channels: self.labelsFM.append([f'aL{profile}', list(self.profile_map[profile])[0], list(self.profile_map[profile])[1]]) # ------------------------------------------------------------------------------------- @@ -297,7 +297,7 @@ def calculate( self.calculateProfileFunctions() # 3. Sources and sinks (populates components and Pe,Pi,...) - relative_error_assumed = self.transport_options["transport_evaluator_options"].get("percentError", [5, 1, 0.5])[-1] + relative_error_assumed = self.target_options["target_evaluator_options"]["percent_error"] self.calculateTargets(relative_error_assumed=relative_error_assumed) # Calculate targets based on powerstate functions (it may be overwritten in next step, if chosen) # 4. Turbulent and neoclassical transport (populates components and Pe_tr,Pi_tr,...) @@ -321,7 +321,7 @@ def modify(self, X): self.Xcurrent = X numeach = self.plasma["rho"].shape[1] - 1 - for c, i in enumerate(self.ProfilesPredicted): + for c, i in enumerate(self.predicted_channels): if X is not None: aLx_before = self.plasma[f"aL{i}"][:, 1:].clone() @@ -406,11 +406,11 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): # Concatenate the input gradients x0 = torch.Tensor().to(self.plasma["aLte"]) - for c, i in enumerate(self.ProfilesPredicted): + for c, i in enumerate(self.predicted_channels): x0 = torch.cat((x0, self.plasma[f"aL{i}"][:, 1:].detach()), dim=1) # Make sure is properly batched - x0 = x0.view((self.plasma["rho"].shape[0],(self.plasma["rho"].shape[1] - 1) * len(self.ProfilesPredicted),)) + x0 = x0.view((self.plasma["rho"].shape[0],(self.plasma["rho"].shape[1] - 1) * len(self.predicted_channels),)) # Optimize x_best,Yopt, Xopt, metric_history = solver_fun(evaluator,x0, bounds=self.bounds_current,solver_options=solver_options) @@ -448,17 +448,17 @@ def plot(self, axs=None, axsRes=None, axsMetrics=None, figs=None, fn=None,c="r", figMain = fn.add_figure(label="PowerState", tab_color='r') # Optimization figOpt = fn.add_figure(label="Optimization", tab_color='r') - grid = plt.GridSpec(2, 1+len(self.ProfilesPredicted), hspace=0.3, wspace=0.3) + grid = plt.GridSpec(2, 1+len(self.predicted_channels), hspace=0.3, wspace=0.3) axsRes = [figOpt.add_subplot(grid[:, 0])] - for i in range(len(self.ProfilesPredicted)): + for i in range(len(self.predicted_channels)): for j in range(2): axsRes.append(figOpt.add_subplot(grid[j, i+1])) # Profiles figs = state_plotting.add_figures(fn, tab_color='b') - axs, axsMetrics = add_axes_powerstate_plot(figMain, num_kp = len(self.ProfilesPredicted)) + axs, axsMetrics = add_axes_powerstate_plot(figMain, num_kp = len(self.predicted_channels)) else: axsNotGiven = False @@ -709,7 +709,7 @@ def calculateTargets(self, relative_error_assumed=1.0): # Merge targets, calculate errors and normalize targets.postprocessing( relative_error_assumed=relative_error_assumed, - forceZeroParticleFlux=self.transport_options["transport_evaluator_options"].get("forceZeroParticleFlux", False)) + forceZeroParticleFlux=self.target_options["target_evaluator_options"]["forceZeroParticleFlux"]) def calculateTransport( self, nameRun="test", folder="~/scratch/", evaluation_number=0): @@ -746,7 +746,7 @@ def _concatenate_flux(plasma, profile_key, flux_key): self.plasma["P"], self.plasma["P_tr"] = torch.Tensor().to(self.plasma["QeMWm2"]), torch.Tensor().to(self.plasma["QeMWm2"]) - for profile in self.ProfilesPredicted: + for profile in self.predicted_channels: _concatenate_flux(self.plasma, *self.profile_map[profile]) self.plasma["S"] = self.plasma["P"] - self.plasma["P_tr"] diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py index f49d41d5..77288306 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -70,7 +70,7 @@ def evaluate(self): """ X = torch.Tensor() - for prof in self.powerstate.ProfilesPredicted: + for prof in self.powerstate.predicted_channels: X = torch.cat((X,self.powerstate.plasma['aL'+prof][:,1:]),axis=1) _, Q, _, _ = self.powerstate.transport_options["transport_evaluator_options"]["flux_fun"](X) @@ -85,6 +85,6 @@ def evaluate(self): "w0": "MtJm2", } - for c, i in enumerate(self.powerstate.ProfilesPredicted): + for c, i in enumerate(self.powerstate.predicted_channels): self.powerstate.plasma[f"{quantities[i]}_tr"] = torch.cat((torch.tensor([[0.0]]),Q[:, numeach * c : numeach * (c + 1)]),dim=1) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 8cc16c1a..4837ba14 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -97,13 +97,13 @@ def _cgyro_trick(self,FolderEvaluation_TGYRO): # ************************************************************************************************************************** evaluateCGYRO( - self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]["PORTALSparameters"], - self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]["folder"], + self.powerstate.transport_options["transport_evaluator_options"]["transport_parameters"], + self.powerstate.transport_options["folder"], self.evaluation_number, FolderEvaluation_TGYRO, self.file_profs, self.powerstate.plasma["roa"][0,1:], - self.powerstate.ProfilesPredicted, + self.powerstate.predicted_channels, ) # Make tensors @@ -140,7 +140,7 @@ def _print_info(self): The CGYRO file must use particle flux. Convective transformation occurs later """ -def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmodified_profiles, radii, ProfilesPredicted): +def evaluateCGYRO(transport, folder, numPORTALS, FolderEvaluation, unmodified_profiles, radii, predicted_channels, impurityPosition): print("\n ** CGYRO evaluation of fluxes has been requested before passing information to the STRATEGY module **",typeMsg="i",) if isinstance(numPORTALS, int): @@ -149,7 +149,7 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod # ------------------------------------------------------------------------------------------------ # Harcoded # ------------------------------------------------------------------------------------------------ - if PORTALSparameters['hardCodedCGYRO'] is not None: + if transport['hardCodedCGYRO'] is not 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 @@ -160,10 +160,10 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod """ - train_sep = PORTALSparameters["hardCodedCGYRO"]["train_sep"] - start_num = PORTALSparameters["hardCodedCGYRO"]["start_num"] - last_one = PORTALSparameters["hardCodedCGYRO"]["last_one"] - trick_hardcoded_f = PORTALSparameters["hardCodedCGYRO"]["trick_hardcoded_f"] + train_sep = transport["hardCodedCGYRO"]["train_sep"] + start_num = transport["hardCodedCGYRO"]["start_num"] + last_one = transport["hardCodedCGYRO"]["last_one"] + trick_hardcoded_f = transport["hardCodedCGYRO"]["trick_hardcoded_f"] else: train_sep = None start_num = None @@ -171,19 +171,10 @@ def evaluateCGYRO(PORTALSparameters, folder, numPORTALS, FolderEvaluation, unmod trick_hardcoded_f = None # ------------------------------------------------------------------------------------------------ - minErrorPercent = PORTALSparameters["percentError_stable"] - Qi_criterion_stable = PORTALSparameters["Qi_criterion_stable"] + minErrorPercent = transport["percentError_stable"] + Qi_criterion_stable = transport["Qi_criterion_stable"] - try: - impurityPosition = MITIMstate.impurity_location(PROFILEStools.gacode_state(unmodified_profiles), PORTALSparameters["ImpurityOfInterest"]) - except ValueError: - if 'nZ' in ProfilesPredicted: - raise ValueError(f"Impurity {PORTALSparameters['ImpurityOfInterest']} not found in the profiles and needed for CGYRO evaluation") - else: - impurityPosition = 0 - print(f'\t- Impurity location not found. Using hardcoded value of {impurityPosition}') - - OriginalFimp = PORTALSparameters["fImp_orig"] + OriginalFimp = portals_parameters["main_parameters"]["fImp_orig"] cgyroing_file = ( lambda file_cgyro, numPORTALS_this=0: cgyroing( diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index f81a19e6..7eda3057 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -12,24 +12,21 @@ def __init__(self, powerstate, **kwargs): # Do not hook here def evaluate_turbulence(self): - + + transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + cold_start = self.powerstate.transport_options["cold_start"] + # Run base TGLF always, to keep track of discrepancies! -------------------------------------- - self.powerstate.transport_options["transport_evaluator_options"]["use_tglf_scan_trick"] = None + transport_evaluator_options["tglf"]["use_scan_trick_for_stds"] = None self._evaluate_tglf() # -------------------------------------------------------------------------------------------- rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] - transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + simulation_options_cgyro = transport_evaluator_options["cgyro"] - cold_start = transport_evaluator_options.get("cold_start", False) + run_type = simulation_options_cgyro["run_type"] - run_type = transport_evaluator_options["MODELparameters"]["transport_model"]["run_type"] - CGYROsettings = transport_evaluator_options["MODELparameters"]["transport_model"]["CGYROsettings"] - extraOptionsCGYRO = transport_evaluator_options["MODELparameters"]["transport_model"]["extraOptionsCGYRO"] - every_n_minutesCGYRO = transport_evaluator_options["MODELparameters"]["transport_model"]["every_n_minutesCGYRO"] - tminCGYRO = transport_evaluator_options["MODELparameters"]["transport_model"]["tminCGYRO"] - # ------------------------------------------------------------------------------------------------------------------------ # Prepare CGYRO object # ------------------------------------------------------------------------------------------------------------------------ @@ -46,21 +43,20 @@ def evaluate_turbulence(self): _ = cgyro.run( 'base_cgyro', run_type = run_type, - code_settings=CGYROsettings, - extraOptions=extraOptionsCGYRO, cold_start=cold_start, forceIfcold_start=True, + **simulation_options_cgyro["run"] ) if run_type in ['normal', 'submit']: if run_type in ['submit']: - cgyro.check(every_n_minutes=every_n_minutesCGYRO) + cgyro.check(every_n_minutes=10) cgyro.fetch() cgyro.read( label='base_cgyro', - tmin = tminCGYRO + **simulation_options_cgyro["read"] ) # ------------------------------------------------------------------------------------------------------------------------ diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 1df7c115..6fbcc80c 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -26,24 +26,23 @@ def evaluate_turbulence(self): # Have it separate such that I can call it from the CGYRO class but without the decorator def _evaluate_tglf(self): + transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + + cold_start = self.powerstate.transport_options["cold_start"] + # ------------------------------------------------------------------------------------------------------------------------ # Grab options from powerstate # ------------------------------------------------------------------------------------------------------------------------ - transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] - - TGLFsettings = transport_evaluator_options["MODELparameters"]["transport_model"]["TGLFsettings"] - extraOptions = transport_evaluator_options["MODELparameters"]["transport_model"]["extraOptionsTGLF"] - - Qi_includes_fast = transport_evaluator_options.get("Qi_includes_fast",False) - launchMODELviaSlurm = transport_evaluator_options.get("launchMODELviaSlurm", False) - cold_start = transport_evaluator_options.get("cold_start", False) - provideTurbulentExchange = transport_evaluator_options.get("TurbulentExchange", False) - percentError = transport_evaluator_options.get("percentError", [5, 1, 0.5]) - use_tglf_scan_trick = transport_evaluator_options.get("use_tglf_scan_trick", None) - cores_per_tglf_instance = transport_evaluator_options.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) - keep_tglf_files = transport_evaluator_options.get("keep_tglf_files", "minimal") + simulation_options_tglf = transport_evaluator_options["tglf"] + Qi_includes_fast = simulation_options_tglf["Qi_includes_fast"] + launchMODELviaSlurm = simulation_options_tglf["launchEvaluationsAsSlurmJobs"] + use_tglf_scan_trick = simulation_options_tglf["use_scan_trick_for_stds"] + cores_per_tglf_instance = simulation_options_tglf["cores_per_tglf_instance"] + keep_tglf_files = simulation_options_tglf["keep_files"] + percent_error = simulation_options_tglf["percent_error"] + # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) impurityPosition = self.powerstate.impurityPosition_transport @@ -67,8 +66,6 @@ def _evaluate_tglf(self): tglf.run( 'base_tglf', - code_settings=TGLFsettings, - extraOptions=extraOptions, ApplyCorrections=False, launchSlurm= launchMODELviaSlurm, cold_start= cold_start, @@ -79,10 +76,14 @@ def _evaluate_tglf(self): "minutes": 2, }, attempts_execution=2, - only_minimal_files=keep_tglf_files in ['minimal'] + only_minimal_files=keep_tglf_files in ['minimal'], + **simulation_options_tglf["run"] ) - tglf.read(label='base',require_all_files=False) + tglf.read( + label='base', + require_all_files=False, + **simulation_options_tglf["read"]) # Grab values Qe = np.array([tglf.results['base']['TGLFout'][i].Qe for i in range(len(rho_locations))]) @@ -113,7 +114,7 @@ def _evaluate_tglf(self): # ******************************************************************* Flux_mean = Flux_base - Flux_std = abs(Flux_mean)*percentError[0]/100.0 + Flux_std = abs(Flux_mean)*percent_error/100.0 else: @@ -124,10 +125,8 @@ def _evaluate_tglf(self): Flux_mean, Flux_std = _run_tglf_uncertainty_model( tglf, rho_locations, - self.powerstate.ProfilesPredicted, + self.powerstate.predicted_channels, Flux_base = Flux_base, - TGLFsettings=TGLFsettings, - extraOptionsTGLF=extraOptions, impurityPosition=impurityPosition, delta = use_tglf_scan_trick, cold_start=cold_start, @@ -135,7 +134,8 @@ def _evaluate_tglf(self): cores_per_tglf_instance=cores_per_tglf_instance, launchMODELviaSlurm=launchMODELviaSlurm, Qi_includes_fast=Qi_includes_fast, - only_minimal_files=keep_tglf_files in ['minimal', 'base'] + only_minimal_files=keep_tglf_files in ['minimal', 'base'], + **simulation_options_tglf["run"] ) self._raise_warnings(tglf, rho_locations, Qi_includes_fast) @@ -159,22 +159,24 @@ def _evaluate_tglf(self): self.MtGB_turb = Flux_mean[4] self.MtGB_turb_stds = Flux_std[4] - self.QieGB_turb = Flux_mean[5] if provideTurbulentExchange else Flux_mean[5]*0.0 - self.QieGB_turb_stds = Flux_std[5] if provideTurbulentExchange else Flux_std[5]*0.0 + self.QieGB_turb = Flux_mean[5] + self.QieGB_turb_stds = Flux_std[5] return tglf @IOtools.hook_method(after=partial(TRANSPORTtools.write_json, file_name = 'fluxes_neoc.json', suffix= 'neoc')) def evaluate_neoclassical(self): + transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + # ------------------------------------------------------------------------------------------------------------------------ # Grab options from powerstate # ------------------------------------------------------------------------------------------------------------------------ - - transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + simulation_options_neo = transport_evaluator_options["neo"] + percent_error = simulation_options_neo["percent_error"] cold_start = transport_evaluator_options.get("cold_start", False) - percentError = transport_evaluator_options.get("percentError", [5, 1, 0.5]) + impurityPosition = self.powerstate.impurityPosition_transport # ------------------------------------------------------------------------------------------------------------------------ @@ -195,9 +197,12 @@ def evaluate_neoclassical(self): 'base_neo', cold_start=cold_start, forceIfcold_start=True, + **simulation_options_neo["run"] ) - neo.read(label='base') + neo.read( + label='base', + **simulation_options_neo["read"]) Qe = np.array([neo.results['base']['NEOout'][i].Qe for i in range(len(rho_locations))]) Qi = np.array([neo.results['base']['NEOout'][i].Qi for i in range(len(rho_locations))]) @@ -216,11 +221,11 @@ def evaluate_neoclassical(self): self.MtGB_neoc = Mt # Uncertainties is just a percent of the value - self.QeGB_neoc_stds = abs(Qe) * percentError[1]/100.0 - self.QiGB_neoc_stds = abs(Qi) * percentError[1]/100.0 - self.GeGB_neoc_stds = abs(Ge) * percentError[1]/100.0 - self.GZGB_neoc_stds = abs(GZ) * percentError[1]/100.0 - self.MtGB_neoc_stds = abs(Mt) * percentError[1]/100.0 + self.QeGB_neoc_stds = abs(Qe) * percent_error/100.0 + self.QiGB_neoc_stds = abs(Qi) * percent_error/100.0 + self.GeGB_neoc_stds = abs(Ge) * percent_error/100.0 + self.GZGB_neoc_stds = abs(GZ) * percent_error/100.0 + self.MtGB_neoc_stds = abs(Mt) * percent_error/100.0 # No neoclassical exchange self.QieGB_neoc = Qe * 0.0 @@ -230,8 +235,8 @@ def evaluate_neoclassical(self): def _profiles_to_store(self): - if "extra_params" in self.powerstate.transport_options["transport_evaluator_options"] and "folder" in self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]: - whereFolder = IOtools.expandPath(self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") + if "folder" in self.powerstate.transport_options["transport_evaluator_options"]: + whereFolder = IOtools.expandPath(self.powerstate.transport_options["folder"] / "Outputs" / "portals_profiles") if not whereFolder.exists(): IOtools.askNewFolder(whereFolder) @@ -266,10 +271,10 @@ def _raise_warnings(self, tglf, rho_locations, Qi_includes_fast): def _run_tglf_uncertainty_model( tglf, rho_locations, - ProfilesPredicted, + predicted_channels, Flux_base = None, - TGLFsettings=None, - extraOptionsTGLF=None, + code_settings=None, + extraOptions=None, impurityPosition=1, delta=0.02, minimum_abs_gradient=0.005, # This is 0.5% of aLx=1.0, to avoid extremely small scans when, for example, having aLn ~ 0.0 @@ -286,7 +291,7 @@ def _run_tglf_uncertainty_model( # Prepare scan variables_to_scan = [] - for i in ProfilesPredicted: + for i in predicted_channels: if i == 'te': variables_to_scan.append('RLTS_1') if i == 'ti': variables_to_scan.append('RLTS_2') if i == 'ne': variables_to_scan.append('RLNS_1') @@ -294,11 +299,11 @@ def _run_tglf_uncertainty_model( if i == 'w0': variables_to_scan.append('VEXB_SHEAR') #TODO: is this correct? or VPAR_SHEAR? #TODO: Only if that parameter is changing at that location - if 'te' in ProfilesPredicted or 'ti' in ProfilesPredicted: + if 'te' in predicted_channels or 'ti' in predicted_channels: variables_to_scan.append('TAUS_2') - if 'te' in ProfilesPredicted or 'ne' in ProfilesPredicted: + if 'te' in predicted_channels or 'ne' in predicted_channels: variables_to_scan.append('XNUE') - if 'te' in ProfilesPredicted or 'ne' in ProfilesPredicted: + if 'te' in predicted_channels or 'ne' in predicted_channels: variables_to_scan.append('BETAE') relative_scan = [1-delta, 1+delta] @@ -328,8 +333,8 @@ def _run_tglf_uncertainty_model( variablesDrives = variables_to_scan, varUpDown = relative_scan, minimum_delta_abs = minimum_delta_abs, - TGLFsettings = TGLFsettings, - extraOptions = extraOptionsTGLF, + TGLFsettings = code_settings, + extraOptions = extraOptions, ApplyCorrections = False, add_baseline_to = 'none', cold_start=cold_start, diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index 5ac4e2d4..bcd0bed7 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -34,16 +34,14 @@ def evaluate_turbulence(self): def _evaluate_tglf_neo(self): - transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + transport = self.powerstate.transport_options["transport_evaluator_options"]["transport_parameters"] + cold_start = self.powerstate.transport_options["transport_evaluator_options"]["cold_start"] - MODELparameters = transport_evaluator_options.get("MODELparameters",None) - Qi_includes_fast = transport_evaluator_options.get("Qi_includes_fast",False) - launchMODELviaSlurm = transport_evaluator_options.get("launchMODELviaSlurm", False) - cold_start = transport_evaluator_options.get("cold_start", False) - provideTurbulentExchange = transport_evaluator_options.get("TurbulentExchange", False) - percentError = transport_evaluator_options.get("percentError", [5, 1, 0.5]) - use_tglf_scan_trick = transport_evaluator_options.get("use_tglf_scan_trick", None) - cores_per_tglf_instance = transport_evaluator_options.get("extra_params", {}).get('PORTALSparameters', {}).get("cores_per_tglf_instance", 1) + Qi_includes_fast = transport["Qi_includes_fast"] + launchMODELviaSlurm = transport["launchMODELviaSlurm"] + percent_error = transport["percent_error"] + cores_per_tglf_instance = transport["cores_per_tglf_instance"] + use_tglf_scan_trick = transport["use_scan_trick_for_stds"] # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) impurityPosition = self.powerstate.impurityPosition_transport #transport_evaluator_options.get("impurityPosition", 1) @@ -69,13 +67,13 @@ def _evaluate_tglf_neo(self): special_radii=rho_locations, iterations=0, PredictionSet=[ - int("te" in self.powerstate.ProfilesPredicted), - int("ti" in self.powerstate.ProfilesPredicted), - int("ne" in self.powerstate.ProfilesPredicted), + int("te" in self.powerstate.predicted_channels), + int("ti" in self.powerstate.predicted_channels), + int("ne" in self.powerstate.predicted_channels), ], - TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], - extraOptionsTGLF=MODELparameters["transport_model"]["extraOptionsTGLF"], - TGYRO_physics_options=MODELparameters["Physics_options"], + TGLFsettings=transport["transport_evaluator_options"]["TGLFsettings"], + extraOptionsTGLF=transport["transport_evaluator_options"]["extraOptionsTGLF"], + TGYRO_physics_options=portals_parameters["model_parameters"]["Physics_options"], launchSlurm=launchMODELviaSlurm, minutesJob=5, forcedName=self.name, @@ -100,12 +98,11 @@ def _evaluate_tglf_neo(self): tgyro, "tglf_neo_original", rho_locations, - self.powerstate.ProfilesPredicted, + self.powerstate.predicted_channels, self.folder / "tglf_neo", - percentError, + percent_error, impurityPosition=impurityPosition, Qi_includes_fast=Qi_includes_fast, - provideTurbulentExchange=provideTurbulentExchange, use_tglf_scan_trick = use_tglf_scan_trick, cold_start=cold_start, extra_name = self.name, @@ -115,44 +112,16 @@ def _evaluate_tglf_neo(self): # Read again to capture errors tgyro.read(label="tglf_neo", folder=self.folder / "tglf_neo") - # $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ - # Run TGLF standalone --> In preparation for the transition - # $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ - - # from mitim_tools.gacode_tools import TGLFtools - # tglf = TGLFtools.TGLF(rhos=rho_locations) - # _ = tglf.prep_using_tgyro( - # self.folder / 'stds', - # inputgacode=self.file_profs, - # recalculate_ptot=False, # Use what's in the input.gacode, which is what PORTALS TGYRO does - # cold_start=cold_start) - - # tglf.run( - # subfolder="tglf_neo_original", - # TGLFsettings=MODELparameters["transport_model"]["TGLFsettings"], - # cold_start=cold_start, - # forceIfcold_start=True, - # extraOptions=MODELparameters["transport_model"]["extraOptionsTGLF"], - # launchSlurm=launchMODELviaSlurm, - # slurm_setup={"cores": 4, "minutes": 1}, - # ) - - # tglf.read(label="tglf_neo_original") - - # $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ - # $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ - return tgyro def _postprocess_tgyro(self, tgyro, label): - transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] - - Qi_includes_fast = transport_evaluator_options.get("Qi_includes_fast",False) - UseFineGridTargets = transport_evaluator_options.get("UseFineGridTargets", False) - provideTurbulentExchange = transport_evaluator_options.get("TurbulentExchange", False) - OriginalFimp = transport_evaluator_options.get("OriginalFimp", 1.0) - forceZeroParticleFlux = transport_evaluator_options.get("forceZeroParticleFlux", False) + portals_parameters = self.powerstate.transport_options["transport_evaluator_options"]["portals_parameters"] + + Qi_includes_fast = portals_parameters["model_parameters"]["transport_parameters"]["Qi_includes_fast"] + UseFineGridTargets = portals_parameters["main_parmaters"]["UseFineGridTargets"] + OriginalFimp = portals_parameters["main_parmaters"]["OriginalFimp"] + forceZeroParticleFlux = portals_parameters["model_parameters"]["target_parameters"]["target_evaluator_options"]["forceZeroParticleFlux"] # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) impurityPosition = self.powerstate.impurityPosition_transport @@ -166,8 +135,7 @@ def _postprocess_tgyro(self, tgyro, label): UseFineGridTargets=UseFineGridTargets, OriginalFimp=OriginalFimp, forceZeroParticleFlux=forceZeroParticleFlux, - provideTurbulentExchange=provideTurbulentExchange, - provideTargets=self.powerstate.target_options['target_evaluator_options']['target_evaluator_method'] == "tgyro", + provideTargets=self.powerstate.target_options['target_evaluator_options']["target_evaluator_method"] == "tgyro", ) tgyro.results["use"] = tgyro.results[label] @@ -188,8 +156,8 @@ def _postprocess_tgyro(self, tgyro, label): def _profiles_to_store(self): - if "extra_params" in self.powerstate.transport_options["transport_evaluator_options"] and "folder" in self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]: - whereFolder = IOtools.expandPath(self.powerstate.transport_options["transport_evaluator_options"]["extra_params"]["folder"] / "Outputs" / "portals_profiles") + if "folder" in self.powerstate.transport_options["transport_evaluator_options"]: + whereFolder = IOtools.expandPath(self.powerstate.transport_options["folder"] / "Outputs" / "portals_profiles") if not whereFolder.exists(): IOtools.askNewFolder(whereFolder) @@ -204,7 +172,7 @@ def _profiles_to_store(self): def tglf_scan_trick( tglf, rho_locations, - ProfilesPredicted, + predicted_channels, impurityPosition=1, Qi_includes_fast=False, delta=0.02, @@ -219,7 +187,7 @@ def tglf_scan_trick( # Prepare scan variables_to_scan = [] - for i in ProfilesPredicted: + for i in predicted_channels: if i == 'te': variables_to_scan.append('RLTS_1') if i == 'ti': variables_to_scan.append('RLTS_2') if i == 'ne': variables_to_scan.append('RLNS_1') @@ -227,11 +195,11 @@ def tglf_scan_trick( if i == 'w0': variables_to_scan.append('VEXB_SHEAR') #TODO: is this correct? or VPAR_SHEAR? #TODO: Only if that parameter is changing at that location - if 'te' in ProfilesPredicted or 'ti' in ProfilesPredicted: + if 'te' in predicted_channels or 'ti' in predicted_channels: variables_to_scan.append('TAUS_2') - if 'te' in ProfilesPredicted or 'ne' in ProfilesPredicted: + if 'te' in predicted_channels or 'ne' in predicted_channels: variables_to_scan.append('XNUE') - if 'te' in ProfilesPredicted or 'ne' in ProfilesPredicted: + if 'te' in predicted_channels or 'ne' in predicted_channels: variables_to_scan.append('BETAE') relative_scan = [1-delta, 1+delta] @@ -343,10 +311,9 @@ def curateTGYROfiles( tgyroObject, label, rho_locations, - ProfilesPredicted, + predicted_channels, folder, - percentError, - provideTurbulentExchange=False, + percent_error, impurityPosition=1, Qi_includes_fast=False, use_tglf_scan_trick=None, @@ -359,8 +326,8 @@ def curateTGYROfiles( tgyro = tgyroObject.results[label] # Determine NEO and Target errors - relativeErrorNEO = percentError[1] / 100.0 - relativeErrorTAR = percentError[2] / 100.0 + relativeErrorNEO = percent_error[1] / 100.0 + relativeErrorTAR = percent_error[2] / 100.0 # Grab fluxes from TGYRO Qe = tgyro.Qe_sim_turb[0, 1:] @@ -380,7 +347,7 @@ def curateTGYROfiles( Flux_base, Flux_mean, Flux_std = tglf_scan_trick( tglfObject, rho_locations, - ProfilesPredicted, + predicted_channels, impurityPosition=impurityPosition, Qi_includes_fast=Qi_includes_fast, delta = use_tglf_scan_trick, @@ -405,27 +372,27 @@ def curateTGYROfiles( Mt_tgyro = tgyro.Mt_sim_turb[0, 1:] Pexch_tgyro = tgyro.EXe_sim_turb[0, 1:] - Qe_err = np.abs( (Qe_base - Qe_tgyro) / Qe_tgyro ) if 'te' in ProfilesPredicted else np.zeros_like(Qe_base) - Qi_err = np.abs( (Qi_base - Qi_tgyro) / Qi_tgyro ) if 'ti' in ProfilesPredicted else np.zeros_like(Qi_base) - Ge_err = np.abs( (Ge_base - Ge_tgyro) / Ge_tgyro ) if 'ne' in ProfilesPredicted else np.zeros_like(Ge_base) - GZ_err = np.abs( (GZ_base - GZ_tgyro) / GZ_tgyro ) if 'nZ' in ProfilesPredicted else np.zeros_like(GZ_base) - Mt_err = np.abs( (Mt_base - Mt_tgyro) / Mt_tgyro ) if 'w0' in ProfilesPredicted else np.zeros_like(Mt_base) - Pexch_err = np.abs( (Pexch - Pexch_tgyro) / Pexch_tgyro ) if provideTurbulentExchange else np.zeros_like(Pexch) + Qe_err = np.abs( (Qe_base - Qe_tgyro) / Qe_tgyro ) if 'te' in predicted_channels else np.zeros_like(Qe_base) + Qi_err = np.abs( (Qi_base - Qi_tgyro) / Qi_tgyro ) if 'ti' in predicted_channels else np.zeros_like(Qi_base) + Ge_err = np.abs( (Ge_base - Ge_tgyro) / Ge_tgyro ) if 'ne' in predicted_channels else np.zeros_like(Ge_base) + GZ_err = np.abs( (GZ_base - GZ_tgyro) / GZ_tgyro ) if 'nZ' in predicted_channels else np.zeros_like(GZ_base) + Mt_err = np.abs( (Mt_base - Mt_tgyro) / Mt_tgyro ) if 'w0' in predicted_channels else np.zeros_like(Mt_base) + Pexch_err = np.abs( (Pexch - Pexch_tgyro) / Pexch_tgyro ) F_err = np.concatenate((Qe_err, Qi_err, Ge_err, GZ_err, Mt_err, Pexch_err)) if F_err.max() > check_coincidence_thr: print(f"\t- TGLF scans are not consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%, in quantity:",typeMsg="w") - if ('te' in ProfilesPredicted) and Qe_err.max() > check_coincidence_thr: + if ('te' in predicted_channels) and Qe_err.max() > check_coincidence_thr: print('\t\t* Qe:',Qe_err) - if ('ti' in ProfilesPredicted) and Qi_err.max() > check_coincidence_thr: + if ('ti' in predicted_channels) and Qi_err.max() > check_coincidence_thr: print('\t\t* Qi:',Qi_err) - if ('ne' in ProfilesPredicted) and Ge_err.max() > check_coincidence_thr: + if ('ne' in predicted_channels) and Ge_err.max() > check_coincidence_thr: print('\t\t* Ge:',Ge_err) - if ('nZ' in ProfilesPredicted) and GZ_err.max() > check_coincidence_thr: + if ('nZ' in predicted_channels) and GZ_err.max() > check_coincidence_thr: print('\t\t* GZ:',GZ_err) - if ('w0' in ProfilesPredicted) and Mt_err.max() > check_coincidence_thr: + if ('w0' in predicted_channels) and Mt_err.max() > check_coincidence_thr: print('\t\t* Mt:',Mt_err) - if provideTurbulentExchange and Pexch_err.max() > check_coincidence_thr: + if Pexch_err.max() > check_coincidence_thr: print('\t\t* Pexch:',Pexch_err) else: print(f"\t- TGLF scans are consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%") @@ -446,7 +413,7 @@ def curateTGYROfiles( # If simply a percentage error provided # -------------------------------------------------------------- - relativeErrorTGLF = [percentError[0] / 100.0]*len(rho_locations) + relativeErrorTGLF = [percent_error[0] / 100.0]*len(rho_locations) QeE = abs(Qe) * relativeErrorTGLF QiE = abs(Qi) * relativeErrorTGLF @@ -907,7 +874,6 @@ def tgyro_to_powerstate( impurityPosition=1, UseFineGridTargets=False, OriginalFimp=1.0, - provideTurbulentExchange=False, provideTargets=False ): """ @@ -990,12 +956,8 @@ def tgyro_to_powerstate( # *********** Energy Exchange # ********************************** - if provideTurbulentExchange: - self.QieGB_turb = TGYROresults.EXeGB_sim_turb[0, 1:nr] - self.QieGB_turb_stds = TGYROresults.EXeGB_sim_turb_stds[0, 1:nr] - else: - self.QieGB_turb = self.QeGB_turb * 0.0 - self.QieGB_turb_stds = self.QeGB_turb * 0.0 + self.QieGB_turb = TGYROresults.EXeGB_sim_turb[0, 1:nr] + self.QieGB_turb_stds = TGYROresults.EXeGB_sim_turb_stds[0, 1:nr] self.QieGB_neoc = self.QeGB_turb * 0.0 self.QieGB_neoc_stds = self.QeGB_turb_stds * 0.0 diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index f0370d58..228351db 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -30,7 +30,7 @@ def calculator( profiles, evolution_options={ "rhoPredicted": rho_vec, - 'fineTargetsResolution': fineTargetsResolution, + "fineTargetsResolution": fineTargetsResolution, }, target_options={ "target_evaluator": targets_analytic.analytical_model, @@ -42,28 +42,32 @@ def calculator( "transport_evaluator": transport_tgyro.tgyro_model, "transport_evaluator_options": { "cold_start": cold_start, - "launchSlurm": True, - "MODELparameters": { - "Physics_options": { - "TypeTarget": 3, - "TurbulentExchange": 0, - "PtotType": 1, - "GradientsType": 0, - "InputType": 1, + "portals_parameters": { + "main_parameters": { + "launchSlurm": True, + "Qi_includes_fast": False, }, - "ProfilesPredicted": ["te", "ti", "ne"], - "RhoLocations": rho_vec, - "applyCorrections": { - "Tfast_ratio": False, - "Ti_thermals": True, - "ni_thermals": True, - "recalculate_ptot": False, + "model_parameters": { + "Physics_options": { + "TypeTarget": 3, + "TurbulentExchange": 0, + "PtotType": 1, + "GradientsType": 0, + "InputType": 1, + }, + "predicted_channels": ["te", "ti", "ne"], + "radii_rho": rho_vec, + "applyCorrections": { + "Tfast_ratio": False, + "Ti_thermals": True, + "ni_thermals": True, + "recalculate_ptot": False, + }, + "transport_model": {"TGLFsettings": 5, "extraOptionsTGLF": {}}, }, - "transport_model": {"TGLFsettings": 5, "extraOptionsTGLF": {}}, }, - "Qi_includes_fast": False, }, - }, + } ) # Calculate using powerstate @@ -72,7 +76,7 @@ def calculator( profiles, evolution_options={ "rhoPredicted": rho_vec, - 'fineTargetsResolution': fineTargetsResolution, + "fineTargetsResolution": fineTargetsResolution, }, target_options={ "target_evaluator": targets_analytic.analytical_model, diff --git a/src/mitim_modules/powertorch/utils/POWERplot.py b/src/mitim_modules/powertorch/utils/POWERplot.py index bf29adaa..009dc9a3 100644 --- a/src/mitim_modules/powertorch/utils/POWERplot.py +++ b/src/mitim_modules/powertorch/utils/POWERplot.py @@ -24,17 +24,17 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co set_plots = [ ] - if "te" in self.ProfilesPredicted: + if "te" in self.predicted_channels: set_plots.append( [ 'te', 'aLte', 'QeMWm2_tr', 'QeMWm2', 'Electron Temperature','$T_e$ (keV)','$a/LT_e$','$Q_e$ (GB)','$Q_e$ ($MW/m^2$)', 1.0,"Qgb"]) - if "ti" in self.ProfilesPredicted: + if "ti" in self.predicted_channels: set_plots.append( [ 'ti', 'aLti', 'QiMWm2_tr', 'QiMWm2', 'Ion Temperature','$T_i$ (keV)','$a/LT_i$','$Q_i$ (GB)','$Q_i$ ($MW/m^2$)', 1.0,"Qgb"]) - if "ne" in self.ProfilesPredicted: + if "ne" in self.predicted_channels: # If this model provides the raw particle flux, go for it if 'Ge1E20m2_tr' in self.plasma: @@ -48,7 +48,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co 'Electron Density','$n_e$ ($10^{20}m^{-3}$)','$a/Ln_e$','$Q_{conv,e}$ (GB)','$Q_{conv,e}$ ($MW/m^2$)', 1E-1,"Qgb"]) - if "nZ" in self.ProfilesPredicted: + if "nZ" in self.predicted_channels: # If this model provides the raw particle flux, go for it if 'GZ1E20m2_tr' in self.plasma: @@ -62,7 +62,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co 'Impurity Density','$n_Z$ ($10^{20}m^{-3}$)','$a/Ln_Z$','$\\widehat{Q}_{conv,Z}$ (GB)','$\\widehat{Q}_{conv,Z}$ ($MW/m^2$)', 1E-1,"Qgb"]) - if "w0" in self.ProfilesPredicted: + if "w0" in self.predicted_channels: set_plots.append( [ 'w0', 'aLw0', 'MtJm2_tr', 'MtJm2', 'Rotation','$\\omega_0$ ($krad/s$)','$-d\\omega_0/dr$ ($krad/s/cm$)','$\\Pi$ (GB)','$\\Pi$ ($J/m^2$)', @@ -104,7 +104,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co colors = GRAPHICStools.listColors() cont = 0 - for i in range(len(self.ProfilesPredicted)): + for i in range(len(self.predicted_channels)): # Plot gradient evolution ax = axsRes[1+cont] @@ -119,7 +119,7 @@ def plot(self, axs, axsRes, figs=None, c="r", label="powerstate",batch_num=0, co ax.set_ylabel(self.labelsFM[i][0]) - if i == len(self.ProfilesPredicted)-1: + if i == len(self.predicted_channels)-1: GRAPHICStools.addLegendApart(ax, ratio=1.0,extraPad=0.05, size=9) # Plot residual evolution diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 133448d2..dd2e7641 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -55,7 +55,7 @@ def fine_grid(self): # ---------------------------------------------------- # Integrate through fine profile constructors # ---------------------------------------------------- - for i in self.powerstate.ProfilesPredicted: + for i in self.powerstate.predicted_channels: _ = self.powerstate.update_var(i,specific_profile_constructor=self.powerstate.profile_constructors_coarse_middle) def flux_integrate(self): diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 60ccd39d..c9dcdada 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -193,7 +193,7 @@ def gacode_to_powerstate(self, rho_vec=None): self.plasma[f"aL{key[0]}"] = aLy_coarse[:-1, 1] # Check that it's not completely zero - if key[0] in self.ProfilesPredicted: + if key[0] in self.predicted_channels: if self.plasma[f"aL{key[0]}"].sum() == 0.0: addT = 1e-15 print(f"\t- All values of {key[0]} detected to be zero, to avoid NaNs, inserting {addT} at the edge",typeMsg="w") @@ -278,7 +278,7 @@ def powerstate_to_gacode( ] for key in quantities: - if key[0] in self.ProfilesPredicted: + if key[0] in self.predicted_channels: print(f"\t- Inserting {key[0]} into input.gacode profiles") # ********************************************************************************************* @@ -319,7 +319,7 @@ def powerstate_to_gacode( print("\t\t* Adjusting ni of thermal ions", typeMsg="i") profiles.scaleAllThermalDensities(scaleFactor=scaleFactor) - if "w0" not in self.ProfilesPredicted and ensureMachNumber is not None: + if "w0" not in self.predicted_channels and ensureMachNumber is not None: # Rotation fixed to ensure Mach number profiles.introduceRotationProfile(Mach_LF=ensureMachNumber) @@ -369,6 +369,8 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): "target_evaluator_options": { "TypeTarget": self.target_options["target_evaluator_options"]["TypeTarget"], # Important to keep the same as in the original "target_evaluator_method": "powerstate", + "forceZeroParticleFlux": self.target_options["target_evaluator_options"]["forceZeroParticleFlux"], + "percent_error": self.target_options["target_evaluator_options"]["percent_error"] } }, increase_profile_resol = False diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 82d8b52c..03187e25 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -127,7 +127,7 @@ def produce_profiles(self): def _produce_profiles(self,derive_quantities=True): - self.applyCorrections = self.powerstate.transport_options["transport_evaluator_options"].get("MODELparameters", {}).get("applyCorrections", {}) + self.applyCorrections = self.powerstate.transport_options["applyCorrections"] # Write this updated profiles class (with parameterized profiles and target powers) self.file_profs = self.folder / "input.gacode" @@ -154,7 +154,7 @@ def _modify_profiles(self): self.file_profs_unmod = self.file_profs.parent / f"{self.file_profs.name}_unmodified" shutil.copy2(self.file_profs, self.file_profs_unmod) - profiles_postprocessing_fun = self.powerstate.transport_options["transport_evaluator_options"].get("profiles_postprocessing_fun", None) + profiles_postprocessing_fun = self.powerstate.transport_options["profiles_postprocessing_fun"] if profiles_postprocessing_fun is not None: print(f"\t- Modifying input.gacode to run transport calculations based on {profiles_postprocessing_fun}",typeMsg="i") diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index 387454f2..417b68d5 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -61,13 +61,13 @@ "portals_beat": { "use_default": false, "portals_namelist" : { - "PORTALSparameters": { + "portals_parameters": { "forceZeroParticleFlux": true, "keep_full_model_folder": false, "cores_per_tglf_instance": 1 }, "MODELparameters": { - "RoaLocations": [0.35,0.45,0.55,0.65,0.75,0.875,0.9], + "radii_roa": [0.35,0.45,0.55,0.65,0.75,0.875,0.9], "Physics_options": {"TypeTarget": 3}, "transport_model": { "TGLFsettings": 100, diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index e3bb56e6..f83e7cf3 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -28,14 +28,17 @@ portals_fun = PORTALSmain.portals(folderWork) portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 portals_fun.optimization_options["initialization_options"]["initial_training"] = 2 -portals_fun.MODELparameters["RhoLocations"] = [0.25, 0.45, 0.65, 0.85] -portals_fun.MODELparameters['ProfilesPredicted'] = ["te", "ti", "ne", "nZ", 'w0'] -portals_fun.PORTALSparameters['ImpurityOfInterest'] = 'N' -portals_fun.PORTALSparameters['turbulent_exchange_as_surrogate'] = True -portals_fun.INITparameters["remove_fast"] = True -portals_fun.INITparameters["quasineutrality"] = True -portals_fun.INITparameters["enforce_same_aLn"] = True -portals_fun.MODELparameters["transport_model"]["TGLFsettings"] = 2 + +portals_fun.portals_parameters["main_parameters"]['turbulent_exchange_as_surrogate'] = True + +portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True +portals_fun.portals_parameters["initialization_parameters"]["quasineutrality"] = True +portals_fun.portals_parameters["initialization_parameters"]["enforce_same_aLn"] = True + +portals_fun.portals_parameters["model_parameters"]["radii_rho"] = [0.25, 0.45, 0.65, 0.85] +portals_fun.portals_parameters["model_parameters"]['predicted_channels'] = ["te", "ti", "ne", "nZ", 'w0'] +portals_fun.portals_parameters["model_parameters"]['ImpurityOfInterest'] = 'N' +portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["tglf"]["run"]["code_settings"] = 2 # Prepare run portals_fun.prep(inputgacode) diff --git a/tests/POWERTORCH_workflow.py b/tests/POWERTORCH_workflow.py index 2306a9a4..6c29c017 100644 --- a/tests/POWERTORCH_workflow.py +++ b/tests/POWERTORCH_workflow.py @@ -15,7 +15,7 @@ 'rhoPredicted': rho }, transport_options = { 'transport_evaluator': transport_analytic.diffusion_model, - 'transport_evaluator_options': { + "transport_evaluator_options": { 'chi_e': torch.ones(rho.shape[0]).to(rho)*0.8, 'chi_i': torch.ones(rho.shape[0]).to(rho)*1.2 } diff --git a/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index 9b6ca805..e0fb7ecd 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -13,25 +13,25 @@ # Initialize PORTALS class portals_fun = PORTALSmain.portals(folder) -# Radial locations (RhoLocations or RoaLocations [last one preceeds]) -portals_fun.MODELparameters["RhoLocations"] = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.85] +# Radial locations (radii_rho or radii_roa [last one preceeds]) +portals_fun.portals_parameters["model_parameters"]["radii_rho"] = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.85] # Profiles to predict -portals_fun.MODELparameters["ProfilesPredicted"] = ["te", "ti", "ne"] +portals_fun.portals_parameters["model_parameters"]["predicted_channels"] = ["te", "ti", "ne"] # Codes to use from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model -portals_fun.PORTALSparameters["transport_evaluator"] = tgyro_model +portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator"] = tgyro_model # TGLF specifications -portals_fun.MODELparameters["transport_model"] = { +portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"] = { "TGLFsettings": 6, # Check out templates/input.tglf.models.json for more options "extraOptionsTGLF": {"USE_BPER": False} # Turn off BPER } # Plasma preparation: remove fast species, adjust quasineutrality -portals_fun.INITparameters["remove_fast"] = True -portals_fun.INITparameters["quasineutrality"] = True +portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True +portals_fun.portals_parameters["initialization_parameters"]["quasineutrality"] = True # Stopping criterion 1: 100x improvement in residual portals_fun.optimization_options['convergence_options']['stopping_criteria_parameters']["maximum_value"] = 1e-2 From bc789e092a767b870bd614d1f4709ca8236cef11 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 14:27:26 -0400 Subject: [PATCH 225/385] Added stable correciton CGYRO and further naming convntions --- regressions/portals_regressions.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 16 ++++++++------ .../portals/utils/PORTALSanalysis.py | 4 ++-- .../portals/utils/PORTALSinit.py | 12 +++++----- .../portals/utils/PORTALSplot.py | 20 ++++++++--------- .../physics_models/transport_cgyro.py | 10 ++++----- .../physics_models/transport_cgyroneo.py | 22 +++++++++++++++++-- .../physics_models/transport_tgyro.py | 6 ++--- .../powertorch/scripts/calculateTargets.py | 2 +- templates/maestro_namelist.json | 2 +- tests/PORTALS_workflow.py | 2 +- tutorials/PORTALS_tutorial.py | 4 ++-- 12 files changed, 61 insertions(+), 41 deletions(-) diff --git a/regressions/portals_regressions.py b/regressions/portals_regressions.py index 7a32ab14..fa05c91e 100644 --- a/regressions/portals_regressions.py +++ b/regressions/portals_regressions.py @@ -93,7 +93,7 @@ def conditions_regressions(variables): portals_fun = PORTALSmain.portals(folderWork) portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - portals_fun.portals_parameters["model_parameters"]["radii_rho"] = [0.25, 0.45, 0.65, 0.85] + portals_fun.portals_parameters["model_parameters"]["predicted_rho"] = [0.25, 0.45, 0.65, 0.85] portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True portals_fun.portals_parameters["initialization_parameters"]["quasineutrality"] = True portals_fun.portals_parameters["initialization_parameters"]["enforce_same_aLn"] = True diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 7fa098b1..1f0056fb 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -182,8 +182,8 @@ def __init__( self.portals_parameters["model_parameters"] = { # Specification of radial locations (roa wins over rho, if provided) - "radii_rho": [0.3, 0.45, 0.6, 0.75, 0.9], - "radii_roa": None, + "predicted_rho": [0.3, 0.45, 0.6, 0.75, 0.9], + "predicted_roa": None, # Channels to be predicted "predicted_channels": ["te", "ti", "ne"], # ['nZ','w0'] @@ -227,8 +227,8 @@ def __init__( "read": { "tmin": 0.0 }, - "Qi_criterion_stable": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable - "percentError_stable": 5.0, # (%) For CGYRO runs, minimum error based on target if case is considered stable + "Qi_stable_criterion": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable + "Qi_stable_percent_error": 5.0, # (%) For CGYRO runs, minimum error based on target if case is considered stable }, }, @@ -496,18 +496,20 @@ def check_flags(self): # Check that I haven't added a deprecated variable that I expect some behavior from - def _check_flags_dictionary(d, d_check): + def _check_flags_dictionary(d, d_check, avoid = ["run", "read"]): for key in d.keys(): if key not in d_check: print(f"\t- {key} is an unexpected variable, prone to errors or misinterpretation",typeMsg="q") elif not isinstance(d[key], dict): continue + elif key in avoid: + continue else: _check_flags_dictionary(d[key], d_check[key]) - + _check_flags_dictionary(self.portals_parameters, self.potential_flags) - key_rhos = "radii_roa" if self.portals_parameters["model_parameters"]["radii_roa"] is not None else "radii_rho" + key_rhos = "predicted_roa" if self.portals_parameters["model_parameters"]["predicted_roa"] is not None else "predicted_rho" return key_rhos diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 8fe72fd4..28c3da7d 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -563,12 +563,12 @@ def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=F """ NOTE on radial location extraction: Two possible options for the rho locations to use: - 1. self.portals_parameters["model_parameters"]["radii_rho"] -> the ones PORTALS sent to TGYRO + 1. self.portals_parameters["model_parameters"]["predicted_rho"] -> 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. #TODO: This should be fixed in the future, we should never send to TGYRO more than 7 decimal places of any variable """ - rhos_considered = self.portals_parameters["model_parameters"]["radii_rho"] + rhos_considered = self.portals_parameters["model_parameters"]["predicted_rho"] if positions is None: rhos = rhos_considered diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index fb64b7a4..a1b7a88f 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -71,13 +71,13 @@ def initializeProblem( profiles = PROFILEStools.gacode_state(initialization_file) # About radial locations - if portals_fun.portals_parameters["model_parameters"]["radii_roa"] is not None: - roa = portals_fun.portals_parameters["model_parameters"]["radii_roa"] + if portals_fun.portals_parameters["model_parameters"]["predicted_roa"] is not None: + roa = portals_fun.portals_parameters["model_parameters"]["predicted_roa"] rho = np.interp(roa, profiles.derived["roa"], profiles.profiles["rho(-)"]) print("\t * r/a provided, transforming to rho:") print(f"\t\t r/a = {roa}") print(f"\t\t rho = {rho}") - portals_fun.portals_parameters["model_parameters"]["radii_rho"] = rho + portals_fun.portals_parameters["model_parameters"]["predicted_rho"] = rho if ( len(portals_parameters["initialization_parameters"]["removeIons"]) > 0 @@ -122,7 +122,7 @@ def initializeProblem( # Prepare and defaults - xCPs = torch.from_numpy(np.array(portals_fun.portals_parameters["model_parameters"]["radii_rho"])).to(dfT) + xCPs = torch.from_numpy(np.array(portals_fun.portals_parameters["model_parameters"]["predicted_rho"])).to(dfT) """ *************************************************************************************************** @@ -260,7 +260,7 @@ def initializeProblem( elif ikey == "w0": var = "Mt" - for i in range(len(portals_fun.portals_parameters["model_parameters"]["radii_rho"])): + for i in range(len(portals_fun.portals_parameters["model_parameters"]["predicted_rho"])): ofs.append(f"{var}_tr_turb_{i+1}") ofs.append(f"{var}_tr_neoc_{i+1}") @@ -269,7 +269,7 @@ def initializeProblem( name_objectives.append(f"{var}Res_{i+1}") if portals_fun.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]: - for i in range(len(portals_fun.portals_parameters["model_parameters"]["radii_rho"])): + for i in range(len(portals_fun.portals_parameters["model_parameters"]["predicted_rho"])): ofs.append(f"Qie_tr_turb_{i+1}") name_transformed_ofs = [] diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index 78f934e4..37956dce 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -1207,7 +1207,7 @@ def PORTALSanalyzer_plotExpected( rho = p.profiles["rho(-)"] roa = p.derived["roa"] - rhoVals = self.portals_parameters["model_parameters"]["radii_rho"] + rhoVals = self.portals_parameters["model_parameters"]["predicted_rho"] roaVals = np.interp(rhoVals, rho, roa) lastX = roaVals[-1] @@ -1395,7 +1395,7 @@ def PORTALSanalyzer_plotExpected( rho = self.profiles_next_new.profiles["rho(-)"] - rhoVals = self.portals_parameters["model_parameters"]["radii_rho"] + rhoVals = self.portals_parameters["model_parameters"]["predicted_rho"] roaVals = np.interp(rhoVals, rho, roa) p0 = self.powerstates[plotPoints[0]].profiles @@ -1892,10 +1892,10 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): axs4, color=colors[i], label=label, - lastRho=self.portals_parameters["model_parameters"]["radii_rho"][-1], + lastRho=self.portals_parameters["model_parameters"]["predicted_rho"][-1], alpha=alpha, useRoa=True, - RhoLocationsPlot=self.portals_parameters["model_parameters"]["radii_rho"], + RhoLocationsPlot=self.portals_parameters["model_parameters"]["predicted_rho"], plotImpurity=self.runWithImpurity, plotRotation=self.runWithRotation, autoscale=i == 3, @@ -1959,7 +1959,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): p.plot_gradients( axsR, color="b", - lastRho=self.portals_parameters["model_parameters"]["radii_rho"][-1], + lastRho=self.portals_parameters["model_parameters"]["predicted_rho"][-1], ms=ms, lw=1.0, label="Initial (#0)", @@ -1976,7 +1976,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): p.plot_gradients( axsR, color="r", - lastRho=self.portals_parameters["model_parameters"]["radii_rho"][-1], + lastRho=self.portals_parameters["model_parameters"]["predicted_rho"][-1], ms=ms, lw=0.3, ls="-o" if self.opt_fun.mitim_model.avoidPoints is not None else "-.o", @@ -1988,7 +1988,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): p.plot_gradients( axsR, color="g", - lastRho=self.portals_parameters["model_parameters"]["radii_rho"][-1], + lastRho=self.portals_parameters["model_parameters"]["predicted_rho"][-1], ms=ms, lw=1.0, label=f"Best (#{self.opt_fun.res.best_absolute_index})", @@ -2420,7 +2420,7 @@ def varToReal(y, mitim_model): Qe, Qi, Ge, GZ, Mt = [], [], [], [], [] Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = [], [], [], [], [] for prof in mitim_model.optimization_object.portals_parameters["model_parameters"]["predicted_channels"]: - for rad in mitim_model.optimization_object.portals_parameters["model_parameters"]["radii_rho"]: + for rad in mitim_model.optimization_object.portals_parameters["model_parameters"]["predicted_rho"]: if prof == "te": Qe.append(of[0, cont]) Qe_tar.append(cal[0, cont]) @@ -2493,7 +2493,7 @@ def plotVars( .plasma["roa"][0, 1:] .cpu() .cpu().numpy() - ) # mitim_model.optimization_object.portals_parameters["model_parameters"]['radii_rho'] + ) # mitim_model.optimization_object.portals_parameters["model_parameters"]['predicted_rho'] try: Qe, Qi, Ge, GZ, Mt, Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = varToReal( @@ -3254,7 +3254,7 @@ def plotFluxComparison( def produceInfoRanges( self_complete, bounds, axsR, label="", color="k", lw=0.2, alpha=0.05 ): - rhos = np.append([0], self_complete.portals_parameters["model_parameters"]["radii_rho"]) + rhos = np.append([0], self_complete.portals_parameters["model_parameters"]["predicted_rho"]) aLTe, aLTi, aLne, aLnZ, aLw0 = ( np.zeros((len(rhos), 2)), np.zeros((len(rhos), 2)), diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 4837ba14..c7977006 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -171,8 +171,8 @@ def evaluateCGYRO(transport, folder, numPORTALS, FolderEvaluation, unmodified_pr trick_hardcoded_f = None # ------------------------------------------------------------------------------------------------ - minErrorPercent = transport["percentError_stable"] - Qi_criterion_stable = transport["Qi_criterion_stable"] + minErrorPercent = transport["Qi_stable_percent_error"] + Qi_stable_criterion = transport["Qi_stable_criterion"] OriginalFimp = portals_parameters["main_parameters"]["fImp_orig"] @@ -182,7 +182,7 @@ def evaluateCGYRO(transport, folder, numPORTALS, FolderEvaluation, unmodified_pr unmodified_profiles, numPORTALS, minErrorPercent, - Qi_criterion_stable, + Qi_stable_criterion, radii, OriginalFimp=OriginalFimp, evaluationsInFile=f"{numPORTALS_this}", @@ -212,7 +212,7 @@ def cgyroing( unmodified_profiles, evaluations, minErrorPercent, - Qi_criterion_stable, + Qi_stable_criterion, radii, OriginalFimp=1.0, file=None, @@ -257,7 +257,7 @@ def cgyroing( tgyro, FolderEvaluation, minErrorPercent=minErrorPercent, - Qi_criterion_stable=Qi_criterion_stable, + Qi_stable_criterion=Qi_stable_criterion, impurityPosition=impurityPosition, OriginalFimp=OriginalFimp, ) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index 7eda3057..c0690b14 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -25,7 +25,7 @@ def evaluate_turbulence(self): simulation_options_cgyro = transport_evaluator_options["cgyro"] - run_type = simulation_options_cgyro["run_type"] + run_type = simulation_options_cgyro["run"]["run_type"] # ------------------------------------------------------------------------------------------------------------------------ # Prepare CGYRO object @@ -42,7 +42,6 @@ def evaluate_turbulence(self): _ = cgyro.run( 'base_cgyro', - run_type = run_type, cold_start=cold_start, forceIfcold_start=True, **simulation_options_cgyro["run"] @@ -107,6 +106,25 @@ def evaluate_turbulence(self): if file_path.exists(): all_good = post_checks(self) + self._stable_correction(simulation_options_cgyro) + + def _stable_correction(self, simulation_options_cgyro): + + Qi_stable_criterion = simulation_options_cgyro["Qi_stable_criterion"] + Qi_stable_percent_error = simulation_options_cgyro["Qi_stable_percent_error"] + + # Check if Qi in MW/m2 < Qi_stable_criterion + QiMWm2 = self.QiGB_turb * self.powerstate.plasma['Qgb'][0,1:].cpu().numpy() + QiGB_target = self.powerstate.plasma['QiGB'][0,1:].cpu().numpy() + + radii_stable = QiMWm2 < Qi_stable_criterion + + for i in range(len(radii_stable)): + if radii_stable[i]: + print(f"\t- Qi considered stable at radius #{i}, ({QiMWm2[i]:.2f} < {Qi_stable_criterion:.2f})", typeMsg='i') + Qi_std = QiGB_target[i] * Qi_stable_percent_error / 100 + print(f"\t\t- Assigning {Qi_stable_percent_error:.1f}% from target as standard deviation: {Qi_std:.2f} instead of {self.QiGB_turb_stds[i]}", typeMsg='i') + self.QiGB_turb_stds[i] = Qi_std def pre_checks(self): diff --git a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py index bcd0bed7..da442852 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tgyro.py @@ -557,7 +557,7 @@ def modifyResults( tgyro, folder_tgyro, minErrorPercent=5.0, - Qi_criterion_stable=0.0025, + Qi_stable_criterion=0.0025, impurityPosition=3, OriginalFimp=1.0, ): @@ -585,9 +585,9 @@ def modifyResults( Mt_min = Mt_target * (minErrorPercent / 100.0) for i in range(Qe.shape[0]): - if Qi[i] < Qi_criterion_stable: + if Qi[i] < Qi_stable_criterion: print( - f"\t- Based on 'Qi_criterion_stable', plasma considered stable (Qi = {Qi[i]:.2e} < {Qi_criterion_stable:.2e} MW/m2) at position #{i}, using minimum errors of {minErrorPercent}% of targets", + f"\t- Based on 'Qi_stable_criterion', plasma considered stable (Qi = {Qi[i]:.2e} < {Qi_stable_criterion:.2e} MW/m2) at position #{i}, using minimum errors of {minErrorPercent}% of targets", typeMsg="w", ) QeE[i] = Qe_min[i] diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 228351db..ea3b88df 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -56,7 +56,7 @@ def calculator( "InputType": 1, }, "predicted_channels": ["te", "ti", "ne"], - "radii_rho": rho_vec, + "predicted_rho": rho_vec, "applyCorrections": { "Tfast_ratio": False, "Ti_thermals": True, diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index 417b68d5..99bd23f9 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -67,7 +67,7 @@ "cores_per_tglf_instance": 1 }, "MODELparameters": { - "radii_roa": [0.35,0.45,0.55,0.65,0.75,0.875,0.9], + "predicted_roa": [0.35,0.45,0.55,0.65,0.75,0.875,0.9], "Physics_options": {"TypeTarget": 3}, "transport_model": { "TGLFsettings": 100, diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index f83e7cf3..9c580f80 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -35,7 +35,7 @@ portals_fun.portals_parameters["initialization_parameters"]["quasineutrality"] = True portals_fun.portals_parameters["initialization_parameters"]["enforce_same_aLn"] = True -portals_fun.portals_parameters["model_parameters"]["radii_rho"] = [0.25, 0.45, 0.65, 0.85] +portals_fun.portals_parameters["model_parameters"]["predicted_rho"] = [0.25, 0.45, 0.65, 0.85] portals_fun.portals_parameters["model_parameters"]['predicted_channels'] = ["te", "ti", "ne", "nZ", 'w0'] portals_fun.portals_parameters["model_parameters"]['ImpurityOfInterest'] = 'N' portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["tglf"]["run"]["code_settings"] = 2 diff --git a/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index e0fb7ecd..4de41bb3 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -13,8 +13,8 @@ # Initialize PORTALS class portals_fun = PORTALSmain.portals(folder) -# Radial locations (radii_rho or radii_roa [last one preceeds]) -portals_fun.portals_parameters["model_parameters"]["radii_rho"] = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.85] +# Radial locations (predicted_rho or predicted_roa [last one preceeds]) +portals_fun.portals_parameters["model_parameters"]["predicted_rho"] = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.85] # Profiles to predict portals_fun.portals_parameters["model_parameters"]["predicted_channels"] = ["te", "ti", "ne"] From 8ebb90d8b8c91df71c0cb2dd1a3b18d397622ac5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 14:33:14 -0400 Subject: [PATCH 226/385] Removal of non-kept-up-to-date TGYRO implementation with PORTALS --- src/mitim_modules/portals/PORTALSmain.py | 3 - .../physics_models/transport_cgyro.py | 476 --------- .../physics_models/transport_tgyro.py | 988 ------------------ .../powertorch/scripts/calculateTargets.py | 4 +- 4 files changed, 2 insertions(+), 1469 deletions(-) delete mode 100644 src/mitim_modules/powertorch/physics_models/transport_cgyro.py delete mode 100644 src/mitim_modules/powertorch/physics_models/transport_tgyro.py diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 1f0056fb..500d31fc 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -577,9 +577,6 @@ def reuseTrainingTabular( if reevaluateTargets == 1: self_copy.powerstate.transport_options["transport_evaluator"] = None self_copy.powerstate.target_options["target_evaluator_options"]["TypeTarget"] = "powerstate" - else: - from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model - self_copy.powerstate.transport_options["transport_evaluator"] = tgyro_model _, dictOFs = runModelEvaluator( self_copy, diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py deleted file mode 100644 index c7977006..00000000 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ /dev/null @@ -1,476 +0,0 @@ -''' -TO BE COMPLETELY REMOVED -''' - -import copy -import shutil -import torch -import numpy as np -from mitim_tools.misc_tools import IOtools -from mitim_tools.gacode_tools import PROFILEStools, TGYROtools -from mitim_tools.plasmastate_tools import MITIMstate -from mitim_modules.powertorch.physics_models import transport_tgyro -from mitim_tools.misc_tools.LOGtools import printMsg as print -from IPython import embed - -class cgyro_model(transport_tgyro.tgyro_model): - def __init__(self, powerstate, **kwargs): - super().__init__(powerstate, **kwargs) - - def evaluate(self): - - # Run original evaluator - tgyro = self._evaluate_tglf_neo() - - # Run CGYRO trick - powerstate_orig = self._trick_cgyro(tgyro) - - # Process results - self._postprocess_tgyro(tgyro, "cgyro_neo") - - # Some checks - print("\t- Checking model modifications:") - for r in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "MtJm2_tr_turb"]: #, "QieMWm3_tr_turb"]: #TODO: FIX - print(f"\t\t{r}(tglf) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(powerstate_orig.plasma[r][0][1:],powerstate_orig.plasma[r+'_stds'][0][1:]) ])}") - print(f"\t\t{r}(cgyro) = {' '.join([f'{k:.1e} (+-{ke:.1e})' for k,ke in zip(self.powerstate.plasma[r][0][1:],self.powerstate.plasma[r+'_stds'][0][1:]) ])}") - - # ************************************************************************************ - # Private functions for CGYRO evaluation - # ************************************************************************************ - - def _trick_cgyro(self, tgyro): - - FolderEvaluation_TGYRO = self.folder / "cgyro_neo" - - print("\t- Checking whether cgyro_neo folder exists and it was written correctly via cgyro_trick...") - - correctly_run = FolderEvaluation_TGYRO.exists() - if correctly_run: - print("\t\t- Folder exists, but was cgyro_trick run?") - with open(FolderEvaluation_TGYRO / "mitim_flag", "r") as f: - correctly_run = bool(float(f.readline())) - - if correctly_run: - print("\t\t\t* Yes, it was", typeMsg="w") - else: - print("\t\t\t* No, it was run, repating process", typeMsg="i") - - # Remove cgyro_neo folder - if FolderEvaluation_TGYRO.exists(): - IOtools.shutil_rmtree(FolderEvaluation_TGYRO) - - # Copy tglf_neo results - shutil.copytree(self.folder / "tglf_neo", FolderEvaluation_TGYRO) - - # ********************************************************** - # CGYRO writter - # ********************************************************** - - # Write a flag indicating this was not performed - with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: - f.write("0") - - self._cgyro_trick(FolderEvaluation_TGYRO) - - # Write a flag indicating this was performed, to avoid an issue that... the script crashes when it has copied tglf_neo, without cgyro_trick modification - with open(FolderEvaluation_TGYRO / "mitim_flag", "w") as f: - f.write("1") - - # Read TGYRO files and construct portals variables - - tgyro.read(label="cgyro_neo", folder=FolderEvaluation_TGYRO) - - powerstate_orig = copy.deepcopy(self.powerstate) - - return powerstate_orig - - def _cgyro_trick(self,FolderEvaluation_TGYRO): - - # Print Information - print(self._print_info()) - - # Copy profiles so that later it is easy to grab all the input.gacodes that were evaluated - self._profiles_to_store() - - # ************************************************************************************************************************** - # Evaluate CGYRO - # ************************************************************************************************************************** - - evaluateCGYRO( - self.powerstate.transport_options["transport_evaluator_options"]["transport_parameters"], - self.powerstate.transport_options["folder"], - self.evaluation_number, - FolderEvaluation_TGYRO, - self.file_profs, - self.powerstate.plasma["roa"][0,1:], - self.powerstate.predicted_channels, - ) - - # Make tensors - for i in ["QeMWm2_tr_turb", "QiMWm2_tr_turb", "Ce_tr_turb", "CZ_tr_turb", "MtJm2_tr_turb"]: - try: - self.powerstate.plasma[i] = torch.from_numpy(self.powerstate.plasma[i]).to(self.powerstate.dfT).unsqueeze(0) - except: - pass - - def _print_info(self): - - txt = "\nFluxes to be matched by CGYRO ( TARGETS - NEO ):" - - for var, varn in zip( - ["r/a ", "rho ", "a/LTe", "a/LTi", "a/Lne", "a/LnZ", "a/Lw0"], - ["roa", "rho", "aLte", "aLti", "aLne", "aLnZ", "aLw0"], - ): - txt += f"\n{var} = " - for j in range(self.powerstate.plasma["rho"].shape[1] - 1): - txt += f"{self.powerstate.plasma[varn][0,j+1]:.6f} " - - for var, varn in zip( - ["Qe (MW/m^2)", "Qi (MW/m^2)", "Ce (MW/m^2)", "CZ (MW/m^2)", "MtJm2 (J/m^2) "], - ["QeMWm2", "QiMWm2", "Ce", "CZ", "MtJm2"], - ): - txt += f"\n{var} = " - for j in range(self.powerstate.plasma["rho"].shape[1] - 1): - txt += f"{self.powerstate.plasma[varn][0,j+1]-self.powerstate.plasma[f'{varn}_tr_neoc'][0,j+1]:.4e} " - - return txt - -""" -The CGYRO file must contain GB units, and the gb unit is MW/m^2, 1E19m^2/s -The CGYRO file must use particle flux. Convective transformation occurs later -""" - -def evaluateCGYRO(transport, folder, numPORTALS, FolderEvaluation, unmodified_profiles, radii, predicted_channels, impurityPosition): - print("\n ** CGYRO evaluation of fluxes has been requested before passing information to the STRATEGY module **",typeMsg="i",) - - if isinstance(numPORTALS, int): - numPORTALS = str(numPORTALS) - - # ------------------------------------------------------------------------------------------------ - # Harcoded - # ------------------------------------------------------------------------------------------------ - if transport['hardCodedCGYRO'] is not 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. 'example_run/Outputs/cgyro_results/iter_rmp_75_' - - e.g.: - train_sep,start_num,last_one,trick_hardcoded_f = 1, 0,100, 'example_run/Outputs/cgyro_results/d3d_5chan_it_' - - """ - - train_sep = transport["hardCodedCGYRO"]["train_sep"] - start_num = transport["hardCodedCGYRO"]["start_num"] - last_one = transport["hardCodedCGYRO"]["last_one"] - trick_hardcoded_f = transport["hardCodedCGYRO"]["trick_hardcoded_f"] - else: - train_sep = None - start_num = None - last_one = None - trick_hardcoded_f = None - # ------------------------------------------------------------------------------------------------ - - minErrorPercent = transport["Qi_stable_percent_error"] - Qi_stable_criterion = transport["Qi_stable_criterion"] - - OriginalFimp = portals_parameters["main_parameters"]["fImp_orig"] - - cgyroing_file = ( - lambda file_cgyro, numPORTALS_this=0: cgyroing( - FolderEvaluation, - unmodified_profiles, - numPORTALS, - minErrorPercent, - Qi_stable_criterion, - radii, - OriginalFimp=OriginalFimp, - evaluationsInFile=f"{numPORTALS_this}", - impurityPosition=impurityPosition, - file=file_cgyro, - ) - ) - print(f"\t- Suggested function call for PORTALS evaluation {numPORTALS} (lambda for cgyroing):",typeMsg="i") - cgyropath = IOtools.expandPath(folder, ensurePathValid=True) / 'Outputs' / 'cgyro_results' / f'cgyro_it_{numPORTALS}.txt' - print(f"\tcgyroing_file('{cgyropath}')") - - print('\t- Then insert "exit" and RETURN', typeMsg="i") - if (trick_hardcoded_f is None) or (int(numPORTALS) > last_one): - embed() - else: - # ------------------------------------------------------------------ - # Hard-coded stuff for quick modifications - # ------------------------------------------------------------------ - if int(numPORTALS) < train_sep: - cgyroing_file(f"{trick_hardcoded_f}{start_num}.txt",numPORTALS_this=numPORTALS) - else: - cgyroing_file(f"{trick_hardcoded_f}{int(numPORTALS)-train_sep+1+start_num}.txt",numPORTALS_this=0) - - -def cgyroing( - FolderEvaluation, - unmodified_profiles, - evaluations, - minErrorPercent, - Qi_stable_criterion, - radii, - OriginalFimp=1.0, - file=None, - evaluationsInFile=0, - impurityPosition=3, -): - """ - Variables need to have dimensions of (evaluation,rho) - """ - - evaluations = np.array([int(i) for i in evaluations.split(",")]) - evaluationsInFile = np.array([int(i) for i in evaluationsInFile.split(",")]) - - aLTe,aLTi,aLne,Q_gb,Qe,Qi,Ge,GZ,Mt,Pexch,QeE,QiE,GeE,GZE,MtE,PexchE,_,_ = readCGYROresults(file, radii) - - cont = 0 - for _ in evaluations: - k = evaluationsInFile[cont] - cont += 1 - - print(f"\t- Modifying {IOtools.clipstr(FolderEvaluation)} with position {k} in CGYRO results file {IOtools.clipstr(file)}") - - # Get TGYRO - tgyro = TGYROtools.TGYROoutput(FolderEvaluation,profiles=PROFILEStools.gacode_state(unmodified_profiles)) - - # Quick checker of correct file - wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro) - - transport_tgyro.modifyResults( - Qe[k, :], - Qi[k, :], - Ge[k, :], - GZ[k, :], - Mt[k, :], - Pexch[k, :], - QeE[k, :], - QiE[k, :], - GeE[k, :], - GZE[k, :], - MtE[k, :], - PexchE[k, :], - tgyro, - FolderEvaluation, - minErrorPercent=minErrorPercent, - Qi_stable_criterion=Qi_stable_criterion, - impurityPosition=impurityPosition, - OriginalFimp=OriginalFimp, - ) - - -def wasThisTheCorrectRun(aLTe, aLTi, aLne, Q_gb, tgyro, ErrorRaised=0.005): - print("\t- Checking that this was the correct run...") - - tgyro_new = copy.deepcopy(tgyro) - tgyro_new.aLti = tgyro_new.aLti[:, 0, :] - - variables = [ - [aLTe, tgyro_new.aLte, "aLTe"], - [aLTi, tgyro_new.aLti, "aLTi"], - [aLne, tgyro_new.aLne, "aLne"], - [Q_gb, tgyro_new.Q_GB, "Qgb"], - ] - - for var in variables: - [c, t, n] = var - - for pos in range(c.shape[0]): - for i in range(c.shape[1]): - error = np.max(abs((t[pos, i + 1] - c[pos, i]) / t[pos, i + 1])) - print( - f"\t\t* Error in {n}[{i}] was {error*100.0:.2f}% (TGYRO {t[pos,i+1]:.3f} vs. CGYRO {c[pos,i]:.3f})", - typeMsg="w" if error > ErrorRaised else "", - ) - - -def readlineNTH(line, full_file=True, unnormalize=True): - s = line.split() - - i = 2 - roa = float(s[i]) - i += 3 - aLne = float(s[i]) - i += 3 - aLTi = float(s[i]) - i += 3 - aLTe = float(s[i]) - i += 3 - - Qi = float(s[i]) - i += 3 - Qi_std = float(s[i]) - i += 3 - Qe = float(s[i]) - i += 3 - Qe_std = float(s[i]) - i += 3 - Ge = float(s[i]) - i += 3 - Ge_std = float(s[i]) - i += 3 - - if full_file: - GZ = float(s[i]) - i += 3 - GZ_std = float(s[i]) - i += 3 - - Mt = float(s[i]) - i += 3 - Mt_std = float(s[i]) - i += 3 - - Pexch = float(s[i]) - i += 3 - Pexch_std = float(s[i]) - i += 3 - - Q_gb = float(s[i]) - i += 3 - G_gb = float(s[i]) * 1e-1 - i += 3 # From 1E19 to 1E20 - - if full_file: - Mt_gb = float(s[i]) - i += 3 - Pexch_gb = float(s[i]) - i += 3 - - tstart = float(s[i]) - i += 3 - tend = float(s[i]) - i += 3 - - if unnormalize: - QiReal = Qi * Q_gb - QiReal_std = Qi_std * Q_gb - QeReal = Qe * Q_gb - QeReal_std = Qe_std * Q_gb - GeReal = Ge * G_gb - GeReal_std = Ge_std * G_gb - else: - QiReal = Qi - QiReal_std = Qi_std - QeReal = Qe - QeReal_std = Qe_std - GeReal = Ge - GeReal_std = Ge_std - - if full_file: - if unnormalize: - GZReal = GZ * G_gb - GZReal_std = GZ_std * G_gb - - MtReal = Mt * Mt_gb - MtReal_std = Mt_std * Mt_gb - - PexchReal = Pexch * Pexch_gb - PexchReal_std = Pexch_std * Pexch_gb - else: - GZReal = GZ - GZReal_std = GZ_std - - MtReal = Mt - MtReal_std = Mt_std - - PexchReal = Pexch - PexchReal_std = Pexch_std - - return roa,aLTe,aLTi,aLne,Q_gb,QeReal,QiReal,GeReal,GZReal,MtReal,PexchReal,QeReal_std,QiReal_std,GeReal_std,GZReal_std,MtReal_std,PexchReal_std,tstart,tend - else: - return roa,aLTe,aLTi,aLne,Q_gb,QeReal,QiReal,GeReal,0.0,0.0,0.0,QeReal_std,QiReal_std,GeReal_std,0.0,0.0,0.0,tstart,tend - - -def readCGYROresults(file, radii, unnormalize=True): - """ - Arrays are in (batch,radii) - MW/m^2 and 1E20 - """ - - with open(file, "r") as f: - lines = f.readlines() - - rad = len(radii) - num = len(lines) // rad - - roa = np.zeros((num, rad)) - aLTe = np.zeros((num, rad)) - aLTi = np.zeros((num, rad)) - aLne = np.zeros((num, rad)) - Q_gb = np.zeros((num, rad)) - - Qe = np.zeros((num, rad)) - Qe_std = np.zeros((num, rad)) - Qi = np.zeros((num, rad)) - Qi_std = np.zeros((num, rad)) - Ge = np.zeros((num, rad)) - Ge_std = np.zeros((num, rad)) - - GZ = np.zeros((num, rad)) - GZ_std = np.zeros((num, rad)) - - Mt = np.zeros((num, rad)) - Mt_std = np.zeros((num, rad)) - - Pexch = np.zeros((num, rad)) - Pexch_std = np.zeros((num, rad)) - - tstart = np.zeros((num, rad)) - tend = np.zeros((num, rad)) - - p = {} - for r in range(len(radii)): - p[r] = 0 - for i in range(len(lines)): - - # -------------------------------------------------------- - # Line not empty - # -------------------------------------------------------- - if len(lines[i].split()) < 10: - continue - - # -------------------------------------------------------- - # Read line - # -------------------------------------------------------- - ( - roa_read, - aLTe_read, - aLTi_read, - aLne_read, - Q_gb_read, - Qe_read, - Qi_read, - Ge_read, - GZ_read, - Mt_read, - Pexch_read, - Qe_std_read, - Qi_std_read, - Ge_std_read, - GZ_std_read, - Mt_std_read, - Pexch_std_read, - tstart_read, - tend_read, - ) = readlineNTH(lines[i], unnormalize=unnormalize) - - # -------------------------------------------------------- - # Radial location position - # -------------------------------------------------------- - threshold_radii = 1E-4 - r = np.where(np.abs(radii-roa_read) powerstate.plasma) - tgyro_to_powerstate( - self, - tgyro.results[label], - Qi_includes_fast=Qi_includes_fast, - impurityPosition=impurityPosition, - UseFineGridTargets=UseFineGridTargets, - OriginalFimp=OriginalFimp, - forceZeroParticleFlux=forceZeroParticleFlux, - provideTargets=self.powerstate.target_options['target_evaluator_options']["target_evaluator_method"] == "tgyro", - ) - - tgyro.results["use"] = tgyro.results[label] - - # Copy profiles to share - self._profiles_to_store() - - # ------------------------------------------------------------------------------------------------------------------------ - # Results class that can be used for further plotting and analysis in PORTALS - # ------------------------------------------------------------------------------------------------------------------------ - - self.model_results = copy.deepcopy(tgyro.results["use"]) # Pass the TGYRO results class that should be use for plotting and analysis - - self.model_results.extra_analysis = {} - for ikey in tgyro.results: - if ikey != "use": - self.model_results.extra_analysis[ikey] = tgyro.results[ikey] - - def _profiles_to_store(self): - - if "folder" in self.powerstate.transport_options["transport_evaluator_options"]: - whereFolder = IOtools.expandPath(self.powerstate.transport_options["folder"] / "Outputs" / "portals_profiles") - if not whereFolder.exists(): - IOtools.askNewFolder(whereFolder) - - fil = whereFolder / f"input.gacode.{self.evaluation_number}" - shutil.copy2(self.file_profs, fil) - shutil.copy2(self.file_profs_unmod, fil.parent / f"{fil.name}_unmodified") - shutil.copy2(self.file_profs_targets, fil.parent / f"{fil.name}.new_fromtgyro_modified") - print(f"\t- Copied profiles to {IOtools.clipstr(fil)}") - else: - print("\t- Could not move files", typeMsg="w") - -def tglf_scan_trick( - tglf, - rho_locations, - predicted_channels, - impurityPosition=1, - Qi_includes_fast=False, - delta=0.02, - minimum_abs_gradient=0.005, # This is 0.5% of aLx=1.0, to avoid extremely small scans when, for example, having aLn ~ 0.0 - cold_start=False, - extra_name="", - remove_folders_out = False, - cores_per_tglf_instance = 4 # e.g. 4 core per radius, since this is going to launch ~ Nr=5 x (Nv=6 x Nd=2 + 1) = 65 TGLFs at once - ): - - print(f"\t- Running TGLF standalone scans ({delta = }) to determine relative errors") - - # Prepare scan - variables_to_scan = [] - for i in predicted_channels: - if i == 'te': variables_to_scan.append('RLTS_1') - if i == 'ti': variables_to_scan.append('RLTS_2') - if i == 'ne': variables_to_scan.append('RLNS_1') - if i == 'nZ': variables_to_scan.append(f'RLNS_{impurityPosition+2}') - if i == 'w0': variables_to_scan.append('VEXB_SHEAR') #TODO: is this correct? or VPAR_SHEAR? - - #TODO: Only if that parameter is changing at that location - if 'te' in predicted_channels or 'ti' in predicted_channels: - variables_to_scan.append('TAUS_2') - if 'te' in predicted_channels or 'ne' in predicted_channels: - variables_to_scan.append('XNUE') - if 'te' in predicted_channels or 'ne' in predicted_channels: - variables_to_scan.append('BETAE') - - relative_scan = [1-delta, 1+delta] - - # Enforce at least "minimum_abs_gradient" in gradient, to avoid zero gradient situations - minimum_delta_abs = {} - for ikey in variables_to_scan: - if 'RL' in ikey: - minimum_delta_abs[ikey] = minimum_abs_gradient - - name = 'turb_drives' - - tglf.rhos = rho_locations # To avoid the case in which TGYRO was run with an extra rho point - - # Estimate job minutes based on cases and cores (mostly IO I think at this moment, otherwise it should be independent on cases) - num_cases = len(rho_locations) * len(variables_to_scan) * len(relative_scan) - if cores_per_tglf_instance == 1: - minutes = 10 * (num_cases / 60) # Ad-hoc formula - else: - minutes = 1 * (num_cases / 60) # Ad-hoc formula - - # Enforce minimum minutes - minutes = max(2, minutes) - - tglf.runScanTurbulenceDrives( - subfolder = name, - variablesDrives = variables_to_scan, - varUpDown = relative_scan, - minimum_delta_abs = minimum_delta_abs, - TGLFsettings = None, - ApplyCorrections = False, - add_baseline_to = 'first', - cold_start=cold_start, - forceIfcold_start=True, - slurm_setup={ - "cores": cores_per_tglf_instance, - "minutes": minutes, - }, - extra_name = f'{extra_name}_{name}', - positionIon=impurityPosition+2, - attempts_execution=2, - only_minimal_files=True, # Since I only care about fluxes here, do not retrieve all the files - ) - - # Remove folders because they are heavy to carry many throughout - if remove_folders_out: - IOtools.shutil_rmtree(tglf.FolderGACODE) - - Qe = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - Qi = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - Qifast = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - Ge = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - GZ = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - Mt = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - S = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - - cont = 0 - for vari in variables_to_scan: - jump = tglf.scans[f'{name}_{vari}']['Qe'].shape[-1] - - Qe[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qe'] - Qi[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qi'] - Qifast[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qifast'] - Ge[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Ge'] - GZ[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Gi'] - Mt[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Mt'] - S[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['S'] - cont += jump - - if Qi_includes_fast: - print(f"\t- Qi includes fast ions, adding their contribution") - Qi += Qifast - - # Calculate the standard deviation of the scans, that's going to be the reported stds - - def calculate_mean_std(Q): - # Assumes Q is [radii, points], with [radii, 0] being the baseline - - Qm = np.mean(Q, axis=1) - Qstd = np.std(Q, axis=1) - - # Qm = Q[:,0] - # Qstd = np.std(Q, axis=1) - - # Qstd = ( Q.max(axis=1)-Q.min(axis=1) )/2 /2 # Such that the range is 2*std - # Qm = Q.min(axis=1) + Qstd*2 # Mean is at the middle of the range - - return Qm, Qstd - - Qe_point, Qe_std = calculate_mean_std(Qe) - Qi_point, Qi_std = calculate_mean_std(Qi) - Ge_point, Ge_std = calculate_mean_std(Ge) - GZ_point, GZ_std = calculate_mean_std(GZ) - Mt_point, Mt_std = calculate_mean_std(Mt) - S_point, S_std = calculate_mean_std(S) - - #TODO: Careful with fast particles - - Flux_base = [Qe[:,0], Qi[:,0], Ge[:,0], GZ[:,0], Mt[:,0], S[:,0]] - Flux_mean = [Qe_point, Qi_point, Ge_point, GZ_point, Mt_point, S_point] - Flux_std = [Qe_std, Qi_std, Ge_std, GZ_std, Mt_std, S_std] - - return Flux_base, Flux_mean, Flux_std -# ************************************************************************************************** -# Functions -# ************************************************************************************************** - -def curateTGYROfiles( - tgyroObject, - label, - rho_locations, - predicted_channels, - folder, - percent_error, - impurityPosition=1, - Qi_includes_fast=False, - use_tglf_scan_trick=None, - cold_start=False, - extra_name="", - cores_per_tglf_instance = 4, - check_coincidence_thr=1E-2, - ): - - tgyro = tgyroObject.results[label] - - # Determine NEO and Target errors - relativeErrorNEO = percent_error[1] / 100.0 - relativeErrorTAR = percent_error[2] / 100.0 - - # Grab fluxes from TGYRO - Qe = tgyro.Qe_sim_turb[0, 1:] - Qi = tgyro.QiIons_sim_turb[0, 1:] if Qi_includes_fast else tgyro.QiIons_sim_turb_thr[0, 1:] - Ge = tgyro.Ge_sim_turb[0, 1:] - GZ = tgyro.Gi_sim_turb[impurityPosition, 0, 1:] - Mt = tgyro.Mt_sim_turb[0, 1:] - Pexch = tgyro.EXe_sim_turb[0, 1:] - - # Determine TGLF standard deviations - if use_tglf_scan_trick is not None: - - # Grab TGLF object - tglfObject = tgyroObject.grab_tglf_objects(fromlabel=label, subfolder = 'tglf_explorations') - - # Run TGLF scan trick - Flux_base, Flux_mean, Flux_std = tglf_scan_trick( - tglfObject, - rho_locations, - predicted_channels, - impurityPosition=impurityPosition, - Qi_includes_fast=Qi_includes_fast, - delta = use_tglf_scan_trick, - cold_start=cold_start, - extra_name=extra_name, - cores_per_tglf_instance=cores_per_tglf_instance - ) - - Qe, Qi, Ge, GZ, Mt, Pexch = Flux_mean - QeE, QiE, GeE, GZE, MtE, PexchE = Flux_std - - # ---------------------------------------------------- - # Do a check that TGLF scans are consistent with TGYRO - - Qe_base, Qi_base, Ge_base, GZ_base, Mt_base, S_base = Flux_base - - # Grab fluxes from TGYRO - Qe_tgyro = tgyro.Qe_sim_turb[0, 1:] - Qi_tgyro = tgyro.QiIons_sim_turb[0, 1:] if Qi_includes_fast else tgyro.QiIons_sim_turb_thr[0, 1:] - Ge_tgyro = tgyro.Ge_sim_turb[0, 1:] - GZ_tgyro = tgyro.Gi_sim_turb[impurityPosition, 0, 1:] - Mt_tgyro = tgyro.Mt_sim_turb[0, 1:] - Pexch_tgyro = tgyro.EXe_sim_turb[0, 1:] - - Qe_err = np.abs( (Qe_base - Qe_tgyro) / Qe_tgyro ) if 'te' in predicted_channels else np.zeros_like(Qe_base) - Qi_err = np.abs( (Qi_base - Qi_tgyro) / Qi_tgyro ) if 'ti' in predicted_channels else np.zeros_like(Qi_base) - Ge_err = np.abs( (Ge_base - Ge_tgyro) / Ge_tgyro ) if 'ne' in predicted_channels else np.zeros_like(Ge_base) - GZ_err = np.abs( (GZ_base - GZ_tgyro) / GZ_tgyro ) if 'nZ' in predicted_channels else np.zeros_like(GZ_base) - Mt_err = np.abs( (Mt_base - Mt_tgyro) / Mt_tgyro ) if 'w0' in predicted_channels else np.zeros_like(Mt_base) - Pexch_err = np.abs( (Pexch - Pexch_tgyro) / Pexch_tgyro ) - - F_err = np.concatenate((Qe_err, Qi_err, Ge_err, GZ_err, Mt_err, Pexch_err)) - if F_err.max() > check_coincidence_thr: - print(f"\t- TGLF scans are not consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%, in quantity:",typeMsg="w") - if ('te' in predicted_channels) and Qe_err.max() > check_coincidence_thr: - print('\t\t* Qe:',Qe_err) - if ('ti' in predicted_channels) and Qi_err.max() > check_coincidence_thr: - print('\t\t* Qi:',Qi_err) - if ('ne' in predicted_channels) and Ge_err.max() > check_coincidence_thr: - print('\t\t* Ge:',Ge_err) - if ('nZ' in predicted_channels) and GZ_err.max() > check_coincidence_thr: - print('\t\t* GZ:',GZ_err) - if ('w0' in predicted_channels) and Mt_err.max() > check_coincidence_thr: - print('\t\t* Mt:',Mt_err) - if Pexch_err.max() > check_coincidence_thr: - print('\t\t* Pexch:',Pexch_err) - else: - print(f"\t- TGLF scans are consistent with TGYRO, maximum error = {F_err.max()*100:.2f}%") - # ---------------------------------------------------- - - min_relative_error = 0.01 # To avoid problems with gpytorch, 1% error minimum - - QeE = QeE.clip(abs(Qe)*min_relative_error) - QiE = QiE.clip(abs(Qi)*min_relative_error) - GeE = GeE.clip(abs(Ge)*min_relative_error) - GZE = GZE.clip(abs(GZ)*min_relative_error) - MtE = MtE.clip(abs(Mt)*min_relative_error) - PexchE = PexchE.clip(abs(Pexch)*min_relative_error) - - else: - - # -------------------------------------------------------------- - # If simply a percentage error provided - # -------------------------------------------------------------- - - relativeErrorTGLF = [percent_error[0] / 100.0]*len(rho_locations) - - QeE = abs(Qe) * relativeErrorTGLF - QiE = abs(Qi) * relativeErrorTGLF - GeE = abs(Ge) * relativeErrorTGLF - GZE = abs(GZ) * relativeErrorTGLF - MtE = abs(Mt) * relativeErrorTGLF - PexchE = abs(Pexch) * relativeErrorTGLF - - # ************************************************************************************************************************** - # Neo - # ************************************************************************************************************************** - - Qe_tr_neoc = tgyro.Qe_sim_neo[0, 1:] - if Qi_includes_fast: - Qi_tr_neoc = tgyro.QiIons_sim_neo[0, 1:] - else: - Qi_tr_neoc = tgyro.QiIons_sim_neo_thr[0, 1:] - Ge_tr_neoc = tgyro.Ge_sim_neo[0, 1:] - GZ_tr_neoc = tgyro.Gi_sim_neo[impurityPosition, 0, 1:] - Mt_tr_neoc = tgyro.Mt_sim_neo[0, 1:] - - Qe_tr_neocE = abs(tgyro.Qe_sim_neo[0, 1:]) * relativeErrorNEO - if Qi_includes_fast: - Qi_tr_neocE = abs(tgyro.QiIons_sim_neo[0, 1:]) * relativeErrorNEO - else: - Qi_tr_neocE = abs(tgyro.QiIons_sim_neo_thr[0, 1:]) * relativeErrorNEO - Ge_tr_neocE = abs(tgyro.Ge_sim_neo[0, 1:]) * relativeErrorNEO - GZ_tr_neocE = abs(tgyro.Gi_sim_neo[impurityPosition, 0, 1:]) * relativeErrorNEO - Mt_tr_neocE = abs(tgyro.Mt_sim_neo[0, 1:]) * relativeErrorNEO - - # Merge - - modifyFLUX( - tgyro, - folder, - Qe, - Qi, - Ge, - GZ, - Mt, - Pexch, - Qe_tr_neoc=Qe_tr_neoc, - Qi_tr_neoc=Qi_tr_neoc, - Ge_tr_neoc=Ge_tr_neoc, - GZ_tr_neoc=GZ_tr_neoc, - Mt_tr_neoc=Mt_tr_neoc, - impurityPosition=impurityPosition, - ) - - modifyFLUX( - tgyro, - folder, - QeE, - QiE, - GeE, - GZE, - MtE, - PexchE, - Qe_tr_neoc=Qe_tr_neocE, - Qi_tr_neoc=Qi_tr_neocE, - Ge_tr_neoc=Ge_tr_neocE, - GZ_tr_neoc=GZ_tr_neocE, - Mt_tr_neoc=Mt_tr_neocE, - impurityPosition=impurityPosition, - special_label="_stds", - ) - - # ************************************************************************************************************************** - # Targets - # ************************************************************************************************************************** - - QeTargetE = abs(tgyro.Qe_tar[0, 1:]) * relativeErrorTAR - QiTargetE = abs(tgyro.Qi_tar[0, 1:]) * relativeErrorTAR - GeTargetE = abs(tgyro.Ge_tar[0, 1:]) * relativeErrorTAR - GZTargetE = GeTargetE * 0.0 - MtTargetE = abs(tgyro.Mt_tar[0, 1:]) * relativeErrorTAR - - modifyEVO( - tgyro, - folder, - QeTargetE * 0.0, - QiTargetE * 0.0, - GeTargetE * 0.0, - GZTargetE * 0.0, - MtTargetE * 0.0, - impurityPosition=impurityPosition, - positionMod=1, - special_label="_stds", - ) - modifyEVO( - tgyro, - folder, - QeTargetE, - QiTargetE, - GeTargetE, - GZTargetE, - MtTargetE, - impurityPosition=impurityPosition, - positionMod=2, - special_label="_stds", - ) - -def dummyCDF(GeneralFolder, FolderEvaluation): - """ - This routine creates path to a dummy CDF file in FolderEvaluation, with the name "simulation_evaluation.CDF" - - GeneralFolder, e.g. ~/runs_portals/run10/ - FolderEvaluation, e.g. ~/runs_portals/run10000/Execution/Evaluation.0/transport_simulation_folder/ - """ - - # ------- Name construction for scratch folders in parallel ---------------- - - GeneralFolder = IOtools.expandPath(GeneralFolder, ensurePathValid=True) - - a, subname = IOtools.reducePathLevel(GeneralFolder, level=1, isItFile=False) - - FolderEvaluation = IOtools.expandPath(FolderEvaluation) - - name = FolderEvaluation.name.split(".")[-1] # 0 (evaluation #) - - if name == "": - name = "0" - - cdf = FolderEvaluation / f"{subname}_ev{name}.CDF" - - return cdf - -def modifyResults( - Qe, - Qi, - Ge, - GZ, - Mt, - Pexch, - QeE, - QiE, - GeE, - GZE, - MtE, - PexchE, - tgyro, - folder_tgyro, - minErrorPercent=5.0, - Qi_stable_criterion=0.0025, - impurityPosition=3, - OriginalFimp=1.0, -): - """ - All in real units, with dimensions of (rho) from axis to edge - """ - - # If a plasma is very close to stable... do something about error - if minErrorPercent is not None: - ( - Qe_target, - Qi_target, - Ge_target_special, - GZ_target_special, - Mt_target, - ) = defineReferenceFluxes( - tgyro, - impurityPosition=impurityPosition, - ) - - Qe_min = Qe_target * (minErrorPercent / 100.0) - Qi_min = Qi_target * (minErrorPercent / 100.0) - Ge_min = Ge_target_special * (minErrorPercent / 100.0) - GZ_min = GZ_target_special * (minErrorPercent / 100.0) - Mt_min = Mt_target * (minErrorPercent / 100.0) - - for i in range(Qe.shape[0]): - if Qi[i] < Qi_stable_criterion: - print( - f"\t- Based on 'Qi_stable_criterion', plasma considered stable (Qi = {Qi[i]:.2e} < {Qi_stable_criterion:.2e} MW/m2) at position #{i}, using minimum errors of {minErrorPercent}% of targets", - typeMsg="w", - ) - QeE[i] = Qe_min[i] - print(f"\t\t* QeE = {QeE[i]}") - QiE[i] = Qi_min[i] - print(f"\t\t* QiE = {QiE[i]}") - GeE[i] = Ge_min[i] - print(f"\t\t* GeE = {GeE[i]}") - GZE[i] = GZ_min[i] - print(f"\t\t* GZE = {GZE[i]}") - MtE[i] = Mt_min[i] - print(f"\t\t* MtE = {MtE[i]}") - - # Heat fluxes - QeTot = Qe + tgyro.Qe_sim_neo[0, 1:] - QiTot = Qi + tgyro.QiIons_sim_neo_thr[0, 1:] - - # Particle fluxes - PeTot = Ge + tgyro.Ge_sim_neo[0, 1:] - PZTot = GZ + tgyro.Gi_sim_neo[impurityPosition, 0, 1:] - - # Momentum fluxes - MtTot = Mt + tgyro.Mt_sim_neo[0, 1:] - - # ************************************************************************************ - # **** Modify complete folder (Division of ion fluxes will be wrong, since I put everything in first ion) - # ************************************************************************************ - - # 1. Modify out.tgyro.evo files (which contain turb+neo summed together) - - print(f"\t- Modifying TGYRO out.tgyro.evo files in {IOtools.clipstr(folder_tgyro)}") - modifyEVO( - tgyro, - folder_tgyro, - QeTot, - QiTot, - PeTot, - PZTot, - MtTot, - impurityPosition=impurityPosition, - ) - - # 2. Modify out.tgyro.flux files (which contain turb and neo separated) - - print(f"\t- Modifying TGYRO out.tgyro.flux files in {folder_tgyro}") - modifyFLUX( - tgyro, - folder_tgyro, - Qe, - Qi, - Ge, - GZ, - Mt, - Pexch, - impurityPosition=impurityPosition, - ) - - # 3. Modify files for errors - - print(f"\t- Modifying TGYRO out.tgyro.flux_stds in {folder_tgyro}") - modifyFLUX( - tgyro, - folder_tgyro, - QeE, - QiE, - GeE, - GZE, - MtE, - PexchE, - impurityPosition=impurityPosition, - special_label="_stds", - ) - - -def modifyEVO( - tgyro, - folder, - QeT, - QiT, - GeT, - GZT, - MtT, - impurityPosition=3, - positionMod=1, - special_label=None, -): - QeTGB = QeT / tgyro.Q_GB[-1, 1:] - QiTGB = QiT / tgyro.Q_GB[-1, 1:] - GeTGB = GeT / tgyro.Gamma_GB[-1, 1:] - GZTGB = GZT / tgyro.Gamma_GB[-1, 1:] - MtTGB = MtT / tgyro.Pi_GB[-1, 1:] - - modTGYROfile(folder / "out.tgyro.evo_te", QeTGB, pos=positionMod, fileN_suffix=special_label) - modTGYROfile(folder / "out.tgyro.evo_ti", QiTGB, pos=positionMod, fileN_suffix=special_label) - modTGYROfile(folder / "out.tgyro.evo_ne", GeTGB, pos=positionMod, fileN_suffix=special_label) - modTGYROfile(folder / "out.tgyro.evo_er", MtTGB, pos=positionMod, fileN_suffix=special_label) - - for i in range(tgyro.Qi_sim_turb.shape[0]): - if i == impurityPosition: - var = GZTGB - else: - var = GZTGB * 0.0 - modTGYROfile( - folder / f"out.tgyro.evo_n{i+1}", - var, - pos=positionMod, - fileN_suffix=special_label, - ) - - -def modifyFLUX( - tgyro, - folder, - Qe, - Qi, - Ge, - GZ, - Mt, - S, - Qe_tr_neoc=None, - Qi_tr_neoc=None, - Ge_tr_neoc=None, - GZ_tr_neoc=None, - Mt_tr_neoc=None, - impurityPosition=3, - special_label=None, -): - folder = IOtools.expandPath(folder) - - QeGB = Qe / tgyro.Q_GB[-1, 1:] - QiGB = Qi / tgyro.Q_GB[-1, 1:] - GeGB = Ge / tgyro.Gamma_GB[-1, 1:] - GZGB = GZ / tgyro.Gamma_GB[-1, 1:] - MtGB = Mt / tgyro.Pi_GB[-1, 1:] - SGB = S / tgyro.S_GB[-1, 1:] - - # ****************************************************************************************** - # Electrons - # ****************************************************************************************** - - # Particle flux: Update - - modTGYROfile(folder / "out.tgyro.flux_e", GeGB, pos=2, fileN_suffix=special_label) - if Ge_tr_neoc is not None: - GeGB_neo = Ge_tr_neoc / tgyro.Gamma_GB[-1, 1:] - modTGYROfile(folder / "out.tgyro.flux_e", GeGB_neo, pos=1, fileN_suffix=special_label) - - # Energy flux: Update - - modTGYROfile(folder / "out.tgyro.flux_e", QeGB, pos=4, fileN_suffix=special_label) - if Qe_tr_neoc is not None: - QeGB_neo = Qe_tr_neoc / tgyro.Q_GB[-1, 1:] - modTGYROfile(folder / "out.tgyro.flux_e", QeGB_neo, pos=3, fileN_suffix=special_label) - - # Rotation: Remove (it will be sum to the first ion) - - modTGYROfile(folder / "out.tgyro.flux_e", GeGB * 0.0, pos=6, fileN_suffix=special_label) - modTGYROfile(folder / "out.tgyro.flux_e", GeGB * 0.0, pos=5, fileN_suffix=special_label) - - # Energy exchange - - modTGYROfile(folder / "out.tgyro.flux_e", SGB, pos=7, fileN_suffix=special_label) - - # SMW = S # S is MW/m^3 - # modTGYROfile(f'{folder}/out.tgyro.power_e',SMW,pos=8,fileN_suffix=special_label) - # print('\t\t- Modified turbulent energy exchange in out.tgyro.power_e') - - # ****************************************************************************************** - # Ions - # ****************************************************************************************** - - # Energy flux: Update - - modTGYROfile(folder / "out.tgyro.flux_i1", QiGB, pos=4, fileN_suffix=special_label) - - if Qi_tr_neoc is not None: - QiGB_neo = Qi_tr_neoc / tgyro.Q_GB[-1, 1:] - modTGYROfile(folder / "out.tgyro.flux_i1", QiGB_neo, pos=3, fileN_suffix=special_label) - - # Particle flux: Make ion particle fluxes zero, because I don't want to mistake TGLF with CGYRO when looking at tgyro results - - for i in range(tgyro.Qi_sim_turb.shape[0]): - if tgyro.profiles.Species[i]["S"] == "therm": - var = QiGB * 0.0 - modTGYROfile(folder / f"out.tgyro.flux_i{i+1}",var,pos=2,fileN_suffix=special_label,) # Gi_turb - modTGYROfile(folder / f"out.tgyro.evo_n{i+1}", var, pos=1, fileN_suffix=special_label) # Gi (Gi_sim) - - if i != impurityPosition: - modTGYROfile(folder / f"out.tgyro.flux_i{i+1}",var,pos=1,fileN_suffix=special_label) # Gi_neo - - # Rotation: Update - - modTGYROfile(folder / "out.tgyro.flux_i1", MtGB, pos=6, fileN_suffix=special_label) - - if Mt_tr_neoc is not None: - MtGB_neo = Mt_tr_neoc / tgyro.Pi_GB[-1, 1:] - modTGYROfile(folder / "out.tgyro.flux_i1", MtGB_neo, pos=5, fileN_suffix=special_label) - - # Energy exchange: Remove (it will be the electrons one) - - modTGYROfile(folder / "out.tgyro.flux_i1", SGB * 0.0, pos=7, fileN_suffix=special_label) - - # ****************************************************************************************** - # Impurities - # ****************************************************************************************** - - # Remove everything from all the rest of non-first ions (except the particles for the impurity chosen) - - for i in range(tgyro.Qi_sim_turb.shape[0] - 1): - if tgyro.profiles.Species[i + 1]["S"] == "therm": - var = QiGB * 0.0 - for pos in [3, 4, 5, 6, 7]: - modTGYROfile(folder / f"out.tgyro.flux_i{i+2}",var,pos=pos,fileN_suffix=special_label) - for pos in [1, 2]: - if i + 2 != impurityPosition: - modTGYROfile(folder / f"out.tgyro.flux_i{i+2}",var,pos=pos,fileN_suffix=special_label) - - modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB,pos=2,fileN_suffix=special_label) - if GZ_tr_neoc is not None: - GZGB_neo = GZ_tr_neoc / tgyro.Gamma_GB[-1, 1:] - modTGYROfile(folder / f"out.tgyro.flux_i{impurityPosition+1}",GZGB_neo,pos=1,fileN_suffix=special_label) - - -def modTGYROfile(file, var, pos=0, fileN_suffix=None): - fileN = file if fileN_suffix is None else file.parent / f"{file.name}{fileN_suffix}" - - if not fileN.exists(): - shutil.copy2(file, fileN) - - with open(fileN, "r") as f: - lines = f.readlines() - - with open(fileN, "w") as f: - f.write(lines[0]) - f.write(lines[1]) - f.write(lines[2]) - for i in range(var.shape[0]): - new_s = [float(k) for k in lines[3 + i].split()] - new_s[pos] = var[i] - - line_new = " " - for k in range(len(new_s)): - line_new += f'{"" if k==0 else " "}{new_s[k]:.6e}' - f.write(line_new + "\n") - -def defineReferenceFluxes( - tgyro, factor_tauptauE=5, impurityPosition=3 -): - Qe_target = abs(tgyro.Qe_tar[0, 1:]) - Qi_target = abs(tgyro.Qi_tar[0, 1:]) - Mt_target = abs(tgyro.Mt_tar[0, 1:]) - - # For particle fluxes, since the targets are often zero... it's more complicated - QeMW_target = abs(tgyro.Qe_tarMW[0, 1:]) - QiMW_target = abs(tgyro.Qi_tarMW[0, 1:]) - We, Wi, Ne, NZ = tgyro.profiles.deriveContentByVolumes( - rhos=tgyro.rho[0, 1:], impurityPosition=impurityPosition - ) - - tau_special = ( - (We + Wi) / (QeMW_target + QiMW_target) * factor_tauptauE - ) # tau_p in seconds - Ge_target_special = (Ne / tau_special) / tgyro.dvoldr[0, 1:] # (1E20/seconds/m^2) - - Ge_target_special = PLASMAtools.convective_flux( - tgyro.Te[0, 1:], Ge_target_special - ) # (1E20/seconds/m^2) - - GZ_target_special = Ge_target_special * NZ / Ne - - return Qe_target, Qi_target, Ge_target_special, GZ_target_special, Mt_target - - - -# ------------------------------------------------------------------------------------------------------------------------------------------------------ -# This is where the definitions for the summation variables happen for mitim and PORTALSplot -# ------------------------------------------------------------------------------------------------------------------------------------------------------ - -def tgyro_to_powerstate( - self, - TGYROresults, - forceZeroParticleFlux=False, - Qi_includes_fast=False, - impurityPosition=1, - UseFineGridTargets=False, - OriginalFimp=1.0, - provideTargets=False - ): - """ - This function is used to extract the TGYRO results and store them in the powerstate object, from numpy arrays to torch tensors. - """ - - if "tgyro_stds" not in TGYROresults.__dict__: - TGYROresults.tgyro_stds = False - - if UseFineGridTargets: - TGYROresults.useFineGridTargets(impurityPosition=impurityPosition) - - nr = self.powerstate.plasma['rho'].shape[-1] - if self.powerstate.plasma['rho'].shape[-1] != TGYROresults.rho.shape[-1]: - print('\t- TGYRO was run with an extra point in the grid, treating it carefully now') - - # ********************************** - # *********** Electron Energy Fluxes - # ********************************** - - self.QeGB_turb = TGYROresults.QeGB_sim_turb[0, 1:nr] - self.QeGB_neoc = TGYROresults.QeGB_sim_neo[0, 1:nr] - - self.QeGB_turb_stds = TGYROresults.QeGB_sim_turb_stds[0, 1:nr] - self.QeGB_neoc_stds = TGYROresults.QeGB_sim_neo_stds[0, 1:nr] - - # ********************************** - # *********** Ion Energy Fluxes - # ********************************** - - if Qi_includes_fast: - - self.QiGB_turb = TGYROresults.QiIons_sim_turb[0, 1:nr] - self.QiGB_neoc = TGYROresults.QiIons_sim_neo[0, 1:nr] - - self.QiGB_turb_stds = TGYROresults.QiGBIons_sim_turb_stds[0, 1:nr] - self.QiGB_neoc_stds = TGYROresults.QiGBIons_sim_neo_stds[0, 1:nr] - - else: - - self.QiGB_turb = TGYROresults.QiGBIons_sim_turb_thr[0, 1:nr] - self.QiGB_neoc = TGYROresults.QiGBIons_sim_neo_thr[0, 1:nr] - - self.QiGB_turb_stds = TGYROresults.QiGBIons_sim_turb_thr_stds[0, 1:nr] - self.QiGB_neoc_stds = TGYROresults.QiGBIons_sim_neo_thr_stds[0, 1:nr] - - # ********************************** - # *********** Momentum Fluxes - # ********************************** - - self.MtGB_turb = TGYROresults.MtGB_sim_turb[0, 1:nr] # So far, let's include fast in momentum - self.MtGB_neoc = TGYROresults.MtGB_sim_neo[0, 1:nr] - - self.MtGB_turb_stds = TGYROresults.MtGB_sim_turb_stds[0, 1:nr] - self.MtGB_neoc_stds = TGYROresults.MtGB_sim_neo_stds[0, 1:nr] - - # ********************************** - # *********** Particle Fluxes - # ********************************** - - # Store raw fluxes for better plotting later - self.GeGB_turb = TGYROresults.GeGB_sim_turb[0, 1:nr] - self.GeGB_neoc = TGYROresults.GeGB_sim_neo[0, 1:nr] - - self.GeGB_turb_stds = TGYROresults.GeGB_sim_turb_stds[0, 1:nr] - self.GeGB_neoc_stds = TGYROresults.GeGB_sim_neo_stds[0, 1:nr] - - # ********************************** - # *********** Impurity Fluxes - # ********************************** - - # Store raw fluxes for better plotting later - self.GZGB_turb = TGYROresults.GiGB_sim_turb[impurityPosition, 0, 1:nr] - self.GZGB_neoc = TGYROresults.GiGB_sim_neo[impurityPosition, 0, 1:nr] - - self.GZGB_turb_stds = TGYROresults.GiGB_sim_turb_stds[impurityPosition, 0, 1:nr] - self.GZGB_neoc_stds = TGYROresults.GiGB_sim_neo_stds[impurityPosition, 0, 1:nr] - - # ********************************** - # *********** Energy Exchange - # ********************************** - - self.QieGB_turb = TGYROresults.EXeGB_sim_turb[0, 1:nr] - self.QieGB_turb_stds = TGYROresults.EXeGB_sim_turb_stds[0, 1:nr] - - self.QieGB_neoc = self.QeGB_turb * 0.0 - self.QieGB_neoc_stds = self.QeGB_turb_stds * 0.0 - - # ********************************** - # *********** Targets - # *********************************** - - if provideTargets: - self.powerstate.plasma["QeMWm2"] = torch.Tensor(TGYROresults.Qe_tar[0, 1:nr]).to(self.powerstate.dfT) - self.powerstate.plasma["QeMWm2_stds"] = torch.Tensor(TGYROresults.Qe_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None - self.powerstate.plasma["QiMWm2"] = torch.Tensor(TGYROresults.Qi_tar[0, 1:nr]).to(self.powerstate.dfT) - self.powerstate.plasma["QiMWm2_stds"] = torch.Tensor(TGYROresults.Qi_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None - self.powerstate.plasma["MtJm2"] = torch.Tensor(TGYROresults.Mt_tar[0, 1:nr]).to(self.powerstate.dfT) - self.powerstate.plasma["MtJm2_stds"] = torch.Tensor(TGYROresults.Mt_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None - self.powerstate.plasma["Ge1E20m2"] = torch.Tensor(TGYROresults.Ge_tar[0, 1:nr]).to(self.powerstate.dfT) - self.powerstate.plasma["Ge1E20m2_stds"] = torch.Tensor(TGYROresults.Ge_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None - self.powerstate.plasma["Ce"] = torch.Tensor(TGYROresults.Ce_tar[0, 1:nr]).to(self.powerstate.dfT) - self.powerstate.plasma["Ce_stds"] = torch.Tensor(TGYROresults.Ce_tar_stds[0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None - self.powerstate.plasma["GZ1E20m2"] = torch.Tensor(TGYROresults.Gi_tar[impurityPosition, 0, 1:nr]).to(self.powerstate.dfT) - self.powerstate.plasma["GZ1E20m2_stds"] = torch.Tensor(TGYROresults.Gi_tar_stds[impurityPosition, 0, 1:nr]).to(self.powerstate.dfT) if TGYROresults.tgyro_stds else None - self.powerstate.plasma["CZ"] = torch.Tensor(TGYROresults.Ci_tar[impurityPosition, 0, 1:nr]).to(self.powerstate.dfT) / OriginalFimp - self.powerstate.plasma["CZ_stds"] = torch.Tensor(TGYROresults.Ci_tar_stds[impurityPosition, 0, 1:nr]).to(self.powerstate.dfT) / OriginalFimp if TGYROresults.tgyro_stds else None - - if forceZeroParticleFlux: - self.powerstate.plasma["Ce"] = self.powerstate.plasma["Ce"] * 0.0 - self.powerstate.plasma["Ce_stds"] = self.powerstate.plasma["Ce_stds"] * 0.0 - diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index ea3b88df..010c7934 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -7,7 +7,7 @@ from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.powertorch import STATEtools -from mitim_modules.powertorch.physics_models import targets_analytic, transport_tgyro +from mitim_modules.powertorch.physics_models import targets_analytic, transport_tglfneo from IPython import embed def calculator( @@ -39,7 +39,7 @@ def calculator( "target_evaluator_method": "tgyro"}, }, transport_options={ - "transport_evaluator": transport_tgyro.tgyro_model, + "transport_evaluator": transport_tglfneo.tglfneo_model, "transport_evaluator_options": { "cold_start": cold_start, "portals_parameters": { From 87bc6af6bfdcd55f380a0a0291a1d309db800d80 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 15:25:12 -0400 Subject: [PATCH 227/385] Allow to provide model to portals as string --- regressions/portals_regressions.py | 26 +-- src/mitim_modules/portals/PORTALSmain.py | 204 +++++++++--------- src/mitim_modules/portals/PORTALStools.py | 30 +-- .../portals/utils/PORTALSanalysis.py | 28 +-- .../portals/utils/PORTALSinit.py | 118 +++++----- .../portals/utils/PORTALSoptimization.py | 2 +- .../portals/utils/PORTALSplot.py | 42 ++-- src/mitim_modules/powertorch/STATEtools.py | 32 +-- .../physics_models/targets_analytic.py | 4 +- .../physics_models/transport_analytic.py | 6 +- .../physics_models/transport_cgyroneo.py | 2 +- .../physics_models/transport_tglfneo.py | 6 +- .../powertorch/scripts/calculateTargets.py | 18 +- .../powertorch/scripts/compareWithTGYRO.py | 2 +- .../powertorch/utils/TARGETStools.py | 8 +- .../powertorch/utils/TRANSFORMtools.py | 20 +- .../powertorch/utils/TRANSPORTtools.py | 2 +- tests/PORTALS_workflow.py | 16 +- tests/POWERTORCH_workflow.py | 2 +- tutorials/PORTALS_tutorial.py | 12 +- 20 files changed, 290 insertions(+), 290 deletions(-) diff --git a/regressions/portals_regressions.py b/regressions/portals_regressions.py index fa05c91e..2a17dfcb 100644 --- a/regressions/portals_regressions.py +++ b/regressions/portals_regressions.py @@ -55,12 +55,12 @@ def conditions_regressions(variables): portals_fun = PORTALSmain.portals(folderWork) portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True + portals_fun.portals_parameters["initialization"]["remove_fast"] = True - portals_fun.portals_parameters["model_parameters"]["predicted_channels"] = ["te", "ti"] + portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti"] portals_fun.optimization_options["acquisition_options"]["optimizers"] = ["botorch"] - portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator"] = TRANSPORTtools.diffusion_model + portals_fun.portals_parameters["transport"]["evaluator"] = TRANSPORTtools.diffusion_model transport_evaluator_options = {'chi_e': torch.ones(5)*0.5,'chi_i': torch.ones(5)*2.0} portals_fun.prep(inputgacode, folderWork, transport_evaluator_options=transport_evaluator_options) @@ -93,13 +93,13 @@ def conditions_regressions(variables): portals_fun = PORTALSmain.portals(folderWork) portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - portals_fun.portals_parameters["model_parameters"]["predicted_rho"] = [0.25, 0.45, 0.65, 0.85] - portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True - portals_fun.portals_parameters["initialization_parameters"]["quasineutrality"] = True - portals_fun.portals_parameters["initialization_parameters"]["enforce_same_aLn"] = True - portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["TGLFsettings"] = 2 + portals_fun.portals_parameters["solution"]["predicted_rho"] = [0.25, 0.45, 0.65, 0.85] + portals_fun.portals_parameters["initialization"]["remove_fast"] = True + portals_fun.portals_parameters["initialization"]["quasineutrality"] = True + portals_fun.portals_parameters["initialization"]["enforce_same_aLn"] = True + portals_fun.portals_parameters["transport"]["options"]["TGLFsettings"] = 2 - portals_fun.portals_parameters["model_parameters"]["predicted_channels"] = ["te", "ti", "ne"] + portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti", "ne"] portals_fun.prep(inputgacode, folderWork) mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=False, askQuestions=False) @@ -134,12 +134,12 @@ def conditions_regressions(variables): # portals_fun = PORTALSmain.portals(folderWork) # portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 # portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - # portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True + # portals_fun.portals_parameters["initialization"]["remove_fast"] = True - # portals_fun.portals_parameters["model_parameters"]["predicted_channels"] = ["te", "ti", "ne",'nZ','w0'] + # portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti", "ne",'nZ','w0'] - # portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"] = 'W' - # portals_fun.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"] = True + # portals_fun.portals_parameters["solution"]["trace_impurity"] = 'W' + # portals_fun.portals_parameters["solution"]["turbulent_exchange_as_surrogate"] = True # portals_fun.prep(inputgacode, folderWork) # mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=False, askQuestions=False) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 500d31fc..c3478915 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -38,7 +38,7 @@ >2 will also plot profiles & gradients comparison (original, initial, best) """ -def default_namelist(optimization_options, CGYROrun=False): +def default_namelist(optimization_options): """ This is to be used after reading the namelist, so self.optimization_options should be completed with main defaults. """ @@ -62,16 +62,16 @@ def default_namelist(optimization_options, CGYROrun=False): optimization_options['acquisition_options']['relative_improvement_for_stopping'] = 1e-2 # Surrogate - optimization_options["surrogate_options"]["surrogate_selection"] = partial(PORTALStools.surrogate_selection_portals, CGYROrun=CGYROrun) + optimization_options["surrogate_options"]["surrogate_selection"] = partial(PORTALStools.surrogate_selection_portals) - if CGYROrun: - # CGYRO runs should prioritize accuracy - optimization_options["acquisition_options"]["type"] = "posterior_mean" - optimization_options["acquisition_options"]["optimizers"] = ["root", "botorch", "ga"] - else: - # TGLF runs should prioritize speed - optimization_options["acquisition_options"]["type"] = "posterior_mean" # "noisy_logei_mc" - optimization_options["acquisition_options"]["optimizers"] = ["sr", "root"] #, "botorch"] + # if CGYROrun: + # # CGYRO runs should prioritize accuracy + # optimization_options["acquisition_options"]["type"] = "posterior_mean" + # optimization_options["acquisition_options"]["optimizers"] = ["root", "botorch", "ga"] + # else: + # # TGLF runs should prioritize speed + optimization_options["acquisition_options"]["type"] = "posterior_mean" # "noisy_logei_mc" + optimization_options["acquisition_options"]["optimizers"] = ["sr", "root"] #, "botorch"] return optimization_options @@ -80,12 +80,12 @@ class portals(STRATEGYtools.opt_evaluator): def __init__( self, folder, # Folder where the PORTALS workflow will be run + transport_model = 'tglf', namelist=None, # If None, default namelist will be used. If not None, it will be read and used tensor_options = { "dtype": torch.double, "device": torch.device("cpu"), }, - CGYROrun=False, # If True, use CGYRO defaults for best optimization practices portals_transformation_variables = None, # If None, use defaults for both main and trace portals_transformation_variables_trace = None, additional_params_in_surrogate = [] # Additional parameters to be used in the surrogate (e.g. ['q']) @@ -105,7 +105,7 @@ def __init__( namelist=namelist, tensor_options=tensor_options, default_namelist_function=( - partial(default_namelist, CGYROrun=CGYROrun) + partial(default_namelist) if (namelist is None) else None ), @@ -132,13 +132,24 @@ def __init__( ---------------------------------------------------------------------------- """ - self.portals_parameters["main_parameters"] = { + self.portals_parameters["solution"] = { + + # Specification of radial locations (roa wins over rho, if provided) + "predicted_roa": [0.35, 0.55, 0.75, 0.875, 0.9], + "predicted_rho": [0.3, 0.45, 0.6, 0.75, 0.9], + + # Channels to be predicted + "predicted_channels": ["te", "ti", "ne"], # ['nZ','w0'] + + "trace_impurity": None, # Impurity to do flux-matching for if nZ enabled (name of first impurity instance AFTER postprocessing), e.g. "W" + + # PORTALS "portals_transformation_variables": portals_transformation_variables, # Physics-informed parameters to fit surrogates "portals_transformation_variables_trace": portals_transformation_variables_trace, # Physics-informed parameters to fit surrogates for trace impurities "turbulent_exchange_as_surrogate": False, # Run turbulent exchange as surrogate? "Pseudo_multipliers": [1.0]*5, # [Qe,Qi,Ge] multipliers to calculate pseudo - "applyImpurityGammaTrick": True, # If True, fit model to GZ/nZ, valid on the trace limit - "UseOriginalImpurityConcentrationAsWeight": 1.0, # If not None, using UseOriginalImpurityConcentrationAsWeight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis + "impurity_trick": True, # If True, fit model to GZ/nZ, valid on the trace limit + "fZ0_as_weight": 1.0, # If not None, using fZ0_as_weight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis "fImp_orig": 1.0, "additional_params_in_surrogate": additional_params_in_surrogate, "keep_full_model_folder": True, # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) @@ -153,7 +164,7 @@ def __init__( to run these corrections. """ - self.portals_parameters["initialization_parameters"] = { + self.portals_parameters["initialization"] = { "recalculate_ptot": True, # Recompute PTOT to match kinetic profiles (after removals) "quasineutrality": False, # Make sure things are quasineutral by changing the *MAIN* ion (D,T or both) (after removals) "removeIons": [], # Remove this ion from the input.gacode (if D,T,Z, eliminate T with [2]) @@ -172,94 +183,83 @@ def __init__( ''' # Selection of model - if CGYROrun: + if transport_model == 'cgyro': from mitim_modules.powertorch.physics_models.transport_cgyroneo import cgyroneo_model as transport_evaluator - else: + elif transport_model == 'tglf': from mitim_modules.powertorch.physics_models.transport_tglfneo import tglfneo_model as transport_evaluator + + # ------------------------- + # Transport model settings + # ------------------------- + self.portals_parameters["transport"] = { - from mitim_modules.powertorch.physics_models.targets_analytic import analytical_model as target_evaluator - - self.portals_parameters["model_parameters"] = { + # Transport model class - # Specification of radial locations (roa wins over rho, if provided) - "predicted_rho": [0.3, 0.45, 0.6, 0.75, 0.9], - "predicted_roa": None, + "evaluator": transport_evaluator, - # Channels to be predicted - "predicted_channels": ["te", "ti", "ne"], # ['nZ','w0'] - - # ------------------------- - # Transport model settings - # ------------------------- - "transport_parameters": { - - # Transport model class - - "transport_evaluator": transport_evaluator, - - # Simulation kwargs to be passed directly to run and read commands - - "transport_evaluator_options": { - "tglf": { - "run": { - "code_settings": 6, - "extraOptions": {}, - }, - "read": {}, - "use_scan_trick_for_stds": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta - "keep_files": "base", # minimal: only retrieve minimal files always; base: retrieve all files when running TGLF base; none: always minimal files - "cores_per_tglf_instance": 1, # Number of cores to use per TGLF instance - "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) - "percent_error": 5, # (%) Error (std, in percent) of model evaluation TGLF if not scan trick - "Qi_includes_fast": False, # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) - }, - "neo": { - "run": {}, - "read": {}, - "percent_error": 10 # (%) Error (std, in percent) of model evaluation - }, - "cgyro": { - "run": { - "code_settings": 1, - "extraOptions": {}, - "run_type": "normal", # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit - }, - "read": { - "tmin": 0.0 - }, - "Qi_stable_criterion": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable - "Qi_stable_percent_error": 5.0, # (%) For CGYRO runs, minimum error based on target if case is considered stable + # Simulation kwargs to be passed directly to run and read commands + + "options": { + "tglf": { + "run": { + "code_settings": 6, + "extraOptions": {}, }, + "read": {}, + "use_scan_trick_for_stds": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta + "keep_files": "base", # minimal: only retrieve minimal files always; base: retrieve all files when running TGLF base; none: always minimal files + "cores_per_tglf_instance": 1, # Number of cores to use per TGLF instance + "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) + "percent_error": 5, # (%) Error (std, in percent) of model evaluation TGLF if not scan trick + "Qi_includes_fast": False, # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) }, - - # Corrections to be applied to each iteration input.gacode file - - "profiles_postprocessing_fun": None, # Function to post-process input.gacode only BEFORE passing to transport codes - - "applyCorrections": { - "Ti_thermals": True, # Keep all thermal ion temperatures equal to the main Ti - "ni_thermals": True, # Adjust for quasineutrality by modifying the thermal ion densities together with ne - "recalculate_ptot": True, # Recompute PTOT to insert in input file each time - "Tfast_ratio": False, # Keep the ratio of Tfast/Te constant throughout the Te evolution - "ensureMachNumber": None, # Change w0 to match this Mach number when Ti varies + "neo": { + "run": {}, + "read": {}, + "percent_error": 10 # (%) Error (std, in percent) of model evaluation + }, + "cgyro": { + "run": { + "code_settings": 1, + "extraOptions": {}, + "run_type": "normal", # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit + }, + "read": { + "tmin": 0.0 + }, + "Qi_stable_criterion": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable + "Qi_stable_percent_error": 5.0, # (%) For CGYRO runs, minimum error based on target if case is considered stable }, }, - # ------------------------- - # Target model settings - # ------------------------- - "target_parameters": { - "target_evaluator": target_evaluator, - "target_evaluator_options": { - "TypeTarget": 3, - "target_evaluator": target_evaluator, - "target_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) - "forceZeroParticleFlux": False, # If True, ignore particle flux profile and assume zero for all radii - "fineTargetsResolution": 20, # If not None, calculate targets with this radial resolution (defaults target_evaluator_method to powerstate) - "percent_error": 1 # (%) Error (std, in percent) of model evaluation - } + # Corrections to be applied to each iteration input.gacode file + + "profiles_postprocessing_fun": None, # Function to post-process input.gacode only BEFORE passing to transport codes + + "applyCorrections": { + "Ti_thermals": True, # Keep all thermal ion temperatures equal to the main Ti + "ni_thermals": True, # Adjust for quasineutrality by modifying the thermal ion densities together with ne + "recalculate_ptot": True, # Recompute PTOT to insert in input file each time + "Tfast_ratio": False, # Keep the ratio of Tfast/Te constant throughout the Te evolution + "ensureMachNumber": None, # Change w0 to match this Mach number when Ti varies }, - "ImpurityOfInterest": None, # Impurity to do flux-matching for if nZ enabled (name of first impurity instance AFTER postprocessing), e.g. "W" + } + + # ------------------------- + # Target model settings + # ------------------------- + + from mitim_modules.powertorch.physics_models.targets_analytic import analytical_model as target_evaluator + + self.portals_parameters["target"] = { + "evaluator": target_evaluator, + "options": { + "TypeTarget": 3, + "target_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) + "forceZeroParticleFlux": False, # If True, ignore particle flux profile and assume zero for all radii + "fineTargetsResolution": 20, # If not None, calculate targets with this radial resolution (defaults target_evaluator_method to powerstate) + "percent_error": 1 # (%) Error (std, in percent) of model evaluation + } } # Grab all the flags here in a way that, after changing the dictionary extenrally, I make sure it's the same flags as PORTALS expects @@ -321,15 +321,15 @@ def prep( ymax_rel0 = copy.deepcopy(ymax_rel) ymax_rel = {} - for prof in self.portals_parameters["model_parameters"]["predicted_channels"]: - ymax_rel[prof] = np.array( [ymax_rel0] * len(self.portals_parameters["model_parameters"][key_rhos]) ) + for prof in self.portals_parameters["solution"]["predicted_channels"]: + ymax_rel[prof] = np.array( [ymax_rel0] * len(self.portals_parameters["solution"][key_rhos]) ) if IOtools.isfloat(ymin_rel): ymin_rel0 = copy.deepcopy(ymin_rel) ymin_rel = {} - for prof in self.portals_parameters["model_parameters"]["predicted_channels"]: - ymin_rel[prof] = np.array( [ymin_rel0] * len(self.portals_parameters["model_parameters"][key_rhos]) ) + for prof in self.portals_parameters["solution"]["predicted_channels"]: + ymin_rel[prof] = np.array( [ymin_rel0] * len(self.portals_parameters["solution"][key_rhos]) ) if enforceFiniteTemperatureGradients is not None: for prof in ['te', 'ti']: @@ -418,7 +418,7 @@ def run(self, paramsfile, resultsfile): name, numPORTALS=numPORTALS, dictOFs=dictOFs, - remove_folder_upon_completion=not self.portals_parameters["main_parameters"]["keep_full_model_folder"], + remove_folder_upon_completion=not self.portals_parameters["solution"]["keep_full_model_folder"], ) # Write results @@ -509,7 +509,7 @@ def _check_flags_dictionary(d, d_check, avoid = ["run", "read"]): _check_flags_dictionary(self.portals_parameters, self.potential_flags) - key_rhos = "predicted_roa" if self.portals_parameters["model_parameters"]["predicted_roa"] is not None else "predicted_rho" + key_rhos = "predicted_roa" if self.portals_parameters["solution"]["predicted_roa"] is not None else "predicted_rho" return key_rhos @@ -575,8 +575,8 @@ def reuseTrainingTabular( self_copy = copy.deepcopy(self) if reevaluateTargets == 1: - self_copy.powerstate.transport_options["transport_evaluator"] = None - self_copy.powerstate.target_options["target_evaluator_options"]["TypeTarget"] = "powerstate" + self_copy.powerstate.transport_options["evaluator"] = None + self_copy.powerstate.target_options["options"]["TypeTarget"] = "powerstate" _, dictOFs = runModelEvaluator( self_copy, @@ -645,7 +645,7 @@ def runModelEvaluator( # --------------------------------------------------------------------------------------------------- # In certain cases, I want to cold_start the model directly from the PORTALS call instead of powerstate - powerstate.transport_options["transport_evaluator_options"]["cold_start"] = cold_start + powerstate.transport_options["options"]["cold_start"] = cold_start # Evaluate X (DVs) through powerstate.calculate(). This will populate .plasma with the results powerstate.calculate(X, nameRun=name, folder=folder_model, evaluation_number=numPORTALS) diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 6538ece0..6731b6b0 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -8,7 +8,7 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -def surrogate_selection_portals(output, surrogate_options, CGYROrun=False): +def surrogate_selection_portals(output, surrogate_options): print(f'\t- Selecting surrogate options for "{output}" to be run') @@ -250,7 +250,7 @@ def ImpurityGammaTrick(x, surrogate_parameters, output, powerstate): pos = int(output.split("_")[-1]) - if ("GZ" in output) and surrogate_parameters["applyImpurityGammaTrick"]: + if ("GZ" in output) and surrogate_parameters["impurity_trick"]: factor = powerstate.plasma["ni"][: x.shape[0],powerstate.indexes_simulation[pos],powerstate.impurityPosition].unsqueeze(-1) else: @@ -303,7 +303,7 @@ def constructEvaluationProfiles(X, surrogate_parameters, recalculateTargets=Fals # Targets only if needed (for speed, GB doesn't need it) if recalculateTargets: - powerstate.target_options["target_evaluator_options"]["target_evaluator_method"] = "powerstate" # For surrogate evaluation, always powerstate, logically. + powerstate.target_options["options"]["target_evaluator_method"] = "powerstate" # For surrogate evaluation, always powerstate, logically. powerstate.calculateTargets() return powerstate @@ -399,7 +399,7 @@ def calculate_residuals(powerstate, portals_parameters, specific_vars=None): # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added # ------------------------------------------------------------------------- - if portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]: + if portals_parameters["solution"]["turbulent_exchange_as_surrogate"]: QieMWm2_tr_turb = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) else: QieMWm2_tr_turb = torch.zeros(dfT.shape).to(dfT) @@ -452,28 +452,28 @@ def calculate_residuals(powerstate, portals_parameters, specific_vars=None): if var == "Qe": of0, cal0 = ( - of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][0], - cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][0], + of0 * portals_parameters["solution"]["Pseudo_multipliers"][0], + cal0 * portals_parameters["solution"]["Pseudo_multipliers"][0], ) elif var == "Qi": of0, cal0 = ( - of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][1], - cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][1], + of0 * portals_parameters["solution"]["Pseudo_multipliers"][1], + cal0 * portals_parameters["solution"]["Pseudo_multipliers"][1], ) elif var == "Ge": of0, cal0 = ( - of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][2], - cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][2], + of0 * portals_parameters["solution"]["Pseudo_multipliers"][2], + cal0 * portals_parameters["solution"]["Pseudo_multipliers"][2], ) elif var == "GZ": of0, cal0 = ( - of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][3], - cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][3], + of0 * portals_parameters["solution"]["Pseudo_multipliers"][3], + cal0 * portals_parameters["solution"]["Pseudo_multipliers"][3], ) elif var == "MtJm2": of0, cal0 = ( - of0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][4], - cal0 * portals_parameters["main_parameters"]["Pseudo_multipliers"][4], + of0 * portals_parameters["solution"]["Pseudo_multipliers"][4], + cal0 * portals_parameters["solution"]["Pseudo_multipliers"][4], ) of, cal = torch.cat((of, of0), dim=-1), torch.cat((cal, cal0), dim=-1) @@ -532,7 +532,7 @@ def calculate_residuals_distributions(powerstate, portals_parameters): # Volume integrate energy exchange from MW/m^3 to a flux MW/m^2 to be added # ------------------------------------------------------------------------- - if portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]: + if portals_parameters["solution"]["turbulent_exchange_as_surrogate"]: QieMWm2_tr_turb = computeTurbExchangeIndividual(var_dict["Qie_tr_turb"], powerstate) QieMWm2_tr_turb_stds = computeTurbExchangeIndividual(var_dict["Qie_tr_turb_stds"], powerstate) else: diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 28c3da7d..0915cf69 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -153,12 +153,12 @@ def prep_metrics(self, ilast=None): self.portals_parameters = self.opt_fun.mitim_model.optimization_object.portals_parameters # Useful flags - self.predicted_channels = self.portals_parameters["model_parameters"]["predicted_channels"] + self.predicted_channels = self.portals_parameters["solution"]["predicted_channels"] self.runWithImpurity = self.powerstate.impurityPosition if "nZ" in self.predicted_channels else None self.runWithRotation = "w0" in self.predicted_channels - self.forceZeroParticleFlux = self.portals_parameters["model_parameters"]["target_parameters"]["target_evaluator_options"]["forceZeroParticleFlux"] + self.forceZeroParticleFlux = self.portals_parameters["target"]["options"]["forceZeroParticleFlux"] # Profiles and tgyro results print("\t- Reading profiles and tgyros for each evaluation") @@ -237,7 +237,7 @@ def prep_metrics(self, ilast=None): GZ_resR = np.zeros(self.rhos.shape[0]) Mt_resR = np.zeros(self.rhos.shape[0]) cont = 0 - for prof in self.portals_parameters["model_parameters"]["predicted_channels"]: + for prof in self.portals_parameters["solution"]["predicted_channels"]: for ix in range(self.rhos.shape[0]): if prof == "te": Qe_resR[ix] = source[0, cont].abs() @@ -328,14 +328,14 @@ def prep_metrics(self, ilast=None): self.resCheck = ( self.resTeM + self.resTiM + self.resneM + self.resnZM + self.resw0M - ) / len(self.portals_parameters["model_parameters"]["predicted_channels"]) + ) / len(self.portals_parameters["solution"]["predicted_channels"]) # --------------------------------------------------------------------------------------------------------------------- # Jacobian # --------------------------------------------------------------------------------------------------------------------- DeltaQ1 = [] - for i in self.portals_parameters["model_parameters"]["predicted_channels"]: + for i in self.portals_parameters["solution"]["predicted_channels"]: if i == "te": DeltaQ1.append(-self.resTe) if i == "ti": @@ -544,12 +544,12 @@ def extractTGYRO(self, folder=None, cold_start=False, evaluation=0, modified_pro folder, profilesclass_custom=profiles, cold_start=cold_start, forceIfcold_start=True ) - TGLFsettings = self.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["TGLFsettings"] - extraOptionsTGLF = self.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["extraOptionsTGLF"] + TGLFsettings = self.portals_parameters["transport"]["options"]["TGLFsettings"] + extraOptionsTGLF = self.portals_parameters["transport"]["options"]["extraOptionsTGLF"] PredictionSet = [ - int("te" in self.portals_parameters["model_parameters"]["predicted_channels"]), - int("ti" in self.portals_parameters["model_parameters"]["predicted_channels"]), - int("ne" in self.portals_parameters["model_parameters"]["predicted_channels"]), + int("te" in self.portals_parameters["solution"]["predicted_channels"]), + int("ti" in self.portals_parameters["solution"]["predicted_channels"]), + int("ne" in self.portals_parameters["solution"]["predicted_channels"]), ] return tgyro, self.rhos, PredictionSet, TGLFsettings, extraOptionsTGLF @@ -563,12 +563,12 @@ def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=F """ NOTE on radial location extraction: Two possible options for the rho locations to use: - 1. self.portals_parameters["model_parameters"]["predicted_rho"] -> the ones PORTALS sent to TGYRO + 1. self.portals_parameters["solution"]["predicted_rho"] -> 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. #TODO: This should be fixed in the future, we should never send to TGYRO more than 7 decimal places of any variable """ - rhos_considered = self.portals_parameters["model_parameters"]["predicted_rho"] + rhos_considered = self.portals_parameters["solution"]["predicted_rho"] if positions is None: rhos = rhos_considered @@ -593,8 +593,8 @@ def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=F tglf = TGLFtools.TGLF(rhos=rhos) _ = tglf.prep_using_tgyro(folder, cold_start=cold_start, inputgacode=inputgacode) - TGLFsettings = self.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["TGLFsettings"] - extraOptions = self.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["extraOptionsTGLF"] + TGLFsettings = self.portals_parameters["transport"]["options"]["TGLFsettings"] + extraOptions = self.portals_parameters["transport"]["options"]["extraOptionsTGLF"] return tglf, TGLFsettings, extraOptions diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index a1b7a88f..1b8ff337 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -71,34 +71,34 @@ def initializeProblem( profiles = PROFILEStools.gacode_state(initialization_file) # About radial locations - if portals_fun.portals_parameters["model_parameters"]["predicted_roa"] is not None: - roa = portals_fun.portals_parameters["model_parameters"]["predicted_roa"] + if portals_fun.portals_parameters["solution"]["predicted_roa"] is not None: + roa = portals_fun.portals_parameters["solution"]["predicted_roa"] rho = np.interp(roa, profiles.derived["roa"], profiles.profiles["rho(-)"]) print("\t * r/a provided, transforming to rho:") print(f"\t\t r/a = {roa}") print(f"\t\t rho = {rho}") - portals_fun.portals_parameters["model_parameters"]["predicted_rho"] = rho + portals_fun.portals_parameters["solution"]["predicted_rho"] = rho if ( - len(portals_parameters["initialization_parameters"]["removeIons"]) > 0 - or portals_parameters["initialization_parameters"]["remove_fast"] - or portals_parameters["initialization_parameters"]["quasineutrality"] - or portals_parameters["initialization_parameters"]["enforce_same_aLn"] - or portals_parameters["initialization_parameters"]["recalculate_ptot"] + len(portals_parameters["initialization"]["removeIons"]) > 0 + or portals_parameters["initialization"]["remove_fast"] + or portals_parameters["initialization"]["quasineutrality"] + or portals_parameters["initialization"]["enforce_same_aLn"] + or portals_parameters["initialization"]["recalculate_ptot"] ): - profiles.correct(options=portals_parameters["initialization_parameters"]) + profiles.correct(options=portals_parameters["initialization"]) - if portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"] is not None: - position_of_impurity = MITIMstate.impurity_location(profiles, portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"]) + if portals_fun.portals_parameters["solution"]["trace_impurity"] is not None: + position_of_impurity = MITIMstate.impurity_location(profiles, portals_fun.portals_parameters["solution"]["trace_impurity"]) else: position_of_impurity = 1 - if portals_fun.portals_parameters["main_parameters"]["UseOriginalImpurityConcentrationAsWeight"] is not None and portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"] is not None: + if portals_fun.portals_parameters["solution"]["fZ0_as_weight"] is not None and portals_fun.portals_parameters["solution"]["trace_impurity"] is not None: f0 = profiles.Species[position_of_impurity]["n0"] / profiles.profiles['ne(10^19/m^3)'][0] - portals_fun.portals_parameters["main_parameters"]["fImp_orig"] = f0/portals_fun.portals_parameters["main_parameters"]["UseOriginalImpurityConcentrationAsWeight"] - print(f'\t- Ion {portals_fun.portals_parameters["model_parameters"]["ImpurityOfInterest"]} has original central concentration of {f0:.2e}, using its inverse multiplied by {portals_fun.portals_parameters["main_parameters"]["UseOriginalImpurityConcentrationAsWeight"]} as scaling factor of GZ -> {portals_fun.portals_parameters["main_parameters"]["fImp_orig"]:.2e}',typeMsg="i") + portals_fun.portals_parameters["solution"]["fImp_orig"] = f0/portals_fun.portals_parameters["solution"]["fZ0_as_weight"] + print(f'\t- Ion {portals_fun.portals_parameters["solution"]["trace_impurity"]} has original central concentration of {f0:.2e}, using its inverse multiplied by {portals_fun.portals_parameters["solution"]["fZ0_as_weight"]} as scaling factor of GZ -> {portals_fun.portals_parameters["solution"]["fImp_orig"]:.2e}',typeMsg="i") else: - portals_fun.portals_parameters["main_parameters"]["fImp_orig"] = 1.0 + portals_fun.portals_parameters["solution"]["fImp_orig"] = 1.0 # Check if I will be able to calculate radiation speciesNotFound = [] @@ -110,7 +110,7 @@ def initializeProblem( # Print warning or question to be careful! if len(speciesNotFound) > 0: - if portals_fun.portals_parameters["model_parameters"]["target_parameters"]["TypeTarget"] == 3: + if portals_fun.portals_parameters["target"]["TypeTarget"] == 3: answerYN = print(f"\t- Species {speciesNotFound} not found in radiation database, radiation will be zero in PORTALS... is this ok for your predictions?",typeMsg="q" if checkForSpecies else "w") if checkForSpecies and (not answerYN): @@ -122,7 +122,7 @@ def initializeProblem( # Prepare and defaults - xCPs = torch.from_numpy(np.array(portals_fun.portals_parameters["model_parameters"]["predicted_rho"])).to(dfT) + xCPs = torch.from_numpy(np.array(portals_fun.portals_parameters["solution"]["predicted_rho"])).to(dfT) """ *************************************************************************************************** @@ -130,17 +130,17 @@ def initializeProblem( *************************************************************************************************** """ - transport_parameters = portals_fun.portals_parameters["model_parameters"]["transport_parameters"] + transport_parameters = portals_fun.portals_parameters["transport"] # Add folder and cold_start to the simulation options transport_options = transport_parameters | {"folder": portals_fun.folder, "cold_start": False} - target_options = portals_fun.portals_parameters["model_parameters"]["target_parameters"] + target_options = portals_fun.portals_parameters["target"] portals_fun.powerstate = STATEtools.powerstate( profiles, evolution_options={ - "ProfilePredicted": portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "ProfilePredicted": portals_fun.portals_parameters["solution"]["predicted_channels"], "rhoPredicted": xCPs, "impurityPosition": position_of_impurity, }, @@ -157,13 +157,13 @@ def initializeProblem( # Store parameterization in dictCPs_base (to define later the relative variations) and modify profiles class with parameterized profiles dictCPs_base = {} - for name in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]: + for name in portals_fun.portals_parameters["solution"]["predicted_channels"]: dictCPs_base[name] = portals_fun.powerstate.update_var(name, var=None)[0, :] # Maybe it was provided from earlier run if start_from_folder is not None: dictCPs_base = grabPrevious(start_from_folder, dictCPs_base) - for name in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]: + for name in portals_fun.portals_parameters["solution"]["predicted_channels"]: _ = portals_fun.powerstate.update_var( name, var=dictCPs_base[name].unsqueeze(0) ) @@ -171,7 +171,7 @@ def initializeProblem( # Write this updated profiles class (with parameterized profiles) _ = portals_fun.powerstate.from_powerstate( write_input_gacode=FolderInitialization / "input.gacode", - postprocess_input_gacode=portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["applyCorrections"], + postprocess_input_gacode=portals_fun.portals_parameters["transport"]["applyCorrections"], ) # Original complete targets @@ -186,25 +186,25 @@ def initializeProblem( powerstate_extra = STATEtools.powerstate( define_ranges_from_profiles, evolution_options={ - "ProfilePredicted": portals_fun.portals_parameters["model_parameters"]["predicted_channels"], + "ProfilePredicted": portals_fun.portals_parameters["solution"]["predicted_channels"], "rhoPredicted": xCPs, "impurityPosition": position_of_impurity, - "fineTargetsResolution": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["fineTargetsResolution"], + "fineTargetsResolution": portals_fun.portals_parameters["target"]["fineTargetsResolution"], }, target_options={ - "target_evaluator": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["target_evaluator"], - "target_evaluator_options": { - "TypeTarget": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["TypeTarget"], - "target_evaluator_method": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["target_evaluator_method"], - "forceZeroParticleFlux": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["target_evaluator_options"]["forceZeroParticleFlux"], - "percent_error": portals_fun.portals_parameters["model_parameters"]["target_parameters"]["percent_error"] + "evaluator": portals_fun.portals_parameters["target"]["evaluator"], + "options": { + "TypeTarget": portals_fun.portals_parameters["target"]["TypeTarget"], + "target_evaluator_method": portals_fun.portals_parameters["target"]["target_evaluator_method"], + "forceZeroParticleFlux": portals_fun.portals_parameters["target"]["options"]["forceZeroParticleFlux"], + "percent_error": portals_fun.portals_parameters["target"]["percent_error"] }, }, tensor_options = tensor_options ) dictCPs_base_extra = {} - for name in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]: + for name in portals_fun.portals_parameters["solution"]["predicted_channels"]: dictCPs_base_extra[name] = powerstate_extra.update_var(name, var=None)[0, :] dictCPs_base = dictCPs_base_extra @@ -260,7 +260,7 @@ def initializeProblem( elif ikey == "w0": var = "Mt" - for i in range(len(portals_fun.portals_parameters["model_parameters"]["predicted_rho"])): + for i in range(len(portals_fun.portals_parameters["solution"]["predicted_rho"])): ofs.append(f"{var}_tr_turb_{i+1}") ofs.append(f"{var}_tr_neoc_{i+1}") @@ -268,13 +268,13 @@ def initializeProblem( name_objectives.append(f"{var}Res_{i+1}") - if portals_fun.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]: - for i in range(len(portals_fun.portals_parameters["model_parameters"]["predicted_rho"])): + if portals_fun.portals_parameters["solution"]["turbulent_exchange_as_surrogate"]: + for i in range(len(portals_fun.portals_parameters["solution"]["predicted_rho"])): ofs.append(f"Qie_tr_turb_{i+1}") name_transformed_ofs = [] for of in ofs: - if ("GZ" in of) and (portals_fun.portals_parameters["main_parameters"]["applyImpurityGammaTrick"]): + if ("GZ" in of) and (portals_fun.portals_parameters["solution"]["impurity_trick"]): lab = f"{of} (GB MOD)" else: lab = f"{of} (GB)" @@ -303,14 +303,14 @@ def initializeProblem( # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Variables = {} - for ikey in portals_fun.portals_parameters["main_parameters"]["portals_transformation_variables"]: + for ikey in portals_fun.portals_parameters["solution"]["portals_transformation_variables"]: Variables[ikey] = prepportals_transformation_variables(portals_fun, ikey) portals_fun.surrogate_parameters = { "transformationInputs": PORTALStools.input_transform_portals, "transformationOutputs": PORTALStools.output_transform_portals, "powerstate": portals_fun.powerstate, - "applyImpurityGammaTrick": portals_fun.portals_parameters["main_parameters"]["applyImpurityGammaTrick"], + "impurity_trick": portals_fun.portals_parameters["solution"]["impurity_trick"], "surrogate_transformation_variables_alltimes": Variables, "surrogate_transformation_variables_lasttime": copy.deepcopy(Variables[list(Variables.keys())[-1]]), "parameters_combined": {}, @@ -318,9 +318,9 @@ def initializeProblem( def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValues=False): allOuts = portals_fun.optimization_options["problem_options"]["ofs"] - portals_transformation_variables = portals_fun.portals_parameters["main_parameters"]["portals_transformation_variables"][ikey] - portals_transformation_variables_trace = portals_fun.portals_parameters["main_parameters"]["portals_transformation_variables_trace"][ikey] - additional_params_in_surrogate = portals_fun.portals_parameters["main_parameters"]["additional_params_in_surrogate"] + portals_transformation_variables = portals_fun.portals_parameters["solution"]["portals_transformation_variables"][ikey] + portals_transformation_variables_trace = portals_fun.portals_parameters["solution"]["portals_transformation_variables_trace"][ikey] + additional_params_in_surrogate = portals_fun.portals_parameters["solution"]["additional_params_in_surrogate"] Variables = {} for output in allOuts: @@ -353,17 +353,17 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue isAbsValFixed = False Variations = { - "aLte": "te" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], - "aLti": "ti" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], - "aLne": "ne" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], - "aLw0": "w0" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], - "te": ("te" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) + "aLte": "te" in portals_fun.portals_parameters["solution"]["predicted_channels"], + "aLti": "ti" in portals_fun.portals_parameters["solution"]["predicted_channels"], + "aLne": "ne" in portals_fun.portals_parameters["solution"]["predicted_channels"], + "aLw0": "w0" in portals_fun.portals_parameters["solution"]["predicted_channels"], + "te": ("te" in portals_fun.portals_parameters["solution"]["predicted_channels"]) and (not isAbsValFixed), - "ti": ("ti" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) + "ti": ("ti" in portals_fun.portals_parameters["solution"]["predicted_channels"]) and (not isAbsValFixed), - "ne": ("ne" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) + "ne": ("ne" in portals_fun.portals_parameters["solution"]["predicted_channels"]) and (not isAbsValFixed), - "w0": ("w0" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) + "w0": ("w0" in portals_fun.portals_parameters["solution"]["predicted_channels"]) and (not isAbsValFixed), } @@ -388,20 +388,20 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue isAbsValFixed = False Variations = { - "aLte": "te" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], - "aLti": "ti" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], - "aLne": "ne" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], - "aLw0": "w0" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], - "aLnZ": "nZ" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"], - "te": ("te" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) + "aLte": "te" in portals_fun.portals_parameters["solution"]["predicted_channels"], + "aLti": "ti" in portals_fun.portals_parameters["solution"]["predicted_channels"], + "aLne": "ne" in portals_fun.portals_parameters["solution"]["predicted_channels"], + "aLw0": "w0" in portals_fun.portals_parameters["solution"]["predicted_channels"], + "aLnZ": "nZ" in portals_fun.portals_parameters["solution"]["predicted_channels"], + "te": ("te" in portals_fun.portals_parameters["solution"]["predicted_channels"]) and (not isAbsValFixed), - "ti": ("ti" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) + "ti": ("ti" in portals_fun.portals_parameters["solution"]["predicted_channels"]) and (not isAbsValFixed), - "ne": ("ne" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) + "ne": ("ne" in portals_fun.portals_parameters["solution"]["predicted_channels"]) and (not isAbsValFixed), - "w0": ("w0" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) + "w0": ("w0" in portals_fun.portals_parameters["solution"]["predicted_channels"]) and (not isAbsValFixed), - "nZ": ("nZ" in portals_fun.portals_parameters["model_parameters"]["predicted_channels"]) + "nZ": ("nZ" in portals_fun.portals_parameters["solution"]["predicted_channels"]) and (not isAbsValFixed), } diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index 05f47e40..8564c0c4 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -140,7 +140,7 @@ def flux_match_surrogate( # Define transport calculation function as a surrogate model transport_options['transport_evaluator'] = transport_analytic.surrogate - transport_options["transport_evaluator_options"] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} + transport_options["options"] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} # Create powerstate with the same options as the original portals but with the new profiles powerstate = STATEtools.powerstate( diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index 37956dce..1dd74f5c 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -1207,7 +1207,7 @@ def PORTALSanalyzer_plotExpected( rho = p.profiles["rho(-)"] roa = p.derived["roa"] - rhoVals = self.portals_parameters["model_parameters"]["predicted_rho"] + rhoVals = self.portals_parameters["solution"]["predicted_rho"] roaVals = np.interp(rhoVals, rho, roa) lastX = roaVals[-1] @@ -1395,7 +1395,7 @@ def PORTALSanalyzer_plotExpected( rho = self.profiles_next_new.profiles["rho(-)"] - rhoVals = self.portals_parameters["model_parameters"]["predicted_rho"] + rhoVals = self.portals_parameters["solution"]["predicted_rho"] roaVals = np.interp(rhoVals, rho, roa) p0 = self.powerstates[plotPoints[0]].profiles @@ -1892,10 +1892,10 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): axs4, color=colors[i], label=label, - lastRho=self.portals_parameters["model_parameters"]["predicted_rho"][-1], + lastRho=self.portals_parameters["solution"]["predicted_rho"][-1], alpha=alpha, useRoa=True, - RhoLocationsPlot=self.portals_parameters["model_parameters"]["predicted_rho"], + RhoLocationsPlot=self.portals_parameters["solution"]["predicted_rho"], plotImpurity=self.runWithImpurity, plotRotation=self.runWithRotation, autoscale=i == 3, @@ -1959,7 +1959,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): p.plot_gradients( axsR, color="b", - lastRho=self.portals_parameters["model_parameters"]["predicted_rho"][-1], + lastRho=self.portals_parameters["solution"]["predicted_rho"][-1], ms=ms, lw=1.0, label="Initial (#0)", @@ -1976,7 +1976,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): p.plot_gradients( axsR, color="r", - lastRho=self.portals_parameters["model_parameters"]["predicted_rho"][-1], + lastRho=self.portals_parameters["solution"]["predicted_rho"][-1], ms=ms, lw=0.3, ls="-o" if self.opt_fun.mitim_model.avoidPoints is not None else "-.o", @@ -1988,7 +1988,7 @@ def PORTALSanalyzer_plotRanges(self, fig=None): p.plot_gradients( axsR, color="g", - lastRho=self.portals_parameters["model_parameters"]["predicted_rho"][-1], + lastRho=self.portals_parameters["solution"]["predicted_rho"][-1], ms=ms, lw=1.0, label=f"Best (#{self.opt_fun.res.best_absolute_index})", @@ -2011,10 +2011,10 @@ def PORTALSanalyzer_plotModelComparison( if (fig is None) and (axs is None): plt.ion() - fig = plt.figure(figsize=(15, 6 if len(self.predicted_channels)+int(self.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]) < 4 else 10)) + fig = plt.figure(figsize=(15, 6 if len(self.predicted_channels)+int(self.portals_parameters["solution"]["turbulent_exchange_as_surrogate"]) < 4 else 10)) if axs is None: - if len(self.predicted_channels)+int(self.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]) < 4: + if len(self.predicted_channels)+int(self.portals_parameters["solution"]["turbulent_exchange_as_surrogate"]) < 4: axs = fig.subplots(ncols=3) else: axs = fig.subplots(ncols=3, nrows=2) @@ -2201,7 +2201,7 @@ def PORTALSanalyzer_plotModelComparison( cont += 1 - if self.portals_parameters["main_parameters"]["turbulent_exchange_as_surrogate"]: + if self.portals_parameters["solution"]["turbulent_exchange_as_surrogate"]: if UseTGLFfull_x is not None: raise Exception("Turbulent exchange plot not implemented yet") # Sexch @@ -2419,8 +2419,8 @@ def varToReal(y, mitim_model): cont = 0 Qe, Qi, Ge, GZ, Mt = [], [], [], [], [] Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = [], [], [], [], [] - for prof in mitim_model.optimization_object.portals_parameters["model_parameters"]["predicted_channels"]: - for rad in mitim_model.optimization_object.portals_parameters["model_parameters"]["predicted_rho"]: + for prof in mitim_model.optimization_object.portals_parameters["solution"]["predicted_channels"]: + for rad in mitim_model.optimization_object.portals_parameters["solution"]["predicted_rho"]: if prof == "te": Qe.append(of[0, cont]) Qe_tar.append(cal[0, cont]) @@ -2493,7 +2493,7 @@ def plotVars( .plasma["roa"][0, 1:] .cpu() .cpu().numpy() - ) # mitim_model.optimization_object.portals_parameters["model_parameters"]['predicted_rho'] + ) try: Qe, Qi, Ge, GZ, Mt, Qe_tar, Qi_tar, Ge_tar, GZ_tar, Mt_tar = varToReal( @@ -3254,7 +3254,7 @@ def plotFluxComparison( def produceInfoRanges( self_complete, bounds, axsR, label="", color="k", lw=0.2, alpha=0.05 ): - rhos = np.append([0], self_complete.portals_parameters["model_parameters"]["predicted_rho"]) + rhos = np.append([0], self_complete.portals_parameters["solution"]["predicted_rho"]) aLTe, aLTi, aLne, aLnZ, aLw0 = ( np.zeros((len(rhos), 2)), np.zeros((len(rhos), 2)), @@ -3275,19 +3275,19 @@ def produceInfoRanges( if f"aLw0_{i+1}" in bounds: aLw0[i + 1, :] = bounds[f"aLw0_{i+1}"] - X = torch.zeros(((len(rhos) - 1) * len(self_complete.portals_parameters["model_parameters"]["predicted_channels"]), 2)) + X = torch.zeros(((len(rhos) - 1) * len(self_complete.portals_parameters["solution"]["predicted_channels"]), 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.portals_parameters["model_parameters"]["predicted_channels"]: + if "ne" in self_complete.portals_parameters["solution"]["predicted_channels"]: X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLne[1:, :]) cont += 1 - if "nZ" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: + if "nZ" in self_complete.portals_parameters["solution"]["predicted_channels"]: X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLnZ[1:, :]) cont += 1 - if "w0" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: + if "w0" in self_complete.portals_parameters["solution"]["predicted_channels"]: X[(2 + cont) * l : (3 + cont) * l, :] = torch.from_numpy(aLw0[1:, :]) cont += 1 @@ -3338,7 +3338,7 @@ def produceInfoRanges( ) cont = 0 - if "ne" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: + if "ne" in self_complete.portals_parameters["solution"]["predicted_channels"]: GRAPHICStools.fillGraph( axsR[3 + cont + 1], powerstate.plasma["rho"][0], @@ -3361,7 +3361,7 @@ def produceInfoRanges( ) cont += 2 - if "nZ" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: + if "nZ" in self_complete.portals_parameters["solution"]["predicted_channels"]: GRAPHICStools.fillGraph( axsR[3 + cont + 1], powerstate.plasma["rho"][0], @@ -3384,7 +3384,7 @@ def produceInfoRanges( ) cont += 2 - if "w0" in self_complete.portals_parameters["model_parameters"]["predicted_channels"]: + if "w0" in self_complete.portals_parameters["solution"]["predicted_channels"]: GRAPHICStools.fillGraph( axsR[3 + cont + 1], powerstate.plasma["rho"][0], diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index df83ddb6..b0bdfa08 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -47,14 +47,14 @@ def __init__( if transport_options is None: transport_options = { - "transport_evaluator": None, - "transport_evaluator_options": {} + "evaluator": None, + "options": {} } if target_options is None: target_options = { - "target_evaluator": targets_analytic.analytical_model, - "target_evaluator_options": { + "evaluator": targets_analytic.analytical_model, + "options": { "TypeTarget": 3, "target_evaluator_method": "powerstate", "forceZeroParticleFlux": False @@ -244,12 +244,12 @@ def combine_states(self, states, includeTransport=True): if includeTransport: for key in ["chi_e", "chi_i"]: - self.transport_options["transport_evaluator_options"][key] = torch.cat( + self.transport_options["options"][key] = torch.cat( ( - self.transport_options["transport_evaluator_options"][key], - state.transport_options["transport_evaluator_options"][key], + self.transport_options["options"][key], + state.transport_options["options"][key], ) - ).to(self.transport_options["transport_evaluator_options"][key]) + ).to(self.transport_options["options"][key]) def copy_state(self): @@ -297,7 +297,7 @@ def calculate( self.calculateProfileFunctions() # 3. Sources and sinks (populates components and Pe,Pi,...) - relative_error_assumed = self.target_options["target_evaluator_options"]["percent_error"] + relative_error_assumed = self.target_options["options"]["percent_error"] self.calculateTargets(relative_error_assumed=relative_error_assumed) # Calculate targets based on powerstate functions (it may be overwritten in next step, if chosen) # 4. Turbulent and neoclassical transport (populates components and Pe_tr,Pi_tr,...) @@ -366,7 +366,7 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): if folder_main is not None: folder = IOtools.expandPath(folder_main) / f"{namingConvention}_{cont}" - if issubclass(self.transport_options["transport_evaluator"], TRANSPORTtools.power_transport): + if issubclass(self.transport_options["evaluator"], TRANSPORTtools.power_transport): (folder / "transport_simulation_folder").mkdir(parents=True, exist_ok=True) # *************************************************************************************************************** @@ -380,7 +380,7 @@ def evaluator(X, y_history=None, x_history=None, metric_history=None): # Save state so that I can check initializations if folder_main is not None: - if issubclass(self.transport_options["transport_evaluator"], TRANSPORTtools.power_transport): + if issubclass(self.transport_options["evaluator"], TRANSPORTtools.power_transport): self.save(folder / "powerstate.pkl") shutil.copy2(folder_run / "input.gacode", folder) @@ -687,10 +687,10 @@ def calculateTargets(self, relative_error_assumed=1.0): """ # If no targets evaluator is given or the targets will come from TGYRO, assume them as zero - if (self.target_options["target_evaluator"] is None) or (self.target_options["target_evaluator_options"]["target_evaluator_method"] == "tgyro"): + if (self.target_options["evaluator"] is None) or (self.target_options["options"]["target_evaluator_method"] == "tgyro"): targets = TARGETStools.power_targets(self) else: - targets = self.target_options["target_evaluator"](self) + targets = self.target_options["evaluator"](self) # [Optional] Calculate local targets and integrals on a fine grid if self.fineTargetsResolution is not None: @@ -709,7 +709,7 @@ def calculateTargets(self, relative_error_assumed=1.0): # Merge targets, calculate errors and normalize targets.postprocessing( relative_error_assumed=relative_error_assumed, - forceZeroParticleFlux=self.target_options["target_evaluator_options"]["forceZeroParticleFlux"]) + forceZeroParticleFlux=self.target_options["options"]["forceZeroParticleFlux"]) def calculateTransport( self, nameRun="test", folder="~/scratch/", evaluation_number=0): @@ -719,10 +719,10 @@ def calculateTransport( folder = IOtools.expandPath(folder) # Select transport evaluator - if self.transport_options["transport_evaluator"] is None: + if self.transport_options["evaluator"] is None: transport = TRANSPORTtools.power_transport( self, name=nameRun, folder=folder, evaluation_number=evaluation_number ) else: - transport = self.transport_options["transport_evaluator"]( self, name=nameRun, folder=folder, evaluation_number=evaluation_number ) + transport = self.transport_options["evaluator"]( self, name=nameRun, folder=folder, evaluation_number=evaluation_number ) # Produce profile object (for certain transport evaluators, this is necessary) transport.produce_profiles() diff --git a/src/mitim_modules/powertorch/physics_models/targets_analytic.py b/src/mitim_modules/powertorch/physics_models/targets_analytic.py index 418065db..d2534282 100644 --- a/src/mitim_modules/powertorch/physics_models/targets_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/targets_analytic.py @@ -28,10 +28,10 @@ def __init__(self,powerstate, **kwargs): def evaluate(self): - if self.powerstate.target_options["target_evaluator_options"]["TypeTarget"] >= 2: + if self.powerstate.target_options["options"]["TypeTarget"] >= 2: self._evaluate_energy_exchange() - if self.powerstate.target_options["target_evaluator_options"]["TypeTarget"] == 3: + if self.powerstate.target_options["options"]["TypeTarget"] == 3: self._evaluate_alpha_heating() self._evaluate_radiation() diff --git a/src/mitim_modules/powertorch/physics_models/transport_analytic.py b/src/mitim_modules/powertorch/physics_models/transport_analytic.py index 77288306..cedce490 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/transport_analytic.py @@ -13,8 +13,8 @@ def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) # Ensure that the provided diffusivities include the zero location - self.chi_e = self.powerstate.transport_options["transport_evaluator_options"]["chi_e"] - self.chi_i = self.powerstate.transport_options["transport_evaluator_options"]["chi_i"] + self.chi_e = self.powerstate.transport_options["options"]["chi_e"] + self.chi_i = self.powerstate.transport_options["options"]["chi_i"] if self.chi_e.shape[0] < self.powerstate.plasma['rho'].shape[-1]: self.chi_e = torch.cat((torch.zeros(1), self.chi_e)) @@ -73,7 +73,7 @@ def evaluate(self): for prof in self.powerstate.predicted_channels: X = torch.cat((X,self.powerstate.plasma['aL'+prof][:,1:]),axis=1) - _, Q, _, _ = self.powerstate.transport_options["transport_evaluator_options"]["flux_fun"](X) + _, Q, _, _ = self.powerstate.transport_options["options"]["flux_fun"](X) numeach = self.powerstate.plasma["rho"].shape[1] - 1 diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index c0690b14..de234a55 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -13,7 +13,7 @@ def __init__(self, powerstate, **kwargs): # Do not hook here def evaluate_turbulence(self): - transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + transport_evaluator_options = self.powerstate.transport_options["options"] cold_start = self.powerstate.transport_options["cold_start"] # Run base TGLF always, to keep track of discrepancies! -------------------------------------- diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index 6fbcc80c..c89604b2 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -26,7 +26,7 @@ def evaluate_turbulence(self): # Have it separate such that I can call it from the CGYRO class but without the decorator def _evaluate_tglf(self): - transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + transport_evaluator_options = self.powerstate.transport_options["options"] cold_start = self.powerstate.transport_options["cold_start"] @@ -167,7 +167,7 @@ def _evaluate_tglf(self): @IOtools.hook_method(after=partial(TRANSPORTtools.write_json, file_name = 'fluxes_neoc.json', suffix= 'neoc')) def evaluate_neoclassical(self): - transport_evaluator_options = self.powerstate.transport_options["transport_evaluator_options"] + transport_evaluator_options = self.powerstate.transport_options["options"] # ------------------------------------------------------------------------------------------------------------------------ # Grab options from powerstate @@ -235,7 +235,7 @@ def evaluate_neoclassical(self): def _profiles_to_store(self): - if "folder" in self.powerstate.transport_options["transport_evaluator_options"]: + if "folder" in self.powerstate.transport_options["options"]: whereFolder = IOtools.expandPath(self.powerstate.transport_options["folder"] / "Outputs" / "portals_profiles") if not whereFolder.exists(): IOtools.askNewFolder(whereFolder) diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 010c7934..68aba57f 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -33,17 +33,17 @@ def calculator( "fineTargetsResolution": fineTargetsResolution, }, target_options={ - "target_evaluator": targets_analytic.analytical_model, - "target_evaluator_options": { + "evaluator": targets_analytic.analytical_model, + "options": { "TypeTarget": TypeTarget, "target_evaluator_method": "tgyro"}, }, transport_options={ - "transport_evaluator": transport_tglfneo.tglfneo_model, - "transport_evaluator_options": { + "evaluator": transport_tglfneo.tglfneo_model, + "options": { "cold_start": cold_start, "portals_parameters": { - "main_parameters": { + "solution": { "launchSlurm": True, "Qi_includes_fast": False, }, @@ -79,14 +79,14 @@ def calculator( "fineTargetsResolution": fineTargetsResolution, }, target_options={ - "target_evaluator": targets_analytic.analytical_model, - "target_evaluator_options": { + "evaluator": targets_analytic.analytical_model, + "options": { "TypeTarget": TypeTarget, "target_evaluator_method": "powerstate"}, }, transport_options={ - "transport_evaluator": None, - "transport_evaluator_options": {} + "evaluator": None, + "options": {} }, ) diff --git a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py index ad798de8..875e1299 100644 --- a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py +++ b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py @@ -30,7 +30,7 @@ # STATE s = STATEtools.powerstate(t.profiles, evolution_options={"rhoPredicted": t.rho[0,1:]}) s.calculateProfileFunctions() -# s.target_options['target_evaluator_options']['TypeTarget'] = 1 +# s.target_options['options']['TypeTarget'] = 1 s.calculateTargets() # diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index dd2e7641..649ae0ed 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -71,11 +71,11 @@ def flux_integrate(self): qe = self.powerstate.plasma["te"]*0.0 qi = self.powerstate.plasma["te"]*0.0 - if self.powerstate.target_options['target_evaluator_options']['TypeTarget'] >= 2: + if self.powerstate.target_options['options']['TypeTarget'] >= 2: qe += -self.powerstate.plasma["qie"] qi += self.powerstate.plasma["qie"] - if self.powerstate.target_options['target_evaluator_options']['TypeTarget'] == 3: + if self.powerstate.target_options['options']['TypeTarget'] == 3: qe += self.powerstate.plasma["qfuse"] - self.powerstate.plasma["qrad"] qi += self.powerstate.plasma["qfusi"] @@ -89,11 +89,11 @@ def coarse_grid(self): # ************************************************************************************************** # Interpolate results from fine to coarse (i.e. whole point is that it is better than integrate interpolated values) - if self.powerstate.target_options['target_evaluator_options']['TypeTarget'] >= 2: + if self.powerstate.target_options['options']['TypeTarget'] >= 2: for i in ["qie"]: self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] - if self.powerstate.target_options['target_evaluator_options']['TypeTarget'] == 3: + if self.powerstate.target_options['options']['TypeTarget'] == 3: for i in [ "qfuse", "qfusi", diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index c9dcdada..edd487a7 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -120,12 +120,12 @@ def gacode_to_powerstate(self, rho_vec=None): quantitites["GZ_fixedtargets"] = input_gacode.derived["ge_10E20"] * 0.0 quantitites["MtJm2_fixedtargets"] = input_gacode.derived["mt_Jmiller"] - if self.target_options["target_evaluator_options"]["TypeTarget"] < 3: + if self.target_options["options"]["TypeTarget"] < 3: # Fusion and radiation fixed if 1,2 quantitites["QeMWm2_fixedtargets"] += input_gacode.derived["qe_fus_MW"] - input_gacode.derived["qrad_MW"] quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qi_fus_MW"] - if self.target_options["target_evaluator_options"]["TypeTarget"] < 2: + if self.target_options["options"]["TypeTarget"] < 2: # Exchange fixed if 1 quantitites["QeMWm2_fixedtargets"] -= input_gacode.derived["qe_exc_MW"] quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qe_exc_MW"] @@ -365,26 +365,26 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): profiles, evolution_options={"rhoPredicted": rhoy}, target_options={ - "target_evaluator": targets_analytic.analytical_model, - "target_evaluator_options": { - "TypeTarget": self.target_options["target_evaluator_options"]["TypeTarget"], # Important to keep the same as in the original + "evaluator": targets_analytic.analytical_model, + "options": { + "TypeTarget": self.target_options["options"]["TypeTarget"], # Important to keep the same as in the original "target_evaluator_method": "powerstate", - "forceZeroParticleFlux": self.target_options["target_evaluator_options"]["forceZeroParticleFlux"], - "percent_error": self.target_options["target_evaluator_options"]["percent_error"] + "forceZeroParticleFlux": self.target_options["options"]["forceZeroParticleFlux"], + "percent_error": self.target_options["options"]["percent_error"] } }, increase_profile_resol = False ) state_temp.calculateProfileFunctions() - state_temp.target_options["target_evaluator_options"]["target_evaluator_method"] = "powerstate" + state_temp.target_options["options"]["target_evaluator_method"] = "powerstate" state_temp.calculateTargets() # ------------------------------------------------------------------------------------------ conversions = {} - if self.target_options["target_evaluator_options"]["TypeTarget"] > 1: + if self.target_options["options"]["TypeTarget"] > 1: conversions['qie'] = "qei(MW/m^3)" - if self.target_options["target_evaluator_options"]["TypeTarget"] > 2: + if self.target_options["options"]["TypeTarget"] > 2: conversions['qrad_bremms'] = "qbrem(MW/m^3)" conversions['qrad_sync'] = "qsync(MW/m^3)" conversions['qrad_line'] = "qline(MW/m^3)" diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 03187e25..335e741a 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -205,7 +205,7 @@ def evaluate(self): def _postprocess(self): - OriginalFimp = self.powerstate.transport_options["transport_evaluator_options"].get("OriginalFimp", 1.0) + OriginalFimp = self.powerstate.transport_options["options"].get("OriginalFimp", 1.0) # ------------------------------------------------------------------------------------------------------------------------ # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index 9c580f80..22baa746 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -29,16 +29,16 @@ portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 portals_fun.optimization_options["initialization_options"]["initial_training"] = 2 -portals_fun.portals_parameters["main_parameters"]['turbulent_exchange_as_surrogate'] = True +portals_fun.portals_parameters["solution"]['turbulent_exchange_as_surrogate'] = True -portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True -portals_fun.portals_parameters["initialization_parameters"]["quasineutrality"] = True -portals_fun.portals_parameters["initialization_parameters"]["enforce_same_aLn"] = True +portals_fun.portals_parameters["initialization"]["remove_fast"] = True +portals_fun.portals_parameters["initialization"]["quasineutrality"] = True +portals_fun.portals_parameters["initialization"]["enforce_same_aLn"] = True -portals_fun.portals_parameters["model_parameters"]["predicted_rho"] = [0.25, 0.45, 0.65, 0.85] -portals_fun.portals_parameters["model_parameters"]['predicted_channels'] = ["te", "ti", "ne", "nZ", 'w0'] -portals_fun.portals_parameters["model_parameters"]['ImpurityOfInterest'] = 'N' -portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"]["tglf"]["run"]["code_settings"] = 2 +portals_fun.portals_parameters["solution"]["predicted_rho"] = [0.25, 0.45, 0.65, 0.85] +portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti", "ne", "nZ", 'w0'] +portals_fun.portals_parameters["solution"]["trace_impurity"] = 'N' +portals_fun.portals_parameters["transport"]["options"]["tglf"]["run"]["code_settings"] = 2 # Prepare run portals_fun.prep(inputgacode) diff --git a/tests/POWERTORCH_workflow.py b/tests/POWERTORCH_workflow.py index 6c29c017..ac093fcf 100644 --- a/tests/POWERTORCH_workflow.py +++ b/tests/POWERTORCH_workflow.py @@ -15,7 +15,7 @@ 'rhoPredicted': rho }, transport_options = { 'transport_evaluator': transport_analytic.diffusion_model, - "transport_evaluator_options": { + "options": { 'chi_e': torch.ones(rho.shape[0]).to(rho)*0.8, 'chi_i': torch.ones(rho.shape[0]).to(rho)*1.2 } diff --git a/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index 4de41bb3..fb190482 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -14,24 +14,24 @@ portals_fun = PORTALSmain.portals(folder) # Radial locations (predicted_rho or predicted_roa [last one preceeds]) -portals_fun.portals_parameters["model_parameters"]["predicted_rho"] = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.85] +portals_fun.portals_parameters["solution"]["predicted_rho"] = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.85] # Profiles to predict -portals_fun.portals_parameters["model_parameters"]["predicted_channels"] = ["te", "ti", "ne"] +portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti", "ne"] # Codes to use from mitim_modules.powertorch.physics_models.transport_tgyro import tgyro_model -portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator"] = tgyro_model +portals_fun.portals_parameters["transport"]["evaluator"] = tgyro_model # TGLF specifications -portals_fun.portals_parameters["model_parameters"]["transport_parameters"]["transport_evaluator_options"] = { +portals_fun.portals_parameters["transport"]["options"] = { "TGLFsettings": 6, # Check out templates/input.tglf.models.json for more options "extraOptionsTGLF": {"USE_BPER": False} # Turn off BPER } # Plasma preparation: remove fast species, adjust quasineutrality -portals_fun.portals_parameters["initialization_parameters"]["remove_fast"] = True -portals_fun.portals_parameters["initialization_parameters"]["quasineutrality"] = True +portals_fun.portals_parameters["initialization"]["remove_fast"] = True +portals_fun.portals_parameters["initialization"]["quasineutrality"] = True # Stopping criterion 1: 100x improvement in residual portals_fun.optimization_options['convergence_options']['stopping_criteria_parameters']["maximum_value"] = 1e-2 From 993a6c3bdbefc4cb271840efa3260f551f9712ce Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 17:19:31 -0400 Subject: [PATCH 228/385] Remove deprecated initialization to PORTALS... always do it outside of portals, modify the mitim_state --- src/mitim_modules/portals/PORTALSmain.py | 32 ++----- .../portals/utils/PORTALSinit.py | 31 ++----- .../portals/utils/PORTALSoptimization.py | 26 +++--- src/mitim_modules/powertorch/STATEtools.py | 1 + .../powertorch/scripts/calculateTargets.py | 91 ++++--------------- .../powertorch/utils/TRANSPORTtools.py | 10 +- tests/PORTALS_workflow.py | 10 +- 7 files changed, 60 insertions(+), 141 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index c3478915..84e1a63d 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -148,34 +148,13 @@ def __init__( "portals_transformation_variables_trace": portals_transformation_variables_trace, # Physics-informed parameters to fit surrogates for trace impurities "turbulent_exchange_as_surrogate": False, # Run turbulent exchange as surrogate? "Pseudo_multipliers": [1.0]*5, # [Qe,Qi,Ge] multipliers to calculate pseudo - "impurity_trick": True, # If True, fit model to GZ/nZ, valid on the trace limit - "fZ0_as_weight": 1.0, # If not None, using fZ0_as_weight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis + "impurity_trick": True, # If True, fit model to GZ/nZ, valid on the trace limit + "fZ0_as_weight": 1.0, # If not None, using fZ0_as_weight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis "fImp_orig": 1.0, "additional_params_in_surrogate": additional_params_in_surrogate, "keep_full_model_folder": True, # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) } - """ - Parameters to initialize files - ------------------------------ - These parameters are used to initialize the input.gacode to work with, before any PORTALS workflow - ( passed directly to profiles.correct() ) - Bear in mind that this is not necessary, you provide an already ready-to-go input.gacode without the need - to run these corrections. - """ - - self.portals_parameters["initialization"] = { - "recalculate_ptot": True, # Recompute PTOT to match kinetic profiles (after removals) - "quasineutrality": False, # Make sure things are quasineutral by changing the *MAIN* ion (D,T or both) (after removals) - "removeIons": [], # Remove this ion from the input.gacode (if D,T,Z, eliminate T with [2]) - "remove_fast": False, # Automatically detect which are fast ions and remove them - "thermalize_fast": False, # Do not remove fast, keep their diluiton effect but make them thermal - "enforce_same_aLn": False, # Make all ion density gradients equal to electrons - "groupQIONE": False, - "ensure_positive_Gamma": False, - "ensureMachNumber": None, - } - ''' model_parameters ---------------- @@ -200,6 +179,8 @@ def __init__( # Simulation kwargs to be passed directly to run and read commands "options": { + + # Defaults for TGLF simulation "tglf": { "run": { "code_settings": 6, @@ -213,11 +194,15 @@ def __init__( "percent_error": 5, # (%) Error (std, in percent) of model evaluation TGLF if not scan trick "Qi_includes_fast": False, # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) }, + + # Defaults for NEO simulation "neo": { "run": {}, "read": {}, "percent_error": 10 # (%) Error (std, in percent) of model evaluation }, + + # Defaults for CGYRO simulation "cgyro": { "run": { "code_settings": 1, @@ -342,7 +327,6 @@ def prep( self, self.folder, fileGACODE, - self.portals_parameters, ymax_rel, ymin_rel, start_from_folder=start_from_folder, diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 1b8ff337..d3b9294b 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -18,7 +18,6 @@ def initializeProblem( portals_fun, folderWork, fileStart, - portals_parameters, RelVar_y_max, RelVar_y_min, limitsAreRelative=True, @@ -58,8 +57,10 @@ def initializeProblem( # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # ---- Copy the file of interest to initialization folder - - shutil.copy2(fileStart, FolderInitialization / "input.gacode") + if isinstance(fileStart, MITIMstate.mitim_state): + fileStart.write_state(file=FolderInitialization / "input.gacode") + else: + shutil.copy2(fileStart, FolderInitialization / "input.gacode") # ---- Make another copy to preserve the original state @@ -79,14 +80,8 @@ def initializeProblem( print(f"\t\t rho = {rho}") portals_fun.portals_parameters["solution"]["predicted_rho"] = rho - if ( - len(portals_parameters["initialization"]["removeIons"]) > 0 - or portals_parameters["initialization"]["remove_fast"] - or portals_parameters["initialization"]["quasineutrality"] - or portals_parameters["initialization"]["enforce_same_aLn"] - or portals_parameters["initialization"]["recalculate_ptot"] - ): - profiles.correct(options=portals_parameters["initialization"]) + # Good approach to ensure this consistency + profiles.correct(options={"recalculate_ptot": True}) if portals_fun.portals_parameters["solution"]["trace_impurity"] is not None: position_of_impurity = MITIMstate.impurity_location(profiles, portals_fun.portals_parameters["solution"]["trace_impurity"]) @@ -134,7 +129,6 @@ def initializeProblem( # Add folder and cold_start to the simulation options transport_options = transport_parameters | {"folder": portals_fun.folder, "cold_start": False} - target_options = portals_fun.portals_parameters["target"] portals_fun.powerstate = STATEtools.powerstate( @@ -143,6 +137,7 @@ def initializeProblem( "ProfilePredicted": portals_fun.portals_parameters["solution"]["predicted_channels"], "rhoPredicted": xCPs, "impurityPosition": position_of_impurity, + "fImp_orig": portals_fun.portals_parameters["solution"]["fImp_orig"] }, transport_options=transport_options, target_options=target_options, @@ -189,17 +184,9 @@ def initializeProblem( "ProfilePredicted": portals_fun.portals_parameters["solution"]["predicted_channels"], "rhoPredicted": xCPs, "impurityPosition": position_of_impurity, - "fineTargetsResolution": portals_fun.portals_parameters["target"]["fineTargetsResolution"], - }, - target_options={ - "evaluator": portals_fun.portals_parameters["target"]["evaluator"], - "options": { - "TypeTarget": portals_fun.portals_parameters["target"]["TypeTarget"], - "target_evaluator_method": portals_fun.portals_parameters["target"]["target_evaluator_method"], - "forceZeroParticleFlux": portals_fun.portals_parameters["target"]["options"]["forceZeroParticleFlux"], - "percent_error": portals_fun.portals_parameters["target"]["percent_error"] - }, + "fImp_orig": portals_fun.portals_parameters["solution"]["fImp_orig"] }, + target_options=portals_fun.portals_parameters["target"], tensor_options = tensor_options ) diff --git a/src/mitim_modules/portals/utils/PORTALSoptimization.py b/src/mitim_modules/portals/utils/PORTALSoptimization.py index 8564c0c4..2b36f026 100644 --- a/src/mitim_modules/portals/utils/PORTALSoptimization.py +++ b/src/mitim_modules/portals/utils/PORTALSoptimization.py @@ -93,16 +93,16 @@ def initialization_simple_relax(self): def flux_match_surrogate( - step, - profiles, - plot_results=False, - fn = None, - file_write_csv=None, - algorithm = None, - solver_options = None, - keep_within_bounds = True, - target_options_use = None, - ): + step, + profiles, + plot_results=False, + fn = None, + file_write_csv=None, + algorithm = None, + solver_options = None, + keep_within_bounds = True, + target_options_use = None, + ): ''' Technique to reutilize flux surrogates to predict new conditions ---------------------------------------------------------------- @@ -139,7 +139,7 @@ def flux_match_surrogate( transport_options = copy.deepcopy(step.surrogate_parameters["powerstate"].transport_options) # Define transport calculation function as a surrogate model - transport_options['transport_evaluator'] = transport_analytic.surrogate + transport_options['evaluator'] = transport_analytic.surrogate transport_options["options"] = {'flux_fun': partial(step.evaluators['residual_function'],outputComponents=True)} # Create powerstate with the same options as the original portals but with the new profiles @@ -149,13 +149,13 @@ def flux_match_surrogate( "ProfilePredicted": step.surrogate_parameters["powerstate"].predicted_channels, "rhoPredicted": step.surrogate_parameters["powerstate"].plasma["rho"][0,1:], "impurityPosition": step.surrogate_parameters["powerstate"].impurityPosition, - "fineTargetsResolution": step.surrogate_parameters["powerstate"].fineTargetsResolution, }, transport_options=transport_options, target_options= step.surrogate_parameters["powerstate"].target_options if target_options_use is None else target_options_use, tensor_options = { "dtype": step.surrogate_parameters["powerstate"].dfT.dtype, - "device": step.surrogate_parameters["powerstate"].dfT.device}, + "device": step.surrogate_parameters["powerstate"].dfT.device + }, ) # Pass powerstate as part of the surrogate_parameters such that transformations now occur with the new profiles diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index b0bdfa08..ef46e56e 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -82,6 +82,7 @@ def __init__( self.impurityPosition_transport = copy.deepcopy(self.impurityPosition) self.fineTargetsResolution = target_options.get("fineTargetsResolution", None) self.scaleIonDensities = evolution_options.get("scaleIonDensities", True) + self.fImp_orig = evolution_options.get("fImp_orig", 1.0) rho_vec = evolution_options.get("rhoPredicted", [0.2, 0.4, 0.6, 0.8]) if rho_vec[0] == 0: diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 68aba57f..63fbc51a 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -1,5 +1,5 @@ """ -calculateTargets.py input.gacode 1 +calculateTargets.py input.gacode """ import sys @@ -12,7 +12,6 @@ def calculator( input_gacode, - typeCalculation=2, TypeTarget=3, folder="~/scratch/", cold_start=True, @@ -20,75 +19,26 @@ def calculator( profProvided=False, fineTargetsResolution = None, ): - profiles = ( - input_gacode if profProvided else PROFILEStools.gacode_state(input_gacode) - ) + profiles = input_gacode if profProvided else PROFILEStools.gacode_state(input_gacode) - # Calculate using TGYRO - if typeCalculation == 1: - p = STATEtools.powerstate( - profiles, - evolution_options={ - "rhoPredicted": rho_vec, - "fineTargetsResolution": fineTargetsResolution, - }, - target_options={ - "evaluator": targets_analytic.analytical_model, - "options": { - "TypeTarget": TypeTarget, - "target_evaluator_method": "tgyro"}, - }, - transport_options={ - "evaluator": transport_tglfneo.tglfneo_model, - "options": { - "cold_start": cold_start, - "portals_parameters": { - "solution": { - "launchSlurm": True, - "Qi_includes_fast": False, - }, - "model_parameters": { - "Physics_options": { - "TypeTarget": 3, - "TurbulentExchange": 0, - "PtotType": 1, - "GradientsType": 0, - "InputType": 1, - }, - "predicted_channels": ["te", "ti", "ne"], - "predicted_rho": rho_vec, - "applyCorrections": { - "Tfast_ratio": False, - "Ti_thermals": True, - "ni_thermals": True, - "recalculate_ptot": False, - }, - "transport_model": {"TGLFsettings": 5, "extraOptionsTGLF": {}}, - }, - }, + p = STATEtools.powerstate( + profiles, + evolution_options={ + "rhoPredicted": rho_vec, + }, + target_options={ + "evaluator": targets_analytic.analytical_model, + "options": { + "TypeTarget": TypeTarget, + "target_evaluator_method": "powerstate", + "fineTargetsResolution": fineTargetsResolution }, - } - ) - - # Calculate using powerstate - elif typeCalculation == 2: - p = STATEtools.powerstate( - profiles, - evolution_options={ - "rhoPredicted": rho_vec, - "fineTargetsResolution": fineTargetsResolution, - }, - target_options={ - "evaluator": targets_analytic.analytical_model, - "options": { - "TypeTarget": TypeTarget, - "target_evaluator_method": "powerstate"}, - }, - transport_options={ - "evaluator": None, - "options": {} - }, - ) + }, + transport_options={ + "evaluator": None, + "options": {} + }, + ) # Determine performance nameRun="test" @@ -150,6 +100,5 @@ def calculator( if __name__ == "__main__": input_gacode = IOtools.expandPath(sys.argv[1]) - typeCalculation = int(sys.argv[2]) - calculator(input_gacode, typeCalculation=typeCalculation) + calculator(input_gacode) diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 335e741a..c6bd9ca5 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -205,8 +205,6 @@ def evaluate(self): def _postprocess(self): - OriginalFimp = self.powerstate.transport_options["options"].get("OriginalFimp", 1.0) - # ------------------------------------------------------------------------------------------------------------------------ # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) # ------------------------------------------------------------------------------------------------------------------------ @@ -234,9 +232,9 @@ def _postprocess(self): for variable in variables: self.powerstate.plasma[f"{variable}_tr"] = self.powerstate.plasma[f"{variable}_tr_turb"] + self.powerstate.plasma[f"{variable}_tr_neoc"] - # ----------------------------------------------------------- + # --------------------------------------------------------------------------------- # Convective fluxes (& Re-scale the GZ flux by the original impurity concentration) - # ----------------------------------------------------------- + # --------------------------------------------------------------------------------- mapper_convective = { 'Ce': 'Ge1E20m2', @@ -246,13 +244,13 @@ def _postprocess(self): for key in mapper_convective.keys(): for tt in ['','_turb', '_turb_stds', '_neoc', '_neoc_stds']: - mult = 1.0 if key == 'Ce' else 1/OriginalFimp + mult = 1/self.powerstate.fImp_orig if key == 'CZ' else 1.0 self.powerstate.plasma[f"{key}_tr{tt}"] = PLASMAtools.convective_flux( self.powerstate.plasma["te"], self.powerstate.plasma[f"{mapper_convective[key]}_tr{tt}"] ) * mult - + def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): ''' diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index 22baa746..04e56b2d 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -31,17 +31,17 @@ portals_fun.portals_parameters["solution"]['turbulent_exchange_as_surrogate'] = True -portals_fun.portals_parameters["initialization"]["remove_fast"] = True -portals_fun.portals_parameters["initialization"]["quasineutrality"] = True -portals_fun.portals_parameters["initialization"]["enforce_same_aLn"] = True - portals_fun.portals_parameters["solution"]["predicted_rho"] = [0.25, 0.45, 0.65, 0.85] portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti", "ne", "nZ", 'w0'] portals_fun.portals_parameters["solution"]["trace_impurity"] = 'N' portals_fun.portals_parameters["transport"]["options"]["tglf"]["run"]["code_settings"] = 2 +# Prepare case to run +plasma_state = PROFILEStools.gacode_state(inputgacode) +plasma_state.correct(options={"recalculate_ptot": True, "remove_fast": True, "quasineutrality": True, "enforce_same_aLn": True}) + # Prepare run -portals_fun.prep(inputgacode) +portals_fun.prep(plasma_state) # -------------------------------------------------------------------------------------------- # Run From 8848bba61e75663066fe19501bc9aa584347d7f1 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 17:30:00 -0400 Subject: [PATCH 229/385] Recover the writer of portals profiles --- src/mitim_modules/portals/PORTALSmain.py | 2 +- .../physics_models/transport_tglfneo.py | 17 +---------------- .../powertorch/utils/TRANSPORTtools.py | 17 +++++++++++++++++ 3 files changed, 19 insertions(+), 17 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 84e1a63d..db8a183f 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -629,7 +629,7 @@ def runModelEvaluator( # --------------------------------------------------------------------------------------------------- # In certain cases, I want to cold_start the model directly from the PORTALS call instead of powerstate - powerstate.transport_options["options"]["cold_start"] = cold_start + powerstate.transport_options["cold_start"] = cold_start # Evaluate X (DVs) through powerstate.calculate(). This will populate .plasma with the results powerstate.calculate(X, nameRun=name, folder=folder_model, evaluation_number=numPORTALS) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index c89604b2..f79a0a52 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -27,7 +27,6 @@ def evaluate_turbulence(self): def _evaluate_tglf(self): transport_evaluator_options = self.powerstate.transport_options["options"] - cold_start = self.powerstate.transport_options["cold_start"] # ------------------------------------------------------------------------------------------------------------------------ @@ -168,6 +167,7 @@ def _evaluate_tglf(self): def evaluate_neoclassical(self): transport_evaluator_options = self.powerstate.transport_options["options"] + cold_start = self.powerstate.transport_options["cold_start"] # ------------------------------------------------------------------------------------------------------------------------ # Grab options from powerstate @@ -175,8 +175,6 @@ def evaluate_neoclassical(self): simulation_options_neo = transport_evaluator_options["neo"] percent_error = simulation_options_neo["percent_error"] - cold_start = transport_evaluator_options.get("cold_start", False) - impurityPosition = self.powerstate.impurityPosition_transport # ------------------------------------------------------------------------------------------------------------------------ @@ -232,20 +230,7 @@ def evaluate_neoclassical(self): self.QieGB_neoc_stds = Qe * 0.0 return neo - - def _profiles_to_store(self): - - if "folder" in self.powerstate.transport_options["options"]: - whereFolder = IOtools.expandPath(self.powerstate.transport_options["folder"] / "Outputs" / "portals_profiles") - if not whereFolder.exists(): - IOtools.askNewFolder(whereFolder) - fil = whereFolder / f"input.gacode.{self.evaluation_number}" - shutil.copy2(self.file_profs, fil) - shutil.copy2(self.file_profs_unmod, fil.parent / f"{fil.name}_unmodified") - print(f"\t- Copied profiles to {IOtools.clipstr(fil)}") - else: - print("\t- Could not move files", typeMsg="w") def _raise_warnings(self, tglf, rho_locations, Qi_includes_fast): diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index c6bd9ca5..c3ea70f2 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -179,6 +179,9 @@ def _modify_profiles(self): def evaluate(self): + # Copy the input.gacode files to the output folder + self._profiles_to_store() + ''' ****************************************************************************************************** Evaluate neoclassical and turbulent transport. @@ -250,6 +253,20 @@ def _postprocess(self): self.powerstate.plasma["te"], self.powerstate.plasma[f"{mapper_convective[key]}_tr{tt}"] ) * mult + + def _profiles_to_store(self): + + if "folder" in self.powerstate.transport_options: + whereFolder = IOtools.expandPath(self.powerstate.transport_options["folder"] / "Outputs" / "portals_profiles") + if not whereFolder.exists(): + IOtools.askNewFolder(whereFolder) + + fil = whereFolder / f"input.gacode.{self.evaluation_number}" + shutil.copy2(self.file_profs, fil) + shutil.copy2(self.file_profs_unmod, fil.parent / f"{fil.name}_unmodified") + print(f"\t- Copied profiles to {IOtools.clipstr(fil)}") + else: + print("\t- Could not move files", typeMsg="w") def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): From 7b9ae828ffdf357ab6a7cd21099614dca506b773 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 18:01:37 -0400 Subject: [PATCH 230/385] misc fixes --- src/mitim_modules/powertorch/STATEtools.py | 9 ++-- .../physics_models/transport_cgyroneo.py | 2 + .../physics_models/transport_tglfneo.py | 5 --- .../powertorch/scripts/calculateTargets.py | 17 ++++---- .../powertorch/utils/TRANSPORTtools.py | 41 +++++++++++-------- 5 files changed, 41 insertions(+), 33 deletions(-) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index ef46e56e..f0c8fc15 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -57,9 +57,11 @@ def __init__( "options": { "TypeTarget": 3, "target_evaluator_method": "powerstate", - "forceZeroParticleFlux": False }, } + + target_options["options"].setdefault("forceZeroParticleFlux", False) + target_options["options"].setdefault("percent_error", 1.0) if tensor_options is None: tensor_options = { @@ -687,7 +689,7 @@ def calculateTargets(self, relative_error_assumed=1.0): Update the targets of the current state """ - # If no targets evaluator is given or the targets will come from TGYRO, assume them as zero + # If no targets evaluator is given or the targets will come from previous calculations (from transport), assume them as zero if (self.target_options["evaluator"] is None) or (self.target_options["options"]["target_evaluator_method"] == "tgyro"): targets = TARGETStools.power_targets(self) else: @@ -710,7 +712,8 @@ def calculateTargets(self, relative_error_assumed=1.0): # Merge targets, calculate errors and normalize targets.postprocessing( relative_error_assumed=relative_error_assumed, - forceZeroParticleFlux=self.target_options["options"]["forceZeroParticleFlux"]) + forceZeroParticleFlux=self.target_options["options"]["forceZeroParticleFlux"] + ) def calculateTransport( self, nameRun="test", folder="~/scratch/", evaluation_number=0): diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py index de234a55..8b1efa7c 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py @@ -87,6 +87,8 @@ def evaluate_turbulence(self): # Wait until the user has placed the json file in the right folder + self.powerstate.profiles_transport.write_state(self.folder / "base_cgyro" / "input.gacode") + pre_checks(self) file_path = self.folder / 'fluxes_turb.json' diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py index f79a0a52..45309fec 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py @@ -225,13 +225,8 @@ def evaluate_neoclassical(self): self.GZGB_neoc_stds = abs(GZ) * percent_error/100.0 self.MtGB_neoc_stds = abs(Mt) * percent_error/100.0 - # No neoclassical exchange - self.QieGB_neoc = Qe * 0.0 - self.QieGB_neoc_stds = Qe * 0.0 - return neo - def _raise_warnings(self, tglf, rho_locations, Qi_includes_fast): for i in range(len(tglf.profiles.Species)): diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 63fbc51a..6b7a03b1 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -44,6 +44,9 @@ def calculator( nameRun="test" folder=IOtools.expandPath(folder) + if not folder.exists(): + folder.mkdir(parents=True) + # ************************************ # Calculate state # ************************************ @@ -80,19 +83,14 @@ def calculator( rederive_profiles=False, ) - p.plasma["QiMWm2n"] = ( - (p.plasma["Paux_e"] + p.plasma["Paux_i"]) * p.plasma["volp"] - )[..., -1] - p.plasma["Q"] = p.plasma["Pfus"] / p.plasma["Pin"] + p.plasma["Q"] = p.profiles.derived["Q"] + p.plasma['Prad'] = p.profiles.derived['Prad'] # ************************************ # Print Info # ************************************ - print( - f"Q = {p.plasma['Q'].item():.2f} (Pfus = {p.plasma['Pfus'].item():.2f}MW, Pin = {p.plasma['Pin'].item():.2f}MW)" - ) - + print(f"Q = {p.plasma['Q'].item():.2f}") print(f"Prad = {p.plasma['Prad'].item():.2f}MW") return p @@ -100,5 +98,6 @@ def calculator( if __name__ == "__main__": input_gacode = IOtools.expandPath(sys.argv[1]) + folder = IOtools.expandPath(sys.argv[2]) - calculator(input_gacode) + calculator(input_gacode, folder=folder) diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index c3ea70f2..3ff4397e 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -46,10 +46,18 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): fluxes_mean = {} fluxes_stds = {} - for var in ['QeGB', 'QiGB', 'GeGB', 'GZGB', 'MtGB', 'QieGB']: + for var in ['QeGB', 'QiGB', 'GeGB', 'GZGB', 'MtGB']: fluxes_mean[var] = self.__dict__[f"{var}_{suffix}"].tolist() fluxes_stds[var] = self.__dict__[f"{var}_{suffix}_stds"].tolist() + try: + var = 'QieGB' + fluxes_mean[var] = self.__dict__[f"{var}_{suffix}"].tolist() + fluxes_stds[var] = self.__dict__[f"{var}_{suffix}_stds"].tolist() + except KeyError: + # NEO file may not have it + pass + json_dict = { 'fluxes_mean': fluxes_mean, 'fluxes_stds': fluxes_stds, @@ -357,16 +365,16 @@ def evaluate_turbulence(self): dim = self.powerstate.plasma['rho'].shape[-1]-1 for var in [ - 'QeMWm2', - 'QiMWm2', - 'Ge1E20m2', - 'GZ1E20m2', - 'MtJm2', - 'QieMWm3' + 'QeGB', + 'QiGB', + 'GeGB', + 'GZGB', + 'MtGB', + 'QieGB' ]: - self.__dict__[f"{var}_tr_turb"] = np.zeros(dim) - self.__dict__[f"{var}_tr_turb_stds"] = np.zeros(dim) + self.__dict__[f"{var}_turb"] = np.zeros(dim) + self.__dict__[f"{var}_turb_stds"] = np.zeros(dim) @IOtools.hook_method(after=partial(write_json, file_name = 'fluxes_neoc.json', suffix= 'neoc')) def evaluate_neoclassical(self): @@ -385,13 +393,14 @@ def evaluate_neoclassical(self): dim = self.powerstate.plasma['rho'].shape[-1]-1 for var in [ - 'QeMWm2', - 'QiMWm2', - 'Ge1E20m2', - 'GZ1E20m2', - 'MtJm2', + 'QeGB', + 'QiGB', + 'GeGB', + 'GZGB', + 'MtGB', + 'QieGB' ]: - self.__dict__[f"{var}_tr_neoc"] = np.zeros(dim) - self.__dict__[f"{var}_tr_neoc_stds"] = np.zeros(dim) + self.__dict__[f"{var}_neoc"] = np.zeros(dim) + self.__dict__[f"{var}_neoc_stds"] = np.zeros(dim) \ No newline at end of file From 89ace37b6b48ff22cbfff0ed027452b97c376b2a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 19:11:34 -0400 Subject: [PATCH 231/385] Beautiful new infrastructure for transport model selection --- src/mitim_modules/portals/PORTALSmain.py | 17 +- src/mitim_modules/powertorch/STATEtools.py | 8 +- ...ansport_cgyroneo.py => transport_cgyro.py} | 73 +++--- .../physics_models/transport_neo.py | 68 ++++++ ...transport_tglfneo.py => transport_tglf.py} | 109 ++------- .../powertorch/utils/TRANSPORTtools.py | 211 ++++++++++-------- 6 files changed, 249 insertions(+), 237 deletions(-) rename src/mitim_modules/powertorch/physics_models/{transport_cgyroneo.py => transport_cgyro.py} (77%) create mode 100644 src/mitim_modules/powertorch/physics_models/transport_neo.py rename src/mitim_modules/powertorch/physics_models/{transport_tglfneo.py => transport_tglf.py} (74%) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index db8a183f..9a64e1de 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -80,7 +80,7 @@ class portals(STRATEGYtools.opt_evaluator): def __init__( self, folder, # Folder where the PORTALS workflow will be run - transport_model = 'tglf', + transport_models = ['tglf', 'neo'], namelist=None, # If None, default namelist will be used. If not None, it will be read and used tensor_options = { "dtype": torch.double, @@ -160,13 +160,9 @@ def __init__( ---------------- These parameters are communicated to the powertorch object. ''' + + from mitim_modules.powertorch.utils.TRANSPORTtools import portals_model as transport_evaluator - # Selection of model - if transport_model == 'cgyro': - from mitim_modules.powertorch.physics_models.transport_cgyroneo import cgyroneo_model as transport_evaluator - elif transport_model == 'tglf': - from mitim_modules.powertorch.physics_models.transport_tglfneo import tglfneo_model as transport_evaluator - # ------------------------- # Transport model settings # ------------------------- @@ -175,7 +171,12 @@ def __init__( # Transport model class "evaluator": transport_evaluator, - + + "evaluator_instance_attributes": { + "turbulence_model": transport_models[0], + "neoclassical_model": transport_models[1], + }, + # Simulation kwargs to be passed directly to run and read commands "options": { diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index f0c8fc15..1f9d97d7 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -50,7 +50,9 @@ def __init__( "evaluator": None, "options": {} } - + + transport_options.setdefault("evaluator_instance_attributes", {}) + if target_options is None: target_options = { "evaluator": targets_analytic.analytical_model, @@ -728,6 +730,10 @@ def calculateTransport( else: transport = self.transport_options["evaluator"]( self, name=nameRun, folder=folder, evaluation_number=evaluation_number ) + # The transport class may have instanciating attributes + for key in self.transport_options["evaluator_instance_attributes"]: + setattr(transport, key, self.transport_options["evaluator_instance_attributes"][key]) + # Produce profile object (for certain transport evaluators, this is necessary) transport.produce_profiles() diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py similarity index 77% rename from src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py rename to src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 8b1efa7c..ee90d45f 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyroneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -1,32 +1,33 @@ import json +from functools import partial import numpy as np +from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import CGYROtools -from mitim_modules.powertorch.physics_models import transport_tglfneo +from mitim_modules.powertorch.physics_models import transport_tglf +from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -# Inherit from transport_tglfneo.tglfneo_model so that I have the NEO evaluator -class cgyroneo_model(transport_tglfneo.tglfneo_model): - def __init__(self, powerstate, **kwargs): - super().__init__(powerstate, **kwargs) - - # Do not hook here +class cgyro_model: + def evaluate_turbulence(self): - transport_evaluator_options = self.powerstate.transport_options["options"] - cold_start = self.powerstate.transport_options["cold_start"] + # ------------------------------------------------------------------------------------------------------------------------ + # Grab options + # ------------------------------------------------------------------------------------------------------------------------ + + simulation_options = self.transport_evaluator_options["cgyro"] + simulation_options_tglf = self.transport_evaluator_options["tglf"] + cold_start = self.cold_start + + rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] + run_type = simulation_options["run"]["run_type"] # Run base TGLF always, to keep track of discrepancies! -------------------------------------- - transport_evaluator_options["tglf"]["use_scan_trick_for_stds"] = None + simulation_options_tglf["use_scan_trick_for_stds"] = None self._evaluate_tglf() # -------------------------------------------------------------------------------------------- - rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] - - simulation_options_cgyro = transport_evaluator_options["cgyro"] - - run_type = simulation_options_cgyro["run"]["run_type"] - # ------------------------------------------------------------------------------------------------------------------------ # Prepare CGYRO object # ------------------------------------------------------------------------------------------------------------------------ @@ -40,11 +41,13 @@ def evaluate_turbulence(self): self.folder, ) + subfolder_name = "base_cgyro" + _ = cgyro.run( - 'base_cgyro', + subfolder_name, cold_start=cold_start, forceIfcold_start=True, - **simulation_options_cgyro["run"] + **simulation_options["run"] ) if run_type in ['normal', 'submit']: @@ -54,22 +57,22 @@ def evaluate_turbulence(self): cgyro.fetch() cgyro.read( - label='base_cgyro', - **simulation_options_cgyro["read"] + label=subfolder_name, + **simulation_options["read"] ) # ------------------------------------------------------------------------------------------------------------------------ # Pass the information to what power_transport expects # ------------------------------------------------------------------------------------------------------------------------ - self.QeGB_turb = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Qe_mean for i in range(len(rho_locations))]) - self.QeGB_turb_stds = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Qe_std for i in range(len(rho_locations))]) + self.QeGB_turb = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Qe_mean for i in range(len(rho_locations))]) + self.QeGB_turb_stds = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Qe_std for i in range(len(rho_locations))]) - self.QiGB_turb = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Qi_mean for i in range(len(rho_locations))]) - self.QiGB_turb_stds = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Qi_std for i in range(len(rho_locations))]) + self.QiGB_turb = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Qi_mean for i in range(len(rho_locations))]) + self.QiGB_turb_stds = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Qi_std for i in range(len(rho_locations))]) - self.GeGB_turb = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Ge_mean for i in range(len(rho_locations))]) - self.GeGB_turb_stds = np.array([cgyro.results['base_cgyro']['CGYROout'][i].Ge_std for i in range(len(rho_locations))]) + self.GeGB_turb = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Ge_mean for i in range(len(rho_locations))]) + self.GeGB_turb_stds = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Ge_std for i in range(len(rho_locations))]) self.GZGB_turb = self.QeGB_turb*0.0 #TODO self.GZGB_turb_stds = self.QeGB_turb*0.0 #TODO @@ -87,7 +90,7 @@ def evaluate_turbulence(self): # Wait until the user has placed the json file in the right folder - self.powerstate.profiles_transport.write_state(self.folder / "base_cgyro" / "input.gacode") + self.powerstate.profiles_transport.write_state(self.folder / subfolder_name / "input.gacode") pre_checks(self) @@ -102,18 +105,18 @@ def evaluate_turbulence(self): print(f" MITIM could not find the file... looping back", typeMsg='i') print(f"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!", typeMsg='i') print(f"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!", typeMsg='i') - logic_to_wait(self.folder) + logic_to_wait(self.folder, self.folder / subfolder_name) attempts += 1 if file_path.exists(): all_good = post_checks(self) - self._stable_correction(simulation_options_cgyro) + self._stable_correction(simulation_options) - def _stable_correction(self, simulation_options_cgyro): + def _stable_correction(self, simulation_options): - Qi_stable_criterion = simulation_options_cgyro["Qi_stable_criterion"] - Qi_stable_percent_error = simulation_options_cgyro["Qi_stable_percent_error"] + Qi_stable_criterion = simulation_options["Qi_stable_criterion"] + Qi_stable_percent_error = simulation_options["Qi_stable_percent_error"] # Check if Qi in MW/m2 < Qi_stable_criterion QiMWm2 = self.QiGB_turb * self.powerstate.plasma['Qgb'][0,1:].cpu().numpy() @@ -154,9 +157,9 @@ def pre_checks(self): print(txt) -def logic_to_wait(folder): - print(f"\n**** CGYRO prepared. Please, run CGYRO from the simulation setup in folder:\n", typeMsg='i') - print(f"\t {folder}/base_cgyro\n", typeMsg='i') +def logic_to_wait(folder, subfolder): + print(f"\n**** Simulation inputs prepared. Please, run it from the simulation setup in folder:\n", typeMsg='i') + print(f"\t {subfolder}\n", typeMsg='i') print(f"**** When finished, the fluxes_turb.json file should be placed in:\n", typeMsg='i') print(f"\t {folder}/fluxes_turb.json\n", typeMsg='i') while not print(f"**** When you have done that, please say yes", typeMsg='q'): diff --git a/src/mitim_modules/powertorch/physics_models/transport_neo.py b/src/mitim_modules/powertorch/physics_models/transport_neo.py new file mode 100644 index 00000000..7fe2252f --- /dev/null +++ b/src/mitim_modules/powertorch/physics_models/transport_neo.py @@ -0,0 +1,68 @@ +import numpy as np +from mitim_tools.gacode_tools import NEOtools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +class neo_model: + + def evaluate_neoclassical(self): + + # ------------------------------------------------------------------------------------------------------------------------ + # Grab options + # ------------------------------------------------------------------------------------------------------------------------ + + simulation_options = self.transport_evaluator_options["neo"] + cold_start = self.cold_start + + percent_error = simulation_options["percent_error"] + impurityPosition = self.powerstate.impurityPosition_transport + + # ------------------------------------------------------------------------------------------------------------------------ + # Run + # ------------------------------------------------------------------------------------------------------------------------ + + rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] + + neo = NEOtools.NEO(rhos=rho_locations) + + _ = neo.prep( + self.powerstate.profiles_transport, + self.folder, + cold_start = cold_start, + ) + + neo.run( + 'base_neo', + cold_start=cold_start, + forceIfcold_start=True, + **simulation_options["run"] + ) + + neo.read( + label='base', + **simulation_options["read"]) + + Qe = np.array([neo.results['base']['NEOout'][i].Qe for i in range(len(rho_locations))]) + Qi = np.array([neo.results['base']['NEOout'][i].Qi for i in range(len(rho_locations))]) + Ge = np.array([neo.results['base']['NEOout'][i].Ge for i in range(len(rho_locations))]) + GZ = np.array([neo.results['base']['NEOout'][i].GiAll[impurityPosition-1] for i in range(len(rho_locations))]) + Mt = np.array([neo.results['base']['NEOout'][i].Mt for i in range(len(rho_locations))]) + + # ------------------------------------------------------------------------------------------------------------------------ + # Pass the information to what power_transport expects + # ------------------------------------------------------------------------------------------------------------------------ + + self.QeGB_neoc = Qe + self.QiGB_neoc = Qi + self.GeGB_neoc = Ge + self.GZGB_neoc = GZ + self.MtGB_neoc = Mt + + # Uncertainties is just a percent of the value + self.QeGB_neoc_stds = abs(Qe) * percent_error/100.0 + self.QiGB_neoc_stds = abs(Qi) * percent_error/100.0 + self.GeGB_neoc_stds = abs(Ge) * percent_error/100.0 + self.GZGB_neoc_stds = abs(GZ) * percent_error/100.0 + self.MtGB_neoc_stds = abs(Mt) * percent_error/100.0 + + return neo diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py similarity index 74% rename from src/mitim_modules/powertorch/physics_models/transport_tglfneo.py rename to src/mitim_modules/powertorch/physics_models/transport_tglf.py index 45309fec..fe322ad9 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglfneo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -1,46 +1,30 @@ - -import shutil import numpy as np from mitim_tools.misc_tools import IOtools -from functools import partial -from mitim_tools.gacode_tools import TGLFtools, NEOtools -from mitim_modules.powertorch.utils import TRANSPORTtools +from mitim_tools.gacode_tools import TGLFtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -class tglfneo_model(TRANSPORTtools.power_transport): - def __init__(self, powerstate, **kwargs): - super().__init__(powerstate, **kwargs) +class tglf_model: - def produce_profiles(self): - self._produce_profiles() - - # ************************************************************************************ - # Private functions for the evaluation - # ************************************************************************************ - - @IOtools.hook_method(after=partial(TRANSPORTtools.write_json, file_name = 'fluxes_turb.json', suffix= 'turb')) def evaluate_turbulence(self): self._evaluate_tglf() # Have it separate such that I can call it from the CGYRO class but without the decorator def _evaluate_tglf(self): - transport_evaluator_options = self.powerstate.transport_options["options"] - cold_start = self.powerstate.transport_options["cold_start"] - # ------------------------------------------------------------------------------------------------------------------------ - # Grab options from powerstate + # Grab options # ------------------------------------------------------------------------------------------------------------------------ - simulation_options_tglf = transport_evaluator_options["tglf"] + simulation_options = self.transport_evaluator_options["tglf"] + cold_start = self.cold_start - Qi_includes_fast = simulation_options_tglf["Qi_includes_fast"] - launchMODELviaSlurm = simulation_options_tglf["launchEvaluationsAsSlurmJobs"] - use_tglf_scan_trick = simulation_options_tglf["use_scan_trick_for_stds"] - cores_per_tglf_instance = simulation_options_tglf["cores_per_tglf_instance"] - keep_tglf_files = simulation_options_tglf["keep_files"] - percent_error = simulation_options_tglf["percent_error"] + Qi_includes_fast = simulation_options["Qi_includes_fast"] + launchMODELviaSlurm = simulation_options["launchEvaluationsAsSlurmJobs"] + use_tglf_scan_trick = simulation_options["use_scan_trick_for_stds"] + cores_per_tglf_instance = simulation_options["cores_per_tglf_instance"] + keep_tglf_files = simulation_options["keep_files"] + percent_error = simulation_options["percent_error"] # Grab impurity from powerstate ( because it may have been modified in produce_profiles() ) impurityPosition = self.powerstate.impurityPosition_transport @@ -76,13 +60,13 @@ def _evaluate_tglf(self): }, attempts_execution=2, only_minimal_files=keep_tglf_files in ['minimal'], - **simulation_options_tglf["run"] + **simulation_options["run"] ) tglf.read( label='base', require_all_files=False, - **simulation_options_tglf["read"]) + **simulation_options["read"]) # Grab values Qe = np.array([tglf.results['base']['TGLFout'][i].Qe for i in range(len(rho_locations))]) @@ -134,7 +118,7 @@ def _evaluate_tglf(self): launchMODELviaSlurm=launchMODELviaSlurm, Qi_includes_fast=Qi_includes_fast, only_minimal_files=keep_tglf_files in ['minimal', 'base'], - **simulation_options_tglf["run"] + **simulation_options["run"] ) self._raise_warnings(tglf, rho_locations, Qi_includes_fast) @@ -163,70 +147,6 @@ def _evaluate_tglf(self): return tglf - @IOtools.hook_method(after=partial(TRANSPORTtools.write_json, file_name = 'fluxes_neoc.json', suffix= 'neoc')) - def evaluate_neoclassical(self): - - transport_evaluator_options = self.powerstate.transport_options["options"] - cold_start = self.powerstate.transport_options["cold_start"] - - # ------------------------------------------------------------------------------------------------------------------------ - # Grab options from powerstate - # ------------------------------------------------------------------------------------------------------------------------ - - simulation_options_neo = transport_evaluator_options["neo"] - percent_error = simulation_options_neo["percent_error"] - impurityPosition = self.powerstate.impurityPosition_transport - - # ------------------------------------------------------------------------------------------------------------------------ - # Run - # ------------------------------------------------------------------------------------------------------------------------ - - rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] - - neo = NEOtools.NEO(rhos=rho_locations) - - _ = neo.prep( - self.powerstate.profiles_transport, - self.folder, - cold_start = cold_start, - ) - - neo.run( - 'base_neo', - cold_start=cold_start, - forceIfcold_start=True, - **simulation_options_neo["run"] - ) - - neo.read( - label='base', - **simulation_options_neo["read"]) - - Qe = np.array([neo.results['base']['NEOout'][i].Qe for i in range(len(rho_locations))]) - Qi = np.array([neo.results['base']['NEOout'][i].Qi for i in range(len(rho_locations))]) - Ge = np.array([neo.results['base']['NEOout'][i].Ge for i in range(len(rho_locations))]) - GZ = np.array([neo.results['base']['NEOout'][i].GiAll[impurityPosition-1] for i in range(len(rho_locations))]) - Mt = np.array([neo.results['base']['NEOout'][i].Mt for i in range(len(rho_locations))]) - - # ------------------------------------------------------------------------------------------------------------------------ - # Pass the information to what power_transport expects - # ------------------------------------------------------------------------------------------------------------------------ - - self.QeGB_neoc = Qe - self.QiGB_neoc = Qi - self.GeGB_neoc = Ge - self.GZGB_neoc = GZ - self.MtGB_neoc = Mt - - # Uncertainties is just a percent of the value - self.QeGB_neoc_stds = abs(Qe) * percent_error/100.0 - self.QiGB_neoc_stds = abs(Qi) * percent_error/100.0 - self.GeGB_neoc_stds = abs(Ge) * percent_error/100.0 - self.GZGB_neoc_stds = abs(GZ) * percent_error/100.0 - self.MtGB_neoc_stds = abs(Mt) * percent_error/100.0 - - return neo - def _raise_warnings(self, tglf, rho_locations, Qi_includes_fast): for i in range(len(tglf.profiles.Species)): @@ -247,7 +167,6 @@ def _raise_warnings(self, tglf, rho_locations, Qi_includes_fast): else: print(f"\t\t\t* The thermal ion considered by TGLF was summed into the Qi", typeMsg="i") - def _run_tglf_uncertainty_model( tglf, rho_locations, diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 3ff4397e..150acdc8 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -76,14 +76,7 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): print(f"\t* Written JSON with {suffix} information to {self.folder / file_name}") class power_transport: - ''' - Default class for power transport models, change "evaluate" method to implement a new model and produce_profiles if the model requires written input.gacode written - - Notes: - - After evaluation, the self.model_results attribute will contain the results of the model, which can be used for plotting and analysis - - model results can have .plot() method that can grab kwargs or be similar to TGYRO plot - ''' def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_number = 0): self.name = name @@ -91,6 +84,9 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ self.evaluation_number = evaluation_number self.powerstate = powerstate + self.transport_evaluator_options = self.powerstate.transport_options["options"] + self.cold_start = self.powerstate.transport_options["cold_start"] + # Allowed fluxes in powerstate so far self.quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2'] @@ -128,65 +124,14 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ "w0": "$M_T$ ($J/m^2$)", } - def produce_profiles(self): - # Only add self._produce_profiles() if it's needed (e.g. full TGLF), otherwise this is somewhat expensive - # (e.g. for flux matching of analytical models) - pass - - def _produce_profiles(self,derive_quantities=True): - - self.applyCorrections = self.powerstate.transport_options["applyCorrections"] - - # Write this updated profiles class (with parameterized profiles and target powers) - self.file_profs = self.folder / "input.gacode" - - powerstate_detached = self.powerstate.copy_state() - - self.powerstate.profiles = powerstate_detached.from_powerstate( - write_input_gacode=self.file_profs, - postprocess_input_gacode=self.applyCorrections, - rederive_profiles = derive_quantities, # Derive quantities so that it's ready for analysis and plotting later - insert_highres_powers = derive_quantities, # Insert powers so that Q, Pfus and all that it's consistent when read later - ) - - self.powerstate.profiles_transport = copy.deepcopy(self.powerstate.profiles) - - self._modify_profiles() - - def _modify_profiles(self): - ''' - Modify the profiles (e.g. lumping) before running the transport model - ''' - - # After producing the profiles, copy for future modifications - self.file_profs_unmod = self.file_profs.parent / f"{self.file_profs.name}_unmodified" - shutil.copy2(self.file_profs, self.file_profs_unmod) - - profiles_postprocessing_fun = self.powerstate.transport_options["profiles_postprocessing_fun"] - - if profiles_postprocessing_fun is not None: - print(f"\t- Modifying input.gacode to run transport calculations based on {profiles_postprocessing_fun}",typeMsg="i") - self.powerstate.profiles_transport = profiles_postprocessing_fun(self.file_profs) - - # Position of impurity ion may have changed - p_old = PROFILEStools.gacode_state(self.file_profs_unmod) - p_new = PROFILEStools.gacode_state(self.file_profs) - - impurity_of_interest = p_old.Species[self.powerstate.impurityPosition] - - try: - impurityPosition_new = p_new.Species.index(impurity_of_interest) - - except ValueError: - print(f"\t- Impurity {impurity_of_interest} not found in new profiles, keeping position {self.powerstate.impurityPosition}",typeMsg="w") - impurityPosition_new = self.powerstate.impurityPosition - - if impurityPosition_new != self.powerstate.impurityPosition: - print(f"\t- Impurity position has changed from {self.powerstate.impurityPosition} to {impurityPosition_new}",typeMsg="i") - self.powerstate.impurityPosition_transport = p_new.Species.index(impurity_of_interest) - def evaluate(self): + # Initialize them as zeros + for var in ['QeGB','QiGB','GeGB','GZGB','MtGB','QieGB']: + for suffix in ['turb', 'neoc']: + for suffix0 in ['', '_stds']: + self.__dict__[f"{var}_{suffix}{suffix0}"] = np.zeros(self.powerstate.plasma['rho'].shape[-1]-1) + # Copy the input.gacode files to the output folder self._profiles_to_store() @@ -213,7 +158,7 @@ def evaluate(self): ****************************************************************************************************** ''' self._postprocess() - + def _postprocess(self): # ------------------------------------------------------------------------------------------------------------------------ @@ -261,7 +206,64 @@ def _postprocess(self): self.powerstate.plasma["te"], self.powerstate.plasma[f"{mapper_convective[key]}_tr{tt}"] ) * mult - + + def produce_profiles(self): + # Only add self._produce_profiles() if it's needed (e.g. full TGLF), otherwise this is somewhat expensive + # (e.g. for flux matching of analytical models) + pass + + def _produce_profiles(self,derive_quantities=True): + + self.applyCorrections = self.powerstate.transport_options["applyCorrections"] + + # Write this updated profiles class (with parameterized profiles and target powers) + self.file_profs = self.folder / "input.gacode" + + powerstate_detached = self.powerstate.copy_state() + + self.powerstate.profiles = powerstate_detached.from_powerstate( + write_input_gacode=self.file_profs, + postprocess_input_gacode=self.applyCorrections, + rederive_profiles = derive_quantities, # Derive quantities so that it's ready for analysis and plotting later + insert_highres_powers = derive_quantities, # Insert powers so that Q, Pfus and all that it's consistent when read later + ) + + self.powerstate.profiles_transport = copy.deepcopy(self.powerstate.profiles) + + self._modify_profiles() + + def _modify_profiles(self): + ''' + Modify the profiles (e.g. lumping) before running the transport model + ''' + + # After producing the profiles, copy for future modifications + self.file_profs_unmod = self.file_profs.parent / f"{self.file_profs.name}_unmodified" + shutil.copy2(self.file_profs, self.file_profs_unmod) + + profiles_postprocessing_fun = self.powerstate.transport_options["profiles_postprocessing_fun"] + + if profiles_postprocessing_fun is not None: + print(f"\t- Modifying input.gacode to run transport calculations based on {profiles_postprocessing_fun}",typeMsg="i") + self.powerstate.profiles_transport = profiles_postprocessing_fun(self.file_profs) + + # Position of impurity ion may have changed + p_old = PROFILEStools.gacode_state(self.file_profs_unmod) + p_new = PROFILEStools.gacode_state(self.file_profs) + + impurity_of_interest = p_old.Species[self.powerstate.impurityPosition] + + try: + impurityPosition_new = p_new.Species.index(impurity_of_interest) + + except ValueError: + print(f"\t- Impurity {impurity_of_interest} not found in new profiles, keeping position {self.powerstate.impurityPosition}",typeMsg="w") + impurityPosition_new = self.powerstate.impurityPosition + + if impurityPosition_new != self.powerstate.impurityPosition: + print(f"\t- Impurity position has changed from {self.powerstate.impurityPosition} to {impurityPosition_new}",typeMsg="i") + self.powerstate.impurityPosition_transport = p_new.Species.index(impurity_of_interest) + def _profiles_to_store(self): if "folder" in self.powerstate.transport_options: @@ -312,7 +314,7 @@ def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): 'GeGB': ['Ggb', 'Ge1E20m2'], 'GZGB': ['Ggb', 'GZ1E20m2'], 'MtGB': ['Pgb', 'MtJm2'], - 'QieGB': ['Sgb', 'QieMWm3'], + 'QieGB': ['Sgb', 'QieMWm3'] } dum = {} @@ -321,6 +323,10 @@ def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): dum[f"{mapper[var][1]}_tr_{suffix}"] = np.array(json_dict['fluxes_mean'][var]) * gb dum[f"{mapper[var][1]}_tr_{suffix}_stds"] = np.array(json_dict['fluxes_stds'][var]) * gb + mapperQ = { + + } + if units == 'GB': print("\t\t- File has fluxes in GB units... using GB units from powerstate to convert to real units") @@ -361,20 +367,6 @@ def evaluate_turbulence(self): ''' print(">> No turbulent fluxes to evaluate", typeMsg="w") - - dim = self.powerstate.plasma['rho'].shape[-1]-1 - - for var in [ - 'QeGB', - 'QiGB', - 'GeGB', - 'GZGB', - 'MtGB', - 'QieGB' - ]: - - self.__dict__[f"{var}_turb"] = np.zeros(dim) - self.__dict__[f"{var}_turb_stds"] = np.zeros(dim) @IOtools.hook_method(after=partial(write_json, file_name = 'fluxes_neoc.json', suffix= 'neoc')) def evaluate_neoclassical(self): @@ -382,25 +374,48 @@ def evaluate_neoclassical(self): This needs to populate the following np.arrays in self.: - QeGB_neoc - QiGB_neoc - - GeGB_tr_neoc - - GZGB_tr_neoc - - MtGB_tr_neoc - and their respective standard deviations, e.g. QeGB_tr_neoc_stds + - GeGB_neoc + - GZGB_neoc + - MtGB_neoc + - QieGB_neoc (zero) + and their respective standard deviations, e.g. QeGB_neoc_stds ''' print(">> No neoclassical fluxes to evaluate", typeMsg="w") - dim = self.powerstate.plasma['rho'].shape[-1]-1 - for var in [ - 'QeGB', - 'QiGB', - 'GeGB', - 'GZGB', - 'MtGB', - 'QieGB' - ]: - - self.__dict__[f"{var}_neoc"] = np.zeros(dim) - self.__dict__[f"{var}_neoc_stds"] = np.zeros(dim) - \ No newline at end of file +# ******************************************************************************************* +# Combinations +# ******************************************************************************************* + +from mitim_modules.powertorch.physics_models.transport_tglf import tglf_model +from mitim_modules.powertorch.physics_models.transport_neo import neo_model +from mitim_modules.powertorch.physics_models.transport_cgyro import cgyro_model + +class portals_model(power_transport, tglf_model, neo_model, cgyro_model): + + def __init__(self, powerstate, **kwargs): + super().__init__(powerstate, **kwargs) + + # Defaults + self.turbulence_model = 'tglf' + self.neoclassical_model = 'neo' + + def produce_profiles(self): + self._produce_profiles() + + @IOtools.hook_method(after=partial(write_json, file_name = 'fluxes_turb.json', suffix= 'turb')) + def evaluate_turbulence(self): + if self.turbulence_model == 'tglf': + return tglf_model.evaluate_turbulence(self) + elif self.turbulence_model == 'cgyro': + return cgyro_model.evaluate_turbulence(self) + else: + raise Exception(f"Unknown turbulence model {self.turbulence_model}") + + @IOtools.hook_method(after=partial(write_json, file_name = 'fluxes_neoc.json', suffix= 'neoc')) + def evaluate_neoclassical(self): + if self.neoclassical_model == 'neo': + return neo_model.evaluate_neoclassical(self) + else: + raise Exception(f"Unknown neoclassical model {self.neoclassical_model}") From 148082af073efc2f454f7af57ba801b09124a875 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 22:20:33 -0400 Subject: [PATCH 232/385] bug fix --- src/mitim_tools/simulation_tools/SIMtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index c0b4fae2..5c4caf59 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -353,7 +353,7 @@ def _run( tmpFolder = self.FolderGACODE / f"tmp_{code}" IOtools.askNewFolder(tmpFolder, force=True) - kkeys = [keys for keys in code_executor.keys()] + kkeys = [keys.strip('/') for keys in code_executor.keys()] log_simulation_file=self.FolderGACODE / f"mitim_simulation_{kkeys[0]}.log" # Refer with the first folder self.simulation_job = FARMINGtools.mitim_job(tmpFolder, log_simulation_file=log_simulation_file) From 00fb5e9a5d1c455232bc824729c5f059dbc471bb Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 22:23:11 -0400 Subject: [PATCH 233/385] First implementaiton of PORTALS-GX working --- src/mitim_modules/portals/PORTALSmain.py | 13 +++- .../physics_models/transport_cgyro.py | 65 ++++++++++--------- .../powertorch/physics_models/transport_gx.py | 10 +++ .../powertorch/utils/TRANSPORTtools.py | 5 +- .../simulation_tools/physics/GXtools.py | 27 ++++++-- 5 files changed, 82 insertions(+), 38 deletions(-) create mode 100644 src/mitim_modules/powertorch/physics_models/transport_gx.py diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 9a64e1de..c2f1c249 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -161,7 +161,7 @@ def __init__( These parameters are communicated to the powertorch object. ''' - from mitim_modules.powertorch.utils.TRANSPORTtools import portals_model as transport_evaluator + from mitim_modules.powertorch.utils.TRANSPORTtools import portals_transport_model as transport_evaluator # ------------------------- # Transport model settings @@ -216,6 +216,17 @@ def __init__( "Qi_stable_criterion": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable "Qi_stable_percent_error": 5.0, # (%) For CGYRO runs, minimum error based on target if case is considered stable }, + # Defaults for GX simulation + "gx": { + "run": { + "code_settings": 1, + "extraOptions": {}, + "run_type": "normal", # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit + }, + "read": { + "tmin": 0.0 + }, + }, }, # Corrections to be applied to each iteration input.gacode file diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index ee90d45f..7fd69c59 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -1,49 +1,36 @@ import json -from functools import partial import numpy as np -from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import CGYROtools -from mitim_modules.powertorch.physics_models import transport_tglf -from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -class cgyro_model: - - def evaluate_turbulence(self): +class gyrokinetic_model: + def _evaluate_gyrokinetic_model(self, code = 'cgyro', gk_object = None, out_name = 'CGYROout'): # ------------------------------------------------------------------------------------------------------------------------ # Grab options # ------------------------------------------------------------------------------------------------------------------------ - simulation_options = self.transport_evaluator_options["cgyro"] - simulation_options_tglf = self.transport_evaluator_options["tglf"] + simulation_options = self.transport_evaluator_options[code] cold_start = self.cold_start rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] run_type = simulation_options["run"]["run_type"] - # Run base TGLF always, to keep track of discrepancies! -------------------------------------- - simulation_options_tglf["use_scan_trick_for_stds"] = None - self._evaluate_tglf() - # -------------------------------------------------------------------------------------------- - # ------------------------------------------------------------------------------------------------------------------------ - # Prepare CGYRO object + # Prepare object # ------------------------------------------------------------------------------------------------------------------------ - - rho_locations = [self.powerstate.plasma["rho"][0, 1:][i].item() for i in range(len(self.powerstate.plasma["rho"][0, 1:]))] - - cgyro = CGYROtools.CGYRO(rhos=rho_locations) + + gk_object = gk_object(rhos=rho_locations) - _ = cgyro.prep( + _ = gk_object.prep( self.powerstate.profiles_transport, self.folder, ) - subfolder_name = "base_cgyro" + subfolder_name = f"base_{code}" - _ = cgyro.run( + _ = gk_object.run( subfolder_name, cold_start=cold_start, forceIfcold_start=True, @@ -53,10 +40,10 @@ def evaluate_turbulence(self): if run_type in ['normal', 'submit']: if run_type in ['submit']: - cgyro.check(every_n_minutes=10) - cgyro.fetch() + gk_object.check(every_n_minutes=10) + gk_object.fetch() - cgyro.read( + gk_object.read( label=subfolder_name, **simulation_options["read"] ) @@ -65,14 +52,14 @@ def evaluate_turbulence(self): # Pass the information to what power_transport expects # ------------------------------------------------------------------------------------------------------------------------ - self.QeGB_turb = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Qe_mean for i in range(len(rho_locations))]) - self.QeGB_turb_stds = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Qe_std for i in range(len(rho_locations))]) + self.QeGB_turb = np.array([gk_object.results[subfolder_name][out_name][i].Qe_mean for i in range(len(rho_locations))]) + self.QeGB_turb_stds = np.array([gk_object.results[subfolder_name][out_name][i].Qe_std for i in range(len(rho_locations))]) - self.QiGB_turb = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Qi_mean for i in range(len(rho_locations))]) - self.QiGB_turb_stds = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Qi_std for i in range(len(rho_locations))]) + self.QiGB_turb = np.array([gk_object.results[subfolder_name][out_name][i].Qi_mean for i in range(len(rho_locations))]) + self.QiGB_turb_stds = np.array([gk_object.results[subfolder_name][out_name][i].Qi_std for i in range(len(rho_locations))]) - self.GeGB_turb = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Ge_mean for i in range(len(rho_locations))]) - self.GeGB_turb_stds = np.array([cgyro.results[subfolder_name]['CGYROout'][i].Ge_std for i in range(len(rho_locations))]) + self.GeGB_turb = np.array([gk_object.results[subfolder_name][out_name][i].Ge_mean for i in range(len(rho_locations))]) + self.GeGB_turb_stds = np.array([gk_object.results[subfolder_name][out_name][i].Ge_std for i in range(len(rho_locations))]) self.GZGB_turb = self.QeGB_turb*0.0 #TODO self.GZGB_turb_stds = self.QeGB_turb*0.0 #TODO @@ -131,11 +118,25 @@ def _stable_correction(self, simulation_options): print(f"\t\t- Assigning {Qi_stable_percent_error:.1f}% from target as standard deviation: {Qi_std:.2f} instead of {self.QiGB_turb_stds[i]}", typeMsg='i') self.QiGB_turb_stds[i] = Qi_std + +class cgyro_model(gyrokinetic_model): + + def evaluate_turbulence(self): + + # Run base TGLF always, to keep track of discrepancies! -------------------------------------- + simulation_options_tglf = self.transport_evaluator_options["tglf"] + simulation_options_tglf["use_scan_trick_for_stds"] = None + self._evaluate_tglf() + # -------------------------------------------------------------------------------------------- + + self._evaluate_gyrokinetic_model(code = 'cgyro', gk_object = CGYROtools.CGYRO, out_name = 'CGYROout') + + def pre_checks(self): plasma = self.powerstate.plasma - txt = "\nFluxes to be matched by CGYRO ( Target - Neoclassical ):" + txt = "\nFluxes to be matched by turbulence ( Target - Neoclassical ):" # Print gradients for var, varn in zip( diff --git a/src/mitim_modules/powertorch/physics_models/transport_gx.py b/src/mitim_modules/powertorch/physics_models/transport_gx.py new file mode 100644 index 00000000..d5bd565e --- /dev/null +++ b/src/mitim_modules/powertorch/physics_models/transport_gx.py @@ -0,0 +1,10 @@ +from mitim_modules.powertorch.physics_models.transport_cgyro import gyrokinetic_model +from mitim_tools.simulation_tools.physics import GXtools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +class gx_model(gyrokinetic_model): + + def evaluate_turbulence(self): + + self._evaluate_gyrokinetic_model(code = 'gx', gk_object = GXtools.GX, out_name = 'CGXout') diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 150acdc8..fd4af7e7 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -391,8 +391,9 @@ def evaluate_neoclassical(self): from mitim_modules.powertorch.physics_models.transport_tglf import tglf_model from mitim_modules.powertorch.physics_models.transport_neo import neo_model from mitim_modules.powertorch.physics_models.transport_cgyro import cgyro_model +from mitim_modules.powertorch.physics_models.transport_gx import gx_model -class portals_model(power_transport, tglf_model, neo_model, cgyro_model): +class portals_transport_model(power_transport, tglf_model, neo_model, cgyro_model, gx_model): def __init__(self, powerstate, **kwargs): super().__init__(powerstate, **kwargs) @@ -410,6 +411,8 @@ def evaluate_turbulence(self): return tglf_model.evaluate_turbulence(self) elif self.turbulence_model == 'cgyro': return cgyro_model.evaluate_turbulence(self) + elif self.turbulence_model == 'gx': + return gx_model.evaluate_turbulence(self) else: raise Exception(f"Unknown turbulence model {self.turbulence_model}") diff --git a/src/mitim_tools/simulation_tools/physics/GXtools.py b/src/mitim_tools/simulation_tools/physics/GXtools.py index 6c6bc225..391ed72e 100644 --- a/src/mitim_tools/simulation_tools/physics/GXtools.py +++ b/src/mitim_tools/simulation_tools/physics/GXtools.py @@ -1,7 +1,7 @@ import netCDF4 import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools, CONFIGread -from mitim_tools.gacode_tools.utils import GACODErun, GACODEdefaults +from mitim_tools.gacode_tools.utils import GACODEdefaults, CGYROutils from mitim_tools.simulation_tools import SIMtools from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __mitimroot__ @@ -163,8 +163,6 @@ def plot( ax1.legend(loc='best', prop={'size': 4}) ax2.set_ylabel("Growth rate") - - ax3 = fig.add_subplot(grid[0, 1]) ax4 = fig.add_subplot(grid[1, 1]) @@ -358,10 +356,12 @@ def _fmt_value(val): return param_written class GXoutput(SIMtools.GACODEoutput): - def __init__(self, FolderGACODE, suffix="", **kwargs): + def __init__(self, FolderGACODE, suffix="", tmin = 0.0, **kwargs): super().__init__() self.FolderGACODE, self.suffix = FolderGACODE, suffix + + self.tmin = tmin if suffix == "": print(f"\t- Reading results from folder {IOtools.clipstr(FolderGACODE)} without suffix") @@ -396,3 +396,22 @@ def read(self): self.GiAll = G[:,:-1] self.Gi = self.GiAll.sum(axis=1) + self._signal_analysis() + + def _signal_analysis(self): + + flags = [ + 'Qe', + 'Qi', + 'Ge', + ] + + for iflag in flags: + self.__dict__[iflag+'_mean'], self.__dict__[iflag+'_std'] = CGYROutils.apply_ac( + self.t, + self.__dict__[iflag], + tmin=self.tmin, + label_print=iflag, + print_msg=True, + ) + From 8d9cbb58848909c9f1e0b0d5f55dbef5a1a144a9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 28 Aug 2025 22:40:19 -0400 Subject: [PATCH 234/385] TargetType now list of strings --- src/mitim_modules/portals/PORTALSmain.py | 8 +- .../portals/utils/PORTALSanalysis.py | 2 +- .../portals/utils/PORTALSinit.py | 2 +- .../portals/utils/PORTALSplot.py | 16 +-- src/mitim_modules/powertorch/STATEtools.py | 16 +-- .../physics_models/targets_analytic.py | 6 +- .../physics_models/transport_cgyro.py | 18 +-- .../powertorch/physics_models/transport_gx.py | 4 +- .../physics_models/transport_neo.py | 10 +- .../physics_models/transport_tglf.py | 14 +-- .../powertorch/scripts/calculateTargets.py | 8 +- .../powertorch/scripts/compareWithTGYRO.py | 2 +- .../powertorch/utils/TARGETStools.py | 26 +++-- .../powertorch/utils/TRANSFORMtools.py | 21 ++-- src/mitim_modules/vitals/VITALSmain.py | 10 +- src/mitim_tools/gacode_tools/CGYROtools.py | 5 +- src/mitim_tools/gacode_tools/NEOtools.py | 11 +- src/mitim_tools/gacode_tools/TGLFtools.py | 107 +++++++++--------- .../gacode_tools/utils/CGYROutils.py | 2 +- src/mitim_tools/simulation_tools/SIMtools.py | 12 +- .../simulation_tools/physics/GXtools.py | 7 +- src/mitim_tools/transp_tools/CDFtools.py | 22 ++-- templates/maestro_namelist.json | 2 +- 23 files changed, 171 insertions(+), 160 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index c2f1c249..841cb8a1 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -251,10 +251,10 @@ def __init__( self.portals_parameters["target"] = { "evaluator": target_evaluator, "options": { - "TypeTarget": 3, + "targets_evolve": ["qie", "qrad", "qfus"], "target_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) - "forceZeroParticleFlux": False, # If True, ignore particle flux profile and assume zero for all radii - "fineTargetsResolution": 20, # If not None, calculate targets with this radial resolution (defaults target_evaluator_method to powerstate) + "force_zero_particle_flux": False, # If True, ignore particle flux profile and assume zero for all radii + "targets_resolution": 20, # If not None, calculate targets with this radial resolution (defaults target_evaluator_method to powerstate) "percent_error": 1 # (%) Error (std, in percent) of model evaluation } } @@ -572,7 +572,7 @@ def reuseTrainingTabular( self_copy = copy.deepcopy(self) if reevaluateTargets == 1: self_copy.powerstate.transport_options["evaluator"] = None - self_copy.powerstate.target_options["options"]["TypeTarget"] = "powerstate" + self_copy.powerstate.target_options["options"]["targets_evolve"] = "target_evaluator_method" _, dictOFs = runModelEvaluator( self_copy, diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 0915cf69..9ed37409 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -158,7 +158,7 @@ def prep_metrics(self, ilast=None): self.runWithImpurity = self.powerstate.impurityPosition if "nZ" in self.predicted_channels else None self.runWithRotation = "w0" in self.predicted_channels - self.forceZeroParticleFlux = self.portals_parameters["target"]["options"]["forceZeroParticleFlux"] + self.force_zero_particle_flux = self.portals_parameters["target"]["options"]["force_zero_particle_flux"] # Profiles and tgyro results print("\t- Reading profiles and tgyros for each evaluation") diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index d3b9294b..a02a2cbb 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -105,7 +105,7 @@ def initializeProblem( # Print warning or question to be careful! if len(speciesNotFound) > 0: - if portals_fun.portals_parameters["target"]["TypeTarget"] == 3: + if "qrad" in portals_fun.portals_parameters["target"]["targets_evolve"]: answerYN = print(f"\t- Species {speciesNotFound} not found in radiation database, radiation will be zero in PORTALS... is this ok for your predictions?",typeMsg="q" if checkForSpecies else "w") if checkForSpecies and (not answerYN): diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index 1dd74f5c..ee3e71f7 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -246,7 +246,7 @@ def PORTALSanalyzer_plotMetrics( "-", c=col, lw=lwt, alpha=alph) axne_f.plot( rho, - power.plasma['Ge1E20m2'].cpu().numpy() * (1 - int(self.forceZeroParticleFlux)), + power.plasma['Ge1E20m2'].cpu().numpy() * (1 - int(self.force_zero_particle_flux)), "--", c=col, lw=lwt, @@ -382,7 +382,7 @@ def PORTALSanalyzer_plotMetrics( col=col, lab=lab, msFlux=msFlux, - forceZeroParticleFlux=self.forceZeroParticleFlux, + force_zero_particle_flux=self.force_zero_particle_flux, maxStore=indexToMaximize == indexUse, decor=self.ibest == indexUse, plotFlows=plotFlows and (self.ibest == indexUse), @@ -2101,7 +2101,7 @@ def PORTALSanalyzer_plotModelComparison( else: val_calc = np.array( [ - self.tglf_full.results["ev0"]["TGLFout"][j].__dict__[ + self.tglf_full.results["ev0"]["output"][j].__dict__[ quantityX.replace("[TGLF]", "") ] for j in range(len(self.rhos)) @@ -2151,7 +2151,7 @@ def PORTALSanalyzer_plotModelComparison( else: val_calc = np.array( [ - self.tglf_full.results["ev0"]["TGLFout"][j].__dict__[ + self.tglf_full.results["ev0"]["output"][j].__dict__[ quantityX.replace("[TGLF]", "") ] for j in range(len(self.rhos)) @@ -2291,7 +2291,7 @@ def plotModelComparison_quantity( if "[TGLF]" in quantityX: X.append( [ - self.tglf_full.results[f"ev{i}"]["TGLFout"][j].__dict__[ + self.tglf_full.results[f"ev{i}"]["output"][j].__dict__[ quantityX.replace("[TGLF]", "") ] for j in range(len(self.rhos)) @@ -2835,7 +2835,7 @@ def plotFluxComparison( axne_f, axnZ_f, axw0_f, - forceZeroParticleFlux=False, + force_zero_particle_flux=False, runWithImpurity=3, labZ="Z", includeFirst=True, @@ -3012,7 +3012,7 @@ def plotFluxComparison( Qe_tar = power.plasma['QeMWm2'].cpu().numpy()[0][ixF:] Qi_tar = power.plasma['QiMWm2'].cpu().numpy()[0][ixF:] - Ge_tar = power.plasma['Ge1E20m2'].cpu().numpy()[0][ixF:] * (1-int(forceZeroParticleFlux)) + Ge_tar = power.plasma['Ge1E20m2'].cpu().numpy()[0][ixF:] * (1-int(force_zero_particle_flux)) GZ_tar = power.plasma['GZ1E20m2'].cpu().numpy()[0][ixF:] Mt_tar = power.plasma['MtJm2'].cpu().numpy()[0][ixF:] @@ -3112,7 +3112,7 @@ def plotFluxComparison( y = tBest.derived[var] * mult if var == "ge_10E20m2": - y *= 1 - int(forceZeroParticleFlux) + y *= 1 - int(force_zero_particle_flux) ax.plot( (tBest.profiles["rho(-)"] if not useRoa else tBest.derived["roa"]), diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 1f9d97d7..6ba01348 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -57,12 +57,12 @@ def __init__( target_options = { "evaluator": targets_analytic.analytical_model, "options": { - "TypeTarget": 3, + "targets_evolve": ["qie", "qrad", "qfus"], "target_evaluator_method": "powerstate", }, } - target_options["options"].setdefault("forceZeroParticleFlux", False) + target_options["options"].setdefault("force_zero_particle_flux", False) target_options["options"].setdefault("percent_error", 1.0) if tensor_options is None: @@ -84,7 +84,7 @@ def __init__( self.predicted_channels = evolution_options.get("ProfilePredicted", ["te", "ti", "ne"]) self.impurityPosition = evolution_options.get("impurityPosition", 1) self.impurityPosition_transport = copy.deepcopy(self.impurityPosition) - self.fineTargetsResolution = target_options.get("fineTargetsResolution", None) + self.targets_resolution = target_options.get("targets_resolution", None) self.scaleIonDensities = evolution_options.get("scaleIonDensities", True) self.fImp_orig = evolution_options.get("fImp_orig", 1.0) rho_vec = evolution_options.get("rhoPredicted", [0.2, 0.4, 0.6, 0.8]) @@ -159,7 +159,7 @@ def _ensure_ne_before_nz(lst): # Fine targets (need to do it here so that it's only once per definition of powerstate) # ------------------------------------------------------------------------------------- - if self.fineTargetsResolution is None: + if self.targets_resolution is None: self.plasma_fine, self.positions_targets = None, None else: self._fine_grid() @@ -184,7 +184,7 @@ def _ensure_ne_before_nz(lst): def _high_res_rho(self): rho_new = torch.linspace( - self.plasma["rho"][0], self.plasma["rho"][-1], self.fineTargetsResolution + self.plasma["rho"][0], self.plasma["rho"][-1], self.targets_resolution ).to(self.plasma["rho"]) for i in self.plasma["rho"]: if not torch.isclose( @@ -698,7 +698,7 @@ def calculateTargets(self, relative_error_assumed=1.0): targets = self.target_options["evaluator"](self) # [Optional] Calculate local targets and integrals on a fine grid - if self.fineTargetsResolution is not None: + if self.targets_resolution is not None: targets.fine_grid() # Evaluate local quantities @@ -708,13 +708,13 @@ def calculateTargets(self, relative_error_assumed=1.0): targets.flux_integrate() # Come back to original grid - if self.fineTargetsResolution is not None: + if self.targets_resolution is not None: targets.coarse_grid() # Merge targets, calculate errors and normalize targets.postprocessing( relative_error_assumed=relative_error_assumed, - forceZeroParticleFlux=self.target_options["options"]["forceZeroParticleFlux"] + force_zero_particle_flux=self.target_options["options"]["force_zero_particle_flux"] ) def calculateTransport( diff --git a/src/mitim_modules/powertorch/physics_models/targets_analytic.py b/src/mitim_modules/powertorch/physics_models/targets_analytic.py index d2534282..45408678 100644 --- a/src/mitim_modules/powertorch/physics_models/targets_analytic.py +++ b/src/mitim_modules/powertorch/physics_models/targets_analytic.py @@ -28,11 +28,13 @@ def __init__(self,powerstate, **kwargs): def evaluate(self): - if self.powerstate.target_options["options"]["TypeTarget"] >= 2: + if "qie" in self.powerstate.target_options["options"]["targets_evolve"]: self._evaluate_energy_exchange() - if self.powerstate.target_options["options"]["TypeTarget"] == 3: + if "qfus" in self.powerstate.target_options["options"]["targets_evolve"]: self._evaluate_alpha_heating() + + if "qrad" in self.powerstate.target_options["options"]["targets_evolve"]: self._evaluate_radiation() def _evaluate_energy_exchange(self): diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 7fd69c59..f98688eb 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -6,7 +6,7 @@ class gyrokinetic_model: - def _evaluate_gyrokinetic_model(self, code = 'cgyro', gk_object = None, out_name = 'CGYROout'): + def _evaluate_gyrokinetic_model(self, code = 'cgyro', gk_object = None): # ------------------------------------------------------------------------------------------------------------------------ # Grab options # ------------------------------------------------------------------------------------------------------------------------ @@ -51,15 +51,15 @@ def _evaluate_gyrokinetic_model(self, code = 'cgyro', gk_object = None, out_name # ------------------------------------------------------------------------------------------------------------------------ # Pass the information to what power_transport expects # ------------------------------------------------------------------------------------------------------------------------ - - self.QeGB_turb = np.array([gk_object.results[subfolder_name][out_name][i].Qe_mean for i in range(len(rho_locations))]) - self.QeGB_turb_stds = np.array([gk_object.results[subfolder_name][out_name][i].Qe_std for i in range(len(rho_locations))]) + + self.QeGB_turb = np.array([gk_object.results[subfolder_name]['output'][i].Qe_mean for i in range(len(rho_locations))]) + self.QeGB_turb_stds = np.array([gk_object.results[subfolder_name]['output'][i].Qe_std for i in range(len(rho_locations))]) - self.QiGB_turb = np.array([gk_object.results[subfolder_name][out_name][i].Qi_mean for i in range(len(rho_locations))]) - self.QiGB_turb_stds = np.array([gk_object.results[subfolder_name][out_name][i].Qi_std for i in range(len(rho_locations))]) + self.QiGB_turb = np.array([gk_object.results[subfolder_name]['output'][i].Qi_mean for i in range(len(rho_locations))]) + self.QiGB_turb_stds = np.array([gk_object.results[subfolder_name]['output'][i].Qi_std for i in range(len(rho_locations))]) - self.GeGB_turb = np.array([gk_object.results[subfolder_name][out_name][i].Ge_mean for i in range(len(rho_locations))]) - self.GeGB_turb_stds = np.array([gk_object.results[subfolder_name][out_name][i].Ge_std for i in range(len(rho_locations))]) + self.GeGB_turb = np.array([gk_object.results[subfolder_name]['output'][i].Ge_mean for i in range(len(rho_locations))]) + self.GeGB_turb_stds = np.array([gk_object.results[subfolder_name]['output'][i].Ge_std for i in range(len(rho_locations))]) self.GZGB_turb = self.QeGB_turb*0.0 #TODO self.GZGB_turb_stds = self.QeGB_turb*0.0 #TODO @@ -129,7 +129,7 @@ def evaluate_turbulence(self): self._evaluate_tglf() # -------------------------------------------------------------------------------------------- - self._evaluate_gyrokinetic_model(code = 'cgyro', gk_object = CGYROtools.CGYRO, out_name = 'CGYROout') + self._evaluate_gyrokinetic_model(code = 'cgyro', gk_object = CGYROtools.CGYRO, out_name = 'output') def pre_checks(self): diff --git a/src/mitim_modules/powertorch/physics_models/transport_gx.py b/src/mitim_modules/powertorch/physics_models/transport_gx.py index d5bd565e..fb9d6a80 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_gx.py +++ b/src/mitim_modules/powertorch/physics_models/transport_gx.py @@ -4,7 +4,5 @@ from IPython import embed class gx_model(gyrokinetic_model): - def evaluate_turbulence(self): - - self._evaluate_gyrokinetic_model(code = 'gx', gk_object = GXtools.GX, out_name = 'CGXout') + self._evaluate_gyrokinetic_model(code = 'gx', gk_object = GXtools.GX) diff --git a/src/mitim_modules/powertorch/physics_models/transport_neo.py b/src/mitim_modules/powertorch/physics_models/transport_neo.py index 7fe2252f..9411f92b 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_neo.py +++ b/src/mitim_modules/powertorch/physics_models/transport_neo.py @@ -42,11 +42,11 @@ def evaluate_neoclassical(self): label='base', **simulation_options["read"]) - Qe = np.array([neo.results['base']['NEOout'][i].Qe for i in range(len(rho_locations))]) - Qi = np.array([neo.results['base']['NEOout'][i].Qi for i in range(len(rho_locations))]) - Ge = np.array([neo.results['base']['NEOout'][i].Ge for i in range(len(rho_locations))]) - GZ = np.array([neo.results['base']['NEOout'][i].GiAll[impurityPosition-1] for i in range(len(rho_locations))]) - Mt = np.array([neo.results['base']['NEOout'][i].Mt for i in range(len(rho_locations))]) + Qe = np.array([neo.results['base']['output'][i].Qe for i in range(len(rho_locations))]) + Qi = np.array([neo.results['base']['output'][i].Qi for i in range(len(rho_locations))]) + Ge = np.array([neo.results['base']['output'][i].Ge for i in range(len(rho_locations))]) + GZ = np.array([neo.results['base']['output'][i].GiAll[impurityPosition-1] for i in range(len(rho_locations))]) + Mt = np.array([neo.results['base']['output'][i].Mt for i in range(len(rho_locations))]) # ------------------------------------------------------------------------------------------------------------------------ # Pass the information to what power_transport expects diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index fe322ad9..5e398bce 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -69,16 +69,16 @@ def _evaluate_tglf(self): **simulation_options["read"]) # Grab values - Qe = np.array([tglf.results['base']['TGLFout'][i].Qe for i in range(len(rho_locations))]) - Qi = np.array([tglf.results['base']['TGLFout'][i].Qi for i in range(len(rho_locations))]) - Ge = np.array([tglf.results['base']['TGLFout'][i].Ge for i in range(len(rho_locations))]) - GZ = np.array([tglf.results['base']['TGLFout'][i].GiAll[impurityPosition] for i in range(len(rho_locations))]) - Mt = np.array([tglf.results['base']['TGLFout'][i].Mt for i in range(len(rho_locations))]) - S = np.array([tglf.results['base']['TGLFout'][i].Se for i in range(len(rho_locations))]) + Qe = np.array([tglf.results['base']['output'][i].Qe for i in range(len(rho_locations))]) + Qi = np.array([tglf.results['base']['output'][i].Qi for i in range(len(rho_locations))]) + Ge = np.array([tglf.results['base']['output'][i].Ge for i in range(len(rho_locations))]) + GZ = np.array([tglf.results['base']['output'][i].GiAll[impurityPosition] for i in range(len(rho_locations))]) + Mt = np.array([tglf.results['base']['output'][i].Mt for i in range(len(rho_locations))]) + S = np.array([tglf.results['base']['output'][i].Se for i in range(len(rho_locations))]) if Qi_includes_fast: - Qifast = [tglf.results['base']['TGLFout'][i].Qifast for i in range(len(rho_locations))] + Qifast = [tglf.results['base']['output'][i].Qifast for i in range(len(rho_locations))] if Qifast.sum() != 0.0: print(f"\t- Qi includes fast ions, adding their contribution") diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index 6b7a03b1..e4d31039 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -12,12 +12,12 @@ def calculator( input_gacode, - TypeTarget=3, + targets_evolve=["qie", "qrad", "qfus"], folder="~/scratch/", cold_start=True, rho_vec=np.linspace(0.1, 0.9, 9), profProvided=False, - fineTargetsResolution = None, + targets_resolution = None, ): profiles = input_gacode if profProvided else PROFILEStools.gacode_state(input_gacode) @@ -29,9 +29,9 @@ def calculator( target_options={ "evaluator": targets_analytic.analytical_model, "options": { - "TypeTarget": TypeTarget, + "targets_evolve": targets_evolve, "target_evaluator_method": "powerstate", - "fineTargetsResolution": fineTargetsResolution + "targets_resolution": targets_resolution }, }, transport_options={ diff --git a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py index 875e1299..011cc2d2 100644 --- a/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py +++ b/src/mitim_modules/powertorch/scripts/compareWithTGYRO.py @@ -30,7 +30,7 @@ # STATE s = STATEtools.powerstate(t.profiles, evolution_options={"rhoPredicted": t.rho[0,1:]}) s.calculateProfileFunctions() -# s.target_options['options']['TypeTarget'] = 1 +# s.target_options['options']['targets_evolve'] = 1 s.calculateTargets() # diff --git a/src/mitim_modules/powertorch/utils/TARGETStools.py b/src/mitim_modules/powertorch/utils/TARGETStools.py index 649ae0ed..8d426864 100644 --- a/src/mitim_modules/powertorch/utils/TARGETStools.py +++ b/src/mitim_modules/powertorch/utils/TARGETStools.py @@ -71,14 +71,17 @@ def flux_integrate(self): qe = self.powerstate.plasma["te"]*0.0 qi = self.powerstate.plasma["te"]*0.0 - if self.powerstate.target_options['options']['TypeTarget'] >= 2: + if "qie" in self.powerstate.target_options['options']['targets_evolve']: qe += -self.powerstate.plasma["qie"] qi += self.powerstate.plasma["qie"] - - if self.powerstate.target_options['options']['TypeTarget'] == 3: - qe += self.powerstate.plasma["qfuse"] - self.powerstate.plasma["qrad"] + + if "qfus" in self.powerstate.target_options['options']['targets_evolve']: + qe += self.powerstate.plasma["qfuse"] qi += self.powerstate.plasma["qfusi"] + if "qrad" in self.powerstate.target_options['options']['targets_evolve']: + qe -= self.powerstate.plasma["qrad"] + q = torch.cat((qe, qi)).to(qe) self.P = self.powerstate.from_density_to_flux(q, force_dim=q.shape[0]) @@ -89,14 +92,19 @@ def coarse_grid(self): # ************************************************************************************************** # Interpolate results from fine to coarse (i.e. whole point is that it is better than integrate interpolated values) - if self.powerstate.target_options['options']['TypeTarget'] >= 2: + if "qie" in self.powerstate.target_options['options']['targets_evolve']: for i in ["qie"]: self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] - if self.powerstate.target_options['options']['TypeTarget'] == 3: + if "qfus" in self.powerstate.target_options['options']['targets_evolve']: for i in [ "qfuse", "qfusi", + ]: + self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] + + if "qrad" in self.powerstate.target_options['options']['targets_evolve']: + for i in [ "qrad", "qrad_bremms", "qrad_line", @@ -104,13 +112,15 @@ def coarse_grid(self): ]: self.powerstate.plasma[i] = self.powerstate.plasma[i][:, self.powerstate.positions_targets] + + self.P = self.P[:, self.powerstate.positions_targets] # Recover variables calculated prior to the fine-targets method for i in self.plasma_original: self.powerstate.plasma[i] = self.plasma_original[i] - def postprocessing(self, forceZeroParticleFlux=False, relative_error_assumed=1.0): + def postprocessing(self, force_zero_particle_flux=False, relative_error_assumed=1.0): # ************************************************************************************************** # Plug-in targets that were fixed @@ -122,7 +132,7 @@ def postprocessing(self, forceZeroParticleFlux=False, relative_error_assumed=1.0 self.powerstate.plasma["GZ1E20m2"] = self.powerstate.plasma["GZ_fixedtargets"] # 1E20/s/m^2 self.powerstate.plasma["MtJm2"] = self.powerstate.plasma["MtJm2_fixedtargets"] # J/m^2 - if forceZeroParticleFlux: + if force_zero_particle_flux: self.powerstate.plasma["Ge1E20m2"] = self.powerstate.plasma["Ge1E20m2"] * 0 # Convective fluxes diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index edd487a7..f71e3195 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -120,12 +120,16 @@ def gacode_to_powerstate(self, rho_vec=None): quantitites["GZ_fixedtargets"] = input_gacode.derived["ge_10E20"] * 0.0 quantitites["MtJm2_fixedtargets"] = input_gacode.derived["mt_Jmiller"] - if self.target_options["options"]["TypeTarget"] < 3: - # Fusion and radiation fixed if 1,2 - quantitites["QeMWm2_fixedtargets"] += input_gacode.derived["qe_fus_MW"] - input_gacode.derived["qrad_MW"] + if 'qfus' not in self.target_options["options"]["targets_evolve"]: + # Fusion fixed + quantitites["QeMWm2_fixedtargets"] += input_gacode.derived["qe_fus_MW"] quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qi_fus_MW"] + + if 'qrad' not in self.target_options["options"]["targets_evolve"]: + # Fusion fixed + quantitites["QeMWm2_fixedtargets"] -= input_gacode.derived["qrad_MW"] - if self.target_options["options"]["TypeTarget"] < 2: + if 'qie' not in self.target_options["options"]["targets_evolve"]: # Exchange fixed if 1 quantitites["QeMWm2_fixedtargets"] -= input_gacode.derived["qe_exc_MW"] quantitites["QiMWm2_fixedtargets"] += input_gacode.derived["qe_exc_MW"] @@ -367,9 +371,9 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): target_options={ "evaluator": targets_analytic.analytical_model, "options": { - "TypeTarget": self.target_options["options"]["TypeTarget"], # Important to keep the same as in the original + "targets_evolve": self.target_options["options"]["targets_evolve"], # Important to keep the same as in the original "target_evaluator_method": "powerstate", - "forceZeroParticleFlux": self.target_options["options"]["forceZeroParticleFlux"], + "force_zero_particle_flux": self.target_options["options"]["force_zero_particle_flux"], "percent_error": self.target_options["options"]["percent_error"] } }, @@ -382,12 +386,13 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): conversions = {} - if self.target_options["options"]["TypeTarget"] > 1: + if 'qie' in self.target_options["options"]["targets_evolve"]: conversions['qie'] = "qei(MW/m^3)" - if self.target_options["options"]["TypeTarget"] > 2: + if 'qrad' in self.target_options["options"]["targets_evolve"]: conversions['qrad_bremms'] = "qbrem(MW/m^3)" conversions['qrad_sync'] = "qsync(MW/m^3)" conversions['qrad_line'] = "qline(MW/m^3)" + if 'qfus' in self.target_options["options"]["targets_evolve"]: conversions['qfuse'] = "qfuse(MW/m^3)" conversions['qfusi'] = "qfusi(MW/m^3)" diff --git a/src/mitim_modules/vitals/VITALSmain.py b/src/mitim_modules/vitals/VITALSmain.py index 11fc6fe1..2ada4f9d 100644 --- a/src/mitim_modules/vitals/VITALSmain.py +++ b/src/mitim_modules/vitals/VITALSmain.py @@ -199,19 +199,19 @@ def run(self, paramsfile, resultsfile): for iquant in dictOFs: if "_exp" not in iquant: if iquant == "Qe": - value = tglf.results["tglf1"]["TGLFout"][0].Qe_unn + value = tglf.results["tglf1"]["output"][0].Qe_unn elif iquant == "Qi": - value = tglf.results["tglf1"]["TGLFout"][0].Qi_unn + value = tglf.results["tglf1"]["output"][0].Qi_unn elif iquant == "TeFluct": - value = tglf.results["tglf1"]["TGLFout"][ + value = tglf.results["tglf1"]["output"][ 0 ].AmplitudeSpectrum_Te_level elif iquant == "neFluct": - value = tglf.results["tglf1"]["TGLFout"][ + value = tglf.results["tglf1"]["output"][ 0 ].AmplitudeSpectrum_ne_level elif iquant == "neTe": - value = tglf.results["tglf1"]["TGLFout"][0].neTeSpectrum_level + value = tglf.results["tglf1"]["output"][0].neTeSpectrum_level dictOFs[iquant]["value"] = value dictOFs[iquant]["error"] = np.abs( diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index bbe9e0a1..b114398c 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -61,8 +61,7 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call 'input_class': CGYROinput, 'complete_variation': None, 'default_cores': 16, # Default cores to use in the simulation - 'output_class': CGYROutils.CGYROout, - 'output_store': 'CGYROout' + 'output_class': CGYROutils.CGYROoutput, } print("\n-----------------------------------------------------------------------------------------") @@ -138,7 +137,7 @@ def _labelize(self, labels): for label in labels: for i,rho in enumerate(self.rhos): labels_with_rho.append(f"{label}_{rho}") - self.results[f'{label}_{rho}'] = self.results_all[label]['CGYROout'][i] + self.results[f'{label}_{rho}'] = self.results_all[label]['output'][i] labels = labels_with_rho # ------------------------------------------------ diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index c32b56b2..7db02cfc 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -51,7 +51,6 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call 'complete_variation': None, 'default_cores': 1, # Default cores to use in the simulation 'output_class': NEOoutput, - 'output_store': 'NEOout' } print("\n-----------------------------------------------------------------------------------------") @@ -88,10 +87,10 @@ def plot( for i,label in enumerate(labels): roa, QeGB, QiGB, GeGB = [], [], [], [] for irho in range(len(self.rhos)): - roa.append(self.results[label]['NEOout'][irho].roa) - QeGB.append(self.results[label]['NEOout'][irho].Qe) - QiGB.append(self.results[label]['NEOout'][irho].Qi) - GeGB.append(self.results[label]['NEOout'][irho].Ge) + roa.append(self.results[label]['output'][irho].roa) + QeGB.append(self.results[label]['output'][irho].Qe) + QiGB.append(self.results[label]['output'][irho].Qi) + GeGB.append(self.results[label]['output'][irho].Ge) axQe.plot(roa, QeGB, label=label, color=colors[i], marker='o', linestyle='-') axQi.plot(roa, QiGB, label=label, color=colors[i], marker='o', linestyle='-') @@ -115,7 +114,7 @@ def read_scan( positionIon=2 ): - output_object = "NEOout" + output_object = "output" variable_mapping = { 'scanned_variable': ["parsed", variable, None], diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 06af82fe..498a6b36 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -161,7 +161,6 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call 'complete_variation': completeVariation_TGLF, 'default_cores': 4, # Default cores to use in the simulation 'output_class': TGLFoutput, - 'output_store': 'TGLFout' } print("\n-----------------------------------------------------------------------------------------") @@ -296,9 +295,9 @@ def _run_wf(self, kys, code_executor, forceClosestUnstableWF=True, **kwargs_TGLF # Only unstable ones kys_n = [] - for j in range(len(self.results[f"ky{ky_single0}"]["TGLFout"][i].ky)): - if self.results[f"ky{ky_single0}"]["TGLFout"][i].g[0, j] > 0.0: - kys_n.append(self.results[f"ky{ky_single0}"]["TGLFout"][i].ky[j]) + for j in range(len(self.results[f"ky{ky_single0}"]["output"][i].ky)): + if self.results[f"ky{ky_single0}"]["output"][i].g[0, j] > 0.0: + kys_n.append(self.results[f"ky{ky_single0}"]["output"][i].ky[j]) kys_n = np.array(kys_n) # ---- @@ -732,13 +731,13 @@ def plot( for irho in range(len(self.rhos)): successful_normalization = ( successful_normalization - and self.results[label]["TGLFout"][irho].unnormalization_successful + and self.results[label]["output"][irho].unnormalization_successful ) max_num_species = np.max( - [max_num_species, self.results[label]["TGLFout"][irho].num_species] + [max_num_species, self.results[label]["output"][irho].num_species] ) - for il in self.results[label]["TGLFout"][irho].fields: + for il in self.results[label]["output"][irho].fields: if il not in max_fields: max_fields.append(il) @@ -1006,10 +1005,10 @@ def plot( # -------------------------------- # Plot Raw TGLF (normalized) # -------------------------------- - self.results[label]["TGLFout"][irho].plotTGLF_Summary( + self.results[label]["output"][irho].plotTGLF_Summary( c=colors[cont], label=full_label, axs=axsTGLF1, irho_cont=irho_cont ) - self.results[label]["TGLFout"][irho].plotTGLF_Contributors( + self.results[label]["output"][irho].plotTGLF_Contributors( c=colors[cont], label=full_label, axs=axsTGLF2, @@ -1017,11 +1016,11 @@ def plot( title_legend=title_legend, cont=cont, ) - self.results[label]["TGLFout"][irho].plotTGLF_Model( + self.results[label]["output"][irho].plotTGLF_Model( axs=axsTGLF3, c=colors[cont], label=full_label ) - self.results[label]["TGLFout"][irho].plotTGLF_Fluctuations( + self.results[label]["output"][irho].plotTGLF_Fluctuations( axs=axsTGLF_flucts, c=colors[cont], label=full_label, @@ -1030,7 +1029,7 @@ def plot( cont=cont, ) - self.results[label]["TGLFout"][irho].plotTGLF_Field( + self.results[label]["output"][irho].plotTGLF_Field( quantity="phi", c=colors[cont], label=full_label, @@ -1041,7 +1040,7 @@ def plot( ) if "a_par" in max_fields: - self.results[label]["TGLFout"][irho].plotTGLF_Field( + self.results[label]["output"][irho].plotTGLF_Field( quantity="a_par", c=colors[cont], label=full_label, @@ -1051,7 +1050,7 @@ def plot( cont=cont, ) if "a_per" in max_fields: - self.results[label]["TGLFout"][irho].plotTGLF_Field( + self.results[label]["output"][irho].plotTGLF_Field( quantity="a_per", c=colors[cont], label=full_label, @@ -1084,9 +1083,9 @@ def plot( if successful_normalization: GACODEplotting.plotTGLFspectrum( [axS00, axS10], - self.results[label]["TGLFout"][irho].ky, - self.results[label]["TGLFout"][irho].g[0, :], - freq=self.results[label]["TGLFout"][irho].f[0, :], + self.results[label]["output"][irho].ky, + self.results[label]["output"][irho].g[0, :], + freq=self.results[label]["output"][irho].f[0, :], coeff=0.0, c=colors[cont], ls="-", @@ -1114,22 +1113,22 @@ def plot( for irho_cont in range(len(self.rhos)): irho = np.where(self.results[label]["x"] == self.rhos[irho_cont])[0][0] - if self.results[label]["TGLFout"][irho].unnormalization_successful: - Qe.append(self.results[label]["TGLFout"][irho].Qe_unn) - Qi.append(self.results[label]["TGLFout"][irho].Qi_unn) - Ge.append(self.results[label]["TGLFout"][irho].Ge_unn) + if self.results[label]["output"][irho].unnormalization_successful: + Qe.append(self.results[label]["output"][irho].Qe_unn) + Qi.append(self.results[label]["output"][irho].Qi_unn) + Ge.append(self.results[label]["output"][irho].Ge_unn) TeF.append( - self.results[label]["TGLFout"][irho].AmplitudeSpectrum_Te_level + self.results[label]["output"][irho].AmplitudeSpectrum_Te_level ) - neTe.append(self.results[label]["TGLFout"][irho].neTeSpectrum_level) + neTe.append(self.results[label]["output"][irho].neTeSpectrum_level) - roas.append(self.results[label]["TGLFout"][irho].roa) + roas.append(self.results[label]["output"][irho].roa) - QeGB.append(self.results[label]["TGLFout"][irho].Qe) - QiGB.append(self.results[label]["TGLFout"][irho].Qi) - GeGB.append(self.results[label]["TGLFout"][irho].Ge) + QeGB.append(self.results[label]["output"][irho].Qe) + QiGB.append(self.results[label]["output"][irho].Qi) + GeGB.append(self.results[label]["output"][irho].Ge) - if self.results[label]["TGLFout"][irho].unnormalization_successful: + if self.results[label]["output"][irho].unnormalization_successful: axT2.plot(self.rhos, Qe, "-o", c=colorLab[0], lw=2, label=full_label) axS01.plot(self.rhos, Qe, "-o", c=colorLab[0], lw=2, label=full_label) @@ -1410,7 +1409,7 @@ def plot( a = normalization["rmin"][-1] * 100 rhosa = rho_s / a - kys = self.results[label]["TGLFout"][irho].ky / rho_s + kys = self.results[label]["output"][irho].ky / rho_s xP = np.linspace(0, kys[-1], 1000) @@ -1431,7 +1430,7 @@ def plot( yP = np.ones(len(xP)) ax = axFluc00 - fluct = self.results[label]["TGLFout"][irho].AmplitudeSpectrum_Te + fluct = self.results[label]["output"][irho].AmplitudeSpectrum_Te ylabel = "$A_{T_e}(k_y)$" GACODEplotting.plotTGLFfluctuations( ax, @@ -1455,7 +1454,7 @@ def plot( axFluc00Sym.plot(xP, yP, ls="-.", lw=0.5, color=colors[cont]) ax = axFluc10e - fluct = self.results[label]["TGLFout"][ + fluct = self.results[label]["output"][ irho ].AmplitudeSpectrum_Te * np.interp(kys, xP, yP) ylabel = "$A_{T_e}(k_y)$*W" @@ -1478,7 +1477,7 @@ def plot( ax.plot(kysPlot, fluctPlot, "--", lw=0.3, color=colors[cont]) ax = axFluc01 - fluct = self.results[label]["TGLFout"][irho].AmplitudeSpectrum_ne + fluct = self.results[label]["output"][irho].AmplitudeSpectrum_ne ylabel = "$A_{n_e}(k_y)$" GACODEplotting.plotTGLFfluctuations( ax, @@ -1502,7 +1501,7 @@ def plot( axFluc01Sym.plot(xP, yP, ls="-.", lw=0.5, color=colors[cont]) ax = axFluc11e - fluct = self.results[label]["TGLFout"][ + fluct = self.results[label]["output"][ irho ].AmplitudeSpectrum_ne * np.interp(kys, xP, yP) ylabel = "$A_{n_e}(k_y)$*W" @@ -1527,10 +1526,10 @@ def plot( # --- for inmode in range( - self.results[label]["TGLFout"][irho].num_nmodes + self.results[label]["output"][irho].num_nmodes ): ax = axFluc02 - fluct = self.results[label]["TGLFout"][irho].neTeSpectrum[ + fluct = self.results[label]["output"][irho].neTeSpectrum[ inmode, : ] ylabel = "$n_eT_e(k_y)$" @@ -1568,7 +1567,7 @@ def plot( ) ax = axFluc12e - fluct = self.results[label]["TGLFout"][irho].neTeSpectrum[ + fluct = self.results[label]["output"][irho].neTeSpectrum[ inmode, : ] * np.interp(kys, xP, yP) ylabel = "$n_eT_e(k_y)$*W" @@ -1647,12 +1646,12 @@ def plot( ][0] T.append( - self.results[label]["TGLFout"][irho].AmplitudeSpectrum_Te_level + self.results[label]["output"][irho].AmplitudeSpectrum_Te_level ) N.append( - self.results[label]["TGLFout"][irho].AmplitudeSpectrum_ne_level + self.results[label]["output"][irho].AmplitudeSpectrum_ne_level ) - NT.append(self.results[label]["TGLFout"][irho].neTeSpectrum_level) + NT.append(self.results[label]["output"][irho].neTeSpectrum_level) TL.append(f"{labZX}$\\rho_N={self.rhos[irho_cont]:.4f}$") C.append(colors[cont]) cont += 1 @@ -1823,15 +1822,15 @@ def plot( # all eigenvalues ax00.plot( - self.results[label]["TGLFout"][irho_cont].ky, - self.results[label]["TGLFout"][irho_cont].g[0], + self.results[label]["output"][irho_cont].ky, + self.results[label]["output"][irho_cont].g[0], "-s", markersize=3, color=colors[cont], ) ax10.plot( - self.results[label]["TGLFout"][irho_cont].ky, - self.results[label]["TGLFout"][irho_cont].f[0], + self.results[label]["output"][irho_cont].ky, + self.results[label]["output"][irho_cont].f[0], "-s", markersize=3, color=colors[cont], @@ -1857,15 +1856,15 @@ def plot( # all eigenvalues ax00.plot( - self.results[label]["TGLFout"][irho_cont].ky, - self.results[label]["TGLFout"][irho_cont].g[i + 1], + self.results[label]["output"][irho_cont].ky, + self.results[label]["output"][irho_cont].g[i + 1], "-s", markersize=1, color=colors[cont], ) ax10.plot( - self.results[label]["TGLFout"][irho_cont].ky, - self.results[label]["TGLFout"][irho_cont].f[i + 1], + self.results[label]["output"][irho_cont].ky, + self.results[label]["output"][irho_cont].f[i + 1], "-s", markersize=1, color=colors[cont], @@ -1996,7 +1995,7 @@ def read_scan( positionIon=2 ): - output_object = "TGLFout" + output_object = "output" variable_mapping = { 'scanned_variable': ["parsed", variable, None], @@ -3845,30 +3844,30 @@ def readTGLFresults( for rho in rhos: # Read full folder - TGLFout = TGLFoutput( + output = TGLFoutput( FolderGACODE_tmp, suffix=f"_{rho:.4f}" if suffix is None else suffix, require_all_files=require_all_files, ) # Unnormalize - TGLFout.unnormalize( + output.unnormalize( NormalizationSets["SELECTED"], rho=rho, convolution_fun_fluct=convolution_fun_fluct, factorTot_to_Perp=factorTot_to_Perp, ) - TGLFstd_TGLFout.append(TGLFout) - inputclasses.append(TGLFout.inputclass) + TGLFstd_TGLFout.append(output) + inputclasses.append(output.inputclass) - parse = SIMtools.buildDictFromInput(TGLFout.inputFile) + parse = SIMtools.buildDictFromInput(output.inputFile) parsed.append(parse) results = { "inputclasses": inputclasses, "parsed": parsed, - "TGLFout": TGLFstd_TGLFout, + "output": TGLFstd_TGLFout, "x": np.array(rhos), } diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index ca480e39..96100ca4 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -50,7 +50,7 @@ def __init__(self, labels, cgyro_data): self.Qe_mean = np.array(self.Qe_mean) self.Qi_mean = np.array(self.Qi_mean) -class CGYROout(SIMtools.GACODEoutput): +class CGYROoutput(SIMtools.GACODEoutput): def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for_linear=True, **kwargs): super().__init__() diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 5c4caf59..7aefa6f7 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -352,8 +352,8 @@ def _run( tmpFolder = self.FolderGACODE / f"tmp_{code}" IOtools.askNewFolder(tmpFolder, force=True) - - kkeys = [keys.strip('/') for keys in code_executor.keys()] + + kkeys = [str(keys).replace('/','') for keys in code_executor.keys()] log_simulation_file=self.FolderGACODE / f"mitim_simulation_{kkeys[0]}.log" # Refer with the first folder self.simulation_job = FARMINGtools.mitim_job(tmpFolder, log_simulation_file=log_simulation_file) @@ -808,20 +808,20 @@ def read( ): print("> Reading simulation results") - class_output = [self.run_specifications['output_class'], self.run_specifications['output_store']] + class_output = self.run_specifications['output_class'] # If no specified folder, check the last one if folder is None: folder = self.FolderSimLast self.results[label] = { - class_output[1]:[], + 'output':[], 'parsed': [], "x": np.array(self.rhos), } for rho in self.rhos: - SIMout = class_output[0]( + SIMout = class_output( folder, suffix=f"_{rho:.4f}" if suffix is None else suffix, **kwargs_to_class_output @@ -836,7 +836,7 @@ def read( else: print("No normalization sets found.") - self.results[label][class_output[1]].append(SIMout) + self.results[label]['output'].append(SIMout) self.results[label]['parsed'].append(buildDictFromInput(SIMout.inputFile) if SIMout.inputFile else None) diff --git a/src/mitim_tools/simulation_tools/physics/GXtools.py b/src/mitim_tools/simulation_tools/physics/GXtools.py index 391ed72e..3d8b84a1 100644 --- a/src/mitim_tools/simulation_tools/physics/GXtools.py +++ b/src/mitim_tools/simulation_tools/physics/GXtools.py @@ -58,7 +58,6 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call 'complete_variation': None, 'default_cores': 4, # Default gpus to use in the simulation 'output_class': GXoutput, - 'output_store': 'GXout' } print("\n-----------------------------------------------------------------------------------------") @@ -104,7 +103,7 @@ def plot( i = 0 for label in labels: for irho in range(len(self.rhos)): - c = self.results[label]['GXout'][irho] + c = self.results[label]['output'][irho] typeLs = '-' if c.t.shape[0]>20 else '-s' @@ -146,7 +145,7 @@ def plot( i = 0 for label in labels: for irho in range(len(self.rhos)): - c = self.results[label]['GXout'][irho] + c = self.results[label]['output'][irho] typeLs = '-' if c.t.shape[0]>20 else '-s' @@ -169,7 +168,7 @@ def plot( i = 0 for label in labels: for irho in range(len(self.rhos)): - c = self.results[label]['GXout'][irho] + c = self.results[label]['output'][irho] ax3.plot(c.ky, c.w[-1, :], '-s', markersize = 5, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) ax4.plot(c.ky, c.g[-1, :], '-s', markersize = 5, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) i += 1 diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index 02910ec4..cb162e76 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -14222,13 +14222,13 @@ def plotStdTRANSP(self, tglfRun="tglf1", fig=None, label="", time=None, leg=True TGLFstd_x = self.TGLFstd[int(time * 1000)].results[tglfRun]["x"] TGLFstd_Qe, TGLFstd_Qi = [], [] for i in range( - len(self.TGLFstd[int(time * 1000)].results[tglfRun]["TGLFout"]) + len(self.TGLFstd[int(time * 1000)].results[tglfRun]["output"]) ): TGLFstd_Qe.append( - self.TGLFstd[int(time * 1000)].results[tglfRun]["TGLFout"][i].Qe_unn + self.TGLFstd[int(time * 1000)].results[tglfRun]["output"][i].Qe_unn ) TGLFstd_Qi.append( - self.TGLFstd[int(time * 1000)].results[tglfRun]["TGLFout"][i].Qi_unn + self.TGLFstd[int(time * 1000)].results[tglfRun]["output"][i].Qi_unn ) TGLFstd_Qe, TGLFstd_Qi = np.array(TGLFstd_Qe), np.array(TGLFstd_Qi) @@ -14361,16 +14361,16 @@ def plotGRTRANSP(self, tglfRun="tglf1", fig=None, label="", time=None): else: TGLFstd_ky, TGLFstd_gamma, TGLFstd_freq = [], [], [] for i in range( - len(self.TGLFstd[int(time * 1000)].results[tglfRun]["TGLFout"]) + len(self.TGLFstd[int(time * 1000)].results[tglfRun]["output"]) ): TGLFstd_ky.append( - self.TGLFstd[int(time * 1000)].results[tglfRun]["TGLFout"][i].ky + self.TGLFstd[int(time * 1000)].results[tglfRun]["output"][i].ky ) TGLFstd_gamma.append( - self.TGLFstd[int(time * 1000)].results[tglfRun]["TGLFout"][i].g[0] + self.TGLFstd[int(time * 1000)].results[tglfRun]["output"][i].g[0] ) TGLFstd_freq.append( - self.TGLFstd[int(time * 1000)].results[tglfRun]["TGLFout"][i].f[0] + self.TGLFstd[int(time * 1000)].results[tglfRun]["output"][i].f[0] ) TGLFstd_ky, TGLFstd_gamma, TGLFstd_freq = ( @@ -14518,19 +14518,19 @@ def plotFLTRANSP(self, tglfRun="tglf1", fig=None, label="", time=None): else: TGLFstd_ky, TGLFstd_te, TGLFstd_ne = [], [], [] for i in range( - len(self.TGLFstd[int(time * 1000)].results[tglfRun]["TGLFout"]) + len(self.TGLFstd[int(time * 1000)].results[tglfRun]["output"]) ): TGLFstd_ky.append( - self.TGLFstd[int(time * 1000)].results[tglfRun]["TGLFout"][i].ky + self.TGLFstd[int(time * 1000)].results[tglfRun]["output"][i].ky ) TGLFstd_te.append( self.TGLFstd[int(time * 1000)] - .results[tglfRun]["TGLFout"][i] + .results[tglfRun]["output"][i] .AmplitudeSpectrum_Te ) TGLFstd_ne.append( self.TGLFstd[int(time * 1000)] - .results[tglfRun]["TGLFout"][i] + .results[tglfRun]["output"][i] .AmplitudeSpectrum_ne ) diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index 99bd23f9..6501a4aa 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -62,7 +62,7 @@ "use_default": false, "portals_namelist" : { "portals_parameters": { - "forceZeroParticleFlux": true, + "force_zero_particle_flux": true, "keep_full_model_folder": false, "cores_per_tglf_instance": 1 }, From 429f4a10c80170617a4b8498909d610bba73e073 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 08:37:29 -0400 Subject: [PATCH 235/385] Removal of launchEvaluationsAsSlurmJobs because I let config handle that --- src/mitim_modules/portals/PORTALSmain.py | 51 +++++++++---------- src/mitim_modules/portals/PORTALStools.py | 20 ++++---- .../portals/utils/PORTALSinit.py | 22 ++++---- .../physics_models/transport_tglf.py | 2 - .../powertorch/scripts/calculateTargets.py | 2 +- .../powertorch/utils/TRANSFORMtools.py | 7 ++- .../plasmastate_tools/MITIMstate.py | 6 +-- src/mitim_tools/simulation_tools/SIMtools.py | 2 +- templates/maestro_namelist.json | 2 +- 9 files changed, 54 insertions(+), 60 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 841cb8a1..8890c877 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -7,6 +7,7 @@ from collections import OrderedDict from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools +from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_modules.portals import PORTALStools from mitim_modules.portals.utils import ( PORTALSinit, @@ -147,9 +148,9 @@ def __init__( "portals_transformation_variables": portals_transformation_variables, # Physics-informed parameters to fit surrogates "portals_transformation_variables_trace": portals_transformation_variables_trace, # Physics-informed parameters to fit surrogates for trace impurities "turbulent_exchange_as_surrogate": False, # Run turbulent exchange as surrogate? - "Pseudo_multipliers": [1.0]*5, # [Qe,Qi,Ge] multipliers to calculate pseudo + "scalar_multipliers": [1.0]*5, # [Qe,Qi,Ge] multipliers to calculate pseudo "impurity_trick": True, # If True, fit model to GZ/nZ, valid on the trace limit - "fZ0_as_weight": 1.0, # If not None, using fZ0_as_weight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis + "fZ0_as_weight": None, # If not None, using fZ0_as_weight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis "fImp_orig": 1.0, "additional_params_in_surrogate": additional_params_in_surrogate, "keep_full_model_folder": True, # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) @@ -161,8 +162,6 @@ def __init__( These parameters are communicated to the powertorch object. ''' - from mitim_modules.powertorch.utils.TRANSPORTtools import portals_transport_model as transport_evaluator - # ------------------------- # Transport model settings # ------------------------- @@ -170,7 +169,7 @@ def __init__( # Transport model class - "evaluator": transport_evaluator, + "evaluator": TRANSPORTtools.portals_transport_model, "evaluator_instance_attributes": { "turbulence_model": transport_models[0], @@ -191,7 +190,6 @@ def __init__( "use_scan_trick_for_stds": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta "keep_files": "base", # minimal: only retrieve minimal files always; base: retrieve all files when running TGLF base; none: always minimal files "cores_per_tglf_instance": 1, # Number of cores to use per TGLF instance - "launchEvaluationsAsSlurmJobs": True, # Launch each evaluation as a batch job (vs just comand line) "percent_error": 5, # (%) Error (std, in percent) of model evaluation TGLF if not scan trick "Qi_includes_fast": False, # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) }, @@ -238,7 +236,7 @@ def __init__( "ni_thermals": True, # Adjust for quasineutrality by modifying the thermal ion densities together with ne "recalculate_ptot": True, # Recompute PTOT to insert in input file each time "Tfast_ratio": False, # Keep the ratio of Tfast/Te constant throughout the Te evolution - "ensureMachNumber": None, # Change w0 to match this Mach number when Ti varies + "force_mach": None, # Change w0 to match this Mach number when Ti varies }, } @@ -271,20 +269,19 @@ def _grab_flags_dictionary(d): def prep( self, - fileGACODE, + mitim_state, cold_start=False, ymax_rel=1.0, ymin_rel=1.0, - limitsAreRelative=True, - dvs_fixed=None, - hardGradientLimits=None, - enforceFiniteTemperatureGradients=None, + limits_are_relative=True, + fixed_gradients=None, + yminymax_atleast=None, + enforce_finite_aLT=None, define_ranges_from_profiles=None, start_from_folder=None, - reevaluateTargets=0, + reevaluate_targets=0, seedInitial=None, askQuestions=True, - transport_evaluator_options=None, ): """ Notes: @@ -294,9 +291,9 @@ def prep( 'ti': [0.5, 0.5, 0.5, 0.5], 'ne': [1.0, 0.5, 0.5, 0.5] } - - enforceFiniteTemperatureGradients is used to be able to select ymin_rel = 2.0 for ne but ensure that te, ti is at, e.g., enforceFiniteTemperatureGradients = 0.95 + - enforce_finite_aLT is used to be able to select ymin_rel = 2.0 for ne but ensure that te, ti is at, e.g., enforce_finite_aLT = 0.95 - start_from_folder is a folder from which to grab optimization_data and optimization_extra - (if used with reevaluateTargets>0, change targets by reevaluating with different parameters) + (if used with reevaluate_targets>0, change targets by reevaluating with different parameters) - seedInitial can be optionally give a seed to randomize the starting profile (useful for developing, paper writing) """ @@ -328,25 +325,25 @@ def prep( for prof in self.portals_parameters["solution"]["predicted_channels"]: ymin_rel[prof] = np.array( [ymin_rel0] * len(self.portals_parameters["solution"][key_rhos]) ) - if enforceFiniteTemperatureGradients is not None: + if enforce_finite_aLT is not None: for prof in ['te', 'ti']: if prof in ymin_rel: - ymin_rel[prof] = np.array(ymin_rel[prof]).clip(min=None,max=enforceFiniteTemperatureGradients) + ymin_rel[prof] = np.array(ymin_rel[prof]).clip(min=None,max=enforce_finite_aLT) # Initialize print(">> PORTALS initalization module (START)", typeMsg="i") PORTALSinit.initializeProblem( self, self.folder, - fileGACODE, + mitim_state, ymax_rel, ymin_rel, start_from_folder=start_from_folder, define_ranges_from_profiles=define_ranges_from_profiles, - dvs_fixed=dvs_fixed, - limitsAreRelative=limitsAreRelative, + fixed_gradients=fixed_gradients, + limits_are_relative=limits_are_relative, cold_start=cold_start, - hardGradientLimits=hardGradientLimits, + yminymax_atleast=yminymax_atleast, tensor_options = self.tensor_options, seedInitial=seedInitial, checkForSpecies=askQuestions, @@ -359,7 +356,7 @@ def prep( if start_from_folder is not None: self.reuseTrainingTabular( - start_from_folder, self.folder, reevaluateTargets=reevaluateTargets + start_from_folder, self.folder, reevaluate_targets=reevaluate_targets ) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -510,9 +507,9 @@ def _check_flags_dictionary(d, d_check, avoid = ["run", "read"]): return key_rhos def reuseTrainingTabular( - self, folderRead, folderNew, reevaluateTargets=0, cold_startIfExists=False): + self, folderRead, folderNew, reevaluate_targets=0, cold_startIfExists=False): """ - reevaluateTargets: + reevaluate_targets: 0: No 1: Quick targets from powerstate with no transport calculation 2: Full original model (either with transport model targets or powerstate targets, but also calculate transport) @@ -538,7 +535,7 @@ def reuseTrainingTabular( typeMsg="i", ) - if reevaluateTargets > 0: + if reevaluate_targets > 0: print("- Re-evaluate targets", typeMsg="i") (folderNew / "TargetsRecalculate").mkdir(parents=True, exist_ok=True) @@ -570,7 +567,7 @@ def reuseTrainingTabular( name = f"portals_{b}_targets_ev{numPORTALS}" self_copy = copy.deepcopy(self) - if reevaluateTargets == 1: + if reevaluate_targets == 1: self_copy.powerstate.transport_options["evaluator"] = None self_copy.powerstate.target_options["options"]["targets_evolve"] = "target_evaluator_method" diff --git a/src/mitim_modules/portals/PORTALStools.py b/src/mitim_modules/portals/PORTALStools.py index 6731b6b0..30c8b2a9 100644 --- a/src/mitim_modules/portals/PORTALStools.py +++ b/src/mitim_modules/portals/PORTALStools.py @@ -452,28 +452,28 @@ def calculate_residuals(powerstate, portals_parameters, specific_vars=None): if var == "Qe": of0, cal0 = ( - of0 * portals_parameters["solution"]["Pseudo_multipliers"][0], - cal0 * portals_parameters["solution"]["Pseudo_multipliers"][0], + of0 * portals_parameters["solution"]["scalar_multipliers"][0], + cal0 * portals_parameters["solution"]["scalar_multipliers"][0], ) elif var == "Qi": of0, cal0 = ( - of0 * portals_parameters["solution"]["Pseudo_multipliers"][1], - cal0 * portals_parameters["solution"]["Pseudo_multipliers"][1], + of0 * portals_parameters["solution"]["scalar_multipliers"][1], + cal0 * portals_parameters["solution"]["scalar_multipliers"][1], ) elif var == "Ge": of0, cal0 = ( - of0 * portals_parameters["solution"]["Pseudo_multipliers"][2], - cal0 * portals_parameters["solution"]["Pseudo_multipliers"][2], + of0 * portals_parameters["solution"]["scalar_multipliers"][2], + cal0 * portals_parameters["solution"]["scalar_multipliers"][2], ) elif var == "GZ": of0, cal0 = ( - of0 * portals_parameters["solution"]["Pseudo_multipliers"][3], - cal0 * portals_parameters["solution"]["Pseudo_multipliers"][3], + of0 * portals_parameters["solution"]["scalar_multipliers"][3], + cal0 * portals_parameters["solution"]["scalar_multipliers"][3], ) elif var == "MtJm2": of0, cal0 = ( - of0 * portals_parameters["solution"]["Pseudo_multipliers"][4], - cal0 * portals_parameters["solution"]["Pseudo_multipliers"][4], + of0 * portals_parameters["solution"]["scalar_multipliers"][4], + cal0 * portals_parameters["solution"]["scalar_multipliers"][4], ) of, cal = torch.cat((of, of0), dim=-1), torch.cat((cal, cal0), dim=-1) diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index a02a2cbb..2b470ae5 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -20,10 +20,10 @@ def initializeProblem( fileStart, RelVar_y_max, RelVar_y_min, - limitsAreRelative=True, - hardGradientLimits=None, + limits_are_relative=True, + yminymax_atleast=None, cold_start=False, - dvs_fixed=None, + fixed_gradients=None, start_from_folder=None, define_ranges_from_profiles=None, seedInitial=None, @@ -201,18 +201,18 @@ def initializeProblem( dictDVs = OrderedDict() for var in dictCPs_base: for conti, i in enumerate(np.arange(1, len(dictCPs_base[var]))): - if limitsAreRelative: + if limits_are_relative: y1 = dictCPs_base[var][i] - abs(dictCPs_base[var][i])*RelVar_y_min[var][conti] y2 = dictCPs_base[var][i] + abs(dictCPs_base[var][i])*RelVar_y_max[var][conti] else: y1 = torch.tensor(RelVar_y_min[var][conti]).to(dfT) y2 = torch.tensor(RelVar_y_max[var][conti]).to(dfT) - if hardGradientLimits is not None: - if hardGradientLimits[0] is not None: - y1 = torch.tensor(np.min([y1, hardGradientLimits[0]])) - if hardGradientLimits[1] is not None: - y2 = torch.tensor(np.max([y2, hardGradientLimits[1]])) + if yminymax_atleast is not None: + if yminymax_atleast[0] is not None: + y1 = torch.tensor(np.min([y1, yminymax_atleast[0]])) + if yminymax_atleast[1] is not None: + y2 = torch.tensor(np.max([y2, yminymax_atleast[1]])) # Check that makes sense if y2-y1 < thr: @@ -225,10 +225,10 @@ def initializeProblem( base_gradient = torch.rand(1)[0] * (y2 - y1) / 4 + (3 * y1 + y2) / 4 name = f"aL{var}_{i}" - if dvs_fixed is None: + if fixed_gradients is None: dictDVs[name] = [y1, base_gradient, y2] else: - dictDVs[name] = [dvs_fixed[name][0], base_gradient, dvs_fixed[name][1]] + dictDVs[name] = [fixed_gradients[name][0], base_gradient, fixed_gradients[name][1]] # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Define output dictionaries diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 5e398bce..9fe83a96 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -20,7 +20,6 @@ def _evaluate_tglf(self): cold_start = self.cold_start Qi_includes_fast = simulation_options["Qi_includes_fast"] - launchMODELviaSlurm = simulation_options["launchEvaluationsAsSlurmJobs"] use_tglf_scan_trick = simulation_options["use_scan_trick_for_stds"] cores_per_tglf_instance = simulation_options["cores_per_tglf_instance"] keep_tglf_files = simulation_options["keep_files"] @@ -50,7 +49,6 @@ def _evaluate_tglf(self): tglf.run( 'base_tglf', ApplyCorrections=False, - launchSlurm= launchMODELviaSlurm, cold_start= cold_start, forceIfcold_start=True, extra_name= self.name, diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index e4d31039..ba869513 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -77,7 +77,7 @@ def calculator( "Ti_thermals": False, "ni_thermals": False, "recalculate_ptot": False, - "ensureMachNumber": None, + "force_mach": None, }, insert_highres_powers=True, rederive_profiles=False, diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index f71e3195..68c68d44 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -264,7 +264,7 @@ def powerstate_to_gacode( Ti_thermals = postprocess_input_gacode.get("Ti_thermals", True) ni_thermals = postprocess_input_gacode.get("ni_thermals", True) recalculate_ptot = postprocess_input_gacode.get("recalculate_ptot", True) - ensureMachNumber = postprocess_input_gacode.get("ensureMachNumber", None) + force_mach = postprocess_input_gacode.get("force_mach", None) # ------------------------------------------------------------------------------------------ # Insert profiles @@ -323,9 +323,9 @@ def powerstate_to_gacode( print("\t\t* Adjusting ni of thermal ions", typeMsg="i") profiles.scaleAllThermalDensities(scaleFactor=scaleFactor) - if "w0" not in self.predicted_channels and ensureMachNumber is not None: + if "w0" not in self.predicted_channels and force_mach is not None: # Rotation fixed to ensure Mach number - profiles.introduceRotationProfile(Mach_LF=ensureMachNumber) + profiles.introduceRotationProfile(Mach_LF=force_mach) # ------------------------------------------------------------------------------------------ # Insert Powers @@ -402,7 +402,6 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): else: profiles.profiles[conversions[ikey]] = np.zeros(len(profiles.profiles["qei(MW/m^3)"])) profiles.profiles[conversions[ikey]][:-extra_points] = state_temp.plasma[ikey][position_in_powerstate_batch,:].cpu().numpy() - def defineIons(self, input_gacode, rho_vec, dfT): """ diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index cb8734af..b245816f 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1623,7 +1623,7 @@ def correct(self, options={}, write=False, new_file=None): enforce_same_aLn = options.get("enforce_same_aLn", False) groupQIONE = options.get("groupQIONE", False) ensure_positive_Gamma = options.get("ensure_positive_Gamma", False) - ensureMachNumber = options.get("ensureMachNumber", None) + force_mach = options.get("force_mach", None) thermalize_fast = options.get("thermalize_fast", False) print("\t- Custom correction of input.gacode file has been requested") @@ -1697,8 +1697,8 @@ def correct(self, options={}, write=False, new_file=None): self.profiles["qpar_wall(1/m^3/s)"] = self.profiles["qpar_wall(1/m^3/s)"].clip(0) # Mach - if ensureMachNumber is not None: - self.introduceRotationProfile(Mach_LF=ensureMachNumber) + if force_mach is not None: + self.introduceRotationProfile(Mach_LF=force_mach) # ---------------------------------------------------------------------- # Re-derive diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 7aefa6f7..48ba6915 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -330,7 +330,7 @@ def _run( # Run simulation # ---------------------------------------------------------------------------------------------------------------- """ - launchSlurm = True -> Launch as a batch job in the machine chosen + launchSlurm = True -> Launch as a batch job in the machine chosen, if partition specified launchSlurm = False -> Launch locally as a bash script """ diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index 6501a4aa..a4188484 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -81,7 +81,7 @@ "exploration_ranges": { "ymax_rel": 1.0, "ymin_rel": 1.0, - "hardGradientLimits": [null, 4] + "yminymax_atleast": [null, 4] }, "change_last_radial_call" : true, "use_previous_surrogate_data" : true, From 1f1b78952d805dadaa03af67e6deac6b3f538aed Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 10:54:02 -0400 Subject: [PATCH 236/385] Implementation of new portals YAML !! --- pyproject.toml | 1 + src/mitim_modules/portals/PORTALSmain.py | 262 +++--------------- .../physics_models/transport_tglf.py | 3 - .../powertorch/utils/TRANSPORTtools.py | 5 +- src/mitim_tools/misc_tools/IOtools.py | 15 + src/mitim_tools/opt_tools/STRATEGYtools.py | 9 +- templates/portals.namelist.yaml | 200 +++++++++++++ 7 files changed, 258 insertions(+), 237 deletions(-) create mode 100644 templates/portals.namelist.yaml diff --git a/pyproject.toml b/pyproject.toml index 865df759..87c521ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dependencies = [ "onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models "tensorflow", "f90nml", + "pyyaml" # "vmecpp", # "quends @ git+https://github.com/sandialabs/quends.git" ] diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 8890c877..5f82ff5b 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -1,99 +1,36 @@ import shutil import torch import copy +import yaml, importlib +from collections import OrderedDict import numpy as np import dill as pickle_dill -from functools import partial from collections import OrderedDict from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools -from mitim_modules.powertorch.utils import TRANSPORTtools from mitim_modules.portals import PORTALStools from mitim_modules.portals.utils import ( PORTALSinit, - PORTALSoptimization, PORTALSanalysis, ) from mitim_tools.opt_tools import STRATEGYtools from mitim_tools.opt_tools.utils import BOgraphics from mitim_tools.misc_tools.LOGtools import printMsg as print +from mitim_tools import __mitimroot__ from IPython import embed -""" -Reading analysis for PORTALS has more options than standard: --------------------------------------------------------------------------------------------------------- - Standard: - ************************************** - -1: Only improvement - 0: Only optimization_results - 1: 0 + Pickle - 2: 1 + Final redone in this machine - - PORTALS-specific: - ************************************** - 3: 1 + PORTALSplot metrics (only works if optimization_extra is provided or Execution exists) - 4: 3 + PORTALSplot expected (only works if optimization_extra is provided or Execution exists) - 5: 2 + 4 (only works if optimization_extra is provided or Execution exists) - - >2 will also plot profiles & gradients comparison (original, initial, best) -""" - -def default_namelist(optimization_options): - """ - This is to be used after reading the namelist, so self.optimization_options should be completed with main defaults. - """ - - # Initialization - optimization_options["initialization_options"]["initial_training"] = 5 - optimization_options["initialization_options"]["initialization_fun"] = PORTALSoptimization.initialization_simple_relax - - # Strategy for stopping - optimization_options["convergence_options"]["maximum_iterations"] = 50 - optimization_options['convergence_options']['stopping_criteria'] = PORTALStools.stopping_criteria_portals - optimization_options['convergence_options']['stopping_criteria_parameters'] = { - "maximum_value": 5e-3, # Reducing residual by 200x is enough - "maximum_value_is_rel": True, - "minimum_dvs_variation": [10, 5, 0.1], # After iteration 10, Check if 5 consecutive DVs are varying less than 0.1% from the rest that has been evaluated - "ricci_value": 0.05, - "ricci_d0": 2.0, - "ricci_lambda": 0.5, - } - - optimization_options['acquisition_options']['relative_improvement_for_stopping'] = 1e-2 - - # Surrogate - optimization_options["surrogate_options"]["surrogate_selection"] = partial(PORTALStools.surrogate_selection_portals) - - # if CGYROrun: - # # CGYRO runs should prioritize accuracy - # optimization_options["acquisition_options"]["type"] = "posterior_mean" - # optimization_options["acquisition_options"]["optimizers"] = ["root", "botorch", "ga"] - # else: - # # TGLF runs should prioritize speed - optimization_options["acquisition_options"]["type"] = "posterior_mean" # "noisy_logei_mc" - optimization_options["acquisition_options"]["optimizers"] = ["sr", "root"] #, "botorch"] - - return optimization_options - - class portals(STRATEGYtools.opt_evaluator): def __init__( self, folder, # Folder where the PORTALS workflow will be run - transport_models = ['tglf', 'neo'], - namelist=None, # If None, default namelist will be used. If not None, it will be read and used + portals_namelist = None, tensor_options = { "dtype": torch.double, "device": torch.device("cpu"), }, - portals_transformation_variables = None, # If None, use defaults for both main and trace - portals_transformation_variables_trace = None, - additional_params_in_surrogate = [] # Additional parameters to be used in the surrogate (e.g. ['q']) ): - ''' - Note that additional_params_in_surrogate They must exist in the plasma dictionary of the powerstate object - ''' + print("\n-----------------------------------------------------------------------------------------") print("\t\t\t PORTALS class module") @@ -101,171 +38,31 @@ def __init__( # Store folder, namelist. Read namelist - super().__init__( - folder, - namelist=namelist, - tensor_options=tensor_options, - default_namelist_function=( - partial(default_namelist) - if (namelist is None) - else None - ), - ) + super().__init__(folder,tensor_options=tensor_options) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Default (please change to your desire after instancing the object) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.portals_parameters = {} - - """ - Physics-informed parameters to fit surrogates - --------------------------------------------- - """ - - ( - portals_transformation_variables, - portals_transformation_variables_trace, - ) = PORTALStools.default_portals_transformation_variables(additional_params = additional_params_in_surrogate) - - """ - PORTALS-specific parameters (i.e. on the transport solver level, not model) - ---------------------------------------------------------------------------- - """ + # Read optimization namelist (always the default, the values to be modified are in the portals one) + namelist = __mitimroot__ / "templates" / "main.namelist.json" + self.optimization_options = IOtools.read_mitim_nml(namelist) - self.portals_parameters["solution"] = { - - # Specification of radial locations (roa wins over rho, if provided) - "predicted_roa": [0.35, 0.55, 0.75, 0.875, 0.9], - "predicted_rho": [0.3, 0.45, 0.6, 0.75, 0.9], - - # Channels to be predicted - "predicted_channels": ["te", "ti", "ne"], # ['nZ','w0'] - - "trace_impurity": None, # Impurity to do flux-matching for if nZ enabled (name of first impurity instance AFTER postprocessing), e.g. "W" - - # PORTALS - "portals_transformation_variables": portals_transformation_variables, # Physics-informed parameters to fit surrogates - "portals_transformation_variables_trace": portals_transformation_variables_trace, # Physics-informed parameters to fit surrogates for trace impurities - "turbulent_exchange_as_surrogate": False, # Run turbulent exchange as surrogate? - "scalar_multipliers": [1.0]*5, # [Qe,Qi,Ge] multipliers to calculate pseudo - "impurity_trick": True, # If True, fit model to GZ/nZ, valid on the trace limit - "fZ0_as_weight": None, # If not None, using fZ0_as_weight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis - "fImp_orig": 1.0, - "additional_params_in_surrogate": additional_params_in_surrogate, - "keep_full_model_folder": True, # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) - } + # Read PORTALS namelist (if not provided, use default) + if portals_namelist is None: + portals_namelist = __mitimroot__ / "templates" / "portals.namelist.yaml" + print(f"\t- No PORTALS namelist provided, using default in {IOtools.clipstr(portals_namelist)}") + else: + print(f"\t- Using provided PORTALS namelist in {IOtools.clipstr(portals_namelist)}") + self.portals_parameters = read_portals_nml(portals_namelist) - ''' - model_parameters - ---------------- - These parameters are communicated to the powertorch object. - ''' - - # ------------------------- - # Transport model settings - # ------------------------- - self.portals_parameters["transport"] = { - - # Transport model class - - "evaluator": TRANSPORTtools.portals_transport_model, - - "evaluator_instance_attributes": { - "turbulence_model": transport_models[0], - "neoclassical_model": transport_models[1], - }, - - # Simulation kwargs to be passed directly to run and read commands - - "options": { - - # Defaults for TGLF simulation - "tglf": { - "run": { - "code_settings": 6, - "extraOptions": {}, - }, - "read": {}, - "use_scan_trick_for_stds": 0.02, # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta - "keep_files": "base", # minimal: only retrieve minimal files always; base: retrieve all files when running TGLF base; none: always minimal files - "cores_per_tglf_instance": 1, # Number of cores to use per TGLF instance - "percent_error": 5, # (%) Error (std, in percent) of model evaluation TGLF if not scan trick - "Qi_includes_fast": False, # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) - }, - - # Defaults for NEO simulation - "neo": { - "run": {}, - "read": {}, - "percent_error": 10 # (%) Error (std, in percent) of model evaluation - }, - - # Defaults for CGYRO simulation - "cgyro": { - "run": { - "code_settings": 1, - "extraOptions": {}, - "run_type": "normal", # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit - }, - "read": { - "tmin": 0.0 - }, - "Qi_stable_criterion": 0.01, # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable - "Qi_stable_percent_error": 5.0, # (%) For CGYRO runs, minimum error based on target if case is considered stable - }, - # Defaults for GX simulation - "gx": { - "run": { - "code_settings": 1, - "extraOptions": {}, - "run_type": "normal", # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit - }, - "read": { - "tmin": 0.0 - }, - }, - }, - - # Corrections to be applied to each iteration input.gacode file - - "profiles_postprocessing_fun": None, # Function to post-process input.gacode only BEFORE passing to transport codes - - "applyCorrections": { - "Ti_thermals": True, # Keep all thermal ion temperatures equal to the main Ti - "ni_thermals": True, # Adjust for quasineutrality by modifying the thermal ion densities together with ne - "recalculate_ptot": True, # Recompute PTOT to insert in input file each time - "Tfast_ratio": False, # Keep the ratio of Tfast/Te constant throughout the Te evolution - "force_mach": None, # Change w0 to match this Mach number when Ti varies - }, - } - - # ------------------------- - # Target model settings - # ------------------------- - - from mitim_modules.powertorch.physics_models.targets_analytic import analytical_model as target_evaluator - - self.portals_parameters["target"] = { - "evaluator": target_evaluator, - "options": { - "targets_evolve": ["qie", "qrad", "qfus"], - "target_evaluator_method": "powerstate", # Method to calculate targets (tgyro or powerstate) - "force_zero_particle_flux": False, # If True, ignore particle flux profile and assume zero for all radii - "targets_resolution": 20, # If not None, calculate targets with this radial resolution (defaults target_evaluator_method to powerstate) - "percent_error": 1 # (%) Error (std, in percent) of model evaluation - } - } + # Apply the optimization options to the proper namelist and drop it from portals_parameters + if 'optimization_options' in self.portals_parameters: + self.optimization_options = IOtools.deep_dict_update(self.optimization_options, self.portals_parameters['optimization_options']) + del self.portals_parameters['optimization_options'] # Grab all the flags here in a way that, after changing the dictionary extenrally, I make sure it's the same flags as PORTALS expects - - def _grab_flags_dictionary(d): - keys = {} - for key in d.keys(): - keys[key] = _grab_flags_dictionary(d[key]) if isinstance(d[key], dict) else None - return keys - - self.potential_flags = _grab_flags_dictionary(self.portals_parameters) + self.potential_flags = IOtools.deep_grab_flags_dict(self.portals_parameters) def prep( self, @@ -736,3 +533,20 @@ def analyze_results( portals_full.runCases(onlyBest=onlyBest, cold_start=cold_start, fn=fn) return portals_full.opt_fun.mitim_model.optimization_object + +def read_portals_nml(path: str): + + def resolve(x): + if isinstance(x, dict): + return {k: resolve(v) for k, v in x.items()} + if isinstance(x, list): + return [resolve(v) for v in x] + if isinstance(x, str) and x.startswith("import::"): + modattr = x[len("import::"):] + module_name, attr = modattr.rsplit(".", 1) + return getattr(importlib.import_module(module_name), attr) + return x + + with open(path, "r") as f: + cfg = yaml.safe_load(f) + return resolve(cfg) \ No newline at end of file diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 9fe83a96..39c2a973 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -113,7 +113,6 @@ def _evaluate_tglf(self): cold_start=cold_start, extra_name=self.name, cores_per_tglf_instance=cores_per_tglf_instance, - launchMODELviaSlurm=launchMODELviaSlurm, Qi_includes_fast=Qi_includes_fast, only_minimal_files=keep_tglf_files in ['minimal', 'base'], **simulation_options["run"] @@ -179,7 +178,6 @@ def _run_tglf_uncertainty_model( extra_name="", remove_folders_out = False, cores_per_tglf_instance = 4, # e.g. 4 core per radius, since this is going to launch ~ Nr=5 x (Nv=6 x Nd=2 + 1) = 65 TGLFs at once - launchMODELviaSlurm=False, Qi_includes_fast=False, only_minimal_files=True, # Since I only care about fluxes here, do not retrieve all the files ): @@ -244,7 +242,6 @@ def _run_tglf_uncertainty_model( positionIon=impurityPosition+2, attempts_execution=2, only_minimal_files=only_minimal_files, - launchSlurm=launchMODELviaSlurm, ) # Remove folders because they are heavy to carry many throughout diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index fd4af7e7..aa649ee7 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -323,10 +323,6 @@ def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): dum[f"{mapper[var][1]}_tr_{suffix}"] = np.array(json_dict['fluxes_mean'][var]) * gb dum[f"{mapper[var][1]}_tr_{suffix}_stds"] = np.array(json_dict['fluxes_stds'][var]) * gb - mapperQ = { - - } - if units == 'GB': print("\t\t- File has fluxes in GB units... using GB units from powerstate to convert to real units") @@ -407,6 +403,7 @@ def produce_profiles(self): @IOtools.hook_method(after=partial(write_json, file_name = 'fluxes_turb.json', suffix= 'turb')) def evaluate_turbulence(self): + if self.turbulence_model == 'tglf': return tglf_model.evaluate_turbulence(self) elif self.turbulence_model == 'cgyro': diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index 46260070..38f185b5 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -312,6 +312,21 @@ def clipstr(txt, chars=40): txt = f"{txt}" return f"{'...' if len(txt) > chars else ''}{txt[-chars:]}" if txt is not None else None + +def deep_dict_update(d, u): + for k, v in u.items(): + if isinstance(v, dict) and isinstance(d.get(k), dict): + deep_dict_update(d[k], v) # recurse into nested dict + else: + d[k] = v # overwrite at lowest level + return d + +def deep_grab_flags_dict(d): + keys = {} + for key in d.keys(): + keys[key] = deep_grab_flags_dict(d[key]) if isinstance(d[key], dict) else None + return keys + def receiveWebsite(url, data=None): NumTriesAfterTimeOut = 60 secWaitTimeOut = 10 diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 5ab13698..f71e478b 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -97,12 +97,12 @@ def __init__( IOtools.askNewFolder(self.folder / "Outputs") if namelist is not None: - print(f"\t- Namelist provided: {namelist}", typeMsg="i") + print(f"\t- Optimizaiton namelist provided: {namelist}", typeMsg="i") self.optimization_options = IOtools.read_mitim_nml(namelist) elif default_namelist_function is not None: - print("\t- Namelist not provided, using MITIM default for this optimization sub-module", typeMsg="i") + print("\t- Optimizaiton namelist not provided, using MITIM default for this optimization sub-module", typeMsg="i") namelist = __mitimroot__ / "templates" / "main.namelist.json" self.optimization_options = IOtools.read_mitim_nml(namelist) @@ -110,10 +110,7 @@ def __init__( self.optimization_options = default_namelist_function(self.optimization_options) else: - print( - "\t- No namelist provided (likely b/c for reading/plotting purposes)", - typeMsg="i", - ) + print("\t- No optimizaiton namelist provided (likely b/c for reading/plotting purposes)",typeMsg="i") self.optimization_options = None self.surrogate_parameters = { diff --git a/templates/portals.namelist.yaml b/templates/portals.namelist.yaml new file mode 100644 index 00000000..7dc3cf39 --- /dev/null +++ b/templates/portals.namelist.yaml @@ -0,0 +1,200 @@ +# ----------------------------------------------------------------- +# Main solver parameters +# ----------------------------------------------------------------- + +solution: + + # Specification of radial locations (roa wins over rho, if provided) + predicted_roa: null + predicted_rho: [0.35, 0.55, 0.75, 0.875, 0.9] + + # Channels to be predicted + predicted_channels: ["te", "ti", "ne"] # ['nZ','w0'] + + # Impurity to do flux-matching for if nZ enabled (name of first impurity instance AFTER postprocessing), e.g. "W" + trace_impurity: null + + # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) + keep_full_model_folder: true + + # Run turbulent exchange as surrogate + turbulent_exchange_as_surrogate: false + + # If True, fit model to GZ/nZ, valid on the trace limit + impurity_trick: true + + # If not None, using fZ0_as_weight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis + fZ0_as_weight: null + fImp_orig: 1.0 + + # [Qe,Qi,Ge,Mt,GZ] multipliers to calculate scalarized function + scalar_multipliers: [1.0, 1.0, 1.0, 1.0, 1.0] + + # Physics-informed parameters to fit surrogates + portals_transformation_variables: + 10: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0]} + 30: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0], nuei: [te, ne], tite: [te, ti], w0_n: [w0]} + 10000: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0], nuei: [te, ne], tite: [te, ti], w0_n: [w0], beta_e: [te, ne]} + + # Physics-informed parameters to fit surrogates for trace impurities + portals_transformation_variables_trace: + 10: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0], aLnZ: [aLnZ]} + 30: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0], nuei: [te, ne], tite: [te, ti], w0_n: [w0], aLnZ: [aLnZ]} + 10000: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0], nuei: [te, ne], tite: [te, ti], w0_n: [w0], beta_e: [te, ne], aLnZ: [aLnZ]} + + additional_params_in_surrogate: [] + +# ----------------------------------------------------------------- +# Transport model parameters +# ----------------------------------------------------------------- + +transport: + + # Transport model class + evaluator: "import::mitim_modules.powertorch.utils.TRANSPORTtools.portals_transport_model" + + # Select transport models (when instantiating the evaluator, assign these parameters) + evaluator_instance_attributes: + turbulence_model: "tglf" + neoclassical_model: "neo" + + # Simulation kwargs to be passed directly to run and read commands (defaults here) + options: + + # *********************************** + # TGLF + # *********************************** + tglf: + + # Kwargs to be passed to the run command + run: + code_settings: 6 + extraOptions: {} + + # Kwargs to be passed to the read command + read: {} + + # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta + use_scan_trick_for_stds: 0.02 + + # Files to keep from simulation (minimal: only retrieve minimal files always, none: always minimal files, base: retrieve all files) + keep_files: "base" + + # Number of cores to use per TGLF instance + cores_per_tglf_instance: 1 + + # (%) Error (std, in percent) of model evaluation TGLF if not scan trick + percent_error: 5 + + # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) + Qi_includes_fast: false + + # *********************************** + # NEO + # *********************************** + neo: + + run: {} + + read: {} + + # (%) Error (std, in percent) of model evaluation + percent_error: 10 + + # *********************************** + # CGYRO + # *********************************** + cgyro: + + run: + code_settings: 1 + extraOptions: {} + + # Run type: normal (submit and wait), submit (submit and do not wait), prep (do not submit) + run_type: "normal" + + read: + tmin: 0.0 + + # For CGYRO runs, MW/m^2 of Qi below which the case is considered stable + Qi_stable_criterion: 0.01 + + # (%) For CGYRO runs, minimum error based on target if case is considered stable + Qi_stable_percent_error: 5.0 + + # *********************************** + # GX + # *********************************** + gx: + + run: + + code_settings: 1 + extraOptions: {} + + # Run type: normal (submit and wait), submit (submit and do not wait), prep (do not submit) + run_type: "normal" + + read: + tmin: 0.0 + + # Function to post-process input.gacode only BEFORE passing to transport codes + profiles_postprocessing_fun: null + + # Corrections to be applied to each iteration input.gacode file + applyCorrections: + Ti_thermals: true # Keep all thermal ion temperatures equal to the main Ti + ni_thermals: true # Adjust for quasineutrality by modifying the thermal ion densities together with ne + recalculate_ptot: true # Recompute PTOT to insert in input file each time + Tfast_ratio: false # Keep the ratio of Tfast/Te constant throughout the Te evolution + force_mach: null # Change w0 to match this Mach number when Ti varies + +# ----------------------------------------------------------------- +# Target model parameters +# ----------------------------------------------------------------- + +target: + + # Target model evaluator + evaluator: "import::mitim_modules.powertorch.physics_models.targets_analytic.analytical_model" + + options: + targets_evolve: ["qie", "qrad", "qfus"] + # Method to calculate targets (tgyro or powerstate) + target_evaluator_method: "powerstate" + # If True, ignore particle flux profile and assume zero for all radii + force_zero_particle_flux: false + # If not None, calculate targets with this radial resolution + targets_resolution: 20 + # (%) Error (std, in percent) of model evaluation + percent_error: 1 + + +# ----------------------------------------------------------------- +# Optimization options (to change the main optimization namelist) +# ----------------------------------------------------------------- + +optimization_options: + + initialization_options: + initial_training: 5 + initialization_fun: "import::mitim_modules.portals.utils.PORTALSoptimization.initialization_simple_relax" + + convergence_options: + maximum_iterations: 50 + stopping_criteria: "import::mitim_modules.portals.PORTALStools.stopping_criteria_portals" + stopping_criteria_parameters: + maximum_value: 5.e-3 # Reducing residual by 200x is enough + maximum_value_is_rel: true + minimum_dvs_variation: [10, 5, 0.1] # After iteration 10, Check if 5 consecutive DVs are varying less than 0.1% from the rest that has been evaluated + ricci_value: 0.05 + ricci_d0: 2.0 + ricci_lambda: 0.5 + + acquisition_options: + relative_improvement_for_stopping: 1.e-2 + type: "posterior_mean" # "noisy_logei_mc" + optimizers: ["sr", "root"] #, "botorch"] + + surrogate_options: + surrogate_selection: "import::mitim_modules.portals.PORTALStools.surrogate_selection_portals" From d349129d494105a8a46ef704d33022647d152950 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 11:17:19 -0400 Subject: [PATCH 237/385] bug fix --- src/mitim_modules/portals/utils/PORTALSinit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 2b470ae5..110785d6 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -105,7 +105,7 @@ def initializeProblem( # Print warning or question to be careful! if len(speciesNotFound) > 0: - if "qrad" in portals_fun.portals_parameters["target"]["targets_evolve"]: + if "qrad" in portals_fun.portals_parameters["target"]["options"]["targets_evolve"]: answerYN = print(f"\t- Species {speciesNotFound} not found in radiation database, radiation will be zero in PORTALS... is this ok for your predictions?",typeMsg="q" if checkForSpecies else "w") if checkForSpecies and (not answerYN): From 0ee09a146e3e79deda88a10bc537f2d191c043a2 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 11:18:49 -0400 Subject: [PATCH 238/385] askQuestions can be provided to the portals init --- src/mitim_modules/portals/PORTALSmain.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 5f82ff5b..a17c4841 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -29,9 +29,9 @@ def __init__( "dtype": torch.double, "device": torch.device("cpu"), }, + askQuestions=True ): - print("\n-----------------------------------------------------------------------------------------") print("\t\t\t PORTALS class module") print("-----------------------------------------------------------------------------------------\n") @@ -98,7 +98,7 @@ def prep( # Make sure that options that are required by good behavior of PORTALS # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - key_rhos = self.check_flags() + key_rhos = self.check_flags(askQuestions=askQuestions) # TO BE REMOVED IN FUTURE if not isinstance(cold_start, bool): @@ -152,9 +152,7 @@ def prep( # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if start_from_folder is not None: - self.reuseTrainingTabular( - start_from_folder, self.folder, reevaluate_targets=reevaluate_targets - ) + self.reuseTrainingTabular(start_from_folder, self.folder, reevaluate_targets=reevaluate_targets) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Ignore targets in surrogate_data.csv @@ -280,7 +278,7 @@ def analyze_results(self, plotYN=True, fn=None, cold_start=False, analysis_level self, plotYN=plotYN, fn=fn, cold_start=cold_start, analysis_level=analysis_level ) - def check_flags(self): + def check_flags(self, askQuestions=True): print(">> PORTALS flags pre-check") @@ -289,7 +287,7 @@ def check_flags(self): def _check_flags_dictionary(d, d_check, avoid = ["run", "read"]): for key in d.keys(): if key not in d_check: - print(f"\t- {key} is an unexpected variable, prone to errors or misinterpretation",typeMsg="q") + print(f"\t- {key} is an unexpected variable, prone to errors or misinterpretation",typeMsg="q" if askQuestions else "w") elif not isinstance(d[key], dict): continue elif key in avoid: From 2c641c2b20c9e99d4c0d0bed0ab2e26f9aff4bcd Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 11:19:11 -0400 Subject: [PATCH 239/385] askQuestions works for portals prep now --- src/mitim_modules/portals/PORTALSmain.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index a17c4841..c0ca0cb7 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -29,7 +29,6 @@ def __init__( "dtype": torch.double, "device": torch.device("cpu"), }, - askQuestions=True ): print("\n-----------------------------------------------------------------------------------------") From b327dad42447644b9792e8c9d1a0258c6744ba7d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 11:23:35 -0400 Subject: [PATCH 240/385] nicer yaml --- src/mitim_tools/opt_tools/STRATEGYtools.py | 20 +++---- templates/main.namelist.json | 2 +- templates/portals.namelist.yaml | 61 +++++++++++++++++++--- 3 files changed, 64 insertions(+), 19 deletions(-) diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index f71e478b..14691071 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -1838,7 +1838,7 @@ def stopping_criteria_default(mitim_bo, parameters = {}): maximum_value_is_rel = parameters["maximum_value_is_rel"] maximum_value_orig = parameters["maximum_value"] - minimum_dvs_variation = parameters["minimum_dvs_variation"] + minimum_inputs_variation = parameters["minimum_inputs_variation"] res_base = -mitim_bo.BOmetrics["overall"]["Residual"][0].item() @@ -1848,8 +1848,8 @@ def stopping_criteria_default(mitim_bo, parameters = {}): # Stopping criteria # ------------------------------------------------------------------------------------ - if minimum_dvs_variation is not None: - converged_by_dvs, yvals = stopping_criteria_by_dvs(mitim_bo, minimum_dvs_variation) + if minimum_inputs_variation is not None: + converged_by_dvs, yvals = stopping_criteria_by_dvs(mitim_bo, minimum_inputs_variation) else: converged_by_dvs = False yvals = None @@ -1880,28 +1880,28 @@ def stopping_criteria_by_value(mitim_bo, maximum_value): return criterion_is_met, -yvals -def stopping_criteria_by_dvs(mitim_bo, minimum_dvs_variation): +def stopping_criteria_by_dvs(mitim_bo, minimum_inputs_variation): print("\t- Checking DV variations...") _, yG_max = TESTtools.DVdistanceMetric(mitim_bo.train_X) criterion_is_met = ( mitim_bo.currentIteration - >= minimum_dvs_variation[0] - + minimum_dvs_variation[1] + >= minimum_inputs_variation[0] + + minimum_inputs_variation[1] ) - for i in range(int(minimum_dvs_variation[1])): + for i in range(int(minimum_inputs_variation[1])): criterion_is_met = criterion_is_met and ( - yG_max[-1 - i] < minimum_dvs_variation[2] + yG_max[-1 - i] < minimum_inputs_variation[2] ) if criterion_is_met: print( - f"\t\t* DVs varied by less than {minimum_dvs_variation[2]}% compared to the rest of individuals for the past {int(minimum_dvs_variation[1])} iterations" + f"\t\t* DVs varied by less than {minimum_inputs_variation[2]}% compared to the rest of individuals for the past {int(minimum_inputs_variation[1])} iterations" ) else: print( - f"\t\t* DVs have varied by more than {minimum_dvs_variation[2]}% compared to the rest of individuals for the past {int(minimum_dvs_variation[1])} iterations" + f"\t\t* DVs have varied by more than {minimum_inputs_variation[2]}% compared to the rest of individuals for the past {int(minimum_inputs_variation[1])} iterations" ) return criterion_is_met, yG_max diff --git a/templates/main.namelist.json b/templates/main.namelist.json index 9dc26580..6276b58a 100644 --- a/templates/main.namelist.json +++ b/templates/main.namelist.json @@ -16,7 +16,7 @@ "stopping_criteria_parameters": { "maximum_value": -1e-3, "maximum_value_is_rel": false, - "minimum_dvs_variation": [10, 3, 0.01] + "minimum_inputs_variation": [10, 3, 0.01] } }, "initialization_options": { diff --git a/templates/portals.namelist.yaml b/templates/portals.namelist.yaml index 7dc3cf39..0138e63f 100644 --- a/templates/portals.namelist.yaml +++ b/templates/portals.namelist.yaml @@ -1,3 +1,21 @@ +# ------------------------------------------------------------------------------------------------------------ +# This is a complete template for a PORTALS-TGLF+NEO profile prediction. +# The user is welcomed to change any of the parameters below to fit their specific use case, +# by copying this template namelist and instantiating the PORTALS object with its path: +# portals_fun = PORTALSmain.portals(folder, portals_namelist=PATH_TO_NAMELIST) +# +# Alternatively, you can simply use the default parameters provided in this template by not +# specifying a namelist and then, in the launch script change the dictionary values: +# portals_fun = PORTALSmain.portals(folderWork) +# portals_fun.portals_parameters["solution"]['predicted_roa'] = [0.25, 0.45, 0.65, 0.85] +# portals_fun.portals_parameters["solution"]['predicted_channels'] = ["te", "ti", "ne", "nZ", 'w0'] +# ... +# +# The dictionary follows the same structure as the YAML namelist, except the optimization_options that must +# be passed as a separate dictionary: +# portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 10 +# + # ----------------------------------------------------------------- # Main solver parameters # ----------------------------------------------------------------- @@ -8,8 +26,8 @@ solution: predicted_roa: null predicted_rho: [0.35, 0.55, 0.75, 0.875, 0.9] - # Channels to be predicted - predicted_channels: ["te", "ti", "ne"] # ['nZ','w0'] + # Channels to be predicted (Options: ["te", "ti", "ne", "nZ", "w0"]) + predicted_channels: ["te", "ti", "ne"] # Impurity to do flux-matching for if nZ enabled (name of first impurity instance AFTER postprocessing), e.g. "W" trace_impurity: null @@ -159,17 +177,22 @@ target: evaluator: "import::mitim_modules.powertorch.physics_models.targets_analytic.analytical_model" options: + + # Targets to evolve (Options: ["qie", "qrad", "qfus"]) targets_evolve: ["qie", "qrad", "qfus"] + # Method to calculate targets (tgyro or powerstate) target_evaluator_method: "powerstate" + # If True, ignore particle flux profile and assume zero for all radii force_zero_particle_flux: false + # If not None, calculate targets with this radial resolution targets_resolution: 20 + # (%) Error (std, in percent) of model evaluation percent_error: 1 - # ----------------------------------------------------------------- # Optimization options (to change the main optimization namelist) # ----------------------------------------------------------------- @@ -177,24 +200,46 @@ target: optimization_options: initialization_options: + + # PORTALS works well with 5 initial training points obtained with simple relaxation initial_training: 5 initialization_fun: "import::mitim_modules.portals.utils.PORTALSoptimization.initialization_simple_relax" + # Convergence options for the optimization process convergence_options: + + # Criterion 1: Maximum iterations maximum_iterations: 50 + + # Stopping criteria function for the optimization process stopping_criteria: "import::mitim_modules.portals.PORTALStools.stopping_criteria_portals" + stopping_criteria_parameters: - maximum_value: 5.e-3 # Reducing residual by 200x is enough + + # Criterion 2: Residual reduction + maximum_value: 5.e-3 # Reducing residual by 200x is enough maximum_value_is_rel: true - minimum_dvs_variation: [10, 5, 0.1] # After iteration 10, Check if 5 consecutive DVs are varying less than 0.1% from the rest that has been evaluated + + # Criterion 3: Variation of input parameters + minimum_inputs_variation: [10, 5, 0.1] # After iteration 10, Check if 5 consecutive DVs are varying less than 0.1% from the rest that has been evaluated + + # Criterion 4: Ricci metric ricci_value: 0.05 ricci_d0: 2.0 ricci_lambda: 0.5 acquisition_options: - relative_improvement_for_stopping: 1.e-2 - type: "posterior_mean" # "noisy_logei_mc" - optimizers: ["sr", "root"] #, "botorch"] + + # Relative improvement for stopping criteria of the acquisition optimization + relative_improvement_for_stopping: 1.e-2 # 100x is enough + + # Type of acquisition function (Options: ["posterior_mean", "noisy_logei_mc", ...]) + type: "posterior_mean" + + # Optimizers to apply in sequence (Options: ["sr", "root", "botorch"]) + optimizers: ["sr", "root"] surrogate_options: + + # Function to select the GP parameters depending on each channel/radius (superseeds the main optimization namelist) surrogate_selection: "import::mitim_modules.portals.PORTALStools.surrogate_selection_portals" From 37ef22554284ccdd0ae60c5f5f408a6c2bf835f5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 13:45:29 -0400 Subject: [PATCH 241/385] MAESTRO working with new portals yml namelist and changes --- .../maestro/scripts/run_maestro.py | 2 +- .../maestro/tmp_tests/maestro_test1.py | 2 +- .../maestro/utils/PORTALSbeat.py | 105 ++++++++---------- .../portals/utils/PORTALSplot.py | 2 +- .../plasmastate_tools/utils/state_plotting.py | 6 +- templates/maestro_namelist.json | 28 +++-- tests/MAESTRO_workflow.py | 2 +- 7 files changed, 73 insertions(+), 74 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index a85f8a1a..b93e9781 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -146,7 +146,7 @@ def profiles_postprocessing_fun(file_profs): p.enforce_same_density_gradients() p.write_state(file=file_profs) return p - beat_namelist['PORTALSparameters']['main_parameters']['profiles_postprocessing_fun'] = profiles_postprocessing_fun + beat_namelist['portals_parameters']['transport']['profiles_postprocessing_fun'] = profiles_postprocessing_fun else: raise ValueError(f"[MITIM] {beat_type} beat not found in the MAESTRO namelist") diff --git a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py b/src/mitim_modules/maestro/tmp_tests/maestro_test1.py index 04c021b9..c84d8579 100644 --- a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py +++ b/src/mitim_modules/maestro/tmp_tests/maestro_test1.py @@ -43,7 +43,7 @@ # To see what values this namelist can take: mitim_modules/portals/PORTALSmain.py: __init__() portals_namelist = { "main_parameters": {"launchEvaluationsAsSlurmJobs": True,"forceZeroParticleFlux": True, 'use_tglf_scan_trick': 0.02}, - "MODELparameters": { "RoaLocations": [0.35,0.55,0.75,0.875,0.9], + "MODELparameters": { "predicted_roa": [0.35,0.55,0.75,0.875,0.9], "ProfilesPredicted": ["te", "ti", "ne"], "Physics_options": {"TypeTarget": 3}, "transport_model": {"TGLFsettings": 6, "extraOptionsTGLF": {'USE_BPER':True}}}, diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 05787210..ba65c7a4 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -31,10 +31,9 @@ def prepare(self, 'ymin_rel': 1.0, 'hardGradientLimits': [0,2] }, - PORTALSparameters = {}, - MODELparameters = {}, + portals_parameters = {}, + initialization_parameters = {}, optimization_options = {}, - INITparameters = {}, enforce_impurity_radiation_existence = True, ): @@ -65,10 +64,9 @@ def prepare(self, self.profiles_current.write_state(file = self.fileGACODE) - self.PORTALSparameters = PORTALSparameters - self.MODELparameters = MODELparameters + self.portals_parameters = portals_parameters self.optimization_options = optimization_options - self.INITparameters = INITparameters + self.initialization_parameters = initialization_parameters self.additional_params_in_surrogate = additional_params_in_surrogate @@ -87,14 +85,19 @@ def run(self, **kwargs): cold_start = kwargs.get('cold_start', False) - portals_fun = PORTALSmain.portals(self.folder, additional_params_in_surrogate = self.additional_params_in_surrogate) + portals_fun = PORTALSmain.portals(self.folder) - modify_dictionary(portals_fun.PORTALSparameters, self.PORTALSparameters) - modify_dictionary(portals_fun.MODELparameters, self.MODELparameters) - modify_dictionary(portals_fun.optimization_options, self.optimization_options) - modify_dictionary(portals_fun.INITparameters, self.INITparameters) + portals_fun.portals_parameters = IOtools.deep_dict_update(portals_fun.portals_parameters, self.portals_parameters) + portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.optimization_options) - portals_fun.prep(self.fileGACODE,askQuestions=False,**self.exploration_ranges) + portals_fun.portals_parameters['solution']['additional_params_in_surrogate'] = self.additional_params_in_surrogate + + # Initialization now happens by the user + from mitim_tools.gacode_tools.PROFILEStools import gacode_state + p = gacode_state(self.fileGACODE) + p.correct(options=self.initialization_parameters) + + portals_fun.prep(p,askQuestions=False,**self.exploration_ranges) self.mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, seed = self.maestro_instance.master_seed, cold_start = cold_start, askQuestions = False) @@ -226,13 +229,14 @@ def merge_parameters(self): # Insert powers opt_fun = PORTALSanalysis.PORTALSanalyzer.from_folder(self.folder) - if opt_fun.MODELparameters['Physics_options']["TypeTarget"] > 1: - # Insert exchange + if 'qie' in opt_fun.portals_parameters['target']['targets_evolve']: self.profiles_output.profiles['qei(MW/m^3)'] = profiles_portals_out.profiles['qei(MW/m^3)'] - if opt_fun.MODELparameters['Physics_options']["TypeTarget"] > 2: - # Insert radiation and fusion - for key in ['qbrem(MW/m^3)', 'qsync(MW/m^3)', 'qline(MW/m^3)', 'qfuse(MW/m^3)', 'qfusi(MW/m^3)']: - self.profiles_output.profiles[key] = profiles_portals_out.profiles[key] + if 'qrad' in opt_fun.portals_parameters['target']['targets_evolve']: + for key in ['qbrem(MW/m^3)', 'qsync(MW/m^3)', 'qline(MW/m^3)']: + self.profiles_output.profiles[key] = profiles_portals_out.profiles[key] + if 'qfus' in opt_fun.portals_parameters['target']['targets_evolve']: + for key in ['qfuse(MW/m^3)', 'qfusi(MW/m^3)']: + self.profiles_output.profiles[key] = profiles_portals_out.profiles[key] # -------------------------------------------------------------------------------------------- # Write to final input.gacode @@ -305,7 +309,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr last_radial_location_moved = False if change_last_radial_call and ('rhotop' in self.maestro_instance.parameters_trans_beat): - if 'RoaLocations' in self.MODELparameters: + if 'predicted_roa' in self.portals_parameters['solution']: print('\t\t- Using EPED pedestal top rho to select last radial location of PORTALS (in r/a)') @@ -317,40 +321,40 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr #roatop = roatop.round(3) # set the last value of the radial locations to the interpolated value - roatop_old = copy.deepcopy(self.MODELparameters["RoaLocations"][-1]) - self.MODELparameters["RoaLocations"][-1] = roatop - print(f'\t\t\t* Last radial location moved from r/a = {roatop_old} to {self.MODELparameters["RoaLocations"][-1]}') - print(f'\t\t\t* RoaLocations: {self.MODELparameters["RoaLocations"]}') + roatop_old = copy.deepcopy(self.portals_parameters['solution']["predicted_roa"][-1]) + self.portals_parameters['solution']["predicted_roa"][-1] = roatop + print(f'\t\t\t* Last radial location moved from r/a = {roatop_old} to {self.portals_parameters['solution']["predicted_roa"][-1]}') + print(f'\t\t\t* predicted_roa: {self.portals_parameters['solution']["predicted_roa"]}') - strKeys = 'RoaLocations' + strKeys = 'predicted_roa' else: print('\t\t- Using EPED pedestal top rho to select last radial location of PORTALS (in rho)') # set the last value of the radial locations to the interpolated value - rhotop_old = copy.deepcopy(self.MODELparameters["RhoLocations"][-1]) - self.MODELparameters["RhoLocations"][-1] = self.maestro_instance.parameters_trans_beat['rhotop'] - print(f'\t\t\t* Last radial location moved from rho = {rhotop_old} to {self.MODELparameters["RhoLocations"][-1]}') + rhotop_old = copy.deepcopy(self.portals_parameters['solution']['predicted_rho'][-1]) + self.portals_parameters['solution']['predicted_rho'][-1] = self.maestro_instance.parameters_trans_beat['rhotop'] + print(f'\t\t\t* Last radial location moved from rho = {rhotop_old} to {self.portals_parameters['solution']['predicted_rho'][-1]}') - strKeys = 'RhoLocations' + strKeys = 'predicted_rho' last_radial_location_moved = True # Check if I changed it previously and it hasn't moved if strKeys in self.maestro_instance.parameters_trans_beat: print(f'\t\t\t* {strKeys} in previous PORTALS beat: {self.maestro_instance.parameters_trans_beat[strKeys]}') - print(f'\t\t\t* {strKeys} in current PORTALS beat: {self.MODELparameters[strKeys]}') + print(f'\t\t\t* {strKeys} in current PORTALS beat: {self.portals_parameters['solution'][strKeys]}') - if abs(self.MODELparameters[strKeys][-1]-self.maestro_instance.parameters_trans_beat[strKeys][-1]) / self.maestro_instance.parameters_trans_beat[strKeys][-1] < minimum_relative_change_in_x: + if abs(self.portals_parameters['solution'][strKeys][-1]-self.maestro_instance.parameters_trans_beat[strKeys][-1]) / self.maestro_instance.parameters_trans_beat[strKeys][-1] < minimum_relative_change_in_x: print('\t\t\t* Last radial location was not moved') last_radial_location_moved = False - self.MODELparameters[strKeys][-1] = self.maestro_instance.parameters_trans_beat[strKeys][-1] + self.portals_parameters['solution'][strKeys][-1] = self.maestro_instance.parameters_trans_beat[strKeys][-1] # In the situation where the last radial location moves, I cannot reuse that surrogate data if last_radial_location_moved and reusing_surrogate_data: print('\t\t- Last radial location was moved, so surrogate data will not be reused for that specific location') - self.optimization_options['surrogate_options']["extrapointsModelsAvoidContent"] = ['_tar',f'_{len(self.MODELparameters[strKeys])}'] + self.optimization_options['surrogate_options']["extrapointsModelsAvoidContent"] = ['_tar',f'_{len(self.portals_parameters['solution'][strKeys])}'] self.try_flux_match_only_for_first_point = False def _inform_save(self): @@ -363,11 +367,11 @@ def _inform_save(self): # Standard PORTALS output try: stepSettings = portals_output.step.stepSettings - MODELparameters = portals_output.MODELparameters + portals_parameters = portals_output.portals_parameters # Converged in training case except AttributeError: stepSettings = portals_output.opt_fun_full.mitim_model.stepSettings - MODELparameters =portals_output.opt_fun_full.mitim_model.optimization_object.MODELparameters + portals_parameters =portals_output.opt_fun_full.mitim_model.optimization_object.portals_parameters max_value_neg_residual = stepSettings['optimization_options']['convergence_options']['stopping_criteria_parameters']['maximum_value'] self.maestro_instance.parameters_trans_beat['portals_neg_residual_obj'] = max_value_neg_residual @@ -379,25 +383,12 @@ def _inform_save(self): self.maestro_instance.parameters_trans_beat['portals_surrogate_data_file'] = fileTraining print(f'\t\t* Surrogate data saved for future beats: {IOtools.clipstr(fileTraining)}') - if 'RoaLocations' in MODELparameters: - self.maestro_instance.parameters_trans_beat['RoaLocations'] = MODELparameters['RoaLocations'] - print(f'\t\t* RoaLocations saved for future beats: {MODELparameters["RoaLocations"]}') - elif 'RhoLocations' in MODELparameters: - self.maestro_instance.parameters_trans_beat['RhoLocations'] = MODELparameters['RhoLocations'] - print(f'\t\t* RhoLocations saved for future beats: {MODELparameters["RhoLocations"]}') - - -def modify_dictionary(original, new): - for key in new: - # If something on the new dictionary is not in the original, add it - if key not in original: - original[key] = new[key] - # If it is a dictionary, go deeper - elif isinstance(new[key], dict): - modify_dictionary(original[key], new[key]) - # If it is not a dictionary, just replace the value - else: - original[key] = new[key] + if 'predicted_roa' in portals_parameters['solution']: + self.maestro_instance.parameters_trans_beat['predicted_roa'] = portals_parameters['solution']['predicted_roa'] + print(f'\t\t* predicted_roa saved for future beats: {portals_parameters['solution']["predicted_roa"]}') + elif 'predicted_rho' in portals_parameters['solution']: + self.maestro_instance.parameters_trans_beat['predicted_rho'] = portals_parameters['solution']['predicted_rho'] + print(f'\t\t* predicted_rho saved for future beats: {portals_parameters['solution']["predicted_rho"]}') # ----------------------------------------------------------------------------------------------------------------------- # Defaults to help MAESTRO @@ -420,9 +411,11 @@ def portals_beat_soft_criteria(portals_namelist): portals_namelist_soft['optimization_options']['convergence_options']["stopping_criteria_parameters"]["minimum_dvs_variation"] = [10, 3, 1.0] portals_namelist_soft['optimization_options']['convergence_options']["stopping_criteria_parameters"]["ricci_value"] = 0.15 - if 'MODELparameters' not in portals_namelist_soft: - portals_namelist_soft['MODELparameters'] = {} + if 'target' not in portals_namelist_soft: + portals_namelist_soft['target'] = {} + if 'options' not in portals_namelist_soft['target']: + portals_namelist_soft['target']['options'] = {} - portals_namelist_soft["MODELparameters"]["Physics_options"] = {"TypeTarget": 2} + portals_namelist_soft["target"]["options"]["targets_evolve"] = ["qie"] return portals_namelist_soft diff --git a/src/mitim_modules/portals/utils/PORTALSplot.py b/src/mitim_modules/portals/utils/PORTALSplot.py index ee3e71f7..11013155 100644 --- a/src/mitim_modules/portals/utils/PORTALSplot.py +++ b/src/mitim_modules/portals/utils/PORTALSplot.py @@ -1895,7 +1895,7 @@ def PORTALSanalyzer_plotSummary(self, fn=None, fn_color=None): lastRho=self.portals_parameters["solution"]["predicted_rho"][-1], alpha=alpha, useRoa=True, - RhoLocationsPlot=self.portals_parameters["solution"]["predicted_rho"], + predicted_rhoPlot=self.portals_parameters["solution"]["predicted_rho"], plotImpurity=self.runWithImpurity, plotRotation=self.runWithRotation, autoscale=i == 3, diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py index 6c9a97de..4c288c7c 100644 --- a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -512,13 +512,13 @@ def plot_gradients( ms=2, alpha=1.0, useRoa=False, - RhoLocationsPlot=None, + predicted_rhoPlot=None, plotImpurity=None, plotRotation=False, autoscale=True, ): - if RhoLocationsPlot is None: RhoLocationsPlot=[] + if predicted_rhoPlot is None: predicted_rhoPlot=[] if axs4 is None: plt.ion() @@ -713,7 +713,7 @@ def plot_gradients( GRAPHICStools.autoscale_y(ax, bottomy=0) cont += 2 - for x0 in RhoLocationsPlot: + for x0 in predicted_rhoPlot: ix = np.argmin(np.abs(self.profiles["rho(-)"] - x0)) for ax in axs4: ax.axvline(x=xcoord[ix], ls="--", lw=0.5, c=color) diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index a4188484..c762f630 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -62,19 +62,25 @@ "use_default": false, "portals_namelist" : { "portals_parameters": { - "force_zero_particle_flux": true, - "keep_full_model_folder": false, - "cores_per_tglf_instance": 1 - }, - "MODELparameters": { - "predicted_roa": [0.35,0.45,0.55,0.65,0.75,0.875,0.9], - "Physics_options": {"TypeTarget": 3}, - "transport_model": { - "TGLFsettings": 100, - "extraOptionsTGLF": {"USE_BPER": true} + "solution": { + "predicted_roa": [0.35,0.45,0.55,0.65,0.75,0.875,0.9], + "keep_full_model_folder": false + }, + "transport": { + "options": { + "run": { + "code_settings": 100, + "extraOptions": {"USE_BPER": true} + } + } + }, + "target": { + "options": { + "force_zero_particle_flux": true + } } }, - "INITparameters": { + "initialization_parameters": { "thermalize_fast": true, "quasineutrality": true }, diff --git a/tests/MAESTRO_workflow.py b/tests/MAESTRO_workflow.py index b206c428..ccb6eb01 100644 --- a/tests/MAESTRO_workflow.py +++ b/tests/MAESTRO_workflow.py @@ -3,7 +3,7 @@ from mitim_tools import __mitimroot__ from mitim_modules.maestro.scripts import run_maestro -cold_start = True +cold_start = False folder = __mitimroot__ / "tests" / "scratch" / "maestro_test" template = __mitimroot__ / "templates" / "maestro_namelist.json" From 9bc696a5428587bce5477a09a266216b0582765b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 13:52:02 -0400 Subject: [PATCH 242/385] Allow eped initializer in MAESTRO to be different than the full EPED --- .../maestro/scripts/run_maestro.py | 6 +++++- templates/maestro_namelist.json | 17 +++++++++++++++-- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index b93e9781..dfc94ee0 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -148,6 +148,10 @@ def profiles_postprocessing_fun(file_profs): return p beat_namelist['portals_parameters']['transport']['profiles_postprocessing_fun'] = profiles_postprocessing_fun + elif beat_type == "eped_initializer" and "eped_beat" in maestro_namelist["maestro"]: + print('Using the eped_beat namelist for the eped_initializer') + beat_namelist = maestro_namelist["maestro"]["eped_beat"]["eped_namelist"] + else: raise ValueError(f"[MITIM] {beat_type} beat not found in the MAESTRO namelist") @@ -212,7 +216,7 @@ def run_maestro_local( # **************************************************************************** if not creator_added: m.define_creator( - 'eped', + 'eped_initializer', BetaN = parameters_initialize["BetaN_initialization"], nu_ne = parameters_initialize["peaking_initialization"], **beat_namelists["eped"], diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index c762f630..ec18fc64 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -55,8 +55,21 @@ }, "ptop_multiplier": 1.0, "TioverTe": 1.0 - } - + } + }, + "eped_initializer_beat": { + "use_default": false, + "eped_namelist":{ + "nn_location": "$MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-MODEL-SPARC.keras", + "norm_location": "$MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-NORMALIZATION-SPARC.txt", + "corrections_set": { + "Bt": 12.2, + "R": 1.85, + "a": 0.57 + }, + "ptop_multiplier": 1.0, + "TioverTe": 1.0 + } }, "portals_beat": { "use_default": false, From 5572eb44a664d9fb0bb953405098898b28e7963b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 14:09:10 -0400 Subject: [PATCH 243/385] code_settings can now receive labels not numbers! --- .../gacode_tools/utils/GACODEdefaults.py | 58 ++++++++----------- templates/input.neo.models.json | 8 +++ templates/input.tglf.models.json | 2 +- templates/maestro_namelist.json | 2 +- templates/portals.namelist.yaml | 9 +-- tests/TGLF_workflow.py | 14 ++--- 6 files changed, 47 insertions(+), 46 deletions(-) create mode 100644 templates/input.neo.models.json diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index 56a6d6a1..6362c6e5 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -41,15 +41,7 @@ def addTGLFcontrol(code_settings, NS=2, minimal=False): ******************************************************************************** """ - with open(__mitimroot__ / "templates" / "input.tglf.models.json", "r") as f: - settings = json.load(f) - - if str(code_settings) in settings: - sett = settings[str(code_settings)] - for ikey in sett["controls"]: - options[ikey] = sett["controls"][ikey] - else: - print(f"\t- {code_settings = } not found in input.tglf.models.json, using defaults",typeMsg="w",) + options = add_code_settings(options, code_settings, models_file="input.tglf.models.json") return options @@ -57,51 +49,51 @@ def addNEOcontrol(code_settings,*args, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.neo.controls", caseInsensitive=False) + options = add_code_settings(options, code_settings, models_file="input.neo.models.json") + return options def addGXcontrol(code_settings,*args, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.gx.controls", caseInsensitive=False) - - """ - ******************************************************************************** - Standard sets of control parameters - (rest of parameters are as defaults) - ******************************************************************************** - """ - with open(__mitimroot__ / "templates" / "input.gx.models.json", "r") as f: - settings = json.load(f) + options = add_code_settings(options, code_settings, models_file="input.gx.models.json") - if str(code_settings) in settings: - sett = settings[str(code_settings)] - for ikey in sett["controls"]: - options[ikey] = sett["controls"][ikey] - else: - print("\t- code_settings not found in input.cgyro.models.json, using defaults",typeMsg="w") - return options def addCGYROcontrol(code_settings, rmin=None, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.cgyro.controls", caseInsensitive=False) - """ - ******************************************************************************** - Standard sets of control parameters - (rest of parameters are as defaults) - ******************************************************************************** - """ + options = add_code_settings(options, code_settings, models_file="input.cgyro.models.json") + + return options + +def add_code_settings(options,code_settings, models_file = "input.tglf.models.json"): - with open(__mitimroot__ / "templates" / "input.cgyro.models.json", "r") as f: + with open(__mitimroot__ / "templates" / models_file, "r") as f: settings = json.load(f) + found = False + + # Search by number first if str(code_settings) in settings: sett = settings[str(code_settings)] for ikey in sett["controls"]: options[ikey] = sett["controls"][ikey] + found = True else: - print("\t- code_settings not found in input.cgyro.models.json, using defaults",typeMsg="w") + # Search by label second + for ikey in settings: + if settings[ikey]["label"] == code_settings: + sett = settings[ikey] + for jkey in sett["controls"]: + options[jkey] = sett["controls"][jkey] + found = True + break + + if not found: + print(f"\t- {code_settings = } not found in {models_file}, using defaults",typeMsg="w") return options diff --git a/templates/input.neo.models.json b/templates/input.neo.models.json new file mode 100644 index 00000000..1e86edac --- /dev/null +++ b/templates/input.neo.models.json @@ -0,0 +1,8 @@ +{ + "0": { + "label": "Sonic", + "controls": { + "ROTATION_MODEL": 2 + } + } +} \ No newline at end of file diff --git a/templates/input.tglf.models.json b/templates/input.tglf.models.json index 56361582..1204e30a 100644 --- a/templates/input.tglf.models.json +++ b/templates/input.tglf.models.json @@ -50,7 +50,7 @@ } }, "100": { - "label": "settings as in ASTRA-TGLF framework", + "label": "SAT2astra", "controls": { "UNITS": "CGYRO", "USE_BPER": true, diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index ec18fc64..f0af4352 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -82,7 +82,7 @@ "transport": { "options": { "run": { - "code_settings": 100, + "code_settings": "SAT2astra", "extraOptions": {"USE_BPER": true} } } diff --git a/templates/portals.namelist.yaml b/templates/portals.namelist.yaml index 0138e63f..22947b07 100644 --- a/templates/portals.namelist.yaml +++ b/templates/portals.namelist.yaml @@ -86,7 +86,7 @@ transport: # Kwargs to be passed to the run command run: - code_settings: 6 + code_settings: "SAT3" extraOptions: {} # Kwargs to be passed to the read command @@ -112,7 +112,8 @@ transport: # *********************************** neo: - run: {} + run: + code_settings: "Sonic" read: {} @@ -125,7 +126,7 @@ transport: cgyro: run: - code_settings: 1 + code_settings: "Nonlinear" extraOptions: {} # Run type: normal (submit and wait), submit (submit and do not wait), prep (do not submit) @@ -147,7 +148,7 @@ transport: run: - code_settings: 1 + code_settings: "Nonlinear" extraOptions: {} # Run type: normal (submit and wait), submit (submit and do not wait), prep (do not submit) diff --git a/tests/TGLF_workflow.py b/tests/TGLF_workflow.py index 4c154920..439ab26a 100644 --- a/tests/TGLF_workflow.py +++ b/tests/TGLF_workflow.py @@ -17,7 +17,7 @@ tglf.run( "run1/", - code_settings=0, + code_settings='SAT1', cold_start=cold_start, runWaveForms = [0.67, 10.0], forceIfcold_start=True, @@ -25,31 +25,31 @@ slurm_setup={"cores": 4, "minutes": 1}, ) -tglf.read(label="ES (0)") +tglf.read(label="ES (SAT1)") tglf.run( "run2/", - code_settings=0, + code_settings='SAT1', cold_start=cold_start, forceIfcold_start=True, extraOptions={"USE_BPER": True, "USE_BPAR": True}, slurm_setup={"cores": 4, "minutes": 1}, ) -tglf.read(label="EM (0)") +tglf.read(label="EM (SAT1)") tglf.run( "run3/", - code_settings=6, + code_settings='SAT3', cold_start=cold_start, forceIfcold_start=True, extraOptions={"USE_BPER": True, "USE_BPAR": True}, slurm_setup={"cores": 4, "minutes": 1}, ) -tglf.read(label="EM (6)") +tglf.read(label="EM (SAT3)") -tglf.plot(labels=["EM (0)","ES (0)", "EM (6)"]) +tglf.plot(labels=["EM (SAT1)","ES (SAT1)", "EM (SAT3)"]) # Required if running in non-interactive mode tglf.fn.show() \ No newline at end of file From c47450c30ffc5dec13ee1e3c13b818512485693f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 14:19:10 -0400 Subject: [PATCH 244/385] Removal of historical additional_params_in_surrogate --- src/mitim_modules/maestro/utils/PORTALSbeat.py | 5 ----- src/mitim_modules/portals/utils/PORTALSinit.py | 7 ------- templates/portals.namelist.yaml | 5 ++--- 3 files changed, 2 insertions(+), 15 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index ba65c7a4..89a93f59 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -25,7 +25,6 @@ def prepare(self, use_previous_surrogate_data = False, try_flux_match_only_for_first_point = True, change_last_radial_call = False, - additional_params_in_surrogate = [], exploration_ranges = { 'ymax_rel': 1.0, 'ymin_rel': 1.0, @@ -68,8 +67,6 @@ def prepare(self, self.optimization_options = optimization_options self.initialization_parameters = initialization_parameters - self.additional_params_in_surrogate = additional_params_in_surrogate - self.exploration_ranges = exploration_ranges self.use_previous_surrogate_data = use_previous_surrogate_data self.change_last_radial_call = change_last_radial_call @@ -90,8 +87,6 @@ def run(self, **kwargs): portals_fun.portals_parameters = IOtools.deep_dict_update(portals_fun.portals_parameters, self.portals_parameters) portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.optimization_options) - portals_fun.portals_parameters['solution']['additional_params_in_surrogate'] = self.additional_params_in_surrogate - # Initialization now happens by the user from mitim_tools.gacode_tools.PROFILEStools import gacode_state p = gacode_state(self.fileGACODE) diff --git a/src/mitim_modules/portals/utils/PORTALSinit.py b/src/mitim_modules/portals/utils/PORTALSinit.py index 110785d6..0e889275 100644 --- a/src/mitim_modules/portals/utils/PORTALSinit.py +++ b/src/mitim_modules/portals/utils/PORTALSinit.py @@ -307,7 +307,6 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue allOuts = portals_fun.optimization_options["problem_options"]["ofs"] portals_transformation_variables = portals_fun.portals_parameters["solution"]["portals_transformation_variables"][ikey] portals_transformation_variables_trace = portals_fun.portals_parameters["solution"]["portals_transformation_variables_trace"][ikey] - additional_params_in_surrogate = portals_fun.portals_parameters["solution"]["additional_params_in_surrogate"] Variables = {} for output in allOuts: @@ -354,9 +353,6 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue and (not isAbsValFixed), } - for kkey in additional_params_in_surrogate: - Variations[kkey] = True - Variables[output] = [] for ikey in portals_transformation_variables: useThisOne = False @@ -392,9 +388,6 @@ def prepportals_transformation_variables(portals_fun, ikey, doNotFitOnFixedValue and (not isAbsValFixed), } - for kkey in additional_params_in_surrogate: - Variations[kkey] = True - Variables[output] = [] for ikey in portals_transformation_variables_trace: useThisOne = False diff --git a/templates/portals.namelist.yaml b/templates/portals.namelist.yaml index 22947b07..030c41c4 100644 --- a/templates/portals.namelist.yaml +++ b/templates/portals.namelist.yaml @@ -48,7 +48,8 @@ solution: # [Qe,Qi,Ge,Mt,GZ] multipliers to calculate scalarized function scalar_multipliers: [1.0, 1.0, 1.0, 1.0, 1.0] - # Physics-informed parameters to fit surrogates + # Physics-informed parameters to fit surrogates (numbers are the last iteration to consider that line; keys are the physics-informed parameters + # and the values are the corresponding variables to check if they have varied, that affect that input) portals_transformation_variables: 10: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0]} 30: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0], nuei: [te, ne], tite: [te, ti], w0_n: [w0]} @@ -60,8 +61,6 @@ solution: 30: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0], nuei: [te, ne], tite: [te, ti], w0_n: [w0], aLnZ: [aLnZ]} 10000: {aLte: [aLte], aLti: [aLti], aLne: [aLne], aLw0_n: [aLw0], nuei: [te, ne], tite: [te, ti], w0_n: [w0], beta_e: [te, ne], aLnZ: [aLnZ]} - additional_params_in_surrogate: [] - # ----------------------------------------------------------------- # Transport model parameters # ----------------------------------------------------------------- From e80202cefa3978cc846df59d8831631fbbedacc7 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 14:31:10 -0400 Subject: [PATCH 245/385] misc --- src/mitim_modules/maestro/MAESTROmain.py | 20 +++++++------------ .../maestro/utils/PORTALSbeat.py | 6 +++--- .../gacode_tools/utils/GACODEdefaults.py | 3 --- 3 files changed, 10 insertions(+), 19 deletions(-) diff --git a/src/mitim_modules/maestro/MAESTROmain.py b/src/mitim_modules/maestro/MAESTROmain.py index 899a6c33..403f2282 100644 --- a/src/mitim_modules/maestro/MAESTROmain.py +++ b/src/mitim_modules/maestro/MAESTROmain.py @@ -124,7 +124,7 @@ def define_creator(self, method, **kwargs_creator): ''' To initialize some profile functional form ''' - if method == 'eped': + if method == 'eped' or method == 'eped' or 'eped_initializer': self.beat.initialize.profile_creator = creator_from_eped(self.beat.initialize,**kwargs_creator) elif method == 'parameterization': self.beat.initialize.profile_creator = creator_from_parameterization(self.beat.initialize,**kwargs_creator) @@ -137,8 +137,7 @@ def define_creator(self, method, **kwargs_creator): # Beat operations # -------------------------------------------------------------------------------------------- - @mitim_timer( - lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Checker') + @mitim_timer(lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Checker') def check(self, beat_check = None, cold_start = False, **kwargs): ''' Note: @@ -173,8 +172,7 @@ def check(self, beat_check = None, cold_start = False, **kwargs): return output_file is not None - @mitim_timer( - lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Initializer', + @mitim_timer(lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Initializer', log_file = lambda self: self.folder_performance / "timing.jsonl") def initialize(self, *args, **kwargs): @@ -201,8 +199,7 @@ def initialize(self, *args, **kwargs): # First initialization, freeze engineering parameters self._freeze_parameters(profiles = PROFILEStools.gacode_state(self.beat.initialize.folder / 'input.gacode')) - @mitim_timer( - lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Preparation', + @mitim_timer(lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Preparation', log_file = lambda self: self.folder_performance / "timing.jsonl") def prepare(self, *args, **kwargs): @@ -223,8 +220,7 @@ def prepare(self, *args, **kwargs): else: print('\t\t- Skipping beat preparation because this beat was already run', typeMsg = 'i') - @mitim_timer( - lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Run + Finalization', + @mitim_timer(lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Run + Finalization', log_file = lambda self: self.folder_performance / "timing.jsonl") def run(self, **kwargs): @@ -271,8 +267,7 @@ def _freeze_parameters(self, profiles = None): self.profiles_with_engineering_parameters = copy.deepcopy(profiles) self.profiles_with_engineering_parameters.write_state(file= (self.folder_output / 'input.gacode_frozen')) - @mitim_timer( - lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Finalizing', + @mitim_timer(lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Finalizing', log_file = lambda self: self.folder_performance / "timing.jsonl") def finalize(self): @@ -291,8 +286,7 @@ def finalize(self): # Plotting operations # -------------------------------------------------------------------------------------------- - @mitim_timer( - lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Plotting') + @mitim_timer(lambda self: f'Beat #{self.counter_current} ({self.beat.name}) - Plotting') def plot(self, fn = None, num_beats = 2, only_beats = None, full_plot = True): print('*** Plotting MAESTRO ******************************************************************** ') diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 89a93f59..72fee92d 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -224,12 +224,12 @@ def merge_parameters(self): # Insert powers opt_fun = PORTALSanalysis.PORTALSanalyzer.from_folder(self.folder) - if 'qie' in opt_fun.portals_parameters['target']['targets_evolve']: + if 'qie' in opt_fun.portals_parameters['target']['options']['targets_evolve']: self.profiles_output.profiles['qei(MW/m^3)'] = profiles_portals_out.profiles['qei(MW/m^3)'] - if 'qrad' in opt_fun.portals_parameters['target']['targets_evolve']: + if 'qrad' in opt_fun.portals_parameters['target']['options']['targets_evolve']: for key in ['qbrem(MW/m^3)', 'qsync(MW/m^3)', 'qline(MW/m^3)']: self.profiles_output.profiles[key] = profiles_portals_out.profiles[key] - if 'qfus' in opt_fun.portals_parameters['target']['targets_evolve']: + if 'qfus' in opt_fun.portals_parameters['target']['options']['targets_evolve']: for key in ['qfuse(MW/m^3)', 'qfusi(MW/m^3)']: self.profiles_output.profiles[key] = profiles_portals_out.profiles[key] # -------------------------------------------------------------------------------------------- diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index 6362c6e5..f0e762e6 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -48,7 +48,6 @@ def addTGLFcontrol(code_settings, NS=2, minimal=False): def addNEOcontrol(code_settings,*args, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.neo.controls", caseInsensitive=False) - options = add_code_settings(options, code_settings, models_file="input.neo.models.json") return options @@ -56,7 +55,6 @@ def addNEOcontrol(code_settings,*args, **kwargs): def addGXcontrol(code_settings,*args, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.gx.controls", caseInsensitive=False) - options = add_code_settings(options, code_settings, models_file="input.gx.models.json") return options @@ -64,7 +62,6 @@ def addGXcontrol(code_settings,*args, **kwargs): def addCGYROcontrol(code_settings, rmin=None, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.cgyro.controls", caseInsensitive=False) - options = add_code_settings(options, code_settings, models_file="input.cgyro.models.json") return options From c5210b2b7f3d4dfc2167837252bd5f6bd1ff7c5f Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 15:06:00 -0400 Subject: [PATCH 246/385] misc --- .../maestro/scripts/run_maestro.py | 4 +- .../maestro/tmp_tests/maestro_test1.py | 136 ------------------ templates/portals.namelist.yaml | 29 ++-- tests/PORTALS_workflow.py | 10 +- 4 files changed, 23 insertions(+), 156 deletions(-) delete mode 100644 src/mitim_modules/maestro/tmp_tests/maestro_test1.py diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index dfc94ee0..51ffa6a4 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -219,7 +219,7 @@ def run_maestro_local( 'eped_initializer', BetaN = parameters_initialize["BetaN_initialization"], nu_ne = parameters_initialize["peaking_initialization"], - **beat_namelists["eped"], + **beat_namelists["eped_initializer"], **parameters_engineering ) m.initialize(BetaN = parameters_initialize["BetaN_initialization"], **geometry, **parameters_engineering) @@ -232,7 +232,7 @@ def run_maestro_local( run_namelist = {} if maestro_beats["beats"][0] in ["transp", "transp_soft"]: run_namelist = {'mpisettings' : {"trmpi": cpus, "toricmpi": cpus, "ptrmpi": 1}} - elif maestro_beats["beats"][0] in ["eped"]: + elif maestro_beats["beats"][0] in ["eped", "eped_initializer"]: run_namelist = {'cold_start': force_cold_start, 'cpus': cpus} m.prepare(**beat_namelists[maestro_beats["beats"][0]]) diff --git a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py b/src/mitim_modules/maestro/tmp_tests/maestro_test1.py deleted file mode 100644 index c84d8579..00000000 --- a/src/mitim_modules/maestro/tmp_tests/maestro_test1.py +++ /dev/null @@ -1,136 +0,0 @@ -from tensorflow._api.v2 import compat -from mitim_modules.maestro.MAESTROmain import maestro - -mfe_im_path = '/Users/pablorf/MFE-IM' -folder = '/Users/pablorf/PROJECTS/project_2024_ARCim/maestro_runs/runs_v2/arcV2B_run15/' - -# ----------------------------------------------------------------------------------------------------------------------- -# Parameters -# ----------------------------------------------------------------------------------------------------------------------- - -parameters = {'Ip_MA': 10.95, 'B_T': 10.8, 'Zeff': 1.5, 'PichT_MW': 18.0, 'neped_20' : 1.8 , 'Tesep_keV': 0.1, 'nesep_20': 2.0/3.0} -parameters_mix = {'DTplasma': True, 'lowZ_impurity': 9.0, 'impurity_ratio_WtoZ': 0.00286*0.5, 'minority': [1,1,0.02]} - -#initializer, geometry = 'freegs', {'R': 4.25, 'a': 1.17, 'kappa_sep': 1.77, 'delta_sep': 0.58, 'zeta_sep': 0.0, 'z0': 0.0} -initializer, geometry = 'geqdsk', {'geqdsk_file': f'{mfe_im_path}/private_data/ARCV2B.geqdsk', 'coeffs_MXH' : 7} - -BetaN_initialization = 1.5 - -# ----------------------------------------------------------------------------------------------------------------------- -# Namelists -# ----------------------------------------------------------------------------------------------------------------------- - -# To see what values this namelist can take: mitim_tools/transp_tools/NMLtools.py: _default_params() -transp_namelist = { - 'flattop_window': 1.0, # <--- To allow stationarity - 'extractAC': True, # <--- To extract TORIC and NUBEAM extra files - 'dtEquilMax_ms': 10.0, # Default - 'dtHeating_ms' : 5.0, # Default - 'dtOut_ms' : 10.0, - 'dtIn_ms' : 10.0, - 'nzones' : 60, - 'nzones_energetic' : 20, # Default but lower than what I used to use - 'nzones_distfun' : 10, # Default but lower than what I used to use - 'MCparticles' : 1e4, - 'toric_ntheta' : 64, # Default values of TORIC, but lower than what I used to use - 'toric_nrho' : 128, # Default values of TORIC, but lower than what I used to use - 'Pich': parameters['PichT_MW']>0.0, - 'DTplasma': parameters_mix['DTplasma'], - 'Minorities': parameters_mix['minority'], - "zlump" :[ [74.0, 184.0, 0.1*parameters_mix['impurity_ratio_WtoZ']], - [parameters_mix['lowZ_impurity'], parameters_mix['lowZ_impurity']*2, 0.1] ], - } - -# To see what values this namelist can take: mitim_modules/portals/PORTALSmain.py: __init__() -portals_namelist = { "main_parameters": {"launchEvaluationsAsSlurmJobs": True,"forceZeroParticleFlux": True, 'use_tglf_scan_trick': 0.02}, - "MODELparameters": { "predicted_roa": [0.35,0.55,0.75,0.875,0.9], - "ProfilesPredicted": ["te", "ti", "ne"], - "Physics_options": {"TypeTarget": 3}, - "transport_model": {"TGLFsettings": 6, "extraOptionsTGLF": {'USE_BPER':True}}}, - "INITparameters": {"thermalize_fast": True, "removeIons": [5,6], "quasineutrality": True}, - "optimization_options": { - "convergence_options": { - "maximum_iterations": 50, - "stopping_criteria_parameters": { - "maximum_value": 1e-3, - "maximum_value_is_rel": True, - }, - }, - "strategy_options": { - "AllowedExcursions":[0.0, 0.0] - }, - }, - "exploration_ranges": { - 'ymax_rel': 1.0, - 'ymin_rel': 0.9, - 'hardGradientLimits': [None,2] - } - } - -# To see what values this namelist can take: mitim_modules/maestro/utils/EPEDbeat.py: prepare() -eped_parameters = { 'nn_location': f'{mfe_im_path}/private_code_mitim/NN_DATA/EPED-NN-ARC/new-EPED-NN-MODEL-ARC.keras', - 'norm_location': f'{mfe_im_path}/private_code_mitim/NN_DATA/EPED-NN-ARC/EPED-NN-NORMALIZATION.txt'} - -# ----------------------------------------------------------------------------------------------------------------------- -# Workflow -# ----------------------------------------------------------------------------------------------------------------------- - -from mitim_tools.misc_tools.IOtools import mitim_timer - -@mitim_timer('\t\t* MAESTRO') -def run_maestro(): - m = maestro(folder, terminal_outputs = False) - - # TRANSP with only current diffusion - transp_namelist['flattop_window'] = 10.0 - transp_namelist['dtEquilMax_ms'] = 50.0 # Let the equilibrium evolve with long steps - transp_namelist['useNUBEAMforAlphas'] = False - transp_namelist['Pich'] = False - - m.define_beat('transp', initializer=initializer) - m.define_creator('eped', BetaN = BetaN_initialization, **eped_parameters,**parameters) - m.initialize(**geometry, **parameters) - m.prepare(**transp_namelist) - m.run() - - # TRANSP for toric and nubeam - transp_namelist['flattop_window'] = 0.5 - transp_namelist['dtEquilMax_ms'] = 10.0 - transp_namelist['useNUBEAMforAlphas'] = True - transp_namelist['Pich'] = True - - m.define_beat('transp') - m.prepare(**transp_namelist) - m.run() - - # EPED - m.define_beat('eped') - m.prepare(**eped_parameters) - m.run() - - # PORTALS - m.define_beat('portals') - m.prepare(**portals_namelist, change_last_radial_call = True) - m.run() - - # TRANSP - m.define_beat('transp') - m.prepare(**transp_namelist) - m.run() - - for i in range(9): - # EPED - m.define_beat('eped') - m.prepare(**eped_parameters) - m.run() - - # PORTALS - m.define_beat('portals') - m.prepare(**portals_namelist,use_previous_surrogate_data=i>0, change_last_radial_call = True) # Reuse the surrogate data if I'm not coming from a TRANSP run - m.run() - - m.finalize() - - return m - -m = run_maestro() \ No newline at end of file diff --git a/templates/portals.namelist.yaml b/templates/portals.namelist.yaml index 030c41c4..229468af 100644 --- a/templates/portals.namelist.yaml +++ b/templates/portals.namelist.yaml @@ -2,18 +2,21 @@ # This is a complete template for a PORTALS-TGLF+NEO profile prediction. # The user is welcomed to change any of the parameters below to fit their specific use case, # by copying this template namelist and instantiating the PORTALS object with its path: -# portals_fun = PORTALSmain.portals(folder, portals_namelist=PATH_TO_NAMELIST) +# +# portals_fun = PORTALSmain.portals(folder, portals_namelist=PATH_TO_NAMELIST) # # Alternatively, you can simply use the default parameters provided in this template by not -# specifying a namelist and then, in the launch script change the dictionary values: -# portals_fun = PORTALSmain.portals(folderWork) -# portals_fun.portals_parameters["solution"]['predicted_roa'] = [0.25, 0.45, 0.65, 0.85] -# portals_fun.portals_parameters["solution"]['predicted_channels'] = ["te", "ti", "ne", "nZ", 'w0'] -# ... +# specifying a namelist and then, in the launching script change the dictionary values: +# +# portals_fun = PORTALSmain.portals(folder) +# portals_fun.portals_parameters["solution"]['predicted_roa'] = [0.25, 0.45, 0.65, 0.85] +# portals_fun.portals_parameters["solution"]['predicted_channels'] = ["te", "ti", "ne", "nZ", 'w0'] +# ... # # The dictionary follows the same structure as the YAML namelist, except the optimization_options that must -# be passed as a separate dictionary: -# portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 10 +# be passed as a separate dictionary (because they are common among optimization problems): +# +# portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 10 # # ----------------------------------------------------------------- @@ -211,14 +214,14 @@ optimization_options: # Criterion 1: Maximum iterations maximum_iterations: 50 - # Stopping criteria function for the optimization process + # Stopping function for the rest of criteria stopping_criteria: "import::mitim_modules.portals.PORTALStools.stopping_criteria_portals" stopping_criteria_parameters: # Criterion 2: Residual reduction - maximum_value: 5.e-3 # Reducing residual by 200x is enough - maximum_value_is_rel: true + maximum_value: 5.e-3 # Reducing residual by 200x is enough for most PORTALS runs + maximum_value_is_rel: true # Indicates that maximum_value is relative to iteration #0 # Criterion 3: Variation of input parameters minimum_inputs_variation: [10, 5, 0.1] # After iteration 10, Check if 5 consecutive DVs are varying less than 0.1% from the rest that has been evaluated @@ -231,12 +234,12 @@ optimization_options: acquisition_options: # Relative improvement for stopping criteria of the acquisition optimization - relative_improvement_for_stopping: 1.e-2 # 100x is enough + relative_improvement_for_stopping: 1.e-2 # Reducing residual by 100x is enough # Type of acquisition function (Options: ["posterior_mean", "noisy_logei_mc", ...]) type: "posterior_mean" - # Optimizers to apply in sequence (Options: ["sr", "root", "botorch"]) + # Optimizers to apply sequentially (Options: ["sr", "root", "botorch"]) optimizers: ["sr", "root"] surrogate_options: diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index 04e56b2d..1a8c20a5 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -20,11 +20,11 @@ # Let's not consume the entire computer resources when running test... limit threads torch.set_num_threads(8) -# -------------------------------------------------------------------------------------------- +# --------------------------------------------------------------------------------------------------------------------- # Optimization Class -# -------------------------------------------------------------------------------------------- +# --------------------------------------------------------------------------------------------------------------------- -# Initialize class +# Initialize class with the default namelist in templates/portals.namelist.yaml but modify some of its parameters portals_fun = PORTALSmain.portals(folderWork) portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 portals_fun.optimization_options["initialization_options"]["initial_training"] = 2 @@ -43,9 +43,9 @@ # Prepare run portals_fun.prep(plasma_state) -# -------------------------------------------------------------------------------------------- +# --------------------------------------------------------------------------------------------------------------------- # Run -# -------------------------------------------------------------------------------------------- +# --------------------------------------------------------------------------------------------------------------------- # Run mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=cold_start, askQuestions=False) From 91f4301f3ffc85d1c5bfb401cd68c6e17a87251a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 15:13:14 -0400 Subject: [PATCH 247/385] TGLFsettings for most of MITIM is now code_settings --- docs/capabilities/tglf_capabilities.rst | 6 +++--- docs/capabilities/vitals_capabilities.rst | 4 ++-- regressions/portals_regressions.py | 2 +- src/mitim_modules/portals/scripts/runTGLF.py | 6 +++--- .../portals/scripts/runTGLFdrivesfromPORTALS.py | 4 ++-- .../portals/utils/PORTALSanalysis.py | 16 ++++++++-------- .../powertorch/physics_models/transport_tglf.py | 2 +- src/mitim_modules/vitals/VITALSmain.py | 4 ++-- src/mitim_tools/gacode_tools/scripts/run_tglf.py | 6 +++--- .../gacode_tools/utils/GACODEdefaults.py | 4 ++-- src/mitim_tools/plasmastate_tools/MITIMstate.py | 4 ++-- src/mitim_tools/transp_tools/CDFtools.py | 16 ++++++++-------- src/mitim_tools/transp_tools/NMLtools.py | 6 +++--- .../transp_tools/utils/TRANSPhelpers.py | 4 ++-- tests/TGLFfull_workflow.py | 6 +++--- tests/TGLFscan_workflow.py | 4 ++-- tests/VITALS_workflow.py | 6 +++--- tutorials/PORTALS_tutorial.py | 2 +- tutorials/TGLF_tutorial.py | 14 +++++++------- 19 files changed, 58 insertions(+), 58 deletions(-) diff --git a/docs/capabilities/tglf_capabilities.rst b/docs/capabilities/tglf_capabilities.rst index 4cdb2f76..e411a833 100644 --- a/docs/capabilities/tglf_capabilities.rst +++ b/docs/capabilities/tglf_capabilities.rst @@ -76,7 +76,7 @@ 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 ``cold_start`` 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`` indicates the base case to start with. +First, the argument ``code_settings`` 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): @@ -84,14 +84,14 @@ For example, the following two commands will run TGLF with saturation rule numbe .. code-block:: python tglf.run( subfolder = 'yes_em_folder', - TGLFsettings = 5, + code_settings = 5, extraOptions = {}, cold_start = False ) tglf.read( label = 'yes_em' ) tglf.run( subfolder = 'no_em_folder', - TGLFsettings = 5, + code_settings = 5, extraOptions = {'USE_BPER':False}, cold_start = False ) diff --git a/docs/capabilities/vitals_capabilities.rst b/docs/capabilities/vitals_capabilities.rst index 7bf157d5..62f08239 100644 --- a/docs/capabilities/vitals_capabilities.rst +++ b/docs/capabilities/vitals_capabilities.rst @@ -39,7 +39,7 @@ As a starting point of VITALS, you need to prepare and run TGLF for the base cas tglf = TGLFtools.TGLF( rhos = [ rho ] ) cdf = tglf.prep( folder, inputgacode = inputgacode_file) - tglf.run( subfolder = 'run_base', TGLFsettings = 5) + tglf.run( subfolder = 'run_base', code_settings = 5) tglf.read( label = 'run_base' ) @@ -120,7 +120,7 @@ Once the VITALS object has been created, parameters such as the TGLF control inp .. code-block:: python - vitals_fun.TGLFparameters['TGLFsettings'] = 5 + vitals_fun.TGLFparameters['code_settings'] = 5 vitals_fun.TGLFparameters['extraOptions'] = {} .. note:: diff --git a/regressions/portals_regressions.py b/regressions/portals_regressions.py index 2a17dfcb..1afbe7b3 100644 --- a/regressions/portals_regressions.py +++ b/regressions/portals_regressions.py @@ -97,7 +97,7 @@ def conditions_regressions(variables): portals_fun.portals_parameters["initialization"]["remove_fast"] = True portals_fun.portals_parameters["initialization"]["quasineutrality"] = True portals_fun.portals_parameters["initialization"]["enforce_same_aLn"] = True - portals_fun.portals_parameters["transport"]["options"]["TGLFsettings"] = 2 + portals_fun.portals_parameters["transport"]["options"]["code_settings"] = 2 portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti", "ne"] diff --git a/src/mitim_modules/portals/scripts/runTGLF.py b/src/mitim_modules/portals/scripts/runTGLF.py index f92a7421..5d6cd8d2 100644 --- a/src/mitim_modules/portals/scripts/runTGLF.py +++ b/src/mitim_modules/portals/scripts/runTGLF.py @@ -43,7 +43,7 @@ # --- Workflow portals = PORTALSanalysis.PORTALSanalyzer.from_folder(folder) -tglf, TGLFsettings, extraOptions = portals.extractTGLF(positions=pos, evaluation=ev, modified_profiles=True, cold_start=cold_start) +tglf, code_settings, extraOptions = portals.extractTGLF(positions=pos, evaluation=ev, modified_profiles=True, cold_start=cold_start) if not drives: varUpDown = np.linspace(1.0 - var, 1.0 + var, num) @@ -54,7 +54,7 @@ subfolder="scan", variable=param, varUpDown=varUpDown, - TGLFsettings=TGLFsettings, + code_settings=code_settings, extraOptions=extraOptions, cold_start=cold_start, runWaveForms=wf, @@ -73,7 +73,7 @@ resolutionPoints=5, variation=var, variablesDrives=["RLTS_1", "RLTS_2", "RLNS_1", "XNUE", "TAUS_2", "BETAE"], - TGLFsettings=TGLFsettings, + code_settings=code_settings, extraOptions=extraOptions, cold_start=cold_start, runWaveForms=wf, diff --git a/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py b/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py index a0812ed8..0a24c090 100644 --- a/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py +++ b/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py @@ -38,14 +38,14 @@ # --- Workflow portals = PORTALSanalysis.PORTALSanalyzer.from_folder(folder) -tglf, TGLFsettings, extraOptions = portals.extractTGLF(positions=pos, evaluation=ev, modified_profiles=True, cold_start=cold_start) +tglf, code_settings, extraOptions = portals.extractTGLF(positions=pos, evaluation=ev, modified_profiles=True, cold_start=cold_start) tglf.runScanTurbulenceDrives( subfolder="turb", resolutionPoints=num, variation=var, variablesDrives=["RLTS_1", "RLTS_2", "RLNS_1", "XNUE", "TAUS_2", "BETAE"], - TGLFsettings=TGLFsettings, + code_settings=code_settings, extraOptions=extraOptions, cold_start=cold_start, runWaveForms=wf, diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 9ed37409..54d6150e 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -544,7 +544,7 @@ def extractTGYRO(self, folder=None, cold_start=False, evaluation=0, modified_pro folder, profilesclass_custom=profiles, cold_start=cold_start, forceIfcold_start=True ) - TGLFsettings = self.portals_parameters["transport"]["options"]["TGLFsettings"] + code_settings = self.portals_parameters["transport"]["options"]["code_settings"] extraOptionsTGLF = self.portals_parameters["transport"]["options"]["extraOptionsTGLF"] PredictionSet = [ int("te" in self.portals_parameters["solution"]["predicted_channels"]), @@ -552,7 +552,7 @@ def extractTGYRO(self, folder=None, cold_start=False, evaluation=0, modified_pro int("ne" in self.portals_parameters["solution"]["predicted_channels"]), ] - return tgyro, self.rhos, PredictionSet, TGLFsettings, extraOptionsTGLF + return tgyro, self.rhos, PredictionSet, code_settings, extraOptionsTGLF def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=False, modified_profiles=False): if evaluation is None: @@ -593,10 +593,10 @@ def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=F tglf = TGLFtools.TGLF(rhos=rhos) _ = tglf.prep_using_tgyro(folder, cold_start=cold_start, inputgacode=inputgacode) - TGLFsettings = self.portals_parameters["transport"]["options"]["TGLFsettings"] + code_settings = self.portals_parameters["transport"]["options"]["code_settings"] extraOptions = self.portals_parameters["transport"]["options"]["extraOptionsTGLF"] - return tglf, TGLFsettings, extraOptions + return tglf, code_settings, extraOptions # **************************************************************************** # UTILITIES for post-analysis @@ -613,7 +613,7 @@ def runTGLFfull( ): """ 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, + This is convenient if I want to re=run TGLF with different settings, e.g. different code_settings, that you can provide as keyword arguments. """ @@ -625,14 +625,14 @@ def runTGLFfull( ranges = [self.ibest] if onlyBest else range(self.ilast + 1) for ev in ranges: - tglf, TGLFsettings, extraOptions = self.extractTGLF( + tglf, code_settings, extraOptions = self.extractTGLF( folder=folder / f"Evaluation.{ev}", evaluation=ev, cold_start=cold_start ) kwargsTGLF_this = copy.deepcopy(kwargsTGLF) - if "TGLFsettings" not in kwargsTGLF_this: - kwargsTGLF_this["TGLFsettings"] = TGLFsettings + if "code_settings" not in kwargsTGLF_this: + kwargsTGLF_this["code_settings"] = code_settings if "extraOptions" not in kwargsTGLF_this: kwargsTGLF_this["extraOptions"] = extraOptions diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 39c2a973..48be9efb 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -228,7 +228,7 @@ def _run_tglf_uncertainty_model( variablesDrives = variables_to_scan, varUpDown = relative_scan, minimum_delta_abs = minimum_delta_abs, - TGLFsettings = code_settings, + code_settings = code_settings, extraOptions = extraOptions, ApplyCorrections = False, add_baseline_to = 'none', diff --git a/src/mitim_modules/vitals/VITALSmain.py b/src/mitim_modules/vitals/VITALSmain.py index 2ada4f9d..e265ce58 100644 --- a/src/mitim_modules/vitals/VITALSmain.py +++ b/src/mitim_modules/vitals/VITALSmain.py @@ -51,7 +51,7 @@ def __init__(self, folder, namelist=None, **kwargs): # Default (please change to your desire after instancing the object) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.TGLFparameters = {"TGLFsettings": 2, "extraOptions": {}} + self.TGLFparameters = {"code_settings": 2, "extraOptions": {}} self.VITALSparameters = { "rel_error": 0.02, # Standard deviation (relative to value) @@ -304,7 +304,7 @@ def runTGLF( tglf.run( subfolder=f"{folder_label}", cold_start=cold_start, - TGLFsettings=self.TGLFparameters["TGLFsettings"], + code_settings=self.TGLFparameters["code_settings"], forceIfcold_start=True, extraOptions=extraOptions, multipliers=multipliers, diff --git a/src/mitim_tools/gacode_tools/scripts/run_tglf.py b/src/mitim_tools/gacode_tools/scripts/run_tglf.py index a2dc3869..c535ff8f 100644 --- a/src/mitim_tools/gacode_tools/scripts/run_tglf.py +++ b/src/mitim_tools/gacode_tools/scripts/run_tglf.py @@ -52,7 +52,7 @@ def main(): # ------------------------------------------------------------------------------ if drives: - tglf.runScanTurbulenceDrives(subfolder="scan_turb", TGLFsettings=None) + tglf.runScanTurbulenceDrives(subfolder="scan_turb", code_settings=None) tglf.plotScanTurbulenceDrives(label="scan_turb") elif scan is not None: @@ -60,14 +60,14 @@ def main(): subfolder="scan1", variable=scan, varUpDown=np.linspace(0.2, 2.0, 5), - TGLFsettings=None, + code_settings=None, cold_start=cold_start, ) tglf.readScan(label="scan1", variable=scan) tglf.plotScan(labels=["scan1"], variableLabel=scan) else: - tglf.run(subfolder="run1", TGLFsettings=None, cold_start=cold_start) + tglf.run(subfolder="run1", code_settings=None, cold_start=cold_start) tglf.read(label="run1") tglf.plot(labels=["run1"]) diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index f0e762e6..84947df8 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -95,9 +95,9 @@ def add_code_settings(options,code_settings, models_file = "input.tglf.models.js return options -def TGLFinTRANSP(TGLFsettings, NS=3): +def TGLFinTRANSP(code_settings, NS=3): - TGLFoptions = addTGLFcontrol(TGLFsettings, NS=NS) + TGLFoptions = addTGLFcontrol(code_settings, NS=NS) """ ------------------------------------------------------------------------------------------------------ diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index b245816f..9eb3995b 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2202,7 +2202,7 @@ def csv(self, file="input.gacode.xlsx"): # Code conversions # ************************************************************************************************************************************************ - def to_tglf(self, r=[0.5], TGLFsettings=1, r_is_rho = True): + def to_tglf(self, r=[0.5], code_settings=1, r_is_rho = True): # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function @@ -2285,7 +2285,7 @@ def interpolator(y): # Controls come from options # --------------------------------------------------------------------------------------------------------------------------------------- - controls = GACODEdefaults.addTGLFcontrol(TGLFsettings) + controls = GACODEdefaults.addTGLFcontrol(code_settings) # --------------------------------------------------------------------------------------------------------------------------------------- # Species come from profiles diff --git a/src/mitim_tools/transp_tools/CDFtools.py b/src/mitim_tools/transp_tools/CDFtools.py index cb162e76..27e69d5d 100644 --- a/src/mitim_tools/transp_tools/CDFtools.py +++ b/src/mitim_tools/transp_tools/CDFtools.py @@ -513,7 +513,7 @@ def evaluateReactorMetrics(self, ReactorTextFile=None, EvaluateExtraAnalysis=Non """ # ****** Settings ********* - TGLFsettings = 5 + code_settings = 5 d_perp_cm = {0.7: 0.757 / np.sqrt(2) / (np.cos(11 * (np.pi / 180)))} # ************************* @@ -540,7 +540,7 @@ def evaluateReactorMetrics(self, ReactorTextFile=None, EvaluateExtraAnalysis=Non quantity=quantity, rho=location, time=self.t[index], - TGLFsettings=TGLFsettings, + code_settings=code_settings, d_perp_cm=d_perp_cm, ) success = success and True @@ -14122,7 +14122,7 @@ def transportAnalysis( rho=0.50, time=None, avTime=0.0, - TGLFsettings=1, + code_settings=1, d_perp_cm=None, ): if time is None: @@ -14147,7 +14147,7 @@ def transportAnalysis( subfolder="chi_per", label="chi_pert", analysisType="e", - TGLFsettings=TGLFsettings, + code_settings=code_settings, ) value = self.TGLFstd[int(time * 1000)].scans["chi_pert"]["chi_inc"][0] @@ -14160,7 +14160,7 @@ def transportAnalysis( subfolder="impurity", label="impurity", analysisType="Z", - TGLFsettings=TGLFsettings, + code_settings=code_settings, trace=addTrace, ) @@ -14174,7 +14174,7 @@ def transportAnalysis( if "FLUC" in typeAnalysis: self.TGLFstd[int(time * 1000)].run( subfolder="fluctuations", - TGLFsettings=TGLFsettings, + code_settings=code_settings, forceIfcold_start=True, ) @@ -14676,7 +14676,7 @@ def compareChiPert( time=None, rhoRange=[0.4, 0.8], timeRange=0.5, - TGLFsettings=1, + code_settings=1, cold_start=False, plotYN=True, ): @@ -14700,7 +14700,7 @@ def compareChiPert( subfolder="chi_per", label="chi_pert", analysisType="e", - TGLFsettings=TGLFsettings, + code_settings=code_settings, cold_start=cold_start, cdf_open=self, ) diff --git a/src/mitim_tools/transp_tools/NMLtools.py b/src/mitim_tools/transp_tools/NMLtools.py index d53d92b1..94cb966b 100644 --- a/src/mitim_tools/transp_tools/NMLtools.py +++ b/src/mitim_tools/transp_tools/NMLtools.py @@ -180,7 +180,7 @@ def _default_params(self,**transp_params): self.grTGLF = transp_params.get("grTGLF",False) self.Te_edge = transp_params.get("Te_edge",80.0) self.Ti_edge = transp_params.get("Ti_edge",80.0) - self.TGLFsettings = transp_params.get("TGLFsettings",5) + self.code_settings = transp_params.get("code_settings",5) def populate(self, **transp_params): @@ -1335,8 +1335,8 @@ def addGLF23(self): self.contents_ptr_glf23 = "\n".join(lines) + "\n" def addTGLF(self): - TGLFoptions = GACODEdefaults.TGLFinTRANSP(self.TGLFsettings) - print(f"\t- Adding TGLF control parameters with TGLFsettings = {self.TGLFsettings}") + TGLFoptions = GACODEdefaults.TGLFinTRANSP(self.code_settings) + print(f"\t- Adding TGLF control parameters with code_settings = {self.code_settings}") lines = [ "!------ TGLF namelist", diff --git a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py index 428e6483..a0f46799 100644 --- a/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py +++ b/src/mitim_tools/transp_tools/utils/TRANSPhelpers.py @@ -1355,7 +1355,7 @@ def default_nml( "Fuel": 2.5, }, pservers=[1, 1, 0], - TGLFsettings=5, + code_settings=5, useMMX = False, isolver = False, grTGLF = False # Disable by default because it takes disk space and time... enable for 2nd preditive outside of this routine @@ -1414,7 +1414,7 @@ def default_nml( AddHe4ifDT=AddHe4ifDT, isolver=isolver, PTsolver=True, - TGLFsettings=TGLFsettings, + code_settings=code_settings, grTGLF=grTGLF, **transp_params ) diff --git a/tests/TGLFfull_workflow.py b/tests/TGLFfull_workflow.py index 3a90abd3..5c7cf628 100644 --- a/tests/TGLFfull_workflow.py +++ b/tests/TGLFfull_workflow.py @@ -22,7 +22,7 @@ tglf.run( subfolder="runSAT2", - TGLFsettings=5, + code_settings=5, runWaveForms=[0.1,0.3], cold_start=cold_start, forceIfcold_start=True, @@ -31,7 +31,7 @@ tglf.run( subfolder="runSAT0", - TGLFsettings=2, + code_settings=2, runWaveForms=[0.5], cold_start=cold_start, forceIfcold_start=True, @@ -40,7 +40,7 @@ tglf.run( subfolder="runSAT3", - TGLFsettings=6, + code_settings=6, runWaveForms=[0.5], cold_start=cold_start, forceIfcold_start=True, diff --git a/tests/TGLFscan_workflow.py b/tests/TGLFscan_workflow.py index 0e93fd90..1fdf7ebe 100644 --- a/tests/TGLFscan_workflow.py +++ b/tests/TGLFscan_workflow.py @@ -17,7 +17,7 @@ tglf.prep(input_gacode,folder, cold_start=cold_start) tglf.run_scan( subfolder = 'scan1', - TGLFsettings = None, + code_settings = None, cold_start = cold_start, runWaveForms = [0.67, 10.0], variable = 'RLTS_1', @@ -30,7 +30,7 @@ tglf.runScanTurbulenceDrives( subfolder = 'turb_drives', - TGLFsettings = None, + code_settings = None, resolutionPoints=3, cold_start = cold_start) diff --git a/tests/VITALS_workflow.py b/tests/VITALS_workflow.py index cd6f3a29..36d72caf 100644 --- a/tests/VITALS_workflow.py +++ b/tests/VITALS_workflow.py @@ -18,7 +18,7 @@ os.system(f"rm -r {folderWork}") rho = 0.5 -TGLFsettings = 2 +code_settings = 2 dvs = ["RLTS_1", "RLTS_2", "RLNS_1", "ZEFF", "TAUS_2"] ofs = ["Qe", "Qi", "TeFluct", "neTe"] @@ -31,7 +31,7 @@ tglf = TGLFtools.TGLF(rhos=[rho]) cdf = tglf.prep_using_tgyro(folderWork, cold_start=cold_start, inputgacode=inputgacode) -tglf.run(subfolder="run_base/", TGLFsettings=TGLFsettings, cold_start=cold_start) +tglf.run(subfolder="run_base/", code_settings=code_settings, cold_start=cold_start) # ******************************************************************************** # Then, add experimental data of fluctuation information and error bars @@ -70,7 +70,7 @@ vitals_fun = VITALSmain.vitals(folderWork) vitals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 -vitals_fun.TGLFparameters["TGLFsettings"] = TGLFsettings +vitals_fun.TGLFparameters["code_settings"] = code_settings vitals_fun.prep(file, rho, ofs, dvs, dvs_min, dvs_max) diff --git a/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index fb190482..732c13e1 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -25,7 +25,7 @@ # TGLF specifications portals_fun.portals_parameters["transport"]["options"] = { - "TGLFsettings": 6, # Check out templates/input.tglf.models.json for more options + "code_settings": 6, # Check out templates/input.tglf.models.json for more options "extraOptionsTGLF": {"USE_BPER": False} # Turn off BPER } diff --git a/tutorials/TGLF_tutorial.py b/tutorials/TGLF_tutorial.py index 12601fb6..6b73f070 100644 --- a/tutorials/TGLF_tutorial.py +++ b/tutorials/TGLF_tutorial.py @@ -21,7 +21,7 @@ # Run TGLF in subfolder tglf.run( subfolder="yes_em_folder", - TGLFsettings=5, + code_settings=5, extraOptions={}, cold_start=False ) @@ -32,7 +32,7 @@ # Run TGLF in a different subfolder with different settings tglf.run( subfolder="no_em_folder", - TGLFsettings=5, + code_settings=5, extraOptions={"USE_BPER": False}, cold_start=False, ) @@ -50,7 +50,7 @@ ''' tglf.runScan( subfolder = 'scan1', - TGLFsettings = 5, + code_settings = 5, cold_start = False, variable = 'RLTS_1', varUpDown = np.linspace(0.5,1.5,3)) @@ -58,7 +58,7 @@ tglf.runScan( subfolder = 'scan2', - TGLFsettings = 5, + code_settings = 5, cold_start = False, variable = 'RLTS_2', varUpDown = np.linspace(0.5,1.5,3)) @@ -75,7 +75,7 @@ tglf.runScanTurbulenceDrives( subfolder = 'turb_drives', - TGLFsettings = 5, + code_settings = 5, cold_start = False) tglf.plotScanTurbulenceDrives(label='turb_drives') @@ -89,7 +89,7 @@ tglf.runAnalysis( subfolder = 'chi_e', analysisType = 'chi_e', - TGLFsettings = 5, + code_settings = 5, cold_start = False, label = 'chi_eu') @@ -105,7 +105,7 @@ tglf.run( subfolder = f'settings{i}', runWaveForms = [0.67], - TGLFsettings = i, + code_settings = i, cold_start = False) tglf.read(label=f'settings{i}') From e2a47122c31d6caee9d7d765b9b8dc674bbe834a Mon Sep 17 00:00:00 2001 From: pabloprf Date: Fri, 29 Aug 2025 15:16:03 -0400 Subject: [PATCH 248/385] bug fix --- src/mitim_modules/maestro/utils/PORTALSbeat.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 72fee92d..e57ded7f 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -407,10 +407,10 @@ def portals_beat_soft_criteria(portals_namelist): portals_namelist_soft['optimization_options']['convergence_options']["stopping_criteria_parameters"]["ricci_value"] = 0.15 if 'target' not in portals_namelist_soft: - portals_namelist_soft['target'] = {} + portals_namelist_soft["portals_parameters"]['target'] = {} if 'options' not in portals_namelist_soft['target']: - portals_namelist_soft['target']['options'] = {} + portals_namelist_soft["portals_parameters"]['target']['options'] = {} - portals_namelist_soft["target"]["options"]["targets_evolve"] = ["qie"] + portals_namelist_soft["portals_parameters"]["target"]["options"]["targets_evolve"] = ["qie"] return portals_namelist_soft From c2f6ed63ebe52c83b899a6429c64cdb718a8c573 Mon Sep 17 00:00:00 2001 From: pabloprf Date: Fri, 29 Aug 2025 15:49:11 -0400 Subject: [PATCH 249/385] bug fixes --- src/mitim_modules/maestro/scripts/run_maestro.py | 2 +- src/mitim_modules/maestro/utils/PORTALSbeat.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index 51ffa6a4..7ba52d64 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -92,7 +92,7 @@ def parse_maestro_nml(file_path): beat_namelists = {} - for beat_type in ["eped", "transp", "transp_soft", "portals", "portals_soft"]: + for beat_type in ["eped","eped_initializer", "transp", "transp_soft", "portals", "portals_soft"]: if f"{beat_type}_beat" in maestro_namelist["maestro"]: diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index e57ded7f..7bc79d84 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -28,7 +28,7 @@ def prepare(self, exploration_ranges = { 'ymax_rel': 1.0, 'ymin_rel': 1.0, - 'hardGradientLimits': [0,2] + 'yminymax_atleast': [0,2] }, portals_parameters = {}, initialization_parameters = {}, @@ -406,9 +406,9 @@ def portals_beat_soft_criteria(portals_namelist): portals_namelist_soft['optimization_options']['convergence_options']["stopping_criteria_parameters"]["minimum_dvs_variation"] = [10, 3, 1.0] portals_namelist_soft['optimization_options']['convergence_options']["stopping_criteria_parameters"]["ricci_value"] = 0.15 - if 'target' not in portals_namelist_soft: + if 'target' not in portals_namelist_soft["portals_parameters"]: portals_namelist_soft["portals_parameters"]['target'] = {} - if 'options' not in portals_namelist_soft['target']: + if 'options' not in portals_namelist_soft["portals_parameters"]['target']: portals_namelist_soft["portals_parameters"]['target']['options'] = {} portals_namelist_soft["portals_parameters"]["target"]["options"]["targets_evolve"] = ["qie"] From b6595bfc0fd8099cc81d3784941eee2cc95dd2ea Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 15:50:05 -0400 Subject: [PATCH 250/385] misc --- templates/maestro_namelist.json | 2 +- tests/EPED_workflow.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index f0af4352..93780736 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -59,7 +59,7 @@ }, "eped_initializer_beat": { "use_default": false, - "eped_namelist":{ + "eped_initializer_namelist":{ "nn_location": "$MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-MODEL-SPARC.keras", "norm_location": "$MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-NORMALIZATION-SPARC.txt", "corrections_set": { diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py index c8740354..93aea967 100644 --- a/tests/EPED_workflow.py +++ b/tests/EPED_workflow.py @@ -30,10 +30,11 @@ scan_param = {'variable': 'neped', 'values': [15.0, 30.0, 45.0, 60.0, 75.0]}, keep_nsep_ratio = 0.4, nproc_per_run = 64, - cold_start = True, + cold_start = cold_start, ) eped.read(subfolder='case1') eped.plot(labels=['case1']) -plt.show() \ No newline at end of file + +eped.fn.show() From 6bdd664e7aabcb24fe3572f68d16f7b704731d6b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 16:00:24 -0400 Subject: [PATCH 251/385] Removed regressions for now, I will come back to it --- regressions/data/input.gacode | 1755 ---------------------------- regressions/portals_regressions.py | 155 --- tutorials/TGLF_tutorial.py | 18 +- 3 files changed, 10 insertions(+), 1918 deletions(-) delete mode 100644 regressions/data/input.gacode delete mode 100644 regressions/portals_regressions.py diff --git a/regressions/data/input.gacode b/regressions/data/input.gacode deleted file mode 100644 index 403d660b..00000000 --- a/regressions/data/input.gacode +++ /dev/null @@ -1,1755 +0,0 @@ -# *original : Thu May 27 12:10:51 EDT 2021 -# *statefile : 10001.cdf -# *gfile : 10001.geq 27May2021 t~ 2.50000 -# *cerfile : null -# *vgen : null -# *tgyro : null -# -# nexp -41 -# nion -4 -# shot -12345 -# name -D C N D -# type -[therm] [therm] [therm] [fast] -# masse - 5.4488741E-04 -# mass - 2.0005209E+00 1.1930011E+01 1.3912254E+01 2.0005209E+00 -# ze --1.0000000E+00 -# z - 1.0000000E+00 6.0000000E+00 7.0000000E+00 1.0000000E+00 -# torfluxa | Wb/radian --4.9977993E-01 -# rcentr | m - 1.6118747E+00 -# bcentr | T --2.4105557E+00 -# current | MA --9.9696344E-01 -# rho | - - 1 0.0000000E+00 - 2 2.5000000E-02 - 3 5.0000000E-02 - 4 7.5000000E-02 - 5 1.0000000E-01 - 6 1.2500000E-01 - 7 1.5000000E-01 - 8 1.7500000E-01 - 9 2.0000000E-01 - 10 2.2500000E-01 - 11 2.5000000E-01 - 12 2.7500000E-01 - 13 3.0000000E-01 - 14 3.2500000E-01 - 15 3.5000000E-01 - 16 3.7500000E-01 - 17 4.0000000E-01 - 18 4.2500000E-01 - 19 4.5000000E-01 - 20 4.7500000E-01 - 21 5.0000000E-01 - 22 5.2500000E-01 - 23 5.5000000E-01 - 24 5.7500000E-01 - 25 6.0000000E-01 - 26 6.2500000E-01 - 27 6.5000000E-01 - 28 6.7500000E-01 - 29 7.0000000E-01 - 30 7.2500000E-01 - 31 7.5000000E-01 - 32 7.7500000E-01 - 33 8.0000000E-01 - 34 8.2500000E-01 - 35 8.5000000E-01 - 36 8.7500000E-01 - 37 9.0000000E-01 - 38 9.2500000E-01 - 39 9.5000000E-01 - 40 9.7500000E-01 - 41 1.0000000E+00 -# rmin | m - 1 0.0000000E+00 - 2 1.4227510E-02 - 3 2.8466934E-02 - 4 4.2718272E-02 - 5 5.6983341E-02 - 6 7.1235625E-02 - 7 8.5486982E-02 - 8 9.9745301E-02 - 9 1.1401303E-01 - 10 1.2828767E-01 - 11 1.4256987E-01 - 12 1.5685875E-01 - 13 1.7115613E-01 - 14 1.8546452E-01 - 15 1.9978907E-01 - 16 2.1413214E-01 - 17 2.2848774E-01 - 18 2.4284758E-01 - 19 2.5717641E-01 - 20 2.7141022E-01 - 21 2.8548256E-01 - 22 2.9936469E-01 - 23 3.1305558E-01 - 24 3.2655537E-01 - 25 3.3986493E-01 - 26 3.5297568E-01 - 27 3.6588018E-01 - 28 3.7857274E-01 - 29 3.9104437E-01 - 30 4.0328613E-01 - 31 4.1529208E-01 - 32 4.2704488E-01 - 33 4.3854427E-01 - 34 4.4977082E-01 - 35 4.6071034E-01 - 36 4.7134821E-01 - 37 4.8165281E-01 - 38 4.9157405E-01 - 39 5.0099974E-01 - 40 5.0972631E-01 - 41 5.1754440E-01 -# polflux | Wb/radian - 1 0.0000000E+00 - 2 -4.3432751E-04 - 3 -1.2679152E-03 - 4 -2.8260758E-03 - 5 -5.0212465E-03 - 6 -7.8402786E-03 - 7 -1.1286703E-02 - 8 -1.5359906E-02 - 9 -2.0057464E-02 - 10 -2.5379258E-02 - 11 -3.1326011E-02 - 12 -3.7898070E-02 - 13 -4.5096351E-02 - 14 -5.2922003E-02 - 15 -6.1375871E-02 - 16 -7.0456348E-02 - 17 -8.0154148E-02 - 18 -9.0432210E-02 - 19 -1.0116711E-01 - 20 -1.1211269E-01 - 21 -1.2302423E-01 - 22 -1.3379089E-01 - 23 -1.4438682E-01 - 24 -1.5479785E-01 - 25 -1.6500794E-01 - 26 -1.7500069E-01 - 27 -1.8476092E-01 - 28 -1.9427599E-01 - 29 -2.0353478E-01 - 30 -2.1252855E-01 - 31 -2.2125020E-01 - 32 -2.2969405E-01 - 33 -2.3785549E-01 - 34 -2.4573046E-01 - 35 -2.5331441E-01 - 36 -2.6060073E-01 - 37 -2.6757807E-01 - 38 -2.7421986E-01 - 39 -2.8046577E-01 - 40 -2.8620030E-01 - 41 -2.9125962E-01 -# q | - - 1 9.9589697E-01 - 2 9.9481734E-01 - 3 9.9365998E-01 - 4 9.9472206E-01 - 5 9.9748836E-01 - 6 9.9714908E-01 - 7 9.9685286E-01 - 8 9.9711845E-01 - 9 9.9766262E-01 - 10 9.9792996E-01 - 11 9.9807024E-01 - 12 9.9811617E-01 - 13 9.9798908E-01 - 14 9.9775859E-01 - 15 9.9752202E-01 - 16 9.9776944E-01 - 17 9.9968513E-01 - 18 1.0077790E+00 - 19 1.0327685E+00 - 20 1.0830403E+00 - 21 1.1518867E+00 - 22 1.2279928E+00 - 23 1.3081780E+00 - 24 1.3932073E+00 - 25 1.4838685E+00 - 26 1.5809555E+00 - 27 1.6849893E+00 - 28 1.7965734E+00 - 29 1.9163970E+00 - 30 2.0450639E+00 - 31 2.1834084E+00 - 32 2.3323207E+00 - 33 2.4929699E+00 - 34 2.6668344E+00 - 35 2.8561567E+00 - 36 3.0647015E+00 - 37 3.2991758E+00 - 38 3.5793870E+00 - 39 3.9449559E+00 - 40 4.4949697E+00 - 41 5.3049685E+00 -# w0 | rad/s - 1 3.7830865E+02 - 2 -9.2485499E+03 - 3 -1.7523088E+04 - 4 -1.5272594E+04 - 5 -1.4187521E+04 - 6 -1.3166442E+04 - 7 -1.2258841E+04 - 8 -1.1556091E+04 - 9 -1.1091054E+04 - 10 -1.0836130E+04 - 11 -1.0721431E+04 - 12 -1.0669172E+04 - 13 -1.0614913E+04 - 14 -1.0522472E+04 - 15 -1.0388046E+04 - 16 -1.0230429E+04 - 17 -1.0077721E+04 - 18 -9.9541613E+03 - 19 -9.8511888E+03 - 20 -9.7436746E+03 - 21 -9.6509289E+03 - 22 -9.5955715E+03 - 23 -9.5478774E+03 - 24 -9.4724843E+03 - 25 -9.3538176E+03 - 26 -9.1826275E+03 - 27 -8.9593418E+03 - 28 -8.6942391E+03 - 29 -8.4033600E+03 - 30 -8.1007805E+03 - 31 -7.7827320E+03 - 32 -7.4123408E+03 - 33 -6.9043116E+03 - 34 -6.1647569E+03 - 35 -5.1821014E+03 - 36 -4.0701056E+03 - 37 -3.0589547E+03 - 38 -2.3900318E+03 - 39 -2.1738929E+03 - 40 -2.3759322E+03 - 41 -2.7268136E+03 -# rmaj | m - 1 1.6484918E+00 - 2 1.6484700E+00 - 3 1.6484046E+00 - 4 1.6482956E+00 - 5 1.6481532E+00 - 6 1.6479721E+00 - 7 1.6477622E+00 - 8 1.6475201E+00 - 9 1.6472411E+00 - 10 1.6469243E+00 - 11 1.6465718E+00 - 12 1.6461826E+00 - 13 1.6457619E+00 - 14 1.6453072E+00 - 15 1.6448203E+00 - 16 1.6442986E+00 - 17 1.6437408E+00 - 18 1.6431463E+00 - 19 1.6425094E+00 - 20 1.6418103E+00 - 21 1.6410369E+00 - 22 1.6401749E+00 - 23 1.6392320E+00 - 24 1.6382086E+00 - 25 1.6371069E+00 - 26 1.6359279E+00 - 27 1.6346751E+00 - 28 1.6333557E+00 - 29 1.6319658E+00 - 30 1.6305141E+00 - 31 1.6289993E+00 - 32 1.6274306E+00 - 33 1.6258139E+00 - 34 1.6241542E+00 - 35 1.6224592E+00 - 36 1.6207333E+00 - 37 1.6189829E+00 - 38 1.6172315E+00 - 39 1.6154874E+00 - 40 1.6138009E+00 - 41 1.6122503E+00 -# zmag | m - 1 2.9382376E-02 - 2 2.9349916E-02 - 3 2.9317448E-02 - 4 2.9284970E-02 - 5 2.9237432E-02 - 6 2.9188873E-02 - 7 2.9125817E-02 - 8 2.9049672E-02 - 9 2.8958796E-02 - 10 2.8857254E-02 - 11 2.8743918E-02 - 12 2.8625477E-02 - 13 2.8494013E-02 - 14 2.8347401E-02 - 15 2.8187795E-02 - 16 2.8015947E-02 - 17 2.7840354E-02 - 18 2.7652074E-02 - 19 2.7433578E-02 - 20 2.7191291E-02 - 21 2.6895180E-02 - 22 2.6544653E-02 - 23 2.6122137E-02 - 24 2.5624563E-02 - 25 2.5033911E-02 - 26 2.4339530E-02 - 27 2.3522709E-02 - 28 2.2558018E-02 - 29 2.1424482E-02 - 30 2.0072877E-02 - 31 1.8478535E-02 - 32 1.6574497E-02 - 33 1.4296897E-02 - 34 1.1553924E-02 - 35 8.2313557E-03 - 36 4.1473295E-03 - 37 -9.4334423E-04 - 38 -7.5108330E-03 - 39 -1.6575151E-02 - 40 -3.0846409E-02 - 41 -5.5729841E-02 -# kappa | - - 1 1.2612226E+00 - 2 1.2603742E+00 - 3 1.2595256E+00 - 4 1.2586768E+00 - 5 1.2578944E+00 - 6 1.2579577E+00 - 7 1.2581564E+00 - 8 1.2582556E+00 - 9 1.2582614E+00 - 10 1.2582558E+00 - 11 1.2582070E+00 - 12 1.2581921E+00 - 13 1.2581933E+00 - 14 1.2581697E+00 - 15 1.2580387E+00 - 16 1.2578513E+00 - 17 1.2577270E+00 - 18 1.2577609E+00 - 19 1.2581936E+00 - 20 1.2594605E+00 - 21 1.2619693E+00 - 22 1.2658061E+00 - 23 1.2708189E+00 - 24 1.2768283E+00 - 25 1.2837512E+00 - 26 1.2915678E+00 - 27 1.3002511E+00 - 28 1.3098026E+00 - 29 1.3202648E+00 - 30 1.3316935E+00 - 31 1.3441186E+00 - 32 1.3576626E+00 - 33 1.3723851E+00 - 34 1.3884792E+00 - 35 1.4061057E+00 - 36 1.4255460E+00 - 37 1.4472327E+00 - 38 1.4720548E+00 - 39 1.5019639E+00 - 40 1.5418564E+00 - 41 1.6012201E+00 -# delta | - - 1 0.0000000E+00 - 2 2.1626237E-03 - 3 3.1039289E-03 - 4 2.8239216E-03 - 5 2.9387609E-03 - 6 3.3220484E-03 - 7 4.1765986E-03 - 8 5.1738285E-03 - 9 6.0758512E-03 - 10 6.8338356E-03 - 11 7.6027395E-03 - 12 8.2974865E-03 - 13 9.0584667E-03 - 14 9.7798816E-03 - 15 1.0524117E-02 - 16 1.1238660E-02 - 17 1.1945644E-02 - 18 1.2663210E-02 - 19 1.3421702E-02 - 20 1.4339771E-02 - 21 1.5584558E-02 - 22 1.7127728E-02 - 23 1.8929971E-02 - 24 2.0924548E-02 - 25 2.3197903E-02 - 26 2.5698506E-02 - 27 2.8476364E-02 - 28 3.1496444E-02 - 29 3.4844808E-02 - 30 3.8475434E-02 - 31 4.2409845E-02 - 32 4.6737469E-02 - 33 5.1387655E-02 - 34 5.6547826E-02 - 35 6.2108405E-02 - 36 6.8257691E-02 - 37 7.5032445E-02 - 38 8.2851371E-02 - 39 9.2288564E-02 - 40 1.0503056E-01 - 41 1.2352707E-01 -# zeta | - - 1 0.0000000E+00 - 2 7.6569916E-05 - 3 4.6449675E-05 - 4 -9.0360723E-05 - 5 5.4174777E-06 - 6 1.8536230E-05 - 7 -3.5086564E-05 - 8 -1.1570194E-04 - 9 -2.2138163E-04 - 10 -3.2814832E-04 - 11 -3.7379754E-04 - 12 -4.3425220E-04 - 13 -5.0977657E-04 - 14 -5.7290889E-04 - 15 -5.6807933E-04 - 16 -5.8939623E-04 - 17 -7.0124722E-04 - 18 -9.6730590E-04 - 19 -1.3502980E-03 - 20 -1.8396865E-03 - 21 -2.4607337E-03 - 22 -3.2861485E-03 - 23 -4.3271366E-03 - 24 -5.5161425E-03 - 25 -6.8910886E-03 - 26 -8.4632225E-03 - 27 -1.0219299E-02 - 28 -1.2188321E-02 - 29 -1.4414920E-02 - 30 -1.6955009E-02 - 31 -1.9834704E-02 - 32 -2.3094533E-02 - 33 -2.6828103E-02 - 34 -3.1125957E-02 - 35 -3.6049755E-02 - 36 -4.1788933E-02 - 37 -4.8545104E-02 - 38 -5.6830254E-02 - 39 -6.7735871E-02 - 40 -8.4186411E-02 - 41 -1.1142245E-01 -# shape_cos0 | - - 1 -5.8662572E-02 - 2 -5.8535242E-02 - 3 -5.8152796E-02 - 4 -5.7515236E-02 - 5 -5.7218122E-02 - 6 -5.7296212E-02 - 7 -5.7392179E-02 - 8 -5.7448972E-02 - 9 -5.7459553E-02 - 10 -5.7456732E-02 - 11 -5.7488482E-02 - 12 -5.7486927E-02 - 13 -5.7473575E-02 - 14 -5.7455085E-02 - 15 -5.7459913E-02 - 16 -5.7474107E-02 - 17 -5.7503650E-02 - 18 -5.7545897E-02 - 19 -5.7646890E-02 - 20 -5.7913652E-02 - 21 -5.8430580E-02 - 22 -5.9201593E-02 - 23 -6.0213227E-02 - 24 -6.1419493E-02 - 25 -6.2815503E-02 - 26 -6.4360743E-02 - 27 -6.6072019E-02 - 28 -6.7939722E-02 - 29 -6.9969876E-02 - 30 -7.2149653E-02 - 31 -7.4530316E-02 - 32 -7.7060199E-02 - 33 -7.9800591E-02 - 34 -8.2768420E-02 - 35 -8.5998217E-02 - 36 -8.9515669E-02 - 37 -9.3438627E-02 - 38 -9.7921957E-02 - 39 -1.0337191E-01 - 40 -1.1084729E-01 - 41 -1.2209109E-01 -# shape_cos1 | - - 1 0.0000000E+00 - 2 -4.4461858E-04 - 3 -1.0439010E-03 - 4 -1.7978473E-03 - 5 -2.6835803E-03 - 6 -3.2842355E-03 - 7 -3.9782791E-03 - 8 -4.6949673E-03 - 9 -5.4175395E-03 - 10 -6.0957064E-03 - 11 -6.7919373E-03 - 12 -7.4900463E-03 - 13 -8.2024793E-03 - 14 -8.9551743E-03 - 15 -9.7004222E-03 - 16 -1.0439161E-02 - 17 -1.1104747E-02 - 18 -1.1787320E-02 - 19 -1.2593190E-02 - 20 -1.3407079E-02 - 21 -1.4381027E-02 - 22 -1.5474662E-02 - 23 -1.6731514E-02 - 24 -1.8117909E-02 - 25 -1.9710195E-02 - 26 -2.1479009E-02 - 27 -2.3437316E-02 - 28 -2.5669795E-02 - 29 -2.8156234E-02 - 30 -3.1013693E-02 - 31 -3.4238694E-02 - 32 -3.7958174E-02 - 33 -4.2241522E-02 - 34 -4.7232231E-02 - 35 -5.3129964E-02 - 36 -6.0200454E-02 - 37 -6.8845062E-02 - 38 -7.9866821E-02 - 39 -9.5045048E-02 - 40 -1.1945163E-01 - 41 -1.6228171E-01 -# shape_cos2 | - - 1 0.0000000E+00 - 2 7.2192197E-05 - 3 7.0735956E-05 - 4 -4.3687241E-06 - 5 -1.0603327E-05 - 6 -4.3275232E-05 - 7 -1.1331768E-04 - 8 -1.5787045E-04 - 9 -1.9175522E-04 - 10 -1.6995380E-04 - 11 -7.0388918E-05 - 12 -1.9137174E-05 - 13 -1.1431546E-05 - 14 -4.2916010E-05 - 15 -8.6273841E-05 - 16 -1.1050112E-04 - 17 -1.2784355E-04 - 18 -7.9642203E-05 - 19 2.7556730E-05 - 20 1.7888252E-04 - 21 4.1168349E-04 - 22 7.2365237E-04 - 23 1.1271256E-03 - 24 1.6286857E-03 - 25 2.1418798E-03 - 26 2.7348310E-03 - 27 3.4120907E-03 - 28 4.1744165E-03 - 29 4.9820726E-03 - 30 5.9271521E-03 - 31 7.0343530E-03 - 32 8.2031070E-03 - 33 9.5264887E-03 - 34 1.1066750E-02 - 35 1.2819949E-02 - 36 1.4737644E-02 - 37 1.7056796E-02 - 38 1.9822086E-02 - 39 2.3467575E-02 - 40 2.8969429E-02 - 41 3.7521910E-02 -# shape_cos3 | - - 1 0.0000000E+00 - 2 -4.7478642E-04 - 3 -4.9121753E-04 - 4 -4.9293312E-05 - 5 6.9868404E-05 - 6 -1.5822842E-05 - 7 -4.5896679E-05 - 8 -5.1235964E-05 - 9 7.3496443E-06 - 10 8.4323699E-05 - 11 1.3785293E-04 - 12 5.8852964E-05 - 13 7.8393094E-06 - 14 -8.6178444E-06 - 15 5.2931422E-06 - 16 3.6944331E-05 - 17 2.3403505E-05 - 18 8.5620459E-06 - 19 3.5628952E-05 - 20 6.6051501E-05 - 21 1.3715763E-04 - 22 2.0531172E-04 - 23 3.1131649E-04 - 24 4.4009304E-04 - 25 5.8884437E-04 - 26 7.6711840E-04 - 27 9.9775281E-04 - 28 1.2587352E-03 - 29 1.5653083E-03 - 30 1.9722043E-03 - 31 2.4324423E-03 - 32 3.0082622E-03 - 33 3.7138507E-03 - 34 4.5857885E-03 - 35 5.6128249E-03 - 36 6.9186619E-03 - 37 8.5270717E-03 - 38 1.0611317E-02 - 39 1.3681252E-02 - 40 1.9260023E-02 - 41 2.9017335E-02 -# shape_cos4 | - - 1 0.0000000E+00 - 2 7.5724635E-05 - 3 8.5411565E-05 - 4 2.9060789E-05 - 5 3.5895487E-05 - 6 -1.8583796E-05 - 7 3.4613923E-05 - 8 7.2591938E-06 - 9 -4.1839144E-06 - 10 1.1194344E-05 - 11 -4.7610705E-05 - 12 -3.9936940E-05 - 13 -9.3037918E-06 - 14 2.5337836E-05 - 15 1.0421580E-07 - 16 -5.6850852E-05 - 17 -1.3691172E-04 - 18 -1.9777334E-04 - 19 -2.3197307E-04 - 20 -2.6133256E-04 - 21 -2.8913506E-04 - 22 -3.1163676E-04 - 23 -3.7554703E-04 - 24 -4.5578823E-04 - 25 -5.8549872E-04 - 26 -6.9946134E-04 - 27 -8.4293226E-04 - 28 -1.0092315E-03 - 29 -1.2043021E-03 - 30 -1.3954823E-03 - 31 -1.6751700E-03 - 32 -1.9080125E-03 - 33 -2.1953776E-03 - 34 -2.5215730E-03 - 35 -2.8875676E-03 - 36 -3.2418906E-03 - 37 -3.6558225E-03 - 38 -4.0924187E-03 - 39 -4.5898400E-03 - 40 -5.4053571E-03 - 41 -5.9925165E-03 -# shape_cos5 | - - 1 0.0000000E+00 - 2 -9.4718188E-06 - 3 -2.8146665E-05 - 4 -5.6024538E-05 - 5 -5.6537558E-05 - 6 1.3122965E-05 - 7 9.0082583E-06 - 8 2.2700256E-06 - 9 1.8172401E-06 - 10 1.8091308E-05 - 11 1.1142437E-05 - 12 1.4416728E-05 - 13 5.2923660E-06 - 14 -4.3541310E-05 - 15 -8.5220517E-05 - 16 -1.1596211E-04 - 17 -6.1875568E-05 - 18 -2.3261463E-05 - 19 -6.0052917E-05 - 20 -4.3640286E-05 - 21 -5.1549240E-05 - 22 -3.5028867E-05 - 23 -3.6163444E-05 - 24 -1.5993548E-05 - 25 -3.2908926E-05 - 26 -4.2214539E-05 - 27 -3.5891594E-05 - 28 -6.5653821E-05 - 29 -6.8239325E-05 - 30 -1.1468567E-04 - 31 -1.2996660E-04 - 32 -1.5916760E-04 - 33 -1.7693671E-04 - 34 -1.7600763E-04 - 35 -1.6731070E-04 - 36 -1.4103138E-04 - 37 -7.9456973E-05 - 38 -7.7212748E-07 - 39 1.0407734E-04 - 40 -9.9575790E-05 - 41 2.4230842E-06 -# shape_sin3 | - - 1 0.0000000E+00 - 2 1.5497501E-04 - 3 1.4466918E-04 - 4 -3.0917479E-05 - 5 2.5706487E-05 - 6 -2.0484797E-05 - 7 -1.6707030E-05 - 8 -2.3101044E-05 - 9 -7.1543558E-06 - 10 1.9399798E-05 - 11 1.5679049E-05 - 12 2.6255181E-06 - 13 -3.9962087E-05 - 14 -4.9897783E-05 - 15 -7.5144997E-05 - 16 -1.1107971E-04 - 17 -1.7225144E-04 - 18 -2.2953545E-04 - 19 -2.3972675E-04 - 20 -2.7833211E-04 - 21 -3.6702691E-04 - 22 -5.7035120E-04 - 23 -7.7020329E-04 - 24 -9.4802967E-04 - 25 -1.2156259E-03 - 26 -1.5121068E-03 - 27 -1.8492729E-03 - 28 -2.1646693E-03 - 29 -2.5964527E-03 - 30 -3.0255678E-03 - 31 -3.5189776E-03 - 32 -4.0936669E-03 - 33 -4.6531585E-03 - 34 -5.3610623E-03 - 35 -6.0647479E-03 - 36 -6.9024497E-03 - 37 -7.8137350E-03 - 38 -8.9456143E-03 - 39 -1.0428333E-02 - 40 -1.2894876E-02 - 41 -1.6209363E-02 -# shape_sin4 | - - 1 0.0000000E+00 - 2 4.3715166E-05 - 3 4.2055509E-05 - 4 -4.9789711E-06 - 5 1.6665040E-05 - 6 2.2864134E-07 - 7 6.8712165E-06 - 8 -1.1076064E-05 - 9 1.7543751E-05 - 10 5.3965458E-05 - 11 1.1959236E-04 - 12 8.8072297E-05 - 13 -1.6889151E-05 - 14 -1.2884948E-04 - 15 -1.1948547E-04 - 16 -4.2024970E-05 - 17 -4.5361861E-05 - 18 -1.2934318E-04 - 19 -1.8281903E-04 - 20 -2.0525205E-04 - 21 -2.3692553E-04 - 22 -2.7650921E-04 - 23 -3.6658018E-04 - 24 -4.7181258E-04 - 25 -5.9445226E-04 - 26 -7.4663843E-04 - 27 -9.1499227E-04 - 28 -1.0941730E-03 - 29 -1.3180811E-03 - 30 -1.5933342E-03 - 31 -1.8887696E-03 - 32 -2.2654706E-03 - 33 -2.6660894E-03 - 34 -3.1580414E-03 - 35 -3.7225598E-03 - 36 -4.3711340E-03 - 37 -5.1009725E-03 - 38 -6.0099862E-03 - 39 -7.2654897E-03 - 40 -9.5401085E-03 - 41 -1.3040626E-02 -# shape_sin5 | - - 1 0.0000000E+00 - 2 1.1874480E-04 - 3 1.0552789E-04 - 4 -3.9650720E-05 - 5 2.6051125E-05 - 6 -2.0642626E-06 - 7 1.0407386E-05 - 8 -2.3711840E-06 - 9 4.6353497E-06 - 10 6.8462924E-07 - 11 3.0937255E-05 - 12 -2.4715696E-05 - 13 -5.1202293E-06 - 14 -1.7059166E-05 - 15 2.4552106E-06 - 16 5.7589577E-07 - 17 -4.1281977E-06 - 18 -1.3802499E-05 - 19 -5.4375305E-05 - 20 -8.7940563E-05 - 21 -5.6754140E-05 - 22 -1.8097531E-05 - 23 1.0520767E-05 - 24 -1.2390009E-05 - 25 -8.1574545E-06 - 26 -2.5886422E-05 - 27 -2.4119318E-05 - 28 -6.1954614E-05 - 29 -6.4755355E-05 - 30 -9.9013735E-05 - 31 -1.7242680E-04 - 32 -2.2414220E-04 - 33 -3.5749457E-04 - 34 -4.1976306E-04 - 35 -5.8923559E-04 - 36 -7.6364100E-04 - 37 -1.0502220E-03 - 38 -1.3569811E-03 - 39 -1.7230707E-03 - 40 -2.0217193E-03 - 41 -2.7780079E-03 -# ne | 10^19/m^3 - 1 2.4293253E+00 - 2 2.4286951E+00 - 3 2.4276761E+00 - 4 2.4257952E+00 - 5 2.4228931E+00 - 6 2.4188349E+00 - 7 2.4135116E+00 - 8 2.4068280E+00 - 9 2.3986810E+00 - 10 2.3889806E+00 - 11 2.3776592E+00 - 12 2.3646434E+00 - 13 2.3498588E+00 - 14 2.3332491E+00 - 15 2.3148026E+00 - 16 2.2945379E+00 - 17 2.2723876E+00 - 18 2.2481282E+00 - 19 2.2215476E+00 - 20 2.1925802E+00 - 21 2.1612886E+00 - 22 2.1278651E+00 - 23 2.0926416E+00 - 24 2.0560173E+00 - 25 2.0183876E+00 - 26 1.9801511E+00 - 27 1.9416858E+00 - 28 1.9033119E+00 - 29 1.8652815E+00 - 30 1.8277941E+00 - 31 1.7910084E+00 - 32 1.7550357E+00 - 33 1.7199429E+00 - 34 1.6857772E+00 - 35 1.6522772E+00 - 36 1.6185600E+00 - 37 1.5836065E+00 - 38 1.5468093E+00 - 39 1.5036345E+00 - 40 1.3755444E+00 - 41 1.1679028E+00 -# ni | 10^19/m^3 - 1 2.0278446E+00 6.5420474E-02 5.6074748E-08 8.7732740E-03 - 2 2.0276837E+00 6.5403494E-02 5.6060194E-08 8.4066287E-03 - 3 2.0271004E+00 6.5376054E-02 5.6036674E-08 8.2315887E-03 - 4 2.0255932E+00 6.5325401E-02 5.5993257E-08 8.2662531E-03 - 5 2.0231756E+00 6.5247246E-02 5.5926267E-08 8.2002150E-03 - 6 2.0198691E+00 6.5137969E-02 5.5832600E-08 8.0746971E-03 - 7 2.0154750E+00 6.4994609E-02 5.5709720E-08 8.0621822E-03 - 8 2.0098094E+00 6.4814627E-02 5.5555450E-08 8.1987158E-03 - 9 2.0028107E+00 6.4595229E-02 5.5367395E-08 8.3981715E-03 - 10 1.9944827E+00 6.4333996E-02 5.5143480E-08 8.5887448E-03 - 11 1.9847597E+00 6.4029133E-02 5.4882169E-08 8.8606636E-03 - 12 1.9734829E+00 6.3678610E-02 5.4581721E-08 9.3158210E-03 - 13 1.9605517E+00 6.3280456E-02 5.4240446E-08 9.9311182E-03 - 14 1.9459545E+00 6.2833185E-02 5.3857069E-08 1.0659186E-02 - 15 1.9296997E+00 6.2336438E-02 5.3431286E-08 1.1509500E-02 - 16 1.9118043E+00 6.1790703E-02 5.2963513E-08 1.2468119E-02 - 17 1.8924705E+00 6.1194204E-02 5.2452228E-08 1.3034327E-02 - 18 1.8723427E+00 6.0540925E-02 5.1892273E-08 1.2045315E-02 - 19 1.8518956E+00 5.9825117E-02 5.1278723E-08 9.3563098E-03 - 20 1.8303081E+00 5.9045037E-02 5.0610082E-08 6.6460702E-03 - 21 1.8061967E+00 5.8202374E-02 4.9887799E-08 5.1082187E-03 - 22 1.7792879E+00 5.7302296E-02 4.9116302E-08 4.4176786E-03 - 23 1.7502952E+00 5.6353742E-02 4.8303255E-08 4.0291828E-03 - 24 1.7199287E+00 5.5367474E-02 4.7457882E-08 3.7381111E-03 - 25 1.6886543E+00 5.4354121E-02 4.6589293E-08 3.4782367E-03 - 26 1.6568529E+00 5.3324432E-02 4.5706701E-08 3.2238446E-03 - 27 1.6248523E+00 5.2288583E-02 4.4818831E-08 2.9794729E-03 - 28 1.5929106E+00 5.1255194E-02 4.3933067E-08 2.7603782E-03 - 29 1.5612327E+00 5.0231045E-02 4.3055225E-08 2.5634039E-03 - 30 1.5300013E+00 4.9221532E-02 4.2189927E-08 2.3624729E-03 - 31 1.4993742E+00 4.8230922E-02 4.1340831E-08 2.1341914E-03 - 32 1.4694559E+00 4.7262198E-02 4.0510496E-08 1.8784234E-03 - 33 1.4402936E+00 4.6317160E-02 3.9700462E-08 1.6129714E-03 - 34 1.4119103E+00 4.5397095E-02 3.8911834E-08 1.3542979E-03 - 35 1.3840749E+00 4.4494958E-02 3.8138574E-08 1.1100722E-03 - 36 1.3560432E+00 4.3586968E-02 3.7360296E-08 8.7914272E-04 - 37 1.3269623E+00 4.2645689E-02 3.6553484E-08 6.5980662E-04 - 38 1.2963197E+00 4.1654756E-02 3.5704113E-08 4.6092023E-04 - 39 1.2602991E+00 4.0492084E-02 3.4707536E-08 3.0363867E-04 - 40 1.1530288E+00 3.7042699E-02 3.1750917E-08 2.1476978E-04 - 41 9.7902665E-01 3.1451034E-02 2.6958056E-08 2.1476978E-04 -# te | keV - 1 3.3836121E+00 - 2 3.3667085E+00 - 3 3.3413738E+00 - 4 3.2995637E+00 - 5 3.2423742E+00 - 6 3.1714978E+00 - 7 3.0890267E+00 - 8 2.9971364E+00 - 9 2.8978224E+00 - 10 2.7929280E+00 - 11 2.6841704E+00 - 12 2.5730409E+00 - 13 2.4608589E+00 - 14 2.3488297E+00 - 15 2.2380601E+00 - 16 2.1295268E+00 - 17 2.0233270E+00 - 18 1.9181176E+00 - 19 1.8123573E+00 - 20 1.7053801E+00 - 21 1.5973095E+00 - 22 1.4889246E+00 - 23 1.3814095E+00 - 24 1.2760906E+00 - 25 1.1743032E+00 - 26 1.0771959E+00 - 27 9.8554524E-01 - 28 8.9985040E-01 - 29 8.2038232E-01 - 30 7.4718859E-01 - 31 6.8013323E-01 - 32 6.1891065E-01 - 33 5.6314616E-01 - 34 5.1242447E-01 - 35 4.6339510E-01 - 36 4.0792724E-01 - 37 3.4053767E-01 - 38 2.6438829E-01 - 39 1.8861037E-01 - 40 1.2534898E-01 - 41 7.1676107E-02 -# ti | keV - 1 1.0530279E+00 1.0607021E+00 1.0607021E+00 1.1101037E+01 - 2 1.0471973E+00 1.0548581E+00 1.0548581E+00 1.1060268E+01 - 3 1.0394862E+00 1.0470875E+00 1.0470875E+00 1.1907518E+01 - 4 1.0288854E+00 1.0363941E+00 1.0363941E+00 1.3444714E+01 - 5 1.0167961E+00 1.0241806E+00 1.0241806E+00 1.4464764E+01 - 6 1.0038177E+00 1.0110295E+00 1.0110295E+00 1.4810818E+01 - 7 9.8986159E-01 9.9688353E-01 9.9688353E-01 1.5153939E+01 - 8 9.7457164E-01 9.8137593E-01 9.8137593E-01 1.5289471E+01 - 9 9.5764313E-01 9.6421567E-01 9.6421567E-01 1.4719245E+01 - 10 9.3906547E-01 9.4542323E-01 9.4542323E-01 1.3808869E+01 - 11 9.1923425E-01 9.2540152E-01 9.2540152E-01 1.3080950E+01 - 12 8.9872573E-01 9.0471349E-01 9.0471349E-01 1.2859517E+01 - 13 8.7795024E-01 8.8376482E-01 8.8376482E-01 1.3261659E+01 - 14 8.5690003E-01 8.6254597E-01 8.6254597E-01 1.3699394E+01 - 15 8.3502370E-01 8.4049864E-01 8.4049864E-01 1.3912696E+01 - 16 8.1139099E-01 8.1668569E-01 8.1668569E-01 1.4461817E+01 - 17 7.8510735E-01 7.9020049E-01 7.9020049E-01 1.5501012E+01 - 18 7.5578493E-01 7.6063024E-01 7.6063024E-01 1.6201223E+01 - 19 7.2384515E-01 7.2837158E-01 7.2837158E-01 1.5661225E+01 - 20 6.9042566E-01 6.9458945E-01 6.9458945E-01 1.3969171E+01 - 21 6.5702436E-01 6.6085572E-01 6.6085572E-01 1.2130061E+01 - 22 6.2505975E-01 6.2862097E-01 6.2862097E-01 1.0939708E+01 - 23 5.9538211E-01 5.9871280E-01 5.9871280E-01 1.0298946E+01 - 24 5.6808217E-01 5.7120076E-01 5.7120076E-01 9.9763451E+00 - 25 5.4261165E-01 5.4553127E-01 5.4553127E-01 9.8521340E+00 - 26 5.1807063E-01 5.2079525E-01 5.2079525E-01 9.6835060E+00 - 27 4.9365491E-01 4.9618140E-01 4.9618140E-01 1.4684507E+01 - 28 4.6901488E-01 4.7134447E-01 4.7134447E-01 8.1738697E+00 - 29 4.4439533E-01 4.4653414E-01 4.4653414E-01 7.9667000E+00 - 30 4.2033147E-01 4.2228503E-01 4.2228503E-01 6.0111928E+00 - 31 3.9675481E-01 3.9852741E-01 3.9852741E-01 5.9075046E+00 - 32 3.7207266E-01 3.7367044E-01 3.7367044E-01 5.7961970E+00 - 33 3.4363739E-01 3.4507277E-01 3.4507277E-01 5.6710201E+00 - 34 3.0937458E-01 3.1066751E-01 3.1066751E-01 5.5306084E+00 - 35 2.6974294E-01 2.7090514E-01 2.7090514E-01 5.3846747E+00 - 36 2.2781668E-01 2.2883461E-01 2.2883461E-01 5.2500185E+00 - 37 1.8696457E-01 1.8781259E-01 1.8781259E-01 5.1416131E+00 - 38 1.4926798E-01 1.4992861E-01 1.4992861E-01 5.0797503E+00 - 39 1.1545574E-01 1.1592265E-01 1.1592265E-01 5.0706745E+00 - 40 8.5627841E-02 8.5922439E-02 8.5922439E-02 5.0892086E+00 - 41 5.7799389E-02 5.7799389E-02 5.7799389E-02 5.0892086E+00 -# ptot | Pa - 1 1.6984435E+04 - 2 1.6921380E+04 - 3 1.6802913E+04 - 4 1.6590480E+04 - 5 1.6309002E+04 - 6 1.5960051E+04 - 7 1.5557551E+04 - 8 1.5115887E+04 - 9 1.4639103E+04 - 10 1.4128929E+04 - 11 1.3592580E+04 - 12 1.3043132E+04 - 13 1.2493270E+04 - 14 1.1947614E+04 - 15 1.1406194E+04 - 16 1.0872117E+04 - 17 1.0337492E+04 - 18 9.7650871E+03 - 19 9.1223969E+03 - 20 8.4354174E+03 - 21 7.7644973E+03 - 22 7.1392322E+03 - 23 6.5549249E+03 - 24 6.0039922E+03 - 25 5.4843084E+03 - 26 4.9956281E+03 - 27 4.5382897E+03 - 28 4.1132578E+03 - 29 3.7215775E+03 - 30 3.3633910E+03 - 31 3.0365410E+03 - 32 2.7348885E+03 - 33 2.4500398E+03 - 34 2.1756497E+03 - 35 1.9042182E+03 - 36 1.6212433E+03 - 37 1.3194779E+03 - 38 1.0118879E+03 - 39 7.2216199E+02 - 40 4.6183686E+02 - 41 2.1372818E+02 -# johm | MA/m^2 - 1 7.1905884E+00 - 2 7.1905884E+00 - 3 3.9773601E+00 - 4 3.3056118E+00 - 5 2.9787690E+00 - 6 2.7980631E+00 - 7 2.7651828E+00 - 8 2.7292832E+00 - 9 2.6896830E+00 - 10 2.6580158E+00 - 11 2.6331764E+00 - 12 2.6156047E+00 - 13 2.6038671E+00 - 14 2.5922532E+00 - 15 2.5751515E+00 - 16 2.5319227E+00 - 17 2.3907021E+00 - 18 1.9749940E+00 - 19 1.3166589E+00 - 20 9.5633239E-01 - 21 8.4210598E-01 - 22 7.6520527E-01 - 23 6.8474791E-01 - 24 6.0713517E-01 - 25 5.3542373E-01 - 26 4.7080751E-01 - 27 4.1350953E-01 - 28 3.6341951E-01 - 29 3.2008348E-01 - 30 2.8293101E-01 - 31 2.5122276E-01 - 32 2.2422882E-01 - 33 2.0128615E-01 - 34 1.8187052E-01 - 35 1.6334421E-01 - 36 1.4193072E-01 - 37 1.1736287E-01 - 38 9.0651664E-02 - 39 6.2838758E-02 - 40 3.6315194E-02 - 41 0.0000000E+00 -# jbs | MA/m^2 - 1 9.7313252E-02 - 2 9.7313252E-02 - 3 6.6863572E-02 - 4 6.3141855E-02 - 5 6.1528603E-02 - 6 6.0929653E-02 - 7 6.0271980E-02 - 8 5.9418384E-02 - 9 5.8250989E-02 - 10 5.6685682E-02 - 11 5.4816889E-02 - 12 5.2723454E-02 - 13 5.0492099E-02 - 14 4.8212133E-02 - 15 4.5963830E-02 - 16 4.3996188E-02 - 17 4.2699463E-02 - 18 4.2297172E-02 - 19 4.2865982E-02 - 20 4.4256392E-02 - 21 4.5573942E-02 - 22 4.6099157E-02 - 23 4.5954726E-02 - 24 4.5328408E-02 - 25 4.4340748E-02 - 26 4.3104757E-02 - 27 4.1650376E-02 - 28 3.9971507E-02 - 29 3.8120231E-02 - 30 3.6292764E-02 - 31 3.4830994E-02 - 32 3.3972872E-02 - 33 3.3666983E-02 - 34 3.4632662E-02 - 35 3.8058929E-02 - 36 4.1950134E-02 - 37 4.1468661E-02 - 38 3.2681635E-02 - 39 1.6909313E-02 - 40 6.6196369E-03 - 41 0.0000000E+00 -# jbstor | MA/m^2 - 1 2.4485066E+00 - 2 2.4485066E+00 - 3 2.4479041E+00 - 4 2.4455887E+00 - 5 2.4316776E+00 - 6 2.4443725E+00 - 7 2.4458027E+00 - 8 2.4423301E+00 - 9 2.4402693E+00 - 10 2.4431179E+00 - 11 2.4452004E+00 - 12 2.4478093E+00 - 13 2.4519972E+00 - 14 2.4560847E+00 - 15 2.4584722E+00 - 16 2.4521170E+00 - 17 2.4185777E+00 - 18 2.2791196E+00 - 19 1.8884881E+00 - 20 1.3398258E+00 - 21 9.8358460E-01 - 22 8.5570169E-01 - 23 7.8085321E-01 - 24 7.0621938E-01 - 25 6.3359853E-01 - 26 5.6417345E-01 - 27 5.0182545E-01 - 28 4.4522103E-01 - 29 3.9533718E-01 - 30 3.5194705E-01 - 31 3.1431180E-01 - 32 2.8218757E-01 - 33 2.5516548E-01 - 34 2.3264400E-01 - 35 2.1360791E-01 - 36 1.9694457E-01 - 37 1.8055293E-01 - 38 1.5360877E-01 - 39 1.1693461E-01 - 40 8.1458002E-02 - 41 -1.8007793E+01 -# z_eff | - - 1 1.8079467E+00 - 2 1.8079463E+00 - 3 1.8079144E+00 - 4 1.8078797E+00 - 5 1.8078965E+00 - 6 1.8079064E+00 - 7 1.8078875E+00 - 8 1.8078625E+00 - 9 1.8078518E+00 - 10 1.8078531E+00 - 11 1.8078391E+00 - 12 1.8078076E+00 - 13 1.8077796E+00 - 14 1.8077593E+00 - 15 1.8077370E+00 - 16 1.8077166E+00 - 17 1.8077847E+00 - 18 1.8080630E+00 - 19 1.8083746E+00 - 20 1.8083850E+00 - 21 1.8081729E+00 - 22 1.8080165E+00 - 23 1.8079603E+00 - 24 1.8079424E+00 - 25 1.8079373E+00 - 26 1.8079372E+00 - 27 1.8079361E+00 - 28 1.8079318E+00 - 29 1.8079279E+00 - 30 1.8079298E+00 - 31 1.8079368E+00 - 32 1.8079442E+00 - 33 1.8079477E+00 - 34 1.8079474E+00 - 35 1.8079450E+00 - 36 1.8079428E+00 - 37 1.8079413E+00 - 38 1.8079373E+00 - 39 1.8079276E+00 - 40 1.8079115E+00 - 41 1.8078547E+00 -# vtor | m/s - 1 0.0000000E+00 -2.7195909E+04 0.0000000E+00 0.0000000E+00 - 2 0.0000000E+00 -2.6806808E+04 0.0000000E+00 0.0000000E+00 - 3 0.0000000E+00 -2.5863262E+04 0.0000000E+00 0.0000000E+00 - 4 0.0000000E+00 -2.4512188E+04 0.0000000E+00 0.0000000E+00 - 5 0.0000000E+00 -2.2979925E+04 0.0000000E+00 0.0000000E+00 - 6 0.0000000E+00 -2.1515148E+04 0.0000000E+00 0.0000000E+00 - 7 0.0000000E+00 -2.0314724E+04 0.0000000E+00 0.0000000E+00 - 8 0.0000000E+00 -1.9489222E+04 0.0000000E+00 0.0000000E+00 - 9 0.0000000E+00 -1.9047102E+04 0.0000000E+00 0.0000000E+00 - 10 0.0000000E+00 -1.8903804E+04 0.0000000E+00 0.0000000E+00 - 11 0.0000000E+00 -1.8930851E+04 0.0000000E+00 0.0000000E+00 - 12 0.0000000E+00 -1.9000560E+04 0.0000000E+00 0.0000000E+00 - 13 0.0000000E+00 -1.9021693E+04 0.0000000E+00 0.0000000E+00 - 14 0.0000000E+00 -1.8958078E+04 0.0000000E+00 0.0000000E+00 - 15 0.0000000E+00 -1.8825341E+04 0.0000000E+00 0.0000000E+00 - 16 0.0000000E+00 -1.8665837E+04 0.0000000E+00 0.0000000E+00 - 17 0.0000000E+00 -1.8526823E+04 0.0000000E+00 0.0000000E+00 - 18 0.0000000E+00 -1.8440053E+04 0.0000000E+00 0.0000000E+00 - 19 0.0000000E+00 -1.8413089E+04 0.0000000E+00 0.0000000E+00 - 20 0.0000000E+00 -1.8434556E+04 0.0000000E+00 0.0000000E+00 - 21 0.0000000E+00 -1.8480962E+04 0.0000000E+00 0.0000000E+00 - 22 0.0000000E+00 -1.8522340E+04 0.0000000E+00 0.0000000E+00 - 23 0.0000000E+00 -1.8525993E+04 0.0000000E+00 0.0000000E+00 - 24 0.0000000E+00 -1.8460698E+04 0.0000000E+00 0.0000000E+00 - 25 0.0000000E+00 -1.8296738E+04 0.0000000E+00 0.0000000E+00 - 26 0.0000000E+00 -1.8020270E+04 0.0000000E+00 0.0000000E+00 - 27 0.0000000E+00 -1.7640929E+04 0.0000000E+00 0.0000000E+00 - 28 0.0000000E+00 -1.7184038E+04 0.0000000E+00 0.0000000E+00 - 29 0.0000000E+00 -1.6683458E+04 0.0000000E+00 0.0000000E+00 - 30 0.0000000E+00 -1.6155762E+04 0.0000000E+00 0.0000000E+00 - 31 0.0000000E+00 -1.5563561E+04 0.0000000E+00 0.0000000E+00 - 32 0.0000000E+00 -1.4785670E+04 0.0000000E+00 0.0000000E+00 - 33 0.0000000E+00 -1.3603553E+04 0.0000000E+00 0.0000000E+00 - 34 0.0000000E+00 -1.1873456E+04 0.0000000E+00 0.0000000E+00 - 35 0.0000000E+00 -9.6897540E+03 0.0000000E+00 0.0000000E+00 - 36 0.0000000E+00 -7.4171597E+03 0.0000000E+00 0.0000000E+00 - 37 0.0000000E+00 -5.6141376E+03 0.0000000E+00 0.0000000E+00 - 38 0.0000000E+00 -4.7027245E+03 0.0000000E+00 0.0000000E+00 - 39 0.0000000E+00 -4.7596468E+03 0.0000000E+00 0.0000000E+00 - 40 0.0000000E+00 -5.6680031E+03 0.0000000E+00 0.0000000E+00 - 41 0.0000000E+00 -6.1323113E+03 0.0000000E+00 0.0000000E+00 -# qohme | MW/m^3 - 1 4.6130765E-02 - 2 4.6130765E-02 - 3 4.7789783E-02 - 4 5.2363318E-02 - 5 5.5364175E-02 - 6 5.9140283E-02 - 7 6.4067318E-02 - 8 7.0906607E-02 - 9 7.8079559E-02 - 10 8.5955397E-02 - 11 9.4737516E-02 - 12 1.0459736E-01 - 13 1.1589246E-01 - 14 1.2884447E-01 - 15 1.4319878E-01 - 16 1.5842548E-01 - 17 1.7175997E-01 - 18 1.7225355E-01 - 19 1.3308279E-01 - 20 6.7613494E-02 - 21 4.1156239E-02 - 22 3.6146986E-02 - 23 3.4006854E-02 - 24 3.1276394E-02 - 25 2.8345691E-02 - 26 2.5470393E-02 - 27 2.2773695E-02 - 28 2.0305332E-02 - 29 1.8090911E-02 - 30 1.6134738E-02 - 31 1.4432887E-02 - 32 1.2975629E-02 - 33 1.1746080E-02 - 34 1.0716319E-02 - 35 9.9026692E-03 - 36 9.2468074E-03 - 37 8.4933782E-03 - 38 7.4053481E-03 - 39 5.8425750E-03 - 40 3.8040674E-03 - 41 2.0009524E-03 -# qbeame | MW/m^3 - 1 5.5986838E-03 - 2 5.5986838E-03 - 3 5.6351768E-03 - 4 5.6550861E-03 - 5 5.6833173E-03 - 6 5.5695400E-03 - 7 5.5575399E-03 - 8 5.4599331E-03 - 9 5.3861793E-03 - 10 5.2581675E-03 - 11 5.1929722E-03 - 12 5.1971953E-03 - 13 5.1775196E-03 - 14 4.9901744E-03 - 15 4.9618998E-03 - 16 5.0845729E-03 - 17 5.0202464E-03 - 18 4.6564555E-03 - 19 6.0302687E-03 - 20 5.6319960E-03 - 21 5.0208355E-03 - 22 4.5896176E-03 - 23 4.4061531E-03 - 24 4.2584543E-03 - 25 4.2434980E-03 - 26 4.4038692E-03 - 27 4.5393166E-03 - 28 4.6001853E-03 - 29 4.7039031E-03 - 30 4.8518295E-03 - 31 4.9373472E-03 - 32 4.9000228E-03 - 33 4.7262114E-03 - 34 4.3806801E-03 - 35 3.9341851E-03 - 36 3.5724695E-03 - 37 3.3381735E-03 - 38 3.1902177E-03 - 39 3.1537835E-03 - 40 3.1630009E-03 - 41 3.1353493E-03 -# qbeami | MW/m^3 - 1 1.4405481E-02 - 2 1.4405481E-02 - 3 1.4317571E-02 - 4 1.4327663E-02 - 5 1.4143115E-02 - 6 1.3960983E-02 - 7 1.3820618E-02 - 8 1.3565788E-02 - 9 1.3189639E-02 - 10 1.2755732E-02 - 11 1.2174514E-02 - 12 1.1618291E-02 - 13 1.0987794E-02 - 14 1.0449563E-02 - 15 1.0153680E-02 - 16 1.0010544E-02 - 17 9.7800973E-03 - 18 9.5159852E-03 - 19 8.7131278E-03 - 20 7.2629089E-03 - 21 5.8477395E-03 - 22 4.8342608E-03 - 23 4.1985347E-03 - 24 3.7099999E-03 - 25 3.3467261E-03 - 26 3.1085471E-03 - 27 2.8696458E-03 - 28 2.6172263E-03 - 29 2.4111043E-03 - 30 2.2243511E-03 - 31 2.0162267E-03 - 32 1.7913977E-03 - 33 1.5603521E-03 - 34 1.3199292E-03 - 35 1.0880217E-03 - 36 8.8893628E-04 - 37 7.0310167E-04 - 38 5.1708801E-04 - 39 3.4922111E-04 - 40 2.1365043E-04 - 41 1.3159068E-04 -# qrfe | MW/m^3 - 1 0.0000000E+00 - 2 0.0000000E+00 - 3 2.0760634E-16 - 4 1.9255264E-03 - 5 1.5408664E-01 - 6 1.7594500E+00 - 7 3.7113208E+00 - 8 1.4050996E+00 - 9 8.3454415E-02 - 10 1.7835978E-04 - 11 1.9043850E-09 - 12 1.9464523E-21 - 13 0.0000000E+00 - 14 0.0000000E+00 - 15 0.0000000E+00 - 16 0.0000000E+00 - 17 0.0000000E+00 - 18 0.0000000E+00 - 19 0.0000000E+00 - 20 0.0000000E+00 - 21 0.0000000E+00 - 22 0.0000000E+00 - 23 0.0000000E+00 - 24 0.0000000E+00 - 25 0.0000000E+00 - 26 0.0000000E+00 - 27 0.0000000E+00 - 28 0.0000000E+00 - 29 0.0000000E+00 - 30 0.0000000E+00 - 31 0.0000000E+00 - 32 0.0000000E+00 - 33 0.0000000E+00 - 34 0.0000000E+00 - 35 0.0000000E+00 - 36 0.0000000E+00 - 37 0.0000000E+00 - 38 0.0000000E+00 - 39 0.0000000E+00 - 40 0.0000000E+00 - 41 0.0000000E+00 -# qbrem | MW/m^3 - 1 1.1407030E-03 - 2 1.1407030E-03 - 3 1.1372623E-03 - 4 1.1302036E-03 - 5 1.1195274E-03 - 6 1.1052900E-03 - 7 1.0877081E-03 - 8 1.0670703E-03 - 9 1.0436777E-03 - 10 1.0177800E-03 - 11 9.8969303E-04 - 12 9.5970617E-04 - 13 9.2808271E-04 - 14 8.9509314E-04 - 15 8.6102423E-04 - 16 8.2620395E-04 - 17 7.9093690E-04 - 18 7.5511073E-04 - 19 7.1828214E-04 - 20 6.8029120E-04 - 21 6.4125602E-04 - 22 6.0150922E-04 - 23 5.6155372E-04 - 24 5.2197510E-04 - 25 4.8332172E-04 - 26 4.4608988E-04 - 27 4.1067157E-04 - 28 3.7731206E-04 - 29 3.4617295E-04 - 30 3.1730024E-04 - 31 2.9069053E-04 - 32 2.6626369E-04 - 33 2.4390307E-04 - 34 2.2346944E-04 - 35 2.0482013E-04 - 36 1.8643430E-04 - 37 1.6579858E-04 - 38 1.4218127E-04 - 39 1.1654208E-04 - 40 8.9846896E-05 - 41 5.3419278E-05 -# qsync | MW/m^3 - 1 3.4265930E-05 - 2 3.4265930E-05 - 3 3.3868640E-05 - 4 3.3080967E-05 - 5 3.1936549E-05 - 6 3.0487198E-05 - 7 2.8800489E-05 - 8 2.6948354E-05 - 9 2.4997433E-05 - 10 2.3003629E-05 - 11 2.1016937E-05 - 12 1.9077037E-05 - 13 1.7213967E-05 - 14 1.5450536E-05 - 15 1.3803203E-05 - 16 1.2282945E-05 - 17 1.0894861E-05 - 18 9.6228222E-06 - 19 8.4396766E-06 - 20 7.3335824E-06 - 21 6.3040395E-06 - 22 5.3567997E-06 - 23 4.4999470E-06 - 24 3.7392995E-06 - 25 3.0769638E-06 - 26 2.5111427E-06 - 27 2.0356874E-06 - 28 1.6413022E-06 - 29 1.3180942E-06 - 30 1.0555725E-06 - 31 8.4399399E-07 - 32 6.7426784E-07 - 33 5.3852648E-07 - 34 4.3017446E-07 - 35 3.4374068E-07 - 36 2.6661549E-07 - 37 1.8758491E-07 - 38 1.1339990E-07 - 39 5.6777628E-08 - 40 2.3049458E-08 - 41 7.9937621E-09 -# qline | MW/m^3 - 1 1.8309857E-04 - 2 1.8309857E-04 - 3 1.8350316E-04 - 4 1.8429378E-04 - 5 1.8544951E-04 - 6 1.8692874E-04 - 7 1.8868706E-04 - 8 1.9067508E-04 - 9 1.9284469E-04 - 10 1.9514846E-04 - 11 1.9754279E-04 - 12 1.9998356E-04 - 13 2.0242588E-04 - 14 2.0482330E-04 - 15 2.0712826E-04 - 16 2.0929799E-04 - 17 2.1128816E-04 - 18 2.1310687E-04 - 19 2.1482555E-04 - 20 2.1651461E-04 - 21 2.1824110E-04 - 22 2.2006086E-04 - 23 2.2201490E-04 - 24 2.2414797E-04 - 25 2.2648244E-04 - 26 2.2901990E-04 - 27 2.3176828E-04 - 28 2.3474486E-04 - 29 2.3794270E-04 - 30 2.4135784E-04 - 31 2.4497734E-04 - 32 2.4880270E-04 - 33 2.5283919E-04 - 34 2.5708747E-04 - 35 2.6156744E-04 - 36 2.6826856E-04 - 37 2.8211328E-04 - 38 3.0861058E-04 - 39 3.5508440E-04 - 40 4.2766988E-04 - 41 4.1913117E-04 -# qei | MW/m^3 - 1 4.6352476E-02 - 2 4.6352476E-02 - 3 4.6449918E-02 - 4 4.6590313E-02 - 5 4.6730615E-02 - 6 4.6851438E-02 - 7 4.6940893E-02 - 8 4.6996802E-02 - 9 4.7022387E-02 - 10 4.7018569E-02 - 11 4.6979334E-02 - 12 4.6885245E-02 - 13 4.6715237E-02 - 14 4.6457306E-02 - 15 4.6117992E-02 - 16 4.5726753E-02 - 17 4.5329079E-02 - 18 4.4966823E-02 - 19 4.4660799E-02 - 20 4.4383728E-02 - 21 4.4045232E-02 - 22 4.3548464E-02 - 23 4.2837095E-02 - 24 4.1892592E-02 - 25 4.0722365E-02 - 26 3.9371524E-02 - 27 3.7913289E-02 - 28 3.6414353E-02 - 29 3.4914721E-02 - 30 3.3400515E-02 - 31 3.1850715E-02 - 32 3.0356975E-02 - 33 2.9279283E-02 - 34 2.9142444E-02 - 35 3.0439143E-02 - 36 3.2773125E-02 - 37 3.5104263E-02 - 38 3.6709147E-02 - 39 3.6878544E-02 - 40 3.4035776E-02 - 41 2.3958948E-02 -# qione | MW/m^3 - 1 5.3910153E-04 - 2 5.3910153E-04 - 3 6.4662612E-04 - 4 5.1594719E-04 - 5 5.4196026E-04 - 6 5.1458619E-04 - 7 4.8159116E-04 - 8 4.4626645E-04 - 9 4.0458568E-04 - 10 3.5937640E-04 - 11 3.1636265E-04 - 12 2.7567996E-04 - 13 2.3617426E-04 - 14 1.9521375E-04 - 15 -2.3638917E-04 - 16 -4.5174482E-03 - 17 -8.7415949E-03 - 18 -1.2308841E-02 - 19 -1.5288069E-02 - 20 -1.7836978E-02 - 21 -1.9886693E-02 - 22 -2.1229683E-02 - 23 -2.1930009E-02 - 24 -2.2214804E-02 - 25 -2.2113272E-02 - 26 -2.1861081E-02 - 27 -2.1553233E-02 - 28 -2.1368320E-02 - 29 -2.1503486E-02 - 30 -2.1947920E-02 - 31 -2.2857179E-02 - 32 -2.3890393E-02 - 33 -2.4594876E-02 - 34 -2.4244476E-02 - 35 -2.2430691E-02 - 36 -1.9646313E-02 - 37 -1.8776171E-02 - 38 -2.5329159E-02 - 39 -4.6660067E-02 - 40 -8.5298847E-02 - 41 -1.3110814E-01 -# qioni | MW/m^3 - 1 -4.3365118E-05 - 2 -4.3365118E-05 - 3 4.6273391E-05 - 4 1.1687450E-04 - 5 2.5538126E-04 - 6 4.0699940E-04 - 7 6.1636439E-04 - 8 8.5328732E-04 - 9 1.0708240E-03 - 10 1.2767844E-03 - 11 1.4871044E-03 - 12 1.6969209E-03 - 13 1.8934763E-03 - 14 2.0604200E-03 - 15 2.1918021E-03 - 16 2.2936575E-03 - 17 2.3538625E-03 - 18 2.3073135E-03 - 19 2.0717009E-03 - 20 1.6262683E-03 - 21 1.1805596E-03 - 22 9.0621738E-04 - 23 7.7980288E-04 - 24 7.1539913E-04 - 25 6.6939431E-04 - 26 6.3212188E-04 - 27 5.9433255E-04 - 28 5.5217150E-04 - 29 5.1022364E-04 - 30 4.4632849E-04 - 31 3.7097058E-04 - 32 2.7705870E-04 - 33 1.8096651E-04 - 34 1.1580443E-04 - 35 5.9330882E-05 - 36 8.4653834E-06 - 37 -1.4957496E-04 - 38 -4.2324867E-04 - 39 -7.9818329E-04 - 40 -1.3273326E-03 - 41 -2.1180075E-03 -# qpar_beam | MW/m^3 - 1 2.4985475E+19 - 2 2.4985475E+19 - 3 2.4909300E+19 - 4 2.4469435E+19 - 5 2.3722559E+19 - 6 2.2542627E+19 - 7 2.1171301E+19 - 8 1.9730667E+19 - 9 1.8298035E+19 - 10 1.6918777E+19 - 11 1.5621143E+19 - 12 1.4446000E+19 - 13 1.3424083E+19 - 14 1.2543931E+19 - 15 1.1783028E+19 - 16 1.1129678E+19 - 17 1.0623925E+19 - 18 1.0318852E+19 - 19 1.0149944E+19 - 20 9.7980021E+18 - 21 9.3636597E+18 - 22 9.1245919E+18 - 23 9.0601580E+18 - 24 9.0636970E+18 - 25 9.1298515E+18 - 26 9.2883475E+18 - 27 9.5562222E+18 - 28 9.9814408E+18 - 29 1.0558753E+19 - 30 1.1457713E+19 - 31 1.2659438E+19 - 32 1.4369589E+19 - 33 1.6856701E+19 - 34 2.0552333E+19 - 35 2.9784198E+19 - 36 4.5386516E+19 - 37 6.6375066E+19 - 38 9.6410362E+19 - 39 1.4417803E+20 - 40 2.2775269E+20 - 41 3.6342333E+20 -# qmom | MW/m^3 - 1 -1.2650594E-02 - 2 -1.2650594E-02 - 3 -1.4097844E-02 - 4 -1.6653485E-02 - 5 -1.7668344E-02 - 6 -1.7228887E-02 - 7 -1.5860580E-02 - 8 -1.4932274E-02 - 9 -1.5723618E-02 - 10 -1.7290906E-02 - 11 -2.0566867E-02 - 12 -2.6171732E-02 - 13 -3.1534842E-02 - 14 -3.5252005E-02 - 15 -3.8457154E-02 - 16 -4.2292588E-02 - 17 -4.6056189E-02 - 18 -4.6690907E-02 - 19 -4.2154801E-02 - 20 -3.6650383E-02 - 21 -3.2884915E-02 - 22 -3.0288787E-02 - 23 -2.8592868E-02 - 24 -2.7244455E-02 - 25 -2.5390866E-02 - 26 -2.3772698E-02 - 27 -2.2751361E-02 - 28 -2.1762852E-02 - 29 -2.0286211E-02 - 30 -1.8369397E-02 - 31 -1.6560286E-02 - 32 -1.4846543E-02 - 33 -1.2943964E-02 - 34 -1.1145865E-02 - 35 -9.6091325E-03 - 36 -7.9979634E-03 - 37 -6.2357745E-03 - 38 -4.4916535E-03 - 39 -2.5427757E-03 - 40 3.7336849E-04 - 41 5.0418548E-03 diff --git a/regressions/portals_regressions.py b/regressions/portals_regressions.py deleted file mode 100644 index 1afbe7b3..00000000 --- a/regressions/portals_regressions.py +++ /dev/null @@ -1,155 +0,0 @@ -import argparse -import torch -import os -from mitim_tools.misc_tools import LOGtools -from mitim_tools.opt_tools import STRATEGYtools -from mitim_modules.portals import PORTALSmain -from mitim_modules.powertorch.physics import TRANSPORTtools -from mitim_tools.misc_tools.LOGtools import printMsg as print -from mitim_tools import __mitimroot__ -from IPython import embed - -# Get test number -parser = argparse.ArgumentParser() -parser.add_argument("test", type=int) -args = parser.parse_args() -test = args.test - -if test == 0: - tests = [1,2] -else: - tests = [test] - -# Set up case -inputgacode = __mitimroot__ + "/regressions/data/input.gacode" - -# --------------------------------------------------------------------------------------------- -# TESTS -# --------------------------------------------------------------------------------------------- - -def conditions_regressions(variables): - - conditions = True - - # Checks - for var in variables: - conditions &= var[0] == var[1] - - # Results - if conditions: - print("\t PASSED") - else: - print("\t FAILED",typeMsg='w') - - -for test in tests: - - folderWork = __mitimroot__ + "/regressions/scratch/portals_regression_{test}/" - - if test == 1: - - print("\n>>>>> Running PORTALS test 1: Standard run with constant diffusivities") - - os.system(f"rm -rf {folderWork} && mkdir {folderWork}") - with LOGtools.redirect_all_output_to_file(f'{folderWork}/regression.log'): - portals_fun = PORTALSmain.portals(folderWork) - portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 - portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - portals_fun.portals_parameters["initialization"]["remove_fast"] = True - - portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti"] - portals_fun.optimization_options["acquisition_options"]["optimizers"] = ["botorch"] - - portals_fun.portals_parameters["transport"]["evaluator"] = TRANSPORTtools.diffusion_model - transport_evaluator_options = {'chi_e': torch.ones(5)*0.5,'chi_i': torch.ones(5)*2.0} - - portals_fun.prep(inputgacode, folderWork, transport_evaluator_options=transport_evaluator_options) - mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=False, askQuestions=False) - mitim_bo.run() - - # Checks - conditions_regressions([ - [mitim_bo.optimization_data.data['QeTurb_1'][0],0.0129878663484079], - [mitim_bo.optimization_data.data['QeTurb_1'][1],0.0174629359509858], - [mitim_bo.optimization_data.data['QeTurb_1'][2],0.0222306543202599], - [mitim_bo.optimization_data.data['QeTurb_1'][3],0.0037220182305746], - [mitim_bo.optimization_data.data['QeTurb_1'][4],0.0301250769357799], - [mitim_bo.optimization_data.data['QeTurb_1'][5],0.0436471750834417], - [mitim_bo.optimization_data.data['QiTurb_5'][0],0.0114099018688661], - [mitim_bo.optimization_data.data['QiTurb_5'][1],0.0103728562456646], - [mitim_bo.optimization_data.data['QiTurb_5'][2],0.0095916319760464], - [mitim_bo.optimization_data.data['QiTurb_5'][3],0.0063868247281859], - [mitim_bo.optimization_data.data['QiTurb_5'][4],0.0062216868661381], - [mitim_bo.optimization_data.data['QiTurb_5'][5],0.0061692702220821], - ]) - - if test == 2: - - print("\n>>>>> Running PORTALS test 2: Standard run with TGLF") - - os.system(f"rm -rf {folderWork} && mkdir {folderWork}") - with LOGtools.redirect_all_output_to_file(f'{folderWork}/regression.log'): - - portals_fun = PORTALSmain.portals(folderWork) - portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 - portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - portals_fun.portals_parameters["solution"]["predicted_rho"] = [0.25, 0.45, 0.65, 0.85] - portals_fun.portals_parameters["initialization"]["remove_fast"] = True - portals_fun.portals_parameters["initialization"]["quasineutrality"] = True - portals_fun.portals_parameters["initialization"]["enforce_same_aLn"] = True - portals_fun.portals_parameters["transport"]["options"]["code_settings"] = 2 - - portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti", "ne"] - - portals_fun.prep(inputgacode, folderWork) - mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=False, askQuestions=False) - mitim_bo.run() - - # Checks - conditions_regressions([ - [mitim_bo.optimization_data.data['QeTar_3'][0],0.0276660734686889], - [mitim_bo.optimization_data.data['QeTar_3'][1],0.026050457428488], - [mitim_bo.optimization_data.data['QeTar_3'][2],0.0245681162983153], - [mitim_bo.optimization_data.data['QeTar_3'][3],0.0225138750256145], - [mitim_bo.optimization_data.data['QeTar_3'][4],0.0238676726307135], - [mitim_bo.optimization_data.data['QiTurb_4'][0],0.01904210194957 ], - [mitim_bo.optimization_data.data['QiTurb_4'][1],0.015054384849328], - [mitim_bo.optimization_data.data['QiTurb_4'][2],0.012453620533174], - [mitim_bo.optimization_data.data['QiTurb_4'][3],0.009167817359775], - [mitim_bo.optimization_data.data['QiTurb_4'][4],0.010592748091966], - [mitim_bo.optimization_data.data['QeTurb_1'][0],0.0008148021791468 ], - [mitim_bo.optimization_data.data['QeTurb_1'][1],0.005048271135896 ], - [mitim_bo.optimization_data.data['QeTurb_1'][2],0.0316597732275 ], - [mitim_bo.optimization_data.data['QeTurb_1'][3],0.4672666906836 ], - [mitim_bo.optimization_data.data['QeTurb_1'][4],-0.0006023859321252], - ]) - - if test == 3: - - print("\n>>>>> Running PORTALS test 3: Run with TGLF multi-channel") - - # os.system(f"rm -rf {folderWork} && mkdir {folderWork}") - # with LOGtools.redirect_all_output_to_file(f'{folderWork}/regression.log'): - - # portals_fun = PORTALSmain.portals(folderWork) - # portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 2 - # portals_fun.optimization_options["initialization_options"]["initial_training"] = 3 - # portals_fun.portals_parameters["initialization"]["remove_fast"] = True - - # portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti", "ne",'nZ','w0'] - - # portals_fun.portals_parameters["solution"]["trace_impurity"] = 'W' - # portals_fun.portals_parameters["solution"]["turbulent_exchange_as_surrogate"] = True - - # portals_fun.prep(inputgacode, folderWork) - # mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=False, askQuestions=False) - # mitim_bo.run() - - # with open(mitim_bo.optimization_object.optimization_extra, "rb") as f: - # mitim_runs = pickle_dill.load(f) - - # # Checks - # conditions_regressions([ - # [mitim_bo.optimization_data.data['QeTurb_1'][5],0.0713711320661], - # [mitim_runs[5]['powerstate'].plasma['QieMWm3_tr_turb'][0,3].item(),-0.0009466626542564001] - # ]) diff --git a/tutorials/TGLF_tutorial.py b/tutorials/TGLF_tutorial.py index 6b73f070..501dbfc3 100644 --- a/tutorials/TGLF_tutorial.py +++ b/tutorials/TGLF_tutorial.py @@ -10,7 +10,8 @@ tglf = TGLFtools.TGLF(rhos=[0.5, 0.7]) # Prepare the TGLF class -cdf = tglf.prep_using_tgyro(folder, inputgacode=inputgacode_file, cold_start=False) +tglf.prep(inputgacode_file,folder, cold_start=False) + ''' *************************************************************************** @@ -21,7 +22,7 @@ # Run TGLF in subfolder tglf.run( subfolder="yes_em_folder", - code_settings=5, + code_settings="SAT2em", extraOptions={}, cold_start=False ) @@ -32,7 +33,7 @@ # Run TGLF in a different subfolder with different settings tglf.run( subfolder="no_em_folder", - code_settings=5, + code_settings="SAT2em", extraOptions={"USE_BPER": False}, cold_start=False, ) @@ -50,7 +51,7 @@ ''' tglf.runScan( subfolder = 'scan1', - code_settings = 5, + code_settings = "SAT2em", cold_start = False, variable = 'RLTS_1', varUpDown = np.linspace(0.5,1.5,3)) @@ -58,7 +59,7 @@ tglf.runScan( subfolder = 'scan2', - code_settings = 5, + code_settings = "SAT2em", cold_start = False, variable = 'RLTS_2', varUpDown = np.linspace(0.5,1.5,3)) @@ -75,21 +76,21 @@ tglf.runScanTurbulenceDrives( subfolder = 'turb_drives', - code_settings = 5, + code_settings = "SAT2em", cold_start = False) tglf.plotScanTurbulenceDrives(label='turb_drives') ''' *************************************************************************** -Automatic scan of turbulence drives +Analysis of chi incremental *************************************************************************** ''' tglf.runAnalysis( subfolder = 'chi_e', analysisType = 'chi_e', - code_settings = 5, + code_settings = "SAT2em", cold_start = False, label = 'chi_eu') @@ -110,3 +111,4 @@ tglf.read(label=f'settings{i}') tglf.plot(labels=[f'settings{i}' for i in range(1,6)]) + From 05abb66df9ff535e64ea8cdb3cd46ec464dcf9f4 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 16:53:56 -0400 Subject: [PATCH 252/385] Corrected maestro template namelist --- templates/maestro_namelist.json | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json index 93780736..b0d712dc 100644 --- a/templates/maestro_namelist.json +++ b/templates/maestro_namelist.json @@ -81,9 +81,11 @@ }, "transport": { "options": { - "run": { - "code_settings": "SAT2astra", - "extraOptions": {"USE_BPER": true} + "tglf": { + "run": { + "code_settings": "SAT2astra", + "extraOptions": {"USE_BPER": true} + } } } }, From 940d953164d5a028a6cf624f7355e434f2763777 Mon Sep 17 00:00:00 2001 From: audreysa Date: Fri, 29 Aug 2025 18:13:37 -0400 Subject: [PATCH 253/385] Fix location of elif statement for adding eped_initializer from eped values in namelist --- src/mitim_modules/maestro/scripts/run_maestro.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index 7ba52d64..3282ce5e 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -121,6 +121,8 @@ def parse_maestro_nml(file_path): beat_namelist = maestro_namelist["maestro"][f"{beat_type}_beat"][f"{beat_type}_namelist"] + + # *************************************************************************** # Nothin yet # *************************************************************************** @@ -148,7 +150,7 @@ def profiles_postprocessing_fun(file_profs): return p beat_namelist['portals_parameters']['transport']['profiles_postprocessing_fun'] = profiles_postprocessing_fun - elif beat_type == "eped_initializer" and "eped_beat" in maestro_namelist["maestro"]: + elif beat_type == "eped_initializer" and "eped_beat" in maestro_namelist["maestro"]: print('Using the eped_beat namelist for the eped_initializer') beat_namelist = maestro_namelist["maestro"]["eped_beat"]["eped_namelist"] From 25339a026808af311d64b7cd312ddabd32f5e655 Mon Sep 17 00:00:00 2001 From: audreysa Date: Fri, 29 Aug 2025 18:27:43 -0400 Subject: [PATCH 254/385] Fix error in last commit --- src/mitim_modules/maestro/scripts/run_maestro.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index 3282ce5e..77292d46 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -121,6 +121,9 @@ def parse_maestro_nml(file_path): beat_namelist = maestro_namelist["maestro"][f"{beat_type}_beat"][f"{beat_type}_namelist"] + elif beat_type == "eped_initializer" and "eped_beat" in maestro_namelist["maestro"]: + print('Using the eped_beat namelist for the eped_initializer') + beat_namelist = maestro_namelist["maestro"]["eped_beat"]["eped_namelist"] # *************************************************************************** @@ -150,9 +153,6 @@ def profiles_postprocessing_fun(file_profs): return p beat_namelist['portals_parameters']['transport']['profiles_postprocessing_fun'] = profiles_postprocessing_fun - elif beat_type == "eped_initializer" and "eped_beat" in maestro_namelist["maestro"]: - print('Using the eped_beat namelist for the eped_initializer') - beat_namelist = maestro_namelist["maestro"]["eped_beat"]["eped_namelist"] else: raise ValueError(f"[MITIM] {beat_type} beat not found in the MAESTRO namelist") From 094eb380c2003d30bb83469fd058a4dae8f6a9eb Mon Sep 17 00:00:00 2001 From: audreysa Date: Fri, 29 Aug 2025 18:33:43 -0400 Subject: [PATCH 255/385] Actually saving the file modified in the last commit --- src/mitim_modules/maestro/scripts/run_maestro.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index 77292d46..3282ce5e 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -121,9 +121,6 @@ def parse_maestro_nml(file_path): beat_namelist = maestro_namelist["maestro"][f"{beat_type}_beat"][f"{beat_type}_namelist"] - elif beat_type == "eped_initializer" and "eped_beat" in maestro_namelist["maestro"]: - print('Using the eped_beat namelist for the eped_initializer') - beat_namelist = maestro_namelist["maestro"]["eped_beat"]["eped_namelist"] # *************************************************************************** @@ -153,6 +150,9 @@ def profiles_postprocessing_fun(file_profs): return p beat_namelist['portals_parameters']['transport']['profiles_postprocessing_fun'] = profiles_postprocessing_fun + elif beat_type == "eped_initializer" and "eped_beat" in maestro_namelist["maestro"]: + print('Using the eped_beat namelist for the eped_initializer') + beat_namelist = maestro_namelist["maestro"]["eped_beat"]["eped_namelist"] else: raise ValueError(f"[MITIM] {beat_type} beat not found in the MAESTRO namelist") From 452a51b2f3d90ce1c007e0a4e5ee1d514d05324a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 19:08:43 -0400 Subject: [PATCH 256/385] Fixed bug with TGLF error bars --- .../powertorch/physics_models/transport_tglf.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 48be9efb..ba2a61eb 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -248,13 +248,13 @@ def _run_tglf_uncertainty_model( if remove_folders_out: IOtools.shutil_rmtree(tglf.FolderGACODE) - Qe = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - Qi = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - Qifast = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - Ge = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - GZ = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - Mt = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) - S = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan)+1 )) + Qe = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) + Qi = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) + Qifast = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) + Ge = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) + GZ = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) + Mt = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) + S = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) cont = 0 for vari in variables_to_scan: From 60b3d9582aa2f512d52ffcc5c58d3000d15efe3a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 20:59:31 -0400 Subject: [PATCH 257/385] Wrote functions to write yaml files to keep track --- src/mitim_modules/portals/PORTALSmain.py | 33 ++---- src/mitim_tools/misc_tools/IOtools.py | 123 +++++++++++++++++++++ src/mitim_tools/opt_tools/STRATEGYtools.py | 3 + 3 files changed, 137 insertions(+), 22 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index c0ca0cb7..574b9e51 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -1,7 +1,6 @@ import shutil import torch import copy -import yaml, importlib from collections import OrderedDict import numpy as np import dill as pickle_dill @@ -49,11 +48,12 @@ def __init__( # Read PORTALS namelist (if not provided, use default) if portals_namelist is None: - portals_namelist = __mitimroot__ / "templates" / "portals.namelist.yaml" - print(f"\t- No PORTALS namelist provided, using default in {IOtools.clipstr(portals_namelist)}") + self.portals_namelist = __mitimroot__ / "templates" / "portals.namelist.yaml" + print(f"\t- No PORTALS namelist provided, using default in {IOtools.clipstr(self.portals_namelist)}") else: - print(f"\t- Using provided PORTALS namelist in {IOtools.clipstr(portals_namelist)}") - self.portals_parameters = read_portals_nml(portals_namelist) + self.portals_namelist = portals_namelist + print(f"\t- Using provided PORTALS namelist in {IOtools.clipstr(self.portals_namelist)}") + self.portals_parameters = IOtools.read_mitim_yaml(self.portals_namelist) # Apply the optimization options to the proper namelist and drop it from portals_parameters if 'optimization_options' in self.portals_parameters: @@ -166,6 +166,12 @@ def prep( else: print("\t- extrapointsModels already defined, not changing") + # Make a copy of the namelist that was imported to the folder + shutil.copy(self.portals_namelist, self.folder / "portals.namelist_original.yaml") + + # Write the parameters (after script modification) to a yaml namelist for tracking purposes + IOtools.write_mitim_yaml(self.portals_parameters, self.folder / "portals.namelist.yaml") + def _define_reuse_models(self): ''' The user can define a list of strings to avoid reusing surrogates. @@ -530,20 +536,3 @@ def analyze_results( portals_full.runCases(onlyBest=onlyBest, cold_start=cold_start, fn=fn) return portals_full.opt_fun.mitim_model.optimization_object - -def read_portals_nml(path: str): - - def resolve(x): - if isinstance(x, dict): - return {k: resolve(v) for k, v in x.items()} - if isinstance(x, list): - return [resolve(v) for v in x] - if isinstance(x, str) and x.startswith("import::"): - modattr = x[len("import::"):] - module_name, attr = modattr.rsplit(".", 1) - return getattr(importlib.import_module(module_name), attr) - return x - - with open(path, "r") as f: - cfg = yaml.safe_load(f) - return resolve(cfg) \ No newline at end of file diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index 38f185b5..fa602265 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -21,6 +21,8 @@ import json import functools import hashlib +import yaml, importlib +from typing import Any, Mapping from collections import OrderedDict from pathlib import Path import platform @@ -2158,3 +2160,124 @@ def find_class(self, module, name): print(f"\t\tModule not found: {module} {name}; returning dummy", typeMsg="i") return super().find_class("torch._utils", name) +def read_mitim_yaml(path: str): + + def resolve(x): + if isinstance(x, dict): + return {k: resolve(v) for k, v in x.items()} + if isinstance(x, list): + return [resolve(v) for v in x] + if isinstance(x, str) and x.startswith("import::"): + modattr = x[len("import::"):] + module_name, attr = modattr.rsplit(".", 1) + return getattr(importlib.import_module(module_name), attr) + return x + + with open(path, "r") as f: + cfg = yaml.safe_load(f) + return resolve(cfg) + + +import yaml +import numpy as np +import inspect +from pathlib import Path +from typing import Any, Mapping + +def _as_import_string(obj: Any) -> str: + """ + Return an 'import::module.qualname' for callables/classes. + Falls back to str(obj) if module/name aren't available. + """ + # Handle bound methods + if inspect.ismethod(obj): + func = obj.__func__ + mod = func.__module__ + qn = getattr(func, "__qualname__", func.__name__) + qn = qn.replace(".", "") + return f"import::{mod}.{qn}" + # Handle functions, classes, other callables with module/name + if inspect.isfunction(obj) or inspect.isbuiltin(obj) or inspect.isclass(obj) or callable(obj): + mod = getattr(obj, "__module__", None) + name = getattr(obj, "__qualname__", getattr(obj, "__name__", None)) + if mod and name: + name = name.replace(".", "") + return f"import::{mod}.{name}" + return f"import::{str(obj)}" + + # Strings: if they already look like import::..., keep; otherwise just return + if isinstance(obj, str): + return obj + + # Fallback + return f"import::{str(obj)}" + +def _normalize_for_yaml(obj: Any) -> Any: + """ + Recursively convert objects into YAML-safe Python builtins. + - NumPy arrays/scalars -> lists/scalars + - Paths -> str + - sets -> lists + - callables/classes/methods -> 'import::module.qualname' + Leaves basic builtins as-is. + """ + # NumPy + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, np.generic): + return obj.item() + + # Simple builtins + if obj is None or isinstance(obj, (bool, int, float, str)): + return obj + + # Path-like + if isinstance(obj, (Path, )): + return str(obj) + + # Sets -> lists + if isinstance(obj, (set, frozenset)): + return [_normalize_for_yaml(v) for v in obj] + + # Mappings + if isinstance(obj, Mapping): + # ensure keys are YAML-safe (coerce to str if needed) + return {str(k): _normalize_for_yaml(v) for k, v in obj.items()} + + # Sequences + if isinstance(obj, (list, tuple)): + return [_normalize_for_yaml(v) for v in obj] + + # Anything callable or class-like -> import string + if inspect.ismethod(obj) or inspect.isfunction(obj) or inspect.isbuiltin(obj) or inspect.isclass(obj) or callable(obj): + return _as_import_string(obj) + + # Fallback to str for unknown objects + return str(obj) + +class _NoAliasDumper(yaml.SafeDumper): + def ignore_aliases(self, data): + return True + def increase_indent(self, flow=False, indentless=False): + return super().increase_indent(flow, indentless=False) + +def write_mitim_yaml(parameters: Mapping[str, Any], path: str) -> None: + """ + General YAML writer: + - No assumptions about keys (works for solution/transport/target and also optimization_options). + - Normalizes everything to YAML-safe types, including function objects. + """ + if not isinstance(parameters, Mapping): + raise TypeError("parameters must be a dict-like mapping") + clean = _normalize_for_yaml(parameters) + + with open(path, "w", encoding="utf-8") as f: + yaml.dump( + clean, + f, + Dumper=_NoAliasDumper, + sort_keys=False, + default_flow_style=False, + allow_unicode=True, + width=1000, + ) \ No newline at end of file diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 14691071..51e06e24 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -359,6 +359,9 @@ def __init__( self.askQuestions = askQuestions self.seed = seed self.avoidPoints = [] + + # Write the optimizaiton parameters stored in the object, into a file + IOtools.write_mitim_yaml(self.optimization_object.optimization_options, self.optimization_object.folder / "optimization.namelist.yaml") if self.optimization_object.name_objectives is None: self.optimization_object.name_objectives = "y" From b0264129163ba512e5d4bf69e87c75850dabf8f1 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 29 Aug 2025 21:34:27 -0400 Subject: [PATCH 258/385] bug fix --- src/mitim_tools/opt_tools/STRATEGYtools.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 51e06e24..1c3b2d8a 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -361,7 +361,8 @@ def __init__( self.avoidPoints = [] # Write the optimizaiton parameters stored in the object, into a file - IOtools.write_mitim_yaml(self.optimization_object.optimization_options, self.optimization_object.folder / "optimization.namelist.yaml") + if (self.optimization_object.optimization_options is not None) and (self.optimization_object.folder is not None): + IOtools.write_mitim_yaml(self.optimization_object.optimization_options, self.optimization_object.folder / "optimization.namelist.yaml") if self.optimization_object.name_objectives is None: self.optimization_object.name_objectives = "y" From 210da3b1699a467496987decbea3bac84f6b0c5d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 30 Aug 2025 00:57:09 -0400 Subject: [PATCH 259/385] misc --- src/mitim_modules/powertorch/utils/TRANSFORMtools.py | 6 ++---- src/mitim_tools/misc_tools/IOtools.py | 7 ++++++- src/mitim_tools/opt_tools/STRATEGYtools.py | 11 ++--------- src/mitim_tools/plasmastate_tools/MITIMstate.py | 4 +--- 4 files changed, 11 insertions(+), 17 deletions(-) diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 68c68d44..1b3fdea4 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -422,9 +422,7 @@ def defineIons(self, input_gacode, rho_vec, dfT): self.plasma["ni"], mi, Zi, c_rad = [], [], [], [] for i in range(len(input_gacode.profiles["mass"])): if input_gacode.profiles["type"][i] == "[therm]": - self.plasma["ni"].append( - interpolation_function(rho_vec, rho_use, input_gacode.profiles["ni(10^19/m^3)"][:, i]) - ) + self.plasma["ni"].append(interpolation_function(rho_vec, rho_use, input_gacode.profiles["ni(10^19/m^3)"][:, i])) mi.append(input_gacode.profiles["mass"][i]) Zi.append(input_gacode.profiles["z"][i]) @@ -433,7 +431,7 @@ def defineIons(self, input_gacode, rho_vec, dfT): try: c = data_df[data_df['Ion'].str.lower()==input_gacode.profiles["name"][i].lower()].to_numpy()[0,2:].astype(float) except IndexError: - print(f'\t- Specie {input_gacode.profiles["name"][i]} not found in ADAS database, assuming zero radiation from it',typeMsg="w") + print(f'\t- Specie {input_gacode.profiles["name"][i]} not found in radiation database, assuming zero radiation from it',typeMsg="w") c = [-1e10, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] c_rad.append(c) diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index fa602265..e6d1601e 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -608,7 +608,7 @@ def read_mitim_nml(json_file): return optimization_options -def curate_mitim_nml(optimization_options, stopping_criteria_default = None): +def curate_mitim_nml(optimization_options, folder, stopping_criteria_default = None): # Optimization criterion if optimization_options['convergence_options']['stopping_criteria'] is None: @@ -634,6 +634,11 @@ def opt_crit(*args,**kwargs): if ikey not in Optim_potential: print(f"\t- Option {ikey} is an unexpected variable, prone to errors", typeMsg="q") + # Write the optimization parameters stored in the object, into a file + if folder is not None: + write_mitim_yaml(optimization_options, folder / "optimization.namelist.yaml") + print(f" --> Optimization namelist written to {folder / 'optimization.namelist.yaml'}") + return optimization_options # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 1c3b2d8a..0ecff21c 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -360,19 +360,11 @@ def __init__( self.seed = seed self.avoidPoints = [] - # Write the optimizaiton parameters stored in the object, into a file - if (self.optimization_object.optimization_options is not None) and (self.optimization_object.folder is not None): - IOtools.write_mitim_yaml(self.optimization_object.optimization_options, self.optimization_object.folder / "optimization.namelist.yaml") - if self.optimization_object.name_objectives is None: self.optimization_object.name_objectives = "y" # Folders and Logger - self.folderExecution = ( - IOtools.expandPath(self.optimization_object.folder) - if (self.optimization_object.folder is not None) - else Path("") - ) + self.folderExecution = IOtools.expandPath(self.optimization_object.folder) if (self.optimization_object.folder is not None) else Path("") self.folderOutputs = self.folderExecution / "Outputs" @@ -426,6 +418,7 @@ def __init__( if self.optimization_options is not None: self.optimization_options = IOtools.curate_mitim_nml( self.optimization_options, + self.optimization_object.folder, stopping_criteria_default = stopping_criteria_default ) # ------------------------------------------------------------------------------------------------- diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 9eb3995b..b8e1d792 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1244,9 +1244,7 @@ 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_ions):.3f}" - ) + print(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_ions * ni_orig From 934bd0d7430f812fa771e11c4a7a55813fb51799 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 30 Aug 2025 10:57:40 -0400 Subject: [PATCH 260/385] Transitioned optimization namelist to YAML --- docs/capabilities/optimization.rst | 4 +- .../maestro/scripts/run_maestro.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 38 ++----- src/mitim_tools/misc_tools/IOtools.py | 43 +++----- src/mitim_tools/opt_tools/STRATEGYtools.py | 35 +++++-- templates/main.namelist.json | 99 ------------------- templates/namelist.optimization.yaml | 89 +++++++++++++++++ ...ls.namelist.yaml => namelist.portals.yaml} | 0 tests/OPT_workflow.py | 4 +- tests/PORTALS_workflow.py | 2 +- 10 files changed, 141 insertions(+), 175 deletions(-) delete mode 100644 templates/main.namelist.json create mode 100644 templates/namelist.optimization.yaml rename templates/{portals.namelist.yaml => namelist.portals.yaml} (100%) diff --git a/docs/capabilities/optimization.rst b/docs/capabilities/optimization.rst index 7ee3c90b..f3c14478 100644 --- a/docs/capabilities/optimization.rst +++ b/docs/capabilities/optimization.rst @@ -44,7 +44,7 @@ Select the location of the MITIM namelist (see :ref:`Understanding the MITIM nam .. code-block:: python folder = Path('MITIM-fusion/tests/scratch/mitim_tut') - namelist = Path('MITIM-fusion/templates/main.namelist.json') + namelist = Path('MITIM-fusion/templates/namelist.optimization.yaml') Then create your custom optimization object as a child of the parent ``STRATEGYtools.opt_evaluator`` class. You only need to modify what operations need to occur inside the ``run()`` (where operations/simulations happen) and ``scalarized_objective()`` (to define what is the target to maximize) methods. @@ -118,7 +118,7 @@ Once finished, we can plot the results easily with: Understanding the MITIM namelist ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Checkout file ``MITIM-fusion/templates/main.namelist.json``, which has comprehensive comments. +Checkout file ``MITIM-fusion/templates/namelist.optimization.yaml``, which has comprehensive comments. *Under development* diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index 3282ce5e..a1c9e582 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -10,7 +10,7 @@ def parse_maestro_nml(file_path): # Extract engineering parameters, initializations, and desired beats to run - maestro_namelist = IOtools.read_mitim_nml(file_path) + maestro_namelist = IOtools.read_mitim_json(file_path) if "seed" in maestro_namelist: seed = maestro_namelist["seed"] diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 574b9e51..9b5134fa 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -43,12 +43,12 @@ def __init__( # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Read optimization namelist (always the default, the values to be modified are in the portals one) - namelist = __mitimroot__ / "templates" / "main.namelist.json" - self.optimization_options = IOtools.read_mitim_nml(namelist) + self.optimization_namelist = __mitimroot__ / "templates" / "namelist.optimization.yaml" + self.optimization_options = IOtools.read_mitim_yaml(self.optimization_namelist) # Read PORTALS namelist (if not provided, use default) if portals_namelist is None: - self.portals_namelist = __mitimroot__ / "templates" / "portals.namelist.yaml" + self.portals_namelist = __mitimroot__ / "templates" / "namelist.portals.yaml" print(f"\t- No PORTALS namelist provided, using default in {IOtools.clipstr(self.portals_namelist)}") else: self.portals_namelist = portals_namelist @@ -97,7 +97,12 @@ def prep( # Make sure that options that are required by good behavior of PORTALS # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - key_rhos = self.check_flags(askQuestions=askQuestions) + print(">> PORTALS flags pre-check") + + # Check that I haven't added a deprecated variable that I expect some behavior from + IOtools.check_flags_dictionary(self.portals_parameters, self.potential_flags, avoid = ["run", "read"], askQuestions=askQuestions) + + key_rhos = "predicted_roa" if self.portals_parameters["solution"]["predicted_roa"] is not None else "predicted_rho" # TO BE REMOVED IN FUTURE if not isinstance(cold_start, bool): @@ -170,7 +175,7 @@ def prep( shutil.copy(self.portals_namelist, self.folder / "portals.namelist_original.yaml") # Write the parameters (after script modification) to a yaml namelist for tracking purposes - IOtools.write_mitim_yaml(self.portals_parameters, self.folder / "portals.namelist.yaml") + IOtools.write_mitim_yaml(self.portals_parameters, self.folder / "namelist.portals.yaml") def _define_reuse_models(self): ''' @@ -283,29 +288,6 @@ def analyze_results(self, plotYN=True, fn=None, cold_start=False, analysis_level self, plotYN=plotYN, fn=fn, cold_start=cold_start, analysis_level=analysis_level ) - def check_flags(self, askQuestions=True): - - print(">> PORTALS flags pre-check") - - # Check that I haven't added a deprecated variable that I expect some behavior from - - def _check_flags_dictionary(d, d_check, avoid = ["run", "read"]): - for key in d.keys(): - if key not in d_check: - print(f"\t- {key} is an unexpected variable, prone to errors or misinterpretation",typeMsg="q" if askQuestions else "w") - elif not isinstance(d[key], dict): - continue - elif key in avoid: - continue - else: - _check_flags_dictionary(d[key], d_check[key]) - - _check_flags_dictionary(self.portals_parameters, self.potential_flags) - - key_rhos = "predicted_roa" if self.portals_parameters["solution"]["predicted_roa"] is not None else "predicted_rho" - - return key_rhos - def reuseTrainingTabular( self, folderRead, folderNew, reevaluate_targets=0, cold_startIfExists=False): """ diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index e6d1601e..471ed530 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -601,45 +601,24 @@ def calculate_size_pickle(file): # MITIM optimization namelist # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -def read_mitim_nml(json_file): +def read_mitim_json(json_file): jpath = Path(json_file).expanduser() with open(jpath, 'r') as file: optimization_options = json.load(file) return optimization_options -def curate_mitim_nml(optimization_options, folder, stopping_criteria_default = None): - - # Optimization criterion - if optimization_options['convergence_options']['stopping_criteria'] is None: - optimization_options['convergence_options']['stopping_criteria'] = stopping_criteria_default - - # Add optimization print - if optimization_options is not None: - unprint_fun = copy.deepcopy(optimization_options['convergence_options']['stopping_criteria']) - def opt_crit(*args,**kwargs): - print('\n') - print('--------------------------------------------------') - print('Convergence criteria') - print('--------------------------------------------------') - v = unprint_fun(*args, **kwargs) - print('--------------------------------------------------\n') - return v - optimization_options['convergence_options']['stopping_criteria'] = opt_crit - - # Check if the optimization options are in the namelist - from mitim_tools import __mitimroot__ - Optim_potential = read_mitim_nml(__mitimroot__ / "templates" / "main.namelist.json") - for ikey in optimization_options: - if ikey not in Optim_potential: - print(f"\t- Option {ikey} is an unexpected variable, prone to errors", typeMsg="q") - - # Write the optimization parameters stored in the object, into a file - if folder is not None: - write_mitim_yaml(optimization_options, folder / "optimization.namelist.yaml") - print(f" --> Optimization namelist written to {folder / 'optimization.namelist.yaml'}") +def check_flags_dictionary(d, d_check, avoid = [], askQuestions=True): + for key in d.keys(): + if key in avoid: + continue + elif key not in d_check: + print(f"\t- {key} is an unexpected variable, prone to errors or misinterpretation",typeMsg="q" if askQuestions else "w") + elif not isinstance(d[key], dict): + continue + else: + check_flags_dictionary(d[key], d_check[key], avoid=avoid, askQuestions=askQuestions) - return optimization_options # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 0ecff21c..de58dee3 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -99,13 +99,13 @@ def __init__( if namelist is not None: print(f"\t- Optimizaiton namelist provided: {namelist}", typeMsg="i") - self.optimization_options = IOtools.read_mitim_nml(namelist) + self.optimization_options = IOtools.read_mitim_yaml(namelist) elif default_namelist_function is not None: print("\t- Optimizaiton namelist not provided, using MITIM default for this optimization sub-module", typeMsg="i") - namelist = __mitimroot__ / "templates" / "main.namelist.json" - self.optimization_options = IOtools.read_mitim_nml(namelist) + namelist = __mitimroot__ / "templates" / "namelist.optimization.yaml" + self.optimization_options = IOtools.read_mitim_yaml(namelist) self.optimization_options = default_namelist_function(self.optimization_options) @@ -345,7 +345,7 @@ def __init__( """ Inputs: - optimization_object : Function that is executed, - with .optimization_options in it (Dictionary with optimization parameters (must be obtained using namelist and read_mitim_nml)) + with .optimization_options in it (Dictionary with optimization parameters (must be obtained using namelist and read_mitim_yaml)) and .folder (Where the function runs) and surrogate_parameters: Parameters to pass to surrogate (e.g. for transformed function), It can be different from function_parameters because of making evaluations fast. - cold_start : If False, try to find the values from Outputs/optimization_data.csv @@ -414,13 +414,22 @@ def __init__( self.surrogate_parameters = self.optimization_object.surrogate_parameters self.optimization_options = self.optimization_object.optimization_options - # Curate namelist --------------------------------------------------------------------------------- if self.optimization_options is not None: - self.optimization_options = IOtools.curate_mitim_nml( - self.optimization_options, - self.optimization_object.folder, - stopping_criteria_default = stopping_criteria_default + + # Check if the optimization options are in the namelist + optimization_options_default = IOtools.read_mitim_yaml(__mitimroot__ / "templates" / "namelist.optimization.yaml") + potential_flags = IOtools.deep_grab_flags_dict(optimization_options_default) + IOtools.check_flags_dictionary( + self.optimization_options, potential_flags, + avoid = ["stopping_criteria_parameters"], # Because they are specific to the stopping criteria + askQuestions=askQuestions ) + + # Write the optimization parameters stored in the object, into a file + if self.optimization_object.folder is not None: + IOtools.write_mitim_yaml(self.optimization_options, self.optimization_object.folder / "optimization.namelist.yaml") + print(f" --> Optimization namelist written to {self.optimization_object.folder / 'optimization.namelist.yaml'}") + # ------------------------------------------------------------------------------------------------- if not onlyInitialize: @@ -1829,6 +1838,12 @@ def max_val(maximum_value_orig, maximum_value_is_rel, res_base): def stopping_criteria_default(mitim_bo, parameters = {}): + + print('\n') + print('--------------------------------------------------') + print('Convergence criteria') + print('--------------------------------------------------') + # ------------------------------------------------------------------------------------ # Determine the stopping criteria # ------------------------------------------------------------------------------------ @@ -1858,7 +1873,7 @@ def stopping_criteria_default(mitim_bo, parameters = {}): yvals = None converged = converged_by_value or converged_by_dvs - + return converged, yvals def stopping_criteria_by_value(mitim_bo, maximum_value): diff --git a/templates/main.namelist.json b/templates/main.namelist.json deleted file mode 100644 index 6276b58a..00000000 --- a/templates/main.namelist.json +++ /dev/null @@ -1,99 +0,0 @@ -{ - "problem_options": { - "ofs": ["y0", "y1"], - "dvs": ["x0", "x1"], - "dvs_min": [0.8, 0.8], - "dvs_base": null, - "dvs_max": [1.2, 1.2] - }, - "evaluation_options": { - "parallel_evaluations": 1, - "train_Ystd": null - }, - "convergence_options": { - "maximum_iterations": 5, - "stopping_criteria": null, - "stopping_criteria_parameters": { - "maximum_value": -1e-3, - "maximum_value_is_rel": false, - "minimum_inputs_variation": [10, 3, 0.01] - } - }, - "initialization_options": { - "initial_training": 5, - "type_initialization": 3, - "read_initial_training_from_csv": false, - "initialization_fun": null, - "ensure_within_bounds": false, - "expand_bounds": true - }, - "acquisition_options": { - "type": "noisy_logei_mc", - "parameters": { - "mc_samples": 1024 - }, - "optimizers": ["botorch"], - "optimizer_options": { - "botorch": { - "num_restarts": 64, - "raw_samples": 4096, - "maxiter": 1000, - "sequential_q": true, - "keep_best": 1 - }, - "root": { - "num_restarts": 5, - "solver": "lm", - "maxiter": 1000, - "relative_improvement_for_stopping": 1e-4, - "keep_best": 1 - }, - "sr": { - "num_restarts": 5, - "maxiter": 1000, - "relative_improvement_for_stopping": 1e-3, - "relax": 0.1, - "relax_dyn": true, - "keep_best": 1 - }, - "ga": { - "num_restarts": 1, - "keep_best": 32 - } - }, - "relative_improvement_for_stopping": null, - "favor_proximity_type": 0, - "ensure_new_points": true, - "points_per_step": 1 - }, - "surrogate_options": { - "TypeKernel": 0, - "TypeMean": 0, - "selectSurrogate": null, - "FixedNoise": true, - "ExtraNoise": false, - "additional_constraints": null, - "ConstrainNoise": -1e-3, - "MinimumRelativeNoise": null, - "stds_outside": null, - "stds_outside_checker": 5, - "extrapointsFile": null, - "extrapointsModels": null, - "extrapointsModelsAvoidContent": null - }, - "strategy_options": { - "AllowedExcursions": [0.0, 0.0], - "HitBoundsIncrease": [1.0, 1.0], - "boundsRefine": null, - "RandomRangeBounds": 0.5, - "ToleranceNiche": 1e-3, - "applyCorrections": true, - "SwitchIterationsReduction": [null, null], - "TURBO_options": { - "apply": false, - "points": 32, - "bounds": [0.75, 1.33], - "metrics": [3, 3] - } - } -} \ No newline at end of file diff --git a/templates/namelist.optimization.yaml b/templates/namelist.optimization.yaml new file mode 100644 index 00000000..45173fe5 --- /dev/null +++ b/templates/namelist.optimization.yaml @@ -0,0 +1,89 @@ +problem_options: + ofs: ["y0", "y1"] + dvs: ["x0", "x1"] + dvs_min: [0.8, 0.8] + dvs_base: null + dvs_max: [1.2, 1.2] + +evaluation_options: + parallel_evaluations: 1 + train_Ystd: null + +# Convergence options for the BO optimization process +convergence_options: + maximum_iterations: 5 + stopping_criteria: "import::mitim_tools.opt_tools.STRATEGYtools.stopping_criteria_default" + stopping_criteria_parameters: + maximum_value: -0.001 + maximum_value_is_rel: false + minimum_inputs_variation: [10, 3, 0.01] + +initialization_options: + initial_training: 5 + type_initialization: 3 + read_initial_training_from_csv: false + initialization_fun: null + ensure_within_bounds: false + expand_bounds: true + +acquisition_options: + type: noisy_logei_mc + parameters: + mc_samples: 1024 + optimizers: ["botorch"] + optimizer_options: + botorch: + num_restarts: 64 + raw_samples: 4096 + maxiter: 1000 + sequential_q: true + keep_best: 1 + root: + num_restarts: 5 + solver: lm + maxiter: 1000 + relative_improvement_for_stopping: 0.0001 + keep_best: 1 + sr: + num_restarts: 5 + maxiter: 1000 + relative_improvement_for_stopping: 0.001 + relax: 0.1 + relax_dyn: true + keep_best: 1 + ga: + num_restarts: 1 + keep_best: 32 + relative_improvement_for_stopping: null + favor_proximity_type: 0 + ensure_new_points: true + points_per_step: 1 + +surrogate_options: + TypeKernel: 0 + TypeMean: 0 + selectSurrogate: null + FixedNoise: true + ExtraNoise: false + additional_constraints: null + ConstrainNoise: -0.001 + MinimumRelativeNoise: null + stds_outside: null + stds_outside_checker: 5 + extrapointsFile: null + extrapointsModels: null + extrapointsModelsAvoidContent: null + +strategy_options: + AllowedExcursions: [0.0, 0.0] + HitBoundsIncrease: [1.0, 1.0] + boundsRefine: null + RandomRangeBounds: 0.5 + ToleranceNiche: 0.001 + applyCorrections: true + SwitchIterationsReduction: [null, null] + TURBO_options: + apply: false + points: 32 + bounds: [0.75, 1.33] + metrics: [3, 3] \ No newline at end of file diff --git a/templates/portals.namelist.yaml b/templates/namelist.portals.yaml similarity index 100% rename from templates/portals.namelist.yaml rename to templates/namelist.portals.yaml diff --git a/tests/OPT_workflow.py b/tests/OPT_workflow.py index d05bcf53..ab77e64d 100644 --- a/tests/OPT_workflow.py +++ b/tests/OPT_workflow.py @@ -62,7 +62,7 @@ def scalarized_objective(self, Y): # ----- Inputs # ----------------------------------------------------------------------------------------------------- -namelist = __mitimroot__ / "templates" / "main.namelist.json" +namelist = __mitimroot__ / "templates" / "namelist.optimization.yamln" folderWork = __mitimroot__ / "tests" / "scratch" / "opt_test" if cold_start and os.path.exists(folderWork): @@ -75,7 +75,7 @@ def scalarized_objective(self, Y): # Initialize class opt_fun1D = opt_class(folderWork, namelist) -# Changes to namelist in templates/main.namelist.json +# Changes to namelist in templates/namelist.optimization.yaml opt_fun1D.optimization_options["initialization_options"]["initial_training"] = 2 # Initialize BO framework diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index 1a8c20a5..760938d9 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -24,7 +24,7 @@ # Optimization Class # --------------------------------------------------------------------------------------------------------------------- -# Initialize class with the default namelist in templates/portals.namelist.yaml but modify some of its parameters +# Initialize class with the default namelist in templates/namelist.portals.yaml but modify some of its parameters portals_fun = PORTALSmain.portals(folderWork) portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 1 portals_fun.optimization_options["initialization_options"]["initial_training"] = 2 From 129d93d0cd4eb23de88654ad2ef1f246d533bdd0 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 30 Aug 2025 11:09:11 -0400 Subject: [PATCH 261/385] Transitioned to yaml for MAESTRO --- .../maestro/scripts/run_maestro.py | 7 +- src/mitim_modules/portals/PORTALSmain.py | 2 +- src/mitim_tools/misc_tools/IOtools.py | 18 +-- src/mitim_tools/opt_tools/STRATEGYtools.py | 2 +- src/mitim_tools/popcon_tools/POPCONtools.py | 2 +- templates/maestro_namelist.json | 126 ----------------- templates/namelist.maestro.yaml | 128 ++++++++++++++++++ ..._popcon.yaml => namelist.plot_popcon.yaml} | 0 tests/MAESTRO_workflow.py | 4 +- 9 files changed, 138 insertions(+), 151 deletions(-) delete mode 100644 templates/maestro_namelist.json create mode 100644 templates/namelist.maestro.yaml rename templates/{plot_popcon.yaml => namelist.plot_popcon.yaml} (100%) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index a1c9e582..be373266 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -1,5 +1,4 @@ import argparse -import shutil from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.maestro.MAESTROmain import maestro @@ -10,7 +9,7 @@ def parse_maestro_nml(file_path): # Extract engineering parameters, initializations, and desired beats to run - maestro_namelist = IOtools.read_mitim_json(file_path) + maestro_namelist = IOtools.read_mitim_yaml(file_path) if "seed" in maestro_namelist: seed = maestro_namelist["seed"] @@ -261,8 +260,8 @@ def main(): if not folder.exists(): folder.mkdir(parents=True, exist_ok=True) - if (folder / 'maestro_namelist.json').exists(): - IOtools.recursive_backup(folder / 'maestro_namelist.json') + if (folder / 'namelist.maestro.yaml').exists(): + IOtools.recursive_backup(folder / 'namelist.maestro.yaml') run_maestro_local(*parse_maestro_nml(file_path),folder=folder,cpus = cpus, terminal_outputs = terminal_outputs) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 9b5134fa..638aeefe 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -100,7 +100,7 @@ def prep( print(">> PORTALS flags pre-check") # Check that I haven't added a deprecated variable that I expect some behavior from - IOtools.check_flags_dictionary(self.portals_parameters, self.potential_flags, avoid = ["run", "read"], askQuestions=askQuestions) + IOtools.check_flags_mitim_namelist(self.portals_parameters, self.potential_flags, avoid = ["run", "read"], askQuestions=askQuestions) key_rhos = "predicted_roa" if self.portals_parameters["solution"]["predicted_roa"] is not None else "predicted_rho" diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index 471ed530..4e340cc2 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -597,18 +597,7 @@ def calculate_size_pickle(file): obj = pickle.load(f) calculate_sizes_obj_recursive(obj, recursion = 20) -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# MITIM optimization namelist -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -def read_mitim_json(json_file): - jpath = Path(json_file).expanduser() - with open(jpath, 'r') as file: - optimization_options = json.load(file) - - return optimization_options - -def check_flags_dictionary(d, d_check, avoid = [], askQuestions=True): +def check_flags_mitim_namelist(d, d_check, avoid = [], askQuestions=True): for key in d.keys(): if key in avoid: continue @@ -617,10 +606,7 @@ def check_flags_dictionary(d, d_check, avoid = [], askQuestions=True): elif not isinstance(d[key], dict): continue else: - check_flags_dictionary(d[key], d_check[key], avoid=avoid, askQuestions=askQuestions) - - -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + check_flags_mitim_namelist(d[key], d_check[key], avoid=avoid, askQuestions=askQuestions) def getpythonversion(): return [ int(i.split("\n")[0].split("+")[0]) for i in sys.version.split()[0].split(".") ] diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index de58dee3..11e708e8 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -419,7 +419,7 @@ def __init__( # Check if the optimization options are in the namelist optimization_options_default = IOtools.read_mitim_yaml(__mitimroot__ / "templates" / "namelist.optimization.yaml") potential_flags = IOtools.deep_grab_flags_dict(optimization_options_default) - IOtools.check_flags_dictionary( + IOtools.check_flags_mitim_namelist( self.optimization_options, potential_flags, avoid = ["stopping_criteria_parameters"], # Because they are specific to the stopping criteria askQuestions=askQuestions diff --git a/src/mitim_tools/popcon_tools/POPCONtools.py b/src/mitim_tools/popcon_tools/POPCONtools.py index 44f7d95c..e4226f94 100644 --- a/src/mitim_tools/popcon_tools/POPCONtools.py +++ b/src/mitim_tools/popcon_tools/POPCONtools.py @@ -377,7 +377,7 @@ def plot(self, # Read template plot options if plot_template is None: - plot_template = __mitimroot__ / "templates" / "plot_popcon.yaml" + plot_template = __mitimroot__ / "templates" / "namelist.plot_popcon.yaml" plot_style = cfspopcon.read_plot_style(plot_template) # Update plot options diff --git a/templates/maestro_namelist.json b/templates/maestro_namelist.json deleted file mode 100644 index b0d712dc..00000000 --- a/templates/maestro_namelist.json +++ /dev/null @@ -1,126 +0,0 @@ -{ - "flag": "SPARC PRD", - "seed": 0, - "machine": { - "Bt": 12.2, - "Ip": 8.7, - "separatrix": { - "type": "freegs", - "parameters": { - "R": 1.85, - "a": 0.57, - "delta_sep": 0.57, - "kappa_sep": 1.97, - "n_mxh": 5, - "geqdsk_file": "" - } - }, - "heating": { - "type": "ICRH", - "parameters": { - "P_icrh": 11.0, - "minority": [2,3], - "fmini": 0.05 - } - } - }, - "assumptions": { - "Zeff": 1.5, - "mix":{ - "fmain": 0.85, - "ZW":50, - "fW": 1.5E-5 - }, - "initialization": { - "BetaN":1.0, - "density_peaking":1.3, - "assume_neped": true, - "neped_20": 2.5, - "nesep_ratio": 0.3 - }, - "Tesep_eV": 75.0 - }, - "maestro": { - "keep_all_files": true, - "beats": ["transp_soft", "transp", "eped", "portals", "eped", "portals"], - "eped_beat": { - "use_default": false, - "eped_namelist":{ - "nn_location": "$MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-MODEL-SPARC.keras", - "norm_location": "$MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-NORMALIZATION-SPARC.txt", - "corrections_set": { - "Bt": 12.2, - "R": 1.85, - "a": 0.57 - }, - "ptop_multiplier": 1.0, - "TioverTe": 1.0 - } - }, - "eped_initializer_beat": { - "use_default": false, - "eped_initializer_namelist":{ - "nn_location": "$MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-MODEL-SPARC.keras", - "norm_location": "$MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-NORMALIZATION-SPARC.txt", - "corrections_set": { - "Bt": 12.2, - "R": 1.85, - "a": 0.57 - }, - "ptop_multiplier": 1.0, - "TioverTe": 1.0 - } - }, - "portals_beat": { - "use_default": false, - "portals_namelist" : { - "portals_parameters": { - "solution": { - "predicted_roa": [0.35,0.45,0.55,0.65,0.75,0.875,0.9], - "keep_full_model_folder": false - }, - "transport": { - "options": { - "tglf": { - "run": { - "code_settings": "SAT2astra", - "extraOptions": {"USE_BPER": true} - } - } - } - }, - "target": { - "options": { - "force_zero_particle_flux": true - } - } - }, - "initialization_parameters": { - "thermalize_fast": true, - "quasineutrality": true - }, - "exploration_ranges": { - "ymax_rel": 1.0, - "ymin_rel": 1.0, - "yminymax_atleast": [null, 4] - }, - "change_last_radial_call" : true, - "use_previous_surrogate_data" : true, - "try_flux_match_only_for_first_point" : true - }, - "transport_preprocessing": { - "lumpImpurities": true, - "enforce_same_density_gradients": true - } - }, - "portals_soft_beat":{ - "use_default": true - }, - "transp_beat":{ - "use_default": true - }, - "transp_soft_beat":{ - "use_default": true - } - } -} diff --git a/templates/namelist.maestro.yaml b/templates/namelist.maestro.yaml new file mode 100644 index 00000000..2be88b90 --- /dev/null +++ b/templates/namelist.maestro.yaml @@ -0,0 +1,128 @@ +# Name of this simulation +flag: SPARC PRD + +# Random seed for reproducibility (master seed to be sent to all beats that receive seed) +seed: 0 + +# Machine parameters +machine: + + Bt: 12.2 + Ip: 8.7 + + # Separatrix specification + separatrix: + type: freegs + parameters: + R: 1.85 + a: 0.57 + delta_sep: 0.57 + kappa_sep: 1.97 + n_mxh: 5 + geqdsk_file: "" + + # Heating parameters + heating: + type: ICRH + parameters: + P_icrh: 11.0 + minority: [2, 3] + fmini: 0.05 + +# Simulation assumptions +assumptions: + + # Impurities + Zeff: 1.5 + mix: + fmain: 0.85 + ZW: 50 + fW: 1.5e-5 + + # Initialization of profiles + initialization: + BetaN: 1.0 + density_peaking: 1.3 + assume_neped: true + neped_20: 2.5 + nesep_ratio: 0.3 + + # Edge parameters + Tesep_eV: 75.0 + +# MAESTRO workflow parameters +maestro: + + # Remove intermediate files + keep_all_files: true + + # Sequence of beats + beats: ["transp_soft", "transp", "eped", "portals", "eped", "portals"] + + # --------------------------------------------------------------------------- + # Each individual beat parameters + # --------------------------------------------------------------------------- + + eped_beat: + use_default: false + eped_namelist: + nn_location: $MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-MODEL-SPARC.keras + norm_location: $MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-NORMALIZATION-SPARC.txt + corrections_set: + Bt: 12.2 + R: 1.85 + a: 0.57 + ptop_multiplier: 1.0 + TioverTe: 1.0 + + eped_initializer_beat: + use_default: false + eped_initializer_namelist: + nn_location: $MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-MODEL-SPARC.keras + norm_location: $MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-NORMALIZATION-SPARC.txt + corrections_set: + Bt: 12.2 + R: 1.85 + a: 0.57 + ptop_multiplier: 1.0 + TioverTe: 1.0 + + portals_beat: + use_default: false + portals_namelist: + portals_parameters: + solution: + predicted_roa: [0.35, 0.45, 0.55, 0.65, 0.75, 0.875, 0.9] + keep_full_model_folder: false + transport: + options: + tglf: + run: + code_settings: SAT2astra + extraOptions: + USE_BPER: true + target: + options: + force_zero_particle_flux: true + initialization_parameters: + thermalize_fast: true + quasineutrality: true + exploration_ranges: + ymax_rel: 1.0 + ymin_rel: 1.0 + yminymax_atleast: [null, 4] + change_last_radial_call: true + use_previous_surrogate_data: true + try_flux_match_only_for_first_point: true + transport_preprocessing: + lumpImpurities: true + enforce_same_density_gradients: true + + portals_soft_beat: + use_default: true + + transp_beat: + use_default: true + + transp_soft_beat: + use_default: true \ No newline at end of file diff --git a/templates/plot_popcon.yaml b/templates/namelist.plot_popcon.yaml similarity index 100% rename from templates/plot_popcon.yaml rename to templates/namelist.plot_popcon.yaml diff --git a/tests/MAESTRO_workflow.py b/tests/MAESTRO_workflow.py index ccb6eb01..ec8857c6 100644 --- a/tests/MAESTRO_workflow.py +++ b/tests/MAESTRO_workflow.py @@ -3,10 +3,10 @@ from mitim_tools import __mitimroot__ from mitim_modules.maestro.scripts import run_maestro -cold_start = False +cold_start = True folder = __mitimroot__ / "tests" / "scratch" / "maestro_test" -template = __mitimroot__ / "templates" / "maestro_namelist.json" +template = __mitimroot__ / "templates" / "namelist.maestro.yaml" if cold_start and os.path.exists(folder): os.system(f"rm -r {folder}") From 5d6f88e53a24065367d47c2168cd6b0a3b735346 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 30 Aug 2025 11:22:27 -0400 Subject: [PATCH 262/385] Transitioned to yaml for all simulation models --- docs/capabilities/tglf_capabilities.rst | 2 +- .../gacode_tools/utils/GACODEdefaults.py | 21 ++--- templates/input.cgyro.models.json | 14 ---- templates/input.cgyro.models.yaml | 9 +++ templates/input.gx.models.json | 19 ----- templates/input.gx.models.yaml | 15 ++++ templates/input.neo.models.json | 8 -- templates/input.neo.models.yaml | 4 + templates/input.tglf.models.json | 78 ------------------- templates/input.tglf.models.yaml | 67 ++++++++++++++++ tutorials/PORTALS_tutorial.py | 2 +- 11 files changed, 108 insertions(+), 131 deletions(-) delete mode 100644 templates/input.cgyro.models.json create mode 100644 templates/input.cgyro.models.yaml delete mode 100644 templates/input.gx.models.json create mode 100644 templates/input.gx.models.yaml delete mode 100644 templates/input.neo.models.json create mode 100644 templates/input.neo.models.yaml delete mode 100644 templates/input.tglf.models.json create mode 100644 templates/input.tglf.models.yaml diff --git a/docs/capabilities/tglf_capabilities.rst b/docs/capabilities/tglf_capabilities.rst index e411a833..41862a74 100644 --- a/docs/capabilities/tglf_capabilities.rst +++ b/docs/capabilities/tglf_capabilities.rst @@ -77,7 +77,7 @@ 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 ``cold_start`` 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 ``code_settings`` 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. +The user is referred to ``templates/input.tglf.models.yaml`` 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): diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index 84947df8..b8b2eb43 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -41,48 +41,49 @@ def addTGLFcontrol(code_settings, NS=2, minimal=False): ******************************************************************************** """ - options = add_code_settings(options, code_settings, models_file="input.tglf.models.json") + options = add_code_settings(options, code_settings, models_file="input.tglf.models.yaml") return options def addNEOcontrol(code_settings,*args, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.neo.controls", caseInsensitive=False) - options = add_code_settings(options, code_settings, models_file="input.neo.models.json") + options = add_code_settings(options, code_settings, models_file="input.neo.models.yaml") return options def addGXcontrol(code_settings,*args, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.gx.controls", caseInsensitive=False) - options = add_code_settings(options, code_settings, models_file="input.gx.models.json") + options = add_code_settings(options, code_settings, models_file="input.gx.models.yaml") return options def addCGYROcontrol(code_settings, rmin=None, **kwargs): options = IOtools.generateMITIMNamelist(__mitimroot__ / "templates" / "input.cgyro.controls", caseInsensitive=False) - options = add_code_settings(options, code_settings, models_file="input.cgyro.models.json") + options = add_code_settings(options, code_settings, models_file="input.cgyro.models.yaml") return options -def add_code_settings(options,code_settings, models_file = "input.tglf.models.json"): +def add_code_settings(options,code_settings, models_file = "input.tglf.models.yaml"): - with open(__mitimroot__ / "templates" / models_file, "r") as f: - settings = json.load(f) + settings = IOtools.read_mitim_yaml(__mitimroot__ / "templates" / models_file) found = False - # Search by number first + code_settings = str(code_settings) + + # Search by label first if str(code_settings) in settings: sett = settings[str(code_settings)] for ikey in sett["controls"]: options[ikey] = sett["controls"][ikey] found = True else: - # Search by label second + # Search by deprecated descriptor for ikey in settings: - if settings[ikey]["label"] == code_settings: + if settings[ikey]["deprecated_descriptor"] == code_settings: sett = settings[ikey] for jkey in sett["controls"]: options[jkey] = sett["controls"][jkey] diff --git a/templates/input.cgyro.models.json b/templates/input.cgyro.models.json deleted file mode 100644 index 421b0ebb..00000000 --- a/templates/input.cgyro.models.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "0": { - "label": "Linear", - "controls": { - "NONLINEAR_FLAG": 0 - } - }, - "1": { - "label": "Nonlinear", - "controls": { - "NONLINEAR_FLAG": 1 - } - } -} \ No newline at end of file diff --git a/templates/input.cgyro.models.yaml b/templates/input.cgyro.models.yaml new file mode 100644 index 00000000..a4d94070 --- /dev/null +++ b/templates/input.cgyro.models.yaml @@ -0,0 +1,9 @@ +"Linear": + deprecated_descriptor: "0" + controls: + NONLINEAR_FLAG: 0 + +"Nonlinear": + deprecated_descriptor: "1" + controls: + NONLINEAR_FLAG: 1 \ No newline at end of file diff --git a/templates/input.gx.models.json b/templates/input.gx.models.json deleted file mode 100644 index 8bd4c885..00000000 --- a/templates/input.gx.models.json +++ /dev/null @@ -1,19 +0,0 @@ -{ - "0": { - "label": "Linear with 10 ky modes", - "controls": { - "nonlinear_mode": false, - "hyper": false, - "HB_hyper": false, - "nperiod": 2, - "nx": 1, - "ny": 28 - } - }, - "1": { - "label": "Nonlinear", - "controls": { - "nonlinear_mode": true - } - } -} \ No newline at end of file diff --git a/templates/input.gx.models.yaml b/templates/input.gx.models.yaml new file mode 100644 index 00000000..ac60ec6c --- /dev/null +++ b/templates/input.gx.models.yaml @@ -0,0 +1,15 @@ +"Linear": + # Linear with 10 ky modes + deprecated_descriptor: "0" + controls: + nonlinear_mode: false + hyper: false + HB_hyper: false + nperiod: 2 + nx: 1 + ny: 28 + +"Nonlinear": + deprecated_descriptor: "1" + controls: + nonlinear_mode: true \ No newline at end of file diff --git a/templates/input.neo.models.json b/templates/input.neo.models.json deleted file mode 100644 index 1e86edac..00000000 --- a/templates/input.neo.models.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "0": { - "label": "Sonic", - "controls": { - "ROTATION_MODEL": 2 - } - } -} \ No newline at end of file diff --git a/templates/input.neo.models.yaml b/templates/input.neo.models.yaml new file mode 100644 index 00000000..6ed703b6 --- /dev/null +++ b/templates/input.neo.models.yaml @@ -0,0 +1,4 @@ +"Sonic": + deprecated_descriptor: "0" + controls: + ROTATION_MODEL: 2 \ No newline at end of file diff --git a/templates/input.tglf.models.json b/templates/input.tglf.models.json deleted file mode 100644 index 1204e30a..00000000 --- a/templates/input.tglf.models.json +++ /dev/null @@ -1,78 +0,0 @@ -{ - "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 - } - }, - "100": { - "label": "SAT2astra", - "controls": { - "UNITS": "CGYRO", - "USE_BPER": true, - "USE_AVE_ION_GRID": true, - "SAT_RULE": 2, - "KYGRID_MODEL": 4, - "XNU_MODEL": 3, - "NBASIS_MAX": 6, - "B_MODEL_SA": 1, - "FT_MODEL_SA": 1 - } - }, - "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/templates/input.tglf.models.yaml b/templates/input.tglf.models.yaml new file mode 100644 index 00000000..94638662 --- /dev/null +++ b/templates/input.tglf.models.yaml @@ -0,0 +1,67 @@ +"SAT1": + deprecated_descriptor: "1" + controls: + SAT_RULE: 1 + UNITS: GYRO + +"SAT0": + deprecated_descriptor: "2" + controls: + SAT_RULE: 0 + UNITS: GYRO + ETG_FACTOR: 1.25 + +"SAT1geo": + deprecated_descriptor: "3" + controls: + SAT_RULE: 1 + UNITS: CGYRO + +"SAT2": + deprecated_descriptor: "4" + controls: + SAT_RULE: 2 + UNITS: CGYRO + XNU_MODEL: 3 + WDIA_TRAPPED: 1.0 + +"SAT2em": + deprecated_descriptor: "5" + controls: + SAT_RULE: 2 + UNITS: CGYRO + XNU_MODEL: 3 + WDIA_TRAPPED: 1.0 + USE_BPER: true + +"SAT3": + deprecated_descriptor: "6" + controls: + SAT_RULE: 3 + UNITS: CGYRO + XNU_MODEL: 3 + WDIA_TRAPPED: 1.0 + +"SAT2astra": + deprecated_descriptor: "100" + controls: + UNITS: CGYRO + USE_BPER: true + USE_AVE_ION_GRID: true + SAT_RULE: 2 + KYGRID_MODEL: 4 + XNU_MODEL: 3 + NBASIS_MAX: 6 + B_MODEL_SA: 1 + FT_MODEL_SA: 1 + +"Experimentation] SAT3em basis": + deprecated_descriptor: "101" + 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/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index 732c13e1..77a74314 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -25,7 +25,7 @@ # TGLF specifications portals_fun.portals_parameters["transport"]["options"] = { - "code_settings": 6, # Check out templates/input.tglf.models.json for more options + "code_settings": 6, # Check out templates/input.tglf.models.yaml for more options "extraOptionsTGLF": {"USE_BPER": False} # Turn off BPER } From 2f70133dfbd5b6bfbe0a453456e0e701377429c9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 30 Aug 2025 11:47:08 -0400 Subject: [PATCH 263/385] Added capability of a simple mitim_run_portals --- pyproject.toml | 3 +- src/mitim_modules/portals/PORTALSmain.py | 22 +++++----- .../portals/scripts/run_portals.py | 40 +++++++++++++++++++ templates/namelist.portals.yaml | 6 +++ 4 files changed, 59 insertions(+), 12 deletions(-) create mode 100644 src/mitim_modules/portals/scripts/run_portals.py diff --git a/pyproject.toml b/pyproject.toml index 87c521ed..ddfced2b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,7 @@ omfit = [ "omfit_classes>3.2024.19.2", # Otherwise, it will need an old version of matplotlib, matplotlib<3.6 "scipy<1.14.0", # As of 08/08/2024, because of https://github.com/gafusion/OMFIT-source/issues/7104 "numpy<2.0.0", # For the xarray requirement below to work - "xarray==2022.6.0", # As of 08/08/2024, because https://github.com/gafusion/OMFIT-source/issues/7104. Froze this version because of the PLASMAstate xr reader (importlib_metadata issues) + #"xarray==2022.6.0", # As of 08/08/2024, because https://github.com/gafusion/OMFIT-source/issues/7104. Froze this version because of the PLASMAstate xr reader (importlib_metadata issues) "omas", "fortranformat", "openpyxl", @@ -89,6 +89,7 @@ mitim_run_tglf = "mitim_tools.gacode_tools.scripts.run_tglf:main" # (fold mitim_plot_opt = "mitim_tools.opt_tools.scripts.read:main" # Not transferred: --type 4 --resolution 20 mitim_plot_portals = "mitim_modules.portals.scripts.read_portals:main" mitim_slurm = "mitim_tools.opt_tools.scripts.slurm:main" +mitim_run_portals = "mitim_modules.portals.scripts.run_portals:main" # folder --input input.gacode --namelist namelist.portals.yaml --cold_start # TRANSP mitim_trcheck = "mitim_tools.transp_tools.scripts.run_check:main" # e.g. mitim_trcheck pablorf diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 638aeefe..832cc3c0 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -34,17 +34,10 @@ def __init__( print("\t\t\t PORTALS class module") print("-----------------------------------------------------------------------------------------\n") - # Store folder, namelist. Read namelist - - super().__init__(folder,tensor_options=tensor_options) - - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Default (please change to your desire after instancing the object) - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - # Read optimization namelist (always the default, the values to be modified are in the portals one) - self.optimization_namelist = __mitimroot__ / "templates" / "namelist.optimization.yaml" - self.optimization_options = IOtools.read_mitim_yaml(self.optimization_namelist) + super().__init__( + folder, + tensor_options=tensor_options + ) # Read PORTALS namelist (if not provided, use default) if portals_namelist is None: @@ -55,6 +48,13 @@ def __init__( print(f"\t- Using provided PORTALS namelist in {IOtools.clipstr(self.portals_namelist)}") self.portals_parameters = IOtools.read_mitim_yaml(self.portals_namelist) + # Read optimization namelist (always the default, the values to be modified are in the portals one) + if self.portals_parameters["optimization_namelist"] is not None: + self.optimization_namelist = self.portals_parameters["optimization_namelist"] + else: + self.optimization_namelist = __mitimroot__ / "templates" / "namelist.optimization.yaml" + self.optimization_options = IOtools.read_mitim_yaml(self.optimization_namelist) + # Apply the optimization options to the proper namelist and drop it from portals_parameters if 'optimization_options' in self.portals_parameters: self.optimization_options = IOtools.deep_dict_update(self.optimization_options, self.portals_parameters['optimization_options']) diff --git a/src/mitim_modules/portals/scripts/run_portals.py b/src/mitim_modules/portals/scripts/run_portals.py new file mode 100644 index 00000000..0b24b498 --- /dev/null +++ b/src/mitim_modules/portals/scripts/run_portals.py @@ -0,0 +1,40 @@ +from pathlib import Path +import argparse +from mitim_tools.opt_tools import STRATEGYtools +from mitim_modules.portals import PORTALSmain +from mitim_tools.misc_tools import IOtools + +def main(): + + parser = argparse.ArgumentParser() + + parser.add_argument("folder", type=str, help="Simulation folder") + parser.add_argument("--namelist", type=str, required=False, default=None) # namelist.portals.yaml file, otherwise what's in the current folder + parser.add_argument("--input", type=str, required=False, default=None) # input.gacode file, otherwise what's in the current folder + parser.add_argument('--cold_start', required=False, default=False, action='store_true') + + args = parser.parse_args() + + folderWork = Path(args.folder) + portals_namelist = args.namelist + inputgacode = args.input + cold_start = args.cold_start + + if portals_namelist is None: + portals_namelist = IOtools.expandPath('.') / "namelist.portals.yaml" + else: + portals_namelist = Path(portals_namelist) + + if inputgacode is None: + inputgacode = IOtools.expandPath('.') / "input.gacode" + else: + inputgacode = Path(inputgacode) + + portals_fun = PORTALSmain.portals(folderWork, portals_namelist=portals_namelist) + portals_fun.prep(inputgacode) + + mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=cold_start, askQuestions=False) + mitim_bo.run() + +if __name__ == "__main__": + main() diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index 229468af..86cb165f 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -196,6 +196,12 @@ target: # (%) Error (std, in percent) of model evaluation percent_error: 1 +# ----------------------------------------------------------------- +# Optimization options namelist +# ----------------------------------------------------------------- + +optimization_namelist: null # If null, it will grab the default at: __mitimroot__ / "templates" / "namelist.optimization.yaml" + # ----------------------------------------------------------------- # Optimization options (to change the main optimization namelist) # ----------------------------------------------------------------- From 1c7a5495179d06e335bd98056f0022f9c85e1e38 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 30 Aug 2025 12:23:17 -0400 Subject: [PATCH 264/385] Transitioned to exploration ranges in YAML, not to prep() --- .../maestro/utils/PORTALSbeat.py | 10 ++----- src/mitim_modules/portals/PORTALSmain.py | 19 +++++++------ .../portals/scripts/run_portals.py | 4 +-- templates/namelist.maestro.yaml | 8 +++--- templates/namelist.portals.yaml | 28 +++++++++++++++++++ 5 files changed, 46 insertions(+), 23 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 7bc79d84..a94c602f 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -25,11 +25,6 @@ def prepare(self, use_previous_surrogate_data = False, try_flux_match_only_for_first_point = True, change_last_radial_call = False, - exploration_ranges = { - 'ymax_rel': 1.0, - 'ymin_rel': 1.0, - 'yminymax_atleast': [0,2] - }, portals_parameters = {}, initialization_parameters = {}, optimization_options = {}, @@ -67,7 +62,6 @@ def prepare(self, self.optimization_options = optimization_options self.initialization_parameters = initialization_parameters - self.exploration_ranges = exploration_ranges self.use_previous_surrogate_data = use_previous_surrogate_data self.change_last_radial_call = change_last_radial_call @@ -92,7 +86,7 @@ def run(self, **kwargs): p = gacode_state(self.fileGACODE) p.correct(options=self.initialization_parameters) - portals_fun.prep(p,askQuestions=False,**self.exploration_ranges) + portals_fun.prep(p,askQuestions=False) self.mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, seed = self.maestro_instance.master_seed, cold_start = cold_start, askQuestions = False) @@ -105,7 +99,7 @@ def run(self, **kwargs): if len(self.mitim_bo.optimization_data.data) == 0: self._flux_match_for_first_point() - portals_fun.prep(self.fileGACODE,askQuestions=False,**self.exploration_ranges) + portals_fun.prep(self.fileGACODE,askQuestions=False) self.mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, seed=self.maestro_instance.master_seed,cold_start = cold_start, askQuestions = False) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 832cc3c0..04ed9c3f 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -67,15 +67,6 @@ def prep( self, mitim_state, cold_start=False, - ymax_rel=1.0, - ymin_rel=1.0, - limits_are_relative=True, - fixed_gradients=None, - yminymax_atleast=None, - enforce_finite_aLT=None, - define_ranges_from_profiles=None, - start_from_folder=None, - reevaluate_targets=0, seedInitial=None, askQuestions=True, ): @@ -93,6 +84,16 @@ def prep( - seedInitial can be optionally give a seed to randomize the starting profile (useful for developing, paper writing) """ + ymax_rel = self.portals_parameters["solution"]["exploration_ranges"]["ymax_rel"] + ymin_rel = self.portals_parameters["solution"]["exploration_ranges"]["ymin_rel"] + limits_are_relative = self.portals_parameters["solution"]["exploration_ranges"]["limits_are_relative"] + fixed_gradients = self.portals_parameters["solution"]["exploration_ranges"]["fixed_gradients"] + yminymax_atleast = self.portals_parameters["solution"]["exploration_ranges"]["yminymax_atleast"] + enforce_finite_aLT = self.portals_parameters["solution"]["exploration_ranges"]["enforce_finite_aLT"] + define_ranges_from_profiles = self.portals_parameters["solution"]["exploration_ranges"]["define_ranges_from_profiles"] + start_from_folder = self.portals_parameters["solution"]["exploration_ranges"]["start_from_folder"] + reevaluate_targets = self.portals_parameters["solution"]["exploration_ranges"]["reevaluate_targets"] + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Make sure that options that are required by good behavior of PORTALS # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/mitim_modules/portals/scripts/run_portals.py b/src/mitim_modules/portals/scripts/run_portals.py index 0b24b498..0df71ca3 100644 --- a/src/mitim_modules/portals/scripts/run_portals.py +++ b/src/mitim_modules/portals/scripts/run_portals.py @@ -11,14 +11,14 @@ def main(): parser.add_argument("folder", type=str, help="Simulation folder") parser.add_argument("--namelist", type=str, required=False, default=None) # namelist.portals.yaml file, otherwise what's in the current folder parser.add_argument("--input", type=str, required=False, default=None) # input.gacode file, otherwise what's in the current folder - parser.add_argument('--cold_start', required=False, default=False, action='store_true') + parser.add_argument('--cold', required=False, default=False, action='store_true') args = parser.parse_args() folderWork = Path(args.folder) portals_namelist = args.namelist inputgacode = args.input - cold_start = args.cold_start + cold_start = args.cold if portals_namelist is None: portals_namelist = IOtools.expandPath('.') / "namelist.portals.yaml" diff --git a/templates/namelist.maestro.yaml b/templates/namelist.maestro.yaml index 2be88b90..8b404622 100644 --- a/templates/namelist.maestro.yaml +++ b/templates/namelist.maestro.yaml @@ -94,6 +94,10 @@ maestro: solution: predicted_roa: [0.35, 0.45, 0.55, 0.65, 0.75, 0.875, 0.9] keep_full_model_folder: false + exploration_ranges: + ymax_rel: 1.0 + ymin_rel: 1.0 + yminymax_atleast: [null, 4] transport: options: tglf: @@ -107,10 +111,6 @@ maestro: initialization_parameters: thermalize_fast: true quasineutrality: true - exploration_ranges: - ymax_rel: 1.0 - ymin_rel: 1.0 - yminymax_atleast: [null, 4] change_last_radial_call: true use_previous_surrogate_data: true try_flux_match_only_for_first_point: true diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index 86cb165f..518dfad1 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -35,6 +35,34 @@ solution: # Impurity to do flux-matching for if nZ enabled (name of first impurity instance AFTER postprocessing), e.g. "W" trace_impurity: null + # Options that define boundaries for the optimization + exploration_ranges: + + # Bounds for the input parameters + # - ymax_rel (and ymin_rel) can be float (common for all radii, channels) or the dictionary directly, e.g.: + # ymax_rel = { + # 'te': [1.0, 0.5, 0.5, 0.5], + # 'ti': [0.5, 0.5, 0.5, 0.5], + # 'ne': [1.0, 0.5, 0.5, 0.5] + # } + limits_are_relative: true + ymax_rel: 1.0 + ymin_rel: 1.0 + yminymax_atleast: null # Cold hardGradientLimits, defines the ranges that will be, at least, covered, e.g. [0,2] + + fixed_gradients: null + + # enforce_finite_aLT is used to be able to select ymin_rel = 2.0 for ne but ensure that te, ti is at, e.g., enforce_finite_aLT = 0.95 + enforce_finite_aLT: null + + define_ranges_from_profiles: null + + # start_from_folder is a folder from which to grab optimization_data and optimization_extra + # (if used with reevaluate_targets>0, change targets by reevaluating with different parameters) + start_from_folder: null + + reevaluate_targets: 0 + # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) keep_full_model_folder: true From 24a58b8012439611f58182d92555d5eb38f6c6e3 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 30 Aug 2025 13:55:06 -0400 Subject: [PATCH 265/385] namelist now refers to namelist location of submodules --- .../maestro/utils/PORTALSbeat.py | 2 +- src/mitim_modules/portals/PORTALSmain.py | 50 +++++---------- .../portals/scripts/run_portals.py | 13 ++-- .../gacode_tools/utils/GACODEdefaults.py | 17 ++--- .../transp_tools/src/TRANSPsingularity.py | 1 + templates/input.cgyro.models.yaml | 2 - templates/input.gx.models.yaml | 2 - templates/input.neo.models.yaml | 1 - templates/input.tglf.models.yaml | 46 +++++++------- templates/namelist.maestro.yaml | 62 ++++++++++++++++--- templates/namelist.portals.yaml | 60 +++++++++--------- tests/MAESTRO_workflow.py | 2 +- tutorials/PORTALS_tutorial.py | 2 +- 13 files changed, 143 insertions(+), 117 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index a94c602f..af5f274a 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -76,7 +76,7 @@ def run(self, **kwargs): cold_start = kwargs.get('cold_start', False) - portals_fun = PORTALSmain.portals(self.folder) + portals_fun = PORTALSmain.portals(self.folder, portals_namelist = self.portals_parameters["portals_namelist_location"]) portals_fun.portals_parameters = IOtools.deep_dict_update(portals_fun.portals_parameters, self.portals_parameters) portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.optimization_options) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index 04ed9c3f..feae2e92 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -49,8 +49,8 @@ def __init__( self.portals_parameters = IOtools.read_mitim_yaml(self.portals_namelist) # Read optimization namelist (always the default, the values to be modified are in the portals one) - if self.portals_parameters["optimization_namelist"] is not None: - self.optimization_namelist = self.portals_parameters["optimization_namelist"] + if self.portals_parameters["optimization_namelist_location"] is not None: + self.optimization_namelist = self.portals_parameters["optimization_namelist_location"] else: self.optimization_namelist = __mitimroot__ / "templates" / "namelist.optimization.yaml" self.optimization_options = IOtools.read_mitim_yaml(self.optimization_namelist) @@ -70,22 +70,10 @@ def prep( seedInitial=None, askQuestions=True, ): - """ - Notes: - - ymax_rel (and ymin_rel) can be float (common for all radii, channels) or the dictionary directly, e.g.: - ymax_rel = { - 'te': [1.0, 0.5, 0.5, 0.5], - 'ti': [0.5, 0.5, 0.5, 0.5], - 'ne': [1.0, 0.5, 0.5, 0.5] - } - - enforce_finite_aLT is used to be able to select ymin_rel = 2.0 for ne but ensure that te, ti is at, e.g., enforce_finite_aLT = 0.95 - - start_from_folder is a folder from which to grab optimization_data and optimization_extra - (if used with reevaluate_targets>0, change targets by reevaluating with different parameters) - - seedInitial can be optionally give a seed to randomize the starting profile (useful for developing, paper writing) - """ - ymax_rel = self.portals_parameters["solution"]["exploration_ranges"]["ymax_rel"] - ymin_rel = self.portals_parameters["solution"]["exploration_ranges"]["ymin_rel"] + # Grab exploration ranges + ymax = self.portals_parameters["solution"]["exploration_ranges"]["ymax"] + ymin = self.portals_parameters["solution"]["exploration_ranges"]["ymin"] limits_are_relative = self.portals_parameters["solution"]["exploration_ranges"]["limits_are_relative"] fixed_gradients = self.portals_parameters["solution"]["exploration_ranges"]["fixed_gradients"] yminymax_atleast = self.portals_parameters["solution"]["exploration_ranges"]["yminymax_atleast"] @@ -105,32 +93,28 @@ def prep( key_rhos = "predicted_roa" if self.portals_parameters["solution"]["predicted_roa"] is not None else "predicted_rho" - # TO BE REMOVED IN FUTURE - if not isinstance(cold_start, bool): - raise Exception("cold_start must be a boolean") - # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialization # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if IOtools.isfloat(ymax_rel): - ymax_rel0 = copy.deepcopy(ymax_rel) + if IOtools.isfloat(ymax): + ymax0 = copy.deepcopy(ymax) - ymax_rel = {} + ymax = {} for prof in self.portals_parameters["solution"]["predicted_channels"]: - ymax_rel[prof] = np.array( [ymax_rel0] * len(self.portals_parameters["solution"][key_rhos]) ) + ymax[prof] = np.array( [ymax0] * len(self.portals_parameters["solution"][key_rhos]) ) - if IOtools.isfloat(ymin_rel): - ymin_rel0 = copy.deepcopy(ymin_rel) + if IOtools.isfloat(ymin): + ymin0 = copy.deepcopy(ymin) - ymin_rel = {} + ymin = {} for prof in self.portals_parameters["solution"]["predicted_channels"]: - ymin_rel[prof] = np.array( [ymin_rel0] * len(self.portals_parameters["solution"][key_rhos]) ) + ymin[prof] = np.array( [ymin0] * len(self.portals_parameters["solution"][key_rhos]) ) if enforce_finite_aLT is not None: for prof in ['te', 'ti']: - if prof in ymin_rel: - ymin_rel[prof] = np.array(ymin_rel[prof]).clip(min=None,max=enforce_finite_aLT) + if prof in ymin: + ymin[prof] = np.array(ymin[prof]).clip(min=None,max=enforce_finite_aLT) # Initialize print(">> PORTALS initalization module (START)", typeMsg="i") @@ -138,8 +122,8 @@ def prep( self, self.folder, mitim_state, - ymax_rel, - ymin_rel, + ymax, + ymin, start_from_folder=start_from_folder, define_ranges_from_profiles=define_ranges_from_profiles, fixed_gradients=fixed_gradients, diff --git a/src/mitim_modules/portals/scripts/run_portals.py b/src/mitim_modules/portals/scripts/run_portals.py index 0df71ca3..0b82d409 100644 --- a/src/mitim_modules/portals/scripts/run_portals.py +++ b/src/mitim_modules/portals/scripts/run_portals.py @@ -20,15 +20,10 @@ def main(): inputgacode = args.input cold_start = args.cold - if portals_namelist is None: - portals_namelist = IOtools.expandPath('.') / "namelist.portals.yaml" - else: - portals_namelist = Path(portals_namelist) - - if inputgacode is None: - inputgacode = IOtools.expandPath('.') / "input.gacode" - else: - inputgacode = Path(inputgacode) + # Actual PORTALS run + + portals_namelist = Path(portals_namelist) if portals_namelist is not None else IOtools.expandPath('.') / "namelist.portals.yaml" + inputgacode = Path(inputgacode) if inputgacode is not None else IOtools.expandPath('.') / "input.gacode" portals_fun = PORTALSmain.portals(folderWork, portals_namelist=portals_namelist) portals_fun.prep(inputgacode) diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index b8b2eb43..f2d83741 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -70,10 +70,10 @@ def add_code_settings(options,code_settings, models_file = "input.tglf.models.ya settings = IOtools.read_mitim_yaml(__mitimroot__ / "templates" / models_file) - found = False - code_settings = str(code_settings) + found = False + # Search by label first if str(code_settings) in settings: sett = settings[str(code_settings)] @@ -81,13 +81,14 @@ def add_code_settings(options,code_settings, models_file = "input.tglf.models.ya options[ikey] = sett["controls"][ikey] found = True else: - # Search by deprecated descriptor + # Search by deprecated descriptor (if available) for ikey in settings: - if settings[ikey]["deprecated_descriptor"] == code_settings: - sett = settings[ikey] - for jkey in sett["controls"]: - options[jkey] = sett["controls"][jkey] - found = True + if "deprecated_descriptor" in settings[ikey]: + if settings[ikey]["deprecated_descriptor"] == code_settings: + sett = settings[ikey] + for jkey in sett["controls"]: + options[jkey] = sett["controls"][jkey] + found = True break if not found: diff --git a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py index 9fcdc89e..1c6f6cf4 100644 --- a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py +++ b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py @@ -514,6 +514,7 @@ def pringLogTail(log_file, howmanylines=100, typeMsg="w"): print(txt, typeMsg=typeMsg) def runSINGULARITY_finish(folderWork, runid, tok, job_name): + embed() transp_job = FARMINGtools.mitim_job(folderWork) transp_job.define_machine( diff --git a/templates/input.cgyro.models.yaml b/templates/input.cgyro.models.yaml index a4d94070..a7a5e1fb 100644 --- a/templates/input.cgyro.models.yaml +++ b/templates/input.cgyro.models.yaml @@ -1,9 +1,7 @@ "Linear": - deprecated_descriptor: "0" controls: NONLINEAR_FLAG: 0 "Nonlinear": - deprecated_descriptor: "1" controls: NONLINEAR_FLAG: 1 \ No newline at end of file diff --git a/templates/input.gx.models.yaml b/templates/input.gx.models.yaml index ac60ec6c..209c0c6e 100644 --- a/templates/input.gx.models.yaml +++ b/templates/input.gx.models.yaml @@ -1,6 +1,5 @@ "Linear": # Linear with 10 ky modes - deprecated_descriptor: "0" controls: nonlinear_mode: false hyper: false @@ -10,6 +9,5 @@ ny: 28 "Nonlinear": - deprecated_descriptor: "1" controls: nonlinear_mode: true \ No newline at end of file diff --git a/templates/input.neo.models.yaml b/templates/input.neo.models.yaml index 6ed703b6..974bc4c5 100644 --- a/templates/input.neo.models.yaml +++ b/templates/input.neo.models.yaml @@ -1,4 +1,3 @@ "Sonic": - deprecated_descriptor: "0" controls: ROTATION_MODEL: 2 \ No newline at end of file diff --git a/templates/input.tglf.models.yaml b/templates/input.tglf.models.yaml index 94638662..de175cfe 100644 --- a/templates/input.tglf.models.yaml +++ b/templates/input.tglf.models.yaml @@ -1,9 +1,3 @@ -"SAT1": - deprecated_descriptor: "1" - controls: - SAT_RULE: 1 - UNITS: GYRO - "SAT0": deprecated_descriptor: "2" controls: @@ -11,6 +5,12 @@ UNITS: GYRO ETG_FACTOR: 1.25 +"SAT1": + deprecated_descriptor: "1" + controls: + SAT_RULE: 1 + UNITS: GYRO + "SAT1geo": deprecated_descriptor: "3" controls: @@ -25,14 +25,18 @@ XNU_MODEL: 3 WDIA_TRAPPED: 1.0 -"SAT2em": - deprecated_descriptor: "5" +"SAT2astra": + deprecated_descriptor: "100" controls: - SAT_RULE: 2 UNITS: CGYRO - XNU_MODEL: 3 - WDIA_TRAPPED: 1.0 USE_BPER: true + USE_AVE_ION_GRID: true + SAT_RULE: 2 + KYGRID_MODEL: 4 + XNU_MODEL: 3 + NBASIS_MAX: 6 + B_MODEL_SA: 1 + FT_MODEL_SA: 1 "SAT3": deprecated_descriptor: "6" @@ -42,26 +46,24 @@ XNU_MODEL: 3 WDIA_TRAPPED: 1.0 -"SAT2astra": - deprecated_descriptor: "100" +# Experimentation +"SAT3em basis": + deprecated_descriptor: "101" controls: + SAT_RULE: 2 UNITS: CGYRO + XNU_MODEL: 3 + WDIA_TRAPPED: 1.0 USE_BPER: true - USE_AVE_ION_GRID: true - SAT_RULE: 2 KYGRID_MODEL: 4 - XNU_MODEL: 3 NBASIS_MAX: 6 - B_MODEL_SA: 1 - FT_MODEL_SA: 1 -"Experimentation] SAT3em basis": - deprecated_descriptor: "101" +#TODO: To be removed in the future +"SAT2em": + deprecated_descriptor: "5" 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/templates/namelist.maestro.yaml b/templates/namelist.maestro.yaml index 8b404622..db0fffce 100644 --- a/templates/namelist.maestro.yaml +++ b/templates/namelist.maestro.yaml @@ -1,12 +1,23 @@ +# (**************************************************************************************************************) +# +# This is a complete template for a MAESTRO simulation +# +# (**************************************************************************************************************) + +# ------------------------------------------------------------------------------------------------------------ +# Plasma +# ------------------------------------------------------------------------------------------------------------ + # Name of this simulation flag: SPARC PRD -# Random seed for reproducibility (master seed to be sent to all beats that receive seed) +# Master random seed for reproducibility, to be sent to all beats that can receive seed seed: 0 # Machine parameters machine: + # Engineering parameters Bt: 12.2 Ip: 8.7 @@ -41,16 +52,25 @@ assumptions: # Initialization of profiles initialization: - BetaN: 1.0 - density_peaking: 1.3 + + # neped as the input (engineering) parameter assume_neped: true neped_20: 2.5 - nesep_ratio: 0.3 + # ne_sep / ne_ped ratio to assume (e.g. for EPED) and construct profiles with + nesep_ratio: 0.3 + + # Profiles will be initialize such that these parameters are matched as close as possible + BetaN: 1.0 + density_peaking: 1.3 + # Edge parameters Tesep_eV: 75.0 +# ------------------------------------------------------------------------------------------------------------ # MAESTRO workflow parameters +# ------------------------------------------------------------------------------------------------------------ + maestro: # Remove intermediate files @@ -68,12 +88,16 @@ maestro: eped_namelist: nn_location: $MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-MODEL-SPARC.keras norm_location: $MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-NORMALIZATION-SPARC.txt + + # Corrections set for EPED indicate the values that I force them to be exact, to avoid NN issues with non-trained parameters corrections_set: Bt: 12.2 R: 1.85 a: 0.57 - ptop_multiplier: 1.0 - TioverTe: 1.0 + + # Operations after EPED pressure predictions + ptop_multiplier: 1.0 # Multiplies ptop + TioverTe: 1.0 # Ti/Te assumption at the pedestal top eped_initializer_beat: use_default: false @@ -88,32 +112,52 @@ maestro: TioverTe: 1.0 portals_beat: + use_default: false + portals_namelist: + + # If null, it will use the default namelist at: __mitimroot__ / "templates" / "namelist.portals.yaml" + portals_namelist_location: null + + # ------------------------------------------------------------------------------------------------ + # PORTALS parameters (only those parameters to change the main PORTALS namelist) + # ------------------------------------------------------------------------------------------------ portals_parameters: solution: predicted_roa: [0.35, 0.45, 0.55, 0.65, 0.75, 0.875, 0.9] keep_full_model_folder: false exploration_ranges: - ymax_rel: 1.0 - ymin_rel: 1.0 + ymax: 1.0 + ymin: 4.0 yminymax_atleast: [null, 4] transport: options: tglf: run: - code_settings: SAT2astra + code_settings: "SAT2astra" extraOptions: USE_BPER: true + keep_files: "none" target: options: force_zero_particle_flux: true + + # Operations to do to mitim_state before passing it to PORTALS initialization_parameters: thermalize_fast: true quasineutrality: true + + # If the width of the pedestal has changed, modify the BC for PORTALS accordingly change_last_radial_call: true + + # Later beats to utilize past beats surrogate data if available use_previous_surrogate_data: true + + # try_flux_match_only_for_first_point: true + + # Operations that will be included as part of the profiles_postprocessing_fun transport_preprocessing: lumpImpurities: true enforce_same_density_gradients: true diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index 518dfad1..357d8ba0 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -1,4 +1,5 @@ -# ------------------------------------------------------------------------------------------------------------ +# (**************************************************************************************************************) +# # This is a complete template for a PORTALS-TGLF+NEO profile prediction. # The user is welcomed to change any of the parameters below to fit their specific use case, # by copying this template namelist and instantiating the PORTALS object with its path: @@ -6,7 +7,8 @@ # portals_fun = PORTALSmain.portals(folder, portals_namelist=PATH_TO_NAMELIST) # # Alternatively, you can simply use the default parameters provided in this template by not -# specifying a namelist and then, in the launching script change the dictionary values: +# specifying a namelist and then, in the launching script change the dictionary values +# (*before* you perform the portals_fun.prep() command): # # portals_fun = PORTALSmain.portals(folder) # portals_fun.portals_parameters["solution"]['predicted_roa'] = [0.25, 0.45, 0.65, 0.85] @@ -18,6 +20,7 @@ # # portals_fun.optimization_options["convergence_options"]["maximum_iterations"] = 10 # +# (**************************************************************************************************************) # ----------------------------------------------------------------- # Main solver parameters @@ -38,21 +41,21 @@ solution: # Options that define boundaries for the optimization exploration_ranges: - # Bounds for the input parameters - # - ymax_rel (and ymin_rel) can be float (common for all radii, channels) or the dictionary directly, e.g.: - # ymax_rel = { + # Are ymax and ymin relative or absolute + limits_are_relative: true + + # Bounds for the input parameters. ymax/ymin can be float (common for all radii, channels) or a dictionary:: + # ymax: # 'te': [1.0, 0.5, 0.5, 0.5], # 'ti': [0.5, 0.5, 0.5, 0.5], - # 'ne': [1.0, 0.5, 0.5, 0.5] - # } - limits_are_relative: true - ymax_rel: 1.0 - ymin_rel: 1.0 - yminymax_atleast: null # Cold hardGradientLimits, defines the ranges that will be, at least, covered, e.g. [0,2] + # ... + ymax: 1.0 + ymin: 1.0 + yminymax_atleast: null # Defines minimum range of exploration, e.g. [0,2], even if not achieved via ymin/ymax fixed_gradients: null - # enforce_finite_aLT is used to be able to select ymin_rel = 2.0 for ne but ensure that te, ti is at, e.g., enforce_finite_aLT = 0.95 + # enforce_finite_aLT is used to be able to select ymin = 2.0 for ne but ensure that te, ti is at, e.g., enforce_finite_aLT = 0.95 enforce_finite_aLT: null define_ranges_from_profiles: null @@ -109,9 +112,9 @@ transport: # Simulation kwargs to be passed directly to run and read commands (defaults here) options: - # *********************************** + # ********************************************************************************************************* # TGLF - # *********************************** + # ********************************************************************************************************* tglf: # Kwargs to be passed to the run command @@ -132,14 +135,14 @@ transport: cores_per_tglf_instance: 1 # (%) Error (std, in percent) of model evaluation TGLF if not scan trick - percent_error: 5 + percent_error: 5.0 # If True, and fast ions have been included, sum fast. This only occurs if the specie is considered fast by TGLF (it could be fast in input.gacode but thermal for TGLF) Qi_includes_fast: false - # *********************************** + # ********************************************************************************************************* # NEO - # *********************************** + # ********************************************************************************************************* neo: run: @@ -148,11 +151,11 @@ transport: read: {} # (%) Error (std, in percent) of model evaluation - percent_error: 10 + percent_error: 10.0 - # *********************************** + # ********************************************************************************************************* # CGYRO - # *********************************** + # ********************************************************************************************************* cgyro: run: @@ -171,9 +174,9 @@ transport: # (%) For CGYRO runs, minimum error based on target if case is considered stable Qi_stable_percent_error: 5.0 - # *********************************** + # ********************************************************************************************************* # GX - # *********************************** + # ********************************************************************************************************* gx: run: @@ -187,16 +190,17 @@ transport: read: tmin: 0.0 - # Function to post-process input.gacode only BEFORE passing to transport codes + # Function to post-process input.gacode only *before* passing to transport codes profiles_postprocessing_fun: null # Corrections to be applied to each iteration input.gacode file applyCorrections: - Ti_thermals: true # Keep all thermal ion temperatures equal to the main Ti - ni_thermals: true # Adjust for quasineutrality by modifying the thermal ion densities together with ne + Ti_thermals: true # Keep all thermal ion temperatures equal to the main Ti + ni_thermals: true # Adjust for quasineutrality by modifying the thermal ion densities together with ne recalculate_ptot: true # Recompute PTOT to insert in input file each time - Tfast_ratio: false # Keep the ratio of Tfast/Te constant throughout the Te evolution - force_mach: null # Change w0 to match this Mach number when Ti varies + Tfast_ratio: false # Keep the ratio of Tfast/Te constant throughout the Te evolution + force_mach: null # Change w0 to match this Mach number when Ti varies + # ----------------------------------------------------------------- # Target model parameters @@ -228,7 +232,7 @@ target: # Optimization options namelist # ----------------------------------------------------------------- -optimization_namelist: null # If null, it will grab the default at: __mitimroot__ / "templates" / "namelist.optimization.yaml" +optimization_namelist_location: null # If null, it will grab the default at: __mitimroot__ / "templates" / "namelist.optimization.yaml" # ----------------------------------------------------------------- # Optimization options (to change the main optimization namelist) diff --git a/tests/MAESTRO_workflow.py b/tests/MAESTRO_workflow.py index ec8857c6..5e5e1178 100644 --- a/tests/MAESTRO_workflow.py +++ b/tests/MAESTRO_workflow.py @@ -3,7 +3,7 @@ from mitim_tools import __mitimroot__ from mitim_modules.maestro.scripts import run_maestro -cold_start = True +cold_start = False folder = __mitimroot__ / "tests" / "scratch" / "maestro_test" template = __mitimroot__ / "templates" / "namelist.maestro.yaml" diff --git a/tutorials/PORTALS_tutorial.py b/tutorials/PORTALS_tutorial.py index 77a74314..36d20804 100644 --- a/tutorials/PORTALS_tutorial.py +++ b/tutorials/PORTALS_tutorial.py @@ -38,7 +38,7 @@ portals_fun.optimization_options['convergence_options']['stopping_criteria_parameters']["maximum_value_is_rel"] = True # Prepare run: search +-100% the original gradients -portals_fun.prep(inputgacode, ymax_rel=1.0, ymin_rel=1.0) +portals_fun.prep(inputgacode) # -------------------------------------------------------------------------------------------- # Run (optimization following namelist: templates/main.namelists.json) From 35a128d0aeb3e892dc05b1e4c6e4f1b2c9459cff Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 30 Aug 2025 15:09:50 -0400 Subject: [PATCH 266/385] Bug fix: TRANSP requires exclusive node --- src/mitim_tools/transp_tools/src/TRANSPsingularity.py | 5 ++++- tests/MAESTRO_workflow.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py index 1c6f6cf4..031eda0b 100644 --- a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py +++ b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py @@ -445,6 +445,9 @@ def runSINGULARITY( shellPreCommands=shellPreCommands, ) + if 'exclusive' not in transp_job.machineSettings["slurm"] or not transp_job.machineSettings["slurm"]["exclusive"]: + print("\tTRANSP typically requires exclusive node allocation, but that has not been requested, prone to failure", typeMsg="q") + transp_job.run(waitYN=False) IOtools.shutil_rmtree(folderWork / 'tmp_inputs') @@ -514,7 +517,7 @@ def pringLogTail(log_file, howmanylines=100, typeMsg="w"): print(txt, typeMsg=typeMsg) def runSINGULARITY_finish(folderWork, runid, tok, job_name): - embed() + transp_job = FARMINGtools.mitim_job(folderWork) transp_job.define_machine( diff --git a/tests/MAESTRO_workflow.py b/tests/MAESTRO_workflow.py index 5e5e1178..ec8857c6 100644 --- a/tests/MAESTRO_workflow.py +++ b/tests/MAESTRO_workflow.py @@ -3,7 +3,7 @@ from mitim_tools import __mitimroot__ from mitim_modules.maestro.scripts import run_maestro -cold_start = False +cold_start = True folder = __mitimroot__ / "tests" / "scratch" / "maestro_test" template = __mitimroot__ / "templates" / "namelist.maestro.yaml" From fadeec89cff4a425657110d418ad363da77112b5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 30 Aug 2025 15:36:12 -0400 Subject: [PATCH 267/385] misc --- src/mitim_modules/maestro/utils/PORTALSbeat.py | 4 +++- templates/namelist.portals.yaml | 4 ++-- tests/MAESTRO_workflow.py | 2 +- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index af5f274a..661a4451 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -26,6 +26,7 @@ def prepare(self, try_flux_match_only_for_first_point = True, change_last_radial_call = False, portals_parameters = {}, + portals_namelist_location = None, initialization_parameters = {}, optimization_options = {}, enforce_impurity_radiation_existence = True, @@ -59,6 +60,7 @@ def prepare(self, self.profiles_current.write_state(file = self.fileGACODE) self.portals_parameters = portals_parameters + self.portals_namelist_location = portals_namelist_location self.optimization_options = optimization_options self.initialization_parameters = initialization_parameters @@ -76,7 +78,7 @@ def run(self, **kwargs): cold_start = kwargs.get('cold_start', False) - portals_fun = PORTALSmain.portals(self.folder, portals_namelist = self.portals_parameters["portals_namelist_location"]) + portals_fun = PORTALSmain.portals(self.folder, portals_namelist = self.portals_namelist_location) portals_fun.portals_parameters = IOtools.deep_dict_update(portals_fun.portals_parameters, self.portals_parameters) portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.optimization_options) diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index 357d8ba0..1d9c8df6 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -46,8 +46,8 @@ solution: # Bounds for the input parameters. ymax/ymin can be float (common for all radii, channels) or a dictionary:: # ymax: - # 'te': [1.0, 0.5, 0.5, 0.5], - # 'ti': [0.5, 0.5, 0.5, 0.5], + # 'te': [1.0, 0.5, 0.5, 0.5] + # 'ti': [0.5, 0.5, 0.5, 0.5] # ... ymax: 1.0 ymin: 1.0 diff --git a/tests/MAESTRO_workflow.py b/tests/MAESTRO_workflow.py index ec8857c6..5e5e1178 100644 --- a/tests/MAESTRO_workflow.py +++ b/tests/MAESTRO_workflow.py @@ -3,7 +3,7 @@ from mitim_tools import __mitimroot__ from mitim_modules.maestro.scripts import run_maestro -cold_start = True +cold_start = False folder = __mitimroot__ / "tests" / "scratch" / "maestro_test" template = __mitimroot__ / "templates" / "namelist.maestro.yaml" From 89e63100a4ce3fc3655bbe545af0d168c1926fa2 Mon Sep 17 00:00:00 2001 From: pabloprf Date: Sun, 31 Aug 2025 15:02:50 -0400 Subject: [PATCH 268/385] Do not raise TRANSP exclusive warning as a question, since it is not clear yet --- src/mitim_tools/opt_tools/scripts/slurm.py | 7 ++++--- src/mitim_tools/transp_tools/src/TRANSPsingularity.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/mitim_tools/opt_tools/scripts/slurm.py b/src/mitim_tools/opt_tools/scripts/slurm.py index dda7006b..2dae0782 100644 --- a/src/mitim_tools/opt_tools/scripts/slurm.py +++ b/src/mitim_tools/opt_tools/scripts/slurm.py @@ -17,7 +17,8 @@ def run_slurm( seed_specific=0, machine="local", exclude=None, - mem=None + mem=None, + exclusive=False ): folder = IOtools.expandPath(folder) @@ -46,13 +47,13 @@ def run_slurm( command, folder, folder_local=folder, - slurm={"partition": partition, 'exclude': exclude}, + slurm={"partition": partition, 'exclude': exclude,'exclusive': exclusive}, slurm_settings = { 'nameJob': nameJob, 'minutes': int(60 * hours), 'ntasks': 1, 'cpuspertask': n, - 'memory_req_by_job': mem, + 'memory_req_by_job': mem } ) diff --git a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py index 031eda0b..44906133 100644 --- a/src/mitim_tools/transp_tools/src/TRANSPsingularity.py +++ b/src/mitim_tools/transp_tools/src/TRANSPsingularity.py @@ -446,7 +446,7 @@ def runSINGULARITY( ) if 'exclusive' not in transp_job.machineSettings["slurm"] or not transp_job.machineSettings["slurm"]["exclusive"]: - print("\tTRANSP typically requires exclusive node allocation, but that has not been requested, prone to failure", typeMsg="q") + print("\tTRANSP typically requires exclusive node allocation, but that has not been requested, prone to failure", typeMsg="w") transp_job.run(waitYN=False) From 57c0b5108a7e866bc67f23c24cb303577eca9943 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 1 Sep 2025 11:37:24 -0400 Subject: [PATCH 269/385] Common CGYRO and GX plotting of time traces --- src/mitim_tools/gacode_tools/CGYROtools.py | 47 +-------------- .../simulation_tools/physics/GXtools.py | 15 ++--- .../simulation_tools/utils/SIMplot.py | 54 ++++++++++++++++++ templates/input.gx.controls | 8 +-- templates/input.gx.models.yaml | 7 ++- templates/namelist.maestro.yaml | 10 +++- tests/CGYRO_workflow.py | 57 ++++++++++--------- tests/GX_workflow.py | 2 +- 8 files changed, 113 insertions(+), 87 deletions(-) create mode 100644 src/mitim_tools/simulation_tools/utils/SIMplot.py diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index b114398c..4b49c20a 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -4,12 +4,13 @@ from mitim_tools import __mitimroot__ from mitim_tools.gacode_tools.utils import GACODEdefaults, CGYROutils from mitim_tools.simulation_tools import SIMtools +from mitim_tools.simulation_tools.utils import SIMplot from mitim_tools.misc_tools import GRAPHICStools, CONFIGread from mitim_tools.gacode_tools.utils import GACODEplotting from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -class CGYRO(SIMtools.mitim_simulation): +class CGYRO(SIMtools.mitim_simulation, SIMplot.GKplotting): def __init__( self, rhos=[0.4, 0.6], # rho locations of interest @@ -361,50 +362,6 @@ def plot( self.results = self.results_all - def _plot_trace(self, ax, label, variable, c="b", lw=1, ls="-", label_plot='', meanstd=True, var_meanstd= None): - - t = self.results[label].t - - if not isinstance(variable, str): - z = variable - if var_meanstd is not None: - z_mean = var_meanstd[0] - z_std = var_meanstd[1] - - else: - z = self.results[label].__dict__[variable] - if meanstd and (f'{variable}_mean' in self.results[label].__dict__): - z_mean = self.results[label].__dict__[variable + '_mean'] - z_std = self.results[label].__dict__[variable + '_std'] - else: - z_mean = None - z_std = None - - ax.plot( - t, - z, - ls=ls, - lw=lw, - c=c, - label=label_plot, - ) - - if meanstd and z_std>0.0: - GRAPHICStools.fillGraph( - ax, - t[t>self.results[label].tmin], - z_mean, - y_down=z_mean - - z_std, - y_up=z_mean - + z_std, - alpha=0.1, - color=c, - lw=0.5, - islwOnlyMean=True, - label=label_plot + f" $\\mathbf{{{z_mean:.2f} \\pm {z_std:.2f}}}$ (1$\\sigma$)", - ) - def plot_inputs(self, ax = None, label="", c="b", ms = 10, normalization_label=None, only_plot_differences=False): if ax is None: diff --git a/src/mitim_tools/simulation_tools/physics/GXtools.py b/src/mitim_tools/simulation_tools/physics/GXtools.py index 3d8b84a1..85b99927 100644 --- a/src/mitim_tools/simulation_tools/physics/GXtools.py +++ b/src/mitim_tools/simulation_tools/physics/GXtools.py @@ -3,12 +3,13 @@ from mitim_tools.misc_tools import GRAPHICStools, IOtools, GUItools, CONFIGread from mitim_tools.gacode_tools.utils import GACODEdefaults, CGYROutils from mitim_tools.simulation_tools import SIMtools +from mitim_tools.simulation_tools.utils import SIMplot from mitim_tools.misc_tools.LOGtools import printMsg as print from mitim_tools import __mitimroot__ from mitim_tools import __version__ as mitim_version from IPython import embed -class GX(SIMtools.mitim_simulation): +class GX(SIMtools.mitim_simulation, SIMplot.GKplotting): def __init__( self, rhos=[0.4, 0.6], # rho locations of interest @@ -92,7 +93,7 @@ def plot( colors = GRAPHICStools.listColors() # Fluxes - fig = self.fn.add_figure(label=f"{extratitle}Fluxes", tab_color=fn_color) + fig = self.fn.add_figure(label=f"{extratitle}Transport Fluxes", tab_color=fn_color) grid = plt.GridSpec(1, 3, hspace=0.7, wspace=0.2) @@ -107,9 +108,9 @@ def plot( typeLs = '-' if c.t.shape[0]>20 else '-s' - ax1.plot(c.t, c.Qe, typeLs, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) - ax2.plot(c.t, c.Qi, typeLs, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) - ax3.plot(c.t, c.Ge, typeLs, label=f"{label} rho={self.rhos[irho]}", color=colors[i]) + self._plot_trace(ax1,self.results[label]['output'][irho],"Qe",c=colors[i],lw=1.0,ls='-',label_plot=f"{label}, Total") + self._plot_trace(ax2,self.results[label]['output'][irho],"Qi",c=colors[i],lw=1.0,ls='-',label_plot=f"{label}, Total") + self._plot_trace(ax3,self.results[label]['output'][irho],"Ge",c=colors[i],lw=1.0,ls='-',label_plot=f"{label}, Total") i += 1 @@ -134,7 +135,7 @@ def plot( # Linear stability - fig = self.fn.add_figure(label=f"{extratitle}Linear stability", tab_color=fn_color) + fig = self.fn.add_figure(label=f"{extratitle}Linear Stability", tab_color=fn_color) grid = plt.GridSpec(2, 2, hspace=0.7, wspace=0.2) @@ -381,7 +382,7 @@ def read(self): ikx = 0 self.ky = data.groups['Grids'].variables['ky'][1:] # (ky) self.w = data.groups['Diagnostics'].variables['omega_kxkyt'][:,1:,ikx,0] # (time, ky) - self.g = data.groups['Diagnostics'].variables['omega_kxkyt'][:,1:,ikx,1] # (time, ky) + self.g = data.groups['Diagnostics'].variables['omega_kxkyt'][:,1:,ikx,1] # (time, ky) # Fluxes Q = data.groups['Diagnostics'].variables['HeatFlux_st'] # (time, species) diff --git a/src/mitim_tools/simulation_tools/utils/SIMplot.py b/src/mitim_tools/simulation_tools/utils/SIMplot.py new file mode 100644 index 00000000..009f610c --- /dev/null +++ b/src/mitim_tools/simulation_tools/utils/SIMplot.py @@ -0,0 +1,54 @@ +from cProfile import label +from mitim_tools.misc_tools import GRAPHICStools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +class GKplotting: + def _plot_trace(self, ax, object_or_label, variable, c="b", lw=1, ls="-", label_plot='', meanstd=True, var_meanstd= None): + + if isinstance(object_or_label, str): + object_grab = self.results[object_or_label] + else: + object_grab = object_or_label + + t = object_grab.t + + if not isinstance(variable, str): + z = variable + if var_meanstd is not None: + z_mean = var_meanstd[0] + z_std = var_meanstd[1] + + else: + z = object_grab.__dict__[variable] + if meanstd and (f'{variable}_mean' in object_grab.__dict__): + z_mean = object_grab.__dict__[variable + '_mean'] + z_std = object_grab.__dict__[variable + '_std'] + else: + z_mean = None + z_std = None + + ax.plot( + t, + z, + ls=ls, + lw=lw, + c=c, + label=label_plot, + ) + + if meanstd and z_std>0.0: + GRAPHICStools.fillGraph( + ax, + t[t>object_grab.tmin], + z_mean, + y_down=z_mean + - z_std, + y_up=z_mean + + z_std, + alpha=0.1, + color=c, + lw=0.5, + islwOnlyMean=True, + label=label_plot + f" $\\mathbf{{{z_mean:.3f} \\pm {z_std:.3f}}}$ (1$\\sigma$)", + ) \ No newline at end of file diff --git a/templates/input.gx.controls b/templates/input.gx.controls index f6e52998..8a22478e 100644 --- a/templates/input.gx.controls +++ b/templates/input.gx.controls @@ -5,10 +5,10 @@ debug = false [Dimensions] - ntheta = 32 # number of points along field line (theta) per 2pi segment + ntheta = 24 # number of points along field line (theta) per 2pi segment nperiod = 1 # number of 2pi segments along field line is 2*nperiod-1 - ny = 106 # number of real-space grid-points in y, nky = 1 + (ny-1)/3 --> 36 ky modes - nx = 382 # number of real-space grid-points in x, nkx = 1 + 2*(nx-1)/3 --> 255 kx modes + ny = 64 # number of real-space grid-points in y, nky = 1 + (ny-1)/3 + nx = 192 # number of real-space grid-points in x, nkx = 1 + 2*(nx-1)/3 nhermite = 48 # number of hermite moments (v_parallel resolution) nlaguerre = 16 # number of laguerre moments (mu B resolution) @@ -22,7 +22,7 @@ ei_colls = false # turn off electron-ion collisions [Time] - t_max = 1000.0 # end time (in units of L_ref/vt_ref) + t_max = 500.0 # end time (in units of L_ref/vt_ref) scheme = "rk3" # use RK3 timestepping scheme (with adaptive timestepping) [Initialization] diff --git a/templates/input.gx.models.yaml b/templates/input.gx.models.yaml index 209c0c6e..b7d24a10 100644 --- a/templates/input.gx.models.yaml +++ b/templates/input.gx.models.yaml @@ -10,4 +10,9 @@ "Nonlinear": controls: - nonlinear_mode: true \ No newline at end of file + nonlinear_mode: true + hyper: true + HB_hyper: true + nperiod: 1 + nx: 192 + ny: 64 diff --git a/templates/namelist.maestro.yaml b/templates/namelist.maestro.yaml index db0fffce..51738d53 100644 --- a/templates/namelist.maestro.yaml +++ b/templates/namelist.maestro.yaml @@ -23,13 +23,21 @@ machine: # Separatrix specification separatrix: + + # Type of separatrix (freegs, geqdsk) type: freegs + + # Parameters for this type of separatrix parameters: + + # If the separatrix type receives boundary parameterization, parameters here R: 1.85 a: 0.57 delta_sep: 0.57 kappa_sep: 1.97 n_mxh: 5 + + # If the separatrix type receives a geqdsk file directly, provide path here geqdsk_file: "" # Heating parameters @@ -73,7 +81,7 @@ assumptions: maestro: - # Remove intermediate files + # Remove intermediate files to avoid heavy MAESTRO simulation folders keep_all_files: true # Sequence of beats diff --git a/tests/CGYRO_workflow.py b/tests/CGYRO_workflow.py index 4f27f67e..ac5f97b1 100644 --- a/tests/CGYRO_workflow.py +++ b/tests/CGYRO_workflow.py @@ -24,7 +24,7 @@ cgyro.run( 'linear', - code_settings=0, + code_settings="Linear", extraOptions={ 'KY':0.5, 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file @@ -45,30 +45,31 @@ cgyro.read(label="cgyro1") cgyro.plot(labels=["cgyro1"]) -# # --------------- -# # Scan of KY -# # --------------- - -# run_type = 'normal' - -# cgyro.run_scan( -# 'scan1', -# extraOptions={ -# 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file -# }, -# variable='KY', -# varUpDown=[0.3,0.4], -# slurm_setup={ -# 'cores':16 -# }, -# cold_start=cold_start, -# forceIfcold_start=True, -# run_type=run_type -# ) - -# cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) - -# fig = cgyro.fn.add_figure(label="Quick linear") -# cgyro.plot_quick_linear(labels=["scan1_KY_0.3","scan1_KY_0.4"], fig = fig) - -# cgyro.fn.show() +# --------------- +# Scan of KY +# --------------- + +run_type = 'normal' + +cgyro.run_scan( + 'scan1', + code_settings="Linear", + extraOptions={ + 'MAX_TIME': 10.0, # Short, I just want to test the run. Enough to get the restart file + }, + variable='KY', + varUpDown=[0.3,0.4], + slurm_setup={ + 'cores':16 + }, + cold_start=cold_start, + forceIfcold_start=True, + run_type=run_type + ) + +cgyro.plot(labels=["scan1_KY_0.3","scan1_KY_0.4"], fn = cgyro.fn) + +fig = cgyro.fn.add_figure(label="Quick linear") +cgyro.plot_quick_linear(labels=["scan1_KY_0.3","scan1_KY_0.4"], fig = fig) + +cgyro.fn.show() diff --git a/tests/GX_workflow.py b/tests/GX_workflow.py index 9c3f4c41..36424da8 100644 --- a/tests/GX_workflow.py +++ b/tests/GX_workflow.py @@ -24,7 +24,7 @@ gx.run( 'gx1/', cold_start=cold_start, - code_settings=0, # Linear + code_settings="Linear", extraOptions={ 't_max':5.0, # Run up to 5 a/c_s (should take ~2min using 8 A100s) 'y0' :5.0, # kymin = 1/y0 = 0.2 From 3760a8f308a8873e7617f0bec3095094d9ab7663 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 1 Sep 2025 11:38:02 -0400 Subject: [PATCH 270/385] Make multipliers also potential lists --- src/mitim_tools/simulation_tools/SIMtools.py | 27 +++++++++++++++----- tests/TGLFscan_workflow.py | 1 + 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 48ba6915..08dee314 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -1048,15 +1048,22 @@ def cold_start_checker( def modifyInputs( input_class, code_settings=None, - extraOptions={}, - multipliers={}, - minimum_delta_abs={}, + extraOptions=None, + multipliers=None, + minimum_delta_abs=None, position_change=0, addControlFunction=None, controls_file = 'input.tglf.controls', **kwargs_to_function, ): + if extraOptions is None: + extraOptions = {} + if multipliers is None: + multipliers = {} + if minimum_delta_abs is None: + minimum_delta_abs = {} + # Check that those are valid flags GACODEdefaults.review_controls(extraOptions, control = controls_file) GACODEdefaults.review_controls(multipliers, control = controls_file) @@ -1123,28 +1130,34 @@ def modifyInputs( if len(multipliers) > 0: print("\t\t- Variables change:") for ikey in multipliers: + + if isinstance(multipliers[ikey], (list, np.ndarray)): + value_to_change_to = multipliers[ikey][position_change] + else: + value_to_change_to = multipliers[ikey] + # is a specie one? if "species" in input_class.__dict__.keys() and ikey.split("_")[0] in input_class.species[1]: specie = int(ikey.split("_")[-1]) varK = "_".join(ikey.split("_")[:-1]) var_orig = input_class.species[specie][varK] - var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + var_new = multiplier_input(var_orig, value_to_change_to, minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.species[specie][varK] = var_new else: if ikey in input_class.controls: var_orig = input_class.controls[ikey] - var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + var_new = multiplier_input(var_orig, value_to_change_to, minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.controls[ikey] = var_new elif ikey in input_class.plasma: var_orig = input_class.plasma[ikey] - var_new = multiplier_input(var_orig, multipliers[ikey], minimum_delta_abs = minimum_delta_abs.get(ikey,None)) + var_new = multiplier_input(var_orig, value_to_change_to, minimum_delta_abs = minimum_delta_abs.get(ikey,None)) input_class.plasma[ikey] = var_new else: print("\t- Variable to scan did not exist in original file, add it as extraOptions first",typeMsg="w",) - print(f"\t\t\t- Changing {ikey} from {var_orig} to {var_new} (x{multipliers[ikey]})") + print(f"\t\t\t- Changing {ikey} from {var_orig} to {var_new} (x{value_to_change_to})") return input_class diff --git a/tests/TGLFscan_workflow.py b/tests/TGLFscan_workflow.py index 1fdf7ebe..f41c044c 100644 --- a/tests/TGLFscan_workflow.py +++ b/tests/TGLFscan_workflow.py @@ -18,6 +18,7 @@ tglf.run_scan( subfolder = 'scan1', code_settings = None, + extraOptions = {"USE_BPER": [False, True]}, # extraOptions can receive a list to provide different values per rho cold_start = cold_start, runWaveForms = [0.67, 10.0], variable = 'RLTS_1', From eb19b0bed792a68b2971261dbfdde962e4a3fca5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 2 Sep 2025 10:51:43 -0400 Subject: [PATCH 271/385] Capability to send additional files, e.g. cgyro and gx --- .../gacode_tools/utils/GACODEdefaults.py | 2 +- .../plasmastate_tools/MITIMstate.py | 14 ++--- src/mitim_tools/simulation_tools/SIMtools.py | 23 ++++++-- .../simulation_tools/physics/GXtools.py | 54 +++++++++++++++++-- .../simulation_tools/utils/SIMplot.py | 1 - templates/input.gx.controls | 8 ++- 6 files changed, 83 insertions(+), 19 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py index f2d83741..73fe2859 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEdefaults.py @@ -73,7 +73,7 @@ def add_code_settings(options,code_settings, models_file = "input.tglf.models.ya code_settings = str(code_settings) found = False - + # Search by label first if str(code_settings) in settings: sett = settings[str(code_settings)] diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index b8e1d792..1ee3a76d 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2200,7 +2200,7 @@ def csv(self, file="input.gacode.xlsx"): # Code conversions # ************************************************************************************************************************************************ - def to_tglf(self, r=[0.5], code_settings=1, r_is_rho = True): + def to_tglf(self, r=[0.5], code_settings='SAT0', r_is_rho = True): # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function @@ -2390,7 +2390,7 @@ def interpolator(y): return input_parameters - def to_neo(self, r=[0.5], r_is_rho = True): + def to_neo(self, r=[0.5], r_is_rho = True, code_settings='Sonic'): # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function @@ -2442,7 +2442,7 @@ def interpolator(y): # Controls come from options # --------------------------------------------------------------------------------------------------------------------------------------- - controls = GACODEdefaults.addNEOcontrol(0) + controls = GACODEdefaults.addNEOcontrol(code_settings) # --------------------------------------------------------------------------------------------------------------------------------------- # Species come from profiles @@ -2538,7 +2538,7 @@ def interpolator(y): return input_parameters - def to_cgyro(self, r=[0.5], r_is_rho = True): + def to_cgyro(self, r=[0.5], r_is_rho = True, code_settings = 'Linear'): # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function @@ -2596,7 +2596,7 @@ def interpolator(y): # Controls come from options # --------------------------------------------------------------------------------------------------------------------------------------- - controls = GACODEdefaults.addCGYROcontrol(0) + controls = GACODEdefaults.addCGYROcontrol(code_settings) controls['PROFILE_MODEL'] = 1 # --------------------------------------------------------------------------------------------------------------------------------------- @@ -2692,7 +2692,7 @@ def interpolator(y): return input_parameters - def to_gx(self, r=[0.5], r_is_rho = True): + def to_gx(self, r=[0.5], r_is_rho = True, code_settings = 'Linear'): # <> Function to interpolate a curve <> from mitim_tools.misc_tools.MATHtools import extrapolateCubicSpline as interpolation_function @@ -2735,7 +2735,7 @@ def interpolator(y): # Controls come from options # --------------------------------------------------------------------------------------------------------------------------------------- - controls = GACODEdefaults.addGXcontrol(0) + controls = GACODEdefaults.addGXcontrol(code_settings) # --------------------------------------------------------------------------------------------------------------------------------------- # Species come from profiles diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 08dee314..bd2e1749 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -116,6 +116,7 @@ def run( attempts_execution=1, only_minimal_files=False, run_type = 'normal', # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit + additional_files_to_send = None, # Dict (rho keys) of files to send along with the run (e.g. for restart) ): if slurm_setup is None: @@ -142,6 +143,8 @@ def run( launchSlurm=launchSlurm, slurm_setup=slurm_setup, # + additional_files_to_send=additional_files_to_send, + # ApplyCorrections=ApplyCorrections, minimum_delta_abs=minimum_delta_abs, Quasineutral=Quasineutral, @@ -195,6 +198,10 @@ def _run_prepare( launchSlurm=True, slurm_setup=None, # ******************************** + # Additional files to send (e.g. restarts). Must be a dictionary with rho keys + # ******************************** + additional_files_to_send = None, + # ******************************** # Additional settings to correct/modify inputs # ******************************** **kwargs_control @@ -277,6 +284,7 @@ def _run_prepare( "inputs": latest_inputsFile[irho], "extraOptions": extraOptions, "multipliers": multipliers, + "additional_files_to_send": additional_files_to_send[irho] if additional_files_to_send is not None else None } if irho in rhosEvaluate: code_executor[subfolder_simulation][irho] = code_executor_full[subfolder_simulation][irho] @@ -383,6 +391,11 @@ def _run( with open(input_file_sim, "w") as f: f.write(code_executor[subfolder_sim][rho]["inputs"]) + # Copy potential additional files to send + if code_executor[subfolder_sim][rho]["additional_files_to_send"] is not None: + for file in code_executor[subfolder_sim][rho]["additional_files_to_send"]: + shutil.copy(file, folder_sim_this / Path(file).name) + # --------------------------------------------- # Prepare command # --------------------------------------------- @@ -425,7 +438,7 @@ def _run( if machineSettings['gpus_per_node'] == 0: max_cores_per_node_compare = max_cores_per_node else: - print(f"\t - Detected {machineSettings['gpus_per_node']} GPUs in machine, using this value as maximum for non-arra execution (vs {max_cores_per_node} specified)",typeMsg="i") + print(f"\t - Detected {machineSettings['gpus_per_node']} GPUs in machine, using this value as maximum for non-array execution (vs {max_cores_per_node} specified)",typeMsg="i") max_cores_per_node_compare = machineSettings['gpus_per_node'] if not (launchSlurm and ("partition" in self.simulation_job.machineSettings["slurm"])): @@ -1180,8 +1193,12 @@ def buildDictFromInput(inputFile): if "=" in line: splits = [i.split()[0] for i in line.split("=")] if ("." in splits[1]) and (splits[1][0].split()[0] != "."): - parsed[splits[0].split()[0]] = float(splits[1].split()[0]) - else: + try: + parsed[splits[0].split()[0]] = float(splits[1].split()[0]) + continue + except: + pass + try: parsed[splits[0].split()[0]] = int(splits[1].split()[0]) except: diff --git a/src/mitim_tools/simulation_tools/physics/GXtools.py b/src/mitim_tools/simulation_tools/physics/GXtools.py index 85b99927..c407aefe 100644 --- a/src/mitim_tools/simulation_tools/physics/GXtools.py +++ b/src/mitim_tools/simulation_tools/physics/GXtools.py @@ -65,16 +65,60 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call print("\t\t\t GX class module") print("-----------------------------------------------------------------------------------------\n") - self.ResultsFiles = self.ResultsFiles_minimal = [ + self.ResultsFiles_minimal = [ + 'gxplasma.out.nc' + ] + + self.ResultsFiles = self.ResultsFiles_minimal + [ 'gxplasma.eik.out', 'gxplasma.eiknc.nc', 'gxplasma.gx_geo.log', - 'gxplasma.restart.nc', 'gxplasma.big.nc', - 'gxplasma.out.nc', - 'gxplasma.mitim.log' + 'gxplasma.mitim.log', + 'gxplasma.restart.nc', ] + ''' + Redefined here so that I handle restart properly and + I can choose numerical setup based on plasma + ''' + def run( + self, + subfolder, + numerics_based_on_plasma = None, # A dictionary with the parameters to match + **kwargs_sim_run + ): + + # ------------------------------------ + # Check about restarts + # ------------------------------------ + + # Assume every template writes a restart file named "gxplasma.restart.nc" + # If extraOptions indicate not to write a restart, remove the file + if not kwargs_sim_run.get('extraOptions', {}).get('save_for_restart', True): + self.ResultsFiles.remove("gxplasma.restart.nc") + print("\t- Not saving restart file") + + # If the name has changed, update the results files list + if kwargs_sim_run.get('extraOptions', {}).get('restart_to_file', None) is not None: + restart_name = kwargs_sim_run['extraOptions']['restart_to_file'] + self.ResultsFiles.remove("gxplasma.restart.nc") + self.ResultsFiles.append(restart_name) + print(f"\t- Saving restart file as {restart_name}") + + # ------------------------------------ + # Add numerical setup based on plasma + # ------------------------------------ + if numerics_based_on_plasma is not None: + pass + #TODO + + # ------------------------------------ + # Run the super run + # ------------------------------------ + + super().run(subfolder, **kwargs_sim_run) + def plot( self, fn=None, @@ -231,7 +275,7 @@ def write_state(self, file=None): '[Dissipation]': [ ['closure_model', 'hypercollisions', 'nu_hyper_m', 'p_hyper_m', 'nu_hyper_l', 'p_hyper_l', 'hyper', 'D_hyper', 'p_hyper', 'D_H', 'w_osc', 'p_HB', 'HB_hyper'], [] ], '[Restart]': - [ ['restart', 'save_for_restart', 'nsave'], [] ], + [ ['save_for_restart', 'nsave','restart_to_file', 'restart', 'restart_from_file'], [] ], '[Diagnostics]': [ ['nwrite', 'omega', 'fluxes', 'fields', 'moments'], [] ] } diff --git a/src/mitim_tools/simulation_tools/utils/SIMplot.py b/src/mitim_tools/simulation_tools/utils/SIMplot.py index 009f610c..b8830a36 100644 --- a/src/mitim_tools/simulation_tools/utils/SIMplot.py +++ b/src/mitim_tools/simulation_tools/utils/SIMplot.py @@ -1,4 +1,3 @@ -from cProfile import label from mitim_tools.misc_tools import GRAPHICStools from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed diff --git a/templates/input.gx.controls b/templates/input.gx.controls index 8a22478e..d6e06340 100644 --- a/templates/input.gx.controls +++ b/templates/input.gx.controls @@ -51,9 +51,13 @@ p_HB = 2 # exponent for the H-B model [Restart] - restart = false save_for_restart = true - nsave = 100 # save restart file every nsave timesteps + nsave = 100 # save restart file every nsave timesteps +#restart_to_file = "" # Specify restart file name, otherwise “[input_stem].restart.nc” + +restart = false +restart_from_file = "none" # Specify restart file name, otherwise “[input_stem].restart.nc” + [Diagnostics] nwrite = 100 # write diagnostics every nwrite timesteps (this is NOT a/c_s) (=1 doesn't work) From a2de81c4be50b5f45145055a26f25a81024b36b6 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 2 Sep 2025 13:22:48 -0400 Subject: [PATCH 272/385] Bug fix --- src/mitim_tools/simulation_tools/SIMtools.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index bd2e1749..46939872 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -1199,10 +1199,10 @@ def buildDictFromInput(inputFile): except: pass - try: - parsed[splits[0].split()[0]] = int(splits[1].split()[0]) - except: - parsed[splits[0].split()[0]] = splits[1].split()[0] + try: + parsed[splits[0].split()[0]] = int(splits[1].split()[0]) + except: + parsed[splits[0].split()[0]] = splits[1].split()[0] for i in parsed: if isinstance(parsed[i], str): From a0fb68a84d376ebec1f422f8511235ab20861719 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 2 Sep 2025 20:02:16 -0400 Subject: [PATCH 273/385] Common dpdr calculator TGLF GX --- .../plasmastate_tools/MITIMstate.py | 67 ++++++++++++++----- tests/OPT_workflow.py | 2 +- tests/TGLFscan_workflow.py | 2 +- 3 files changed, 54 insertions(+), 17 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 1ee3a76d..9d4ea17e 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -2200,6 +2200,28 @@ def csv(self, file="input.gacode.xlsx"): # Code conversions # ************************************************************************************************************************************************ + def _print_gb_normalizations(self,L_label,Z_label,A_label,n_label,T_label, B_label, L, Z, A): + print(f'\t- GB normalizations, such that Q_gb = n_ref * T_ref^5/2 * m_ref^0.5 / (Z_ref * L_ref * B_ref)^2') + print(f'\t\t* L_ref = {L_label} = {L:.3f}') + print(f'\t\t* Z_ref = {Z_label} = {Z:.3f}') + print(f'\t\t* A_ref = {A_label} = {A:.3f}') + print(f'\t\t* B_ref = {B_label}') + print(f'\t\t* n_ref = {n_label}') + print(f'\t\t* T_ref = {T_label}') + print(f'') + + def _calculate_pressure_gradient_from_aLx(self, pe, pi, aLTe, aLTi, aLne, aLni, a): + ''' + pe and pi in MPa. pi two dimensional + ''' + + adpedr = - pe * 1E6 * (aLTe + aLne) + adpjdr = - pi * 1E6 * (aLTi + aLni) + + dpdr = ( adpedr + adpjdr.sum(axis=-1)) / a + + return dpdr + def to_tglf(self, r=[0.5], code_settings='SAT0', r_is_rho = True): # <> Function to interpolate a curve <> @@ -2219,8 +2241,10 @@ def to_tglf(self, r=[0.5], code_settings='SAT0', r_is_rho = True): else: tglf_ions_num = len(self.Species) - # Determine the mass reference - mass_ref = 2.0 # TODO: This is the only way to make it consistent with TGYRO (derivations with mD_u but mass in tglf with 2.0... https://github.com/gafusion/gacode/issues/398 + # Determine the mass reference for TGLF (use 2.0 for D-mass normalization; derivatives use mD_u elsewhere) + mass_ref = 2.0 + + self._print_gb_normalizations('a', 'Z_D', 'A_D', 'n_e', 'T_e', 'B_unit', self.derived["a"], 1.0, mass_ref) # ----------------------------------------------------------------------- # Derived profiles @@ -2238,11 +2262,13 @@ def to_tglf(self, r=[0.5], code_settings='SAT0', r_is_rho = True): -------------------------------------------------------- Recompute pprime with those species that belong to this run #TODO not exact? ''' - - adpedr = - self.derived['pe'] * (self.derived['aLTe'] + self.derived['aLne']) - adpjdr = - self.derived['pi_all'][:,:tglf_ions_num] * (self.derived['aLTi'][:,:tglf_ions_num] + self.derived['aLni'][:,:tglf_ions_num]) - - dpdr = ( adpedr + adpjdr.sum(axis=-1)) / self.derived['a'] * 1E6 + + dpdr = self._calculate_pressure_gradient_from_aLx( + self.derived['pe'], self.derived['pi_all'][:,:tglf_ions_num], + self.derived['aLTe'], self.derived['aLTi'][:,:tglf_ions_num], + self.derived['aLne'], self.derived['aLni'][:,:tglf_ions_num], + self.derived['a'] + ) pprime = 1E-7 * abs(self.profiles["q(-)"])*self.derived['a']**2/self.derived["r"]/self.derived["B_unit"]**2*dpdr pprime[0] = 0 # infinite in first location @@ -2292,7 +2318,7 @@ def interpolator(y): species = { 1: { 'ZS': -1.0, - 'MASS': 0.000272445, + 'MASS': PLASMAtools.me_u / mass_ref, 'RLNS': interpolator(self.derived['aLne']), 'RLTS': interpolator(self.derived['aLTe']), 'TAUS': 1.0, @@ -2428,6 +2454,8 @@ def to_neo(self, r=[0.5], r_is_rho = True, code_settings='Sonic'): # NEO definition: 'OMEGA_ROT_DERIV=',-gamma_p_loc*a/cs_loc/rmaj_loc omega_rot_deriv = gamma_p * a / cs / rmaj # Equivalent to: self._deriv_gacode(self.profiles["w0(rad/s)"])/ self.derived['c_s'] * self.derived['a']**2 + self._print_gb_normalizations('a', 'Z_D', 'A_D', 'n_e', 'T_e', 'B_unit', self.derived["a"], 1.0, mass_ref) + input_parameters = {} for rho in r: @@ -2582,6 +2610,8 @@ def to_cgyro(self, r=[0.5], r_is_rho = True, code_settings = 'Linear'): # Because in MITIMstate I keep Bunit always positive, but CGYRO routines may need it negative? #TODO sign_Bunit = np.sign(self.profiles['torfluxa(Wb/radian)'][0]) + self._print_gb_normalizations('a', 'Z_D', 'A_D', 'n_e', 'T_e', 'B_unit', self.derived["a"], 1.0, mass_ref) + input_parameters = {} for rho in r: @@ -2709,17 +2739,20 @@ def to_gx(self, r=[0.5], r_is_rho = True, code_settings = 'Linear'): # Determine the mass reference mass_ref = 2.0 - - import scipy.constants as const - p_normalized = self.profiles['ptot(Pa)'] / (8*np.pi) * (2*const.mu_0) - betaprim = -(8*np.pi / self.derived['B_unit']**2) * np.gradient(p_normalized, self.derived['roa']) + dpdr = self._calculate_pressure_gradient_from_aLx( + self.derived['pe'], self.derived['pi_all'][:,:], + self.derived['aLTe'], self.derived['aLTi'][:,:], + self.derived['aLne'], self.derived['aLni'][:,:], + self.derived['a'] + ) + betaprim = -(8*np.pi*1E-7) * self.derived['a'] / self.derived['B_unit']**2 * dpdr #TODO #to check s_kappa = self.derived["r"] / self.profiles["kappa(-)"] * self._deriv_gacode(self.profiles["kappa(-)"]) s_delta = self.derived["r"] * self._deriv_gacode(self.profiles["delta(-)"]) - + self._print_gb_normalizations('a', 'Z_D', 'A_D', 'n_e', 'T_e', 'B_unit', self.derived["a"], 1.0, mass_ref) input_parameters = {} for rho in r: @@ -2746,7 +2779,11 @@ def interpolator(y): # Ions for i in range(len(self.Species)): - nu_ii = self.derived['xnue'] * (self.Species[i]['Z']/self.profiles['ze'][0])**4 * (self.profiles['ni(10^19/m^3)'][:,0]/self.profiles['ne(10^19/m^3)']) * (self.profiles['mass'][i]/self.profiles['masse'][0])**-0.5 * (self.profiles['ti(keV)'][:,0]/self.profiles['te(keV)'])**-1.5 + nu_ii = self.derived['xnue'] * \ + (self.Species[i]['Z']/self.profiles['ze'][0])**4 * \ + (self.profiles['ni(10^19/m^3)'][:,0]/self.profiles['ne(10^19/m^3)']) * \ + (self.profiles['mass'][i]/self.profiles['masse'][0])**-0.5 * \ + (self.profiles['ti(keV)'][:,0]/self.profiles['te(keV)'])**-1.5 species[i+1] = { 'z': self.Species[i]['Z'], @@ -3055,4 +3092,4 @@ def impurity_location(profiles, impurity_of_interest): if position_of_impurity is None: raise ValueError(f"[MITIM] Species {impurity_of_interest} not found in profiles") - return position_of_impurity \ No newline at end of file + return position_of_impurity diff --git a/tests/OPT_workflow.py b/tests/OPT_workflow.py index ab77e64d..f6118a9d 100644 --- a/tests/OPT_workflow.py +++ b/tests/OPT_workflow.py @@ -62,7 +62,7 @@ def scalarized_objective(self, Y): # ----- Inputs # ----------------------------------------------------------------------------------------------------- -namelist = __mitimroot__ / "templates" / "namelist.optimization.yamln" +namelist = __mitimroot__ / "templates" / "namelist.optimization.yaml" folderWork = __mitimroot__ / "tests" / "scratch" / "opt_test" if cold_start and os.path.exists(folderWork): diff --git a/tests/TGLFscan_workflow.py b/tests/TGLFscan_workflow.py index f41c044c..1b991ffc 100644 --- a/tests/TGLFscan_workflow.py +++ b/tests/TGLFscan_workflow.py @@ -1,6 +1,6 @@ import os import numpy as np -from mitim_tools.gacode_tools import TGLFtools, PROFILEStools +from mitim_tools.gacode_tools import TGLFtools from mitim_tools import __mitimroot__ cold_start = True From f141b5bc17bc286c894e135d6c945e613b951a7e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 3 Sep 2025 11:14:34 -0400 Subject: [PATCH 274/385] bug fix -> fine targets --- src/mitim_modules/portals/PORTALSmain.py | 10 ++---- src/mitim_modules/powertorch/STATEtools.py | 21 ++++++----- .../powertorch/scripts/calculateTargets.py | 6 ++-- .../powertorch/utils/TRANSFORMtools.py | 2 +- .../plasmastate_tools/MITIMstate.py | 35 ++++++++++--------- 5 files changed, 38 insertions(+), 36 deletions(-) diff --git a/src/mitim_modules/portals/PORTALSmain.py b/src/mitim_modules/portals/PORTALSmain.py index feae2e92..d5cbee3f 100644 --- a/src/mitim_modules/portals/PORTALSmain.py +++ b/src/mitim_modules/portals/PORTALSmain.py @@ -216,14 +216,10 @@ def run(self, paramsfile, resultsfile): # Extra operations: Store data that will be useful to store and interpret in a machine were this was not run if self.optimization_extra is not None: - dictStore = IOtools.unpickle_mitim(self.optimization_extra) #TODO: This will fail in future versions of torch + dictStore = IOtools.unpickle_mitim(self.optimization_extra) #TODO: This will fail in future versions of torch dictStore[int(numPORTALS)] = {"powerstate": powerstate} - dictStore["profiles_modified"] = PROFILEStools.gacode_state( - self.folder / "Initialization" / "input.gacode_modified" - ) - dictStore["profiles_original"] = PROFILEStools.gacode_state( - self.folder / "Initialization" / "input.gacode_original" - ) + dictStore["profiles_modified"] = PROFILEStools.gacode_state(self.folder / "Initialization" / "input.gacode_modified") + dictStore["profiles_original"] = PROFILEStools.gacode_state( self.folder / "Initialization" / "input.gacode_original") with open(self.optimization_extra, "wb") as handle: pickle_dill.dump(dictStore, handle, protocol=4) diff --git a/src/mitim_modules/powertorch/STATEtools.py b/src/mitim_modules/powertorch/STATEtools.py index 6ba01348..02c43345 100644 --- a/src/mitim_modules/powertorch/STATEtools.py +++ b/src/mitim_modules/powertorch/STATEtools.py @@ -52,18 +52,20 @@ def __init__( } transport_options.setdefault("evaluator_instance_attributes", {}) + transport_options.setdefault("cold_start", False) + # Target options defaults -------------------- if target_options is None: - target_options = { - "evaluator": targets_analytic.analytical_model, - "options": { - "targets_evolve": ["qie", "qrad", "qfus"], - "target_evaluator_method": "powerstate", - }, - } - + target_options = {} + if "options" not in target_options: + target_options["options"] = {} + + target_options.setdefault("evaluator", targets_analytic.analytical_model) + target_options["options"].setdefault("target_evaluator_method", "powerstate") + target_options["options"].setdefault("targets_evolve", ["qie", "qrad", "qfus"]) target_options["options"].setdefault("force_zero_particle_flux", False) target_options["options"].setdefault("percent_error", 1.0) + # --------------------------------------------- if tensor_options is None: tensor_options = { @@ -84,10 +86,11 @@ def __init__( self.predicted_channels = evolution_options.get("ProfilePredicted", ["te", "ti", "ne"]) self.impurityPosition = evolution_options.get("impurityPosition", 1) self.impurityPosition_transport = copy.deepcopy(self.impurityPosition) - self.targets_resolution = target_options.get("targets_resolution", None) self.scaleIonDensities = evolution_options.get("scaleIonDensities", True) self.fImp_orig = evolution_options.get("fImp_orig", 1.0) rho_vec = evolution_options.get("rhoPredicted", [0.2, 0.4, 0.6, 0.8]) + + self.targets_resolution = target_options["options"].get("targets_resolution", None) if rho_vec[0] == 0: raise ValueError("[MITIM] The radial grid must not contain the initial zero") diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index ba869513..b5bbae89 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -1,5 +1,5 @@ """ -calculateTargets.py input.gacode +calculateTargets.py input.gacode run1 """ import sys @@ -7,7 +7,7 @@ from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.powertorch import STATEtools -from mitim_modules.powertorch.physics_models import targets_analytic, transport_tglfneo +from mitim_modules.powertorch.physics_models import targets_analytic from IPython import embed def calculator( @@ -15,7 +15,7 @@ def calculator( targets_evolve=["qie", "qrad", "qfus"], folder="~/scratch/", cold_start=True, - rho_vec=np.linspace(0.1, 0.9, 9), + rho_vec=np.linspace(0.01, 0.94, 50), profProvided=False, targets_resolution = None, ): diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 1b3fdea4..5f78f95b 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -353,7 +353,7 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): profiles.derive_quantities(rederiveGeometry=False) - print("\t- Insering powers") + print("\t- Inserting powers") state_temp = self.copy_state() diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 9d4ea17e..173fecdc 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1738,7 +1738,7 @@ def selfconsistentPTOT(self): print(f"\t\t* Recomputing ptot and inserting it as ptot(Pa), changed from p0 = {self.profiles['ptot(Pa)'][0] * 1e-3:.1f} to {self.derived['ptot_manual'][0]*1e+3:.1f} kPa",typeMsg="i") self.profiles["ptot(Pa)"] = self.derived["ptot_manual"] * 1e6 - def enforceQuasineutrality(self): + def enforceQuasineutrality(self, using_ion = None): print(f"\t\t- Enforcing quasineutrality (error = {self.derived['QN_Error']:.1e})",typeMsg="i",) # What's the lack of quasineutrality? @@ -1748,23 +1748,26 @@ def enforceQuasineutrality(self): ne_missing = self.profiles["ne(10^19/m^3)"] - ni # What ion to modify? - if self.DTplasmaBool: - print("\t\t\t* Enforcing quasineutrality by modifying D and T equally") - prev_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Dion]) - self.profiles["ni(10^19/m^3)"][:, self.Dion] += ne_missing / 2 - self.profiles["ni(10^19/m^3)"][:, self.Tion] += ne_missing / 2 - new_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Dion]) + if using_ion is None: + if self.DTplasmaBool: + print("\t\t\t* Enforcing quasineutrality by modifying D and T equally") + prev_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Dion]) + self.profiles["ni(10^19/m^3)"][:, self.Dion] += ne_missing / 2 + self.profiles["ni(10^19/m^3)"][:, self.Tion] += ne_missing / 2 + new_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Dion]) + else: + print(f"\t\t\t* Enforcing quasineutrality by modifying main ion (position #{self.Mion})") + prev_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Mion]) + self.profiles["ni(10^19/m^3)"][:, self.Mion] += ne_missing + new_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Mion]) else: - print( - f"\t\t\t* Enforcing quasineutrality by modifying main ion (position #{self.Mion})" - ) - prev_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Mion]) - self.profiles["ni(10^19/m^3)"][:, self.Mion] += ne_missing - new_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, self.Mion]) + print(f"\t\t\t* Enforcing quasineutrality by modifying ion (position #{using_ion})") + prev_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, using_ion]) + self.profiles["ni(10^19/m^3)"][:, using_ion] += ne_missing + new_on_axis = copy.deepcopy(self.profiles["ni(10^19/m^3)"][0, using_ion]) - print( - f"\t\t\t\t- Changed on-axis density from n0 = {prev_on_axis:.2f} to {new_on_axis:.2f} ({100*(new_on_axis-prev_on_axis)/prev_on_axis:.1f}%)" - ) + + print(f"\t\t\t\t- Changed on-axis density from n0 = {prev_on_axis:.2f} to {new_on_axis:.2f} ({100*(new_on_axis-prev_on_axis)/prev_on_axis:.1f}%)") self.derive_quantities(rederiveGeometry=False) From 4e3ef6b33a80660761db16c28dfd8590db28f268 Mon Sep 17 00:00:00 2001 From: audreysa Date: Wed, 3 Sep 2025 17:23:30 -0400 Subject: [PATCH 275/385] Add job array limit option to EPEDtools.py --- src/mitim_tools/eped_tools/EPEDtools.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index fb087025..0d6d08b2 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -36,6 +36,7 @@ def run( nproc_per_run = 64, minutes_slurm = 30, cold_start = False, + job_array_limit = 5, ): # ------------------------------------ @@ -64,7 +65,8 @@ def run( 'name': 'mitim_eped', 'minutes': minutes_slurm, 'ntasks': nproc_per_run, - 'job_array': job_array + 'job_array': job_array, + 'job_array_limit': job_array_limit, } ) From ceac4b32ecb62e71064016653c2456aa4723f7d6 Mon Sep 17 00:00:00 2001 From: audreysa Date: Wed, 3 Sep 2025 17:30:45 -0400 Subject: [PATCH 276/385] Add ob_array_limit to EPED_workflow.py as self documentation --- tests/EPED_workflow.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py index 93aea967..944af578 100644 --- a/tests/EPED_workflow.py +++ b/tests/EPED_workflow.py @@ -31,6 +31,7 @@ keep_nsep_ratio = 0.4, nproc_per_run = 64, cold_start = cold_start, + job_array_limit=5, ) eped.read(subfolder='case1') From 7f05b1b6ad514dcd6bd91309e377ebbc6b5fef6c Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Thu, 4 Sep 2025 11:47:38 -0400 Subject: [PATCH 277/385] linear checking for cgyro analysis --- src/mitim_tools/gacode_tools/utils/CGYROutils.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 96100ca4..3435994c 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -81,10 +81,13 @@ def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for # -------------------------------------------------------------- # Check for linear run - if 'phib' in self.cgyrodata.__dict__ and last_tmin_for_linear: - print('\t- Forcing tmin to the last time point because this is a linear run', typeMsg='i') - self.tmin = self.cgyrodata.t[-1] + if 'phib' in self.cgyrodata.__dict__: + print('\t- This is a linear run', typeMsg='i') self.linear = True + if last_tmin_for_linear: + print('\t- Forcing tmin to the last time point', typeMsg='i') + self.tmin = self.cgyrodata.t[-1] + else: self.linear = False @@ -128,7 +131,7 @@ def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for self._process_linear() - if not minimal: # or not self.linear: + if not minimal and self.linear == False: self.cgyrodata.getbigfield() if 'kxky_phi' in self.cgyrodata.__dict__: From 617f0ee5540fe84a19ec370e079fdbae5418da36 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 4 Sep 2025 17:25:06 -0400 Subject: [PATCH 278/385] misc --- src/mitim_modules/maestro/scripts/run_maestro.py | 15 ++++++++------- src/mitim_tools/misc_tools/FARMINGtools.py | 8 ++------ src/mitim_tools/plasmastate_tools/MITIMstate.py | 12 +++--------- 3 files changed, 13 insertions(+), 22 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index be373266..8b99f7b9 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -1,4 +1,5 @@ import argparse +from pathlib import Path from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.maestro.MAESTROmain import maestro @@ -248,22 +249,22 @@ def run_maestro_local( def main(): parser = argparse.ArgumentParser(description='Parse MAESTRO namelist') parser.add_argument('folder', type=str, help='Folder to run MAESTRO') - parser.add_argument('file_path', type=str, help='Path to MAESTRO namelist file') - parser.add_argument('cpus', type=int, help='Number of CPUs to use') + parser.add_argument("--namelist", type=str, required=False, default=None) # namelist.maestro.yaml file, otherwise what's in the current folder + parser.add_argument('--cpus', type=int, required=False, default=8, help='Number of CPUs to use') parser.add_argument('--terminal', action='store_true', help='Print terminal outputs') args = parser.parse_args() + folder = IOtools.expandPath(args.folder) - file_path = args.file_path + maestro_namelist = args.namelist cpus = args.cpus terminal_outputs = args.terminal + maestro_namelist = Path(maestro_namelist) if maestro_namelist is not None else IOtools.expandPath('.') / "namelist.maestro.yaml" + if not folder.exists(): folder.mkdir(parents=True, exist_ok=True) - if (folder / 'namelist.maestro.yaml').exists(): - IOtools.recursive_backup(folder / 'namelist.maestro.yaml') - - run_maestro_local(*parse_maestro_nml(file_path),folder=folder,cpus = cpus, terminal_outputs = terminal_outputs) + run_maestro_local(*parse_maestro_nml(maestro_namelist),folder=folder,cpus = cpus, terminal_outputs = terminal_outputs) if __name__ == "__main__": diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 503fc002..7f546924 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -798,9 +798,7 @@ def interpret_status(self, file_output = "slurm_output.dat"): else: self.infoSLURM = {} for i in range(len(output_squeue[0].split())): - self.infoSLURM[output_squeue[0].split()[i]] = output_squeue[1].split()[ - i - ] + self.infoSLURM[output_squeue[0].split()[i]] = output_squeue[1].split()[i] self.jobid_found = self.infoSLURM["JOBID"] @@ -810,9 +808,7 @@ def interpret_status(self, file_output = "slurm_output.dat"): if self.infoSLURM["STATE"] == "PENDING": self.status = 0 - elif (self.infoSLURM["STATE"] == "RUNNING") or ( - self.infoSLURM["STATE"] == "COMPLETING" - ): + elif (self.infoSLURM["STATE"] == "RUNNING") or (self.infoSLURM["STATE"] == "COMPLETING"): self.status = 1 elif self.infoSLURM["STATE"] == "NOT FOUND": self.status = 2 diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 173fecdc..49e9c2c5 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -473,17 +473,11 @@ def derive_quantities_full(self, mi_ref=None, rederiveGeometry=True): self.derived["QiQe"] = self.derived["qi_MWm2"] / np.where(self.derived["qe_MWm2"] == 0, 1e-10, self.derived["qe_MWm2"]) # to avoid division by zero # "Convective" flux - self.derived["ce_MW"] = PLASMAtools.convective_flux( - self.profiles["te(keV)"], self.derived["ge_10E20"] - ) - self.derived["ce_MWm2"] = PLASMAtools.convective_flux( - self.profiles["te(keV)"], self.derived["ge_10E20m2"] - ) + self.derived["ce_MW"] = PLASMAtools.convective_flux(self.profiles["te(keV)"], self.derived["ge_10E20"]) + self.derived["ce_MWm2"] = PLASMAtools.convective_flux(self.profiles["te(keV)"], self.derived["ge_10E20m2"]) # qmom - self.derived["mt_Jmiller"] = CALCtools.volume_integration( - self.profiles["qmom(N/m^2)"], r, volp - ) + self.derived["mt_Jmiller"] = CALCtools.volume_integration(self.profiles["qmom(N/m^2)"], r, volp) self.derived["mt_Jm2"] = self.derived["mt_Jmiller"] / (volp) # Extras for plotting in TGYRO for comparison From a883344641c2cf9de95544f03206daf044018955 Mon Sep 17 00:00:00 2001 From: audreysa Date: Fri, 5 Sep 2025 10:46:24 -0400 Subject: [PATCH 279/385] Add ability to use squareness to EPEDtools.py, added pass of removeScratchFolders variable in run() --- src/mitim_tools/eped_tools/EPEDtools.py | 7 ++++--- tests/EPED_workflow.py | 4 ++++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index 0d6d08b2..778a1617 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -25,18 +25,19 @@ def __init__( self.results = {} - self.inputs_potential = ['ip', 'bt', 'r', 'a', 'kappa', 'delta', 'neped', 'betan', 'zeffped', 'nesep', 'tesep'] + self.inputs_potential = ['ip', 'bt', 'r', 'a', 'kappa', 'delta', 'neped', 'betan', 'zeffped', 'nesep', 'tesep', 'zeta'] def run( self, subfolder = 'run1', - input_params = None, # {'ip': 12.0, 'bt': 12.16, 'r': 1.85, 'a': 0.57, 'kappa': 1.9, 'delta': 0.5, 'neped': 30.0, 'betan': 1.0, 'zeffped': 1.5, 'nesep': 10.0, 'tesep': 100.0}, + input_params = None, # {'ip': 12.0, 'bt': 12.16, 'r': 1.85, 'a': 0.57, 'kappa': 1.9, 'delta': 0.5, 'neped': 30.0, 'betan': 1.0, 'zeffped': 1.5, 'nesep': 10.0, 'tesep': 100.0, 'zeta': 0}, scan_param = None, # {'variable': 'neped', 'values': [10.0, 20.0, 30.0]} keep_nsep_ratio = None, # Ratio of neped to nesep nproc_per_run = 64, minutes_slurm = 30, cold_start = False, job_array_limit = 5, + removeScratchFolders = True, #ONLY CHANGE THIS FOR DEBUGGING, if you make this False, your EPED runs will be saved and they are enormous ): # ------------------------------------ @@ -149,7 +150,7 @@ def run( self.eped_job.prep(EPEDcommand,input_folders=folder_cases,output_files=copy.deepcopy(output_files),shellPreCommands=shellPreCommands) # Run the job - self.eped_job.run() + self.eped_job.run(removeScratchFolders=removeScratchFolders) # ------------------------------------- # Postprocessing diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py index 944af578..dc8d6273 100644 --- a/tests/EPED_workflow.py +++ b/tests/EPED_workflow.py @@ -26,12 +26,15 @@ 'zeffped': 1.5, 'nesep': 10.0, 'tesep': 100.0, + # Can add zeta if your implementation of EPED supports it }, scan_param = {'variable': 'neped', 'values': [15.0, 30.0, 45.0, 60.0, 75.0]}, keep_nsep_ratio = 0.4, nproc_per_run = 64, cold_start = cold_start, job_array_limit=5, + removeScratchFolders = True, #ONLY CHANGE THIS FOR DEBUGGING, if you make this False, your EPED runs will be saved and they are enormous + ) eped.read(subfolder='case1') @@ -39,3 +42,4 @@ eped.plot(labels=['case1']) eped.fn.show() + From 4413d5130b2c2ed8710c81d8d7cfee86b1c34540 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 5 Sep 2025 10:51:36 -0400 Subject: [PATCH 280/385] roa can be list or np.array in CGYRO json writer --- .../powertorch/physics_models/transport_cgyro.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index f98688eb..903ea529 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -228,8 +228,8 @@ def write_json_CGYRO(roa, fluxes_mean, fluxes_stds, additional_info = None, file additional_info = {} with open(file, 'w') as f: - - additional_info_extended = additional_info | {'roa': roa.tolist()} + + additional_info_extended = additional_info | {'roa': roa.tolist() if not isinstance(roa, list) else roa} json_dict = { 'fluxes_mean': fluxes_mean, From 146769b6e58ccb912783c7a5a01796c2b2dda5d7 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 5 Sep 2025 15:29:30 -0400 Subject: [PATCH 281/385] misc defaults --- templates/config_user_example.json | 24 ++++++++++++------------ templates/namelist.portals.yaml | 5 +++++ tests/PORTALS_workflow.py | 2 +- 3 files changed, 18 insertions(+), 13 deletions(-) diff --git a/templates/config_user_example.json b/templates/config_user_example.json index 42f9f150..6dd6f325 100644 --- a/templates/config_user_example.json +++ b/templates/config_user_example.json @@ -26,6 +26,18 @@ "cores_per_node": 8, "gpus_per_node": 0 }, + "engaging": { + "machine": "orcd-login001.mit.edu", + "username": "exampleusername", + "slurm": { + "partition": "sched_mit_psfc", + "exclusive": false + }, + "scratch": "/orcd/pool/003/exampleusername/scratch/", + "modules": "", + "cores_per_node": 64, + "gpus_per_node": 0 + }, "perlmutter": { "machine": "perlmutter.nersc.gov", "username": "exampleusername", @@ -42,18 +54,6 @@ "cores_per_node": 32, "gpus_per_node": 0 }, - "engaging": { - "machine": "eofe7.mit.edu", - "username": "exampleusername", - "slurm": { - "partition": "sched_mit_psfc", - "exclusive": true - }, - "scratch": "/nobackup1/exampleusername/", - "modules": "", - "cores_per_node": 32, - "gpus_per_node": 0 - }, "mfews": { "machine": "mfews02.psfc.mit.edu", "username": "exampleusername", diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index 1d9c8df6..e4b4478e 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -284,3 +284,8 @@ optimization_options: # Function to select the GP parameters depending on each channel/radius (superseeds the main optimization namelist) surrogate_selection: "import::mitim_modules.portals.PORTALStools.surrogate_selection_portals" + + strategy_options: + + # Allow excursions from the bounds + AllowedExcursions: [0.0, 0.0] \ No newline at end of file diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index 760938d9..e25b2e7f 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -34,7 +34,7 @@ portals_fun.portals_parameters["solution"]["predicted_rho"] = [0.25, 0.45, 0.65, 0.85] portals_fun.portals_parameters["solution"]["predicted_channels"] = ["te", "ti", "ne", "nZ", 'w0'] portals_fun.portals_parameters["solution"]["trace_impurity"] = 'N' -portals_fun.portals_parameters["transport"]["options"]["tglf"]["run"]["code_settings"] = 2 +portals_fun.portals_parameters["transport"]["options"]["tglf"]["run"]["code_settings"] = "SAT0" # Prepare case to run plasma_state = PROFILEStools.gacode_state(inputgacode) From 655a89b4cce342eef081e88dfe020d3b74e52765 Mon Sep 17 00:00:00 2001 From: audreysa Date: Fri, 5 Sep 2025 16:06:38 -0400 Subject: [PATCH 282/385] Close xr files once done using them to prevent errors - changes suggested but chatGPT 4.1 --- tests/EPED_workflow.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/EPED_workflow.py b/tests/EPED_workflow.py index dc8d6273..66a38662 100644 --- a/tests/EPED_workflow.py +++ b/tests/EPED_workflow.py @@ -5,7 +5,7 @@ cold_start = True -folder = __mitimroot__ / "tests" / "scratch" / "eped_test" +folder = __mitimroot__ / "tests" / "scratch" / "eped_test2" if cold_start and os.path.exists(folder): os.system(f"rm -r {folder}") @@ -26,6 +26,7 @@ 'zeffped': 1.5, 'nesep': 10.0, 'tesep': 100.0, + 'zeta': 0.01 # Can add zeta if your implementation of EPED supports it }, scan_param = {'variable': 'neped', 'values': [15.0, 30.0, 45.0, 60.0, 75.0]}, From d5d64f4a9dc7bbae7d2d55bd63f894aafa2bdb9b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 5 Sep 2025 16:22:32 -0400 Subject: [PATCH 283/385] pygacode not necessary unless CGYRO --- src/mitim_tools/gacode_tools/utils/CGYROutils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 3435994c..60ecf4c0 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -7,8 +7,11 @@ from mitim_tools.misc_tools import IOtools from mitim_tools.simulation_tools import SIMtools from mitim_tools.misc_tools.LOGtools import printMsg as print -from pygacode.cgyro.data_plot import cgyrodata_plot -from pygacode import gacodefuncs +try: + from pygacode.cgyro.data_plot import cgyrodata_plot + from pygacode import gacodefuncs +except ModuleNotFoundError: + print("\t- Could not find pygacode module in this environment. Please install it if you need CGYRO capabilities", typeMsg='w') from IPython import embed import pandas as pd From f26b8a1f1c229418c0b8aa18f828c50ed1ad6cd5 Mon Sep 17 00:00:00 2001 From: audreysa Date: Tue, 9 Sep 2025 14:46:31 -0400 Subject: [PATCH 284/385] Add labels to EPED graphs and close the xr to avoid errors --- src/mitim_tools/eped_tools/EPEDtools.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index 778a1617..0dee1070 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -219,8 +219,8 @@ def read( for output_file in output_files: - data = xr.open_dataset(f'{output_file.resolve()}', engine='netcdf4') - data = postprocess_eped(data, 'G', 0.03) + with xr.open_dataset(f'{output_file.resolve()}', engine='netcdf4') as ds: + data = postprocess_eped(ds, 'G', 0.03) sublabel = output_file.name.split('_')[-1].split('.')[0] @@ -280,12 +280,15 @@ def plot( ax.set_xlabel('neped ($10^{19}m^{-3}$)') ax.set_ylabel('ptop (kPa)') ax.set_ylim(bottom=0) + ax.legend(labels) GRAPHICStools.addDenseAxis(ax) ax = axs[1] + ax.set_xlabel('neped ($10^{19}m^{-3}$)') ax.set_ylabel('wptop (psi_pol)') ax.set_ylim(bottom=0) + ax.legend(labels) GRAPHICStools.addDenseAxis(ax) plt.tight_layout() @@ -433,7 +436,8 @@ def read_eped_file(ipaths): dummy_vars = {k: (['dim_one'], [v]) for k, v in set_inputs['eped_input'].items() if k in invars} data = xr.Dataset(coords=dummy_coords, data_vars=dummy_vars) if ipath.is_file(): - data = xr.open_dataset(f'{ipath.resolve()}', engine='netcdf4') + with xr.open_dataset(f'{ipath.resolve()}', engine='netcdf4') as ds: + data = ds.load() data = postprocess_eped(data, 'G', 0.03) data_arrays.append(data.expand_dims({'filename': [ipath.parent.parent.parent.name]})) From 2ce47d463b11b7191b49a4991d1d19ca6a4aaf57 Mon Sep 17 00:00:00 2001 From: audreysa Date: Wed, 10 Sep 2025 11:49:25 -0400 Subject: [PATCH 285/385] Add ability to pass wait to run_slurm --- src/mitim_tools/opt_tools/scripts/slurm.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/mitim_tools/opt_tools/scripts/slurm.py b/src/mitim_tools/opt_tools/scripts/slurm.py index 2dae0782..15430926 100644 --- a/src/mitim_tools/opt_tools/scripts/slurm.py +++ b/src/mitim_tools/opt_tools/scripts/slurm.py @@ -18,7 +18,8 @@ def run_slurm( machine="local", exclude=None, mem=None, - exclusive=False + exclusive=False, + wait=False, ): folder = IOtools.expandPath(folder) @@ -40,7 +41,6 @@ def run_slurm( folder.mkdir(parents=True, exist_ok=True) command = [venv,script + (f" --seed {seed}" if seed is not None else "")] - nameJob = f"mitim_{folder.name}{extra_name}" _, fileSBATCH, _ = FARMINGtools.create_slurm_execution_files( @@ -55,9 +55,12 @@ def run_slurm( 'cpuspertask': n, 'memory_req_by_job': mem } - ) - command_execution = f"sbatch {fileSBATCH}" + if wait == True: + print('* Waiting for job to complete...') + command_execution = f"sbatch --wait {fileSBATCH}" + else: + command_execution = f"sbatch {fileSBATCH}" if machine == "local": os.system(command_execution) From 7a1378f570d441b75ae19b74bb177b6ad752f7e9 Mon Sep 17 00:00:00 2001 From: audreysa Date: Wed, 10 Sep 2025 11:57:53 -0400 Subject: [PATCH 286/385] Fix bug where all submisions through run_slurm are called mitim_job, slurm_settings dictionary needed key 'name' not 'nameJob' --- src/mitim_tools/opt_tools/scripts/slurm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/opt_tools/scripts/slurm.py b/src/mitim_tools/opt_tools/scripts/slurm.py index 15430926..0683238c 100644 --- a/src/mitim_tools/opt_tools/scripts/slurm.py +++ b/src/mitim_tools/opt_tools/scripts/slurm.py @@ -49,12 +49,13 @@ def run_slurm( folder_local=folder, slurm={"partition": partition, 'exclude': exclude,'exclusive': exclusive}, slurm_settings = { - 'nameJob': nameJob, + 'name': nameJob, 'minutes': int(60 * hours), 'ntasks': 1, 'cpuspertask': n, 'memory_req_by_job': mem } + ) if wait == True: print('* Waiting for job to complete...') From b586b22716444f9ef4bca5c4d4edd536fcb7add9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 11 Sep 2025 09:12:35 +0200 Subject: [PATCH 287/385] powerstate to gacode power flows only works with 0 btach for now --- src/mitim_modules/powertorch/utils/TRANSFORMtools.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py index 5f78f95b..18327bc7 100644 --- a/src/mitim_modules/powertorch/utils/TRANSFORMtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSFORMtools.py @@ -210,6 +210,7 @@ def to_gacode( postprocess_input_gacode={}, insert_highres_powers=False, rederive_profiles=True, + debugPlot=False, ): ''' Notes: @@ -223,6 +224,7 @@ def to_gacode( postprocess_input_gacode=postprocess_input_gacode, insert_highres_powers=insert_highres_powers, rederive=rederive_profiles, + debugPlot=debugPlot, ) # Write input.gacode @@ -332,7 +334,7 @@ def powerstate_to_gacode( # ------------------------------------------------------------------------------------------ if insert_highres_powers: - powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch) + powerstate_to_gacode_powers(self, profiles) # ------------------------------------------------------------------------------------------ # Recalculate and change ptot to make it consistent? @@ -349,7 +351,7 @@ def powerstate_to_gacode( return profiles -def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): +def powerstate_to_gacode_powers(self, profiles): profiles.derive_quantities(rederiveGeometry=False) @@ -396,6 +398,8 @@ def powerstate_to_gacode_powers(self, profiles, position_in_powerstate_batch=0): conversions['qfuse'] = "qfuse(MW/m^3)" conversions['qfusi'] = "qfusi(MW/m^3)" + position_in_powerstate_batch = 0 + for ikey in conversions: if conversions[ikey] in profiles.profiles: profiles.profiles[conversions[ikey]][:-extra_points] = state_temp.plasma[ikey][position_in_powerstate_batch,:].cpu().numpy() From 9ea3fe293d4586273a93ade939f506d66f705f11 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 11 Sep 2025 09:13:12 +0200 Subject: [PATCH 288/385] raise species warning in TGLF only if specie is there --- .../powertorch/physics_models/transport_tglf.py | 6 +++++- src/mitim_modules/powertorch/scripts/calculateTargets.py | 3 ++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index ba2a61eb..03529e00 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -149,7 +149,11 @@ def _raise_warnings(self, tglf, rho_locations, Qi_includes_fast): for i in range(len(tglf.profiles.Species)): gacode_type = tglf.profiles.Species[i]['S'] for rho in rho_locations: - tglf_type = tglf.inputs_files[rho].ions_info[i+2]['type'] + try: + tglf_type = tglf.inputs_files[rho].ions_info[i+2]['type'] + except KeyError: + print(f"\t\t\t* Could not determine ion type from TGLF inputs because ion {i+2} was not there for {rho =}, skipping consistency check", typeMsg="w") + continue if gacode_type[:5] != tglf_type[:5]: print(f"\t- For location {rho=:.2f}, ion specie #{i+1} ({tglf.profiles.Species[i]['N']}) is considered '{gacode_type}' by gacode but '{tglf_type}' by TGLF. Make sure this is consistent with your use case", typeMsg="w") diff --git a/src/mitim_modules/powertorch/scripts/calculateTargets.py b/src/mitim_modules/powertorch/scripts/calculateTargets.py index b5bbae89..d82cd4ce 100644 --- a/src/mitim_modules/powertorch/scripts/calculateTargets.py +++ b/src/mitim_modules/powertorch/scripts/calculateTargets.py @@ -15,6 +15,7 @@ def calculator( targets_evolve=["qie", "qrad", "qfus"], folder="~/scratch/", cold_start=True, + file_name = "input.gacode.new.powerstate", rho_vec=np.linspace(0.01, 0.94, 50), profProvided=False, targets_resolution = None, @@ -70,7 +71,7 @@ def calculator( p.profiles.derive_quantities() p.from_powerstate( - write_input_gacode=folder / "input.gacode.new.powerstate", + write_input_gacode=folder / file_name, position_in_powerstate_batch=0, postprocess_input_gacode={ "Tfast_ratio": False, From c802434e7d75f594a8ca39d9e69fbd2c7c313c3b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 11 Sep 2025 11:55:28 +0200 Subject: [PATCH 289/385] Do not force threads in pytorch tests --- src/mitim_tools/misc_tools/IOtools.py | 33 +++++++++++++++++- src/mitim_tools/opt_tools/STRATEGYtools.py | 13 +++++--- templates/namelist.maestro.yaml | 39 +++++++++++++++------- tests/PORTALS_workflow.py | 3 -- 4 files changed, 68 insertions(+), 20 deletions(-) diff --git a/src/mitim_tools/misc_tools/IOtools.py b/src/mitim_tools/misc_tools/IOtools.py index 4e340cc2..c431b689 100644 --- a/src/mitim_tools/misc_tools/IOtools.py +++ b/src/mitim_tools/misc_tools/IOtools.py @@ -2,7 +2,6 @@ import re import shutil import psutil -import copy from typing import Callable import dill as pickle_dill import pandas as pd @@ -21,6 +20,8 @@ import json import functools import hashlib +import io +from contextlib import redirect_stdout import yaml, importlib from typing import Any, Mapping from collections import OrderedDict @@ -1964,6 +1965,36 @@ def print_machine_info(output_file=None): info_lines.append(f"OpenMP Enabled in PyTorch: {openmp_enabled.is_available() if openmp_enabled else 'N/A'}") info_lines.append(f"MKL Enabled in PyTorch: {mkl_enabled.is_available() if mkl_enabled else 'N/A'}") + for var in ["OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", + "NUMEXPR_NUM_THREADS", "SLURM_CPUS_PER_TASK"]: + info_lines.append(f"{var}: {os.environ.get(var, 'Not set')}") + + f = io.StringIO() + with redirect_stdout(f): + torch.__config__.show() + info_lines.append("\n=== PyTorch Build Config ===") + info_lines.append(f.getvalue()) + + info_lines.append("\n=== Package Versions ===") + for pkg in ["torch", "gpytorch", "botorch"]: + try: + mod = __import__(pkg) + info_lines.append(f"{pkg}: {mod.__version__}") + except Exception: + info_lines.append(f"{pkg}: not available") + + try: + import psutil + proc = psutil.Process() + if hasattr(proc, "cpu_affinity"): + info_lines.append(f"Process affinity (cpus): {proc.cpu_affinity()}") + else: + info_lines.append("CPU affinity not supported on this platform/psutil build") + except ImportError: + info_lines.append("psutil not installed (skipping affinity check)") + + info_lines.append("=============================\n\n") + # Output to screen or file output = '\n'.join(info_lines) if output_file: diff --git a/src/mitim_tools/opt_tools/STRATEGYtools.py b/src/mitim_tools/opt_tools/STRATEGYtools.py index 11e708e8..c45de257 100644 --- a/src/mitim_tools/opt_tools/STRATEGYtools.py +++ b/src/mitim_tools/opt_tools/STRATEGYtools.py @@ -3,6 +3,7 @@ import datetime import array import traceback +from typing import IO import torch from pathlib import Path from collections import OrderedDict @@ -433,10 +434,7 @@ def __init__( # ------------------------------------------------------------------------------------------------- if not onlyInitialize: - print("\n-----------------------------------------------------------------------------------------") - print("\t\t\t BO class module") - print("-----------------------------------------------------------------------------------------\n") - + """ ------------------------------------------------------------------------------ Grab variables @@ -448,6 +446,13 @@ def __init__( # Logger sys.stdout = LOGtools.Logger(logFile=self.folderOutputs / "optimization_log.txt", writeAlsoTerminal=True) + print("\n-----------------------------------------------------------------------------------------") + print("\t\t\t BO class module") + print("-----------------------------------------------------------------------------------------\n") + + # Print machine resources + IOtools.print_machine_info() + # Meta self.numIterations = self.optimization_options["convergence_options"]["maximum_iterations"] self.strategy_options = self.optimization_options["strategy_options"] diff --git a/templates/namelist.maestro.yaml b/templates/namelist.maestro.yaml index 51738d53..e5920e72 100644 --- a/templates/namelist.maestro.yaml +++ b/templates/namelist.maestro.yaml @@ -17,7 +17,7 @@ seed: 0 # Machine parameters machine: - # Engineering parameters + # Engineering parameters (if using values in geqdsk, these must be null, otherwise they will override the geqdsk values if separatrix type is geqdsk) Bt: 12.2 Ip: 8.7 @@ -35,17 +35,26 @@ machine: a: 0.57 delta_sep: 0.57 kappa_sep: 1.97 - n_mxh: 5 # If the separatrix type receives a geqdsk file directly, provide path here geqdsk_file: "" + # For both geqdsk and freegs, number of MXH coefficients to parametrize the boundary + n_mxh: 5 + # Heating parameters heating: + + # Only ICRH heating supported for now type: ICRH + parameters: + + # ICRF input power P_icrh: 11.0 + # Minority species [Z,A] minority: [2, 3] + # Minority fraction to be added to the mix fmini: 0.05 # Simulation assumptions @@ -57,6 +66,9 @@ assumptions: fmain: 0.85 ZW: 50 fW: 1.5e-5 + + # Edge parameters + Tesep_eV: 75.0 # Initialization of profiles initialization: @@ -68,12 +80,9 @@ assumptions: # ne_sep / ne_ped ratio to assume (e.g. for EPED) and construct profiles with nesep_ratio: 0.3 - # Profiles will be initialize such that these parameters are matched as close as possible + # Profiles will be initialized such that these parameters are matched as close as possible BetaN: 1.0 density_peaking: 1.3 - - # Edge parameters - Tesep_eV: 75.0 # ------------------------------------------------------------------------------------------------------------ # MAESTRO workflow parameters @@ -81,19 +90,23 @@ assumptions: maestro: - # Remove intermediate files to avoid heavy MAESTRO simulation folders - keep_all_files: true - # Sequence of beats beats: ["transp_soft", "transp", "eped", "portals", "eped", "portals"] + # Remove intermediate files to avoid heavy MAESTRO simulation folders + keep_all_files: true + # --------------------------------------------------------------------------- # Each individual beat parameters # --------------------------------------------------------------------------- eped_beat: + use_default: false + eped_namelist: + + # Location of EPED NN files (if null, use full EPED) nn_location: $MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-MODEL-SPARC.keras norm_location: $MFEIM_PATH/private_code_mitim/NN_DATA/EPED-NN-SPARC/EPED-NN-NORMALIZATION-SPARC.txt @@ -156,18 +169,20 @@ maestro: thermalize_fast: true quasineutrality: true - # If the width of the pedestal has changed, modify the BC for PORTALS accordingly + # If the width of the pedestal has changed, modify the BC (last in predicted_rX) for PORTALS accordingly change_last_radial_call: true - # Later beats to utilize past beats surrogate data if available + # Later beats to utilize past beats surrogate data if available (e.g. if position has not changed) use_previous_surrogate_data: true - # + # If surrogates exist, use them to find the first training point (and limit to only one point) try_flux_match_only_for_first_point: true # Operations that will be included as part of the profiles_postprocessing_fun transport_preprocessing: + # By default, run PORTALS with all impurities lumped into one species lumpImpurities: true + # By default, enforce that all species have the same density gradient enforce_same_density_gradients: true portals_soft_beat: diff --git a/tests/PORTALS_workflow.py b/tests/PORTALS_workflow.py index e25b2e7f..79f5f296 100644 --- a/tests/PORTALS_workflow.py +++ b/tests/PORTALS_workflow.py @@ -17,9 +17,6 @@ if cold_start and folderWork.exists(): os.system(f"rm -r {folderWork.resolve()}") -# Let's not consume the entire computer resources when running test... limit threads -torch.set_num_threads(8) - # --------------------------------------------------------------------------------------------------------------------- # Optimization Class # --------------------------------------------------------------------------------------------------------------------- From 8726b285b7d80a77370fabfcc8a70129fc4163dc Mon Sep 17 00:00:00 2001 From: Audrey Saltzman Date: Mon, 15 Sep 2025 16:26:47 -0400 Subject: [PATCH 290/385] Improvements to EPED plot function --- src/mitim_tools/eped_tools/EPEDtools.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/src/mitim_tools/eped_tools/EPEDtools.py b/src/mitim_tools/eped_tools/EPEDtools.py index 0dee1070..31248fb5 100644 --- a/src/mitim_tools/eped_tools/EPEDtools.py +++ b/src/mitim_tools/eped_tools/EPEDtools.py @@ -250,6 +250,9 @@ def plot( self, labels = ['run1'], axs = None, + plot_labels = None, + legend_title = None, + legend_location = 'best', ): if axs is None: @@ -259,20 +262,27 @@ def plot( colors = GRAPHICStools.listColors() + + for i,name in enumerate(labels): data = self.results[name] neped, ptop, wtop = [], [], [] - for sublabel in data: + sublabels = data.keys() + try: + sublabels = sorted(sublabels, key=lambda x: int(x.split('run')[1])) + except: + print('\t> Warning: sublabels could not be sorted numerically.', typeMsg='w') + for sublabel in sublabels: neped.append(float(data[sublabel]['neped'])) if 'ptop' in data[sublabel].data_vars: ptop.append(float(data[sublabel]['ptop'])) wtop.append(float(data[sublabel]['wptop'])) else: - ptop.append(0.0) - wtop.append(0.0) - + ptop.append(np.nan) + wtop.append(np.nan) + axs[0].plot(neped,ptop,'-s', c = colors[i], ms = 10) axs[1].plot(neped,wtop,'-s', c = colors[i], ms = 10) @@ -280,7 +290,9 @@ def plot( ax.set_xlabel('neped ($10^{19}m^{-3}$)') ax.set_ylabel('ptop (kPa)') ax.set_ylim(bottom=0) - ax.legend(labels) + plot_labels = plot_labels if plot_labels is not None else labels + ax.legend(plot_labels, loc=legend_location, title =legend_title) + GRAPHICStools.addDenseAxis(ax) ax = axs[1] @@ -288,7 +300,7 @@ def plot( ax.set_xlabel('neped ($10^{19}m^{-3}$)') ax.set_ylabel('wptop (psi_pol)') ax.set_ylim(bottom=0) - ax.legend(labels) + ax.legend(plot_labels, loc=legend_location, title=legend_title) GRAPHICStools.addDenseAxis(ax) plt.tight_layout() From aa06f7bc7f343e0bb9872b4237bf6c9352b6d9e3 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 16 Sep 2025 13:35:03 -0400 Subject: [PATCH 291/385] Changing defaults for no-transport model evaluation, generalizable --- .../powertorch/utils/TRANSPORTtools.py | 155 +++++++++--------- .../misc_tools/utils/remote_tools.py | 12 +- src/mitim_tools/opt_tools/SURROGATEtools.py | 12 +- 3 files changed, 98 insertions(+), 81 deletions(-) diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index aa649ee7..1a2ae3e8 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -41,39 +41,45 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): } ''' - with open(self.folder / file_name, 'w') as f: - - fluxes_mean = {} - fluxes_stds = {} - - for var in ['QeGB', 'QiGB', 'GeGB', 'GZGB', 'MtGB']: - fluxes_mean[var] = self.__dict__[f"{var}_{suffix}"].tolist() - fluxes_stds[var] = self.__dict__[f"{var}_{suffix}_stds"].tolist() - - try: - var = 'QieGB' - fluxes_mean[var] = self.__dict__[f"{var}_{suffix}"].tolist() - fluxes_stds[var] = self.__dict__[f"{var}_{suffix}_stds"].tolist() - except KeyError: - # NEO file may not have it - pass - - json_dict = { - 'fluxes_mean': fluxes_mean, - 'fluxes_stds': fluxes_stds, - 'additional_info': { - 'rho': self.powerstate.plasma["rho"][0, 1:].cpu().numpy().tolist(), - 'roa': self.powerstate.plasma["roa"][0, 1:].cpu().numpy().tolist(), - 'Qgb': self.powerstate.plasma["Qgb"][0, 1:].cpu().numpy().tolist(), - 'aLte': self.powerstate.plasma["aLte"][0, 1:].cpu().numpy().tolist(), - 'aLti': self.powerstate.plasma["aLti"][0, 1:].cpu().numpy().tolist(), - 'aLne': self.powerstate.plasma["aLne"][0, 1:].cpu().numpy().tolist(), + if self.folder.exists(): + + with open(self.folder / file_name, 'w') as f: + + fluxes_mean = {} + fluxes_stds = {} + + for var in ['QeGB', 'QiGB', 'GeGB', 'GZGB', 'MtGB']: + fluxes_mean[var] = self.__dict__[f"{var}_{suffix}"].tolist() + fluxes_stds[var] = self.__dict__[f"{var}_{suffix}_stds"].tolist() + + try: + var = 'QieGB' + fluxes_mean[var] = self.__dict__[f"{var}_{suffix}"].tolist() + fluxes_stds[var] = self.__dict__[f"{var}_{suffix}_stds"].tolist() + except KeyError: + # NEO file may not have it + pass + + json_dict = { + 'fluxes_mean': fluxes_mean, + 'fluxes_stds': fluxes_stds, + 'additional_info': { + 'rho': self.powerstate.plasma["rho"][0, 1:].cpu().numpy().tolist(), + 'roa': self.powerstate.plasma["roa"][0, 1:].cpu().numpy().tolist(), + 'Qgb': self.powerstate.plasma["Qgb"][0, 1:].cpu().numpy().tolist(), + 'aLte': self.powerstate.plasma["aLte"][0, 1:].cpu().numpy().tolist(), + 'aLti': self.powerstate.plasma["aLti"][0, 1:].cpu().numpy().tolist(), + 'aLne': self.powerstate.plasma["aLne"][0, 1:].cpu().numpy().tolist(), + } } - } - json.dump(json_dict, f, indent=4) + json.dump(json_dict, f, indent=4) - print(f"\t* Written JSON with {suffix} information to {self.folder / file_name}") + print(f"\t* Written JSON with {suffix} information to {self.folder / file_name}") + + else: + + print(f"\t* Folder {self.folder} does not exist, cannot write {file_name}", typeMsg='w') class power_transport: @@ -87,31 +93,9 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ self.transport_evaluator_options = self.powerstate.transport_options["options"] self.cold_start = self.powerstate.transport_options["cold_start"] - # Allowed fluxes in powerstate so far - self.quantities = ['QeMWm2', 'QiMWm2', 'Ce', 'CZ', 'MtJm2'] - - # Each flux has a turbulent and neoclassical component - self.variables = [f'{i}_tr_turb' for i in self.quantities] + [f'{i}_tr_neoc' for i in self.quantities] - - # Each flux component has a standard deviation - self.variables += [f'{i}_stds' for i in self.variables] - - # There is also turbulent exchange - self.variables += ['QieMWm3_tr_turb', 'QieMWm3_tr_turb_stds'] - - # And total transport flux - self.variables += [f'{i}_tr' for i in self.quantities] - # Model results is None by default, but can be assigned in evaluate self.model_results = None - # Assign zeros to transport ones if not evaluated - for i in self.variables: - self.powerstate.plasma[i] = self.powerstate.plasma["te"] * 0.0 - - # There is also target components - self.variables += [f'{i}' for i in self.quantities] + [f'{i}_stds' for i in self.quantities] - # ---------------------------------------------------------------------------------------- # labels for plotting # ---------------------------------------------------------------------------------------- @@ -126,27 +110,28 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ def evaluate(self): - # Initialize them as zeros - for var in ['QeGB','QiGB','GeGB','GZGB','MtGB','QieGB']: - for suffix in ['turb', 'neoc']: - for suffix0 in ['', '_stds']: - self.__dict__[f"{var}_{suffix}{suffix0}"] = np.zeros(self.powerstate.plasma['rho'].shape[-1]-1) - # Copy the input.gacode files to the output folder self._profiles_to_store() ''' ****************************************************************************************************** - Evaluate neoclassical and turbulent transport. + Evaluate neoclassical and turbulent transport (*in GB units*). These functions use a hook to write the .json files to communicate the results to powerstate.plasma ****************************************************************************************************** ''' + + # Initialize them as zeros + for var in ['QeGB','QiGB','GeGB','GZGB','MtGB','QieGB']: + for suffix in ['turb', 'neoc']: + for suffix0 in ['', '_stds']: + self.__dict__[f"{var}_{suffix}{suffix0}"] = torch.zeros(self.powerstate.plasma['rho'].shape[-1]-1) + neoclassical = self.evaluate_neoclassical() turbulence = self.evaluate_turbulence() ''' ****************************************************************************************************** - From the json to powerstate.plasma + From the json to powerstate.plasma and GB to real units transformation ****************************************************************************************************** ''' self._populate_from_json(file_name = 'fluxes_turb.json', suffix= 'turb') @@ -160,11 +145,14 @@ def evaluate(self): self._postprocess() def _postprocess(self): + ''' + Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) + Before calling this function, the powerstate.plasma should have the following variables: + 'QeMWm2_tr_X', 'QiMWm2_tr_X', 'Ge1E20m2_tr_X', 'GZ1E20m2_tr_X', 'MtJm2_tr_X', 'QieMWm3_tr_X' + where X = 'turb' or 'neoc' + and also the corresponding _stds versions + ''' - # ------------------------------------------------------------------------------------------------------------------------ - # Curate information for the powerstate (e.g. add models, add batch dimension, rho=0.0, and tensorize) - # ------------------------------------------------------------------------------------------------------------------------ - variables = ['QeMWm2', 'QiMWm2', 'Ge1E20m2', 'GZ1E20m2', 'MtJm2', 'QieMWm3'] for variable in variables: @@ -279,11 +267,39 @@ def _profiles_to_store(self): print("\t- Could not move files", typeMsg="w") def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): - ''' Populate the powerstate.plasma with the results from the json file ''' + mapper = { + 'QeGB': ['Qgb', 'QeMWm2'], + 'QiGB': ['Qgb', 'QiMWm2'], + 'GeGB': ['Ggb', 'Ge1E20m2'], + 'GZGB': ['Ggb', 'GZ1E20m2'], + 'MtGB': ['Pgb', 'MtJm2'], + 'QieGB': ['Sgb', 'QieMWm3'] + } + + ''' + ********************************************************************************************** + If no population file exists, I only convert from GB to real units and return + ********************************************************************************************** + ''' + if not (self.folder / file_name).exists(): + print(f"\t* File {self.folder / file_name} does not exist, cannot populate powerstate.plasma", typeMsg='w') + print(f"\t- Tranforming from GB to real units:") + + for var in mapper: + self.powerstate.plasma[f"{mapper[var][1]}_tr_{suffix}"] = self.__dict__[f"{var}_{suffix}"] * self.powerstate.plasma[f"{mapper[var][0]}"][0,1:] + self.powerstate.plasma[f"{mapper[var][1]}_tr_{suffix}_stds"] = self.__dict__[f"{var}_{suffix}_stds"] * self.powerstate.plasma[f"{mapper[var][0]}"][0,1:] + + return + + ''' + ********************************************************************************************** + Populate the powerstate.plasma from the json file + ********************************************************************************************** + ''' print(f"\t* Populating powerstate.plasma with JSON data from {self.folder / file_name}") with open(self.folder / file_name, 'r') as f: @@ -308,15 +324,6 @@ def _populate_from_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): elif units == 'GB' or units == 'both': - mapper = { - 'QeGB': ['Qgb', 'QeMWm2'], - 'QiGB': ['Qgb', 'QiMWm2'], - 'GeGB': ['Ggb', 'Ge1E20m2'], - 'GZGB': ['Ggb', 'GZ1E20m2'], - 'MtGB': ['Pgb', 'MtJm2'], - 'QieGB': ['Sgb', 'QieMWm3'] - } - dum = {} for var in mapper: gb = self.powerstate.plasma[f"{mapper[var][0]}"][0,1:].cpu().numpy() diff --git a/src/mitim_tools/misc_tools/utils/remote_tools.py b/src/mitim_tools/misc_tools/utils/remote_tools.py index 1e183044..45179651 100644 --- a/src/mitim_tools/misc_tools/utils/remote_tools.py +++ b/src/mitim_tools/misc_tools/utils/remote_tools.py @@ -1,5 +1,5 @@ -import os -from mitim_tools.misc_tools import IOtools, FARMINGtools +import os, shutil +from mitim_tools.misc_tools import IOtools, FARMINGtools, CONFIGread from IPython import embed def retrieve_remote_folders(folders_local, remote, remote_folder_parent, remote_folders, only_folder_structure_with_files): @@ -15,10 +15,13 @@ def retrieve_remote_folders(folders_local, remote, remote_folder_parent, remote_ folders_remote = folders_local # Retrieve remote + s = CONFIGread.load_settings() + scratch_local_folder = s['local']['scratch'] + if remote is not None: _, folders = FARMINGtools.retrieve_files_from_remote( - IOtools.expandPath('./'), + scratch_local_folder, remote, folders_remote = folders_remote, purge_tmp_files = True, @@ -35,7 +38,8 @@ def retrieve_remote_folders(folders_local, remote, remote_folder_parent, remote_ if folder_orig.exists(): IOtools.shutil_rmtree(folder_orig) - os.rename(folder, folder_orig) + shutil.copytree(folder, folder_orig) + IOtools.shutil_rmtree(folder) return folders_local \ No newline at end of file diff --git a/src/mitim_tools/opt_tools/SURROGATEtools.py b/src/mitim_tools/opt_tools/SURROGATEtools.py index 38d5fdfd..5e60a6de 100644 --- a/src/mitim_tools/opt_tools/SURROGATEtools.py +++ b/src/mitim_tools/opt_tools/SURROGATEtools.py @@ -436,13 +436,19 @@ def predict(self, X, produceFundamental=False, nSamples=None): - Samples if nSamples not None """ - # Fast - # with gpytorch.settings.fast_computations(), gpytorch.settings.fast_pred_samples(), \ - # gpytorch.settings.fast_pred_var(), gpytorch.settings.lazily_evaluate_kernels(): + # Accurate # with gpytorch.settings.fast_computations(log_prob=False, solves=False, covar_root_decomposition=False), \ # gpytorch.settings.eval_cg_tolerance(1E-6), gpytorch.settings.fast_pred_samples(state=False), gpytorch.settings.num_trace_samples(0): + # # Fast + # with gpytorch.settings.fast_computations(), \ + # gpytorch.settings.fast_pred_samples(), \ + # gpytorch.settings.fast_pred_var(), \ + # gpytorch.settings.lazily_evaluate_kernels(True), \ + # (fundamental_model_context(self) if produceFundamental else contextlib.nullcontext(self)) as surrogate_model: + # posterior = surrogate_model.gpmodel.posterior(X) + with ( fundamental_model_context(self) if produceFundamental From 950d9ee9e26609a22fdf421c52721389889dbe0f Mon Sep 17 00:00:00 2001 From: pabloprf Date: Tue, 16 Sep 2025 16:34:31 -0400 Subject: [PATCH 292/385] Fixed folder removal with proper shutil --- src/mitim_tools/simulation_tools/SIMtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 46939872..99c1b39b 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -671,7 +671,7 @@ def _organize_results(self, code_executor, tmpFolder, filesToRetrieve): print("\t\t- All files were successfully retrieved") # Remove temporary folder - shutil.rmtree(tmpFolder) + IOtools.shutil_rmtree(tmpFolder) else: print("\t\t- Some files were not retrieved", typeMsg="w") From bc55dfb3a187684db233519129d4c6dd48ae0022 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 16 Sep 2025 17:26:32 -0400 Subject: [PATCH 293/385] bug fix surrogate_selection --- templates/namelist.optimization.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/namelist.optimization.yaml b/templates/namelist.optimization.yaml index 45173fe5..83f10e29 100644 --- a/templates/namelist.optimization.yaml +++ b/templates/namelist.optimization.yaml @@ -62,7 +62,7 @@ acquisition_options: surrogate_options: TypeKernel: 0 TypeMean: 0 - selectSurrogate: null + surrogate_selection: null FixedNoise: true ExtraNoise: false additional_constraints: null From 6ccf2aaead4aba31f74b6ec44bf9f37b72fc7dec Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 17 Sep 2025 11:15:59 -0400 Subject: [PATCH 294/385] Corrected bug that read twice the gacode sim folders --- src/mitim_tools/gacode_tools/CGYROtools.py | 6 +++--- src/mitim_tools/simulation_tools/SIMtools.py | 5 +++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 4b49c20a..f7d61c43 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -13,10 +13,10 @@ class CGYRO(SIMtools.mitim_simulation, SIMplot.GKplotting): def __init__( self, - rhos=[0.4, 0.6], # rho locations of interest + **kwargs, ): - - super().__init__(rhos=rhos) + + super().__init__(**kwargs) def code_call(folder, p, n = 1, nomp = 1, additional_command="", **kwargs): return f"cgyro -e {folder} -n {n} -nomp {nomp} -p {p} {additional_command}" diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 99c1b39b..1a4da3c2 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -20,7 +20,7 @@ class mitim_simulation: ''' def __init__( self, - rhos=[0.4, 0.6], # rho locations of interest + rhos=[None], # rho locations of interest, e.g. [0.4,0.6,0.8] ): self.rhos = np.array(rhos) if rhos is not None else None @@ -832,11 +832,12 @@ def read( 'parsed': [], "x": np.array(self.rhos), } + for rho in self.rhos: SIMout = class_output( folder, - suffix=f"_{rho:.4f}" if suffix is None else suffix, + suffix=(f"_{rho:.4f}" if rho is not None else "") if suffix is None else suffix, **kwargs_to_class_output ) From 5442ebb2fea91d663e46b0781d322d673cd8bc08 Mon Sep 17 00:00:00 2001 From: Howard Date: Thu, 18 Sep 2025 09:45:30 -0400 Subject: [PATCH 295/385] Fixed bug CGYRO evaluation in PORTALS --- src/mitim_modules/powertorch/physics_models/transport_cgyro.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 903ea529..85157e19 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -129,7 +129,7 @@ def evaluate_turbulence(self): self._evaluate_tglf() # -------------------------------------------------------------------------------------------- - self._evaluate_gyrokinetic_model(code = 'cgyro', gk_object = CGYROtools.CGYRO, out_name = 'output') + self._evaluate_gyrokinetic_model(code = 'cgyro', gk_object = CGYROtools.CGYRO) def pre_checks(self): From 1b9b683acbbe5c9597247812ce654685787893f9 Mon Sep 17 00:00:00 2001 From: nthoward Date: Thu, 18 Sep 2025 09:50:40 -0400 Subject: [PATCH 296/385] Fixed bug with remote minimal maestro --- src/mitim_modules/maestro/utils/MAESTROplot.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index 11595dd5..7a502caf 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -65,9 +65,11 @@ def plot_results(self, fn): # ******************************************************************************************************** # Collect initialization - ini = {'geqdsk': None, 'profiles': PROFILEStools.gacode_state(f'{self.beats[1].initialize.folder}/input.gacode')} + ini = {'geqdsk': None, 'profiles': None} if (self.beats[1].initialize.folder / 'input.geqdsk').exists(): ini['geqdsk'] = GEQtools.MITIMgeqdsk(self.beats[1].initialize.folder / 'input.geqdsk') + if Path(f'{self.beats[1].initialize.folder}/input.gacode').exists(): + ini['profiles'] = PROFILEStools.gacode_state(f'{self.beats[1].initialize.folder}/input.gacode') # Collect PORTALS profiles and TRANSP cdfs translated to profiles objs = OrderedDict() From 3b69d78de1f1e4534be4f7ddbfa5f59147468f88 Mon Sep 17 00:00:00 2001 From: nthoward Date: Thu, 18 Sep 2025 10:03:57 -0400 Subject: [PATCH 297/385] Fixed bug pathlib remote folders --- src/mitim_tools/misc_tools/utils/remote_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/misc_tools/utils/remote_tools.py b/src/mitim_tools/misc_tools/utils/remote_tools.py index 45179651..a900d440 100644 --- a/src/mitim_tools/misc_tools/utils/remote_tools.py +++ b/src/mitim_tools/misc_tools/utils/remote_tools.py @@ -8,7 +8,7 @@ def retrieve_remote_folders(folders_local, remote, remote_folder_parent, remote_ folders_local = [IOtools.expandPath(folder).resolve() for folder in folders_local] if remote_folder_parent is not None: - folders_remote = [remote_folder_parent + '/' + folder.split('/')[-1] for folder in folders_local] + folders_remote = [remote_folder_parent + '/' + folder.name for folder in folders_local] elif remote_folders is not None: folders_remote = remote_folders else: From 4f1531ac3f556f799b63780379a999e76eedfee5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 18 Sep 2025 11:11:11 -0400 Subject: [PATCH 298/385] Bug fixes linear scan plotting in CGYRO module after SIMtools transition --- src/mitim_tools/gacode_tools/CGYROtools.py | 95 +++++++++---------- .../gacode_tools/scripts/read_cgyro.py | 6 +- .../gacode_tools/utils/CGYROutils.py | 32 +++++-- 3 files changed, 73 insertions(+), 60 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index f7d61c43..26139103 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -1,4 +1,5 @@ import copy +from pathlib import Path import numpy as np import matplotlib.pyplot as plt from mitim_tools import __mitimroot__ @@ -129,6 +130,31 @@ def read( last_tmin_for_linear = last_tmin_for_linear, **kwargs) + def read_linear_scan( + self, + folder=None, + preffix="scan", + **kwargs + ): + ''' + Useful utility for when a folder contains subfolders like... scan0, scan1, scan2... with different ky + ''' + + main_label = kwargs.get('label', 'run1') + del kwargs['label'] + + # Get all folders inside "folder" that start with "preffix" + subfolders = [subfolder for subfolder in Path(folder).glob(f"{preffix}*") if subfolder.is_dir()] + + if len(subfolders) == 0: + print(f"No subfolders found in {folder} with preffix {preffix}. Reading the folder directly.") + self.read(label=f'{main_label}_KY_scan0', folder=folder, **kwargs) + return + + for subfolder in subfolders: + self.read(label=f'{main_label}_KY_{subfolder.name}', folder=folder / subfolder, **kwargs) + + def _labelize(self, labels): # If it has radii, we need to correct the labels @@ -142,17 +168,6 @@ def _labelize(self, labels): labels = labels_with_rho # ------------------------------------------------ - # If it has scans, we need to correct the labels - labels_corrected = [] - for i in range(len(labels)): - if isinstance(self.results[labels[i]], CGYROutils.CGYROlinear_scan): - for scan_label in self.results[labels[i]].labels: - labels_corrected.append(scan_label) - else: - labels_corrected.append(labels[i]) - labels = labels_corrected - # ------------------------------------------------ - return labels def plot( @@ -1469,7 +1484,7 @@ def plot_quick_linear(self, labels=["cgyro1"], fig=None): def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): for cont, label in enumerate(labels): - c = self.results[label] + c = self.results[label]['output'][0] baseColor = colors[cont+start_cont+1] colorsC, _ = GRAPHICStools.colorTableFade( len(c.ky), @@ -1516,19 +1531,12 @@ def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): return cont - - - labels = self._labelize(labels) - - # Make it linear object for nice plotting - labels = self._kyfy(labels) + # Convert the raw labels into the linear scan for nicer plotting + self._kyfy(labels) co = -1 for i,label0 in enumerate(labels): - if isinstance(self.results[label0], CGYROutils.CGYROlinear_scan): - co = _plot_linear_stability(axs, self.results[label0].labels, label0, start_cont=co, col_lin=colors[i]) - else: - co = _plot_linear_stability(axs, [label0], label0, start_cont=co, col_lin=colors[i]) + co = _plot_linear_stability(axs, self.results[label0].labels, label0, start_cont=co, col_lin=colors[i]) ax = axs['1'] ax.set_xlabel("Time $(a/c_s)$") @@ -1544,44 +1552,31 @@ def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): ax.axhline(y=0, lw=0.5, ls="--", c="k") ax.set_xlim(left=0) - def _kyfy(self,labels_original): + def _kyfy(self,labels_new): ''' This function transforms the original labels into the linear scan e.g. from labels: - ['scan1_KY_0.3_0.5', - 'scan1_KY_0.3_0.7', - 'scan1_KY_0.4_0.5', - 'scan1_KY_0.4_0.7'] + ['scan1_KY_0.3', + 'scan1_KY_0.4'] to: - ['scan1_0.5', - 'scan1_0.7'] + ['scan1'] where these are the CGYROlinear_scan object ''' - + + labels_in_results = list(self.results.keys()) labelsD = {} - for label in labels_original: + for label in labels_new: + labelsD[label] = [] + for label in labels_in_results: parts = label.split('_') - if len(parts) >= 4 and parts[1] == "KY": - # Extract the base name (scan1), middle value (0.3/0.4), and last value (0.5/0.7) - base_name = parts[0] - middle_value = float(parts[2]) - last_value = parts[3] - - # Create the new key format: base_name + "_" + last_value - new_key = f"{base_name}_{last_value}" - - # Add the middle value to the list for this key - if new_key not in labelsD: - labelsD[new_key] = [] - labelsD[new_key].append(label) - - labels = [] + if len(parts) >= 3 and parts[1] == "KY": + # Extract the base name (scan1) and middle value (0.3/0.4) + base_name = parts[0] + labelsD[base_name].append(label) + for label in labelsD: - self.results[label] = CGYROutils.CGYROlinear_scan(labelsD[label], self.results - ) - labels.append(label) + self.results[label] = CGYROutils.CGYROlinear_scan(labelsD[label], self.results) - return labels class CGYROinput(SIMtools.GACODEinput): def __init__(self, file=None): diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 779bdb77..79830128 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -34,7 +34,11 @@ def main(): labels = [] for i, folder in enumerate(folders): labels.append(f"case {i + 1}") - c.read(label=labels[-1], folder=folder, tmin=tmin[i], last_tmin_for_linear=last_tmin_for_linear) + + if linear: + c.read_linear_scan(label=labels[-1], folder=folder) + else: + c.read(label=labels[-1], folder=folder, tmin=tmin[i], last_tmin_for_linear=last_tmin_for_linear) if linear: # Plot linear spectrum diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 60ecf4c0..6b4ea701 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -16,7 +16,7 @@ import pandas as pd class CGYROlinear_scan: - def __init__(self, labels, cgyro_data): + def __init__(self, labels, results): self.labels = labels @@ -32,16 +32,16 @@ def __init__(self, labels, cgyro_data): self.Qi_mean = [] for label in labels: - self.ky.append(cgyro_data[label].ky[0]) - self.aLTi.append(cgyro_data[label].aLTi) - self.g_mean.append(cgyro_data[label].g_mean[0]) - self.f_mean.append(cgyro_data[label].f_mean[0]) + self.ky.append(results[label]['output'][0].ky[0]) + self.aLTi.append(results[label]['output'][0].aLTi) + self.g_mean.append(results[label]['output'][0].g_mean[0]) + self.f_mean.append(results[label]['output'][0].f_mean[0]) - self.Qe_mean.append(cgyro_data[label].Qe_mean) - self.Qi_mean.append(cgyro_data[label].Qi_mean) + self.Qe_mean.append(results[label]['output'][0].Qe_mean) + self.Qi_mean.append(results[label]['output'][0].Qi_mean) try: - self.neTe_mean.append(cgyro_data[label].neTe_kx0_mean[0]) + self.neTe_mean.append(results[label]['output'][0].neTe_kx0_mean[0]) except: self.neTe_mean.append(np.nan) @@ -52,6 +52,20 @@ def __init__(self, labels, cgyro_data): self.neTe_mean = np.array(self.neTe_mean) self.Qe_mean = np.array(self.Qe_mean) self.Qi_mean = np.array(self.Qi_mean) + + + # Organize them by ky + order = np.argsort(self.ky) + self.ky = self.ky[order] + self.aLTi = self.aLTi[order] + self.g_mean = self.g_mean[order] + self.f_mean = self.f_mean[order] + self.neTe_mean = self.neTe_mean[order] + self.Qe_mean = self.Qe_mean[order] + self.Qi_mean = self.Qi_mean[order] + self.labels = [self.labels[i] for i in order] + self.results = {label: results[label] for label in self.labels} + class CGYROoutput(SIMtools.GACODEoutput): def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for_linear=True, **kwargs): @@ -134,7 +148,7 @@ def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for self._process_linear() - if not minimal and self.linear == False: + if (not minimal): # and (self.linear == False): self.cgyrodata.getbigfield() if 'kxky_phi' in self.cgyrodata.__dict__: From 3bf707e20e77fb559c6760ba41965c13f9094f8b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Thu, 18 Sep 2025 16:41:55 -0400 Subject: [PATCH 299/385] Some touches to fix cgyro linear scan plotting --- src/mitim_tools/gacode_tools/CGYROtools.py | 78 +++++++++------------- 1 file changed, 30 insertions(+), 48 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 26139103..f8055e27 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -146,17 +146,39 @@ def read_linear_scan( # Get all folders inside "folder" that start with "preffix" subfolders = [subfolder for subfolder in Path(folder).glob(f"{preffix}*") if subfolder.is_dir()] + labels_in_results = [] if len(subfolders) == 0: print(f"No subfolders found in {folder} with preffix {preffix}. Reading the folder directly.") - self.read(label=f'{main_label}_KY_scan0', folder=folder, **kwargs) - return - - for subfolder in subfolders: - self.read(label=f'{main_label}_KY_{subfolder.name}', folder=folder / subfolder, **kwargs) - + labels_in_results.append(f'{main_label}_KY_scan0') + self.read(label=labels_in_results[-1], folder=folder, **kwargs) + else: + for subfolder in subfolders: + labels_in_results.append(f'{main_label}_KY_{subfolder.name}') + self.read(label=labels_in_results[-1], folder=subfolder, **kwargs) + + # ---------------------------------------------------------- + # Make it a linear scan for the main label + # ---------------------------------------------------------- + labelsD = {} + for label in [main_label]: + labelsD[label] = [] + for label in labels_in_results: + parts = label.split('_') + if len(parts) >= 3 and parts[-2] == "KY": + # Extract the base name (scan1) and middle value (0.3/0.4) + base_name = '_'.join(parts[0:-2]) + labelsD[base_name].append(label) - def _labelize(self, labels): + for label in labelsD: + self.results[label] = CGYROutils.CGYROlinear_scan(labelsD[label], self.results) + def plot( + self, + labels=[""], + fn=None, + include_2D=True, + common_colorbar=True): + # If it has radii, we need to correct the labels self.results_all = copy.deepcopy(self.results) self.results = {} @@ -167,17 +189,6 @@ def _labelize(self, labels): self.results[f'{label}_{rho}'] = self.results_all[label]['output'][i] labels = labels_with_rho # ------------------------------------------------ - - return labels - - def plot( - self, - labels=[""], - fn=None, - include_2D=True, - common_colorbar=True): - - labels = self._labelize(labels) if fn is None: from mitim_tools.misc_tools.GUItools import FigureNotebook @@ -1530,10 +1541,7 @@ def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): ) return cont - - # Convert the raw labels into the linear scan for nicer plotting - self._kyfy(labels) - + co = -1 for i,label0 in enumerate(labels): co = _plot_linear_stability(axs, self.results[label0].labels, label0, start_cont=co, col_lin=colors[i]) @@ -1552,32 +1560,6 @@ def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): ax.axhline(y=0, lw=0.5, ls="--", c="k") ax.set_xlim(left=0) - def _kyfy(self,labels_new): - ''' - This function transforms the original labels into the linear scan - e.g. from labels: - ['scan1_KY_0.3', - 'scan1_KY_0.4'] - to: - ['scan1'] - where these are the CGYROlinear_scan object - ''' - - labels_in_results = list(self.results.keys()) - labelsD = {} - for label in labels_new: - labelsD[label] = [] - for label in labels_in_results: - parts = label.split('_') - if len(parts) >= 3 and parts[1] == "KY": - # Extract the base name (scan1) and middle value (0.3/0.4) - base_name = parts[0] - labelsD[base_name].append(label) - - for label in labelsD: - self.results[label] = CGYROutils.CGYROlinear_scan(labelsD[label], self.results) - - class CGYROinput(SIMtools.GACODEinput): def __init__(self, file=None): super().__init__( From 8eb61ffec75cb07cd5c8e0eb5571e66db382c7af Mon Sep 17 00:00:00 2001 From: nthoward Date: Fri, 19 Sep 2025 10:31:46 -0400 Subject: [PATCH 300/385] makes sense to enforce density gradients to fast too by default --- src/mitim_tools/plasmastate_tools/MITIMstate.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 49e9c2c5..7e7ef637 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1705,13 +1705,13 @@ def correct(self, options={}, write=False, new_file=None): self.write_state(file=new_file) self.printInfo() - def enforce_same_density_gradients(self): + def enforce_same_density_gradients(self, onlyThermal=False): txt = "" for sp in range(len(self.Species)): - if self.Species[sp]["S"] == "therm": + if (not onlyThermal) or (self.Species[sp]["S"] == "therm"): self.profiles["ni(10^19/m^3)"][:, sp] = self.derived["fi_vol"][sp] * self.profiles["ne(10^19/m^3)"] txt += f"{self.Species[sp]['N']} " - print(f"\t\t- Making all thermal ions ({txt}) have the same a/Ln as electrons (making them an exact flat fraction)",typeMsg="i",) + print(f"\t\t- Making all {'thermal ' if onlyThermal else ''}ions ({txt}) have the same a/Ln as electrons (making them an exact flat fraction)",typeMsg="i",) self.derive_quantities(rederiveGeometry=False) def make_fast_ions_thermal(self): From 843b995301463b8bf3bff2dfaee2649c825578c5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 19 Sep 2025 14:29:20 -0400 Subject: [PATCH 301/385] To avoid potentially overwritting an original input.gacode, rename --- src/mitim_tools/simulation_tools/SIMtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 1a4da3c2..2cc890fa 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -69,7 +69,7 @@ def prep( self.profiles = mitim_state # Keep a copy of the file - self.profiles.write_state(file=self.FolderGACODE / "input.gacode") + self.profiles.write_state(file=self.FolderGACODE / "input.gacode_torun") self.profiles.derive_quantities(mi_ref=md_u) From 4cacd649c07956b7fa56fc368caecf35784da9c4 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 19 Sep 2025 14:29:33 -0400 Subject: [PATCH 302/385] Pe/Pi to include fusion too --- .../plasmastate_tools/utils/state_plotting.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py index 4c288c7c..77678371 100644 --- a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -803,7 +803,21 @@ def plot_other(self, axs6, color="b", lw=1.0, extralab="", fs=6): c=color, lw=lw, ls="--", - label=extralab + "$P_i/P_e$", + label=extralab + "$P_{aux,i}/P_{aux,e}$", + ) + safe_division = np.divide( + self.derived["qi_aux_MW"]+self.derived['qi_fus_MW'], + self.derived["qe_aux_MW"]+self.derived['qe_fus_MW'], + where=(self.derived["qe_aux_MW"]+self.derived['qe_fus_MW']) != 0, + out=np.full_like(self.derived["qi_aux_MW"], np.nan), + ) + ax.plot( + self.profiles["rho(-)"], + safe_division, + c=color, + lw=lw, + ls="-.", + label=extralab + "$(P_{aux,i}+P_{fus,i})/(P_{aux,e}+P_{fus,e})$", ) ax.set_ylabel("Power ratios") ax.set_xlabel("$\\rho$") From 94abae3a06380c9b2f8a1b0d4ed9d0a83539e5e9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 20 Sep 2025 20:37:15 -0400 Subject: [PATCH 303/385] Capability to only retrieve files (cases with connection lost but simulation finished, and heavy) --- src/mitim_tools/misc_tools/FARMINGtools.py | 34 ++++++++++++++----- .../plasmastate_tools/MITIMstate.py | 5 +-- src/mitim_tools/simulation_tools/SIMtools.py | 12 +++++-- 3 files changed, 37 insertions(+), 14 deletions(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 7f546924..942cefca 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -163,7 +163,13 @@ def run( removeScratchFolders_goingIn=None, check_if_files_received=True, attempts_execution=1, + helper_lostconnection=False, ): + + ''' + if helper_lostconnection is True, it means that the connection to the remote machine was lost, but the files are there, + so I just want to retrieve them. In that case, I do not remove the scratch folder going in, and I do not execute the commands. + ''' removeScratchFolders_goingOut = removeScratchFolders if removeScratchFolders_goingIn is None: @@ -208,12 +214,13 @@ def run( # Process self.full_process( comm, - removeScratchFolders_goingIn=removeScratchFolders_goingIn, + removeScratchFolders_goingIn=removeScratchFolders_goingIn and (not helper_lostconnection), removeScratchFolders_goingOut=removeScratchFolders_goingOut, timeoutSecs=timeoutSecs, check_if_files_received=waitYN and check_if_files_received, check_files_in_folder=self.check_files_in_folder, attempts_execution=attempts_execution, + execute_flag=not helper_lostconnection ) # Get jobid @@ -242,11 +249,17 @@ def full_process( check_if_files_received=True, check_files_in_folder={}, attempts_execution = 1, + execute_flag=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 finish the sbatch launch process. + + Notes: + - If execute_flag is False, the commands will not be executed. This is useful, + together with removeScratchFolders_goingIn=False if the results exist in the remote + but the connection failed with your local machine. You can then just retrieve the results. """ wait_for_all_commands = True @@ -268,13 +281,18 @@ def full_process( execution_counter = 0 while execution_counter < attempts_execution: - output, error = self.execute( - comm, - wait_for_all_commands=wait_for_all_commands, - printYN=True, - timeoutSecs=timeoutSecs if timeoutSecs < 1e6 else None, - log_file=self.log_simulation_file - ) + + if execute_flag: + output, error = self.execute( + comm, + wait_for_all_commands=wait_for_all_commands, + printYN=True, + timeoutSecs=timeoutSecs if timeoutSecs < 1e6 else None, + log_file=self.log_simulation_file + ) + else: + output, error = b"", b"" + print("\t* Not executing commands, just retrieving files (execute_flag=False)", typeMsg="i") # ~~~~~~ Retrieve received = self.retrieve( diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 7e7ef637..652d2ce4 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1569,10 +1569,7 @@ def moveSpecie(self, pos=2, pos_new=1): self.remove([position_to_moveFROM_in_profiles + 2]) def addSpecie(self, Z=5.0, mass=10.0, fi_vol=0.1, forcename=None): - print( - f"\t\t- Creating new specie with Z={Z}, mass={mass}, fi_vol={fi_vol}", - typeMsg="i", - ) + print(f"\t\t- Creating new specie with Z={Z}, mass={mass}, fi_vol={fi_vol}",typeMsg="i",) if forcename is None: forcename = "LUMPED" diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index 2cc890fa..fa249a7f 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -117,6 +117,7 @@ def run( only_minimal_files=False, run_type = 'normal', # 'normal': submit and wait; 'submit': submit and do not wait; 'prep': do not submit additional_files_to_send = None, # Dict (rho keys) of files to send along with the run (e.g. for restart) + helper_lostconnection=False, # If True, it means that the connection to the remote machine was lost, but the files are there, so I just want to retrieve them not execute the commands ): if slurm_setup is None: @@ -167,7 +168,8 @@ def run( slurm_setup=slurm_setup, only_minimal_files=only_minimal_files, attempts_execution=attempts_execution, - run_type=run_type + run_type=run_type, + helper_lostconnection=helper_lostconnection, ) return code_executor_full @@ -554,7 +556,8 @@ def _run( self.simulation_job.run( removeScratchFolders=True, - attempts_execution=attempts_execution + attempts_execution=attempts_execution, + helper_lostconnection=kwargs_run.get("helper_lostconnection", False) ) self._organize_results(code_executor, tmpFolder, filesToRetrieve) @@ -869,6 +872,11 @@ def read_scan( if subfolder is None: subfolder = self.subfolder_scan + + if variable_mapping is None: + variable_mapping = {} + if variable_mapping_unn is None: + variable_mapping_unn = {} self.scans[label] = {} self.scans[label]["variable"] = variable From 7f5f6d6eb46131939bb465b106943b40e6ed5cd5 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 20 Sep 2025 20:53:33 -0400 Subject: [PATCH 304/385] Common way to read CGYRO ky scans --- src/mitim_tools/gacode_tools/CGYROtools.py | 33 ++++++++++++++----- .../gacode_tools/scripts/read_cgyro.py | 30 +++++++++++++++-- 2 files changed, 52 insertions(+), 11 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index f8055e27..523bcf85 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -130,6 +130,20 @@ def read( last_tmin_for_linear = last_tmin_for_linear, **kwargs) + # Re-defined to make allowing reading a scan of KY linear runs easily + def read_scan( + self, + label="scan1", + cgyro_linear_scan = False, + **kwargs + ): + + super().read_scan(label=label,**kwargs) + + if cgyro_linear_scan: + self.results[label] = CGYROutils.CGYROlinear_scan(list(self.results.keys()), self.results) + print(f"\t- Created a linear scan object with label {label} from all the read cases", typeMsg='i') + def read_linear_scan( self, folder=None, @@ -144,7 +158,7 @@ def read_linear_scan( del kwargs['label'] # Get all folders inside "folder" that start with "preffix" - subfolders = [subfolder for subfolder in Path(folder).glob(f"{preffix}*") if subfolder.is_dir()] + subfolders = [subfolder for subfolder in Path(folder).glob(f"*{preffix}*") if subfolder.is_dir()] labels_in_results = [] if len(subfolders) == 0: @@ -159,18 +173,15 @@ def read_linear_scan( # ---------------------------------------------------------- # Make it a linear scan for the main label # ---------------------------------------------------------- - labelsD = {} - for label in [main_label]: - labelsD[label] = [] + labelsD = [] for label in labels_in_results: parts = label.split('_') if len(parts) >= 3 and parts[-2] == "KY": # Extract the base name (scan1) and middle value (0.3/0.4) base_name = '_'.join(parts[0:-2]) - labelsD[base_name].append(label) + labelsD.append(label) - for label in labelsD: - self.results[label] = CGYROutils.CGYROlinear_scan(labelsD[label], self.results) + self.results[main_label] = CGYROutils.CGYROlinear_scan(labelsD, self.results) def plot( self, @@ -1552,13 +1563,19 @@ def _plot_linear_stability(axs, labels, label_base,col_lin ='b', start_cont=0): ax.set_ylabel("$\\gamma$ $(c_s/a)$") ax.set_title("Growth Rate") ax.set_xlim(left=0) - ax.legend() + ax.legend(loc='best', prop={'size': 8},) + ax = axs['2'] ax.set_xlabel("Time $(a/c_s)$") ax.set_ylabel("$\\omega$ $(c_s/a)$") ax.set_title("Real Frequency") ax.axhline(y=0, lw=0.5, ls="--", c="k") ax.set_xlim(left=0) + + for ax in [axs['1'], axs['2'], axs['3'], axs['4']]: + GRAPHICStools.addDenseAxis(ax) + + plt.tight_layout() class CGYROinput(SIMtools.GACODEinput): def __init__(self, file=None): diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 79830128..26f8fd08 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -12,9 +12,11 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("folders", type=str, nargs="*") + parser.add_argument("--suffixes", required=False, type=str, nargs="*", default=None) parser.add_argument("--two", action="store_true", help="Include 2D plots") parser.add_argument("--linear", action="store_true", help="Just a plot of the linear spectra") parser.add_argument("--tmin", type=float, nargs="*", default=None, help="Minimum time to calculate mean and std") + parser.add_argument("--scan_subfolder_id" , type=str, default="KY", help="If reading a linear scan, the subfolders contain this common identifier") args = parser.parse_args() folders = args.folders @@ -22,6 +24,17 @@ def main(): tmin = args.tmin include_2D = args.two + suffixes = args.suffixes + + scan_subfolder_id = args.scan_subfolder_id + + if suffixes is None: + suffixes = ["" for _ in range(len(folders))] + + for i in range(len(suffixes)): + if suffixes[i] == "_": + suffixes[i] = "" + if tmin is None: tmin = [0.0] * len(folders) last_tmin_for_linear = True @@ -36,9 +49,21 @@ def main(): labels.append(f"case {i + 1}") if linear: - c.read_linear_scan(label=labels[-1], folder=folder) + c.read_linear_scan( + label=labels[-1], + folder=folder, + suffix=suffixes[i], + preffix=scan_subfolder_id + ) + else: - c.read(label=labels[-1], folder=folder, tmin=tmin[i], last_tmin_for_linear=last_tmin_for_linear) + c.read( + label=labels[-1], + folder=folder, + tmin=tmin[i], + last_tmin_for_linear=last_tmin_for_linear, + suffix=suffixes[i] + ) if linear: # Plot linear spectrum @@ -47,7 +72,6 @@ def main(): else: c.plot(labels=labels, include_2D=include_2D, common_colorbar=True) c.fn.show() - embed() From 43dcabd00b3e48a0d001c2d8208a69e8cd7c8b95 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 22 Sep 2025 08:49:46 -0400 Subject: [PATCH 305/385] CGYRO standalone reader now takes prefix for scans --- src/mitim_tools/gacode_tools/scripts/read_cgyro.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 26f8fd08..3e2ba014 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -16,7 +16,7 @@ def main(): parser.add_argument("--two", action="store_true", help="Include 2D plots") parser.add_argument("--linear", action="store_true", help="Just a plot of the linear spectra") parser.add_argument("--tmin", type=float, nargs="*", default=None, help="Minimum time to calculate mean and std") - parser.add_argument("--scan_subfolder_id" , type=str, default="KY", help="If reading a linear scan, the subfolders contain this common identifier") + parser.add_argument("--scan_subfolder_id" , type=str, nargs="*", default="KY", help="If reading a linear scan, the subfolders contain this common identifier") args = parser.parse_args() folders = args.folders @@ -53,7 +53,7 @@ def main(): label=labels[-1], folder=folder, suffix=suffixes[i], - preffix=scan_subfolder_id + preffix=scan_subfolder_id[i] ) else: @@ -62,7 +62,8 @@ def main(): folder=folder, tmin=tmin[i], last_tmin_for_linear=last_tmin_for_linear, - suffix=suffixes[i] + suffix=suffixes[i], + preffix=scan_subfolder_id[i] ) if linear: From 1041dbe944f13c605cd6112302baacc58310d954 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 22 Sep 2025 08:50:00 -0400 Subject: [PATCH 306/385] _SCALE_ variables are not recommended as extraOptions --- src/mitim_tools/gacode_tools/CGYROtools.py | 34 ++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index 523bcf85..d030d3b3 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -115,6 +115,40 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call "out.cgyro.rotation", ] + # Redefine to raise warning + def _run_prepare( + self, + subfolder_simulation, + extraOptions=None, + multipliers=None, + **kwargs, + ): + + # --------------------------------------------- + # Check if any *_SCALE_* variable is being used + # --------------------------------------------- + dictionary_check = {} + if extraOptions is not None: + if multipliers is not None: + dictionary_check = {**extraOptions, **multipliers} + else: + dictionary_check = extraOptions + elif multipliers is not None: + dictionary_check = multipliers + + for key in dictionary_check: + if '_SCALE_' in key: + print(f"The use of *_SCALE_* is discouraged, please use the appropriate variable instead.", typeMsg='q') + + # --------------------------------------------- + + return super()._run_prepare( + subfolder_simulation, + extraOptions=extraOptions, + multipliers=multipliers, + **kwargs, + ) + # Re-defined to make specific arguments explicit def read( self, From 1635db047b2adee027132577545c2bb27189ca61 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 22 Sep 2025 11:06:13 -0400 Subject: [PATCH 307/385] When doing PORTALS-CGYRO, also run TGLF depending on namelist flag --- .../powertorch/physics_models/transport_cgyro.py | 11 ++++++----- templates/namelist.portals.yaml | 3 +++ 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index 85157e19..fee17d4f 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -123,11 +123,12 @@ class cgyro_model(gyrokinetic_model): def evaluate_turbulence(self): - # Run base TGLF always, to keep track of discrepancies! -------------------------------------- - simulation_options_tglf = self.transport_evaluator_options["tglf"] - simulation_options_tglf["use_scan_trick_for_stds"] = None - self._evaluate_tglf() - # -------------------------------------------------------------------------------------------- + if self.transport_evaluator_options["cgyro"].get("run_base_tglf", True): + # Run base TGLF, to keep track of discrepancies! --------------------------------------------- + simulation_options_tglf = self.transport_evaluator_options["tglf"] + simulation_options_tglf["use_scan_trick_for_stds"] = None + self._evaluate_tglf() + # -------------------------------------------------------------------------------------------- self._evaluate_gyrokinetic_model(code = 'cgyro', gk_object = CGYROtools.CGYRO) diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index e4b4478e..3e5b0a31 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -174,6 +174,9 @@ transport: # (%) For CGYRO runs, minimum error based on target if case is considered stable Qi_stable_percent_error: 5.0 + # If True, always run base TGLF to keep track of discrepancies + run_base_tglf: True + # ********************************************************************************************************* # GX # ********************************************************************************************************* From 5af4521e1011fb43de4fd7f40f531ffbbaa6c265 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 22 Sep 2025 11:33:41 -0400 Subject: [PATCH 308/385] Fixed bug that prevented to read the optimization_options from the MAESTRO namelist --- .../maestro/utils/PORTALSbeat.py | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 661a4451..c4e6501b 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -25,12 +25,20 @@ def prepare(self, use_previous_surrogate_data = False, try_flux_match_only_for_first_point = True, change_last_radial_call = False, - portals_parameters = {}, portals_namelist_location = None, - initialization_parameters = {}, - optimization_options = {}, + portals_parameters = None, + initialization_parameters = None, + optimization_options = None, enforce_impurity_radiation_existence = True, ): + + if portals_parameters is None: + portals_parameters = {} + if initialization_parameters is None: + initialization_parameters = {} + if optimization_options is None: + optimization_options = {} + self.fileGACODE = self.initialize.folder / 'input.gacode' @@ -78,10 +86,15 @@ def run(self, **kwargs): cold_start = kwargs.get('cold_start', False) + # Read the namelist if explicitly given in the MAESTRO namelist (variable: portals_namelist_location) portals_fun = PORTALSmain.portals(self.folder, portals_namelist = self.portals_namelist_location) + # Update the namelist with the parameters in the MAESTRO namelist (variable: portals_parameters) portals_fun.portals_parameters = IOtools.deep_dict_update(portals_fun.portals_parameters, self.portals_parameters) - portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.optimization_options) + portals_fun.portals_parameters['optimization_options'] = portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.portals_parameters['optimization_options']) + + # MAESTRO beat may receive optimization options changes from previous beats, so allow that too + portals_fun.portals_parameters['optimization_options'] = portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.optimization_options) # Initialization now happens by the user from mitim_tools.gacode_tools.PROFILEStools import gacode_state From f96ca9fb6473e22161e8e0dc6b0856996d6aa083 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 22 Sep 2025 11:35:38 -0400 Subject: [PATCH 309/385] MAESTRO defaults should use 30 portals iterations and allow lower excursions --- templates/namelist.maestro.yaml | 6 ++++++ tests/MAESTRO_workflow.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/templates/namelist.maestro.yaml b/templates/namelist.maestro.yaml index e5920e72..5d600baa 100644 --- a/templates/namelist.maestro.yaml +++ b/templates/namelist.maestro.yaml @@ -163,6 +163,12 @@ maestro: target: options: force_zero_particle_flux: true + + optimization_options: + convergence_options: + maximum_iterations: 30 + strategy_options: + AllowedExcursions: [0.25, 0.0] # Operations to do to mitim_state before passing it to PORTALS initialization_parameters: diff --git a/tests/MAESTRO_workflow.py b/tests/MAESTRO_workflow.py index 5e5e1178..ec8857c6 100644 --- a/tests/MAESTRO_workflow.py +++ b/tests/MAESTRO_workflow.py @@ -3,7 +3,7 @@ from mitim_tools import __mitimroot__ from mitim_modules.maestro.scripts import run_maestro -cold_start = False +cold_start = True folder = __mitimroot__ / "tests" / "scratch" / "maestro_test" template = __mitimroot__ / "templates" / "namelist.maestro.yaml" From 5514f1d44c6c21c0e2504b52f7487e84a954e90b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 22 Sep 2025 14:46:01 -0400 Subject: [PATCH 310/385] Capability to expand surrogate_data to include other channels --- .../portals/utils/PORTALSanalysis.py | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 54d6150e..7182d8c8 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -1,5 +1,6 @@ import copy import torch +import json import numpy as np import pandas as pd import matplotlib.pyplot as plt @@ -1074,3 +1075,99 @@ def plotMetrics(self, extra_lab="", **kwargs): axsGrads[0].legend(prop={"size": 8}) +def surrogate_file_expansion( + portals_folder, + file_new, + variables = ['aLte', 'aLti', 'aLne', 'nuei', 'tite', 'beta_e'], + output_mapper ={ + 'Qe_tr_turb': ['QeMWm2_tr_turb', 'Qgb'], + 'Qi_tr_turb': ['QiMWm2_tr_turb', 'Qgb'], + 'Ge_tr_turb': ['Ge1E20m2_tr_turb', 'Ggb'], + 'Qe_tr_neoc': ['QeMWm2_tr_neoc', 'Qgb'], + 'Qi_tr_neoc': ['QiMWm2_tr_neoc', 'Qgb'], + 'Ge_tr_neoc': ['Ge1E20m2_tr_neoc', 'Ggb'], + } + ): + ''' + This function reads a PORTALS folder and extracts the inputs and outputs using variables and output_mapper. + It then writes a surrogate_data.csv file that can be used as extrapointsFile + + This is useful when you have a PORTALS simulation with [te, ti] and now you want to create a surrogate model with [te, ti, ne] + + ''' + + # ---------------------------------------------------------------------------- + # Grab powerstates and turb_files + # ---------------------------------------------------------------------------- + + portals = PORTALSanalyzer.from_folder(portals_folder) + mitim_runs = IOtools.unpickle_mitim(portals.opt_fun.mitim_model.optimization_object.optimization_extra) + + powerstates = [mitim_runs[i]['powerstate'] for i in range(0, portals.ilast+1)] + turb_files = [powerstates[0].transport_options['folder'] / 'Execution' / f'Evaluation.{i}' / 'transport_simulation_folder' / 'fluxes_turb.json' for i in range(0, portals.ilast+1)] + + turb_info = [] + for file_name in turb_files: + with open(file_name, 'r') as f: + turb_info.append(json.load(f)) + + # ---------------------------------------------------------------------------- + # Prepare dictionary with new inputs and outputs + # ---------------------------------------------------------------------------- + + df_new = [] + + for i in range(0, portals.ilast+1): + + df_helper = {} + for var in output_mapper: + df_helper[var] = { + 'y': ( mitim_runs[i]['powerstate'].plasma[output_mapper[var][0]][0,1:] / mitim_runs[i]['powerstate'].plasma[output_mapper[var][1]][0,1:] ).cpu().numpy(), + 'yvar': ( (mitim_runs[i]['powerstate'].plasma[output_mapper[var][0]+'_stds'][0,1:] / mitim_runs[i]['powerstate'].plasma[output_mapper[var][1]][0,1:])**2 ).cpu().numpy(), + 'x_names': variables, + } + for ix,x in enumerate(df_helper[var]['x_names']): + df_helper[var][f'x{ix}'] = mitim_runs[i]['powerstate'].plasma[x][0,1:].cpu().numpy() + + # Make it per radius + df_helper_new = {} + for ir in range(len(df_helper['Qe_tr_turb']['y'])): + for var in output_mapper: + new_name = var+f'_{ir+1}' + df_helper_new[new_name] = {} + df_helper_new[new_name]['y'] = df_helper[var]['y'][ir] + df_helper_new[new_name]['yvar'] = df_helper[var]['yvar'][ir] + df_helper_new[new_name]['x_names'] = df_helper[var]['x_names'] + for ix,x in enumerate(df_helper[var]['x_names']): + df_helper_new[new_name][f'x{ix}'] = df_helper[var][f'x{ix}'][ir] + + df_new.append(df_helper_new) + + # ---------------------------------------------------------------------------- + # Insert in new dataframe + # ---------------------------------------------------------------------------- + + # Flatten df_new into rows + rows = [] + for d in df_new: + for model, vals in d.items(): + row = { + "Model": model, + "y": vals["y"], + "yvar": vals["yvar"], + "x_names": vals["x_names"], + } + # Add all x0, x1, ... keys + for k, v in vals.items(): + if k.startswith("x"): + row[k] = v + rows.append(row) + + # Build dataframe + df_new_flat = pd.DataFrame(rows) + + # ---------------------------------------------------------------------------- + # Grab those that have not been updated (targets) + # ---------------------------------------------------------------------------- + + df_new_flat.to_csv(file_new, index=False) From 72e58339141529db4dddea6b0082553cc3649812 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 23 Sep 2025 11:15:12 -0400 Subject: [PATCH 311/385] Bug fix that prevented PORTALS to use CGYRO's result in prep mode (json file external) --- .../physics_models/transport_cgyro.py | 8 ++--- .../physics_models/transport_tglf.py | 34 ++++++++++--------- .../powertorch/utils/TRANSPORTtools.py | 9 +++-- 3 files changed, 29 insertions(+), 22 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index fee17d4f..f2b49a9b 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -69,12 +69,12 @@ def _evaluate_gyrokinetic_model(self, code = 'cgyro', gk_object = None): self.QieGB_turb = self.QeGB_turb*0.0 #TODO self.QieGB_turb_stds = self.QeGB_turb*0.0 #TODO - - from mitim_modules.powertorch.utils.TRANSPORTtools import write_json - write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb') elif run_type == 'prep': + # Prevent writing the json file from variables, as we will wait for the user to run CGYRO externally and provide the json themselves + self._write_json_from_variables_turb = False + # Wait until the user has placed the json file in the right folder self.powerstate.profiles_transport.write_state(self.folder / subfolder_name / "input.gacode") @@ -127,7 +127,7 @@ def evaluate_turbulence(self): # Run base TGLF, to keep track of discrepancies! --------------------------------------------- simulation_options_tglf = self.transport_evaluator_options["tglf"] simulation_options_tglf["use_scan_trick_for_stds"] = None - self._evaluate_tglf() + self._evaluate_tglf(pass_info = False) # -------------------------------------------------------------------------------------------- self._evaluate_gyrokinetic_model(code = 'cgyro', gk_object = CGYROtools.CGYRO) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 03529e00..3c63bd3a 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -10,7 +10,7 @@ def evaluate_turbulence(self): self._evaluate_tglf() # Have it separate such that I can call it from the CGYRO class but without the decorator - def _evaluate_tglf(self): + def _evaluate_tglf(self, pass_info = True): # ------------------------------------------------------------------------------------------------------------------------ # Grab options @@ -124,23 +124,25 @@ def _evaluate_tglf(self): # Pass the information to what power_transport expects # ------------------------------------------------------------------------------------------------------------------------ - self.QeGB_turb = Flux_mean[0] - self.QeGB_turb_stds = Flux_std[0] - - self.QiGB_turb = Flux_mean[1] - self.QiGB_turb_stds = Flux_std[1] - - self.GeGB_turb = Flux_mean[2] - self.GeGB_turb_stds = Flux_std[2] - - self.GZGB_turb = Flux_mean[3] - self.GZGB_turb_stds = Flux_std[3] + if pass_info: + + self.QeGB_turb = Flux_mean[0] + self.QeGB_turb_stds = Flux_std[0] + + self.QiGB_turb = Flux_mean[1] + self.QiGB_turb_stds = Flux_std[1] + + self.GeGB_turb = Flux_mean[2] + self.GeGB_turb_stds = Flux_std[2] + + self.GZGB_turb = Flux_mean[3] + self.GZGB_turb_stds = Flux_std[3] - self.MtGB_turb = Flux_mean[4] - self.MtGB_turb_stds = Flux_std[4] + self.MtGB_turb = Flux_mean[4] + self.MtGB_turb_stds = Flux_std[4] - self.QieGB_turb = Flux_mean[5] - self.QieGB_turb_stds = Flux_std[5] + self.QieGB_turb = Flux_mean[5] + self.QieGB_turb_stds = Flux_std[5] return tglf diff --git a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py index 1a2ae3e8..b7c4438b 100644 --- a/src/mitim_modules/powertorch/utils/TRANSPORTtools.py +++ b/src/mitim_modules/powertorch/utils/TRANSPORTtools.py @@ -16,7 +16,6 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): containing the simulation results. JSON should look like: { - 'fluxes_mean': { 'QeGB': ... @@ -41,7 +40,9 @@ def write_json(self, file_name = 'fluxes_turb.json', suffix= 'turb'): } ''' - if self.folder.exists(): + write_json_from_variables = self._write_json_from_variables_turb if suffix == 'turb' else self._write_json_from_variables_neoc + + if self.folder.exists() and write_json_from_variables: with open(self.folder / file_name, 'w') as f: @@ -95,6 +96,10 @@ def __init__(self, powerstate, name = "test", folder = "~/scratch/", evaluation_ # Model results is None by default, but can be assigned in evaluate self.model_results = None + + # By default, write the json files after evaluating the variables (will be changed in gyrokinetic "prep" run mode) + self._write_json_from_variables_turb = True + self._write_json_from_variables_neoc = True # ---------------------------------------------------------------------------------------- # labels for plotting From 4d0eee29a0786df79af8d557eb230531d04a6f8e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 23 Sep 2025 11:21:22 -0400 Subject: [PATCH 312/385] Default PORTALS-CGYRO is prep --- templates/namelist.portals.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index 3e5b0a31..b127927f 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -163,7 +163,7 @@ transport: extraOptions: {} # Run type: normal (submit and wait), submit (submit and do not wait), prep (do not submit) - run_type: "normal" + run_type: "prep" read: tmin: 0.0 From b699b7514a63057e6fb9a760a11ac2528e8a2e0d Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Tue, 23 Sep 2025 14:19:40 -0400 Subject: [PATCH 313/385] idiot-proofed strange parser error --- src/mitim_tools/simulation_tools/SIMtools.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index fa249a7f..b88bf7d9 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -476,7 +476,8 @@ def _run( # Loop over each folder and launch code, waiting if we've reached max_parallel_execution GACODEcommand += "for folder in \"${folders[@]}\"; do\n" - GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = self.simulation_job.folderExecution)} &\n' + folder_arg = '"$folder"' + GACODEcommand += f' {code_call(folder = {folder_arg}, n = cores_per_code_call, p = self.simulation_job.folderExecution)} &\n' GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" GACODEcommand += "done\n\n" GACODEcommand += "wait\n" From 40d04efbc79cb3c7ef06dd72d8a85ddd6a5a4543 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Tue, 23 Sep 2025 14:57:28 -0400 Subject: [PATCH 314/385] generalize process_fluctuations for cases without some fields --- .../gacode_tools/utils/CGYROutils.py | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 6b4ea701..dc75cc2f 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -114,7 +114,10 @@ def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for self.cgyrodata.getxflux() # Understand positions - self.electron_flag = np.where(self.cgyrodata.z == -1)[0][0] + if -1 in self.cgyrodata.z: + self.electron_flag = np.where(self.cgyrodata.z == -1)[0][0] + else: + self.electron_flag = None self.all_flags = np.arange(0, len(self.cgyrodata.z), 1) self.ions_flags = self.all_flags[self.all_flags != self.electron_flag] @@ -262,25 +265,25 @@ def _process_fluctuations(self): self.bpar, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) self.tmax_fluct = _detect_exploiding_signal(self.t, self.phi**2) - - moment, species, field = 'n', self.electron_flag, 0 - self.ne, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) + if 'kxky_n' in self.cgyrodata.__dict__: + moment, species, field = 'n', self.electron_flag, 0 + self.ne, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) species = self.ions_flags self.ni_all, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,nions,ntoroidal,time) self.ni = self.ni_all.sum(axis=1) # [COMPLEX] (nradial,ntoroidal,time) - moment, species, field = 'e', self.electron_flag, 0 - Ee, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) - - species = self.ions_flags - Ei_all, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,nions,ntoroidal,time) - Ei = Ei_all.sum(axis=1) - - # Transform to temperature - self.Te = 2/3 * Ee - self.ne - self.Ti_all = 2/3 * Ei_all - self.ni_all - self.Ti = 2/3 * Ei - self.ni + if 'kxky_e' in self.cgyrodata.__dict__: + moment, species, field = 'e', self.electron_flag, 0 + Ee, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) + + species = self.ions_flags + Ei_all, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,nions,ntoroidal,time) + Ei = Ei_all.sum(axis=1) # [COMPLEX] (nradial,ntoroidal,time) + # Transform to temperature + self.Te = 2/3 * Ee - self.ne + self.Ti_all = 2/3 * Ei_all - self.ni_all + self.Ti = 2/3 * Ei - self.ni # Sum over radial modes and divide between n=0 and n>0 modes, RMS variables = ['phi', 'apar', 'bpar', 'ne', 'ni_all', 'ni', 'Te', 'Ti', 'Ti_all'] From 391bc0e1674d07a6713322480d972c8f0be9d3a5 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Tue, 23 Sep 2025 15:30:54 -0400 Subject: [PATCH 315/385] more generalizations to process fluxes and flucts without EM fields --- .../gacode_tools/utils/CGYROutils.py | 151 ++++++++++-------- 1 file changed, 88 insertions(+), 63 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index dc75cc2f..1fc587bd 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -163,8 +163,11 @@ def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for print(f'\t- No fluctuations found in CGYRO data ({IOtools.clipstr(self.folder)}), skipping fluctuation processing and will not be able to plot default Notebook', typeMsg='w') else: print('\t- Minimal mode, skipping fluctuations processing', typeMsg='i') - - self._process_fluxes() + try: + self._process_fluxes() + except ValueError as e: + print(f'\t- Error processing fluxes: {e}', typeMsg='w') + #self._process_fluxes() self._saturate_signals() self.remove_symlinks() @@ -269,9 +272,9 @@ def _process_fluctuations(self): moment, species, field = 'n', self.electron_flag, 0 self.ne, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,ntoroidal,time) - species = self.ions_flags - self.ni_all, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,nions,ntoroidal,time) - self.ni = self.ni_all.sum(axis=1) # [COMPLEX] (nradial,ntoroidal,time) + species = self.ions_flags + self.ni_all, _ = self.cgyrodata.kxky_select(theta,field,moment,species,gbnorm=gbnorm) # [COMPLEX] (nradial,nions,ntoroidal,time) + self.ni = self.ni_all.sum(axis=1) # [COMPLEX] (nradial,ntoroidal,time) if 'kxky_e' in self.cgyrodata.__dict__: moment, species, field = 'e', self.electron_flag, 0 @@ -289,7 +292,7 @@ def _process_fluctuations(self): variables = ['phi', 'apar', 'bpar', 'ne', 'ni_all', 'ni', 'Te', 'Ti', 'Ti_all'] for var in variables: if var in self.__dict__: - + print(var) # Make sure I go to the real units for all of them ******************* self.__dict__[var] = self.__dict__[var] * self.artificial_rhos_factor # ******************************************************************** @@ -323,35 +326,43 @@ def _process_fluctuations(self): self.__dict__[var+'_rms_sumn'] = (abs(self.__dict__[var])**2).sum(axis=(axis_toroidal))**0.5 # (nradial,time) # Cross-phases - self.neTe = _cross_phase(self.t, self.ne, self.Te) * 180/ np.pi # (nradial, ntoroidal, time) - self.neTe_kx0 = self.neTe[np.argmin(np.abs(self.kx)),:,:] # (ntoroidal, time) - - self.niTi = _cross_phase(self.t, self.ni, self.Ti) * 180/ np.pi # (nradial, ntoroidal, time) - self.niTi_kx0 = self.niTi[np.argmin(np.abs(self.kx)),:,:] + if 'ne' in self.__dict__ and 'Te' in self.__dict__: + self.neTe = _cross_phase(self.t, self.ne, self.Te) * 180/ np.pi # (nradial, ntoroidal, time) + self.neTe_kx0 = self.neTe[np.argmin(np.abs(self.kx)),:,:] # (ntoroidal, time) - self.phiTe = _cross_phase(self.t, self.phi, self.Te) * 180/ np.pi # (nradial, ntoroidal, time) - self.phiTe_kx0 = self.phiTe[np.argmin(np.abs(self.kx)),:,:] + if 'ni' in self.__dict__ and 'Ti' in self.__dict__: + self.niTi = _cross_phase(self.t, self.ni, self.Ti) * 180/ np.pi # (nradial, ntoroidal, time) + self.niTi_kx0 = self.niTi[np.argmin(np.abs(self.kx)),:,:] - self.phiTi = _cross_phase(self.t, self.phi, self.Ti) * 180/ np.pi # (nradial, ntoroidal, time) - self.phiTi_kx0 = self.phiTi[np.argmin(np.abs(self.kx)),:,:] - + if 'phi' in self.__dict__ and 'Te' in self.__dict__: + self.phiTe = _cross_phase(self.t, self.phi, self.Te) * 180/ np.pi # (nradial, ntoroidal, time) + self.phiTe_kx0 = self.phiTe[np.argmin(np.abs(self.kx)),:,:] + + if 'phi' in self.__dict__ and 'Ti' in self.__dict__: + self.phiTi = _cross_phase(self.t, self.phi, self.Ti) * 180/ np.pi # (nradial, ntoroidal, time) + self.phiTi_kx0 = self.phiTi[np.argmin(np.abs(self.kx)),:,:] + self.phiTi_all = [] - for ion in self.ions_flags: - self.phiTi_all.append(_cross_phase(self.t, self.phi, self.Ti_all[:,ion,:]) * 180/ np.pi) - self.phiTi_all = np.array(self.phiTi_all) - self.phiTi_all_kx0 = self.phiTi_all[:,np.argmin(np.abs(self.kx)),:,:] - - self.phine = _cross_phase(self.t, self.phi, self.ne) * 180/ np.pi # (nradial, ntoroidal, time) - self.phine_kx0 = self.phine[np.argmin(np.abs(self.kx)),:,:] + if 'phi' in self.__dict__ and 'Ti_all' in self.__dict__: + for ion in self.ions_flags: + self.phiTi_all.append(_cross_phase(self.t, self.phi, self.Ti_all[:,ion,:]) * 180/ np.pi) + self.phiTi_all = np.array(self.phiTi_all) + self.phiTi_all_kx0 = self.phiTi_all[:,np.argmin(np.abs(self.kx)),:,:] + + if 'ne' in self.__dict__ and 'phi' in self.__dict__: + self.phine = _cross_phase(self.t, self.phi, self.ne) * 180/ np.pi # (nradial, ntoroidal, time) + self.phine_kx0 = self.phine[np.argmin(np.abs(self.kx)),:,:] - self.phini = _cross_phase(self.t, self.phi, self.ni) * 180/ np.pi # (nradial, ntoroidal, time) - self.phini_kx0 = self.phini[np.argmin(np.abs(self.kx)),:,:] + if 'ni' in self.__dict__ and 'phi' in self.__dict__: + self.phini = _cross_phase(self.t, self.phi, self.ni) * 180/ np.pi # (nradial, ntoroidal, time) + self.phini_kx0 = self.phini[np.argmin(np.abs(self.kx)),:,:] self.phini_all = [] - for ion in self.ions_flags: - self.phini_all.append(_cross_phase(self.t, self.phi, self.ni_all[:,ion,:]) * 180/ np.pi) - self.phini_all = np.array(self.phini_all) - self.phini_all_kx0 = self.phini_all[:,np.argmin(np.abs(self.kx)),:,:] + if 'phi' in self.__dict__ and 'ni_all' in self.__dict__: + for ion in self.ions_flags: + self.phini_all.append(_cross_phase(self.t, self.phi, self.ni_all[:,ion,:]) * 180/ np.pi) + self.phini_all = np.array(self.phini_all) + self.phini_all_kx0 = self.phini_all[:,np.argmin(np.abs(self.kx)),:,:] # Correlation length phi = (abs(self.phi[:,self.ky>0,:])).sum(axis=1) # Sum over toroidal modes n>0 @@ -371,50 +382,62 @@ def _process_fluxes(self): ky_flux = self.cgyrodata.ky_flux # (species, moments, fields, ntoroidal, time) + fields = ['phi','apar','bpar'] + # Electron energy flux i_species, i_moment = -1, 1 - i_fields = 0 - self.Qe_ES_ky = ky_flux[i_species, i_moment, i_fields, :, :] - i_fields = 1 - self.Qe_EM_apar_ky = ky_flux[i_species, i_moment, i_fields, :, :] - i_fields = 2 - self.Qe_EM_aper_ky = ky_flux[i_species, i_moment, i_fields, :, :] - - self.Qe_EM_ky = self.Qe_EM_apar_ky + self.Qe_EM_aper_ky - self.Qe_ky = self.Qe_ES_ky + self.Qe_EM_ky + for i_field, field in enumerate(fields): + if field in self.__dict__: + if field == 'phi': + self.Qe_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] + else: + self.Qe_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] + + if 'Qe_EM_ky' in self.__dict__: + self.Qe_ky = self.Qe_ES_ky + self.Qe_EM_ky + else: + self.Qe_ky = self.Qe_ES_ky # Electron particle flux i_species, i_moment = -1, 0 - i_fields = 0 - self.Ge_ES_ky = ky_flux[i_species, i_moment, i_fields, :] - i_fields = 1 - self.Ge_EM_apar_ky = ky_flux[i_species, i_moment, i_fields, :] - i_fields = 2 - self.Ge_EM_aper_ky = ky_flux[i_species, i_moment, i_fields, :] - - self.Ge_EM_ky = self.Ge_EM_apar_ky + self.Ge_EM_aper_ky - self.Ge_ky = self.Ge_ES_ky + self.Ge_EM_ky + for i_field, field in enumerate(fields): + if field in self.__dict__: + if field == 'phi': + self.Ge_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] + else: + self.Ge_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] + + if 'Ge_EM_ky' in self.__dict__: + self.Ge_ky = self.Ge_ES_ky + self.Ge_EM_ky + else: + self.Ge_ky = self.Ge_ES_ky # Ions energy flux i_species, i_moment = self.ions_flags, 1 - i_fields = 0 - self.Qi_all_ES_ky = ky_flux[i_species, i_moment, i_fields, :] - i_fields = 1 - self.Qi_all_EM_apar_ky = ky_flux[i_species, i_moment, i_fields, :] - i_fields = 2 - self.Qi_all_EM_aper_ky = ky_flux[i_species, i_moment, i_fields, :] - - self.Qi_all_EM_ky = self.Qi_all_EM_apar_ky + self.Qi_all_EM_aper_ky - self.Qi_all_ky = self.Qi_all_ES_ky + self.Qi_all_EM_ky - - self.Qi_ky = self.Qi_all_ky.sum(axis=0) - self.Qi_EM_ky = self.Qi_all_EM_ky.sum(axis=0) - self.Qi_EM_apar_ky = self.Qi_all_EM_apar_ky.sum(axis=0) - self.Qi_EM_aper_ky = self.Qi_all_EM_aper_ky.sum(axis=0) - self.Qi_ES_ky = self.Qi_all_ES_ky.sum(axis=0) + for i_field, field in enumerate(fields): + if field in self.__dict__: + if field == 'phi': + self.Qi_all_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] + else: + self.Qi_all_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] + + if 'Qi_all_EM_ky' in self.__dict__: + self.Qi_all_ky = self.Qi_all_ES_ky + self.Qi_all_EM_ky + self.Qi_ky = self.Qi_all_ky.sum(axis=0) + self.Qi_EM_ky = self.Qi_all_EM_ky.sum(axis=0) + self.Qi_EM_apar_ky = self.Qi_all_EM_apar_ky.sum(axis=0) + self.Qi_EM_aper_ky = self.Qi_all_EM_aper_ky.sum(axis=0) + self.Qi_ES_ky = self.Qi_all_ES_ky.sum(axis=0) + else: + self.Qi_all_ky = self.Qi_all_ES_ky + self.Qi_ky = self.Qi_all_ky.sum(axis=0) + self.Qi_ES_ky = self.Qi_all_ES_ky.sum(axis=0) + + + # ************************ # Sum total @@ -422,7 +445,9 @@ def _process_fluxes(self): variables = ['Qe','Ge','Qi','Qi_all'] for var in variables: for i in ['', '_ES', '_EM_apar', '_EM_aper', '_EM']: - self.__dict__[var+i] = self.__dict__[var+i+'_ky'].sum(axis=-2) # (time) + if var+i+'_ky' in self.__dict__: + print(var+i) + self.__dict__[var+i] = self.__dict__[var+i+'_ky'].sum(axis=-2) # (time) # Convert to MW/m^2 self.QeMWm2 = self.Qe * self.Qgb From db2a6be5b853562a47dfcb81fe6510da2cb03e7e Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Tue, 23 Sep 2025 15:43:25 -0400 Subject: [PATCH 316/385] removed print statement and mixup between moment and field --- .../gacode_tools/utils/CGYROutils.py | 31 +++++++------------ 1 file changed, 12 insertions(+), 19 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 1fc587bd..32247229 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -388,11 +388,10 @@ def _process_fluxes(self): i_species, i_moment = -1, 1 for i_field, field in enumerate(fields): - if field in self.__dict__: - if field == 'phi': - self.Qe_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] - else: - self.Qe_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] + if field == 'phi': + self.Qe_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] + else: + self.Qe_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] if 'Qe_EM_ky' in self.__dict__: self.Qe_ky = self.Qe_ES_ky + self.Qe_EM_ky @@ -403,11 +402,10 @@ def _process_fluxes(self): i_species, i_moment = -1, 0 for i_field, field in enumerate(fields): - if field in self.__dict__: - if field == 'phi': - self.Ge_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] - else: - self.Ge_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] + if field == 'phi': + self.Ge_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] + else: + self.Ge_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] if 'Ge_EM_ky' in self.__dict__: self.Ge_ky = self.Ge_ES_ky + self.Ge_EM_ky @@ -418,11 +416,10 @@ def _process_fluxes(self): i_species, i_moment = self.ions_flags, 1 for i_field, field in enumerate(fields): - if field in self.__dict__: - if field == 'phi': - self.Qi_all_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] - else: - self.Qi_all_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] + if field == 'phi': + self.Qi_all_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] + else: + self.Qi_all_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] if 'Qi_all_EM_ky' in self.__dict__: self.Qi_all_ky = self.Qi_all_ES_ky + self.Qi_all_EM_ky @@ -435,9 +432,6 @@ def _process_fluxes(self): self.Qi_all_ky = self.Qi_all_ES_ky self.Qi_ky = self.Qi_all_ky.sum(axis=0) self.Qi_ES_ky = self.Qi_all_ES_ky.sum(axis=0) - - - # ************************ # Sum total @@ -446,7 +440,6 @@ def _process_fluxes(self): for var in variables: for i in ['', '_ES', '_EM_apar', '_EM_aper', '_EM']: if var+i+'_ky' in self.__dict__: - print(var+i) self.__dict__[var+i] = self.__dict__[var+i+'_ky'].sum(axis=-2) # (time) # Convert to MW/m^2 From 6942b6dd7450e18956f26e483e0f7eb7c0bc3b4c Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Tue, 23 Sep 2025 16:07:22 -0400 Subject: [PATCH 317/385] fixed issue for cases with only apar --- .../gacode_tools/utils/CGYROutils.py | 37 +++++++++++++------ 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 32247229..15fd77d6 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -124,7 +124,6 @@ def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for self.all_names = [f"{gacodefuncs.specmap(self.cgyrodata.mass[i],self.cgyrodata.z[i])}({self.cgyrodata.z[i]},{self.cgyrodata.mass[i]:.1f})" for i in self.all_flags] self.fields = np.arange(self.cgyrodata.n_field) - self.aLTi = self.cgyrodata.dlntdr[0] self.aLTe = self.cgyrodata.dlntdr[self.electron_flag] self.aLne = self.cgyrodata.dlnndr[self.electron_flag] @@ -382,7 +381,7 @@ def _process_fluxes(self): ky_flux = self.cgyrodata.ky_flux # (species, moments, fields, ntoroidal, time) - fields = ['phi','apar','bpar'] + fields = ['phi','apar','bpar'][:self.cgyrodata.n_field] # Electron energy flux @@ -390,8 +389,12 @@ def _process_fluxes(self): for i_field, field in enumerate(fields): if field == 'phi': self.Qe_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] - else: - self.Qe_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] + elif field == 'apar': + self.Qe_EM_apar_ky = ky_flux[i_species, i_moment, i_field, :, :] + self.Qe_EM_ky = self.Qe_EM_apar_ky.copy() + elif field == 'bpar': + self.Qe_EM_aper_ky = ky_flux[i_species, i_moment, i_field, :, :] + self.Qe_EM_ky += self.Qe_EM_aper_ky if 'Qe_EM_ky' in self.__dict__: self.Qe_ky = self.Qe_ES_ky + self.Qe_EM_ky @@ -404,8 +407,12 @@ def _process_fluxes(self): for i_field, field in enumerate(fields): if field == 'phi': self.Ge_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] - else: - self.Ge_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] + elif field == 'apar': + self.Ge_EM_apar_ky = ky_flux[i_species, i_moment, i_field, :, :] + self.Ge_EM_ky = self.Ge_EM_apar_ky.copy() + elif field == 'bpar': + self.Ge_EM_aper_ky = ky_flux[i_species, i_moment, i_field, :, :] + self.Ge_EM_ky += self.Ge_EM_aper_ky if 'Ge_EM_ky' in self.__dict__: self.Ge_ky = self.Ge_ES_ky + self.Ge_EM_ky @@ -418,16 +425,24 @@ def _process_fluxes(self): for i_field, field in enumerate(fields): if field == 'phi': self.Qi_all_ES_ky = ky_flux[i_species, i_moment, i_field, :, :] - else: - self.Qi_all_EM_ky += ky_flux[i_species, i_moment, i_field, :, :] + # sum over species + self.Qi_ES_ky = self.Qi_all_ES_ky.sum(axis=0) + elif field == 'apar': + self.Qi_all_EM_apar_ky = ky_flux[i_species, i_moment, i_field, :, :] + self.Qi_all_EM_ky = self.Qi_all_EM_apar_ky.copy() + # sum over species + self.Qi_EM_apar_ky = self.Qi_all_EM_apar_ky.sum(axis=0) + elif field == 'bpar': + self.Qi_all_EM_aper_ky = ky_flux[i_species, i_moment, i_field, :, :] + self.Qi_all_EM_ky += self.Qi_all_EM_aper_ky + # sum over species + self.Qi_EM_aper_ky = self.Qi_all_EM_aper_ky.sum(axis=0) + if 'Qi_all_EM_ky' in self.__dict__: self.Qi_all_ky = self.Qi_all_ES_ky + self.Qi_all_EM_ky self.Qi_ky = self.Qi_all_ky.sum(axis=0) self.Qi_EM_ky = self.Qi_all_EM_ky.sum(axis=0) - self.Qi_EM_apar_ky = self.Qi_all_EM_apar_ky.sum(axis=0) - self.Qi_EM_aper_ky = self.Qi_all_EM_aper_ky.sum(axis=0) - self.Qi_ES_ky = self.Qi_all_ES_ky.sum(axis=0) else: self.Qi_all_ky = self.Qi_all_ES_ky self.Qi_ky = self.Qi_all_ky.sum(axis=0) From c4714b6acb32a6441aa0b20fa88140a03680eaac Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Tue, 23 Sep 2025 16:45:09 -0400 Subject: [PATCH 318/385] added read_cgyro option to save output to pickle --- .../gacode_tools/scripts/read_cgyro.py | 36 ++++++++++++++----- .../gacode_tools/utils/CGYROutils.py | 16 ++++++++- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 3e2ba014..9a3d6e6d 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -1,4 +1,6 @@ import argparse +import pickle +from mitim_tools.gacode_tools.utils.CGYROutils import CGYROoutput from xml.etree.ElementInclude import include import matplotlib.pyplot as plt from IPython import embed @@ -17,13 +19,17 @@ def main(): parser.add_argument("--linear", action="store_true", help="Just a plot of the linear spectra") parser.add_argument("--tmin", type=float, nargs="*", default=None, help="Minimum time to calculate mean and std") parser.add_argument("--scan_subfolder_id" , type=str, nargs="*", default="KY", help="If reading a linear scan, the subfolders contain this common identifier") + parser.add_argument("--noplot", action="store_true", help="If set, it will not plot anything, just read the data.") + parser.add_argument("--pickle", action="store_true", help="If set, it will save the read data in a pickle file for faster reading next time.") args = parser.parse_args() folders = args.folders linear = args.linear tmin = args.tmin include_2D = args.two - + skip_plotting = args.noplot + pkl = args.pickle + suffixes = args.suffixes scan_subfolder_id = args.scan_subfolder_id @@ -45,6 +51,7 @@ def main(): c = CGYROtools.CGYRO() labels = [] + output_pickle = {} for i, folder in enumerate(folders): labels.append(f"case {i + 1}") @@ -66,15 +73,26 @@ def main(): preffix=scan_subfolder_id[i] ) - if linear: - # Plot linear spectrum - c.plot_quick_linear(labels=labels) - plt.show() - else: - c.plot(labels=labels, include_2D=include_2D, common_colorbar=True) - c.fn.show() + if pkl: + print("Pickling data...") + print(c.results[labels[-1]]['output']) + with open(folder + "/data.pkl", "wb") as f: + pickle.dump(c.results[labels[-1]]['output'], f) + print("Pickling done.") + + if not skip_plotting: + if linear: + # Plot linear spectrum + c.plot_quick_linear(labels=labels) + plt.show() + else: + c.plot(labels=labels, include_2D=include_2D, common_colorbar=True) + c.fn.show() - embed() + embed() + + + if __name__ == "__main__": main() diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 15fd77d6..74c9531e 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -291,7 +291,6 @@ def _process_fluctuations(self): variables = ['phi', 'apar', 'bpar', 'ne', 'ni_all', 'ni', 'Te', 'Ti', 'Ti_all'] for var in variables: if var in self.__dict__: - print(var) # Make sure I go to the real units for all of them ******************* self.__dict__[var] = self.__dict__[var] * self.artificial_rhos_factor # ******************************************************************** @@ -726,3 +725,18 @@ def quends_analysis(t, S, debug = False): embed() return mean, std, stats + +def fetch_CGYROoutput(folder, remote): + '''This is a helper function to bring back only the python object from a remote CGYRO run + this is useful when nonlinear runs are too large to be transfered back. It is important to + make sure MITIM is the same version in the remote and local machine. for now I'm writing the commit hash to a file + and checking if they are the same and raising a warning if not.''' + + # execute pickle_cgyro remotely + + # retrieve remote folder + + # read pickle file + + c=None + return c From 5361e1235a0598712841d3b3d97d341411148df8 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Tue, 23 Sep 2025 16:52:14 -0400 Subject: [PATCH 319/385] don't pickle fluctuation data unless --two is used --- .../gacode_tools/scripts/read_cgyro.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 9a3d6e6d..bd8617e3 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -61,9 +61,8 @@ def main(): folder=folder, suffix=suffixes[i], preffix=scan_subfolder_id[i] - ) - - else: + ) + elif include_2D: c.read( label=labels[-1], folder=folder, @@ -72,6 +71,16 @@ def main(): suffix=suffixes[i], preffix=scan_subfolder_id[i] ) + else: + c.read( + label=labels[-1], + folder=folder, + tmin=tmin[i], + last_tmin_for_linear=last_tmin_for_linear, + suffix=suffixes[i], + preffix=scan_subfolder_id[i], + minimal=True + ) if pkl: print("Pickling data...") From 9cf9b40a450f2c9fcdaec68528e4a8a4297d8ed3 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 23 Sep 2025 18:35:46 -0400 Subject: [PATCH 320/385] --two should not be required if one wants fluctution information. Pass minimal as a flag --- src/mitim_tools/gacode_tools/scripts/read_cgyro.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index bd8617e3..799a746f 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -21,6 +21,8 @@ def main(): parser.add_argument("--scan_subfolder_id" , type=str, nargs="*", default="KY", help="If reading a linear scan, the subfolders contain this common identifier") parser.add_argument("--noplot", action="store_true", help="If set, it will not plot anything, just read the data.") parser.add_argument("--pickle", action="store_true", help="If set, it will save the read data in a pickle file for faster reading next time.") + parser.add_argument("--minimal", action="store_true") + args = parser.parse_args() folders = args.folders @@ -29,10 +31,14 @@ def main(): include_2D = args.two skip_plotting = args.noplot pkl = args.pickle + minimal = args.minimal suffixes = args.suffixes scan_subfolder_id = args.scan_subfolder_id + + if isinstance(scan_subfolder_id, str): + scan_subfolder_id = [scan_subfolder_id for _ in range(len(folders))] if suffixes is None: suffixes = ["" for _ in range(len(folders))] @@ -60,7 +66,8 @@ def main(): label=labels[-1], folder=folder, suffix=suffixes[i], - preffix=scan_subfolder_id[i] + preffix=scan_subfolder_id[i], + minimal=minimal ) elif include_2D: c.read( @@ -69,7 +76,8 @@ def main(): tmin=tmin[i], last_tmin_for_linear=last_tmin_for_linear, suffix=suffixes[i], - preffix=scan_subfolder_id[i] + preffix=scan_subfolder_id[i], + minimal=minimal ) else: c.read( @@ -79,7 +87,7 @@ def main(): last_tmin_for_linear=last_tmin_for_linear, suffix=suffixes[i], preffix=scan_subfolder_id[i], - minimal=True + minimal=minimal ) if pkl: From d4ae0d0f3fe51faa72728052897e499c3067a37d Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Wed, 24 Sep 2025 10:40:06 -0400 Subject: [PATCH 321/385] MITIMstate: added mass weighted velocity profiles when lumping species --- .../plasmastate_tools/MITIMstate.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 652d2ce4..a9f634fd 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1320,9 +1320,10 @@ def remove(self, ions_list): break if not fail: - var_changes = ["ni(10^19/m^3)", "ti(keV)"] + var_changes = ["ni(10^19/m^3)", "ti(keV)", "vpol(m/s)", "vtor(m/s)"] for i in var_changes: - self.profiles[i] = np.delete(self.profiles[i], ions_list, axis=1) + if i in self.profiles: + self.profiles[i] = np.delete(self.profiles[i], ions_list, axis=1) if not fail: # Ensure we extract the scalar value from the array @@ -1378,6 +1379,13 @@ def lumpSpecies( A = Z * 2 if force_mass is None else force_mass nZ = fZ1 / Z * self.profiles["ne(10^19/m^3)"] + mass_density = A * self.derived["fi"] + + # Compute the mass weighted average velocity profiles + if "vpol(m/s)" in self.profiles: + vpol = np.sum((mass_density * self.profiles["vpol(m/s)"])[:,np.array(ions_list)-1],axis=1) / np.sum(mass_density[:,np.array(ions_list)-1],axis=1) + if "vtor(m/s)" in self.profiles: + vtor = np.sum((mass_density * self.profiles["vtor(m/s)"])[:,np.array(ions_list)-1],axis=1) / np.sum(mass_density[:,np.array(ions_list)-1],axis=1) print(f"\t\t\t* New lumped impurity has Z={Z:.2f}, A={A:.2f} (calculated as 2*Z)") @@ -1395,6 +1403,14 @@ def lumpSpecies( np.transpose(np.atleast_2d(self.profiles["ti(keV)"][:, 0])), axis=1, ) + if "vpol(m/s)" in self.profiles: + self.profiles["vpol(m/s)"] = np.append( + self.profiles["vpol(m/s)"], np.transpose(np.atleast_2d(vpol)), axis=1 + ) + if "vtor(m/s)" in self.profiles: + self.profiles["vtor(m/s)"] = np.append( + self.profiles["vtor(m/s)"], np.transpose(np.atleast_2d(vtor)), axis=1 + ) self.readSpecies() self.derive_quantities(rederiveGeometry=False) @@ -1549,7 +1565,7 @@ def moveSpecie(self, pos=2, pos_new=1): self.profiles["nion"] = np.array([f"{int(self.profiles['nion'][0])+1}"]) - for ikey in ["name", "mass", "z", "type", "ni(10^19/m^3)", "ti(keV)"]: + for ikey in ["name", "mass", "z", "type", "ni(10^19/m^3)", "ti(keV)", "vpol(m/s)", "vtor(m/s)"]: if len(self.profiles[ikey].shape) > 1: axis = 1 newly = self.profiles[ikey][:, position_to_moveFROM_in_profiles] From bb6bd1c4e1a50710e7ea3e260320cd12a85fbbc6 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Wed, 24 Sep 2025 10:42:44 -0400 Subject: [PATCH 322/385] name pickle file descriptively --- src/mitim_tools/gacode_tools/scripts/read_cgyro.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 799a746f..7d4e1e88 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -5,6 +5,7 @@ import matplotlib.pyplot as plt from IPython import embed from mitim_tools.gacode_tools import CGYROtools +import os """ e.g. read_cgyro.py folder @@ -93,7 +94,11 @@ def main(): if pkl: print("Pickling data...") print(c.results[labels[-1]]['output']) - with open(folder + "/data.pkl", "wb") as f: + folder_abs = os.path.abspath(folder) + simname = folder_abs.rstrip("/").split("/")[-1] + print(f"Pickling to {simname}.pkl", folder_abs.rstrip("/").split("/")) + + with open(f"{simname}_data.pkl", "wb") as f: pickle.dump(c.results[labels[-1]]['output'], f) print("Pickling done.") From 53c1ae70c2847acb0623f2115e90092be4f662cc Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Wed, 24 Sep 2025 10:43:26 -0400 Subject: [PATCH 323/385] added initial retrieval function --- .../gacode_tools/utils/CGYROutils.py | 37 ++++++++++++++++--- 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 74c9531e..9a09c332 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -162,6 +162,7 @@ def __init__(self, folder, suffix = None, tmin=0.0, minimal=False, last_tmin_for print(f'\t- No fluctuations found in CGYRO data ({IOtools.clipstr(self.folder)}), skipping fluctuation processing and will not be able to plot default Notebook', typeMsg='w') else: print('\t- Minimal mode, skipping fluctuations processing', typeMsg='i') + try: self._process_fluxes() except ValueError as e: @@ -726,17 +727,41 @@ def quends_analysis(t, S, debug = False): return mean, std, stats -def fetch_CGYROoutput(folder, remote): +def fetch_CGYROoutput(folder_local, folders_remote, machine, minimal=True): '''This is a helper function to bring back only the python object from a remote CGYRO run this is useful when nonlinear runs are too large to be transfered back. It is important to make sure MITIM is the same version in the remote and local machine. for now I'm writing the commit hash to a file and checking if they are the same and raising a warning if not.''' - + from mitim_tools.misc_tools import FARMINGtools + import pickle # execute pickle_cgyro remotely + folders_string = " ".join(folders_remote) + command = f"mitim_plot_cgyro --noplot --pickle {"--minimal" if minimal else ""} {folders_string}" + print(f"Executing remotely: {command}") + FARMINGtools.perform_quick_remote_execution( + folder_local, + machine, + command + ) + + # retrieve remote file + remote_files = [f"{folder_remote}/{folder_remote.rstrip("/").split("/")[-1]}_data.pkl" for folder_remote in folders_remote] + FARMINGtools.retrieve_files_from_remote( + folder_local, + machine, + remote_files + ) + + # read pickle file as cgyroOutput object + c={} + for i, folder_remote in enumerate(folders_remote): + with open(f"{folder_local}/{folder_remote.rstrip('/').split('/')[-1]}_data.pkl", "rb") as f: + c[folder_remote] = pickle.load(f) + print(f"Retrieved CGYRO output from {folder_remote}") - # retrieve remote folder - - # read pickle file - c=None return c + +if __name__ == "__main__": + c = fetch_CGYROoutput(folder_local=".", folders_remote=["/cosmos/vast/scratch/hallefkt/arc_low_current_v3a_n8_fast_+20%_alne_sugama", "/cosmos/vast/scratch/hallefkt/arc_low_current_v3a_n8_fast_-20%_alne_sugama"], machine="cosmos", minimal=False) + c \ No newline at end of file From 5b167897160175cc6f0b07ae0f14eee79f3adcf4 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Wed, 24 Sep 2025 10:47:32 -0400 Subject: [PATCH 324/385] yet another parser issue on cosmos :( --- src/mitim_tools/gacode_tools/utils/CGYROutils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 9a09c332..3adc0115 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -736,7 +736,8 @@ def fetch_CGYROoutput(folder_local, folders_remote, machine, minimal=True): import pickle # execute pickle_cgyro remotely folders_string = " ".join(folders_remote) - command = f"mitim_plot_cgyro --noplot --pickle {"--minimal" if minimal else ""} {folders_string}" + minimal_flag = "--minimal" if minimal else "" + command = f"mitim_plot_cgyro --noplot --pickle {minimal_flag} {folders_string}" print(f"Executing remotely: {command}") FARMINGtools.perform_quick_remote_execution( folder_local, From 8f7d190f0b9028286b0a2b147aa8dcf8fb10ca75 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Wed, 24 Sep 2025 10:53:19 -0400 Subject: [PATCH 325/385] more parser errors yay sorry --- src/mitim_tools/gacode_tools/utils/CGYROutils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index 3adc0115..d608dec0 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -746,7 +746,7 @@ def fetch_CGYROoutput(folder_local, folders_remote, machine, minimal=True): ) # retrieve remote file - remote_files = [f"{folder_remote}/{folder_remote.rstrip("/").split("/")[-1]}_data.pkl" for folder_remote in folders_remote] + remote_files = [f"{folder_remote}/{folder_remote.split('/')[-1]}_data.pkl" for folder_remote in folders_remote] FARMINGtools.retrieve_files_from_remote( folder_local, machine, From ee5167f412e18f4227d9dd2a508fbdb826260acb Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Wed, 24 Sep 2025 11:17:40 -0400 Subject: [PATCH 326/385] correct location for pickle output --- src/mitim_tools/gacode_tools/scripts/read_cgyro.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py index 7d4e1e88..e908af39 100644 --- a/src/mitim_tools/gacode_tools/scripts/read_cgyro.py +++ b/src/mitim_tools/gacode_tools/scripts/read_cgyro.py @@ -96,9 +96,9 @@ def main(): print(c.results[labels[-1]]['output']) folder_abs = os.path.abspath(folder) simname = folder_abs.rstrip("/").split("/")[-1] - print(f"Pickling to {simname}.pkl", folder_abs.rstrip("/").split("/")) + print(f"Pickling to {simname}.pkl") - with open(f"{simname}_data.pkl", "wb") as f: + with open(f"{folder}/{simname}_data.pkl", "wb") as f: pickle.dump(c.results[labels[-1]]['output'], f) print("Pickling done.") From 2d4629153559ffd79a6e454b425a9676351fd008 Mon Sep 17 00:00:00 2001 From: josephb-hall Date: Wed, 24 Sep 2025 11:29:28 -0400 Subject: [PATCH 327/385] working remote pickling and retrieval function for CGYRO! --- .../gacode_tools/utils/CGYROutils.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/CGYROutils.py b/src/mitim_tools/gacode_tools/utils/CGYROutils.py index d608dec0..960df58d 100644 --- a/src/mitim_tools/gacode_tools/utils/CGYROutils.py +++ b/src/mitim_tools/gacode_tools/utils/CGYROutils.py @@ -727,7 +727,7 @@ def quends_analysis(t, S, debug = False): return mean, std, stats -def fetch_CGYROoutput(folder_local, folders_remote, machine, minimal=True): +def fetch_CGYROoutput(folder_local, folders_remote, machine, minimal=True, delete_local=False): '''This is a helper function to bring back only the python object from a remote CGYRO run this is useful when nonlinear runs are too large to be transfered back. It is important to make sure MITIM is the same version in the remote and local machine. for now I'm writing the commit hash to a file @@ -736,7 +736,13 @@ def fetch_CGYROoutput(folder_local, folders_remote, machine, minimal=True): import pickle # execute pickle_cgyro remotely folders_string = " ".join(folders_remote) - minimal_flag = "--minimal" if minimal else "" + + if minimal: + minimal_flag = "--minimal" + else: + minimal_flag = "" + + print("MINIMAL FLAG", minimal_flag) command = f"mitim_plot_cgyro --noplot --pickle {minimal_flag} {folders_string}" print(f"Executing remotely: {command}") FARMINGtools.perform_quick_remote_execution( @@ -760,9 +766,15 @@ def fetch_CGYROoutput(folder_local, folders_remote, machine, minimal=True): c[folder_remote] = pickle.load(f) print(f"Retrieved CGYRO output from {folder_remote}") + if delete_local: + import os + for i, folder_remote in enumerate(folders_remote): + os.remove(f"{folder_local}/{folder_remote.rstrip('/').split('/')[-1]}_data.pkl") + print(f"Deleted local pickle file {folder_local}/{folder_remote.rstrip('/').split('/')[-1]}_data.pkl") return c if __name__ == "__main__": - c = fetch_CGYROoutput(folder_local=".", folders_remote=["/cosmos/vast/scratch/hallefkt/arc_low_current_v3a_n8_fast_+20%_alne_sugama", "/cosmos/vast/scratch/hallefkt/arc_low_current_v3a_n8_fast_-20%_alne_sugama"], machine="cosmos", minimal=False) - c \ No newline at end of file + c = fetch_CGYROoutput(folder_local=".", folders_remote=["/cosmos/vast/scratch/hallefkt/arc_low_current_v3a_n8_fast_+20%_alne_sugama", "/cosmos/vast/scratch/hallefkt/arc_low_current_v3a_n8_fast_-20%_alne_sugama"], machine="cosmos", minimal=True) + c + embed() \ No newline at end of file From 5cfaafd6a2b6a8ada9d29c5b0f2c5e3d00ee3d50 Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Wed, 24 Sep 2025 10:53:16 -0400 Subject: [PATCH 328/385] MITIMstate: bugfix key check in moveSpecie() --- .../plasmastate_tools/MITIMstate.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index a9f634fd..4cfc7f19 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1566,15 +1566,16 @@ def moveSpecie(self, pos=2, pos_new=1): self.profiles["nion"] = np.array([f"{int(self.profiles['nion'][0])+1}"]) for ikey in ["name", "mass", "z", "type", "ni(10^19/m^3)", "ti(keV)", "vpol(m/s)", "vtor(m/s)"]: - if len(self.profiles[ikey].shape) > 1: - axis = 1 - newly = self.profiles[ikey][:, position_to_moveFROM_in_profiles] - else: - axis = 0 - newly = self.profiles[ikey][position_to_moveFROM_in_profiles] - self.profiles[ikey] = np.insert( - self.profiles[ikey], position_to_moveTO_in_profiles, newly, axis=axis - ) + if ikey in self.profiles: + if len(self.profiles[ikey].shape) > 1: + axis = 1 + newly = self.profiles[ikey][:, position_to_moveFROM_in_profiles] + else: + axis = 0 + newly = self.profiles[ikey][position_to_moveFROM_in_profiles] + self.profiles[ikey] = np.insert( + self.profiles[ikey], position_to_moveTO_in_profiles, newly, axis=axis + ) self.readSpecies() self.derive_quantities(rederiveGeometry=False) From 33f01605be386996d7e750997e5f7ec7d59720ad Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 24 Sep 2025 16:57:06 -0400 Subject: [PATCH 329/385] Back to \"$folder\"' until further notice --- src/mitim_tools/simulation_tools/SIMtools.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index b88bf7d9..fa249a7f 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -476,8 +476,7 @@ def _run( # Loop over each folder and launch code, waiting if we've reached max_parallel_execution GACODEcommand += "for folder in \"${folders[@]}\"; do\n" - folder_arg = '"$folder"' - GACODEcommand += f' {code_call(folder = {folder_arg}, n = cores_per_code_call, p = self.simulation_job.folderExecution)} &\n' + GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = self.simulation_job.folderExecution)} &\n' GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" GACODEcommand += "done\n\n" GACODEcommand += "wait\n" From 90a174a83990c40baefae8ec18a1d3dcf1dfb4fb Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 26 Sep 2025 10:51:35 -0400 Subject: [PATCH 330/385] Plot L-mode functionals --- .../popcon_tools/FunctionalForms.py | 33 +++++++++++-------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/src/mitim_tools/popcon_tools/FunctionalForms.py b/src/mitim_tools/popcon_tools/FunctionalForms.py index c34d67b9..b0d8352c 100644 --- a/src/mitim_tools/popcon_tools/FunctionalForms.py +++ b/src/mitim_tools/popcon_tools/FunctionalForms.py @@ -1,3 +1,4 @@ +from matplotlib import markers import torch import numpy as np import matplotlib.pyplot as plt @@ -65,11 +66,28 @@ def PRFfunctionals_Lmode(T_avol, n_avol, nu_n, rho=None, debug=False): print(f">> Gradient aLT outside of search range ({g1})", typeMsg="w") if debug: - fig, ax = plt.subplots() + fig, axs = plt.subplots(nrows=3, figsize=(6, 10)) + ax = axs[0] ax.plot(gs, Tvol, "-s") ax.axhline(y=T_avol) ax.axvline(x=g1) + ax = axs[1] + ax.plot(x[0], T, "-s", markersize=1) + ax.set_xlabel(r"$\rho$") + ax.set_ylabel(r"$T$") + GRAPHICStools.addDenseAxis(ax) + ax = axs[2] + aLT_reconstructed = CALCtools.derivation_into_Lx( + torch.from_numpy(x[0]), + torch.from_numpy(T), + ).cpu().numpy() + ax.plot(x[0], aLT_reconstructed, "-s", markersize=1) + ax.set_xlabel(r"$\rho$") + ax.set_ylabel(r"$1/L_T$") + GRAPHICStools.addDenseAxis(ax) + plt.tight_layout() plt.show() + embed() # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # ~~~ Density @@ -79,19 +97,6 @@ def PRFfunctionals_Lmode(T_avol, n_avol, nu_n, rho=None, debug=False): _, n = parabolic(Tbar=n_avol, nu=nu_n, rho=x[0], Tedge=n_roa1) - # g2_n = 1 - # g1Range_n = [0.1,10] - # xtransition_n = 0.8 - - # gs,nvol = np.linspace(g1Range_n[0],g1Range_n[1],points_search), [] - - # n = FUNCTIONALScalc.doubleLinear_aLT(x,gs,g2_n,xtransition_n,n_roa1) - # nvol = FUNCTIONALScalc.calculate_simplified_volavg(x,n) - # g1 = MATHtools.extrapolateCubicSpline(n_avol,nvol,gs) - # n = FUNCTIONALScalc.doubleLinear_aLT(np.atleast_2d(x[0]),g1,g2_n,xtransition_n,n_roa1)[0] - - # if g1g1Range_n[1]: print(f'>> Gradient aLn outside of search range ({g1})',typeMsg='w') - return x[0], T, n # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ From a29f6171fc64bc09ab43268c3631354b83aa44d9 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 29 Sep 2025 17:43:25 -0400 Subject: [PATCH 331/385] Do not force the user to install onnx2pytorch, not required --- .gitignore | 1 - pyproject.toml | 34 ++++++++++++--------------------- templates/namelist.portals.yaml | 1 + 3 files changed, 13 insertions(+), 23 deletions(-) diff --git a/.gitignore b/.gitignore index 117cd399..12bfc4b8 100644 --- a/.gitignore +++ b/.gitignore @@ -4,7 +4,6 @@ */scratch/* *.DS_Store* config_user.json -.DS_Store *.egg-info/ docs/build/ *.bat diff --git a/pyproject.toml b/pyproject.toml index ddfced2b..23d135ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "MITIM" version = "4.0.0" description = "MIT Integrated Modeling Suite for Fusion Applications" readme = "README.md" -requires-python = ">=3.10, <3.13" # Notes: 3.9 has issues with the latest BOTORCH, 3.13 has issues with tensorflow (nn) and omfit_classesv (omfit_new) +requires-python = ">=3.10, <3.13" # Notes: 3.9 has issues with the latest BOTORCH, 3.13 has issues with tensorflow (nn) license = {file = "LICENSE"} authors = [ {name = "P. Rodriguez-Fernandez", email = "pablorf@mit.edu"}, @@ -42,7 +42,7 @@ dependencies = [ "botorch", "scikit-image", # Stricly not for MITIM, but good to have for pygacode "psutil", - "onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models + #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models "tensorflow", "f90nml", "pyyaml" @@ -74,7 +74,7 @@ test = [ "Source" = "https://github.com/pabloprf/MITIM-fusion" [project.scripts] -# mitim_tools interfaces: read, run, plot + mitim_plot_gacode = "mitim_tools.gacode_tools.scripts.read_gacode:main" mitim_plot_tgyro = "mitim_tools.gacode_tools.scripts.read_tgyro:main" mitim_plot_tglf = "mitim_tools.gacode_tools.scripts.read_tglf:main" # [--suffix _0.55] [--gacode input.gacode] @@ -82,14 +82,18 @@ mitim_plot_cgyro = "mitim_tools.gacode_tools.scripts.read_cgyro:main" mitim_plot_eq = "mitim_tools.gs_tools.scripts.read_eq:main" mitim_plot_transp = "mitim_tools.transp_tools.scripts.read_transp:main" mitim_plot_eped = "mitim_tools.eped_tools.scripts.plot_eped:main" - -mitim_run_tglf = "mitim_tools.gacode_tools.scripts.run_tglf:main" # (folder input.tglf) [--gacode input.gacode] [--scan RLTS_2] [--drives True] - -# Optimizations +mitim_plot_maestro = "mitim_modules.maestro.scripts.plot_maestro:main" # --beats 2 (for the last two beats) --only transp mitim_plot_opt = "mitim_tools.opt_tools.scripts.read:main" # Not transferred: --type 4 --resolution 20 mitim_plot_portals = "mitim_modules.portals.scripts.read_portals:main" -mitim_slurm = "mitim_tools.opt_tools.scripts.slurm:main" + mitim_run_portals = "mitim_modules.portals.scripts.run_portals:main" # folder --input input.gacode --namelist namelist.portals.yaml --cold_start +mitim_run_maestro = "mitim_modules.maestro.scripts.run_maestro:main" # add file argument +mitim_run_tglf = "mitim_tools.gacode_tools.scripts.run_tglf:main" # (folder input.tglf) [--gacode input.gacode] [--scan RLTS_2] [--drives True] +mitim_run_transp = "mitim_tools.transp_tools.scripts.run_transp:main" # To run TRANSP (in folder with required files): mitim_run_transp 88664 P01 CMOD --version tshare --trmpi 32 --toricmpi 32 --ptrmpi 32 +mitim_slurm = "mitim_tools.opt_tools.scripts.slurm:main" +mitim_compare_nml = "mitim_tools.misc_tools.scripts.compare_namelist:main" +mitim_scp = "mitim_tools.misc_tools.scripts.retrieve_files:main" # e.g. mitim_scp mfews15 --files /home/pablorf/file1 --folders /home/pablorf/folder1 +mitim_check_maestro = "mitim_modules.maestro.scripts.check_maestro:main" # TRANSP mitim_trcheck = "mitim_tools.transp_tools.scripts.run_check:main" # e.g. mitim_trcheck pablorf @@ -97,20 +101,6 @@ mitim_trcheck_p = "mitim_tools.transp_tools.scripts.run_check_periodic:main" # e mitim_trclean = "mitim_tools.transp_tools.scripts.run_clean:main" # e.g. mitim_trclean 88664P CMOD --numbers 1,2,3 mitim_trlook = "mitim_tools.transp_tools.scripts.run_look:main" # e.g. mitim_trlook 152895P01 CMOD --nofull --plot --remove -# MAESTRO -mitim_plot_maestro = "mitim_modules.maestro.scripts.plot_maestro:main" # --beats 2 (for the last two beats) --only transp -mitim_run_maestro = "mitim_modules.maestro.scripts.run_maestro:main" # add file argument - -# To run TRANSP (in folder with required files): mitim_run_transp 88664 P01 CMOD --version tshare --trmpi 32 --toricmpi 32 --ptrmpi 32 -mitim_run_transp = "mitim_tools.transp_tools.scripts.run_transp:main" - -# Others -mitim_compare_nml = "mitim_tools.misc_tools.scripts.compare_namelist:main" -mitim_scp = "mitim_tools.misc_tools.scripts.retrieve_files:main" # e.g. mitim_scp mfews15 --files /home/pablorf/file1 --folders /home/pablorf/folder1 -#eff_job="mitim_tools.misc_tools.PARALLELtools.py $1" # Give mitim.out or slurm_output.dat -mitim_check_maestro = "mitim_modules.maestro.scripts.check_maestro:main" - - [tool.pytest.ini_options] markers = [ ] diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index b127927f..c00a0608 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -53,6 +53,7 @@ solution: ymin: 1.0 yminymax_atleast: null # Defines minimum range of exploration, e.g. [0,2], even if not achieved via ymin/ymax + # DEPCRECATED fixed_gradients: null # enforce_finite_aLT is used to be able to select ymin = 2.0 for ne but ensure that te, ti is at, e.g., enforce_finite_aLT = 0.95 From 63d0ca2cf3f32e692e1ea840f9eb13a12e71b9e0 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 29 Sep 2025 17:43:38 -0400 Subject: [PATCH 332/385] Problematic folder quotes fixed --- src/mitim_tools/simulation_tools/SIMtools.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index fa249a7f..cba68503 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -476,7 +476,8 @@ def _run( # Loop over each folder and launch code, waiting if we've reached max_parallel_execution GACODEcommand += "for folder in \"${folders[@]}\"; do\n" - GACODEcommand += f' {code_call(folder = '\"$folder\"', n = cores_per_code_call, p = self.simulation_job.folderExecution)} &\n' + folder_str = '"$folder"' # literal double quotes around $folder + GACODEcommand += f' {code_call(folder=folder_str, n=cores_per_code_call, p=self.simulation_job.folderExecution)} &\n' GACODEcommand += " while (( $(jobs -r | wc -l) >= max_parallel_execution )); do sleep 1; done\n" GACODEcommand += "done\n\n" GACODEcommand += "wait\n" From 29bb0e5a4b14ee548df0e062b2565c898ba29d37 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 29 Sep 2025 17:47:06 -0400 Subject: [PATCH 333/385] GIve full location of farming errors --- src/mitim_tools/misc_tools/FARMINGtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 942cefca..63f517f2 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -316,7 +316,7 @@ def full_process( # If not received, write output and error to files self._write_debugging_files(output, error) - cont = print("\t* Not all expected files received, not removing scratch folder (mitim_farming.out and mitim_farming.err written)",typeMsg="q") + cont = print(f"\t* Not all expected files received, not removing scratch folder (mitim_farming.out and mitim_farming.err written in '{self.folder_local / 'mitim_farming.err'}')",typeMsg="q") if not cont: print("[MITIM] Stopped with embed(), you can look at output and error",typeMsg="w",) embed() From 9836fb238e2adcbd6dcf4b15830d2b6bd480b08e Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 29 Sep 2025 17:58:35 -0400 Subject: [PATCH 334/385] Bug fix quotes portals beat --- src/mitim_modules/maestro/utils/PORTALSbeat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index c4e6501b..2636abf9 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -327,7 +327,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr # set the last value of the radial locations to the interpolated value roatop_old = copy.deepcopy(self.portals_parameters['solution']["predicted_roa"][-1]) self.portals_parameters['solution']["predicted_roa"][-1] = roatop - print(f'\t\t\t* Last radial location moved from r/a = {roatop_old} to {self.portals_parameters['solution']["predicted_roa"][-1]}') + print(f'\t\t\t* Last radial location moved from r/a = {roatop_old} to {self.portals_parameters["solution"]["predicted_roa"][-1]}') print(f'\t\t\t* predicted_roa: {self.portals_parameters['solution']["predicted_roa"]}') strKeys = 'predicted_roa' From b8f3f4efa2d654a901911e6dac4ee8e3df1d383a Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 29 Sep 2025 18:00:10 -0400 Subject: [PATCH 335/385] Various double quotes bug fixes --- src/mitim_modules/maestro/utils/PORTALSbeat.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 2636abf9..902d8cc1 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -328,7 +328,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr roatop_old = copy.deepcopy(self.portals_parameters['solution']["predicted_roa"][-1]) self.portals_parameters['solution']["predicted_roa"][-1] = roatop print(f'\t\t\t* Last radial location moved from r/a = {roatop_old} to {self.portals_parameters["solution"]["predicted_roa"][-1]}') - print(f'\t\t\t* predicted_roa: {self.portals_parameters['solution']["predicted_roa"]}') + print(f'\t\t\t* predicted_roa: {self.portals_parameters["solution"]["predicted_roa"]}') strKeys = 'predicted_roa' @@ -339,7 +339,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr # set the last value of the radial locations to the interpolated value rhotop_old = copy.deepcopy(self.portals_parameters['solution']['predicted_rho'][-1]) self.portals_parameters['solution']['predicted_rho'][-1] = self.maestro_instance.parameters_trans_beat['rhotop'] - print(f'\t\t\t* Last radial location moved from rho = {rhotop_old} to {self.portals_parameters['solution']['predicted_rho'][-1]}') + print(f'\t\t\t* Last radial location moved from rho = {rhotop_old} to {self.portals_parameters["solution"]['predicted_rho'][-1]}') strKeys = 'predicted_rho' @@ -348,7 +348,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr # Check if I changed it previously and it hasn't moved if strKeys in self.maestro_instance.parameters_trans_beat: print(f'\t\t\t* {strKeys} in previous PORTALS beat: {self.maestro_instance.parameters_trans_beat[strKeys]}') - print(f'\t\t\t* {strKeys} in current PORTALS beat: {self.portals_parameters['solution'][strKeys]}') + print(f'\t\t\t* {strKeys} in current PORTALS beat: {self.portals_parameters["solution"][strKeys]}') if abs(self.portals_parameters['solution'][strKeys][-1]-self.maestro_instance.parameters_trans_beat[strKeys][-1]) / self.maestro_instance.parameters_trans_beat[strKeys][-1] < minimum_relative_change_in_x: print('\t\t\t* Last radial location was not moved') From fc1aded3297ff83b3a16fa09ad64d873f4248ed4 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Mon, 29 Sep 2025 18:00:46 -0400 Subject: [PATCH 336/385] oops, missed it... --- src/mitim_modules/maestro/utils/PORTALSbeat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 902d8cc1..a1c50cfe 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -339,7 +339,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr # set the last value of the radial locations to the interpolated value rhotop_old = copy.deepcopy(self.portals_parameters['solution']['predicted_rho'][-1]) self.portals_parameters['solution']['predicted_rho'][-1] = self.maestro_instance.parameters_trans_beat['rhotop'] - print(f'\t\t\t* Last radial location moved from rho = {rhotop_old} to {self.portals_parameters["solution"]['predicted_rho'][-1]}') + print(f'\t\t\t* Last radial location moved from rho = {rhotop_old} to {self.portals_parameters["solution"]["predicted_rho"][-1]}') strKeys = 'predicted_rho' From d6ee44630858fa94a0913daf2e81b3ee442acf74 Mon Sep 17 00:00:00 2001 From: pabloprf Date: Mon, 29 Sep 2025 18:02:14 -0400 Subject: [PATCH 337/385] And... once again... sorry --- src/mitim_modules/maestro/utils/PORTALSbeat.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index a1c50cfe..30603fef 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -358,7 +358,7 @@ def _inform(self, use_previous_residual = True, use_previous_surrogate_data = Tr # In the situation where the last radial location moves, I cannot reuse that surrogate data if last_radial_location_moved and reusing_surrogate_data: print('\t\t- Last radial location was moved, so surrogate data will not be reused for that specific location') - self.optimization_options['surrogate_options']["extrapointsModelsAvoidContent"] = ['_tar',f'_{len(self.portals_parameters['solution'][strKeys])}'] + self.optimization_options['surrogate_options']["extrapointsModelsAvoidContent"] = ['_tar',f"_{len(self.portals_parameters['solution'][strKeys])}"] self.try_flux_match_only_for_first_point = False def _inform_save(self): @@ -389,10 +389,10 @@ def _inform_save(self): if 'predicted_roa' in portals_parameters['solution']: self.maestro_instance.parameters_trans_beat['predicted_roa'] = portals_parameters['solution']['predicted_roa'] - print(f'\t\t* predicted_roa saved for future beats: {portals_parameters['solution']["predicted_roa"]}') + print(f'\t\t* predicted_roa saved for future beats: {portals_parameters["solution"]["predicted_roa"]}') elif 'predicted_rho' in portals_parameters['solution']: self.maestro_instance.parameters_trans_beat['predicted_rho'] = portals_parameters['solution']['predicted_rho'] - print(f'\t\t* predicted_rho saved for future beats: {portals_parameters['solution']["predicted_rho"]}') + print(f'\t\t* predicted_rho saved for future beats: {portals_parameters["solution"]["predicted_rho"]}') # ----------------------------------------------------------------------------------------------------------------------- # Defaults to help MAESTRO From d779f90265c5813b811e569474cfed3cc2b0e306 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 30 Sep 2025 13:29:22 -0400 Subject: [PATCH 338/385] PORTALS debugger with standalone and scans of TGLF now is one tool --- .../scripts/runTGLFdrivesfromPORTALS.py | 55 ------------------- ...runTGLF.py => run_tglf_todebug_portals.py} | 2 +- 2 files changed, 1 insertion(+), 56 deletions(-) delete mode 100644 src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py rename src/mitim_modules/portals/scripts/{runTGLF.py => run_tglf_todebug_portals.py} (99%) diff --git a/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py b/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py deleted file mode 100644 index 0a24c090..00000000 --- a/src/mitim_modules/portals/scripts/runTGLFdrivesfromPORTALS.py +++ /dev/null @@ -1,55 +0,0 @@ -import argparse -from mitim_tools.misc_tools import IOtools -from mitim_modules.portals.utils import PORTALSanalysis - -""" -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 5 --pos 0 2 --var 0.05 --wf 0.2 1.0 --num 5 - -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=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("--num",type=int, required=False, default=5) -parser.add_argument("--cold_start", "-r", required=False, default=False, action="store_true") -parser.add_argument("--ion", type=int, required=False, default=2) - - -args = parser.parse_args() -folder = IOtools.expandPath(args.folder) -ev = args.ev -pos = args.pos -wf = args.wf -var = args.var -num = args.num -cold_start = args.cold_start -ion = args.ion - -# --- Workflow - -portals = PORTALSanalysis.PORTALSanalyzer.from_folder(folder) -tglf, code_settings, extraOptions = portals.extractTGLF(positions=pos, evaluation=ev, modified_profiles=True, cold_start=cold_start) - -tglf.runScanTurbulenceDrives( - subfolder="turb", - resolutionPoints=num, - variation=var, - variablesDrives=["RLTS_1", "RLTS_2", "RLNS_1", "XNUE", "TAUS_2", "BETAE"], - code_settings=code_settings, - extraOptions=extraOptions, - cold_start=cold_start, - runWaveForms=wf, - positionIon=ion, -) - -tglf.plotScanTurbulenceDrives(label="turb") diff --git a/src/mitim_modules/portals/scripts/runTGLF.py b/src/mitim_modules/portals/scripts/run_tglf_todebug_portals.py similarity index 99% rename from src/mitim_modules/portals/scripts/runTGLF.py rename to src/mitim_modules/portals/scripts/run_tglf_todebug_portals.py index 5d6cd8d2..f606aad7 100644 --- a/src/mitim_modules/portals/scripts/runTGLF.py +++ b/src/mitim_modules/portals/scripts/run_tglf_todebug_portals.py @@ -30,7 +30,7 @@ args = parser.parse_args() folder = IOtools.expandPath(args.folder) -ev = args.ev +ev = int(args.ev) params = args.params pos = args.pos wf = args.wf From 1f487da40bb843cded31e33ab8e99b9c049ac909 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 30 Sep 2025 13:29:41 -0400 Subject: [PATCH 339/385] Comparing namelists now organizes values per relative differences --- .../misc_tools/scripts/compare_namelist.py | 55 ++++++++++++------- 1 file changed, 35 insertions(+), 20 deletions(-) diff --git a/src/mitim_tools/misc_tools/scripts/compare_namelist.py b/src/mitim_tools/misc_tools/scripts/compare_namelist.py index 878a8216..7efb0d79 100644 --- a/src/mitim_tools/misc_tools/scripts/compare_namelist.py +++ b/src/mitim_tools/misc_tools/scripts/compare_namelist.py @@ -138,31 +138,46 @@ def compareDictionaries(d1, d2, precision_of=None, close_enough=1e-7): return different -def printTable(diff, printing_percent = 1e-5, warning_percent=1e-1): +def printTable(diff, warning_percent=1e-1): + # Compute percent differences first so we can sort by them + percs = {} 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 = 100 * np.abs( - (diff[key][0] - diff[key][1]) / diff[key][0] - ) - except: - perc = np.nan + v0, v1 = diff[key] + if v0 is None or v1 is None: + # Treat missing values as 100% difference for sorting + percs[key] = 100.0 + else: + if v0 != 0.0: + try: + percs[key] = 100 * np.abs((v0 - v1) / v0) + except Exception: + percs[key] = np.nan + else: + percs[key] = np.nan + + # Sort keys by descending percent; NaNs go last + def sort_key(k): + p = percs[k] + return (1, 0) if (p is None or (isinstance(p, float) and np.isnan(p))) else (0, -p) + + for key in sorted(diff.keys(), key=sort_key): + v0, v1 = diff[key] + if v0 is not None: + if v1 is not None: + perc = percs[key] + if perc<1e-2: + perc_str = f"{perc:.2e}" + elif perc<1.0: + perc_str = f"{perc:.3f}" else: - perc = np.nan - print( - f"{key:>15}{str(diff[key][0]):>25}{str(diff[key][1]):>25} (~{perc:.0e}%)", - typeMsg="i" if perc > warning_percent else "", - ) + perc_str = f"{perc:.1f}" + print(f"{key:>15}{str(v0):>25}{str(v1):>25} ({perc_str} %)",typeMsg="i" if perc > warning_percent else "",) else: - print(f"{key:>15}{str(diff[key][0]):>25}{'':>25} (100%)", typeMsg="i") + print(f"{key:>15}{str(v0):>25}{'':>25} (100%)", typeMsg="i") else: - print(f"{key:>15}{'':>25}{str(diff[key][1]):>25} (100%)", typeMsg="i") - print( - "--------------------------------------------------------------------------------" - ) + print(f"{key:>15}{'':>25}{str(v1):>25} (100%)", typeMsg="i") + print("--------------------------------------------------------------------------------") def main(): From bb7dbefd6ba5dafd2762c67aab4af867249383a8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 30 Sep 2025 13:30:07 -0400 Subject: [PATCH 340/385] Fixed TGLF grabber in new portals namelist --- .../portals/utils/PORTALSanalysis.py | 13 ++++++---- templates/namelist.portals.yaml | 26 +++++++++++-------- 2 files changed, 23 insertions(+), 16 deletions(-) diff --git a/src/mitim_modules/portals/utils/PORTALSanalysis.py b/src/mitim_modules/portals/utils/PORTALSanalysis.py index 7182d8c8..ecf85bd5 100644 --- a/src/mitim_modules/portals/utils/PORTALSanalysis.py +++ b/src/mitim_modules/portals/utils/PORTALSanalysis.py @@ -431,12 +431,14 @@ def extractProfiles(self, evaluation=None, modified_profiles=False): evaluation = self.ibest elif evaluation < 0: evaluation = self.ilast + else: + evaluation = int(evaluation) powerstate = self.mitim_runs[evaluation]["powerstate"] try: - p0 = powerstate.profiles if not modified_profiles else powerstate.model_results.profiles - except TypeError: + p0 = powerstate.profiles if not modified_profiles else powerstate.profiles_transport + except (TypeError, AttributeError): raise Exception(f"[MITIM] Could not extract profiles from evaluation {evaluation}, are you sure you have the right index?") p = copy.deepcopy(p0) @@ -592,10 +594,11 @@ def extractTGLF(self, folder=None, positions=None, evaluation=None, cold_start=F p.write_state(file=inputgacode) tglf = TGLFtools.TGLF(rhos=rhos) - _ = tglf.prep_using_tgyro(folder, cold_start=cold_start, inputgacode=inputgacode) + + _ = tglf.prep(p,folder,cold_start = cold_start) - code_settings = self.portals_parameters["transport"]["options"]["code_settings"] - extraOptions = self.portals_parameters["transport"]["options"]["extraOptionsTGLF"] + code_settings = self.portals_parameters["transport"]["options"]["tglf"]["run"]["code_settings"] + extraOptions = self.portals_parameters["transport"]["options"]["tglf"]["run"]["extraOptions"] return tglf, code_settings, extraOptions diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index c00a0608..8c9bb827 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -28,13 +28,16 @@ solution: - # Specification of radial locations (roa wins over rho, if provided) + # Specification of radial locations (if both provided, predicted_roa is used instead of predicted_rho) predicted_roa: null predicted_rho: [0.35, 0.55, 0.75, 0.875, 0.9] # Channels to be predicted (Options: ["te", "ti", "ne", "nZ", "w0"]) predicted_channels: ["te", "ti", "ne"] + # Run turbulent exchange as surrogate + turbulent_exchange_as_surrogate: false + # Impurity to do flux-matching for if nZ enabled (name of first impurity instance AFTER postprocessing), e.g. "W" trace_impurity: null @@ -51,10 +54,7 @@ solution: # ... ymax: 1.0 ymin: 1.0 - yminymax_atleast: null # Defines minimum range of exploration, e.g. [0,2], even if not achieved via ymin/ymax - - # DEPCRECATED - fixed_gradients: null + yminymax_atleast: null # Defines minimum range of exploration, e.g. [0,2] means the gradient range will be at least from 0 to 2.0 even if not achieved via the ymin/ymax specification # enforce_finite_aLT is used to be able to select ymin = 2.0 for ne but ensure that te, ti is at, e.g., enforce_finite_aLT = 0.95 enforce_finite_aLT: null @@ -67,16 +67,16 @@ solution: reevaluate_targets: 0 - # If False, remove full model folder after evaluation, to avoid large folders (e.g. in MAESTRO runs) - keep_full_model_folder: true + # DEPCRECATED, #TOREMOVE + fixed_gradients: null - # Run turbulent exchange as surrogate - turbulent_exchange_as_surrogate: false + # If False, remove full model folder after evaluation, to avoid carrying large folders (e.g. in MAESTRO runs) + keep_full_model_folder: true - # If True, fit model to GZ/nZ, valid on the trace limit + # If True, fit surrogate model to GZ/nZ instead of GZ, valid on the trace limit impurity_trick: true - # If not None, using fZ0_as_weight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis + # If provided, using fZ0_as_weight/fZ_0 as scaling factor for GZ, where fZ_0 is the original impurity concentration on axis fZ0_as_weight: null fImp_orig: 1.0 @@ -107,7 +107,11 @@ transport: # Select transport models (when instantiating the evaluator, assign these parameters) evaluator_instance_attributes: + + # Turbulent transport model turbulence_model: "tglf" + + # Neoclassical transport model neoclassical_model: "neo" # Simulation kwargs to be passed directly to run and read commands (defaults here) From c460c0281e88655a6114728f6fa1f9144554c245 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 30 Sep 2025 19:19:19 -0400 Subject: [PATCH 341/385] Extract postprocessing of MAESTRO-PORTALS such that it can be part of namelist --- .../maestro/scripts/run_maestro.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index 8b99f7b9..c96e0337 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -1,5 +1,6 @@ import argparse from pathlib import Path +from functools import partial from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import PROFILEStools from mitim_modules.maestro.MAESTROmain import maestro @@ -8,6 +9,15 @@ from mitim_tools.misc_tools import PLASMAtools from IPython import embed +def profiles_postprocessing_fun(file_profs, lumpImpurities = True, enforce_same_density_gradients = True): + p = PROFILEStools.gacode_state(file_profs) + if lumpImpurities: + p.lumpImpurities() + if enforce_same_density_gradients: + p.enforce_same_density_gradients() + p.write_state(file=file_profs) + return p + def parse_maestro_nml(file_path): # Extract engineering parameters, initializations, and desired beats to run maestro_namelist = IOtools.read_mitim_yaml(file_path) @@ -140,15 +150,7 @@ def parse_maestro_nml(file_path): enforce_same_density_gradients = maestro_namelist["maestro"]["portals_beat"]["transport_preprocessing"]["enforce_same_density_gradients"] # add postprocessing function - def profiles_postprocessing_fun(file_profs): - p = PROFILEStools.gacode_state(file_profs) - if lumpImpurities: - p.lumpImpurities() - if enforce_same_density_gradients: - p.enforce_same_density_gradients() - p.write_state(file=file_profs) - return p - beat_namelist['portals_parameters']['transport']['profiles_postprocessing_fun'] = profiles_postprocessing_fun + beat_namelist['portals_parameters']['transport']['profiles_postprocessing_fun'] = partial(profiles_postprocessing_fun, lumpImpurities=lumpImpurities, enforce_same_density_gradients=enforce_same_density_gradients) elif beat_type == "eped_initializer" and "eped_beat" in maestro_namelist["maestro"]: print('Using the eped_beat namelist for the eped_initializer') From 4de7e769f6c3a2e72dcc2b339e328889262cd625 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 30 Sep 2025 19:19:46 -0400 Subject: [PATCH 342/385] Bug fix of str vs int in variables --- src/mitim_modules/portals/scripts/run_portals.py | 2 +- src/mitim_tools/opt_tools/SURROGATEtools.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/portals/scripts/run_portals.py b/src/mitim_modules/portals/scripts/run_portals.py index 0b82d409..cf007c03 100644 --- a/src/mitim_modules/portals/scripts/run_portals.py +++ b/src/mitim_modules/portals/scripts/run_portals.py @@ -28,7 +28,7 @@ def main(): portals_fun = PORTALSmain.portals(folderWork, portals_namelist=portals_namelist) portals_fun.prep(inputgacode) - mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=cold_start, askQuestions=False) + mitim_bo = STRATEGYtools.MITIM_BO(portals_fun, cold_start=cold_start) mitim_bo.run() if __name__ == "__main__": diff --git a/src/mitim_tools/opt_tools/SURROGATEtools.py b/src/mitim_tools/opt_tools/SURROGATEtools.py index 5e60a6de..1aa726e9 100644 --- a/src/mitim_tools/opt_tools/SURROGATEtools.py +++ b/src/mitim_tools/opt_tools/SURROGATEtools.py @@ -298,6 +298,11 @@ def _select_transition_physics_based_params(self, ): self.surrogate_transformation_variables = None if ("surrogate_transformation_variables_alltimes" in self.surrogate_parameters) and (self.surrogate_parameters["surrogate_transformation_variables_alltimes"] is not None): + # Make sure that I can read both int or str as keys to surrogate_transformation_variables_alltimes + # Change the dictionary keys to be always integers + if not isinstance(list(self.surrogate_parameters["surrogate_transformation_variables_alltimes"].keys())[0], int): + self.surrogate_parameters["surrogate_transformation_variables_alltimes"] = {int(k): v for k, v in self.surrogate_parameters["surrogate_transformation_variables_alltimes"].items()} + transition_position = list(self.surrogate_parameters["surrogate_transformation_variables_alltimes"].keys())[ np.where(self.num_training_points < np.array(list(self.surrogate_parameters["surrogate_transformation_variables_alltimes"].keys())))[0][0]] From ee1c84fbd3bc7ab07a07db843078b3451fbbdf46 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 30 Sep 2025 20:29:10 -0400 Subject: [PATCH 343/385] [Experimental feature] Implementation of "ball" idea to reuse of TGLF scans --- .../physics_models/transport_tglf.py | 145 +++++++++++++++++- templates/namelist.portals.yaml | 3 + 2 files changed, 141 insertions(+), 7 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 3c63bd3a..00958c5b 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -1,3 +1,4 @@ +from pathlib import Path import numpy as np from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import TGLFtools @@ -21,6 +22,7 @@ def _evaluate_tglf(self, pass_info = True): Qi_includes_fast = simulation_options["Qi_includes_fast"] use_tglf_scan_trick = simulation_options["use_scan_trick_for_stds"] + reuse_scan_ball_file = self.powerstate.transport_options['folder'] / 'Outputs' / 'tglf_ball.npz' if simulation_options.get("reuse_scan_ball", False) else None cores_per_tglf_instance = simulation_options["cores_per_tglf_instance"] keep_tglf_files = simulation_options["keep_files"] percent_error = simulation_options["percent_error"] @@ -115,6 +117,7 @@ def _evaluate_tglf(self, pass_info = True): cores_per_tglf_instance=cores_per_tglf_instance, Qi_includes_fast=Qi_includes_fast, only_minimal_files=keep_tglf_files in ['minimal', 'base'], + reuse_scan_ball_file=reuse_scan_ball_file, **simulation_options["run"] ) @@ -186,6 +189,7 @@ def _run_tglf_uncertainty_model( cores_per_tglf_instance = 4, # e.g. 4 core per radius, since this is going to launch ~ Nr=5 x (Nv=6 x Nd=2 + 1) = 65 TGLFs at once Qi_includes_fast=False, only_minimal_files=True, # Since I only care about fluxes here, do not retrieve all the files + reuse_scan_ball_file=None, # If not None, it will reuse previous evaluations within the delta ball (to capture combinations) ): print(f"\t- Running TGLF standalone scans ({delta = }) to determine relative errors") @@ -256,7 +260,6 @@ def _run_tglf_uncertainty_model( Qe = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) Qi = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) - Qifast = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) Ge = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) GZ = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) Mt = np.zeros((len(rho_locations), len(variables_to_scan)*len(relative_scan) )) @@ -266,18 +269,18 @@ def _run_tglf_uncertainty_model( for vari in variables_to_scan: jump = tglf.scans[f'{name}_{vari}']['Qe'].shape[-1] + # Outputs Qe[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qe_gb'] - Qi[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qi_gb'] - Qifast[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qifast_gb'] + Qi[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Qi_gb'] + (0 if not Qi_includes_fast else tglf.scans[f'{name}_{vari}']['Qifast_gb']) Ge[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Ge_gb'] GZ[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Gi_gb'] Mt[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['Mt_gb'] S[:,cont:cont+jump] = tglf.scans[f'{name}_{vari}']['S_gb'] - cont += jump + cont += jump + if Qi_includes_fast: print(f"\t- Qi includes fast ions, adding their contribution") - Qi += Qifast # Add the base that was calculated earlier if Flux_base is not None: @@ -288,13 +291,16 @@ def _run_tglf_uncertainty_model( Mt = np.append(np.atleast_2d(Flux_base[4]).T, Mt, axis=1) S = np.append(np.atleast_2d(Flux_base[5]).T, S, axis=1) + if reuse_scan_ball_file is not None: + Qe, Qi, Ge, GZ, Mt, S = _ball_workflow(reuse_scan_ball_file, variables_to_scan, rho_locations, tglf, impurityPosition, Qi_includes_fast, Qe, Qi, Ge, GZ, Mt, S, delta_ball=delta) + # Calculate the standard deviation of the scans, that's going to be the reported stds def calculate_mean_std(Q): # Assumes Q is [radii, points], with [radii, 0] being the baseline - Qm = np.mean(Q, axis=1) - Qstd = np.std(Q, axis=1) + Qm = np.nanmean(Q, axis=1) + Qstd = np.nanstd(Q, axis=1) # Qm = Q[:,0] # Qstd = np.std(Q, axis=1) @@ -316,3 +322,128 @@ def calculate_mean_std(Q): Flux_std = [Qe_std, Qi_std, Ge_std, GZ_std, Mt_std, S_std] return Flux_mean, Flux_std + + +def _ball_workflow(file, variables_to_scan, rho_locations, tglf, impurityPosition, Qi_includes_fast, Qe, Qi, Ge, GZ, Mt, S, delta_ball=0.02): + ''' + Workflow to reuse previous TGLF evaluations within a delta ball to capture combinations + around the current base case. + ''' + + # Grab all inputs and outputs of the current run + input_params_keys = variables_to_scan + output_params_keys = ['Qe', 'Qi', 'Ge', 'Gi', 'Mt', 'S'] + input_params = np.zeros((len(rho_locations), len(input_params_keys), len(tglf.results))) + output_params = np.zeros((len(rho_locations), len(output_params_keys), len(tglf.results))) + input_params_base = np.zeros((len(rho_locations), len(input_params_keys))) + for i, key in enumerate(tglf.results.keys()): + for irho in range(len(rho_locations)): + + # Grab all inputs + for ikey in range(len(input_params_keys)): + input_params[irho, ikey, i] = tglf.results[key]['parsed'][irho][input_params_keys[ikey]] + if key == 'base': + input_params_base[irho, ikey] = tglf.results[key]['parsed'][irho][input_params_keys[ikey]] + + # Grab all outputs + output_params[irho, 0, i] = tglf.results[key]['output'][irho].Qe + output_params[irho, 1, i] = tglf.results[key]['output'][irho].Qi + (0 if not Qi_includes_fast else tglf.results[key]['output'][irho].Qifast) + output_params[irho, 2, i] = tglf.results[key]['output'][irho].Ge + output_params[irho, 3, i] = tglf.results[key]['output'][irho].GiAll[impurityPosition] + output_params[irho, 4, i] = tglf.results[key]['output'][irho].Mt + output_params[irho, 5, i] = tglf.results[key]['output'][irho].Se + + # -------------------------------------------------------------------------------------------------------- + # Read previous ball and append + # -------------------------------------------------------------------------------------------------------- + + if Path(file).exists(): + + print(f"\t- Reusing previous TGLF scan evaluations within the delta ball to capture combinations", typeMsg="i") + + # Grab ball + with np.load(file) as data: + rho_ball = data['rho'] + input_ball = data['input_params'] + output_ball = data['output_params'] + + precision_check = 1E-5 # I needed to add a small number to avoid numerical issues because TGLF input files have limited precision + + # Get the indeces of the points within the delta ball (condition in which all inputs are within the delta of the base case for that specific radius) + indices_to_grab = {} + for irho in range(len(rho_locations)): + indices_to_grab[irho] = [] + inputs_base = input_params_base[irho, :] + for icase in range(input_ball.shape[-1]): + inputs_case = input_ball[irho, :, icase] + + # Check if all inputs are within the delta ball (but not exactly equal, in case the ball has been run at the wrong time) + is_this_within_ball = True + for ikey in range(len(input_params_keys)): + val_current = inputs_base[ikey] + val_ball = inputs_case[ikey] + + # I need to have all inputs within the delta ball + is_this_within_ball = is_this_within_ball and ( abs(val_current-val_ball) <= abs(val_current*delta_ball) + precision_check ) + + if is_this_within_ball: + indices_to_grab[irho].append(icase) + + print(f"\t\t- Out of {input_ball.shape[-1]} points in file, found {len(indices_to_grab[irho])} at location {irho} within the delta ball ({delta_ball*100}%)", typeMsg="i") + + # Make an output_ball_select array equivalent to output_ball but only with the points within the delta ball (rest make them NaN) + output_ball_select = np.full_like(output_ball, np.nan) + for irho in range(len(rho_locations)): + for icase in indices_to_grab[irho]: + output_ball_select[irho, :, icase] = output_ball[irho, :, icase] + + # Append those points to the current run + Qe = np.append(Qe, output_ball_select[:, 0, :], axis=1) + Qi = np.append(Qi, output_ball_select[:, 1, :], axis=1) + Ge = np.append(Ge, output_ball_select[:, 2, :], axis=1) + GZ = np.append(GZ, output_ball_select[:, 3, :], axis=1) + Mt = np.append(Mt, output_ball_select[:, 4, :], axis=1) + S = np.append(S, output_ball_select[:, 5, :], axis=1) + + # Remove repeat points (for example when transitioning from simple relaxation initialization to full optimization) + def remove_duplicate_cases(*arrays): + """Remove duplicate cases (columns) from arrays of shape (rho_size, cases_size)""" + if len(arrays) == 0: + return arrays + + # Stack all arrays to create a combined signature for each case + combined = np.vstack(arrays) # Shape: (total_channels, cases_size) + + # Find unique cases, handling NaN values properly + # Use pandas for robust duplicate detection with NaN support + import pandas as pd + df = pd.DataFrame(combined.T) # Transpose so each row is a case + unique_indices = df.drop_duplicates().index.values + + print(f"\t\t- Removed {combined.shape[1] - len(unique_indices)} duplicate cases, keeping {len(unique_indices)} unique cases", typeMsg="i") + + # Return arrays with only unique cases + return tuple(arr[:, unique_indices] for arr in arrays) + + Qe, Qi, Ge, GZ, Mt, S = remove_duplicate_cases(Qe, Qi, Ge, GZ, Mt, S) + + + else: + rho_ball = np.array([]) + input_ball = np.array([]) + output_ball = np.array([]) + + # -------------------------------------------------------------------------------------------------------- + # Save new ball + # -------------------------------------------------------------------------------------------------------- + + # Append to the values read from previous ball + if rho_ball.shape[0] != 0: + input_params = np.append(input_ball, input_params, axis=2) + output_params = np.append(output_ball, output_params, axis=2) + + # Save the new ball + np.savez(file, rho=rho_locations, input_params=input_params, output_params=output_params) + print(f"\t- Saved updated ball with {input_params.shape[-1]} points to {IOtools.clipstr(file)}", typeMsg="i") + + return Qe, Qi, Ge, GZ, Mt, S \ No newline at end of file diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index 8c9bb827..442ef5de 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -133,6 +133,9 @@ transport: # If not None, use TGLF scan trick to calculate TGLF errors with this maximum delta use_scan_trick_for_stds: 0.02 + # [EXPERIMENTAL] If True, store previous evaluations and reuse them if they are within the delta of all inputs (to capture combinations) + reuse_scan_ball: false + # Files to keep from simulation (minimal: only retrieve minimal files always, none: always minimal files, base: retrieve all files) keep_files: "base" From 2e38d2877f5adabcf6c4ee9dec53c847fd8e8098 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Tue, 30 Sep 2025 20:30:43 -0400 Subject: [PATCH 344/385] Bug fix PORTALS --- src/mitim_modules/maestro/utils/PORTALSbeat.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/utils/PORTALSbeat.py b/src/mitim_modules/maestro/utils/PORTALSbeat.py index 30603fef..11276430 100644 --- a/src/mitim_modules/maestro/utils/PORTALSbeat.py +++ b/src/mitim_modules/maestro/utils/PORTALSbeat.py @@ -91,7 +91,8 @@ def run(self, **kwargs): # Update the namelist with the parameters in the MAESTRO namelist (variable: portals_parameters) portals_fun.portals_parameters = IOtools.deep_dict_update(portals_fun.portals_parameters, self.portals_parameters) - portals_fun.portals_parameters['optimization_options'] = portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.portals_parameters['optimization_options']) + if 'optimization_options' in self.portals_parameters: + portals_fun.portals_parameters['optimization_options'] = portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.portals_parameters['optimization_options']) # MAESTRO beat may receive optimization options changes from previous beats, so allow that too portals_fun.portals_parameters['optimization_options'] = portals_fun.optimization_options = IOtools.deep_dict_update(portals_fun.optimization_options, self.optimization_options) From 6bac6931d82a2ff9ae43dc3778983daa3854abc8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 1 Oct 2025 10:55:09 -0400 Subject: [PATCH 345/385] Cleaned up ball_workflow for more descriptive prints --- .../physics_models/transport_tglf.py | 82 ++++++++++++------- 1 file changed, 52 insertions(+), 30 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 00958c5b..2e9ae17f 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -1,5 +1,6 @@ from pathlib import Path import numpy as np +import pandas as pd from mitim_tools.misc_tools import IOtools from mitim_tools.gacode_tools import TGLFtools from mitim_tools.misc_tools.LOGtools import printMsg as print @@ -324,7 +325,7 @@ def calculate_mean_std(Q): return Flux_mean, Flux_std -def _ball_workflow(file, variables_to_scan, rho_locations, tglf, impurityPosition, Qi_includes_fast, Qe, Qi, Ge, GZ, Mt, S, delta_ball=0.02): +def _ball_workflow(file, variables_to_scan, rho_locations, tglf, impurityPosition, Qi_includes_fast, Qe_orig, Qi_orig, Ge_orig, GZ_orig, Mt_orig, S_orig, delta_ball=0.02): ''' Workflow to reuse previous TGLF evaluations within a delta ball to capture combinations around the current base case. @@ -332,20 +333,26 @@ def _ball_workflow(file, variables_to_scan, rho_locations, tglf, impurityPositio # Grab all inputs and outputs of the current run input_params_keys = variables_to_scan - output_params_keys = ['Qe', 'Qi', 'Ge', 'Gi', 'Mt', 'S'] input_params = np.zeros((len(rho_locations), len(input_params_keys), len(tglf.results))) - output_params = np.zeros((len(rho_locations), len(output_params_keys), len(tglf.results))) + input_params_base = np.zeros((len(rho_locations), len(input_params_keys))) + + output_params_keys = ['Qe', 'Qi', 'Ge', 'Gi', 'Mt', 'S'] + output_params = np.zeros((len(rho_locations), len(output_params_keys), len(tglf.results))) + for i, key in enumerate(tglf.results.keys()): for irho in range(len(rho_locations)): - # Grab all inputs + # Grab all inputs in array with shape (Nr, Ninputs, Ncases) for ikey in range(len(input_params_keys)): input_params[irho, ikey, i] = tglf.results[key]['parsed'][irho][input_params_keys[ikey]] - if key == 'base': + + # Grab base inputs in array with shape (Nr, Ninputs) + if key == 'base': + for ikey in range(len(input_params_keys)): input_params_base[irho, ikey] = tglf.results[key]['parsed'][irho][input_params_keys[ikey]] - # Grab all outputs + # Grab all outputs in array with shape (Nr, Noutputs, Ncases) output_params[irho, 0, i] = tglf.results[key]['output'][irho].Qe output_params[irho, 1, i] = tglf.results[key]['output'][irho].Qi + (0 if not Qi_includes_fast else tglf.results[key]['output'][irho].Qifast) output_params[irho, 2, i] = tglf.results[key]['output'][irho].Ge @@ -359,9 +366,9 @@ def _ball_workflow(file, variables_to_scan, rho_locations, tglf, impurityPositio if Path(file).exists(): - print(f"\t- Reusing previous TGLF scan evaluations within the delta ball to capture combinations", typeMsg="i") + print(f"\t- Reusing previous TGLF scan evaluations within the delta ball to capture combinations") - # Grab ball + # Grab ball contents with np.load(file) as data: rho_ball = data['rho'] input_ball = data['input_params'] @@ -389,49 +396,64 @@ def _ball_workflow(file, variables_to_scan, rho_locations, tglf, impurityPositio if is_this_within_ball: indices_to_grab[irho].append(icase) - print(f"\t\t- Out of {input_ball.shape[-1]} points in file, found {len(indices_to_grab[irho])} at location {irho} within the delta ball ({delta_ball*100}%)", typeMsg="i") + print(f"\t\t- Out of {input_ball.shape[-1]} points in file, found {len(indices_to_grab[irho])} at location {irho} within the delta ball ({delta_ball*100}%)", typeMsg="i" if len(indices_to_grab[irho]) > 0 else "") # Make an output_ball_select array equivalent to output_ball but only with the points within the delta ball (rest make them NaN) output_ball_select = np.full_like(output_ball, np.nan) for irho in range(len(rho_locations)): for icase in indices_to_grab[irho]: output_ball_select[irho, :, icase] = output_ball[irho, :, icase] - - # Append those points to the current run - Qe = np.append(Qe, output_ball_select[:, 0, :], axis=1) - Qi = np.append(Qi, output_ball_select[:, 1, :], axis=1) - Ge = np.append(Ge, output_ball_select[:, 2, :], axis=1) - GZ = np.append(GZ, output_ball_select[:, 3, :], axis=1) - Mt = np.append(Mt, output_ball_select[:, 4, :], axis=1) - S = np.append(S, output_ball_select[:, 5, :], axis=1) + + # Append those points to the current run (these will always have shape (Nr, Ncases+original) but those cases that were not in the ball will be NaN) + # The reason to do it this way is that I want to keep it as a uniform shape to be able to calculate stds later, and I would risk otherwise having different shapes per radius + Qe = np.append(Qe_orig, output_ball_select[:, 0, :], axis=1) + Qi = np.append(Qi_orig, output_ball_select[:, 1, :], axis=1) + Ge = np.append(Ge_orig, output_ball_select[:, 2, :], axis=1) + GZ = np.append(GZ_orig, output_ball_select[:, 3, :], axis=1) + Mt = np.append(Mt_orig, output_ball_select[:, 4, :], axis=1) + S = np.append(S_orig, output_ball_select[:, 5, :], axis=1) + print(f"\t\t>>> Flux arrays have shape {Qe.shape} after appending ball points (NaNs are added to those locations and cases that did not fall within delta ball)") # Remove repeat points (for example when transitioning from simple relaxation initialization to full optimization) def remove_duplicate_cases(*arrays): - """Remove duplicate cases (columns) from arrays of shape (rho_size, cases_size)""" - if len(arrays) == 0: - return arrays + """Remove duplicate cases (columns) from arrays of shape (rho_size, cases_size) + Returns: + tuple: (unique_arrays, duplicate_arrays) where each is a tuple of arrays + """ + # Stack all arrays to create a combined signature for each case combined = np.vstack(arrays) # Shape: (total_channels, cases_size) - - # Find unique cases, handling NaN values properly - # Use pandas for robust duplicate detection with NaN support - import pandas as pd + + # Find unique cases, handling NaN values properly (Use pandas for robust duplicate detection with NaN support) df = pd.DataFrame(combined.T) # Transpose so each row is a case unique_indices = df.drop_duplicates().index.values + nan_indices = np.where(np.all(np.isnan(combined), axis=0))[0] + unique_notnan_indeces = [idx for idx in unique_indices if idx not in nan_indices] + all_indices = np.arange(combined.shape[1]) + duplicate_indices = np.setdiff1d(all_indices, unique_notnan_indeces) + + print(f"\t\t* Removed {len(duplicate_indices)} duplicate / all-nan cases, keeping {len(unique_notnan_indeces)} unique cases", typeMsg="i") - print(f"\t\t- Removed {combined.shape[1] - len(unique_indices)} duplicate cases, keeping {len(unique_indices)} unique cases", typeMsg="i") + # Return arrays with unique cases and duplicate cases + unique_arrays = tuple(arr[:, unique_notnan_indeces] for arr in arrays) + duplicate_arrays = tuple(arr[:, duplicate_indices] for arr in arrays) - # Return arrays with only unique cases - return tuple(arr[:, unique_indices] for arr in arrays) + return unique_arrays, duplicate_arrays - Qe, Qi, Ge, GZ, Mt, S = remove_duplicate_cases(Qe, Qi, Ge, GZ, Mt, S) + unique_results, duplicate_results = remove_duplicate_cases(Qe, Qi, Ge, GZ, Mt, S) + + Qe, Qi, Ge, GZ, Mt, S = unique_results + + print(f"\t\t>>> Flux arrays have shape {Qe.shape} after finding unique points") - else: + rho_ball = np.array([]) input_ball = np.array([]) output_ball = np.array([]) + + Qe, Qi, Ge, GZ, Mt, S = Qe_orig, Qi_orig, Ge_orig, GZ_orig, Mt_orig, S_orig # -------------------------------------------------------------------------------------------------------- # Save new ball @@ -446,4 +468,4 @@ def remove_duplicate_cases(*arrays): np.savez(file, rho=rho_locations, input_params=input_params, output_params=output_params) print(f"\t- Saved updated ball with {input_params.shape[-1]} points to {IOtools.clipstr(file)}", typeMsg="i") - return Qe, Qi, Ge, GZ, Mt, S \ No newline at end of file + return Qe, Qi, Ge, GZ, Mt, S From 0907212ca8f58ca4db294e689d4fc1454072f238 Mon Sep 17 00:00:00 2001 From: Xavior Date: Wed, 1 Oct 2025 15:41:57 -0400 Subject: [PATCH 346/385] GX should be, by default, only prepared, for now --- templates/namelist.portals.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index 442ef5de..1fa03cd5 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -196,7 +196,7 @@ transport: extraOptions: {} # Run type: normal (submit and wait), submit (submit and do not wait), prep (do not submit) - run_type: "normal" + run_type: "prep" read: tmin: 0.0 From 6a4bb016e4993ffaec328f5f20563bd971f11974 Mon Sep 17 00:00:00 2001 From: Xavior Date: Wed, 1 Oct 2025 15:42:25 -0400 Subject: [PATCH 347/385] Do not raise warnings about GPUs if I'm only preparing GX inputs, not running it --- src/mitim_tools/gacode_tools/CGYROtools.py | 2 +- src/mitim_tools/gacode_tools/NEOtools.py | 2 +- src/mitim_tools/gacode_tools/TGLFtools.py | 2 +- src/mitim_tools/simulation_tools/SIMtools.py | 3 ++- src/mitim_tools/simulation_tools/physics/GXtools.py | 9 ++++++--- 5 files changed, 11 insertions(+), 7 deletions(-) diff --git a/src/mitim_tools/gacode_tools/CGYROtools.py b/src/mitim_tools/gacode_tools/CGYROtools.py index d030d3b3..8a52d7a1 100644 --- a/src/mitim_tools/gacode_tools/CGYROtools.py +++ b/src/mitim_tools/gacode_tools/CGYROtools.py @@ -22,7 +22,7 @@ def __init__( def code_call(folder, p, n = 1, nomp = 1, additional_command="", **kwargs): return f"cgyro -e {folder} -n {n} -nomp {nomp} -p {p} {additional_command}" - def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): + def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None, **kwargs_slurm): slurm_settings = { "name": name, diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index 7db02cfc..ddf9d46f 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -18,7 +18,7 @@ def __init__( def code_call(folder, n, p, additional_command="", **kwargs): return f"neo -e {folder} -n {n} -p {p} {additional_command}" - def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): + def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None, **kwargs_slurm): slurm_settings = { "name": name, diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 498a6b36..002dc9c9 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -128,7 +128,7 @@ def __init__( def code_call(folder, p, n = 1, additional_command="", **kwargs): return f"tglf -e {folder} -n {n} -p {p} {additional_command}" - def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): + def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None, **kwargs_slurm): slurm_settings = { "name": name, diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index cba68503..a22530b8 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -527,7 +527,8 @@ def _run( total_cores_required=total_cores_required, cores_per_code_call=cores_per_code_call, type_of_submission=type_of_submission, - array_list=array_list if type_of_submission == "slurm_array" else None + array_list=array_list if type_of_submission == "slurm_array" else None, + raise_warning= run_type == 'normal' ) self.simulation_job.define_machine( diff --git a/src/mitim_tools/simulation_tools/physics/GXtools.py b/src/mitim_tools/simulation_tools/physics/GXtools.py index c407aefe..8739b5c0 100644 --- a/src/mitim_tools/simulation_tools/physics/GXtools.py +++ b/src/mitim_tools/simulation_tools/physics/GXtools.py @@ -20,7 +20,7 @@ def __init__( def code_call(folder, n, p, additional_command="", **kwargs): return f"gx -n {n} {folder}/gxplasma.in > {folder}/gxplasma.mitim.log" - def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, array_list=None): + def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call, type_of_submission, raise_warning=True,array_list=None): slurm_settings = { "name": name, @@ -31,8 +31,11 @@ def code_slurm_settings(name, minutes, total_cores_required, cores_per_code_call machineSettings = CONFIGread.machineSettings(code='gx') if machineSettings['gpus_per_node'] == 0: - raise Exception("[MITIM] GX needs GPUs to run, but the selected machine does not have any GPU configured. Please select another machine in the config file with gpus_per_node>0.") - + if raise_warning: + raise Exception("[MITIM] GX needs GPUs to run, but the selected machine does not have any GPU configured. Please select another machine in the config file with gpus_per_node>0.") + else: + print("[MITIM] Warning: GX needs GPUs to run, but the selected machine does not have any GPU configured. Running without GPUs, but this will likely fail.", typeMsg="w") + if type_of_submission == "slurm_standard": slurm_settings['ntasks'] = total_cores_required From 93e29860b4d1458c6b0697b33733325b6ed95d90 Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Wed, 1 Oct 2025 17:32:25 -0400 Subject: [PATCH 348/385] pyproject: added megpy as dependency --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 23d135ba..7184a5f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ dependencies = [ "pyyaml" # "vmecpp", # "quends @ git+https://github.com/sandialabs/quends.git" + "megpy", ] [project.optional-dependencies] From c81e0fd99b7dd26c4c9fb475403a38f4889f359c Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Wed, 1 Oct 2025 17:33:20 -0400 Subject: [PATCH 349/385] mitim_tools: updated gs_tools/GEQtools to use megpy equilibrium processing --- src/mitim_tools/gs_tools/GEQtools.py | 175 ++++++++++----------------- 1 file changed, 66 insertions(+), 109 deletions(-) diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 5c71f12d..a4401efa 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -8,68 +8,28 @@ from mitim_tools.gacode_tools import PROFILEStools from mitim_tools.gs_tools.utils import GEQplotting from shapely.geometry import LineString -from scipy.integrate import quad +from scipy.integrate import quad, cumulative_trapezoid +import megpy import freegs from freegs import geqdsk from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed """ -Note that this module relies on OMFIT classes (https://omfit.io/classes.html) procedures to intrepret the content of g-eqdsk files. -Modifications are made in MITINM for visualizations and a few extra derivations. +Note that this module relies on megpy to intrepret the content of g-eqdsk files. +Modifications are made in MITIM for visualizations and a few extra derivations. """ -def fix_file(filename): - - with open(filename, "r") as f: - lines = f.readlines() - - # ----------------------------------------------------------------------- - # Remove coils (chatGPT 4o as of 08/24/24) - # ----------------------------------------------------------------------- - # Use StringIO to simulate the file writing - noCoils_file = io.StringIO() - for cont, line in enumerate(lines): - if cont > 0 and line[:2] == " ": - break - noCoils_file.write(line) - - # Reset cursor to the start of StringIO - noCoils_file.seek(0) - - # Write the StringIO content to a temporary file - with tempfile.NamedTemporaryFile(delete=False) as tmp_file: - tmp_file.write(noCoils_file.getvalue().encode('utf-8')) - noCoils_file = tmp_file.name - # ----------------------------------------------------------------------- - - with open(filename, 'r') as file1, open(noCoils_file, 'r') as file2: - file1_content = file1.read() - file2_content = file2.read() - - if file1_content != file2_content: - print(f"\t- geqdsk file {IOtools.clipstr(filename)} had coils, I have removed them") - - filename = noCoils_file - - return filename - class MITIMgeqdsk: def __init__(self, filename): - # Fix file by removing coils if it has them - filename = fix_file(filename) - - # Read GEQDSK file using OMFIT - import omfit_classes.omfit_eqdsk - self.g = omfit_classes.omfit_eqdsk.OMFITgeqdsk(filename, forceFindSeparatrix=True) + self.g = megpy.Equilibrium() + self.g.read_geqdsk(f_path=filename) + self.g.add_derived(incl_fluxsurfaces=True, analytic_shape=True, incl_B=True) # Extra derivations in MITIM self.derive() - # Remove temporary file - os.remove(filename) - @classmethod def timeslices(cls, filename, **kwargs): print("\n...Opening GEQ file with several time slices") @@ -111,61 +71,65 @@ def timeslices(cls, filename, **kwargs): def derive(self, debug=False): - self.Jt = self.g.surfAvg("Jt") * 1e-6 - self.Jt_fb = self.g.surfAvg("Jt_fb") * 1e-6 + zero_vector = np.zeros(self.g.derived["rho_pol"].shape) + + self.rho_pol = self.g.derived["rho_pol"].copy() + self.rho_tor = self.g.derived["rho_tor"].copy() + self.psi_pol_norm = self.rho_pol ** 2 + self.psi_tor_norm = self.rho_tor ** 2 + + self.Jt = self.g.derived["j_tor"] * 1e-6 + self.Jt_fb = zero_vector.copy() self.Jerror = np.abs(self.Jt - self.Jt_fb) - self.Ip = self.g["CURRENT"] + self.Ip = self.g.raw["current"] # Parameterizations of LCFS - self.kappa = self.g["fluxSurfaces"]["geo"]["kap"][-1] - self.kappaU = self.g["fluxSurfaces"]["geo"]["kapu"][-1] - self.kappaL = self.g["fluxSurfaces"]["geo"]["kapl"][-1] - - self.delta = self.g["fluxSurfaces"]["geo"]["delta"][-1] - self.deltaU = self.g["fluxSurfaces"]["geo"]["dell"][-1] - self.deltaL = self.g["fluxSurfaces"]["geo"]["dell"][-1] - - self.zeta = self.g["fluxSurfaces"]["geo"]["zeta"][-1] - - self.a = self.g["fluxSurfaces"]["geo"]["a"][-1] - self.Rmag = self.g["fluxSurfaces"]["geo"]["R"][0] - self.Zmag = self.g["fluxSurfaces"]["geo"]["Z"][0] - self.Rmajor = np.mean( - [ - self.g["fluxSurfaces"]["geo"]["Rmin_centroid"][-1], - self.g["fluxSurfaces"]["geo"]["Rmax_centroid"][-1], - ] - ) + self.kappa = self.g.derived["miller_geo"]["kappa"][-1] + self.kappaU = self.g.derived["miller_geo"]["kappa_u"][-1] + self.kappaL = self.g.derived["miller_geo"]["kappa_l"][-1] + + self.delta = self.g.derived["miller_geo"]["delta"][-1] + self.deltaU = self.g.derived["miller_geo"]["delta_u"][-1] + self.deltaL = self.g.derived["miller_geo"]["delta_l"][-1] + + self.zeta = self.g.derived["miller_geo"]["zeta"][-1] + + self.a = self.g.derived["r"][-1] + self.Rmag = self.g.derived["Ro"][0] + self.Zmag = self.g.derived["Zo"][0] - self.Zmajor = self.Zmag + self.Rmajor = self.g.derived["Ro"][-1] + self.Zmajor = self.Zmag #self.g.derived["Zo"][-1] TODO: check which Z0 to use, perhaps switch to MXH definition self.eps = self.a / self.Rmajor # Core values - - self.kappa_a = self.g["fluxSurfaces"]["geo"]["cxArea"][-1] / (np.pi * self.a**2) + vp = np.array(self.g.fluxsurfaces["Vprime"]).flatten() + ir = np.array(self.g.fluxsurfaces["1/R"]).flatten() + self.cx_area = cumulative_trapezoid(vp * ir, self.g.derived["psi"], initial=0.0) + self.kappa_a = self.cx_area[-1] / (np.pi * self.a**2) self.kappa995 = np.interp( 0.995, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["kap"], + self.psi_pol_norm, + self.g.derived["miller_geo"]["kappa"], ) self.kappa95 = np.interp( 0.95, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["kap"], + self.psi_pol_norm, + self.g.derived["miller_geo"]["kappa"], ) self.delta995 = np.interp( 0.995, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["delta"], + self.psi_pol_norm, + self.g.derived["miller_geo"]["delta"], ) self.delta95 = np.interp( 0.95, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["delta"], + self.psi_pol_norm, + self.g.derived["miller_geo"]["delta"], ) """ @@ -173,26 +137,23 @@ def derive(self, debug=False): Boundary -------------------------------------------------------------------------------------------------------------------------------------- Note that the RBBS and ZBBS values in the gfile are often too scattered and do not reproduce the boundary near x-points. - The shaping parameters calculated using fluxSurfaces are correct though. + The shaping parameters calculated using fluxsurfaces are correct though. """ - self.Rb_gfile, self.Yb_gfile = self.g["RBBBS"], self.g["ZBBBS"] - self.Rb, self.Yb = self.g["fluxSurfaces"].sep.transpose() + self.Rb_gfile, self.Yb_gfile = self.g.raw["rbbbs"].copy(), self.g.raw["zbbbs"].copy() + self.Rb, self.Yb = self.g.fluxsurfaces["R"][-1], self.g.fluxsurfaces["Z"][-1] if len(self.Rb) == 0: - print("\t- MITIM > No separatrix found in the OMFIT fluxSurfaces, increasing resolution and going all in!",typeMsg='i') + print("\t- MITIM > No separatrix found in the megpy fluxsurfaces, using explicit boundary in g-eqdsk file!",typeMsg='i') - flx = copy.deepcopy(self.g['fluxSurfaces']) - flx._changeResolution(6) - flx.findSurfaces([0.0,0.5,1.0]) - fs = flx['flux'][list(flx['flux'].keys())[-1]] - self.Rb, self.Yb = fs['R'], fs['Z'] + self.Rb = self.Rb_gfile.copy() + self.Yb = self.Yb_gfile.copy() if debug: fig, ax = plt.subplots() - # OMFIT - ax.plot(self.Rb, self.Yb, "-s", c="r", label="OMFIT") + # megpy + ax.plot(self.Rb, self.Yb, "-s", c="r", label="megpy") # GFILE ax.plot(self.Rb_gfile, self.Yb_gfile, "-s", c="y", label="GFILE") @@ -245,10 +206,7 @@ def write(self, filename=None): If filename is None, use the original one """ - if filename is not None: - self.g.filename = filename - - self.g.save() + self.g.write_geqdsk(f_path=filename) # ----------------------------------------------------------------------------- # Parameterizations @@ -256,8 +214,8 @@ def write(self, filename=None): def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): - psis = self.g["AuxQuantities"]["PSI_NORM"] - flux_surfaces = self.g['fluxSurfaces']['flux'] + psis = self.psi_pol_norm + flux_surfaces = self.g.fluxsurfaces["psi"] # Cannot parallelize because different number of points? kappa, rmin, rmaj, zmag, sn, cn = [],[],[],[],[],[] @@ -266,7 +224,7 @@ def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): if flux == len(flux_surfaces)-1: Rf, Zf = self.Rb, self.Yb else: - Rf, Zf = flux_surfaces[flux]['R'],flux_surfaces[flux]['Z'] + Rf, Zf = self.g.fluxsurfaces["R"][flux], self.g.fluxsurfaces["Z"][flux] # Perform the MXH decompositionusing the MITIM surface class surfaces = mitim_flux_surfaces() @@ -311,24 +269,22 @@ def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): # ----------------------------------------------------------------------------- # For MAESTRO and TRANSP converstions # ----------------------------------------------------------------------------- - def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH = 7, plotYN = False): # ------------------------------------------------------------------------------------------------------- # Quantities from the equilibrium # ------------------------------------------------------------------------------------------------------- - rhotor = self.g['RHOVN'] - psi = self.g['AuxQuantities']['PSI'] # Wb/rad - torfluxa = self.g['AuxQuantities']['PHI'][-1] / (2*np.pi) # Wb/rad - q = self.g['QPSI'] - pressure = self.g['PRES'] # Pa - Ip = self.g['CURRENT']*1E-6 # MA + rhotor = self.g.derived['rho_tor'] + psi = self.g.derived['psi'] # Wb/rad + torfluxa = self.g.derived['phi'][-1] / (2*np.pi) # Wb/rad + q = self.g.raw['qpsi'] + pressure = self.g.raw['pres'] # Pa + Ip = self.g.raw['current']*1E-6 # MA RZ = np.array([self.Rb,self.Yb]).T R0 = (RZ.max(axis=0)[0] + RZ.min(axis=0)[0])/2 - - B0 = self.g['RCENTR']*self.g['BCENTR'] / R0 + B0 = self.g.raw['rcentr']*self.g.raw['bcentr'] / R0 # Ensure positive quantities #TODO: Check if this is necessary, pass directions rhotor = np.array([np.abs(i) for i in rhotor]) @@ -543,6 +499,7 @@ def _to_mxh(self, n_coeff=6): self.cn = np.zeros((self.R.shape[0],n_coeff)) self.sn = np.zeros((self.R.shape[0],n_coeff)) self.gn = np.zeros((self.R.shape[0],4)) + self.gn[-1] = 1.0 for i in range(self.R.shape[0]): self.cn[i,:], self.sn[i,:], self.gn[i,:] = from_RZ_to_mxh(self.R[i,:], self.Z[i,:], n_coeff=n_coeff) @@ -562,8 +519,8 @@ def _to_miller(self): # Elongations - self.kappa_u = (Zmax - self.Z0) / self.a - self.kappa_l = (self.Z0 - Zmin) / self.a + self.kappa_u = (Zmax - self.Z0) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) + self.kappa_l = (self.Z0 - Zmin) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) self.kappa = (self.kappa_u + self.kappa_l) / 2 # Triangularities From 1086453738f10073b5249ca75730846361949297 Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Wed, 1 Oct 2025 17:53:26 -0400 Subject: [PATCH 350/385] mitim_tools: changed gs_tools/utils/GEQplotting to use megpy values --- src/mitim_tools/gs_tools/utils/GEQplotting.py | 341 +++++++++--------- 1 file changed, 174 insertions(+), 167 deletions(-) diff --git a/src/mitim_tools/gs_tools/utils/GEQplotting.py b/src/mitim_tools/gs_tools/utils/GEQplotting.py index fe6531b4..5c4c950b 100644 --- a/src/mitim_tools/gs_tools/utils/GEQplotting.py +++ b/src/mitim_tools/gs_tools/utils/GEQplotting.py @@ -4,6 +4,7 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +#TODO: add current profiles and flux-surface average fields to megpy and restore plots def compareGeqdsk(geqdsks, fn=None, extraLabel="", plotAll=True, labelsGs=None): @@ -229,8 +230,8 @@ def plotFS(self, axs=None, color="b", label=""): ax.set_ylabel("Z (m)") ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["AuxQuantities"]["RHO"] + x = self.psi_pol_norm + y = self.rho_tor ax.plot(x, y, lw=2, ls="-", c=color, label=label) ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) @@ -240,8 +241,8 @@ def plotFS(self, axs=None, color="b", label=""): ax.set_ylim([0, 1]) ax = axs[3] - x = self.g["AuxQuantities"]["RHO"] - y = self.g["AuxQuantities"]["RHOp"] + x = self.rho_tor + y = self.rho_pol ax.plot(x, y, lw=2, ls="-", c=color) ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) @@ -256,8 +257,8 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): fig, axs = plt.subplots(ncols=10) ax = axs[0] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jr") * 1e-6 + x = self.psi_pol_norm + y = np.zeros(x.shape) ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_ylabel("FSA $\\langle J\\rangle$ ($MA/m^2$)") @@ -267,14 +268,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[1] - plot2Dquantity(self, - ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 - ) + #ax = axs[1] + #plot2Dquantity(self, + # ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 + #) ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jz") * 1e-6 + x = self.psi_pol_norm + y = np.zeros(x.shape) ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -282,14 +283,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[3] - plot2Dquantity(self, - ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 - ) + #ax = axs[3] + #plot2Dquantity(self, + # ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 + #) ax = axs[4] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jt") * 1e-6 + x = self.psi_pol_norm + y = self.Jt ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -297,14 +298,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[5] - plot2Dquantity(self, - ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 - ) + #ax = axs[5] + #plot2Dquantity(self, + # ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 + #) ax = axs[6] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jp") * 1e-6 + x = self.psi_pol_norm + y = np.zeros(x.shape) ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -312,14 +313,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[7] - plot2Dquantity(self, - ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 - ) + #ax = axs[7] + #plot2Dquantity(self, + # ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 + #) ax = axs[8] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jpar") * 1e-6 + x = self.psi_pol_norm + y = np.zeros(x.shape) ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -327,10 +328,10 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[9] - plot2Dquantity(self, - ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 - ) + #ax = axs[9] + #plot2Dquantity(self, + # ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 + #) def plotFields(self, axs=None, zlims_thr=[-1, 1]): if axs is None: @@ -338,8 +339,8 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): fig, axs = plt.subplots(ncols=10) ax = axs[0] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Br") + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Br") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_ylabel("FSA $\\langle B\\rangle$ ($T$)") @@ -351,12 +352,12 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[1] plot2Dquantity(self, - ax=ax, var="Br", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="B_r", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" ) ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Bz") + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Bz") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -366,12 +367,12 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[3] plot2Dquantity(self, - ax=ax, var="Bz", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="B_z", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" ) ax = axs[4] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Bt") + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Bt") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -380,14 +381,14 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): # zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[5] - plot2Dquantity(self, - ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" - ) + #ax = axs[5] + #plot2Dquantity(self, + # ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" + #) ax = axs[6] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Bp") + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Bp") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -397,32 +398,32 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[7] plot2Dquantity(self, - ax=ax, var="Bp", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="B_pol_rz", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" ) ax = axs[8] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["avg"]["Bp**2"] + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g["fluxSurfaces"]["avg"]["Bp**2"] ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) ax.set_ylabel("$\\langle B_{\\theta}^2\\rangle$") - ax = axs[9] - x = self.g["fluxSurfaces"]["midplane"]["R"] - y = self.g["fluxSurfaces"]["midplane"]["Bt"] - ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") - y = self.g["fluxSurfaces"]["midplane"]["Bp"] - ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") - y = self.g["fluxSurfaces"]["midplane"]["Bz"] - ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") - y = self.g["fluxSurfaces"]["midplane"]["Br"] - ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") - y = self.g["fluxSurfaces"]["geo"]["bunit"] - ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") - ax.set_xlabel("$R$ LF midplane") - ax.set_ylabel("$B$ (T)") - ax.legend() + #ax = axs[9] + #x = self.g["fluxSurfaces"]["midplane"]["R"] + #y = self.g["fluxSurfaces"]["midplane"]["Bt"] + #ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") + #y = self.g["fluxSurfaces"]["midplane"]["Bp"] + #ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") + #y = self.g["fluxSurfaces"]["midplane"]["Bz"] + #ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") + #y = self.g["fluxSurfaces"]["midplane"]["Br"] + #ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") + #y = self.g["fluxSurfaces"]["geo"]["bunit"] + #ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") + #ax.set_xlabel("$R$ LF midplane") + #ax.set_ylabel("$B$ (T)") + #ax.legend() def plotChecks(self, axs=None): if axs is None: @@ -430,7 +431,7 @@ def plotChecks(self, axs=None): fig, axs = plt.subplots(ncols=8) ax = axs[0] - x = self.g["AuxQuantities"]["PSI_NORM"] + x = self.psi_pol_norm y1 = self.Jt ax.plot(x, np.abs(y1), lw=2, ls="-", c="b", label="$\\langle Jt\\rangle$") zmax = y1.max() @@ -463,10 +464,10 @@ def plotChecks(self, axs=None): ax.legend() ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y1 = self.g["FFPRIM"] + x = self.psi_pol_norm + y1 = self.g.raw["ffprim"] ax.plot(x, y1, lw=2, ls="-", c="r", label="$FF'$") - y2 = self.g["PPRIME"] * (4 * np.pi * 1e-7) + y2 = self.g.raw["pprime"] * (4 * np.pi * 1e-7) ax.plot(x, y2, lw=2, ls="-", c="b", label="$p'*\\mu_0$") ax.set_ylabel("") @@ -474,40 +475,40 @@ def plotChecks(self, axs=None): ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) - ax = axs[3] - plot2Dquantity(self, - ax=ax, - var="Jt", - title="Toroidal Current Jt", - zlims=[zmin, zmax], - cmap="viridis", - factor=1e-6, - ) - - ax = axs[4] - plot2Dquantity(self, - ax=ax, - var="Jt_fb", - title="Toroidal Current Jt (FB)", - zlims=[zmin, zmax], - cmap="viridis", - factor=1e-6, - ) - - ax = axs[5] - z = ( - np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) - * 1e-6 - ) - zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) - plot2Dquantity(self, - ax=ax, - var=z, - title="Absolute Error", - zlims=[0, zmaxx], - cmap="viridis", - direct=True, - ) + #ax = axs[3] + #plot2Dquantity(self, + # ax=ax, + # var="Jt", + # title="Toroidal Current Jt", + # zlims=[zmin, zmax], + # cmap="viridis", + # factor=1e-6, + #) + + #ax = axs[4] + #plot2Dquantity(self, + # ax=ax, + # var="Jt_fb", + # title="Toroidal Current Jt (FB)", + # zlims=[zmin, zmax], + # cmap="viridis", + # factor=1e-6, + #) + + #ax = axs[5] + #z = ( + # np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) + # * 1e-6 + #) + #zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) + #plot2Dquantity(self, + # ax=ax, + # var=z, + # title="Absolute Error", + # zlims=[0, zmaxx], + # cmap="viridis", + # direct=True, + #) def plotParameterization(self, axs=None): if axs is None: @@ -520,15 +521,15 @@ def plotParameterization(self, axs=None): ) # Boundary, axis and limiter ax.plot(self.Rb, self.Yb, lw=1, c="r") - ax.plot(self.g["RMAXIS"], self.g["ZMAXIS"], "+", markersize=10, c="r") + ax.plot(self.g.raw["rmaxis"], self.g.raw["zmaxis"], "+", markersize=10, c="r") ax.plot([self.Rmag], [self.Zmag], "o", markersize=5, c="m") ax.plot([self.Rmajor], [self.Zmag], "+", markersize=10, c="k") - ax.plot(self.g["RLIM"], self.g["ZLIM"], lw=1, c="k") + ax.plot(self.g.raw["rlim"], self.g.raw["zlim"], lw=1, c="k") import matplotlib path = matplotlib.path.Path( - np.transpose(np.array([self.g["RLIM"], self.g["ZLIM"]])) + np.transpose(np.array([self.g.raw["rlim"], self.g.raw["zlim"]])) ) patch = matplotlib.patches.PathPatch(path, facecolor="none") ax.add_patch(patch) @@ -545,12 +546,12 @@ def plotParameterization(self, axs=None): ax.set_ylabel("Z (m)") ax = axs[1] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["geo"]["kap"] + x = self.psi_pol_norm + y = self.g.derived["miller_geo"]["kappa"].copy() ax.plot(x, y, label="$\\kappa$") - y = self.g["fluxSurfaces"]["geo"]["kapl"] + y = self.g.derived["miller_geo"]["kappa_l"].copy() ax.plot(x, y, ls="--", label="$\\kappa_L$") - y = self.g["fluxSurfaces"]["geo"]["kapu"] + y = self.g.derived["miller_geo"]["kappa_u"].copy() ax.plot(x, y, ls="--", label="$\\kappa_U$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -558,12 +559,12 @@ def plotParameterization(self, axs=None): ax.legend() ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["geo"]["delta"] + x = self.psi_pol_norm + y = self.g.derived["miller_geo"]["delta"].copy() ax.plot(x, y, label="$\\delta$") - y = self.g["fluxSurfaces"]["geo"]["dell"] + y = self.g.derived["miller_geo"]["delta_l"].copy() ax.plot(x, y, ls="--", label="$\\delta_L$") - y = self.g["fluxSurfaces"]["geo"]["delu"] + y = self.g.derived["miller_geo"]["delta_u"].copy() ax.plot(x, y, ls="--", label="$\\delta_U$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -571,16 +572,16 @@ def plotParameterization(self, axs=None): ax.legend() ax = axs[3] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["geo"]["zeta"] + x = self.psi_pol_norm + y = self.g.derived["miller_geo"]["zeta"].copy() ax.plot(x, y, label="$\\zeta$") - y = self.g["fluxSurfaces"]["geo"]["zetail"] + y = self.g.derived["miller_geo"]["zeta_li"].copy() ax.plot(x, y, ls="--", label="$\\zeta_{IL}$") - y = self.g["fluxSurfaces"]["geo"]["zetaiu"] + y = self.g.derived["miller_geo"]["zeta_ui"].copy() ax.plot(x, y, ls="--", label="$\\zeta_{IU}$") - y = self.g["fluxSurfaces"]["geo"]["zetaol"] + y = self.g.derived["miller_geo"]["zeta_lo"].copy() ax.plot(x, y, ls="--", label="$\\zeta_{OL}$") - y = self.g["fluxSurfaces"]["geo"]["zetaou"] + y = self.g.derived["miller_geo"]["zeta_uo"].copy() ax.plot(x, y, ls="--", label="$\\zeta_{OU}$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -697,8 +698,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[0] ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["PRES"] * 1e-6, + self.rho_tor, + self.g.raw["pres"] * 1e-6, "-s", c=color, lw=2, @@ -712,8 +713,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[1] ax.plot( - self.g["AuxQuantities"]["RHO"], - -self.g["PPRIME"] * 1e-6, + self.rho_tor, + -self.g.raw["pprime"] * 1e-6, c=color, lw=2, ls="-", @@ -724,13 +725,13 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax.axhline(y=0.0, ls="--", lw=0.5, c="k") ax = ax_plasma[2] - ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FPOL"], c=color, lw=2, ls="-") + ax.plot(self.rho_tor, self.g.raw["fpol"], c=color, lw=2, ls="-") ax.set_xlim([0, 1]) ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") ax.set_ylabel("$F = RB_{\\phi}$ (T*m)") ax = ax_plasma[3] - ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FFPRIM"], c=color, lw=2, ls="-") + ax.plot(self.rho_tor, self.g.raw["ffprim"], c=color, lw=2, ls="-") ax.set_xlim([0, 1]) ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") ax.set_ylabel("FF' (T*m/[])") @@ -738,8 +739,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[4] ax.plot( - self.g["AuxQuantities"]["RHO"], - np.abs(self.g["QPSI"]), + self.rho_tor, + np.abs(self.g.raw["qpsi"]), "-s", c=color, lw=2, @@ -754,8 +755,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[5] ax.plot( - self.g["AuxQuantities"]["RHO"], - np.abs(self.g.surfAvg("Jt") * 1e-6), + self.rho_tor, + np.abs(self.Jt), "-s", c=color, lw=2, @@ -763,8 +764,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): label=label + "geqdsk Jt", ) ax.plot( - self.g["AuxQuantities"]["RHO"], - np.abs(self.g.surfAvg("Jt_fb") * 1e-6), + self.rho_tor, + np.abs(self.Jt_fb), "--o", c=color, lw=2, @@ -779,27 +780,27 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): if legendYN: ax.legend() - ax = ax_plasma[6] - ax.plot( - self.g["fluxSurfaces"]["midplane"]["R"], - np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), - "-s", - c=color, - lw=2, - markersize=3, - label=label + "geqdsk Bt", - ) - ax.plot( - self.g["fluxSurfaces"]["midplane"]["R"], - np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), - "--o", - c=color, - lw=2, - markersize=3, - label=label + "geqdsk Bp", - ) - ax.set_xlabel("R (m) midplane") - ax.set_ylabel("Midplane fields (abs())") + #ax = ax_plasma[6] + #ax.plot( + # self.g["fluxSurfaces"]["midplane"]["R"], + # np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), + # "-s", + # c=color, + # lw=2, + # markersize=3, + # label=label + "geqdsk Bt", + #) + #ax.plot( + # self.g["fluxSurfaces"]["midplane"]["R"], + # np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), + # "--o", + # c=color, + # lw=2, + # markersize=3, + # label=label + "geqdsk Bp", + #) + #ax.set_xlabel("R (m) midplane") + #ax.set_ylabel("Midplane fields (abs())") if legendYN: ax.legend() @@ -812,9 +813,11 @@ def plotGeometry(self, axs=None, color="r"): fig, axs = plt.subplots(ncols=4) ax = axs[0] + x = self.rho_tor + y = self.cx_area ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["fluxSurfaces"]["geo"]["cxArea"], + x, + y, "-", c=color, lw=2, @@ -825,9 +828,11 @@ def plotGeometry(self, axs=None, color="r"): ax.set_ylabel("CX Area ($m^2$)") ax = axs[1] + x = self.rho_tor + y = np.zeros(x.shape) ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["fluxSurfaces"]["geo"]["surfArea"], + x, # self.rho_tor, + y, # self.g["fluxSurfaces"]["geo"]["surfArea"], "-", c=color, lw=2, @@ -838,9 +843,11 @@ def plotGeometry(self, axs=None, color="r"): ax.set_ylabel("Surface Area ($m^2$)") ax = axs[2] + x = self.rho_tor + y = np.zeros(x.shape) ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["fluxSurfaces"]["geo"]["vol"], + x, # self.rho_tor, + y, # self.g["fluxSurfaces"]["geo"]["vol"], "-", c=color, lw=2, @@ -863,13 +870,13 @@ def plotFluxSurfaces( plot1=True, label = '', ): - x = self.g["AuxQuantities"]["R"] - y = self.g["AuxQuantities"]["Z"] + x = self.g.derived["R"] + y = self.g.derived["Z"] if rhoPol: - z = self.g["AuxQuantities"]["RHOpRZ"] + z = self.g.derived["rhorz_pol"] else: - z = self.g["AuxQuantities"]["RHORZ"] + z = self.g.derived["rhorz_tor"] if not sqrt: z = z**2 @@ -905,10 +912,10 @@ def plot2Dquantity( if ax is None: fig, ax = plt.subplots() - x = self.g["AuxQuantities"]["R"] - y = self.g["AuxQuantities"]["Z"] + x = self.g.derived["R"] + y = self.g.derived["Z"] if not direct: - z = self.g["AuxQuantities"][var] * factor + z = self.g.derived[var] * factor else: z = var From 6f79922631f33b876191faf16d7a375b21bcabee Mon Sep 17 00:00:00 2001 From: Xavior Date: Wed, 1 Oct 2025 17:58:32 -0400 Subject: [PATCH 351/385] fixed write_json_CGYRO by converting np objects to native python objects --- .../powertorch/physics_models/transport_cgyro.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index f2b49a9b..c736f592 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -238,4 +238,16 @@ def write_json_CGYRO(roa, fluxes_mean, fluxes_stds, additional_info = None, file 'additional_info': additional_info_extended } - json.dump(json_dict, f, indent=4) + def convert_numpy(obj): + if isinstance(obj, dict): + return {k: convert_numpy(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [convert_numpy(v) for v in obj] + elif isinstance(obj, np.ndarray): + return obj.tolist() + elif isinstance(obj, (np.generic,)): + return obj.item() + else: + return obj + + json.dump(convert_numpy(json_dict), f, indent=4) From d431819c996489d00977f7e61d38de5b5a85ce87 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 1 Oct 2025 18:28:01 -0400 Subject: [PATCH 352/385] Only check for stable criterion if exists --- src/mitim_modules/powertorch/physics_models/transport_cgyro.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py index c736f592..a9065d1c 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_cgyro.py +++ b/src/mitim_modules/powertorch/physics_models/transport_cgyro.py @@ -98,7 +98,8 @@ def _evaluate_gyrokinetic_model(self, code = 'cgyro', gk_object = None): if file_path.exists(): all_good = post_checks(self) - self._stable_correction(simulation_options) + if 'Qi_stable_criterion' in simulation_options: + self._stable_correction(simulation_options) def _stable_correction(self, simulation_options): From 863f0d32eb30eb34e6567ac31a17055e36a8c5ea Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 1 Oct 2025 18:28:11 -0400 Subject: [PATCH 353/385] Added vertical lines for fG in RAPIDS --- src/mitim_tools/popcon_tools/RAPIDStools.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/mitim_tools/popcon_tools/RAPIDStools.py b/src/mitim_tools/popcon_tools/RAPIDStools.py index 6ee0dd8f..fae4690f 100644 --- a/src/mitim_tools/popcon_tools/RAPIDStools.py +++ b/src/mitim_tools/popcon_tools/RAPIDStools.py @@ -229,6 +229,13 @@ def scan_parameter(nn,p_base, xparam, x, nominal_parameters, core, xparamlab='', if vertical_at_nominal: axs[0,0].axvline(x=nominal_parameters[xparam],ls='-.',lw=1.0,c=c) axs[1,0].axvline(x=nominal_parameters[xparam],ls='-.',lw=1.0,c=c) + + fG_nominal = results1['fG'][np.argmin(np.abs(results1['x']- (nominal_parameters[xparam] if not relative else nominal_parameters[xparam])))] + axs[0,1].axvline(x=fG_nominal,ls='-.',lw=1.0,c=c) + axs[1,1].axvline(x=fG_nominal,ls='-.',lw=1.0,c=c) + + + axs[0,1].axvspan(1.0, 1.5, facecolor="k", alpha=0.1, edgecolor="none") axs[1,1].axvspan(1.0, 1.5, facecolor="k", alpha=0.1, edgecolor="none") From ab86cf5faae19c9bde45821715a01f1ccf0bb963 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 1 Oct 2025 18:46:34 -0400 Subject: [PATCH 354/385] Removal of omfit_classes dependency, megpy as substitute geqdsk reader --- pyproject.toml | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7184a5f4..c680bdfb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,28 +42,19 @@ dependencies = [ "botorch", "scikit-image", # Stricly not for MITIM, but good to have for pygacode "psutil", - #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models "tensorflow", "f90nml", - "pyyaml" + "pyyaml", + "megpy @ git+https://github.com/gsnoep/megpy.git@feature/tracer-rebuild", + #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models # "vmecpp", # "quends @ git+https://github.com/sandialabs/quends.git" - "megpy", ] [project.optional-dependencies] pyqt = [ "PyQt6", ] -omfit = [ - "omfit_classes>3.2024.19.2", # Otherwise, it will need an old version of matplotlib, matplotlib<3.6 - "scipy<1.14.0", # As of 08/08/2024, because of https://github.com/gafusion/OMFIT-source/issues/7104 - "numpy<2.0.0", # For the xarray requirement below to work - #"xarray==2022.6.0", # As of 08/08/2024, because https://github.com/gafusion/OMFIT-source/issues/7104. Froze this version because of the PLASMAstate xr reader (importlib_metadata issues) - "omas", - "fortranformat", - "openpyxl", -] test = [ "pytest", "coverage", From 752537d12c41025f18853365c4c7d6746e1458cf Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Thu, 2 Oct 2025 11:40:15 -0400 Subject: [PATCH 355/385] pyproject: changed megpy dependency to unstable --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c680bdfb..b3832fbd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ dependencies = [ "tensorflow", "f90nml", "pyyaml", - "megpy @ git+https://github.com/gsnoep/megpy.git@feature/tracer-rebuild", + "megpy @ git+https://github.com/gsnoep/megpy.git@feature/unstable", #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models # "vmecpp", # "quends @ git+https://github.com/sandialabs/quends.git" From d6b681634980592a375890f7f243850a7458efd6 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 3 Oct 2025 11:30:30 -0400 Subject: [PATCH 356/385] Add fibe but not implemented yet --- pyproject.toml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b3832fbd..2a3fba09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "MITIM" version = "4.0.0" description = "MIT Integrated Modeling Suite for Fusion Applications" readme = "README.md" -requires-python = ">=3.10, <3.13" # Notes: 3.9 has issues with the latest BOTORCH, 3.13 has issues with tensorflow (nn) +requires-python = ">=3.10, <3.13" # Notes: 3.9 has issues with the latest botorch, 3.13 has issues with tensorflow license = {file = "LICENSE"} authors = [ {name = "P. Rodriguez-Fernandez", email = "pablorf@mit.edu"}, @@ -40,12 +40,13 @@ dependencies = [ "shapely", "freegs", "botorch", - "scikit-image", # Stricly not for MITIM, but good to have for pygacode "psutil", "tensorflow", "f90nml", "pyyaml", - "megpy @ git+https://github.com/gsnoep/megpy.git@feature/unstable", + "megpy @ git+https://github.com/gsnoep/megpy.git@unstable", + "fibe @ git+https://github.com/aaronkho/fibe.git", + "scikit-image", # Stricly not for MITIM, but good to have for pygacode #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models # "vmecpp", # "quends @ git+https://github.com/sandialabs/quends.git" From 8ae98c54374a9016bce30b921f386d775eb3c30b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 3 Oct 2025 11:39:58 -0400 Subject: [PATCH 357/385] Do not crash geqdsk crashing if no limiters --- src/mitim_tools/gs_tools/utils/GEQplotting.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/mitim_tools/gs_tools/utils/GEQplotting.py b/src/mitim_tools/gs_tools/utils/GEQplotting.py index 5c4c950b..db545acd 100644 --- a/src/mitim_tools/gs_tools/utils/GEQplotting.py +++ b/src/mitim_tools/gs_tools/utils/GEQplotting.py @@ -524,19 +524,20 @@ def plotParameterization(self, axs=None): ax.plot(self.g.raw["rmaxis"], self.g.raw["zmaxis"], "+", markersize=10, c="r") ax.plot([self.Rmag], [self.Zmag], "o", markersize=5, c="m") ax.plot([self.Rmajor], [self.Zmag], "+", markersize=10, c="k") - ax.plot(self.g.raw["rlim"], self.g.raw["zlim"], lw=1, c="k") + if 'rlim' in self.g.raw and 'zlim' in self.g.raw: + ax.plot(self.g.raw["rlim"], self.g.raw["zlim"], lw=1, c="k") - import matplotlib + import matplotlib - path = matplotlib.path.Path( - np.transpose(np.array([self.g.raw["rlim"], self.g.raw["zlim"]])) - ) - patch = matplotlib.patches.PathPatch(path, facecolor="none") - ax.add_patch(patch) - # for col in cs.collections: - # col.set_clip_path(patch) - # for col in csA.collections: - # col.set_clip_path(patch) + path = matplotlib.path.Path( + np.transpose(np.array([self.g.raw["rlim"], self.g.raw["zlim"]])) + ) + patch = matplotlib.patches.PathPatch(path, facecolor="none") + ax.add_patch(patch) + # for col in cs.collections: + # col.set_clip_path(patch) + # for col in csA.collections: + # col.set_clip_path(patch) self.plotEnclosingBox(ax=ax) From 403780f9570a718872d1071f460cb0fa9084615b Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Fri, 3 Oct 2025 12:00:47 -0400 Subject: [PATCH 358/385] MITIMstate: filter out species with close to zero vtor profiles in mass weighted average velocity profiles in lumpSpecies() --- src/mitim_tools/plasmastate_tools/MITIMstate.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/MITIMstate.py b/src/mitim_tools/plasmastate_tools/MITIMstate.py index 4cfc7f19..28d0c65e 100644 --- a/src/mitim_tools/plasmastate_tools/MITIMstate.py +++ b/src/mitim_tools/plasmastate_tools/MITIMstate.py @@ -1385,7 +1385,8 @@ def lumpSpecies( if "vpol(m/s)" in self.profiles: vpol = np.sum((mass_density * self.profiles["vpol(m/s)"])[:,np.array(ions_list)-1],axis=1) / np.sum(mass_density[:,np.array(ions_list)-1],axis=1) if "vtor(m/s)" in self.profiles: - vtor = np.sum((mass_density * self.profiles["vtor(m/s)"])[:,np.array(ions_list)-1],axis=1) / np.sum(mass_density[:,np.array(ions_list)-1],axis=1) + mask = np.isclose(np.mean(self.profiles["vtor(m/s)"][:,np.array(ions_list)-1],axis=0),1e-12) + vtor = np.sum((mass_density * self.profiles["vtor(m/s)"])[:,np.array(ions_list)-1][:,~mask],axis=1) / np.sum(mass_density[:,np.array(ions_list)-1][:,~mask],axis=1) print(f"\t\t\t* New lumped impurity has Z={Z:.2f}, A={A:.2f} (calculated as 2*Z)") @@ -1787,7 +1788,7 @@ def introduceRotationProfile(self, Mach_LF=1.0, new_file=None): ) # m/s self.profiles["w0(rad/s)"] = Vtor_LF / (self.derived["R_LF"]) # rad/s - + self.derive_quantities() if new_file is not None: From 61dd358b6008be682741cdc6b1d427a7c09904b6 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 3 Oct 2025 16:40:23 -0400 Subject: [PATCH 359/385] Bug fix: NEO now properly searches for electron index in output files --- src/mitim_tools/gacode_tools/NEOtools.py | 35 +++++++++++++++--------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/src/mitim_tools/gacode_tools/NEOtools.py b/src/mitim_tools/gacode_tools/NEOtools.py index ddf9d46f..ce9dfbce 100644 --- a/src/mitim_tools/gacode_tools/NEOtools.py +++ b/src/mitim_tools/gacode_tools/NEOtools.py @@ -296,26 +296,35 @@ def read(self): # Found the header line, now process the data break - line = lines[i+2] - self.Ge, self.Qe, self.Me = [float(x) for x in line.split()[1:]] - - self.GiAll, self.QiAll, self.MiAll = [], [], [] - for i in range(i+3, len(lines)): + Z, G, Q, M = [], [], [], [] + for i in range(i+2, len(lines)): line = lines[i] - self.GiAll.append(float(line.split()[1])) - self.QiAll.append(float(line.split()[2])) - self.MiAll.append(float(line.split()[3])) - - self.GiAll = np.array(self.GiAll) - self.QiAll = np.array(self.QiAll) - self.MiAll = np.array(self.MiAll) + Z.append(float(line.split()[0])) + G.append(float(line.split()[1])) + Q.append(float(line.split()[2])) + M.append(float(line.split()[3])) + Z = np.array(Z) + G = np.array(G) + Q = np.array(Q) + M = np.array(M) + + + # Find electron line (Z= -1) + ie = int(np.where(Z == -1)[0][0]) + + self.Ge = G[ie] + self.Qe = Q[ie] + self.Me = M[ie] + + self.GiAll = np.delete(G, ie) + self.QiAll = np.delete(Q, ie) + self.MiAll = np.delete(M, ie) self.Qi = self.QiAll.sum() self.Mt = self.Me + self.MiAll.sum() self.roa = float(lines[0].split()[-1]) - # ------------------------------------------------------------------------ # Input file # ------------------------------------------------------------------------ From 6fce8c06408121241bde0e47749fc830c32d4813 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 3 Oct 2025 18:46:37 -0400 Subject: [PATCH 360/385] Bring back omfit_classes until megpy is ready to handle freegs equilibria (to be worked in branch development_equilibrium) --- pyproject.toml | 9 + src/mitim_tools/gs_tools/GEQtools.py | 177 +++++---- src/mitim_tools/gs_tools/utils/GEQplotting.py | 360 +++++++++--------- 3 files changed, 295 insertions(+), 251 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2a3fba09..1f2979a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,15 @@ dependencies = [ pyqt = [ "PyQt6", ] +omfit = [ + "omfit_classes>3.2024.19.2", # Otherwise, it will need an old version of matplotlib, matplotlib<3.6 + "scipy<1.14.0", # As of 08/08/2024, because of https://github.com/gafusion/OMFIT-source/issues/7104 + "numpy<2.0.0", # For the xarray requirement below to work + #"xarray==2022.6.0", # As of 08/08/2024, because https://github.com/gafusion/OMFIT-source/issues/7104. Froze this version because of the PLASMAstate xr reader (importlib_metadata issues) + "omas", + "fortranformat", + "openpyxl", +] test = [ "pytest", "coverage", diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index a4401efa..74370d31 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -6,30 +6,70 @@ import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, PLASMAtools from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.gs_tools.utils import GEQplotting +from mitim_tools.gs_tools.utils import GEQplotting_old as GEQplotting from shapely.geometry import LineString -from scipy.integrate import quad, cumulative_trapezoid -import megpy +from scipy.integrate import quad import freegs from freegs import geqdsk from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed """ -Note that this module relies on megpy to intrepret the content of g-eqdsk files. -Modifications are made in MITIM for visualizations and a few extra derivations. +Note that this module relies on OMFIT classes (https://omfit.io/classes.html) procedures to intrepret the content of g-eqdsk files. +Modifications are made in MITINM for visualizations and a few extra derivations. """ +def fix_file(filename): + + with open(filename, "r") as f: + lines = f.readlines() + + # ----------------------------------------------------------------------- + # Remove coils (chatGPT 4o as of 08/24/24) + # ----------------------------------------------------------------------- + # Use StringIO to simulate the file writing + noCoils_file = io.StringIO() + for cont, line in enumerate(lines): + if cont > 0 and line[:2] == " ": + break + noCoils_file.write(line) + + # Reset cursor to the start of StringIO + noCoils_file.seek(0) + + # Write the StringIO content to a temporary file + with tempfile.NamedTemporaryFile(delete=False) as tmp_file: + tmp_file.write(noCoils_file.getvalue().encode('utf-8')) + noCoils_file = tmp_file.name + # ----------------------------------------------------------------------- + + with open(filename, 'r') as file1, open(noCoils_file, 'r') as file2: + file1_content = file1.read() + file2_content = file2.read() + + if file1_content != file2_content: + print(f"\t- geqdsk file {IOtools.clipstr(filename)} had coils, I have removed them") + + filename = noCoils_file + + return filename + class MITIMgeqdsk: def __init__(self, filename): - self.g = megpy.Equilibrium() - self.g.read_geqdsk(f_path=filename) - self.g.add_derived(incl_fluxsurfaces=True, analytic_shape=True, incl_B=True) + # Fix file by removing coils if it has them + filename = fix_file(filename) + + # Read GEQDSK file using OMFIT + import omfit_classes.omfit_eqdsk + self.g = omfit_classes.omfit_eqdsk.OMFITgeqdsk(filename, forceFindSeparatrix=True) # Extra derivations in MITIM self.derive() + # Remove temporary file + os.remove(filename) + @classmethod def timeslices(cls, filename, **kwargs): print("\n...Opening GEQ file with several time slices") @@ -71,65 +111,61 @@ def timeslices(cls, filename, **kwargs): def derive(self, debug=False): - zero_vector = np.zeros(self.g.derived["rho_pol"].shape) - - self.rho_pol = self.g.derived["rho_pol"].copy() - self.rho_tor = self.g.derived["rho_tor"].copy() - self.psi_pol_norm = self.rho_pol ** 2 - self.psi_tor_norm = self.rho_tor ** 2 - - self.Jt = self.g.derived["j_tor"] * 1e-6 - self.Jt_fb = zero_vector.copy() + self.Jt = self.g.surfAvg("Jt") * 1e-6 + self.Jt_fb = self.g.surfAvg("Jt_fb") * 1e-6 self.Jerror = np.abs(self.Jt - self.Jt_fb) - self.Ip = self.g.raw["current"] + self.Ip = self.g["CURRENT"] # Parameterizations of LCFS - self.kappa = self.g.derived["miller_geo"]["kappa"][-1] - self.kappaU = self.g.derived["miller_geo"]["kappa_u"][-1] - self.kappaL = self.g.derived["miller_geo"]["kappa_l"][-1] - - self.delta = self.g.derived["miller_geo"]["delta"][-1] - self.deltaU = self.g.derived["miller_geo"]["delta_u"][-1] - self.deltaL = self.g.derived["miller_geo"]["delta_l"][-1] - - self.zeta = self.g.derived["miller_geo"]["zeta"][-1] - - self.a = self.g.derived["r"][-1] - self.Rmag = self.g.derived["Ro"][0] - self.Zmag = self.g.derived["Zo"][0] + self.kappa = self.g["fluxSurfaces"]["geo"]["kap"][-1] + self.kappaU = self.g["fluxSurfaces"]["geo"]["kapu"][-1] + self.kappaL = self.g["fluxSurfaces"]["geo"]["kapl"][-1] + + self.delta = self.g["fluxSurfaces"]["geo"]["delta"][-1] + self.deltaU = self.g["fluxSurfaces"]["geo"]["dell"][-1] + self.deltaL = self.g["fluxSurfaces"]["geo"]["dell"][-1] + + self.zeta = self.g["fluxSurfaces"]["geo"]["zeta"][-1] + + self.a = self.g["fluxSurfaces"]["geo"]["a"][-1] + self.Rmag = self.g["fluxSurfaces"]["geo"]["R"][0] + self.Zmag = self.g["fluxSurfaces"]["geo"]["Z"][0] + self.Rmajor = np.mean( + [ + self.g["fluxSurfaces"]["geo"]["Rmin_centroid"][-1], + self.g["fluxSurfaces"]["geo"]["Rmax_centroid"][-1], + ] + ) - self.Rmajor = self.g.derived["Ro"][-1] - self.Zmajor = self.Zmag #self.g.derived["Zo"][-1] TODO: check which Z0 to use, perhaps switch to MXH definition + self.Zmajor = self.Zmag self.eps = self.a / self.Rmajor # Core values - vp = np.array(self.g.fluxsurfaces["Vprime"]).flatten() - ir = np.array(self.g.fluxsurfaces["1/R"]).flatten() - self.cx_area = cumulative_trapezoid(vp * ir, self.g.derived["psi"], initial=0.0) - self.kappa_a = self.cx_area[-1] / (np.pi * self.a**2) + + self.kappa_a = self.g["fluxSurfaces"]["geo"]["cxArea"][-1] / (np.pi * self.a**2) self.kappa995 = np.interp( 0.995, - self.psi_pol_norm, - self.g.derived["miller_geo"]["kappa"], + self.g["AuxQuantities"]["PSI_NORM"], + self.g["fluxSurfaces"]["geo"]["kap"], ) self.kappa95 = np.interp( 0.95, - self.psi_pol_norm, - self.g.derived["miller_geo"]["kappa"], + self.g["AuxQuantities"]["PSI_NORM"], + self.g["fluxSurfaces"]["geo"]["kap"], ) self.delta995 = np.interp( 0.995, - self.psi_pol_norm, - self.g.derived["miller_geo"]["delta"], + self.g["AuxQuantities"]["PSI_NORM"], + self.g["fluxSurfaces"]["geo"]["delta"], ) self.delta95 = np.interp( 0.95, - self.psi_pol_norm, - self.g.derived["miller_geo"]["delta"], + self.g["AuxQuantities"]["PSI_NORM"], + self.g["fluxSurfaces"]["geo"]["delta"], ) """ @@ -137,23 +173,26 @@ def derive(self, debug=False): Boundary -------------------------------------------------------------------------------------------------------------------------------------- Note that the RBBS and ZBBS values in the gfile are often too scattered and do not reproduce the boundary near x-points. - The shaping parameters calculated using fluxsurfaces are correct though. + The shaping parameters calculated using fluxSurfaces are correct though. """ - self.Rb_gfile, self.Yb_gfile = self.g.raw["rbbbs"].copy(), self.g.raw["zbbbs"].copy() - self.Rb, self.Yb = self.g.fluxsurfaces["R"][-1], self.g.fluxsurfaces["Z"][-1] + self.Rb_gfile, self.Yb_gfile = self.g["RBBBS"], self.g["ZBBBS"] + self.Rb, self.Yb = self.g["fluxSurfaces"].sep.transpose() if len(self.Rb) == 0: - print("\t- MITIM > No separatrix found in the megpy fluxsurfaces, using explicit boundary in g-eqdsk file!",typeMsg='i') + print("\t- MITIM > No separatrix found in the OMFIT fluxSurfaces, increasing resolution and going all in!",typeMsg='i') - self.Rb = self.Rb_gfile.copy() - self.Yb = self.Yb_gfile.copy() + flx = copy.deepcopy(self.g['fluxSurfaces']) + flx._changeResolution(6) + flx.findSurfaces([0.0,0.5,1.0]) + fs = flx['flux'][list(flx['flux'].keys())[-1]] + self.Rb, self.Yb = fs['R'], fs['Z'] if debug: fig, ax = plt.subplots() - # megpy - ax.plot(self.Rb, self.Yb, "-s", c="r", label="megpy") + # OMFIT + ax.plot(self.Rb, self.Yb, "-s", c="r", label="OMFIT") # GFILE ax.plot(self.Rb_gfile, self.Yb_gfile, "-s", c="y", label="GFILE") @@ -206,7 +245,10 @@ def write(self, filename=None): If filename is None, use the original one """ - self.g.write_geqdsk(f_path=filename) + if filename is not None: + self.g.filename = filename + + self.g.save() # ----------------------------------------------------------------------------- # Parameterizations @@ -214,8 +256,8 @@ def write(self, filename=None): def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): - psis = self.psi_pol_norm - flux_surfaces = self.g.fluxsurfaces["psi"] + psis = self.g["AuxQuantities"]["PSI_NORM"] + flux_surfaces = self.g['fluxSurfaces']['flux'] # Cannot parallelize because different number of points? kappa, rmin, rmaj, zmag, sn, cn = [],[],[],[],[],[] @@ -224,7 +266,7 @@ def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): if flux == len(flux_surfaces)-1: Rf, Zf = self.Rb, self.Yb else: - Rf, Zf = self.g.fluxsurfaces["R"][flux], self.g.fluxsurfaces["Z"][flux] + Rf, Zf = flux_surfaces[flux]['R'],flux_surfaces[flux]['Z'] # Perform the MXH decompositionusing the MITIM surface class surfaces = mitim_flux_surfaces() @@ -269,22 +311,24 @@ def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): # ----------------------------------------------------------------------------- # For MAESTRO and TRANSP converstions # ----------------------------------------------------------------------------- + def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH = 7, plotYN = False): # ------------------------------------------------------------------------------------------------------- # Quantities from the equilibrium # ------------------------------------------------------------------------------------------------------- - rhotor = self.g.derived['rho_tor'] - psi = self.g.derived['psi'] # Wb/rad - torfluxa = self.g.derived['phi'][-1] / (2*np.pi) # Wb/rad - q = self.g.raw['qpsi'] - pressure = self.g.raw['pres'] # Pa - Ip = self.g.raw['current']*1E-6 # MA + rhotor = self.g['RHOVN'] + psi = self.g['AuxQuantities']['PSI'] # Wb/rad + torfluxa = self.g['AuxQuantities']['PHI'][-1] / (2*np.pi) # Wb/rad + q = self.g['QPSI'] + pressure = self.g['PRES'] # Pa + Ip = self.g['CURRENT']*1E-6 # MA RZ = np.array([self.Rb,self.Yb]).T R0 = (RZ.max(axis=0)[0] + RZ.min(axis=0)[0])/2 - B0 = self.g.raw['rcentr']*self.g.raw['bcentr'] / R0 + + B0 = self.g['RCENTR']*self.g['BCENTR'] / R0 # Ensure positive quantities #TODO: Check if this is necessary, pass directions rhotor = np.array([np.abs(i) for i in rhotor]) @@ -499,7 +543,6 @@ def _to_mxh(self, n_coeff=6): self.cn = np.zeros((self.R.shape[0],n_coeff)) self.sn = np.zeros((self.R.shape[0],n_coeff)) self.gn = np.zeros((self.R.shape[0],4)) - self.gn[-1] = 1.0 for i in range(self.R.shape[0]): self.cn[i,:], self.sn[i,:], self.gn[i,:] = from_RZ_to_mxh(self.R[i,:], self.Z[i,:], n_coeff=n_coeff) @@ -519,8 +562,8 @@ def _to_miller(self): # Elongations - self.kappa_u = (Zmax - self.Z0) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) - self.kappa_l = (self.Z0 - Zmin) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) + self.kappa_u = (Zmax - self.Z0) / self.a + self.kappa_l = (self.Z0 - Zmin) / self.a self.kappa = (self.kappa_u + self.kappa_l) / 2 # Triangularities diff --git a/src/mitim_tools/gs_tools/utils/GEQplotting.py b/src/mitim_tools/gs_tools/utils/GEQplotting.py index db545acd..fe6531b4 100644 --- a/src/mitim_tools/gs_tools/utils/GEQplotting.py +++ b/src/mitim_tools/gs_tools/utils/GEQplotting.py @@ -4,7 +4,6 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -#TODO: add current profiles and flux-surface average fields to megpy and restore plots def compareGeqdsk(geqdsks, fn=None, extraLabel="", plotAll=True, labelsGs=None): @@ -230,8 +229,8 @@ def plotFS(self, axs=None, color="b", label=""): ax.set_ylabel("Z (m)") ax = axs[2] - x = self.psi_pol_norm - y = self.rho_tor + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["AuxQuantities"]["RHO"] ax.plot(x, y, lw=2, ls="-", c=color, label=label) ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) @@ -241,8 +240,8 @@ def plotFS(self, axs=None, color="b", label=""): ax.set_ylim([0, 1]) ax = axs[3] - x = self.rho_tor - y = self.rho_pol + x = self.g["AuxQuantities"]["RHO"] + y = self.g["AuxQuantities"]["RHOp"] ax.plot(x, y, lw=2, ls="-", c=color) ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) @@ -257,8 +256,8 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): fig, axs = plt.subplots(ncols=10) ax = axs[0] - x = self.psi_pol_norm - y = np.zeros(x.shape) + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jr") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_ylabel("FSA $\\langle J\\rangle$ ($MA/m^2$)") @@ -268,14 +267,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[1] - #plot2Dquantity(self, - # ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 - #) + ax = axs[1] + plot2Dquantity(self, + ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 + ) ax = axs[2] - x = self.psi_pol_norm - y = np.zeros(x.shape) + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jz") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -283,14 +282,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[3] - #plot2Dquantity(self, - # ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 - #) + ax = axs[3] + plot2Dquantity(self, + ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 + ) ax = axs[4] - x = self.psi_pol_norm - y = self.Jt + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jt") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -298,14 +297,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[5] - #plot2Dquantity(self, - # ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 - #) + ax = axs[5] + plot2Dquantity(self, + ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 + ) ax = axs[6] - x = self.psi_pol_norm - y = np.zeros(x.shape) + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jp") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -313,14 +312,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[7] - #plot2Dquantity(self, - # ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 - #) + ax = axs[7] + plot2Dquantity(self, + ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 + ) ax = axs[8] - x = self.psi_pol_norm - y = np.zeros(x.shape) + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jpar") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -328,10 +327,10 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[9] - #plot2Dquantity(self, - # ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 - #) + ax = axs[9] + plot2Dquantity(self, + ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 + ) def plotFields(self, axs=None, zlims_thr=[-1, 1]): if axs is None: @@ -339,8 +338,8 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): fig, axs = plt.subplots(ncols=10) ax = axs[0] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g.surfAvg("Br") + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Br") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_ylabel("FSA $\\langle B\\rangle$ ($T$)") @@ -352,12 +351,12 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[1] plot2Dquantity(self, - ax=ax, var="B_r", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="Br", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" ) ax = axs[2] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g.surfAvg("Bz") + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Bz") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -367,12 +366,12 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[3] plot2Dquantity(self, - ax=ax, var="B_z", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="Bz", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" ) ax = axs[4] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g.surfAvg("Bt") + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Bt") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -381,14 +380,14 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): # zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[5] - #plot2Dquantity(self, - # ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" - #) + ax = axs[5] + plot2Dquantity(self, + ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" + ) ax = axs[6] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g.surfAvg("Bp") + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Bp") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -398,32 +397,32 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[7] plot2Dquantity(self, - ax=ax, var="B_pol_rz", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="Bp", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" ) ax = axs[8] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g["fluxSurfaces"]["avg"]["Bp**2"] + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["fluxSurfaces"]["avg"]["Bp**2"] ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) ax.set_ylabel("$\\langle B_{\\theta}^2\\rangle$") - #ax = axs[9] - #x = self.g["fluxSurfaces"]["midplane"]["R"] - #y = self.g["fluxSurfaces"]["midplane"]["Bt"] - #ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") - #y = self.g["fluxSurfaces"]["midplane"]["Bp"] - #ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") - #y = self.g["fluxSurfaces"]["midplane"]["Bz"] - #ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") - #y = self.g["fluxSurfaces"]["midplane"]["Br"] - #ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") - #y = self.g["fluxSurfaces"]["geo"]["bunit"] - #ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") - #ax.set_xlabel("$R$ LF midplane") - #ax.set_ylabel("$B$ (T)") - #ax.legend() + ax = axs[9] + x = self.g["fluxSurfaces"]["midplane"]["R"] + y = self.g["fluxSurfaces"]["midplane"]["Bt"] + ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") + y = self.g["fluxSurfaces"]["midplane"]["Bp"] + ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") + y = self.g["fluxSurfaces"]["midplane"]["Bz"] + ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") + y = self.g["fluxSurfaces"]["midplane"]["Br"] + ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") + y = self.g["fluxSurfaces"]["geo"]["bunit"] + ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") + ax.set_xlabel("$R$ LF midplane") + ax.set_ylabel("$B$ (T)") + ax.legend() def plotChecks(self, axs=None): if axs is None: @@ -431,7 +430,7 @@ def plotChecks(self, axs=None): fig, axs = plt.subplots(ncols=8) ax = axs[0] - x = self.psi_pol_norm + x = self.g["AuxQuantities"]["PSI_NORM"] y1 = self.Jt ax.plot(x, np.abs(y1), lw=2, ls="-", c="b", label="$\\langle Jt\\rangle$") zmax = y1.max() @@ -464,10 +463,10 @@ def plotChecks(self, axs=None): ax.legend() ax = axs[2] - x = self.psi_pol_norm - y1 = self.g.raw["ffprim"] + x = self.g["AuxQuantities"]["PSI_NORM"] + y1 = self.g["FFPRIM"] ax.plot(x, y1, lw=2, ls="-", c="r", label="$FF'$") - y2 = self.g.raw["pprime"] * (4 * np.pi * 1e-7) + y2 = self.g["PPRIME"] * (4 * np.pi * 1e-7) ax.plot(x, y2, lw=2, ls="-", c="b", label="$p'*\\mu_0$") ax.set_ylabel("") @@ -475,40 +474,40 @@ def plotChecks(self, axs=None): ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) - #ax = axs[3] - #plot2Dquantity(self, - # ax=ax, - # var="Jt", - # title="Toroidal Current Jt", - # zlims=[zmin, zmax], - # cmap="viridis", - # factor=1e-6, - #) - - #ax = axs[4] - #plot2Dquantity(self, - # ax=ax, - # var="Jt_fb", - # title="Toroidal Current Jt (FB)", - # zlims=[zmin, zmax], - # cmap="viridis", - # factor=1e-6, - #) - - #ax = axs[5] - #z = ( - # np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) - # * 1e-6 - #) - #zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) - #plot2Dquantity(self, - # ax=ax, - # var=z, - # title="Absolute Error", - # zlims=[0, zmaxx], - # cmap="viridis", - # direct=True, - #) + ax = axs[3] + plot2Dquantity(self, + ax=ax, + var="Jt", + title="Toroidal Current Jt", + zlims=[zmin, zmax], + cmap="viridis", + factor=1e-6, + ) + + ax = axs[4] + plot2Dquantity(self, + ax=ax, + var="Jt_fb", + title="Toroidal Current Jt (FB)", + zlims=[zmin, zmax], + cmap="viridis", + factor=1e-6, + ) + + ax = axs[5] + z = ( + np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) + * 1e-6 + ) + zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) + plot2Dquantity(self, + ax=ax, + var=z, + title="Absolute Error", + zlims=[0, zmaxx], + cmap="viridis", + direct=True, + ) def plotParameterization(self, axs=None): if axs is None: @@ -521,23 +520,22 @@ def plotParameterization(self, axs=None): ) # Boundary, axis and limiter ax.plot(self.Rb, self.Yb, lw=1, c="r") - ax.plot(self.g.raw["rmaxis"], self.g.raw["zmaxis"], "+", markersize=10, c="r") + ax.plot(self.g["RMAXIS"], self.g["ZMAXIS"], "+", markersize=10, c="r") ax.plot([self.Rmag], [self.Zmag], "o", markersize=5, c="m") ax.plot([self.Rmajor], [self.Zmag], "+", markersize=10, c="k") - if 'rlim' in self.g.raw and 'zlim' in self.g.raw: - ax.plot(self.g.raw["rlim"], self.g.raw["zlim"], lw=1, c="k") + ax.plot(self.g["RLIM"], self.g["ZLIM"], lw=1, c="k") - import matplotlib + import matplotlib - path = matplotlib.path.Path( - np.transpose(np.array([self.g.raw["rlim"], self.g.raw["zlim"]])) - ) - patch = matplotlib.patches.PathPatch(path, facecolor="none") - ax.add_patch(patch) - # for col in cs.collections: - # col.set_clip_path(patch) - # for col in csA.collections: - # col.set_clip_path(patch) + path = matplotlib.path.Path( + np.transpose(np.array([self.g["RLIM"], self.g["ZLIM"]])) + ) + patch = matplotlib.patches.PathPatch(path, facecolor="none") + ax.add_patch(patch) + # for col in cs.collections: + # col.set_clip_path(patch) + # for col in csA.collections: + # col.set_clip_path(patch) self.plotEnclosingBox(ax=ax) @@ -547,12 +545,12 @@ def plotParameterization(self, axs=None): ax.set_ylabel("Z (m)") ax = axs[1] - x = self.psi_pol_norm - y = self.g.derived["miller_geo"]["kappa"].copy() + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["fluxSurfaces"]["geo"]["kap"] ax.plot(x, y, label="$\\kappa$") - y = self.g.derived["miller_geo"]["kappa_l"].copy() + y = self.g["fluxSurfaces"]["geo"]["kapl"] ax.plot(x, y, ls="--", label="$\\kappa_L$") - y = self.g.derived["miller_geo"]["kappa_u"].copy() + y = self.g["fluxSurfaces"]["geo"]["kapu"] ax.plot(x, y, ls="--", label="$\\kappa_U$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -560,12 +558,12 @@ def plotParameterization(self, axs=None): ax.legend() ax = axs[2] - x = self.psi_pol_norm - y = self.g.derived["miller_geo"]["delta"].copy() + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["fluxSurfaces"]["geo"]["delta"] ax.plot(x, y, label="$\\delta$") - y = self.g.derived["miller_geo"]["delta_l"].copy() + y = self.g["fluxSurfaces"]["geo"]["dell"] ax.plot(x, y, ls="--", label="$\\delta_L$") - y = self.g.derived["miller_geo"]["delta_u"].copy() + y = self.g["fluxSurfaces"]["geo"]["delu"] ax.plot(x, y, ls="--", label="$\\delta_U$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -573,16 +571,16 @@ def plotParameterization(self, axs=None): ax.legend() ax = axs[3] - x = self.psi_pol_norm - y = self.g.derived["miller_geo"]["zeta"].copy() + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["fluxSurfaces"]["geo"]["zeta"] ax.plot(x, y, label="$\\zeta$") - y = self.g.derived["miller_geo"]["zeta_li"].copy() + y = self.g["fluxSurfaces"]["geo"]["zetail"] ax.plot(x, y, ls="--", label="$\\zeta_{IL}$") - y = self.g.derived["miller_geo"]["zeta_ui"].copy() + y = self.g["fluxSurfaces"]["geo"]["zetaiu"] ax.plot(x, y, ls="--", label="$\\zeta_{IU}$") - y = self.g.derived["miller_geo"]["zeta_lo"].copy() + y = self.g["fluxSurfaces"]["geo"]["zetaol"] ax.plot(x, y, ls="--", label="$\\zeta_{OL}$") - y = self.g.derived["miller_geo"]["zeta_uo"].copy() + y = self.g["fluxSurfaces"]["geo"]["zetaou"] ax.plot(x, y, ls="--", label="$\\zeta_{OU}$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -699,8 +697,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[0] ax.plot( - self.rho_tor, - self.g.raw["pres"] * 1e-6, + self.g["AuxQuantities"]["RHO"], + self.g["PRES"] * 1e-6, "-s", c=color, lw=2, @@ -714,8 +712,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[1] ax.plot( - self.rho_tor, - -self.g.raw["pprime"] * 1e-6, + self.g["AuxQuantities"]["RHO"], + -self.g["PPRIME"] * 1e-6, c=color, lw=2, ls="-", @@ -726,13 +724,13 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax.axhline(y=0.0, ls="--", lw=0.5, c="k") ax = ax_plasma[2] - ax.plot(self.rho_tor, self.g.raw["fpol"], c=color, lw=2, ls="-") + ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FPOL"], c=color, lw=2, ls="-") ax.set_xlim([0, 1]) ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") ax.set_ylabel("$F = RB_{\\phi}$ (T*m)") ax = ax_plasma[3] - ax.plot(self.rho_tor, self.g.raw["ffprim"], c=color, lw=2, ls="-") + ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FFPRIM"], c=color, lw=2, ls="-") ax.set_xlim([0, 1]) ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") ax.set_ylabel("FF' (T*m/[])") @@ -740,8 +738,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[4] ax.plot( - self.rho_tor, - np.abs(self.g.raw["qpsi"]), + self.g["AuxQuantities"]["RHO"], + np.abs(self.g["QPSI"]), "-s", c=color, lw=2, @@ -756,8 +754,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[5] ax.plot( - self.rho_tor, - np.abs(self.Jt), + self.g["AuxQuantities"]["RHO"], + np.abs(self.g.surfAvg("Jt") * 1e-6), "-s", c=color, lw=2, @@ -765,8 +763,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): label=label + "geqdsk Jt", ) ax.plot( - self.rho_tor, - np.abs(self.Jt_fb), + self.g["AuxQuantities"]["RHO"], + np.abs(self.g.surfAvg("Jt_fb") * 1e-6), "--o", c=color, lw=2, @@ -781,27 +779,27 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): if legendYN: ax.legend() - #ax = ax_plasma[6] - #ax.plot( - # self.g["fluxSurfaces"]["midplane"]["R"], - # np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), - # "-s", - # c=color, - # lw=2, - # markersize=3, - # label=label + "geqdsk Bt", - #) - #ax.plot( - # self.g["fluxSurfaces"]["midplane"]["R"], - # np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), - # "--o", - # c=color, - # lw=2, - # markersize=3, - # label=label + "geqdsk Bp", - #) - #ax.set_xlabel("R (m) midplane") - #ax.set_ylabel("Midplane fields (abs())") + ax = ax_plasma[6] + ax.plot( + self.g["fluxSurfaces"]["midplane"]["R"], + np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), + "-s", + c=color, + lw=2, + markersize=3, + label=label + "geqdsk Bt", + ) + ax.plot( + self.g["fluxSurfaces"]["midplane"]["R"], + np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), + "--o", + c=color, + lw=2, + markersize=3, + label=label + "geqdsk Bp", + ) + ax.set_xlabel("R (m) midplane") + ax.set_ylabel("Midplane fields (abs())") if legendYN: ax.legend() @@ -814,11 +812,9 @@ def plotGeometry(self, axs=None, color="r"): fig, axs = plt.subplots(ncols=4) ax = axs[0] - x = self.rho_tor - y = self.cx_area ax.plot( - x, - y, + self.g["AuxQuantities"]["RHO"], + self.g["fluxSurfaces"]["geo"]["cxArea"], "-", c=color, lw=2, @@ -829,11 +825,9 @@ def plotGeometry(self, axs=None, color="r"): ax.set_ylabel("CX Area ($m^2$)") ax = axs[1] - x = self.rho_tor - y = np.zeros(x.shape) ax.plot( - x, # self.rho_tor, - y, # self.g["fluxSurfaces"]["geo"]["surfArea"], + self.g["AuxQuantities"]["RHO"], + self.g["fluxSurfaces"]["geo"]["surfArea"], "-", c=color, lw=2, @@ -844,11 +838,9 @@ def plotGeometry(self, axs=None, color="r"): ax.set_ylabel("Surface Area ($m^2$)") ax = axs[2] - x = self.rho_tor - y = np.zeros(x.shape) ax.plot( - x, # self.rho_tor, - y, # self.g["fluxSurfaces"]["geo"]["vol"], + self.g["AuxQuantities"]["RHO"], + self.g["fluxSurfaces"]["geo"]["vol"], "-", c=color, lw=2, @@ -871,13 +863,13 @@ def plotFluxSurfaces( plot1=True, label = '', ): - x = self.g.derived["R"] - y = self.g.derived["Z"] + x = self.g["AuxQuantities"]["R"] + y = self.g["AuxQuantities"]["Z"] if rhoPol: - z = self.g.derived["rhorz_pol"] + z = self.g["AuxQuantities"]["RHOpRZ"] else: - z = self.g.derived["rhorz_tor"] + z = self.g["AuxQuantities"]["RHORZ"] if not sqrt: z = z**2 @@ -913,10 +905,10 @@ def plot2Dquantity( if ax is None: fig, ax = plt.subplots() - x = self.g.derived["R"] - y = self.g.derived["Z"] + x = self.g["AuxQuantities"]["R"] + y = self.g["AuxQuantities"]["Z"] if not direct: - z = self.g.derived[var] * factor + z = self.g["AuxQuantities"][var] * factor else: z = var From 5ad7e6785eb638358b101d52cfb321c6471d6655 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 3 Oct 2025 18:49:39 -0400 Subject: [PATCH 361/385] New megpy implementations as separate files for now --- src/mitim_tools/gs_tools/GEQtools_megpy.py | 1154 +++++++++++++++++ .../gs_tools/utils/GEQplotting_megpy.py | 984 ++++++++++++++ 2 files changed, 2138 insertions(+) create mode 100644 src/mitim_tools/gs_tools/GEQtools_megpy.py create mode 100644 src/mitim_tools/gs_tools/utils/GEQplotting_megpy.py diff --git a/src/mitim_tools/gs_tools/GEQtools_megpy.py b/src/mitim_tools/gs_tools/GEQtools_megpy.py new file mode 100644 index 00000000..a4401efa --- /dev/null +++ b/src/mitim_tools/gs_tools/GEQtools_megpy.py @@ -0,0 +1,1154 @@ +import os +import io +import tempfile +import copy +import numpy as np +import matplotlib.pyplot as plt +from mitim_tools.misc_tools import GRAPHICStools, IOtools, PLASMAtools +from mitim_tools.gacode_tools import PROFILEStools +from mitim_tools.gs_tools.utils import GEQplotting +from shapely.geometry import LineString +from scipy.integrate import quad, cumulative_trapezoid +import megpy +import freegs +from freegs import geqdsk +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +""" +Note that this module relies on megpy to intrepret the content of g-eqdsk files. +Modifications are made in MITIM for visualizations and a few extra derivations. +""" + +class MITIMgeqdsk: + def __init__(self, filename): + + self.g = megpy.Equilibrium() + self.g.read_geqdsk(f_path=filename) + self.g.add_derived(incl_fluxsurfaces=True, analytic_shape=True, incl_B=True) + + # Extra derivations in MITIM + self.derive() + + @classmethod + def timeslices(cls, filename, **kwargs): + print("\n...Opening GEQ file with several time slices") + + with open(filename, "rb") as f: + lines_full = f.readlines() + + resolutions = [int(a) for a in lines_full[0].split()[-3:]] + + lines_files = [] + lines_files0 = [] + for i in range(len(lines_full)): + line = lines_full[i] + resols = [] + try: + resols = [int(a) for a in line.split()[-3:]] + except: + pass + if (resols == resolutions) and (i != 0): + lines_files.append(lines_files0) + lines_files0 = [] + lines_files0.append(line) + lines_files.append(lines_files0) + + # Write files + gs = [] + for i in range(len(lines_files)): + with open(f"{filename}_time{i}.geqdsk", "wb") as f: + f.writelines(lines_files[i]) + + gs.append( + cls( + f"{filename}_time{i}.geqdsk",**kwargs, + ) + ) + os.remove(f"{filename}_time{i}.geqdsk") + + return gs + + def derive(self, debug=False): + + zero_vector = np.zeros(self.g.derived["rho_pol"].shape) + + self.rho_pol = self.g.derived["rho_pol"].copy() + self.rho_tor = self.g.derived["rho_tor"].copy() + self.psi_pol_norm = self.rho_pol ** 2 + self.psi_tor_norm = self.rho_tor ** 2 + + self.Jt = self.g.derived["j_tor"] * 1e-6 + self.Jt_fb = zero_vector.copy() + + self.Jerror = np.abs(self.Jt - self.Jt_fb) + + self.Ip = self.g.raw["current"] + + # Parameterizations of LCFS + self.kappa = self.g.derived["miller_geo"]["kappa"][-1] + self.kappaU = self.g.derived["miller_geo"]["kappa_u"][-1] + self.kappaL = self.g.derived["miller_geo"]["kappa_l"][-1] + + self.delta = self.g.derived["miller_geo"]["delta"][-1] + self.deltaU = self.g.derived["miller_geo"]["delta_u"][-1] + self.deltaL = self.g.derived["miller_geo"]["delta_l"][-1] + + self.zeta = self.g.derived["miller_geo"]["zeta"][-1] + + self.a = self.g.derived["r"][-1] + self.Rmag = self.g.derived["Ro"][0] + self.Zmag = self.g.derived["Zo"][0] + + self.Rmajor = self.g.derived["Ro"][-1] + self.Zmajor = self.Zmag #self.g.derived["Zo"][-1] TODO: check which Z0 to use, perhaps switch to MXH definition + + self.eps = self.a / self.Rmajor + + # Core values + vp = np.array(self.g.fluxsurfaces["Vprime"]).flatten() + ir = np.array(self.g.fluxsurfaces["1/R"]).flatten() + self.cx_area = cumulative_trapezoid(vp * ir, self.g.derived["psi"], initial=0.0) + self.kappa_a = self.cx_area[-1] / (np.pi * self.a**2) + + self.kappa995 = np.interp( + 0.995, + self.psi_pol_norm, + self.g.derived["miller_geo"]["kappa"], + ) + self.kappa95 = np.interp( + 0.95, + self.psi_pol_norm, + self.g.derived["miller_geo"]["kappa"], + ) + self.delta995 = np.interp( + 0.995, + self.psi_pol_norm, + self.g.derived["miller_geo"]["delta"], + ) + self.delta95 = np.interp( + 0.95, + self.psi_pol_norm, + self.g.derived["miller_geo"]["delta"], + ) + + """ + -------------------------------------------------------------------------------------------------------------------------------------- + Boundary + -------------------------------------------------------------------------------------------------------------------------------------- + Note that the RBBS and ZBBS values in the gfile are often too scattered and do not reproduce the boundary near x-points. + The shaping parameters calculated using fluxsurfaces are correct though. + """ + + self.Rb_gfile, self.Yb_gfile = self.g.raw["rbbbs"].copy(), self.g.raw["zbbbs"].copy() + self.Rb, self.Yb = self.g.fluxsurfaces["R"][-1], self.g.fluxsurfaces["Z"][-1] + + if len(self.Rb) == 0: + print("\t- MITIM > No separatrix found in the megpy fluxsurfaces, using explicit boundary in g-eqdsk file!",typeMsg='i') + + self.Rb = self.Rb_gfile.copy() + self.Yb = self.Yb_gfile.copy() + + if debug: + fig, ax = plt.subplots() + + # megpy + ax.plot(self.Rb, self.Yb, "-s", c="r", label="megpy") + + # GFILE + ax.plot(self.Rb_gfile, self.Yb_gfile, "-s", c="y", label="GFILE") + + # Extras + self.plotFluxSurfaces( + ax=ax, fluxes=[0.99999, 1.0], rhoPol=True, sqrt=False, color="m" + ) + self.plotXpointEnvelope( + ax=ax, color="c", alpha=1.0, rhoPol=True, sqrt=False + ) + self.plotEnclosingBox(ax=ax) + + ax.legend() + ax.set_aspect("equal") + ax.set_xlabel("R [m]") + ax.set_ylabel("Z [m]") + + plt.show() + + def plotEnclosingBox(self, ax=None, c= "k"): + if ax is None: + fig, ax = plt.subplots() + + Rmajor, a, Zmajor, kappaU, kappaL, deltaU, deltaL = ( + self.Rmajor, + self.a, + self.Zmajor, + self.kappaU, + self.kappaL, + self.deltaU, + self.deltaL, + ) + + ax.axhline(y=Zmajor, ls="--", c=c, lw=0.5) + ax.axvline(x=Rmajor, ls="--", c=c, lw=0.5) + ax.axvline(x=Rmajor - a, ls="--", c=c, lw=0.5) + ax.axvline(x=Rmajor + a, ls="--", c=c, lw=0.5) + Rtop = Zmajor + a * kappaU + ax.axhline(y=Rtop, ls="--", c=c, lw=0.5) + Rbot = Zmajor - a * kappaL + ax.axhline(y=Rbot, ls="--", c=c, lw=0.5) + ax.axvline(x=Rmajor - a * deltaU, ls="--", c=c, lw=0.5) + ax.axvline(x=Rmajor - a * deltaL, ls="--", c=c, lw=0.5) + + ax.plot([self.Rmajor, self.Rmag], [self.Zmajor, self.Zmag], "o", markersize=5) + + def write(self, filename=None): + """ + If filename is None, use the original one + """ + + self.g.write_geqdsk(f_path=filename) + + # ----------------------------------------------------------------------------- + # Parameterizations + # ----------------------------------------------------------------------------- + + def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): + + psis = self.psi_pol_norm + flux_surfaces = self.g.fluxsurfaces["psi"] + + # Cannot parallelize because different number of points? + kappa, rmin, rmaj, zmag, sn, cn = [],[],[],[],[],[] + + for flux in range(len(flux_surfaces)): + if flux == len(flux_surfaces)-1: + Rf, Zf = self.Rb, self.Yb + else: + Rf, Zf = self.g.fluxsurfaces["R"][flux], self.g.fluxsurfaces["Z"][flux] + + # Perform the MXH decompositionusing the MITIM surface class + surfaces = mitim_flux_surfaces() + surfaces.reconstruct_from_RZ(Rf,Zf) + surfaces._to_mxh(n_coeff=n_coeff) + + kappa.append(surfaces.kappa[0]) + rmin.append(surfaces.a[0]) + rmaj.append(surfaces.R0[0]) + zmag.append(surfaces.Z0[0]) + + sn.append(surfaces.sn[0,:]) + cn.append(surfaces.cn[0,:]) + + kappa = np.array(kappa) + rmin = np.array(rmin) + rmaj = np.array(rmaj) + zmag = np.array(zmag) + sn = np.array(sn) + cn = np.array(cn) + + if plotYN: + fig, ax = plt.subplots() + ax.plot(self.Rb, self.Yb, 'o-', c = 'b') + + surfaces = mitim_flux_surfaces() + surfaces.reconstruct_from_RZ(self.Rb, self.Yb) + surfaces._to_mxh(n_coeff=n_coeff) + surfaces._from_mxh() + + ax.plot(surfaces.R[0], surfaces.Z[0], 'o-', c = 'r') + + plt.show() + + ''' + Reminder that + sn = [0.0, np.arcsin(delta), -zeta, ...] + ''' + + return psis, rmaj, rmin, zmag, kappa, cn, sn + + # ----------------------------------------------------------------------------- + # For MAESTRO and TRANSP converstions + # ----------------------------------------------------------------------------- + def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH = 7, plotYN = False): + + # ------------------------------------------------------------------------------------------------------- + # Quantities from the equilibrium + # ------------------------------------------------------------------------------------------------------- + + rhotor = self.g.derived['rho_tor'] + psi = self.g.derived['psi'] # Wb/rad + torfluxa = self.g.derived['phi'][-1] / (2*np.pi) # Wb/rad + q = self.g.raw['qpsi'] + pressure = self.g.raw['pres'] # Pa + Ip = self.g.raw['current']*1E-6 # MA + + RZ = np.array([self.Rb,self.Yb]).T + R0 = (RZ.max(axis=0)[0] + RZ.min(axis=0)[0])/2 + B0 = self.g.raw['rcentr']*self.g.raw['bcentr'] / R0 + + # Ensure positive quantities #TODO: Check if this is necessary, pass directions + rhotor = np.array([np.abs(i) for i in rhotor]) + psi = np.array([np.abs(i) for i in psi]) + q = np.array([np.abs(i) for i in q]) + pressure = np.array([np.abs(i) for i in pressure]) + + torfluxa = np.abs(torfluxa) + Ip = np.abs(Ip) + B0 = np.abs(B0) + # ------------------------------------------ + + _, rmaj, rmin, zmag, kappa, cn, sn = self.get_MXH_coeff_new(n_coeff=coeffs_MXH) + + delta = np.sin(sn[:,1]) + zeta = -sn[:,2] + + # ------------------------------------------------------------------------------------------------------- + # Pass to profiles + # ------------------------------------------------------------------------------------------------------- + + profiles = {} + + profiles['nexp'] = np.array([f'{rhotor.shape[0]}']) + profiles['nion'] = np.array(['2']) + profiles['shot'] = np.array(['12345']) + + # Just one specie + profiles['name'] = np.array(['D','F']) + profiles['type'] = np.array(['[therm]','[therm]']) + profiles['masse'] = np.array([5.4488748e-04]) + profiles['mass'] = np.array([2.0, Z*2]) + profiles['ze'] = np.array([-1.0]) + profiles['z'] = np.array([1.0, Z]) + + profiles['torfluxa(Wb/radian)'] = np.array([torfluxa]) + profiles['rcentr(m)'] = np.array([R0]) + profiles['bcentr(T)'] = np.array([B0]) + profiles['current(MA)'] = np.array([Ip]) + + profiles['rho(-)'] = rhotor + profiles['polflux(Wb/radian)'] = psi + profiles['q(-)'] = q + + # ------------------------------------------------------------------------------------------------------- + # Flux surfaces + # ------------------------------------------------------------------------------------------------------- + + profiles['kappa(-)'] = kappa + profiles['delta(-)'] = delta + profiles['zeta(-)'] = zeta + profiles['rmin(m)'] = rmin + profiles['rmaj(m)'] = rmaj + profiles['zmag(m)'] = zmag + + sn, cn = np.array(sn), np.array(cn) + for i in range(coeffs_MXH): + profiles[f'shape_cos{i}(-)'] = cn[:,i] + for i in range(coeffs_MXH-3): + profiles[f'shape_sin{i+3}(-)'] = sn[:,i+3] + + ''' + ------------------------------------------------------------------------------------------------------- + Kinetic profiles + ------------------------------------------------------------------------------------------------------- + Pressure division into temperature and density + p_Pa = p_e + p_i = Te_eV * e_J * ne_20 * 1e20 + Ti_eV * e_J * ni_20 * 1e20 + if T=Te=Ti and ne=ni + p_Pa = 2 * T_eV * e_J * ne_20 * 1e20 + T_eV = p_Pa / (2 * e_J * ne_20 * 1e20) + ''' + + C = 1 / (2 * 1.60217662e-19 * 1e20) + _, ne_20 = PLASMAtools.parabolicProfile(Tbar=ne0_20/1.25,nu=1.25,rho=rhotor,Tedge=ne0_20/5) + T_keV = C * (pressure / ne_20) * 1E-3 + + fZ = (Zeff-1) / (Z**2-Z) # One-impurity model to give desired Zeff + + profiles['te(keV)'] = T_keV + profiles['ti(keV)'] = np.array([T_keV]*2).T + profiles['ne(10^19/m^3)'] = ne_20*10.0 + profiles['ni(10^19/m^3)'] = np.array([profiles['ne(10^19/m^3)']*(1-Z*fZ),profiles['ne(10^19/m^3)']*fZ]).T + + # ------------------------------------------------------------------------------------------------------- + # Power: insert parabolic and use PROFILES volume integration to find desired power + # ------------------------------------------------------------------------------------------------------- + + _, profiles["qrfe(MW/m^3)"] = PLASMAtools.parabolicProfile(Tbar=1.0,nu=5.0,rho=rhotor,Tedge=0.0) + + p = PROFILEStools.gacode_state.scratch(profiles) + + p.profiles["qrfe(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] * PichT/p.derived['qRF_MW'][-1] /2 + p.profiles["qrfi(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] + + # ------------------------------------------------------------------------------------------------------- + # Ready to go + # ------------------------------------------------------------------------------------------------------- + + p.derive_quantities() + + # ------------------------------------------------------------------------------------------------------- + # Plotting + # ------------------------------------------------------------------------------------------------------- + + if plotYN: + + fig, ax = plt.subplots() + ff = np.linspace(0, 1, 11) + self.plotFluxSurfaces(ax=ax, fluxes=ff, rhoPol=False, sqrt=True, color="r", plot1=False) + p.plot_state_flux_surfaces(ax=ax, surfaces_rho=ff, color="b") + plt.show() + + return p + + def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', ne0_20 = 1E19, Vsurf = 0.0, Zeff = 1.5, PichT_MW = 11.0, times = [0.0,1.0]): + + print("\t- Converting to TRANSP") + folder = IOtools.expandPath(folder) + folder.mkdir(parents=True, exist_ok=True) + + p = self.to_profiles(ne0_20 = ne0_20, Zeff = Zeff, PichT = PichT_MW) + p.write_state(folder / 'input.gacode') + + transp = p.to_transp(folder = folder, shot = shot, runid = runid, times = times, Vsurf = Vsurf) + + return transp + + # --------------------------------------------------------------------------------------------------------------------------------------- + # Plotting + # --------------------------------------------------------------------------------------------------------------------------------------- + + def plot(self, fn=None, extraLabel=""): + GEQplotting.plot(self, fn=fn, extraLabel=extraLabel) + + def plotFS(self, axs=None, color="b", label=""): + GEQplotting.plotFS(self, axs=axs, color=color, label=label) + + def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): + GEQplotting.plotPlasma(self, axs=axs, legendYN=legendYN, color=color, label=label) + + def plotGeometry(self, axs=None, color="r"): + GEQplotting.plotGeometry(self, axs=axs, color=color) + + def plotFluxSurfaces(self, ax=None, fluxes=[1.0], color="b", alpha=1.0, rhoPol=True, sqrt=False, lw=1, lwB=2, plot1=True, label = ''): + return GEQplotting.plotFluxSurfaces(self, ax=ax, fluxes=fluxes, color=color, alpha=alpha, rhoPol=rhoPol, sqrt=sqrt, lw=lw, lwB=lwB, plot1=plot1, label = label) + + def plotXpointEnvelope(self, ax=None, color="b", alpha=1.0, rhoPol=True, sqrt=False): + GEQplotting.plotXpointEnvelope(self, ax=ax, color=color, alpha=alpha, rhoPol=rhoPol, sqrt=sqrt) + +# ----------------------------------------------------------------------------- +# Tools to handle flux surface definitions +# ----------------------------------------------------------------------------- +class mitim_flux_surfaces: + + def reconstruct_from_mxh_moments(self,R0, a, kappa, Z0, cn, sn, thetas = None): + ''' + sn = [0.0, np.arcsin(delta), -zeta, ...] + cn = [...] + + You can provide a multi-dim array of (radii, ) + ''' + + self.R0 = R0 + self.a = a + self.kappa = kappa + self.Z0 = Z0 + self.cn = cn + self.sn = sn + + self.delta = np.sin(self.sn[...,1]) + self.zeta = -self.sn[...,2] + + self._from_mxh(thetas = thetas) + + def reconstruct_from_miller(self,R0, a, kappa, Z0, delta, zeta, thetas = None): + ''' + sn = [0.0, np.arcsin(delta), -zeta, ...] + cn = [...] + + You can provide a multi-dim array of (radii, ) or not + ''' + + self.R0 = R0 + self.a = a + self.kappa = kappa + self.Z0 = Z0 + self.delta = delta + self.zeta = zeta + + self.cn = np.array([0.0,0.0,0.0]) + self.sn = np.array([0.0,np.arcsin(self.delta),-self.zeta]) + + self._from_mxh(thetas = thetas) + + def _from_mxh(self, thetas = None): + + self.R, self.Z = from_mxh_to_RZ(self.R0, self.a, self.kappa, self.Z0, self.cn, self.sn, thetas = thetas) + + def reconstruct_from_RZ(self, R, Z): + + self.R = R + self.Z = Z + + if len(self.R.shape) == 1: + self.R = self.R[np.newaxis,:] + self.Z = self.Z[np.newaxis,:] + + self._to_miller() + + def _to_mxh(self, n_coeff=6): + + self.cn = np.zeros((self.R.shape[0],n_coeff)) + self.sn = np.zeros((self.R.shape[0],n_coeff)) + self.gn = np.zeros((self.R.shape[0],4)) + self.gn[-1] = 1.0 + for i in range(self.R.shape[0]): + self.cn[i,:], self.sn[i,:], self.gn[i,:] = from_RZ_to_mxh(self.R[i,:], self.Z[i,:], n_coeff=n_coeff) + + [self.R0, self.a, self.Z0, self.kappa] = self.gn.T + + def _to_miller(self): + + Rmin = np.min(self.R, axis=-1) + Rmax = np.max(self.R, axis=-1) + Zmax = np.max(self.Z, axis=-1) + Zmin = np.min(self.Z, axis=-1) + + self.R0 = 0.5* (Rmax + Rmin) + self.Z0 = 0.5* (Zmax + Zmin) + + self.a = (Rmax - Rmin) / 2 + + # Elongations + + self.kappa_u = (Zmax - self.Z0) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) + self.kappa_l = (self.Z0 - Zmin) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) + self.kappa = (self.kappa_u + self.kappa_l) / 2 + + # Triangularities + + RatZmax = self.R[np.arange(self.R.shape[0]),np.argmax(self.Z,axis=-1)] + self.delta_u = (self.R0-RatZmax) / self.a + + RatZmin = self.R[np.arange(self.R.shape[0]),np.argmin(self.Z,axis=-1)] + self.delta_l = (self.R0-RatZmin) / self.a + + self.delta = (self.delta_u + self.delta_l) / 2 + + # Squareness (not parallel for the time being) + self.zeta = np.zeros(self.R0.shape) + for i in range(self.R0.shape[0]): + try: + Ri, Zi, zeta_uo = find_squareness_points(self.R[i,:], self.Z[i,:]) + except AttributeError: + zeta_uo = np.nan + self.zeta[i] = zeta_uo + + def plot(self, ax = None, color = 'r', label = None, plot_extremes=False): + + if ax is None: + fig, ax = plt.subplots() + + for i in range(self.R.shape[0]): + ax.plot(self.R[i,:], self.Z[i,:], color = color, label = label) + + if plot_extremes: + ax.plot([self.R[i,self.Z[i,:].argmax()]], [self.Z[i,:].max()], 'o', color=color, markersize=5) + ax.plot([self.R[i,self.Z[i,:].argmin()]], [self.Z[i,:].min()], 'o', color=color, markersize=5) + + ax.set_aspect('equal') + ax.set_xlabel('R [m]') + ax.set_ylabel('Z [m]') + if label is not None: + ax.legend() + GRAPHICStools.addDenseAxis(ax) + +def find_squareness_points(R, Z, debug = False): + + # Reference point (A) + A_r = R_of_maxZ = R[Z.argmax()] + A_z = Z_of_maxR = Z[R.argmax()] + + # Upper Outer Squareness point (D) + C_r = R_of_maxR = R.max() + C_z = Z_of_maxZ = Z.max() + + # Find intersection with separatrix (C) + Ri_uo, Zi_uo = find_intersection_squareness(R, Z, R_of_maxZ, Z_of_maxR, R_of_maxR, Z_of_maxZ) + C_r, C_z = Ri_uo, Zi_uo + + # Find point B + B_r = Rs_uo = R_of_maxZ + (R_of_maxR-R_of_maxZ)/2 + B_z = Zs_uo = Z_of_maxR + (Z_of_maxZ-Z_of_maxR)/2 + + # Squareness is defined as the distance BC divided by the distance AB + zeta_uo = np.sqrt((C_r-B_r)**2 + (C_z-B_z)**2) / np.sqrt((A_r-B_r)**2 + (A_z-B_z)**2) + #zeta_uo = np.sqrt((Ri_uo-Rs_uo)**2 + (Zi_uo-Zs_uo)**2) / np.sqrt((R_of_maxZ-R_of_maxR)**2 + (Z_of_maxZ-Z_of_maxR)**2) + + if debug: + plt.ion() + fig, ax = plt.subplots() + ax.plot(R, Z, 'o-', markersize=3, color='k') + ax.plot([Ri_uo], [Zi_uo], 'o', color='k', label = 'C', markersize=6) + + ax.plot([R_of_maxZ], [Z_of_maxR], 'o', color='b', label = 'A') + ax.axvline(x=R_of_maxZ, ls='--', color='b') + ax.axhline(y=Z_of_maxR, ls='--', color='b') + + ax.plot([R_of_maxR], [Z_of_maxZ], 'o', color='r', label = 'D') + ax.axhline(y=Z_of_maxZ, ls='--', color='r') + ax.axvline(x=R_of_maxR, ls='--', color='r') + + ax.plot([Rs_uo], [Zs_uo], 'o', color='g', label = 'B') + + # Connect A and D + ax.plot([R_of_maxZ, R_of_maxR], [Z_of_maxR, Z_of_maxZ], 'm--') + + # Connect maxZ with maxR + ax.plot([R_of_maxZ, R_of_maxR], [Z_of_maxZ, Z_of_maxR], 'm--') + + ax.set_aspect('equal') + ax.legend() + ax.set_xlabel('R [m]') + ax.set_ylabel('Z [m]') + + return Ri_uo, Zi_uo, zeta_uo + +def find_intersection_squareness(R, Z, Ax, Az, Dx, Dz): + + R_line = np.linspace(Ax, Dx, 100) + Z_line = np.linspace(Az, Dz, 100) + line1 = LineString(zip(R_line, Z_line)) + line2 = LineString(zip(R, Z)) + intersection = line1.intersection(line2) + + return intersection.x, intersection.y + +# ----------------------------------------------------------------------------- +# Utilities for parameterizations +# ----------------------------------------------------------------------------- + +def from_RZ_to_mxh(R, Z, n_coeff=3): + """ + Calculates MXH Coefficients for a flux surface + """ + Z = np.roll(Z, -np.argmax(R)) + R = np.roll(R, -np.argmax(R)) + if Z[1] < Z[0]: # reverses array so that theta increases + Z = np.flip(Z) + R = np.flip(R) + + # compute bounding box for each flux surface + r = 0.5*(np.max(R)-np.min(R)) + kappa = 0.5*(np.max(Z) - np.min(Z))/r + R0 = 0.5*(np.max(R)+np.min(R)) + Z0 = 0.5*(np.max(Z)+np.min(Z)) + bbox = [R0, r, Z0, kappa] + + # solve for polar angles + # need to use np.clip to avoid floating-point precision errors + theta_r = np.arccos(np.clip(((R - R0) / r), -1, 1)) + theta = np.arcsin(np.clip(((Z - Z0) / r / kappa),-1,1)) + + # Find the continuation of theta and theta_r to [0,2pi] + theta_r_cont = np.copy(theta_r) + theta_cont = np.copy(theta) + + max_theta = np.argmax(theta) + min_theta = np.argmin(theta) + max_theta_r = np.argmax(theta_r) + min_theta_r = np.argmin(theta_r) + + theta_cont[:max_theta] = theta_cont[:max_theta] + theta_cont[max_theta:max_theta_r] = np.pi-theta[max_theta:max_theta_r] + theta_cont[max_theta_r:min_theta] = np.pi-theta[max_theta_r:min_theta] + theta_cont[min_theta:] = 2*np.pi+theta[min_theta:] + + theta_r_cont[:max_theta] = theta_r_cont[:max_theta] + theta_r_cont[max_theta:max_theta_r] = theta_r[max_theta:max_theta_r] + theta_r_cont[max_theta_r:min_theta] = 2*np.pi - theta_r[max_theta_r:min_theta] + theta_r_cont[min_theta:] = 2*np.pi - theta_r[min_theta:] + + theta_r_cont = theta_r_cont - theta_cont ; theta_r_cont[-1] = theta_r_cont[0] + + # Fourier decompose to find coefficients + + c, s = np.zeros(n_coeff), np.zeros(n_coeff) + + def f_theta_r(theta): + return np.interp(theta, theta_cont, theta_r_cont) + + for i in np.arange(n_coeff): + s[i] = quad(f_theta_r,0,2*np.pi, weight="sin", wvar=i)[0]/np.pi + c[i] = quad(f_theta_r,0,2*np.pi, weight="cos", wvar=i)[0]/np.pi + + c[0] /= 2 + + return c, s, bbox + +def from_mxh_to_RZ(R0, a, kappa, Z0, cn, sn, thetas = None): + + if thetas is None: + thetas = np.linspace(0, 2 * np.pi, 100) + + # Prepare data to always have the first dimension a batch (e.g. a radius) for parallel computation + if IOtools.isfloat(R0): + R0 = [R0] + a = [a] + kappa = [kappa] + Z0 = [Z0] + cn = np.array(cn)[np.newaxis,:] + sn = np.array(sn)[np.newaxis,:] + + R0 = np.array(R0) + a = np.array(a) + kappa = np.array(kappa) + Z0 = np.array(Z0) + + R = np.zeros((R0.shape[0],len(thetas))) + Z = np.zeros((R0.shape[0],len(thetas))) + n = np.arange(1, sn.shape[1]) + for i,theta in enumerate(thetas): + theta_R = theta + cn[:,0] + np.sum( cn[:,1:]*np.cos(n*theta) + sn[:,1:]*np.sin(n*theta), axis=-1 ) + R[:,i] = R0 + a*np.cos(theta_R) + Z[:,i] = Z0 + kappa*a*np.sin(theta) + + return R, Z + +# -------------------------------------------------------------------------------------------------------------- +# Fixed boundary stuff +# -------------------------------------------------------------------------------------------------------------- + +class freegs_millerized: + + def __init__(self, R, a, kappa_sep, delta_sep, zeta_sep, z0): + + print("> Fixed-boundary equilibrium with FREEGS") + + print("\t- Initializing miller geometry") + print(f"\t\t* R={R} m, a={a} m, kappa_sep={kappa_sep}, delta_sep={delta_sep}, zeta_sep={zeta_sep}, z0={z0} m") + + self.R0 = R + self.a = a + self.kappa_sep = kappa_sep + self.delta_sep = delta_sep + self.zeta_sep = zeta_sep + self.Z0 = z0 + + thetas = np.linspace(0, 2*np.pi, 1000, endpoint=False) + + self.mitim_separatrix = mitim_flux_surfaces() + self.mitim_separatrix.reconstruct_from_miller(self.R0, self.a, self.kappa_sep, self.Z0, self.delta_sep, self.zeta_sep, thetas = thetas) + self.R_sep, self.Z_sep = self.mitim_separatrix.R[0,:], self.mitim_separatrix.Z[0,:] + + def prep(self, p0_MPa, Ip_MA, B_T, + beta_pol = None, n_coils = 10, resol_eq = 2**8+1, + parameters_profiles = {'alpha_m':2.0, 'alpha_n':2.0, 'Raxis':1.0}, + constraint_miller_squareness_point = False): + + print("\t- Initializing plasma parameters") + if beta_pol is not None: + print(f"\t\t* beta_pol={beta_pol:.5f}, Ip={Ip_MA:.5f} MA, B={B_T:.5f} T") + else: + print(f"\t\t* p0={p0_MPa:.5f} MPa, Ip={Ip_MA:.5f} MA, B={B_T:.5f} T") + + self.p0_MPa = p0_MPa + self.Ip_MA = Ip_MA + self.B_T = B_T + self.beta_pol = beta_pol + self.parameters_profiles = parameters_profiles + + print(f"\t- Preparing equilibrium with FREEGS, with a resolution of {resol_eq}x{resol_eq}") + self._define_coils(n_coils) + self._define_eq(resol = resol_eq) + self._define_gs() + + # Define xpoints + print("\t\t* Defining upper and lower x-points") + self.xpoints = [ + (self.R0 - self.a*self.delta_sep, self.Z0+self.a*self.kappa_sep), + (self.R0 - self.a*self.delta_sep, self.Z0-self.a*self.kappa_sep), + ] + + # Define isoflux + print("\t\t* Defining midplane separatrix (isoflux)") + self.isoflux = [ + (self.xpoints[0][0], self.xpoints[0][1], self.R0 + self.a, self.Z0), # Upper x-point with outer midplane + (self.xpoints[0][0], self.xpoints[0][1], self.R0 - self.a, self.Z0), # Upper x-point with inner midplane + (self.xpoints[0][0], self.xpoints[0][1], self.xpoints[1][0], self.xpoints[1][1]), # Between x-points + ] + + print("\t\t* Defining squareness isoflux point") + + # Find squareness point + if constraint_miller_squareness_point: + Rsq, Zsq, _ = find_squareness_points(self.R_sep, self.Z_sep) + + self.isoflux.append( + (self.xpoints[0][0], self.xpoints[0][1], Rsq, Zsq) # Upper x-point with squareness point + ) + self.isoflux.append( + (self.xpoints[0][0], self.xpoints[0][1], Rsq, -Zsq) # Upper x-point with squareness point + ) + + # Combine + self.constrain = freegs.control.constrain( + isoflux=self.isoflux, + xpoints=self.xpoints, + ) + + def _define_coils(self, n, rel_distance_coils = 0.5, updown_coils = True): + + print(f"\t- Defining {n} coils{' (up-down symmetric)' if updown_coils else ''} at a distance of {rel_distance_coils}*a from the separatrix") + + self.distance_coils = self.a*rel_distance_coils + self.updown_coils = updown_coils + + if self.updown_coils: + thetas = np.linspace(0, np.pi, n) + else: + thetas = np.linspace(0, 2*np.pi, n, endpoint=False) + + self.mitim_coils_surface = mitim_flux_surfaces() + self.mitim_coils_surface.reconstruct_from_miller(self.R0, (self.a+self.distance_coils), self.kappa_sep, self.Z0, self.delta_sep, self.zeta_sep, thetas = thetas) + self.Rcoils, self.Zcoils = self.mitim_coils_surface.R[0,:], self.mitim_coils_surface.Z[0,:] + + self.coils = [] + for num, (R, Z) in enumerate(zip(self.Rcoils, self.Zcoils)): + + if self.updown_coils and Z > 0: + coilU = freegs.machine.Coil( + R, + Z + ) + coilL = freegs.machine.Coil( + R, + -Z + ) + coil = freegs.machine.Circuit( [ ('U', coilU, 1.0 ), ('L', coilL, 1.0 ) ] ) + else: + + coil = freegs.machine.Coil( + R, + Z + ) + + self.coils.append( + (f"coil_{num}", coil) + ) + + def _define_eq(self, resol=2**9+1): + + print("\t- Defining equilibrium") + self.tokamak = freegs.machine.Machine(self.coils) + + a = self.a + self.distance_coils + + Rmin = (self.R0-a) - a*0.25 + Rmax = (self.R0+a) + a*0.25 + + b = a*self.kappa_sep + Zmin = (self.Z0 - b) - b*0.25 + Zmax = (self.Z0 + b) + b*0.25 + + self.eq = freegs.Equilibrium(tokamak=self.tokamak, + Rmin=Rmin, Rmax=Rmax, + Zmin=Zmin, Zmax=Zmax, + nx=resol, ny=resol, + boundary=freegs.boundary.freeBoundaryHagenow) + + def _define_gs(self): + + if self.beta_pol is None: + print("\t- Defining Grad-Shafranov equilibrium: p0, Ip and vaccum R*Bt") + + self.profiles = freegs.jtor.ConstrainPaxisIp(self.eq, + self.p0_MPa*1E6, self.Ip_MA*1E6, self.R0*self.B_T, + alpha_m=self.parameters_profiles['alpha_m'], alpha_n=self.parameters_profiles['alpha_n'], Raxis=self.parameters_profiles['Raxis']) + + else: + print("\t- Defining Grad-Shafranov equilibrium: beta_pol, Ip and vaccum R*Bt") + + self.profiles = freegs.jtor.ConstrainBetapIp(self.eq, + self.beta_pol, self.Ip_MA*1E6, self.R0*self.B_T, + alpha_m=self.parameters_profiles['alpha_m'], alpha_n=self.parameters_profiles['alpha_n'], Raxis=self.parameters_profiles['Raxis']) + + def solve(self, show = False, rtol=1e-6): + + print("\t- Solving equilibrium with FREEGS") + with IOtools.timer(): + self.x,self.y = freegs.solve(self.eq, # The equilibrium to adjust + self.profiles, # The toroidal current profile function + constrain=self.constrain, # Constraint function to set coil currents + show=show, + rtol=rtol, # Default is 1e-3 + atol=1e-10, + maxits=100, # Default is 50 + convergenceInfo=True) + print("\t\t * Done!") + + self.check() + + def check(self, warning_error = 0.01, plotYN = False): + + print("\t- Checking separatrix quality (Miller vs FREEGS)") + RZ = self.eq.separatrix() + + self.mitim_separatrix_eq = mitim_flux_surfaces() + self.mitim_separatrix_eq.reconstruct_from_RZ(RZ[:,0], RZ[:,1]) + + # -------------------------------------------------------------- + # Check errors + # -------------------------------------------------------------- + + max_error = 0.0 + for key in ['R0', 'a', 'kappa_sep', 'delta_sep', 'zeta_sep']: + miller_value = getattr(self, key) + sep_value = getattr(self.mitim_separatrix_eq, key.replace('_sep', ''))[0] + error = abs( (miller_value-sep_value)/miller_value ) + print(f"\t\t* {key}: {miller_value:.3f} vs {sep_value:.3f} ({100*error:.2f}%)") + + max_error = np.max([max_error, error]) + + if max_error > warning_error: + print(f"\t\t- Maximum error in equilibrium quantities is {100*max_error:.2f}%", typeMsg='w') + else: + print(f"\t\t- Maximum error in equilibrium quantities is {100*max_error:.2f}%") + + # -------------------------------------------------------------- + # Plotting + # -------------------------------------------------------------- + + if plotYN: + + fig = plt.figure(figsize=(12,8)) + axs = fig.subplot_mosaic( + """ + AB + AB + CB + """ + ) + + # Plot direct FreeGS output + + ax = axs['A'] + self.eq.plot(axis=ax,show=False) + self.constrain.plot(axis=ax, show=False) + + for coil in self.coils: + if isinstance(coil[1],freegs.machine.Circuit): + ax.plot([coil[1]['U'].R], [coil[1]['U'].Z], 's', c='k', markersize=2) + ax.plot([coil[1]['L'].R], [coil[1]['L'].Z], 's', c='k', markersize=2) + else: + ax.plot([coil[1].R], [coil[1].Z], 's', c='k', markersize=2) + + GRAPHICStools.addLegendApart(ax,ratio=0.9,size=10) + + ax = axs['C'] + ax.plot(self.x,'-o', markersize=3, color='b', label = '$\\psi$ max change') + ax.set_xlabel('Iteration') + ax.set_ylabel('$\\psi$ max change') + ax.set_yscale('log') + ax.legend(loc='lower left',prop={'size': 10}) + + ax = axs['C'].twinx() + ax.plot(self.y,'-o', markersize=3, color='r', label = '$\\psi$ max relative change') + ax.set_ylabel('$\\psi$ max relative change') + ax.set_yscale('log') + ax.legend(loc='upper right',prop={'size': 10}) + + # Plot comparison of equilibria + + ax = axs['B'] + + self.mitim_separatrix.plot(ax=ax, color = 'b', label = 'Miller (original)', plot_extremes=True) + self.mitim_separatrix_eq.plot(ax=ax, color = 'r', label = 'Separatrix (freegs)', plot_extremes=True) + + ax.legend(prop={'size': 10}) + + plt.show() + + def derive(self, psi_surfaces = np.linspace(0,1.0,10), psi_profiles = np.linspace(0,1.0,100)): + + # Grab surfaces + Rs, Zs = [], [] + for psi_norm in psi_surfaces: + R, Z = self.find_surface(psi_norm) + Rs.append(R) + Zs.append(Z) + Rs = np.array(Rs) + Zs = np.array(Zs) + + # Calculate surface stuff in parallel + self.surfaces = mitim_flux_surfaces() + self.surfaces.reconstruct_from_RZ(Rs, Zs) + self.surfaces.psi = psi_surfaces + + # Grab profiles + self.profile_psi_norm = psi_profiles + self.profile_pressure = self.eq.pressure(psinorm =psi_profiles)*1E-6 + self.profile_q = self.eq.q(psinorm = psi_profiles) + self.profile_RB = self.eq.fpol(psinorm = psi_profiles) + + # Grab quantities + self.profile_q95 = self.eq.q(psinorm = 0.95) + self.profile_q0 = self.eq.q(psinorm = 0.0) + self.profile_betaN = self.eq.betaN() + self.profile_Li2 = self.eq.internalInductance2() + self.profile_pave = self.eq.pressure_ave() + self.profile_beta_pol = self.eq.poloidalBeta() + self.profile_Ashaf = self.eq.shafranovShift + + def find_surface(self, psi_norm = 0.5, thetas = None): + + if psi_norm == 0.0: + psi_norm = 1E-6 + + if psi_norm == 1.0: + RZ = self.eq.separatrix(npoints= 1000 if thetas is None else len(thetas)) + R, Z = RZ[:,0], RZ[:,1] + else: + if thetas is None: + thetas = np.linspace(0, 2*np.pi, 1000, endpoint=False) + + from freegs.critical import find_psisurface, find_critical + from scipy import interpolate + + psi = self.eq.psi() + opoint, xpoint = find_critical(self.eq.R, self.eq.Z, psi) + psinorm = (psi - opoint[0][2]) / (self.eq.psi_bndry - opoint[0][2]) + psifunc = interpolate.RectBivariateSpline(self.eq.R[:, 0], self.eq.Z[0, :], psinorm) + r0, z0 = opoint[0][0:2] + + R = np.zeros(len(thetas)) + Z = np.zeros(len(thetas)) + for i,theta in enumerate(thetas): + R[i],Z[i] = find_psisurface( + self.eq, + psifunc, + r0, + z0, + r0 + 10.0 * np.sin(theta), + z0 + 10.0 * np.cos(theta), + psival=psi_norm, + n=1000, + ) + return R,Z + + # -------------------------------------------------------------- + # Plotting + # -------------------------------------------------------------- + + def plot(self, axs = None, color = 'b', label = ''): + + if axs is None: + plt.ion() + fig = plt.figure(figsize=(16,7)) + axs = fig.subplot_mosaic( + """ + A12 + A34 + """) + axs = [axs['A'], axs['1'], axs['2'], axs['3'], axs['4']] + + self.plot_flux_surfaces(ax = axs[0], color = color) + self.plot_profiles(axs = axs[1:], color = color, label = label) + + def plot_flux_surfaces(self, ax = None, color = 'b'): + + if ax is None: + plt.ion() + fig, ax = plt.subplots(figsize=(12,8)) + + for i in range(self.surfaces.R.shape[0]): + ax.plot(self.surfaces.R[i],self.surfaces.Z[i], '-', label = f'$\\psi$ = {self.surfaces.psi[i]:.2f}', color = color, markersize=3) + + ax.set_aspect('equal') + ax.set_xlabel('R [m]') + ax.set_ylabel('Z [m]') + GRAPHICStools.addDenseAxis(ax) + + def plot_profiles(self, axs = None, color = 'b', label = ''): + + if axs is None: + plt.ion() + fig, axs = plt.subplots(nrows=2,ncols=2,figsize=(8,8)) + axs = axs.flatten() + + ax = axs[0] + ax.plot(self.profile_psi_norm,self.profile_pressure,'-',color=color, label = label) + ax.set_xlabel('$\\psi$') + ax.set_xlim([0,1]) + ax.set_ylabel('Pressure (MPa)') + GRAPHICStools.addDenseAxis(ax) + + ax = axs[1] + ax.plot(self.profile_psi_norm,self.profile_q,'-',color=color) + ax.axhline(y=1, color='k', ls='--', lw=0.5) + ax.set_xlabel('$\\psi$') + ax.set_ylabel('q') + ax.set_xlim([0,1]) + GRAPHICStools.addDenseAxis(ax) + + ax = axs[2] + ax.plot(self.profile_psi_norm,self.profile_RB,'-',color=color) + ax.axhline(y=self.R0*self.B_T, color=color, ls='--', lw=0.5) + ax.set_xlabel('$\\psi$') + ax.set_ylabel('$R\\cdot B_t$ ($T\\cdot m$)') + ax.set_xlim([0,1]) + GRAPHICStools.addDenseAxis(ax) + + def plot_flux_surfaces_characteristics(self, axs = None, color = 'b', label = ''): + + if axs is None: + plt.ion() + fig, axs = plt.subplots(nrows=2,ncols=2,figsize=(8,8)) + axs = axs.flatten() + + ax = axs[0] + ax.plot(self.surfaces.psi, self.surfaces.kappa, '-o', color=color, label = label, markersize=3) + ax.set_xlabel('$\\psi$') + ax.set_ylabel('$\\kappa$') + ax.set_xlim([0,1]) + GRAPHICStools.addDenseAxis(ax) + + ax = axs[1] + ax.plot(self.surfaces.psi, self.surfaces.delta, '-o', color=color, label = label, markersize=3) + ax.set_xlabel('$\\psi$') + ax.set_ylabel('$\\delta$') + ax.set_xlim([0,1]) + GRAPHICStools.addDenseAxis(ax) + + # -------------------------------------------------------------- + # Writing + # -------------------------------------------------------------- + + def write(self, filename): + + print(f"\t- Writing equilibrium to {IOtools.clipstr(filename)}") + + with open(filename, "w") as f: + geqdsk.write(self.eq, f) + + def to_profiles(self, scratch_folder = '~/scratch/'): + + # Produce geqdsk object + scratch_folder = IOtools.expandPath(scratch_folder) + file_scratch = scratch_folder / 'mitim_freegs.geqdsk' + self.write(file_scratch) + g = MITIMgeqdsk(file_scratch) + + os.remove(file_scratch) + + # From geqdsk to profiles + return g.to_profiles() + + def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', ne0_20 = 1E19, Vsurf = 0.0, Zeff = 1.5, PichT_MW = 11.0, times = [0.0,1.0]): + + # Produce geqdsk object + scratch_folder = IOtools.expandPath(folder) + scratch_folder.mkdir(parents=True, exist_ok=True) + file_scratch = scratch_folder / 'mitim_freegs.geqdsk' + self.write(file_scratch) + g = MITIMgeqdsk(file_scratch) + + return g.to_transp(folder=folder, shot=shot, runid=runid, ne0_20=ne0_20, Vsurf=Vsurf, Zeff=Zeff, PichT_MW=PichT_MW, times=times) diff --git a/src/mitim_tools/gs_tools/utils/GEQplotting_megpy.py b/src/mitim_tools/gs_tools/utils/GEQplotting_megpy.py new file mode 100644 index 00000000..db545acd --- /dev/null +++ b/src/mitim_tools/gs_tools/utils/GEQplotting_megpy.py @@ -0,0 +1,984 @@ +import numpy as np +import matplotlib.pyplot as plt +from mitim_tools.misc_tools import GRAPHICStools +from mitim_tools.misc_tools.LOGtools import printMsg as print +from IPython import embed + +#TODO: add current profiles and flux-surface average fields to megpy and restore plots + +def compareGeqdsk(geqdsks, fn=None, extraLabel="", plotAll=True, labelsGs=None): + + if fn is None: + from mitim_tools.misc_tools.GUItools import FigureNotebook + fn = FigureNotebook("GEQDSK Notebook", geometry="1600x1000") + + if labelsGs is None: + labelsGs = [] + for i, g in enumerate(geqdsks): + labelsGs.append(f"#{i + 1}") + + # ----------------------------------------------------------------------------- + # Plot All + # ----------------------------------------------------------------------------- + if plotAll: + for i, g in enumerate(geqdsks): + _ = g.plot(fn=fn, extraLabel=f"{labelsGs[i]} - ") + + # ----------------------------------------------------------------------------- + # Compare in same plot - Surfaces + # ----------------------------------------------------------------------------- + fig = fn.add_figure(label=extraLabel + "Comp. - Surfaces") + + grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) + ax1 = fig.add_subplot(grid[:, 0]) + ax2 = fig.add_subplot(grid[:, 1]) + ax3 = fig.add_subplot(grid[0, 2]) + ax4 = fig.add_subplot(grid[1, 2]) + axs = [ax1, ax2, ax3, ax4] + + cols = GRAPHICStools.listColors() + + for i, g in enumerate(geqdsks): + g.plotFS(axs=axs, color=cols[i], label=f"#{i + 1}") + + ax3.legend() + + # ----------------------------------------------------------------------------- + # Compare in same plot - Surfaces + # ----------------------------------------------------------------------------- + fig = fn.add_figure(label=extraLabel + "Comp. - Plasma") + grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) + + ax_plasma = [ + fig.add_subplot(grid[0, 0]), + fig.add_subplot(grid[1, 0]), + fig.add_subplot(grid[0, 1]), + fig.add_subplot(grid[1, 1]), + fig.add_subplot(grid[0, 2]), + fig.add_subplot(grid[1, 2]), + fig.add_subplot(grid[0, 3]), + ] # , + # fig.add_subplot(grid[1,3])] + + cols = GRAPHICStools.listColors() + + for i, g in enumerate(geqdsks): + g.plotFS(axs=[ax1, ax2, ax3, ax4], color=cols[i], label=f"{labelsGs[i]} ") + g.plotPlasma( + axs=ax_plasma, + legendYN=False, + color=cols[i], + label=f"{labelsGs[i]} ", + ) + + return ax_plasma, fn + +# ----------------------------------------------------------------------------- +# Plot of GEQ class +# ----------------------------------------------------------------------------- + +def plot(self, fn=None, extraLabel=""): + if fn is None: + wasProvided = False + + from mitim_tools.misc_tools.GUItools import FigureNotebook + + self.fn = FigureNotebook("GEQDSK Notebook", geometry="1600x1000") + else: + wasProvided = True + self.fn = fn + # ----------------------------------------------------------------------------- + # OMFIT Summary + # ----------------------------------------------------------------------------- + # fig = self.fn.add_figure(label=extraLabel+'OMFIT Summ') + # self.g.plot() + + # ----------------------------------------------------------------------------- + # Flux + # ----------------------------------------------------------------------------- + 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]) + ax3 = fig.add_subplot(grid[0, 2]) + ax4 = fig.add_subplot(grid[1, 2]) + + self.plotFS(axs=[ax1, ax2, ax3, ax4]) + + # ----------------------------------------------------------------------------- + # 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]) + ax3 = fig.add_subplot(grid[2, 1]) + ax4 = fig.add_subplot(grid[:2, 1]) + ax5 = fig.add_subplot(grid[2, 2]) + ax6 = fig.add_subplot(grid[:2, 2]) + ax7 = fig.add_subplot(grid[2, 3]) + ax8 = fig.add_subplot(grid[:2, 3]) + ax9 = fig.add_subplot(grid[2, 4]) + ax10 = fig.add_subplot(grid[:2, 4]) + + plotCurrents(self, + axs=[ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10], zlims_thr=[-1, 1] + ) + + # ----------------------------------------------------------------------------- + # 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]) + ax3 = fig.add_subplot(grid[2, 1]) + ax4 = fig.add_subplot(grid[:2, 1]) + ax5 = fig.add_subplot(grid[2, 2]) + ax6 = fig.add_subplot(grid[:2, 2]) + ax7 = fig.add_subplot(grid[2, 3]) + ax8 = fig.add_subplot(grid[:2, 3]) + ax9 = fig.add_subplot(grid[2, 4]) + ax10 = fig.add_subplot(grid[:2, 4]) + + plotFields(self, + axs=[ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10], zlims_thr=[-1, 1] + ) + + # ----------------------------------------------------------------------------- + # Checks + # ----------------------------------------------------------------------------- + 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() + ax2 = fig.add_subplot(grid[1, 0]) + ax3 = fig.add_subplot(grid[:, 1]) + ax4 = fig.add_subplot(grid[:, 2]) + ax5 = fig.add_subplot(grid[:, 3]) + + plotChecks(self,axs=[ax1, ax1E, ax2, ax3, ax4, ax5]) + + # ----------------------------------------------------------------------------- + # Parameterization + # ----------------------------------------------------------------------------- + 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]) + ax3 = fig.add_subplot(grid[1, 1]) + ax4 = fig.add_subplot(grid[2, 1]) + ax5 = fig.add_subplot(grid[:, 2]) + + plotParameterization(self,axs=[ax1, ax2, ax3, ax4, ax5]) + + # ----------------------------------------------------------------------------- + # Plasma + # ----------------------------------------------------------------------------- + fig = self.fn.add_figure(label=extraLabel + "Plasma") + grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) + + ax_plasma = [ + fig.add_subplot(grid[0, 0]), + fig.add_subplot(grid[1, 0]), + fig.add_subplot(grid[0, 1]), + fig.add_subplot(grid[1, 1]), + fig.add_subplot(grid[0, 2]), + fig.add_subplot(grid[1, 2]), + fig.add_subplot(grid[0, 3]), + fig.add_subplot(grid[1, 3]), + ] + ax_plasma = self.plotPlasma(axs=ax_plasma, legendYN=not wasProvided) + + # ----------------------------------------------------------------------------- + # 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]) + ax3 = fig.add_subplot(grid[1, 0]) + ax4 = fig.add_subplot(grid[1, 1]) + + self.plotGeometry(axs=[ax1, ax2, ax3, ax4]) + + return ax_plasma + +def plotFS(self, axs=None, color="b", label=""): + if axs is None: + plt.ion() + fig, axs = plt.subplots(ncols=4) + + ax = axs[0] + self.plotFluxSurfaces( + ax=ax, fluxes=np.linspace(0, 1, 21), rhoPol=True, sqrt=False, color=color + ) + ax.plot(self.Rb, self.Yb, lw=1, c="r") + ax.set_title("Poloidal Flux") + ax.set_aspect("equal") + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + + ax = axs[1] + self.plotFluxSurfaces( + ax=ax, fluxes=np.linspace(0, 1, 21), rhoPol=False, sqrt=True, color=color + ) + ax.plot(self.Rb, self.Yb, lw=1, c="r") + ax.set_title("Sqrt Toroidal Flux") + ax.set_aspect("equal") + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + + ax = axs[2] + x = self.psi_pol_norm + y = self.rho_tor + ax.plot(x, y, lw=2, ls="-", c=color, label=label) + ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) + + ax.set_xlabel("$\\Psi_n$ (PSI_NORM)") + ax.set_ylabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_xlim([0, 1]) + ax.set_ylim([0, 1]) + + ax = axs[3] + x = self.rho_tor + y = self.rho_pol + ax.plot(x, y, lw=2, ls="-", c=color) + ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) + + ax.set_ylabel("$\\sqrt{\\Psi_n}$ (RHOp)") + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_xlim([0, 1]) + ax.set_ylim([0, 1]) + +def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): + if axs is None: + plt.ion() + fig, axs = plt.subplots(ncols=10) + + ax = axs[0] + x = self.psi_pol_norm + y = np.zeros(x.shape) + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_ylabel("FSA $\\langle J\\rangle$ ($MA/m^2$)") + ax.set_xlim([0, 1]) + + zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] + zlims = GRAPHICStools.aroundZeroLims(zlims) + ax.set_ylim(zlims) + + #ax = axs[1] + #plot2Dquantity(self, + # ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 + #) + + ax = axs[2] + x = self.psi_pol_norm + y = np.zeros(x.shape) + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] + zlims = GRAPHICStools.aroundZeroLims(zlims) + ax.set_ylim(zlims) + + #ax = axs[3] + #plot2Dquantity(self, + # ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 + #) + + ax = axs[4] + x = self.psi_pol_norm + y = self.Jt + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] + zlims = GRAPHICStools.aroundZeroLims(zlims) + ax.set_ylim(zlims) + + #ax = axs[5] + #plot2Dquantity(self, + # ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 + #) + + ax = axs[6] + x = self.psi_pol_norm + y = np.zeros(x.shape) + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] + zlims = GRAPHICStools.aroundZeroLims(zlims) + ax.set_ylim(zlims) + + #ax = axs[7] + #plot2Dquantity(self, + # ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 + #) + + ax = axs[8] + x = self.psi_pol_norm + y = np.zeros(x.shape) + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] + zlims = GRAPHICStools.aroundZeroLims(zlims) + ax.set_ylim(zlims) + + #ax = axs[9] + #plot2Dquantity(self, + # ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 + #) + +def plotFields(self, axs=None, zlims_thr=[-1, 1]): + if axs is None: + plt.ion() + fig, axs = plt.subplots(ncols=10) + + ax = axs[0] + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Br") + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_ylabel("FSA $\\langle B\\rangle$ ($T$)") + ax.set_xlim([0, 1]) + + zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] + zlims = GRAPHICStools.aroundZeroLims(zlims) + ax.set_ylim(zlims) + + ax = axs[1] + plot2Dquantity(self, + ax=ax, var="B_r", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" + ) + + ax = axs[2] + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Bz") + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] + zlims = GRAPHICStools.aroundZeroLims(zlims) + ax.set_ylim(zlims) + + ax = axs[3] + plot2Dquantity(self, + ax=ax, var="B_z", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" + ) + + ax = axs[4] + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Bt") + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + zlims = [y.min(), y.max()] + # zlims = [np.min([zlims_thr[0],y.min()]),np.max([zlims_thr[1],y.max()])] + # zlims = GRAPHICStools.aroundZeroLims(zlims) + ax.set_ylim(zlims) + + #ax = axs[5] + #plot2Dquantity(self, + # ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" + #) + + ax = axs[6] + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Bp") + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] + zlims = GRAPHICStools.aroundZeroLims(zlims) + ax.set_ylim(zlims) + + ax = axs[7] + plot2Dquantity(self, + ax=ax, var="B_pol_rz", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" + ) + + ax = axs[8] + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g["fluxSurfaces"]["avg"]["Bp**2"] + ax.plot(x, y, lw=2, ls="-", c="r") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + ax.set_ylabel("$\\langle B_{\\theta}^2\\rangle$") + + #ax = axs[9] + #x = self.g["fluxSurfaces"]["midplane"]["R"] + #y = self.g["fluxSurfaces"]["midplane"]["Bt"] + #ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") + #y = self.g["fluxSurfaces"]["midplane"]["Bp"] + #ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") + #y = self.g["fluxSurfaces"]["midplane"]["Bz"] + #ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") + #y = self.g["fluxSurfaces"]["midplane"]["Br"] + #ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") + #y = self.g["fluxSurfaces"]["geo"]["bunit"] + #ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") + #ax.set_xlabel("$R$ LF midplane") + #ax.set_ylabel("$B$ (T)") + #ax.legend() + +def plotChecks(self, axs=None): + if axs is None: + plt.ion() + fig, axs = plt.subplots(ncols=8) + + ax = axs[0] + x = self.psi_pol_norm + y1 = self.Jt + ax.plot(x, np.abs(y1), lw=2, ls="-", c="b", label="$\\langle Jt\\rangle$") + zmax = y1.max() + zmin = y1.min() + y2 = self.Jt_fb + ax.plot(x, np.abs(y2), lw=2, ls="-", c="g", label="$\\langle Jt_{FB}\\rangle$") + + y3 = self.Jerror + ax.plot( + x, + y3, + lw=1, + ls="-", + c="r", + label="$|\\langle Jt\\rangle-\\langle Jt_{FB}\\rangle|$", + ) + + ax.set_ylabel("Current Density ($MA/m^2$)") + + axE = axs[1] + yErr = np.abs(self.Jerror / self.Jt) * 100.0 + axE.plot(x, yErr, lw=0.5, ls="--", c="b") + axE.set_ylim([0, 50]) + axE.set_ylabel("Relative Error (%)") + + ax.set_title("$|\\langle Jt\\rangle|$") + ax.set_ylim(bottom=0) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\Psi_n$") + ax.legend() + + ax = axs[2] + x = self.psi_pol_norm + y1 = self.g.raw["ffprim"] + ax.plot(x, y1, lw=2, ls="-", c="r", label="$FF'$") + y2 = self.g.raw["pprime"] * (4 * np.pi * 1e-7) + ax.plot(x, y2, lw=2, ls="-", c="b", label="$p'*\\mu_0$") + + ax.set_ylabel("") + ax.legend() + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + + #ax = axs[3] + #plot2Dquantity(self, + # ax=ax, + # var="Jt", + # title="Toroidal Current Jt", + # zlims=[zmin, zmax], + # cmap="viridis", + # factor=1e-6, + #) + + #ax = axs[4] + #plot2Dquantity(self, + # ax=ax, + # var="Jt_fb", + # title="Toroidal Current Jt (FB)", + # zlims=[zmin, zmax], + # cmap="viridis", + # factor=1e-6, + #) + + #ax = axs[5] + #z = ( + # np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) + # * 1e-6 + #) + #zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) + #plot2Dquantity(self, + # ax=ax, + # var=z, + # title="Absolute Error", + # zlims=[0, zmaxx], + # cmap="viridis", + # direct=True, + #) + +def plotParameterization(self, axs=None): + if axs is None: + plt.ion() + fig, axs = plt.subplots(ncols=5) + + ax = axs[0] + cs, csA = self.plotFluxSurfaces( + ax=ax, fluxes=np.linspace(0, 1, 21), rhoPol=True, sqrt=False + ) + # Boundary, axis and limiter + ax.plot(self.Rb, self.Yb, lw=1, c="r") + ax.plot(self.g.raw["rmaxis"], self.g.raw["zmaxis"], "+", markersize=10, c="r") + ax.plot([self.Rmag], [self.Zmag], "o", markersize=5, c="m") + ax.plot([self.Rmajor], [self.Zmag], "+", markersize=10, c="k") + if 'rlim' in self.g.raw and 'zlim' in self.g.raw: + ax.plot(self.g.raw["rlim"], self.g.raw["zlim"], lw=1, c="k") + + import matplotlib + + path = matplotlib.path.Path( + np.transpose(np.array([self.g.raw["rlim"], self.g.raw["zlim"]])) + ) + patch = matplotlib.patches.PathPatch(path, facecolor="none") + ax.add_patch(patch) + # for col in cs.collections: + # col.set_clip_path(patch) + # for col in csA.collections: + # col.set_clip_path(patch) + + self.plotEnclosingBox(ax=ax) + + ax.set_aspect("equal") + ax.set_title("Poloidal Flux") + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + + ax = axs[1] + x = self.psi_pol_norm + y = self.g.derived["miller_geo"]["kappa"].copy() + ax.plot(x, y, label="$\\kappa$") + y = self.g.derived["miller_geo"]["kappa_l"].copy() + ax.plot(x, y, ls="--", label="$\\kappa_L$") + y = self.g.derived["miller_geo"]["kappa_u"].copy() + ax.plot(x, y, ls="--", label="$\\kappa_U$") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + ax.set_ylabel("Elongation $\\kappa$") + ax.legend() + + ax = axs[2] + x = self.psi_pol_norm + y = self.g.derived["miller_geo"]["delta"].copy() + ax.plot(x, y, label="$\\delta$") + y = self.g.derived["miller_geo"]["delta_l"].copy() + ax.plot(x, y, ls="--", label="$\\delta_L$") + y = self.g.derived["miller_geo"]["delta_u"].copy() + ax.plot(x, y, ls="--", label="$\\delta_U$") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + ax.set_ylabel("Triangularity $\\delta$") + ax.legend() + + ax = axs[3] + x = self.psi_pol_norm + y = self.g.derived["miller_geo"]["zeta"].copy() + ax.plot(x, y, label="$\\zeta$") + y = self.g.derived["miller_geo"]["zeta_li"].copy() + ax.plot(x, y, ls="--", label="$\\zeta_{IL}$") + y = self.g.derived["miller_geo"]["zeta_ui"].copy() + ax.plot(x, y, ls="--", label="$\\zeta_{IU}$") + y = self.g.derived["miller_geo"]["zeta_lo"].copy() + ax.plot(x, y, ls="--", label="$\\zeta_{OL}$") + y = self.g.derived["miller_geo"]["zeta_uo"].copy() + ax.plot(x, y, ls="--", label="$\\zeta_{OU}$") + ax.set_xlabel("$\\Psi_n$") + ax.set_xlim([0, 1]) + ax.set_ylabel("Squareness $\\zeta$") + ax.legend() + + ax = axs[4] + ax.text( + 0.0, + 11.0, + "Rmajor = {0:.3f}m, Rmag = {1:.3f}m (Zmag = {2:.3f}m)".format( + self.Rmajor, self.Rmag, self.Zmag + ), + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + ) + ax.text( + 0.0, + 10.0, + f"a = {self.a:.3f}m, eps = {self.eps:.3f}", + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + ) + ax.text( + 0.0, + 9.0, + "kappa = {0:.3f} (kU = {1:.3f}, kL = {2:.3f})".format( + self.kappa, self.kappaU, self.kappaL + ), + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + ) + ax.text( + 0.0, + 8.0, + f" kappa95 = {self.kappa95:.3f}, kappa995 = {self.kappa995:.3f}", + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + ) + ax.text( + 0.0, + 7.0, + f" kappa_areal = {self.kappa_a:.3f}", + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + ) + ax.text( + 0.0, + 6.0, + "delta = {0:.3f} (dU = {1:.3f}, dL = {2:.3f})".format( + self.delta, self.deltaU, self.deltaL + ), + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + ) + ax.text( + 0.0, + 5.0, + f" delta95 = {self.delta95:.3f}, delta995 = {self.delta995:.3f}", + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + ) + ax.text( + 0.0, + 4.0, + f"zeta = {self.zeta:.3f}", + color="k", + fontsize=10, + fontweight="normal", + horizontalalignment="left", + verticalalignment="bottom", + rotation=0, + ) + + ax.set_ylim([0, 12]) + ax.set_xlim([-1, 1]) + + ax.set_axis_off() + +def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): + if axs is None: + plt.ion() + fig, axs = plt.subplots(ncols=7) + + ax_plasma = axs + + ax = ax_plasma[0] + ax.plot( + self.rho_tor, + self.g.raw["pres"] * 1e-6, + "-s", + c=color, + lw=2, + markersize=3, + label="geqdsk p", + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_ylim(bottom=0) + ax.set_ylabel("pressure (MPa)") + + ax = ax_plasma[1] + ax.plot( + self.rho_tor, + -self.g.raw["pprime"] * 1e-6, + c=color, + lw=2, + ls="-", + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_ylabel("pressure gradient -p' (MPa/[])") + ax.axhline(y=0.0, ls="--", lw=0.5, c="k") + + ax = ax_plasma[2] + ax.plot(self.rho_tor, self.g.raw["fpol"], c=color, lw=2, ls="-") + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_ylabel("$F = RB_{\\phi}$ (T*m)") + + ax = ax_plasma[3] + ax.plot(self.rho_tor, self.g.raw["ffprim"], c=color, lw=2, ls="-") + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_ylabel("FF' (T*m/[])") + ax.axhline(y=0.0, ls="--", lw=0.5, c="k") + + ax = ax_plasma[4] + ax.plot( + self.rho_tor, + np.abs(self.g.raw["qpsi"]), + "-s", + c=color, + lw=2, + markersize=3, + label=label + "geqdsk q", + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_ylim(bottom=0) + ax.set_ylabel("safety factor q") + ax.axhline(y=1.0, ls="--", lw=0.5, c="k") + + ax = ax_plasma[5] + ax.plot( + self.rho_tor, + np.abs(self.Jt), + "-s", + c=color, + lw=2, + markersize=3, + label=label + "geqdsk Jt", + ) + ax.plot( + self.rho_tor, + np.abs(self.Jt_fb), + "--o", + c=color, + lw=2, + markersize=3, + label=label + "geqdsk Jt(fb)", + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_ylabel("FSA toroidal current density ($MA/m^2$)") + ax.axhline(y=0.0, ls="--", lw=0.5, c="k") + + if legendYN: + ax.legend() + + #ax = ax_plasma[6] + #ax.plot( + # self.g["fluxSurfaces"]["midplane"]["R"], + # np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), + # "-s", + # c=color, + # lw=2, + # markersize=3, + # label=label + "geqdsk Bt", + #) + #ax.plot( + # self.g["fluxSurfaces"]["midplane"]["R"], + # np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), + # "--o", + # c=color, + # lw=2, + # markersize=3, + # label=label + "geqdsk Bp", + #) + #ax.set_xlabel("R (m) midplane") + #ax.set_ylabel("Midplane fields (abs())") + + if legendYN: + ax.legend() + + return ax_plasma + +def plotGeometry(self, axs=None, color="r"): + if axs is None: + plt.ion() + fig, axs = plt.subplots(ncols=4) + + ax = axs[0] + x = self.rho_tor + y = self.cx_area + ax.plot( + x, + y, + "-", + c=color, + lw=2, + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_ylim(bottom=0) + ax.set_ylabel("CX Area ($m^2$)") + + ax = axs[1] + x = self.rho_tor + y = np.zeros(x.shape) + ax.plot( + x, # self.rho_tor, + y, # self.g["fluxSurfaces"]["geo"]["surfArea"], + "-", + c=color, + lw=2, + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_ylim(bottom=0) + ax.set_ylabel("Surface Area ($m^2$)") + + ax = axs[2] + x = self.rho_tor + y = np.zeros(x.shape) + ax.plot( + x, # self.rho_tor, + y, # self.g["fluxSurfaces"]["geo"]["vol"], + "-", + c=color, + lw=2, + ) + ax.set_xlim([0, 1]) + ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") + ax.set_ylim(bottom=0) + ax.set_ylabel("Volume ($m^3$)") + +def plotFluxSurfaces( + self, + ax=None, + fluxes=[1.0], + color="b", + alpha=1.0, + rhoPol=True, + sqrt=False, + lw=1, + lwB=2, + plot1=True, + label = '', +): + x = self.g.derived["R"] + y = self.g.derived["Z"] + + if rhoPol: + z = self.g.derived["rhorz_pol"] + else: + z = self.g.derived["rhorz_tor"] + + if not sqrt: + z = z**2 + + cs, csA = plotSurfaces( + x, y, z, fluxes=fluxes, ax=ax, color=color, alpha=alpha, lw=lw, lwB=lwB, plot1=plot1, label = label + ) + + return cs, csA + +def plotXpointEnvelope( + self, ax=None, color="b", alpha=1.0, rhoPol=True, sqrt=False +): + flx = 0.001 + fluxes = [1.0 - flx, 1.0 - flx / 2, 1.0, 1.0 + flx / 2, 1.0 + flx] + + self.plotFluxSurfaces( + fluxes=fluxes, ax=ax, color=color, alpha=alpha, rhoPol=rhoPol, sqrt=sqrt + ) + +def plot2Dquantity( + self, + ax=None, + var="Jr", + zlims=None, + title="Radial Current", + cmap="seismic", + direct=False, + titlebar="J ($MA/m^2$)", + factor=1.0, + includeSurfs=True, +): + if ax is None: + fig, ax = plt.subplots() + + x = self.g.derived["R"] + y = self.g.derived["Z"] + if not direct: + z = self.g.derived[var] * factor + else: + z = var + + if zlims is None: + am = np.amax(np.abs(z[:, :])) + ming = -am + maxg = am + else: + ming = zlims[0] + maxg = zlims[1] + levels = np.linspace(ming, maxg, 100) + colticks = np.linspace(ming, maxg, 5) + + cs = ax.contourf(x, y, z, levels=levels, extend="both", cmap=cmap) + + cbar = GRAPHICStools.addColorbarSubplot( + ax, + cs, + barfmt="%3.1f", + title=titlebar, + fontsize=10, + fontsizeTitle=8, + ylabel="", + ticks=colticks, + orientation="bottom", + drawedges=False, + padCB="25%", + ) + + ax.set_aspect("equal") + ax.set_title(title) + ax.set_xlabel("R (m)") + ax.set_ylabel("Z (m)") + + if includeSurfs: + self.plotFluxSurfaces( + ax=ax, + fluxes=np.linspace(0, 1, 6), + rhoPol=False, + sqrt=True, + color="k", + lw=0.5, + ) + + return cs + + +def plotSurfaces(R, Z, F, fluxes=[1.0], ax=None, color="b", alpha=1.0, lw=1, lwB=2, plot1=True, label = ''): + if ax is None: + fig, ax = plt.subplots() + + [Rg, Yg] = np.meshgrid(R, Z) + + if plot1: + csA = ax.contour( + Rg, Yg, F, 1000, levels=[1.0], colors=color, alpha=alpha, linewidths=lwB + ) + else: + csA = None + cs = ax.contour( + Rg, Yg, F, 1000, levels=fluxes, colors=color, alpha=alpha, linewidths=lw, label = label + ) + + return cs, csA + From 8a7b9681156cfb57aeca9220cec711384d56cec1 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 8 Oct 2025 18:53:11 +0200 Subject: [PATCH 362/385] megpy as the main GEQ, leaving omfit_classes as the old --- src/mitim_tools/gs_tools/GEQtools.py | 177 ++++----- .../{GEQtools_megpy.py => GEQtools_old.py} | 177 +++++---- src/mitim_tools/gs_tools/utils/GEQplotting.py | 360 +++++++++--------- ...EQplotting_megpy.py => GEQplotting_old.py} | 360 +++++++++--------- 4 files changed, 537 insertions(+), 537 deletions(-) rename src/mitim_tools/gs_tools/{GEQtools_megpy.py => GEQtools_old.py} (89%) rename src/mitim_tools/gs_tools/utils/{GEQplotting_megpy.py => GEQplotting_old.py} (77%) diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 74370d31..a4401efa 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -6,70 +6,30 @@ import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, PLASMAtools from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.gs_tools.utils import GEQplotting_old as GEQplotting +from mitim_tools.gs_tools.utils import GEQplotting from shapely.geometry import LineString -from scipy.integrate import quad +from scipy.integrate import quad, cumulative_trapezoid +import megpy import freegs from freegs import geqdsk from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed """ -Note that this module relies on OMFIT classes (https://omfit.io/classes.html) procedures to intrepret the content of g-eqdsk files. -Modifications are made in MITINM for visualizations and a few extra derivations. +Note that this module relies on megpy to intrepret the content of g-eqdsk files. +Modifications are made in MITIM for visualizations and a few extra derivations. """ -def fix_file(filename): - - with open(filename, "r") as f: - lines = f.readlines() - - # ----------------------------------------------------------------------- - # Remove coils (chatGPT 4o as of 08/24/24) - # ----------------------------------------------------------------------- - # Use StringIO to simulate the file writing - noCoils_file = io.StringIO() - for cont, line in enumerate(lines): - if cont > 0 and line[:2] == " ": - break - noCoils_file.write(line) - - # Reset cursor to the start of StringIO - noCoils_file.seek(0) - - # Write the StringIO content to a temporary file - with tempfile.NamedTemporaryFile(delete=False) as tmp_file: - tmp_file.write(noCoils_file.getvalue().encode('utf-8')) - noCoils_file = tmp_file.name - # ----------------------------------------------------------------------- - - with open(filename, 'r') as file1, open(noCoils_file, 'r') as file2: - file1_content = file1.read() - file2_content = file2.read() - - if file1_content != file2_content: - print(f"\t- geqdsk file {IOtools.clipstr(filename)} had coils, I have removed them") - - filename = noCoils_file - - return filename - class MITIMgeqdsk: def __init__(self, filename): - # Fix file by removing coils if it has them - filename = fix_file(filename) - - # Read GEQDSK file using OMFIT - import omfit_classes.omfit_eqdsk - self.g = omfit_classes.omfit_eqdsk.OMFITgeqdsk(filename, forceFindSeparatrix=True) + self.g = megpy.Equilibrium() + self.g.read_geqdsk(f_path=filename) + self.g.add_derived(incl_fluxsurfaces=True, analytic_shape=True, incl_B=True) # Extra derivations in MITIM self.derive() - # Remove temporary file - os.remove(filename) - @classmethod def timeslices(cls, filename, **kwargs): print("\n...Opening GEQ file with several time slices") @@ -111,61 +71,65 @@ def timeslices(cls, filename, **kwargs): def derive(self, debug=False): - self.Jt = self.g.surfAvg("Jt") * 1e-6 - self.Jt_fb = self.g.surfAvg("Jt_fb") * 1e-6 + zero_vector = np.zeros(self.g.derived["rho_pol"].shape) + + self.rho_pol = self.g.derived["rho_pol"].copy() + self.rho_tor = self.g.derived["rho_tor"].copy() + self.psi_pol_norm = self.rho_pol ** 2 + self.psi_tor_norm = self.rho_tor ** 2 + + self.Jt = self.g.derived["j_tor"] * 1e-6 + self.Jt_fb = zero_vector.copy() self.Jerror = np.abs(self.Jt - self.Jt_fb) - self.Ip = self.g["CURRENT"] + self.Ip = self.g.raw["current"] # Parameterizations of LCFS - self.kappa = self.g["fluxSurfaces"]["geo"]["kap"][-1] - self.kappaU = self.g["fluxSurfaces"]["geo"]["kapu"][-1] - self.kappaL = self.g["fluxSurfaces"]["geo"]["kapl"][-1] - - self.delta = self.g["fluxSurfaces"]["geo"]["delta"][-1] - self.deltaU = self.g["fluxSurfaces"]["geo"]["dell"][-1] - self.deltaL = self.g["fluxSurfaces"]["geo"]["dell"][-1] - - self.zeta = self.g["fluxSurfaces"]["geo"]["zeta"][-1] - - self.a = self.g["fluxSurfaces"]["geo"]["a"][-1] - self.Rmag = self.g["fluxSurfaces"]["geo"]["R"][0] - self.Zmag = self.g["fluxSurfaces"]["geo"]["Z"][0] - self.Rmajor = np.mean( - [ - self.g["fluxSurfaces"]["geo"]["Rmin_centroid"][-1], - self.g["fluxSurfaces"]["geo"]["Rmax_centroid"][-1], - ] - ) + self.kappa = self.g.derived["miller_geo"]["kappa"][-1] + self.kappaU = self.g.derived["miller_geo"]["kappa_u"][-1] + self.kappaL = self.g.derived["miller_geo"]["kappa_l"][-1] + + self.delta = self.g.derived["miller_geo"]["delta"][-1] + self.deltaU = self.g.derived["miller_geo"]["delta_u"][-1] + self.deltaL = self.g.derived["miller_geo"]["delta_l"][-1] + + self.zeta = self.g.derived["miller_geo"]["zeta"][-1] + + self.a = self.g.derived["r"][-1] + self.Rmag = self.g.derived["Ro"][0] + self.Zmag = self.g.derived["Zo"][0] - self.Zmajor = self.Zmag + self.Rmajor = self.g.derived["Ro"][-1] + self.Zmajor = self.Zmag #self.g.derived["Zo"][-1] TODO: check which Z0 to use, perhaps switch to MXH definition self.eps = self.a / self.Rmajor # Core values - - self.kappa_a = self.g["fluxSurfaces"]["geo"]["cxArea"][-1] / (np.pi * self.a**2) + vp = np.array(self.g.fluxsurfaces["Vprime"]).flatten() + ir = np.array(self.g.fluxsurfaces["1/R"]).flatten() + self.cx_area = cumulative_trapezoid(vp * ir, self.g.derived["psi"], initial=0.0) + self.kappa_a = self.cx_area[-1] / (np.pi * self.a**2) self.kappa995 = np.interp( 0.995, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["kap"], + self.psi_pol_norm, + self.g.derived["miller_geo"]["kappa"], ) self.kappa95 = np.interp( 0.95, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["kap"], + self.psi_pol_norm, + self.g.derived["miller_geo"]["kappa"], ) self.delta995 = np.interp( 0.995, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["delta"], + self.psi_pol_norm, + self.g.derived["miller_geo"]["delta"], ) self.delta95 = np.interp( 0.95, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["delta"], + self.psi_pol_norm, + self.g.derived["miller_geo"]["delta"], ) """ @@ -173,26 +137,23 @@ def derive(self, debug=False): Boundary -------------------------------------------------------------------------------------------------------------------------------------- Note that the RBBS and ZBBS values in the gfile are often too scattered and do not reproduce the boundary near x-points. - The shaping parameters calculated using fluxSurfaces are correct though. + The shaping parameters calculated using fluxsurfaces are correct though. """ - self.Rb_gfile, self.Yb_gfile = self.g["RBBBS"], self.g["ZBBBS"] - self.Rb, self.Yb = self.g["fluxSurfaces"].sep.transpose() + self.Rb_gfile, self.Yb_gfile = self.g.raw["rbbbs"].copy(), self.g.raw["zbbbs"].copy() + self.Rb, self.Yb = self.g.fluxsurfaces["R"][-1], self.g.fluxsurfaces["Z"][-1] if len(self.Rb) == 0: - print("\t- MITIM > No separatrix found in the OMFIT fluxSurfaces, increasing resolution and going all in!",typeMsg='i') + print("\t- MITIM > No separatrix found in the megpy fluxsurfaces, using explicit boundary in g-eqdsk file!",typeMsg='i') - flx = copy.deepcopy(self.g['fluxSurfaces']) - flx._changeResolution(6) - flx.findSurfaces([0.0,0.5,1.0]) - fs = flx['flux'][list(flx['flux'].keys())[-1]] - self.Rb, self.Yb = fs['R'], fs['Z'] + self.Rb = self.Rb_gfile.copy() + self.Yb = self.Yb_gfile.copy() if debug: fig, ax = plt.subplots() - # OMFIT - ax.plot(self.Rb, self.Yb, "-s", c="r", label="OMFIT") + # megpy + ax.plot(self.Rb, self.Yb, "-s", c="r", label="megpy") # GFILE ax.plot(self.Rb_gfile, self.Yb_gfile, "-s", c="y", label="GFILE") @@ -245,10 +206,7 @@ def write(self, filename=None): If filename is None, use the original one """ - if filename is not None: - self.g.filename = filename - - self.g.save() + self.g.write_geqdsk(f_path=filename) # ----------------------------------------------------------------------------- # Parameterizations @@ -256,8 +214,8 @@ def write(self, filename=None): def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): - psis = self.g["AuxQuantities"]["PSI_NORM"] - flux_surfaces = self.g['fluxSurfaces']['flux'] + psis = self.psi_pol_norm + flux_surfaces = self.g.fluxsurfaces["psi"] # Cannot parallelize because different number of points? kappa, rmin, rmaj, zmag, sn, cn = [],[],[],[],[],[] @@ -266,7 +224,7 @@ def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): if flux == len(flux_surfaces)-1: Rf, Zf = self.Rb, self.Yb else: - Rf, Zf = flux_surfaces[flux]['R'],flux_surfaces[flux]['Z'] + Rf, Zf = self.g.fluxsurfaces["R"][flux], self.g.fluxsurfaces["Z"][flux] # Perform the MXH decompositionusing the MITIM surface class surfaces = mitim_flux_surfaces() @@ -311,24 +269,22 @@ def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): # ----------------------------------------------------------------------------- # For MAESTRO and TRANSP converstions # ----------------------------------------------------------------------------- - def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH = 7, plotYN = False): # ------------------------------------------------------------------------------------------------------- # Quantities from the equilibrium # ------------------------------------------------------------------------------------------------------- - rhotor = self.g['RHOVN'] - psi = self.g['AuxQuantities']['PSI'] # Wb/rad - torfluxa = self.g['AuxQuantities']['PHI'][-1] / (2*np.pi) # Wb/rad - q = self.g['QPSI'] - pressure = self.g['PRES'] # Pa - Ip = self.g['CURRENT']*1E-6 # MA + rhotor = self.g.derived['rho_tor'] + psi = self.g.derived['psi'] # Wb/rad + torfluxa = self.g.derived['phi'][-1] / (2*np.pi) # Wb/rad + q = self.g.raw['qpsi'] + pressure = self.g.raw['pres'] # Pa + Ip = self.g.raw['current']*1E-6 # MA RZ = np.array([self.Rb,self.Yb]).T R0 = (RZ.max(axis=0)[0] + RZ.min(axis=0)[0])/2 - - B0 = self.g['RCENTR']*self.g['BCENTR'] / R0 + B0 = self.g.raw['rcentr']*self.g.raw['bcentr'] / R0 # Ensure positive quantities #TODO: Check if this is necessary, pass directions rhotor = np.array([np.abs(i) for i in rhotor]) @@ -543,6 +499,7 @@ def _to_mxh(self, n_coeff=6): self.cn = np.zeros((self.R.shape[0],n_coeff)) self.sn = np.zeros((self.R.shape[0],n_coeff)) self.gn = np.zeros((self.R.shape[0],4)) + self.gn[-1] = 1.0 for i in range(self.R.shape[0]): self.cn[i,:], self.sn[i,:], self.gn[i,:] = from_RZ_to_mxh(self.R[i,:], self.Z[i,:], n_coeff=n_coeff) @@ -562,8 +519,8 @@ def _to_miller(self): # Elongations - self.kappa_u = (Zmax - self.Z0) / self.a - self.kappa_l = (self.Z0 - Zmin) / self.a + self.kappa_u = (Zmax - self.Z0) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) + self.kappa_l = (self.Z0 - Zmin) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) self.kappa = (self.kappa_u + self.kappa_l) / 2 # Triangularities diff --git a/src/mitim_tools/gs_tools/GEQtools_megpy.py b/src/mitim_tools/gs_tools/GEQtools_old.py similarity index 89% rename from src/mitim_tools/gs_tools/GEQtools_megpy.py rename to src/mitim_tools/gs_tools/GEQtools_old.py index a4401efa..74370d31 100644 --- a/src/mitim_tools/gs_tools/GEQtools_megpy.py +++ b/src/mitim_tools/gs_tools/GEQtools_old.py @@ -6,30 +6,70 @@ import matplotlib.pyplot as plt from mitim_tools.misc_tools import GRAPHICStools, IOtools, PLASMAtools from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.gs_tools.utils import GEQplotting +from mitim_tools.gs_tools.utils import GEQplotting_old as GEQplotting from shapely.geometry import LineString -from scipy.integrate import quad, cumulative_trapezoid -import megpy +from scipy.integrate import quad import freegs from freegs import geqdsk from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed """ -Note that this module relies on megpy to intrepret the content of g-eqdsk files. -Modifications are made in MITIM for visualizations and a few extra derivations. +Note that this module relies on OMFIT classes (https://omfit.io/classes.html) procedures to intrepret the content of g-eqdsk files. +Modifications are made in MITINM for visualizations and a few extra derivations. """ +def fix_file(filename): + + with open(filename, "r") as f: + lines = f.readlines() + + # ----------------------------------------------------------------------- + # Remove coils (chatGPT 4o as of 08/24/24) + # ----------------------------------------------------------------------- + # Use StringIO to simulate the file writing + noCoils_file = io.StringIO() + for cont, line in enumerate(lines): + if cont > 0 and line[:2] == " ": + break + noCoils_file.write(line) + + # Reset cursor to the start of StringIO + noCoils_file.seek(0) + + # Write the StringIO content to a temporary file + with tempfile.NamedTemporaryFile(delete=False) as tmp_file: + tmp_file.write(noCoils_file.getvalue().encode('utf-8')) + noCoils_file = tmp_file.name + # ----------------------------------------------------------------------- + + with open(filename, 'r') as file1, open(noCoils_file, 'r') as file2: + file1_content = file1.read() + file2_content = file2.read() + + if file1_content != file2_content: + print(f"\t- geqdsk file {IOtools.clipstr(filename)} had coils, I have removed them") + + filename = noCoils_file + + return filename + class MITIMgeqdsk: def __init__(self, filename): - self.g = megpy.Equilibrium() - self.g.read_geqdsk(f_path=filename) - self.g.add_derived(incl_fluxsurfaces=True, analytic_shape=True, incl_B=True) + # Fix file by removing coils if it has them + filename = fix_file(filename) + + # Read GEQDSK file using OMFIT + import omfit_classes.omfit_eqdsk + self.g = omfit_classes.omfit_eqdsk.OMFITgeqdsk(filename, forceFindSeparatrix=True) # Extra derivations in MITIM self.derive() + # Remove temporary file + os.remove(filename) + @classmethod def timeslices(cls, filename, **kwargs): print("\n...Opening GEQ file with several time slices") @@ -71,65 +111,61 @@ def timeslices(cls, filename, **kwargs): def derive(self, debug=False): - zero_vector = np.zeros(self.g.derived["rho_pol"].shape) - - self.rho_pol = self.g.derived["rho_pol"].copy() - self.rho_tor = self.g.derived["rho_tor"].copy() - self.psi_pol_norm = self.rho_pol ** 2 - self.psi_tor_norm = self.rho_tor ** 2 - - self.Jt = self.g.derived["j_tor"] * 1e-6 - self.Jt_fb = zero_vector.copy() + self.Jt = self.g.surfAvg("Jt") * 1e-6 + self.Jt_fb = self.g.surfAvg("Jt_fb") * 1e-6 self.Jerror = np.abs(self.Jt - self.Jt_fb) - self.Ip = self.g.raw["current"] + self.Ip = self.g["CURRENT"] # Parameterizations of LCFS - self.kappa = self.g.derived["miller_geo"]["kappa"][-1] - self.kappaU = self.g.derived["miller_geo"]["kappa_u"][-1] - self.kappaL = self.g.derived["miller_geo"]["kappa_l"][-1] - - self.delta = self.g.derived["miller_geo"]["delta"][-1] - self.deltaU = self.g.derived["miller_geo"]["delta_u"][-1] - self.deltaL = self.g.derived["miller_geo"]["delta_l"][-1] - - self.zeta = self.g.derived["miller_geo"]["zeta"][-1] - - self.a = self.g.derived["r"][-1] - self.Rmag = self.g.derived["Ro"][0] - self.Zmag = self.g.derived["Zo"][0] + self.kappa = self.g["fluxSurfaces"]["geo"]["kap"][-1] + self.kappaU = self.g["fluxSurfaces"]["geo"]["kapu"][-1] + self.kappaL = self.g["fluxSurfaces"]["geo"]["kapl"][-1] + + self.delta = self.g["fluxSurfaces"]["geo"]["delta"][-1] + self.deltaU = self.g["fluxSurfaces"]["geo"]["dell"][-1] + self.deltaL = self.g["fluxSurfaces"]["geo"]["dell"][-1] + + self.zeta = self.g["fluxSurfaces"]["geo"]["zeta"][-1] + + self.a = self.g["fluxSurfaces"]["geo"]["a"][-1] + self.Rmag = self.g["fluxSurfaces"]["geo"]["R"][0] + self.Zmag = self.g["fluxSurfaces"]["geo"]["Z"][0] + self.Rmajor = np.mean( + [ + self.g["fluxSurfaces"]["geo"]["Rmin_centroid"][-1], + self.g["fluxSurfaces"]["geo"]["Rmax_centroid"][-1], + ] + ) - self.Rmajor = self.g.derived["Ro"][-1] - self.Zmajor = self.Zmag #self.g.derived["Zo"][-1] TODO: check which Z0 to use, perhaps switch to MXH definition + self.Zmajor = self.Zmag self.eps = self.a / self.Rmajor # Core values - vp = np.array(self.g.fluxsurfaces["Vprime"]).flatten() - ir = np.array(self.g.fluxsurfaces["1/R"]).flatten() - self.cx_area = cumulative_trapezoid(vp * ir, self.g.derived["psi"], initial=0.0) - self.kappa_a = self.cx_area[-1] / (np.pi * self.a**2) + + self.kappa_a = self.g["fluxSurfaces"]["geo"]["cxArea"][-1] / (np.pi * self.a**2) self.kappa995 = np.interp( 0.995, - self.psi_pol_norm, - self.g.derived["miller_geo"]["kappa"], + self.g["AuxQuantities"]["PSI_NORM"], + self.g["fluxSurfaces"]["geo"]["kap"], ) self.kappa95 = np.interp( 0.95, - self.psi_pol_norm, - self.g.derived["miller_geo"]["kappa"], + self.g["AuxQuantities"]["PSI_NORM"], + self.g["fluxSurfaces"]["geo"]["kap"], ) self.delta995 = np.interp( 0.995, - self.psi_pol_norm, - self.g.derived["miller_geo"]["delta"], + self.g["AuxQuantities"]["PSI_NORM"], + self.g["fluxSurfaces"]["geo"]["delta"], ) self.delta95 = np.interp( 0.95, - self.psi_pol_norm, - self.g.derived["miller_geo"]["delta"], + self.g["AuxQuantities"]["PSI_NORM"], + self.g["fluxSurfaces"]["geo"]["delta"], ) """ @@ -137,23 +173,26 @@ def derive(self, debug=False): Boundary -------------------------------------------------------------------------------------------------------------------------------------- Note that the RBBS and ZBBS values in the gfile are often too scattered and do not reproduce the boundary near x-points. - The shaping parameters calculated using fluxsurfaces are correct though. + The shaping parameters calculated using fluxSurfaces are correct though. """ - self.Rb_gfile, self.Yb_gfile = self.g.raw["rbbbs"].copy(), self.g.raw["zbbbs"].copy() - self.Rb, self.Yb = self.g.fluxsurfaces["R"][-1], self.g.fluxsurfaces["Z"][-1] + self.Rb_gfile, self.Yb_gfile = self.g["RBBBS"], self.g["ZBBBS"] + self.Rb, self.Yb = self.g["fluxSurfaces"].sep.transpose() if len(self.Rb) == 0: - print("\t- MITIM > No separatrix found in the megpy fluxsurfaces, using explicit boundary in g-eqdsk file!",typeMsg='i') + print("\t- MITIM > No separatrix found in the OMFIT fluxSurfaces, increasing resolution and going all in!",typeMsg='i') - self.Rb = self.Rb_gfile.copy() - self.Yb = self.Yb_gfile.copy() + flx = copy.deepcopy(self.g['fluxSurfaces']) + flx._changeResolution(6) + flx.findSurfaces([0.0,0.5,1.0]) + fs = flx['flux'][list(flx['flux'].keys())[-1]] + self.Rb, self.Yb = fs['R'], fs['Z'] if debug: fig, ax = plt.subplots() - # megpy - ax.plot(self.Rb, self.Yb, "-s", c="r", label="megpy") + # OMFIT + ax.plot(self.Rb, self.Yb, "-s", c="r", label="OMFIT") # GFILE ax.plot(self.Rb_gfile, self.Yb_gfile, "-s", c="y", label="GFILE") @@ -206,7 +245,10 @@ def write(self, filename=None): If filename is None, use the original one """ - self.g.write_geqdsk(f_path=filename) + if filename is not None: + self.g.filename = filename + + self.g.save() # ----------------------------------------------------------------------------- # Parameterizations @@ -214,8 +256,8 @@ def write(self, filename=None): def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): - psis = self.psi_pol_norm - flux_surfaces = self.g.fluxsurfaces["psi"] + psis = self.g["AuxQuantities"]["PSI_NORM"] + flux_surfaces = self.g['fluxSurfaces']['flux'] # Cannot parallelize because different number of points? kappa, rmin, rmaj, zmag, sn, cn = [],[],[],[],[],[] @@ -224,7 +266,7 @@ def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): if flux == len(flux_surfaces)-1: Rf, Zf = self.Rb, self.Yb else: - Rf, Zf = self.g.fluxsurfaces["R"][flux], self.g.fluxsurfaces["Z"][flux] + Rf, Zf = flux_surfaces[flux]['R'],flux_surfaces[flux]['Z'] # Perform the MXH decompositionusing the MITIM surface class surfaces = mitim_flux_surfaces() @@ -269,22 +311,24 @@ def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): # ----------------------------------------------------------------------------- # For MAESTRO and TRANSP converstions # ----------------------------------------------------------------------------- + def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH = 7, plotYN = False): # ------------------------------------------------------------------------------------------------------- # Quantities from the equilibrium # ------------------------------------------------------------------------------------------------------- - rhotor = self.g.derived['rho_tor'] - psi = self.g.derived['psi'] # Wb/rad - torfluxa = self.g.derived['phi'][-1] / (2*np.pi) # Wb/rad - q = self.g.raw['qpsi'] - pressure = self.g.raw['pres'] # Pa - Ip = self.g.raw['current']*1E-6 # MA + rhotor = self.g['RHOVN'] + psi = self.g['AuxQuantities']['PSI'] # Wb/rad + torfluxa = self.g['AuxQuantities']['PHI'][-1] / (2*np.pi) # Wb/rad + q = self.g['QPSI'] + pressure = self.g['PRES'] # Pa + Ip = self.g['CURRENT']*1E-6 # MA RZ = np.array([self.Rb,self.Yb]).T R0 = (RZ.max(axis=0)[0] + RZ.min(axis=0)[0])/2 - B0 = self.g.raw['rcentr']*self.g.raw['bcentr'] / R0 + + B0 = self.g['RCENTR']*self.g['BCENTR'] / R0 # Ensure positive quantities #TODO: Check if this is necessary, pass directions rhotor = np.array([np.abs(i) for i in rhotor]) @@ -499,7 +543,6 @@ def _to_mxh(self, n_coeff=6): self.cn = np.zeros((self.R.shape[0],n_coeff)) self.sn = np.zeros((self.R.shape[0],n_coeff)) self.gn = np.zeros((self.R.shape[0],4)) - self.gn[-1] = 1.0 for i in range(self.R.shape[0]): self.cn[i,:], self.sn[i,:], self.gn[i,:] = from_RZ_to_mxh(self.R[i,:], self.Z[i,:], n_coeff=n_coeff) @@ -519,8 +562,8 @@ def _to_miller(self): # Elongations - self.kappa_u = (Zmax - self.Z0) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) - self.kappa_l = (self.Z0 - Zmin) / self.a if self.Z.shape[1] > 1 else np.ones(self.Z0.shape) + self.kappa_u = (Zmax - self.Z0) / self.a + self.kappa_l = (self.Z0 - Zmin) / self.a self.kappa = (self.kappa_u + self.kappa_l) / 2 # Triangularities diff --git a/src/mitim_tools/gs_tools/utils/GEQplotting.py b/src/mitim_tools/gs_tools/utils/GEQplotting.py index fe6531b4..db545acd 100644 --- a/src/mitim_tools/gs_tools/utils/GEQplotting.py +++ b/src/mitim_tools/gs_tools/utils/GEQplotting.py @@ -4,6 +4,7 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed +#TODO: add current profiles and flux-surface average fields to megpy and restore plots def compareGeqdsk(geqdsks, fn=None, extraLabel="", plotAll=True, labelsGs=None): @@ -229,8 +230,8 @@ def plotFS(self, axs=None, color="b", label=""): ax.set_ylabel("Z (m)") ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["AuxQuantities"]["RHO"] + x = self.psi_pol_norm + y = self.rho_tor ax.plot(x, y, lw=2, ls="-", c=color, label=label) ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) @@ -240,8 +241,8 @@ def plotFS(self, axs=None, color="b", label=""): ax.set_ylim([0, 1]) ax = axs[3] - x = self.g["AuxQuantities"]["RHO"] - y = self.g["AuxQuantities"]["RHOp"] + x = self.rho_tor + y = self.rho_pol ax.plot(x, y, lw=2, ls="-", c=color) ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) @@ -256,8 +257,8 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): fig, axs = plt.subplots(ncols=10) ax = axs[0] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jr") * 1e-6 + x = self.psi_pol_norm + y = np.zeros(x.shape) ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_ylabel("FSA $\\langle J\\rangle$ ($MA/m^2$)") @@ -267,14 +268,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[1] - plot2Dquantity(self, - ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 - ) + #ax = axs[1] + #plot2Dquantity(self, + # ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 + #) ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jz") * 1e-6 + x = self.psi_pol_norm + y = np.zeros(x.shape) ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -282,14 +283,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[3] - plot2Dquantity(self, - ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 - ) + #ax = axs[3] + #plot2Dquantity(self, + # ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 + #) ax = axs[4] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jt") * 1e-6 + x = self.psi_pol_norm + y = self.Jt ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -297,14 +298,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[5] - plot2Dquantity(self, - ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 - ) + #ax = axs[5] + #plot2Dquantity(self, + # ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 + #) ax = axs[6] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jp") * 1e-6 + x = self.psi_pol_norm + y = np.zeros(x.shape) ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -312,14 +313,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[7] - plot2Dquantity(self, - ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 - ) + #ax = axs[7] + #plot2Dquantity(self, + # ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 + #) ax = axs[8] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jpar") * 1e-6 + x = self.psi_pol_norm + y = np.zeros(x.shape) ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -327,10 +328,10 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[9] - plot2Dquantity(self, - ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 - ) + #ax = axs[9] + #plot2Dquantity(self, + # ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 + #) def plotFields(self, axs=None, zlims_thr=[-1, 1]): if axs is None: @@ -338,8 +339,8 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): fig, axs = plt.subplots(ncols=10) ax = axs[0] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Br") + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Br") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_ylabel("FSA $\\langle B\\rangle$ ($T$)") @@ -351,12 +352,12 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[1] plot2Dquantity(self, - ax=ax, var="Br", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="B_r", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" ) ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Bz") + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Bz") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -366,12 +367,12 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[3] plot2Dquantity(self, - ax=ax, var="Bz", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="B_z", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" ) ax = axs[4] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Bt") + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Bt") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -380,14 +381,14 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): # zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - ax = axs[5] - plot2Dquantity(self, - ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" - ) + #ax = axs[5] + #plot2Dquantity(self, + # ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" + #) ax = axs[6] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Bp") + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g.surfAvg("Bp") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -397,32 +398,32 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[7] plot2Dquantity(self, - ax=ax, var="Bp", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="B_pol_rz", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" ) ax = axs[8] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["avg"]["Bp**2"] + x = self.psi_pol_norm + y = np.zeros(x.shape) # self.g["fluxSurfaces"]["avg"]["Bp**2"] ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) ax.set_ylabel("$\\langle B_{\\theta}^2\\rangle$") - ax = axs[9] - x = self.g["fluxSurfaces"]["midplane"]["R"] - y = self.g["fluxSurfaces"]["midplane"]["Bt"] - ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") - y = self.g["fluxSurfaces"]["midplane"]["Bp"] - ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") - y = self.g["fluxSurfaces"]["midplane"]["Bz"] - ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") - y = self.g["fluxSurfaces"]["midplane"]["Br"] - ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") - y = self.g["fluxSurfaces"]["geo"]["bunit"] - ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") - ax.set_xlabel("$R$ LF midplane") - ax.set_ylabel("$B$ (T)") - ax.legend() + #ax = axs[9] + #x = self.g["fluxSurfaces"]["midplane"]["R"] + #y = self.g["fluxSurfaces"]["midplane"]["Bt"] + #ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") + #y = self.g["fluxSurfaces"]["midplane"]["Bp"] + #ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") + #y = self.g["fluxSurfaces"]["midplane"]["Bz"] + #ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") + #y = self.g["fluxSurfaces"]["midplane"]["Br"] + #ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") + #y = self.g["fluxSurfaces"]["geo"]["bunit"] + #ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") + #ax.set_xlabel("$R$ LF midplane") + #ax.set_ylabel("$B$ (T)") + #ax.legend() def plotChecks(self, axs=None): if axs is None: @@ -430,7 +431,7 @@ def plotChecks(self, axs=None): fig, axs = plt.subplots(ncols=8) ax = axs[0] - x = self.g["AuxQuantities"]["PSI_NORM"] + x = self.psi_pol_norm y1 = self.Jt ax.plot(x, np.abs(y1), lw=2, ls="-", c="b", label="$\\langle Jt\\rangle$") zmax = y1.max() @@ -463,10 +464,10 @@ def plotChecks(self, axs=None): ax.legend() ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y1 = self.g["FFPRIM"] + x = self.psi_pol_norm + y1 = self.g.raw["ffprim"] ax.plot(x, y1, lw=2, ls="-", c="r", label="$FF'$") - y2 = self.g["PPRIME"] * (4 * np.pi * 1e-7) + y2 = self.g.raw["pprime"] * (4 * np.pi * 1e-7) ax.plot(x, y2, lw=2, ls="-", c="b", label="$p'*\\mu_0$") ax.set_ylabel("") @@ -474,40 +475,40 @@ def plotChecks(self, axs=None): ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) - ax = axs[3] - plot2Dquantity(self, - ax=ax, - var="Jt", - title="Toroidal Current Jt", - zlims=[zmin, zmax], - cmap="viridis", - factor=1e-6, - ) - - ax = axs[4] - plot2Dquantity(self, - ax=ax, - var="Jt_fb", - title="Toroidal Current Jt (FB)", - zlims=[zmin, zmax], - cmap="viridis", - factor=1e-6, - ) - - ax = axs[5] - z = ( - np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) - * 1e-6 - ) - zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) - plot2Dquantity(self, - ax=ax, - var=z, - title="Absolute Error", - zlims=[0, zmaxx], - cmap="viridis", - direct=True, - ) + #ax = axs[3] + #plot2Dquantity(self, + # ax=ax, + # var="Jt", + # title="Toroidal Current Jt", + # zlims=[zmin, zmax], + # cmap="viridis", + # factor=1e-6, + #) + + #ax = axs[4] + #plot2Dquantity(self, + # ax=ax, + # var="Jt_fb", + # title="Toroidal Current Jt (FB)", + # zlims=[zmin, zmax], + # cmap="viridis", + # factor=1e-6, + #) + + #ax = axs[5] + #z = ( + # np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) + # * 1e-6 + #) + #zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) + #plot2Dquantity(self, + # ax=ax, + # var=z, + # title="Absolute Error", + # zlims=[0, zmaxx], + # cmap="viridis", + # direct=True, + #) def plotParameterization(self, axs=None): if axs is None: @@ -520,22 +521,23 @@ def plotParameterization(self, axs=None): ) # Boundary, axis and limiter ax.plot(self.Rb, self.Yb, lw=1, c="r") - ax.plot(self.g["RMAXIS"], self.g["ZMAXIS"], "+", markersize=10, c="r") + ax.plot(self.g.raw["rmaxis"], self.g.raw["zmaxis"], "+", markersize=10, c="r") ax.plot([self.Rmag], [self.Zmag], "o", markersize=5, c="m") ax.plot([self.Rmajor], [self.Zmag], "+", markersize=10, c="k") - ax.plot(self.g["RLIM"], self.g["ZLIM"], lw=1, c="k") + if 'rlim' in self.g.raw and 'zlim' in self.g.raw: + ax.plot(self.g.raw["rlim"], self.g.raw["zlim"], lw=1, c="k") - import matplotlib + import matplotlib - path = matplotlib.path.Path( - np.transpose(np.array([self.g["RLIM"], self.g["ZLIM"]])) - ) - patch = matplotlib.patches.PathPatch(path, facecolor="none") - ax.add_patch(patch) - # for col in cs.collections: - # col.set_clip_path(patch) - # for col in csA.collections: - # col.set_clip_path(patch) + path = matplotlib.path.Path( + np.transpose(np.array([self.g.raw["rlim"], self.g.raw["zlim"]])) + ) + patch = matplotlib.patches.PathPatch(path, facecolor="none") + ax.add_patch(patch) + # for col in cs.collections: + # col.set_clip_path(patch) + # for col in csA.collections: + # col.set_clip_path(patch) self.plotEnclosingBox(ax=ax) @@ -545,12 +547,12 @@ def plotParameterization(self, axs=None): ax.set_ylabel("Z (m)") ax = axs[1] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["geo"]["kap"] + x = self.psi_pol_norm + y = self.g.derived["miller_geo"]["kappa"].copy() ax.plot(x, y, label="$\\kappa$") - y = self.g["fluxSurfaces"]["geo"]["kapl"] + y = self.g.derived["miller_geo"]["kappa_l"].copy() ax.plot(x, y, ls="--", label="$\\kappa_L$") - y = self.g["fluxSurfaces"]["geo"]["kapu"] + y = self.g.derived["miller_geo"]["kappa_u"].copy() ax.plot(x, y, ls="--", label="$\\kappa_U$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -558,12 +560,12 @@ def plotParameterization(self, axs=None): ax.legend() ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["geo"]["delta"] + x = self.psi_pol_norm + y = self.g.derived["miller_geo"]["delta"].copy() ax.plot(x, y, label="$\\delta$") - y = self.g["fluxSurfaces"]["geo"]["dell"] + y = self.g.derived["miller_geo"]["delta_l"].copy() ax.plot(x, y, ls="--", label="$\\delta_L$") - y = self.g["fluxSurfaces"]["geo"]["delu"] + y = self.g.derived["miller_geo"]["delta_u"].copy() ax.plot(x, y, ls="--", label="$\\delta_U$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -571,16 +573,16 @@ def plotParameterization(self, axs=None): ax.legend() ax = axs[3] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["geo"]["zeta"] + x = self.psi_pol_norm + y = self.g.derived["miller_geo"]["zeta"].copy() ax.plot(x, y, label="$\\zeta$") - y = self.g["fluxSurfaces"]["geo"]["zetail"] + y = self.g.derived["miller_geo"]["zeta_li"].copy() ax.plot(x, y, ls="--", label="$\\zeta_{IL}$") - y = self.g["fluxSurfaces"]["geo"]["zetaiu"] + y = self.g.derived["miller_geo"]["zeta_ui"].copy() ax.plot(x, y, ls="--", label="$\\zeta_{IU}$") - y = self.g["fluxSurfaces"]["geo"]["zetaol"] + y = self.g.derived["miller_geo"]["zeta_lo"].copy() ax.plot(x, y, ls="--", label="$\\zeta_{OL}$") - y = self.g["fluxSurfaces"]["geo"]["zetaou"] + y = self.g.derived["miller_geo"]["zeta_uo"].copy() ax.plot(x, y, ls="--", label="$\\zeta_{OU}$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -697,8 +699,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[0] ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["PRES"] * 1e-6, + self.rho_tor, + self.g.raw["pres"] * 1e-6, "-s", c=color, lw=2, @@ -712,8 +714,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[1] ax.plot( - self.g["AuxQuantities"]["RHO"], - -self.g["PPRIME"] * 1e-6, + self.rho_tor, + -self.g.raw["pprime"] * 1e-6, c=color, lw=2, ls="-", @@ -724,13 +726,13 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax.axhline(y=0.0, ls="--", lw=0.5, c="k") ax = ax_plasma[2] - ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FPOL"], c=color, lw=2, ls="-") + ax.plot(self.rho_tor, self.g.raw["fpol"], c=color, lw=2, ls="-") ax.set_xlim([0, 1]) ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") ax.set_ylabel("$F = RB_{\\phi}$ (T*m)") ax = ax_plasma[3] - ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FFPRIM"], c=color, lw=2, ls="-") + ax.plot(self.rho_tor, self.g.raw["ffprim"], c=color, lw=2, ls="-") ax.set_xlim([0, 1]) ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") ax.set_ylabel("FF' (T*m/[])") @@ -738,8 +740,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[4] ax.plot( - self.g["AuxQuantities"]["RHO"], - np.abs(self.g["QPSI"]), + self.rho_tor, + np.abs(self.g.raw["qpsi"]), "-s", c=color, lw=2, @@ -754,8 +756,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[5] ax.plot( - self.g["AuxQuantities"]["RHO"], - np.abs(self.g.surfAvg("Jt") * 1e-6), + self.rho_tor, + np.abs(self.Jt), "-s", c=color, lw=2, @@ -763,8 +765,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): label=label + "geqdsk Jt", ) ax.plot( - self.g["AuxQuantities"]["RHO"], - np.abs(self.g.surfAvg("Jt_fb") * 1e-6), + self.rho_tor, + np.abs(self.Jt_fb), "--o", c=color, lw=2, @@ -779,27 +781,27 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): if legendYN: ax.legend() - ax = ax_plasma[6] - ax.plot( - self.g["fluxSurfaces"]["midplane"]["R"], - np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), - "-s", - c=color, - lw=2, - markersize=3, - label=label + "geqdsk Bt", - ) - ax.plot( - self.g["fluxSurfaces"]["midplane"]["R"], - np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), - "--o", - c=color, - lw=2, - markersize=3, - label=label + "geqdsk Bp", - ) - ax.set_xlabel("R (m) midplane") - ax.set_ylabel("Midplane fields (abs())") + #ax = ax_plasma[6] + #ax.plot( + # self.g["fluxSurfaces"]["midplane"]["R"], + # np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), + # "-s", + # c=color, + # lw=2, + # markersize=3, + # label=label + "geqdsk Bt", + #) + #ax.plot( + # self.g["fluxSurfaces"]["midplane"]["R"], + # np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), + # "--o", + # c=color, + # lw=2, + # markersize=3, + # label=label + "geqdsk Bp", + #) + #ax.set_xlabel("R (m) midplane") + #ax.set_ylabel("Midplane fields (abs())") if legendYN: ax.legend() @@ -812,9 +814,11 @@ def plotGeometry(self, axs=None, color="r"): fig, axs = plt.subplots(ncols=4) ax = axs[0] + x = self.rho_tor + y = self.cx_area ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["fluxSurfaces"]["geo"]["cxArea"], + x, + y, "-", c=color, lw=2, @@ -825,9 +829,11 @@ def plotGeometry(self, axs=None, color="r"): ax.set_ylabel("CX Area ($m^2$)") ax = axs[1] + x = self.rho_tor + y = np.zeros(x.shape) ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["fluxSurfaces"]["geo"]["surfArea"], + x, # self.rho_tor, + y, # self.g["fluxSurfaces"]["geo"]["surfArea"], "-", c=color, lw=2, @@ -838,9 +844,11 @@ def plotGeometry(self, axs=None, color="r"): ax.set_ylabel("Surface Area ($m^2$)") ax = axs[2] + x = self.rho_tor + y = np.zeros(x.shape) ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["fluxSurfaces"]["geo"]["vol"], + x, # self.rho_tor, + y, # self.g["fluxSurfaces"]["geo"]["vol"], "-", c=color, lw=2, @@ -863,13 +871,13 @@ def plotFluxSurfaces( plot1=True, label = '', ): - x = self.g["AuxQuantities"]["R"] - y = self.g["AuxQuantities"]["Z"] + x = self.g.derived["R"] + y = self.g.derived["Z"] if rhoPol: - z = self.g["AuxQuantities"]["RHOpRZ"] + z = self.g.derived["rhorz_pol"] else: - z = self.g["AuxQuantities"]["RHORZ"] + z = self.g.derived["rhorz_tor"] if not sqrt: z = z**2 @@ -905,10 +913,10 @@ def plot2Dquantity( if ax is None: fig, ax = plt.subplots() - x = self.g["AuxQuantities"]["R"] - y = self.g["AuxQuantities"]["Z"] + x = self.g.derived["R"] + y = self.g.derived["Z"] if not direct: - z = self.g["AuxQuantities"][var] * factor + z = self.g.derived[var] * factor else: z = var diff --git a/src/mitim_tools/gs_tools/utils/GEQplotting_megpy.py b/src/mitim_tools/gs_tools/utils/GEQplotting_old.py similarity index 77% rename from src/mitim_tools/gs_tools/utils/GEQplotting_megpy.py rename to src/mitim_tools/gs_tools/utils/GEQplotting_old.py index db545acd..fe6531b4 100644 --- a/src/mitim_tools/gs_tools/utils/GEQplotting_megpy.py +++ b/src/mitim_tools/gs_tools/utils/GEQplotting_old.py @@ -4,7 +4,6 @@ from mitim_tools.misc_tools.LOGtools import printMsg as print from IPython import embed -#TODO: add current profiles and flux-surface average fields to megpy and restore plots def compareGeqdsk(geqdsks, fn=None, extraLabel="", plotAll=True, labelsGs=None): @@ -230,8 +229,8 @@ def plotFS(self, axs=None, color="b", label=""): ax.set_ylabel("Z (m)") ax = axs[2] - x = self.psi_pol_norm - y = self.rho_tor + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["AuxQuantities"]["RHO"] ax.plot(x, y, lw=2, ls="-", c=color, label=label) ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) @@ -241,8 +240,8 @@ def plotFS(self, axs=None, color="b", label=""): ax.set_ylim([0, 1]) ax = axs[3] - x = self.rho_tor - y = self.rho_pol + x = self.g["AuxQuantities"]["RHO"] + y = self.g["AuxQuantities"]["RHOp"] ax.plot(x, y, lw=2, ls="-", c=color) ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) @@ -257,8 +256,8 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): fig, axs = plt.subplots(ncols=10) ax = axs[0] - x = self.psi_pol_norm - y = np.zeros(x.shape) + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jr") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_ylabel("FSA $\\langle J\\rangle$ ($MA/m^2$)") @@ -268,14 +267,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[1] - #plot2Dquantity(self, - # ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 - #) + ax = axs[1] + plot2Dquantity(self, + ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 + ) ax = axs[2] - x = self.psi_pol_norm - y = np.zeros(x.shape) + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jz") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -283,14 +282,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[3] - #plot2Dquantity(self, - # ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 - #) + ax = axs[3] + plot2Dquantity(self, + ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 + ) ax = axs[4] - x = self.psi_pol_norm - y = self.Jt + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jt") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -298,14 +297,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[5] - #plot2Dquantity(self, - # ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 - #) + ax = axs[5] + plot2Dquantity(self, + ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 + ) ax = axs[6] - x = self.psi_pol_norm - y = np.zeros(x.shape) + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jp") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -313,14 +312,14 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[7] - #plot2Dquantity(self, - # ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 - #) + ax = axs[7] + plot2Dquantity(self, + ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 + ) ax = axs[8] - x = self.psi_pol_norm - y = np.zeros(x.shape) + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Jpar") * 1e-6 ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -328,10 +327,10 @@ def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[9] - #plot2Dquantity(self, - # ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 - #) + ax = axs[9] + plot2Dquantity(self, + ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 + ) def plotFields(self, axs=None, zlims_thr=[-1, 1]): if axs is None: @@ -339,8 +338,8 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): fig, axs = plt.subplots(ncols=10) ax = axs[0] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g.surfAvg("Br") + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Br") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_ylabel("FSA $\\langle B\\rangle$ ($T$)") @@ -352,12 +351,12 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[1] plot2Dquantity(self, - ax=ax, var="B_r", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="Br", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" ) ax = axs[2] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g.surfAvg("Bz") + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Bz") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -367,12 +366,12 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[3] plot2Dquantity(self, - ax=ax, var="B_z", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="Bz", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" ) ax = axs[4] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g.surfAvg("Bt") + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Bt") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -381,14 +380,14 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): # zlims = GRAPHICStools.aroundZeroLims(zlims) ax.set_ylim(zlims) - #ax = axs[5] - #plot2Dquantity(self, - # ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" - #) + ax = axs[5] + plot2Dquantity(self, + ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" + ) ax = axs[6] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g.surfAvg("Bp") + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g.surfAvg("Bp") ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -398,32 +397,32 @@ def plotFields(self, axs=None, zlims_thr=[-1, 1]): ax = axs[7] plot2Dquantity(self, - ax=ax, var="B_pol_rz", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" + ax=ax, var="Bp", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" ) ax = axs[8] - x = self.psi_pol_norm - y = np.zeros(x.shape) # self.g["fluxSurfaces"]["avg"]["Bp**2"] + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["fluxSurfaces"]["avg"]["Bp**2"] ax.plot(x, y, lw=2, ls="-", c="r") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) ax.set_ylabel("$\\langle B_{\\theta}^2\\rangle$") - #ax = axs[9] - #x = self.g["fluxSurfaces"]["midplane"]["R"] - #y = self.g["fluxSurfaces"]["midplane"]["Bt"] - #ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") - #y = self.g["fluxSurfaces"]["midplane"]["Bp"] - #ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") - #y = self.g["fluxSurfaces"]["midplane"]["Bz"] - #ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") - #y = self.g["fluxSurfaces"]["midplane"]["Br"] - #ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") - #y = self.g["fluxSurfaces"]["geo"]["bunit"] - #ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") - #ax.set_xlabel("$R$ LF midplane") - #ax.set_ylabel("$B$ (T)") - #ax.legend() + ax = axs[9] + x = self.g["fluxSurfaces"]["midplane"]["R"] + y = self.g["fluxSurfaces"]["midplane"]["Bt"] + ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") + y = self.g["fluxSurfaces"]["midplane"]["Bp"] + ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") + y = self.g["fluxSurfaces"]["midplane"]["Bz"] + ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") + y = self.g["fluxSurfaces"]["midplane"]["Br"] + ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") + y = self.g["fluxSurfaces"]["geo"]["bunit"] + ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") + ax.set_xlabel("$R$ LF midplane") + ax.set_ylabel("$B$ (T)") + ax.legend() def plotChecks(self, axs=None): if axs is None: @@ -431,7 +430,7 @@ def plotChecks(self, axs=None): fig, axs = plt.subplots(ncols=8) ax = axs[0] - x = self.psi_pol_norm + x = self.g["AuxQuantities"]["PSI_NORM"] y1 = self.Jt ax.plot(x, np.abs(y1), lw=2, ls="-", c="b", label="$\\langle Jt\\rangle$") zmax = y1.max() @@ -464,10 +463,10 @@ def plotChecks(self, axs=None): ax.legend() ax = axs[2] - x = self.psi_pol_norm - y1 = self.g.raw["ffprim"] + x = self.g["AuxQuantities"]["PSI_NORM"] + y1 = self.g["FFPRIM"] ax.plot(x, y1, lw=2, ls="-", c="r", label="$FF'$") - y2 = self.g.raw["pprime"] * (4 * np.pi * 1e-7) + y2 = self.g["PPRIME"] * (4 * np.pi * 1e-7) ax.plot(x, y2, lw=2, ls="-", c="b", label="$p'*\\mu_0$") ax.set_ylabel("") @@ -475,40 +474,40 @@ def plotChecks(self, axs=None): ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) - #ax = axs[3] - #plot2Dquantity(self, - # ax=ax, - # var="Jt", - # title="Toroidal Current Jt", - # zlims=[zmin, zmax], - # cmap="viridis", - # factor=1e-6, - #) - - #ax = axs[4] - #plot2Dquantity(self, - # ax=ax, - # var="Jt_fb", - # title="Toroidal Current Jt (FB)", - # zlims=[zmin, zmax], - # cmap="viridis", - # factor=1e-6, - #) - - #ax = axs[5] - #z = ( - # np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) - # * 1e-6 - #) - #zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) - #plot2Dquantity(self, - # ax=ax, - # var=z, - # title="Absolute Error", - # zlims=[0, zmaxx], - # cmap="viridis", - # direct=True, - #) + ax = axs[3] + plot2Dquantity(self, + ax=ax, + var="Jt", + title="Toroidal Current Jt", + zlims=[zmin, zmax], + cmap="viridis", + factor=1e-6, + ) + + ax = axs[4] + plot2Dquantity(self, + ax=ax, + var="Jt_fb", + title="Toroidal Current Jt (FB)", + zlims=[zmin, zmax], + cmap="viridis", + factor=1e-6, + ) + + ax = axs[5] + z = ( + np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) + * 1e-6 + ) + zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) + plot2Dquantity(self, + ax=ax, + var=z, + title="Absolute Error", + zlims=[0, zmaxx], + cmap="viridis", + direct=True, + ) def plotParameterization(self, axs=None): if axs is None: @@ -521,23 +520,22 @@ def plotParameterization(self, axs=None): ) # Boundary, axis and limiter ax.plot(self.Rb, self.Yb, lw=1, c="r") - ax.plot(self.g.raw["rmaxis"], self.g.raw["zmaxis"], "+", markersize=10, c="r") + ax.plot(self.g["RMAXIS"], self.g["ZMAXIS"], "+", markersize=10, c="r") ax.plot([self.Rmag], [self.Zmag], "o", markersize=5, c="m") ax.plot([self.Rmajor], [self.Zmag], "+", markersize=10, c="k") - if 'rlim' in self.g.raw and 'zlim' in self.g.raw: - ax.plot(self.g.raw["rlim"], self.g.raw["zlim"], lw=1, c="k") + ax.plot(self.g["RLIM"], self.g["ZLIM"], lw=1, c="k") - import matplotlib + import matplotlib - path = matplotlib.path.Path( - np.transpose(np.array([self.g.raw["rlim"], self.g.raw["zlim"]])) - ) - patch = matplotlib.patches.PathPatch(path, facecolor="none") - ax.add_patch(patch) - # for col in cs.collections: - # col.set_clip_path(patch) - # for col in csA.collections: - # col.set_clip_path(patch) + path = matplotlib.path.Path( + np.transpose(np.array([self.g["RLIM"], self.g["ZLIM"]])) + ) + patch = matplotlib.patches.PathPatch(path, facecolor="none") + ax.add_patch(patch) + # for col in cs.collections: + # col.set_clip_path(patch) + # for col in csA.collections: + # col.set_clip_path(patch) self.plotEnclosingBox(ax=ax) @@ -547,12 +545,12 @@ def plotParameterization(self, axs=None): ax.set_ylabel("Z (m)") ax = axs[1] - x = self.psi_pol_norm - y = self.g.derived["miller_geo"]["kappa"].copy() + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["fluxSurfaces"]["geo"]["kap"] ax.plot(x, y, label="$\\kappa$") - y = self.g.derived["miller_geo"]["kappa_l"].copy() + y = self.g["fluxSurfaces"]["geo"]["kapl"] ax.plot(x, y, ls="--", label="$\\kappa_L$") - y = self.g.derived["miller_geo"]["kappa_u"].copy() + y = self.g["fluxSurfaces"]["geo"]["kapu"] ax.plot(x, y, ls="--", label="$\\kappa_U$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -560,12 +558,12 @@ def plotParameterization(self, axs=None): ax.legend() ax = axs[2] - x = self.psi_pol_norm - y = self.g.derived["miller_geo"]["delta"].copy() + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["fluxSurfaces"]["geo"]["delta"] ax.plot(x, y, label="$\\delta$") - y = self.g.derived["miller_geo"]["delta_l"].copy() + y = self.g["fluxSurfaces"]["geo"]["dell"] ax.plot(x, y, ls="--", label="$\\delta_L$") - y = self.g.derived["miller_geo"]["delta_u"].copy() + y = self.g["fluxSurfaces"]["geo"]["delu"] ax.plot(x, y, ls="--", label="$\\delta_U$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -573,16 +571,16 @@ def plotParameterization(self, axs=None): ax.legend() ax = axs[3] - x = self.psi_pol_norm - y = self.g.derived["miller_geo"]["zeta"].copy() + x = self.g["AuxQuantities"]["PSI_NORM"] + y = self.g["fluxSurfaces"]["geo"]["zeta"] ax.plot(x, y, label="$\\zeta$") - y = self.g.derived["miller_geo"]["zeta_li"].copy() + y = self.g["fluxSurfaces"]["geo"]["zetail"] ax.plot(x, y, ls="--", label="$\\zeta_{IL}$") - y = self.g.derived["miller_geo"]["zeta_ui"].copy() + y = self.g["fluxSurfaces"]["geo"]["zetaiu"] ax.plot(x, y, ls="--", label="$\\zeta_{IU}$") - y = self.g.derived["miller_geo"]["zeta_lo"].copy() + y = self.g["fluxSurfaces"]["geo"]["zetaol"] ax.plot(x, y, ls="--", label="$\\zeta_{OL}$") - y = self.g.derived["miller_geo"]["zeta_uo"].copy() + y = self.g["fluxSurfaces"]["geo"]["zetaou"] ax.plot(x, y, ls="--", label="$\\zeta_{OU}$") ax.set_xlabel("$\\Psi_n$") ax.set_xlim([0, 1]) @@ -699,8 +697,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[0] ax.plot( - self.rho_tor, - self.g.raw["pres"] * 1e-6, + self.g["AuxQuantities"]["RHO"], + self.g["PRES"] * 1e-6, "-s", c=color, lw=2, @@ -714,8 +712,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[1] ax.plot( - self.rho_tor, - -self.g.raw["pprime"] * 1e-6, + self.g["AuxQuantities"]["RHO"], + -self.g["PPRIME"] * 1e-6, c=color, lw=2, ls="-", @@ -726,13 +724,13 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax.axhline(y=0.0, ls="--", lw=0.5, c="k") ax = ax_plasma[2] - ax.plot(self.rho_tor, self.g.raw["fpol"], c=color, lw=2, ls="-") + ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FPOL"], c=color, lw=2, ls="-") ax.set_xlim([0, 1]) ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") ax.set_ylabel("$F = RB_{\\phi}$ (T*m)") ax = ax_plasma[3] - ax.plot(self.rho_tor, self.g.raw["ffprim"], c=color, lw=2, ls="-") + ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FFPRIM"], c=color, lw=2, ls="-") ax.set_xlim([0, 1]) ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") ax.set_ylabel("FF' (T*m/[])") @@ -740,8 +738,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[4] ax.plot( - self.rho_tor, - np.abs(self.g.raw["qpsi"]), + self.g["AuxQuantities"]["RHO"], + np.abs(self.g["QPSI"]), "-s", c=color, lw=2, @@ -756,8 +754,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): ax = ax_plasma[5] ax.plot( - self.rho_tor, - np.abs(self.Jt), + self.g["AuxQuantities"]["RHO"], + np.abs(self.g.surfAvg("Jt") * 1e-6), "-s", c=color, lw=2, @@ -765,8 +763,8 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): label=label + "geqdsk Jt", ) ax.plot( - self.rho_tor, - np.abs(self.Jt_fb), + self.g["AuxQuantities"]["RHO"], + np.abs(self.g.surfAvg("Jt_fb") * 1e-6), "--o", c=color, lw=2, @@ -781,27 +779,27 @@ def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): if legendYN: ax.legend() - #ax = ax_plasma[6] - #ax.plot( - # self.g["fluxSurfaces"]["midplane"]["R"], - # np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), - # "-s", - # c=color, - # lw=2, - # markersize=3, - # label=label + "geqdsk Bt", - #) - #ax.plot( - # self.g["fluxSurfaces"]["midplane"]["R"], - # np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), - # "--o", - # c=color, - # lw=2, - # markersize=3, - # label=label + "geqdsk Bp", - #) - #ax.set_xlabel("R (m) midplane") - #ax.set_ylabel("Midplane fields (abs())") + ax = ax_plasma[6] + ax.plot( + self.g["fluxSurfaces"]["midplane"]["R"], + np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), + "-s", + c=color, + lw=2, + markersize=3, + label=label + "geqdsk Bt", + ) + ax.plot( + self.g["fluxSurfaces"]["midplane"]["R"], + np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), + "--o", + c=color, + lw=2, + markersize=3, + label=label + "geqdsk Bp", + ) + ax.set_xlabel("R (m) midplane") + ax.set_ylabel("Midplane fields (abs())") if legendYN: ax.legend() @@ -814,11 +812,9 @@ def plotGeometry(self, axs=None, color="r"): fig, axs = plt.subplots(ncols=4) ax = axs[0] - x = self.rho_tor - y = self.cx_area ax.plot( - x, - y, + self.g["AuxQuantities"]["RHO"], + self.g["fluxSurfaces"]["geo"]["cxArea"], "-", c=color, lw=2, @@ -829,11 +825,9 @@ def plotGeometry(self, axs=None, color="r"): ax.set_ylabel("CX Area ($m^2$)") ax = axs[1] - x = self.rho_tor - y = np.zeros(x.shape) ax.plot( - x, # self.rho_tor, - y, # self.g["fluxSurfaces"]["geo"]["surfArea"], + self.g["AuxQuantities"]["RHO"], + self.g["fluxSurfaces"]["geo"]["surfArea"], "-", c=color, lw=2, @@ -844,11 +838,9 @@ def plotGeometry(self, axs=None, color="r"): ax.set_ylabel("Surface Area ($m^2$)") ax = axs[2] - x = self.rho_tor - y = np.zeros(x.shape) ax.plot( - x, # self.rho_tor, - y, # self.g["fluxSurfaces"]["geo"]["vol"], + self.g["AuxQuantities"]["RHO"], + self.g["fluxSurfaces"]["geo"]["vol"], "-", c=color, lw=2, @@ -871,13 +863,13 @@ def plotFluxSurfaces( plot1=True, label = '', ): - x = self.g.derived["R"] - y = self.g.derived["Z"] + x = self.g["AuxQuantities"]["R"] + y = self.g["AuxQuantities"]["Z"] if rhoPol: - z = self.g.derived["rhorz_pol"] + z = self.g["AuxQuantities"]["RHOpRZ"] else: - z = self.g.derived["rhorz_tor"] + z = self.g["AuxQuantities"]["RHORZ"] if not sqrt: z = z**2 @@ -913,10 +905,10 @@ def plot2Dquantity( if ax is None: fig, ax = plt.subplots() - x = self.g.derived["R"] - y = self.g.derived["Z"] + x = self.g["AuxQuantities"]["R"] + y = self.g["AuxQuantities"]["Z"] if not direct: - z = self.g.derived[var] * factor + z = self.g["AuxQuantities"][var] * factor else: z = var From b23f7fbcaf9d3fa868a70fbf5dd9b51dee5c0c76 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 8 Oct 2025 19:07:56 +0200 Subject: [PATCH 363/385] Removed omfit classes --- pyproject.toml | 9 - src/mitim_tools/gs_tools/GEQtools_old.py | 1197 ----------------- .../gs_tools/utils/GEQplotting_old.py | 976 -------------- 3 files changed, 2182 deletions(-) delete mode 100644 src/mitim_tools/gs_tools/GEQtools_old.py delete mode 100644 src/mitim_tools/gs_tools/utils/GEQplotting_old.py diff --git a/pyproject.toml b/pyproject.toml index 1f2979a8..2a3fba09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,15 +56,6 @@ dependencies = [ pyqt = [ "PyQt6", ] -omfit = [ - "omfit_classes>3.2024.19.2", # Otherwise, it will need an old version of matplotlib, matplotlib<3.6 - "scipy<1.14.0", # As of 08/08/2024, because of https://github.com/gafusion/OMFIT-source/issues/7104 - "numpy<2.0.0", # For the xarray requirement below to work - #"xarray==2022.6.0", # As of 08/08/2024, because https://github.com/gafusion/OMFIT-source/issues/7104. Froze this version because of the PLASMAstate xr reader (importlib_metadata issues) - "omas", - "fortranformat", - "openpyxl", -] test = [ "pytest", "coverage", diff --git a/src/mitim_tools/gs_tools/GEQtools_old.py b/src/mitim_tools/gs_tools/GEQtools_old.py deleted file mode 100644 index 74370d31..00000000 --- a/src/mitim_tools/gs_tools/GEQtools_old.py +++ /dev/null @@ -1,1197 +0,0 @@ -import os -import io -import tempfile -import copy -import numpy as np -import matplotlib.pyplot as plt -from mitim_tools.misc_tools import GRAPHICStools, IOtools, PLASMAtools -from mitim_tools.gacode_tools import PROFILEStools -from mitim_tools.gs_tools.utils import GEQplotting_old as GEQplotting -from shapely.geometry import LineString -from scipy.integrate import quad -import freegs -from freegs import geqdsk -from mitim_tools.misc_tools.LOGtools import printMsg as print -from IPython import embed - -""" -Note that this module relies on OMFIT classes (https://omfit.io/classes.html) procedures to intrepret the content of g-eqdsk files. -Modifications are made in MITINM for visualizations and a few extra derivations. -""" - -def fix_file(filename): - - with open(filename, "r") as f: - lines = f.readlines() - - # ----------------------------------------------------------------------- - # Remove coils (chatGPT 4o as of 08/24/24) - # ----------------------------------------------------------------------- - # Use StringIO to simulate the file writing - noCoils_file = io.StringIO() - for cont, line in enumerate(lines): - if cont > 0 and line[:2] == " ": - break - noCoils_file.write(line) - - # Reset cursor to the start of StringIO - noCoils_file.seek(0) - - # Write the StringIO content to a temporary file - with tempfile.NamedTemporaryFile(delete=False) as tmp_file: - tmp_file.write(noCoils_file.getvalue().encode('utf-8')) - noCoils_file = tmp_file.name - # ----------------------------------------------------------------------- - - with open(filename, 'r') as file1, open(noCoils_file, 'r') as file2: - file1_content = file1.read() - file2_content = file2.read() - - if file1_content != file2_content: - print(f"\t- geqdsk file {IOtools.clipstr(filename)} had coils, I have removed them") - - filename = noCoils_file - - return filename - -class MITIMgeqdsk: - def __init__(self, filename): - - # Fix file by removing coils if it has them - filename = fix_file(filename) - - # Read GEQDSK file using OMFIT - import omfit_classes.omfit_eqdsk - self.g = omfit_classes.omfit_eqdsk.OMFITgeqdsk(filename, forceFindSeparatrix=True) - - # Extra derivations in MITIM - self.derive() - - # Remove temporary file - os.remove(filename) - - @classmethod - def timeslices(cls, filename, **kwargs): - print("\n...Opening GEQ file with several time slices") - - with open(filename, "rb") as f: - lines_full = f.readlines() - - resolutions = [int(a) for a in lines_full[0].split()[-3:]] - - lines_files = [] - lines_files0 = [] - for i in range(len(lines_full)): - line = lines_full[i] - resols = [] - try: - resols = [int(a) for a in line.split()[-3:]] - except: - pass - if (resols == resolutions) and (i != 0): - lines_files.append(lines_files0) - lines_files0 = [] - lines_files0.append(line) - lines_files.append(lines_files0) - - # Write files - gs = [] - for i in range(len(lines_files)): - with open(f"{filename}_time{i}.geqdsk", "wb") as f: - f.writelines(lines_files[i]) - - gs.append( - cls( - f"{filename}_time{i}.geqdsk",**kwargs, - ) - ) - os.remove(f"{filename}_time{i}.geqdsk") - - return gs - - def derive(self, debug=False): - - self.Jt = self.g.surfAvg("Jt") * 1e-6 - self.Jt_fb = self.g.surfAvg("Jt_fb") * 1e-6 - - self.Jerror = np.abs(self.Jt - self.Jt_fb) - - self.Ip = self.g["CURRENT"] - - # Parameterizations of LCFS - self.kappa = self.g["fluxSurfaces"]["geo"]["kap"][-1] - self.kappaU = self.g["fluxSurfaces"]["geo"]["kapu"][-1] - self.kappaL = self.g["fluxSurfaces"]["geo"]["kapl"][-1] - - self.delta = self.g["fluxSurfaces"]["geo"]["delta"][-1] - self.deltaU = self.g["fluxSurfaces"]["geo"]["dell"][-1] - self.deltaL = self.g["fluxSurfaces"]["geo"]["dell"][-1] - - self.zeta = self.g["fluxSurfaces"]["geo"]["zeta"][-1] - - self.a = self.g["fluxSurfaces"]["geo"]["a"][-1] - self.Rmag = self.g["fluxSurfaces"]["geo"]["R"][0] - self.Zmag = self.g["fluxSurfaces"]["geo"]["Z"][0] - self.Rmajor = np.mean( - [ - self.g["fluxSurfaces"]["geo"]["Rmin_centroid"][-1], - self.g["fluxSurfaces"]["geo"]["Rmax_centroid"][-1], - ] - ) - - self.Zmajor = self.Zmag - - self.eps = self.a / self.Rmajor - - # Core values - - self.kappa_a = self.g["fluxSurfaces"]["geo"]["cxArea"][-1] / (np.pi * self.a**2) - - self.kappa995 = np.interp( - 0.995, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["kap"], - ) - self.kappa95 = np.interp( - 0.95, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["kap"], - ) - self.delta995 = np.interp( - 0.995, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["delta"], - ) - self.delta95 = np.interp( - 0.95, - self.g["AuxQuantities"]["PSI_NORM"], - self.g["fluxSurfaces"]["geo"]["delta"], - ) - - """ - -------------------------------------------------------------------------------------------------------------------------------------- - Boundary - -------------------------------------------------------------------------------------------------------------------------------------- - Note that the RBBS and ZBBS values in the gfile are often too scattered and do not reproduce the boundary near x-points. - The shaping parameters calculated using fluxSurfaces are correct though. - """ - - self.Rb_gfile, self.Yb_gfile = self.g["RBBBS"], self.g["ZBBBS"] - self.Rb, self.Yb = self.g["fluxSurfaces"].sep.transpose() - - if len(self.Rb) == 0: - print("\t- MITIM > No separatrix found in the OMFIT fluxSurfaces, increasing resolution and going all in!",typeMsg='i') - - flx = copy.deepcopy(self.g['fluxSurfaces']) - flx._changeResolution(6) - flx.findSurfaces([0.0,0.5,1.0]) - fs = flx['flux'][list(flx['flux'].keys())[-1]] - self.Rb, self.Yb = fs['R'], fs['Z'] - - if debug: - fig, ax = plt.subplots() - - # OMFIT - ax.plot(self.Rb, self.Yb, "-s", c="r", label="OMFIT") - - # GFILE - ax.plot(self.Rb_gfile, self.Yb_gfile, "-s", c="y", label="GFILE") - - # Extras - self.plotFluxSurfaces( - ax=ax, fluxes=[0.99999, 1.0], rhoPol=True, sqrt=False, color="m" - ) - self.plotXpointEnvelope( - ax=ax, color="c", alpha=1.0, rhoPol=True, sqrt=False - ) - self.plotEnclosingBox(ax=ax) - - ax.legend() - ax.set_aspect("equal") - ax.set_xlabel("R [m]") - ax.set_ylabel("Z [m]") - - plt.show() - - def plotEnclosingBox(self, ax=None, c= "k"): - if ax is None: - fig, ax = plt.subplots() - - Rmajor, a, Zmajor, kappaU, kappaL, deltaU, deltaL = ( - self.Rmajor, - self.a, - self.Zmajor, - self.kappaU, - self.kappaL, - self.deltaU, - self.deltaL, - ) - - ax.axhline(y=Zmajor, ls="--", c=c, lw=0.5) - ax.axvline(x=Rmajor, ls="--", c=c, lw=0.5) - ax.axvline(x=Rmajor - a, ls="--", c=c, lw=0.5) - ax.axvline(x=Rmajor + a, ls="--", c=c, lw=0.5) - Rtop = Zmajor + a * kappaU - ax.axhline(y=Rtop, ls="--", c=c, lw=0.5) - Rbot = Zmajor - a * kappaL - ax.axhline(y=Rbot, ls="--", c=c, lw=0.5) - ax.axvline(x=Rmajor - a * deltaU, ls="--", c=c, lw=0.5) - ax.axvline(x=Rmajor - a * deltaL, ls="--", c=c, lw=0.5) - - ax.plot([self.Rmajor, self.Rmag], [self.Zmajor, self.Zmag], "o", markersize=5) - - def write(self, filename=None): - """ - If filename is None, use the original one - """ - - if filename is not None: - self.g.filename = filename - - self.g.save() - - # ----------------------------------------------------------------------------- - # Parameterizations - # ----------------------------------------------------------------------------- - - def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): - - psis = self.g["AuxQuantities"]["PSI_NORM"] - flux_surfaces = self.g['fluxSurfaces']['flux'] - - # Cannot parallelize because different number of points? - kappa, rmin, rmaj, zmag, sn, cn = [],[],[],[],[],[] - - for flux in range(len(flux_surfaces)): - if flux == len(flux_surfaces)-1: - Rf, Zf = self.Rb, self.Yb - else: - Rf, Zf = flux_surfaces[flux]['R'],flux_surfaces[flux]['Z'] - - # Perform the MXH decompositionusing the MITIM surface class - surfaces = mitim_flux_surfaces() - surfaces.reconstruct_from_RZ(Rf,Zf) - surfaces._to_mxh(n_coeff=n_coeff) - - kappa.append(surfaces.kappa[0]) - rmin.append(surfaces.a[0]) - rmaj.append(surfaces.R0[0]) - zmag.append(surfaces.Z0[0]) - - sn.append(surfaces.sn[0,:]) - cn.append(surfaces.cn[0,:]) - - kappa = np.array(kappa) - rmin = np.array(rmin) - rmaj = np.array(rmaj) - zmag = np.array(zmag) - sn = np.array(sn) - cn = np.array(cn) - - if plotYN: - fig, ax = plt.subplots() - ax.plot(self.Rb, self.Yb, 'o-', c = 'b') - - surfaces = mitim_flux_surfaces() - surfaces.reconstruct_from_RZ(self.Rb, self.Yb) - surfaces._to_mxh(n_coeff=n_coeff) - surfaces._from_mxh() - - ax.plot(surfaces.R[0], surfaces.Z[0], 'o-', c = 'r') - - plt.show() - - ''' - Reminder that - sn = [0.0, np.arcsin(delta), -zeta, ...] - ''' - - return psis, rmaj, rmin, zmag, kappa, cn, sn - - # ----------------------------------------------------------------------------- - # For MAESTRO and TRANSP converstions - # ----------------------------------------------------------------------------- - - def to_profiles(self, ne0_20 = 1.0, Zeff = 1.5, PichT = 1.0, Z = 9, coeffs_MXH = 7, plotYN = False): - - # ------------------------------------------------------------------------------------------------------- - # Quantities from the equilibrium - # ------------------------------------------------------------------------------------------------------- - - rhotor = self.g['RHOVN'] - psi = self.g['AuxQuantities']['PSI'] # Wb/rad - torfluxa = self.g['AuxQuantities']['PHI'][-1] / (2*np.pi) # Wb/rad - q = self.g['QPSI'] - pressure = self.g['PRES'] # Pa - Ip = self.g['CURRENT']*1E-6 # MA - - RZ = np.array([self.Rb,self.Yb]).T - R0 = (RZ.max(axis=0)[0] + RZ.min(axis=0)[0])/2 - - B0 = self.g['RCENTR']*self.g['BCENTR'] / R0 - - # Ensure positive quantities #TODO: Check if this is necessary, pass directions - rhotor = np.array([np.abs(i) for i in rhotor]) - psi = np.array([np.abs(i) for i in psi]) - q = np.array([np.abs(i) for i in q]) - pressure = np.array([np.abs(i) for i in pressure]) - - torfluxa = np.abs(torfluxa) - Ip = np.abs(Ip) - B0 = np.abs(B0) - # ------------------------------------------ - - _, rmaj, rmin, zmag, kappa, cn, sn = self.get_MXH_coeff_new(n_coeff=coeffs_MXH) - - delta = np.sin(sn[:,1]) - zeta = -sn[:,2] - - # ------------------------------------------------------------------------------------------------------- - # Pass to profiles - # ------------------------------------------------------------------------------------------------------- - - profiles = {} - - profiles['nexp'] = np.array([f'{rhotor.shape[0]}']) - profiles['nion'] = np.array(['2']) - profiles['shot'] = np.array(['12345']) - - # Just one specie - profiles['name'] = np.array(['D','F']) - profiles['type'] = np.array(['[therm]','[therm]']) - profiles['masse'] = np.array([5.4488748e-04]) - profiles['mass'] = np.array([2.0, Z*2]) - profiles['ze'] = np.array([-1.0]) - profiles['z'] = np.array([1.0, Z]) - - profiles['torfluxa(Wb/radian)'] = np.array([torfluxa]) - profiles['rcentr(m)'] = np.array([R0]) - profiles['bcentr(T)'] = np.array([B0]) - profiles['current(MA)'] = np.array([Ip]) - - profiles['rho(-)'] = rhotor - profiles['polflux(Wb/radian)'] = psi - profiles['q(-)'] = q - - # ------------------------------------------------------------------------------------------------------- - # Flux surfaces - # ------------------------------------------------------------------------------------------------------- - - profiles['kappa(-)'] = kappa - profiles['delta(-)'] = delta - profiles['zeta(-)'] = zeta - profiles['rmin(m)'] = rmin - profiles['rmaj(m)'] = rmaj - profiles['zmag(m)'] = zmag - - sn, cn = np.array(sn), np.array(cn) - for i in range(coeffs_MXH): - profiles[f'shape_cos{i}(-)'] = cn[:,i] - for i in range(coeffs_MXH-3): - profiles[f'shape_sin{i+3}(-)'] = sn[:,i+3] - - ''' - ------------------------------------------------------------------------------------------------------- - Kinetic profiles - ------------------------------------------------------------------------------------------------------- - Pressure division into temperature and density - p_Pa = p_e + p_i = Te_eV * e_J * ne_20 * 1e20 + Ti_eV * e_J * ni_20 * 1e20 - if T=Te=Ti and ne=ni - p_Pa = 2 * T_eV * e_J * ne_20 * 1e20 - T_eV = p_Pa / (2 * e_J * ne_20 * 1e20) - ''' - - C = 1 / (2 * 1.60217662e-19 * 1e20) - _, ne_20 = PLASMAtools.parabolicProfile(Tbar=ne0_20/1.25,nu=1.25,rho=rhotor,Tedge=ne0_20/5) - T_keV = C * (pressure / ne_20) * 1E-3 - - fZ = (Zeff-1) / (Z**2-Z) # One-impurity model to give desired Zeff - - profiles['te(keV)'] = T_keV - profiles['ti(keV)'] = np.array([T_keV]*2).T - profiles['ne(10^19/m^3)'] = ne_20*10.0 - profiles['ni(10^19/m^3)'] = np.array([profiles['ne(10^19/m^3)']*(1-Z*fZ),profiles['ne(10^19/m^3)']*fZ]).T - - # ------------------------------------------------------------------------------------------------------- - # Power: insert parabolic and use PROFILES volume integration to find desired power - # ------------------------------------------------------------------------------------------------------- - - _, profiles["qrfe(MW/m^3)"] = PLASMAtools.parabolicProfile(Tbar=1.0,nu=5.0,rho=rhotor,Tedge=0.0) - - p = PROFILEStools.gacode_state.scratch(profiles) - - p.profiles["qrfe(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] * PichT/p.derived['qRF_MW'][-1] /2 - p.profiles["qrfi(MW/m^3)"] = p.profiles["qrfe(MW/m^3)"] - - # ------------------------------------------------------------------------------------------------------- - # Ready to go - # ------------------------------------------------------------------------------------------------------- - - p.derive_quantities() - - # ------------------------------------------------------------------------------------------------------- - # Plotting - # ------------------------------------------------------------------------------------------------------- - - if plotYN: - - fig, ax = plt.subplots() - ff = np.linspace(0, 1, 11) - self.plotFluxSurfaces(ax=ax, fluxes=ff, rhoPol=False, sqrt=True, color="r", plot1=False) - p.plot_state_flux_surfaces(ax=ax, surfaces_rho=ff, color="b") - plt.show() - - return p - - def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', ne0_20 = 1E19, Vsurf = 0.0, Zeff = 1.5, PichT_MW = 11.0, times = [0.0,1.0]): - - print("\t- Converting to TRANSP") - folder = IOtools.expandPath(folder) - folder.mkdir(parents=True, exist_ok=True) - - p = self.to_profiles(ne0_20 = ne0_20, Zeff = Zeff, PichT = PichT_MW) - p.write_state(folder / 'input.gacode') - - transp = p.to_transp(folder = folder, shot = shot, runid = runid, times = times, Vsurf = Vsurf) - - return transp - - # --------------------------------------------------------------------------------------------------------------------------------------- - # Plotting - # --------------------------------------------------------------------------------------------------------------------------------------- - - def plot(self, fn=None, extraLabel=""): - GEQplotting.plot(self, fn=fn, extraLabel=extraLabel) - - def plotFS(self, axs=None, color="b", label=""): - GEQplotting.plotFS(self, axs=axs, color=color, label=label) - - def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): - GEQplotting.plotPlasma(self, axs=axs, legendYN=legendYN, color=color, label=label) - - def plotGeometry(self, axs=None, color="r"): - GEQplotting.plotGeometry(self, axs=axs, color=color) - - def plotFluxSurfaces(self, ax=None, fluxes=[1.0], color="b", alpha=1.0, rhoPol=True, sqrt=False, lw=1, lwB=2, plot1=True, label = ''): - return GEQplotting.plotFluxSurfaces(self, ax=ax, fluxes=fluxes, color=color, alpha=alpha, rhoPol=rhoPol, sqrt=sqrt, lw=lw, lwB=lwB, plot1=plot1, label = label) - - def plotXpointEnvelope(self, ax=None, color="b", alpha=1.0, rhoPol=True, sqrt=False): - GEQplotting.plotXpointEnvelope(self, ax=ax, color=color, alpha=alpha, rhoPol=rhoPol, sqrt=sqrt) - -# ----------------------------------------------------------------------------- -# Tools to handle flux surface definitions -# ----------------------------------------------------------------------------- -class mitim_flux_surfaces: - - def reconstruct_from_mxh_moments(self,R0, a, kappa, Z0, cn, sn, thetas = None): - ''' - sn = [0.0, np.arcsin(delta), -zeta, ...] - cn = [...] - - You can provide a multi-dim array of (radii, ) - ''' - - self.R0 = R0 - self.a = a - self.kappa = kappa - self.Z0 = Z0 - self.cn = cn - self.sn = sn - - self.delta = np.sin(self.sn[...,1]) - self.zeta = -self.sn[...,2] - - self._from_mxh(thetas = thetas) - - def reconstruct_from_miller(self,R0, a, kappa, Z0, delta, zeta, thetas = None): - ''' - sn = [0.0, np.arcsin(delta), -zeta, ...] - cn = [...] - - You can provide a multi-dim array of (radii, ) or not - ''' - - self.R0 = R0 - self.a = a - self.kappa = kappa - self.Z0 = Z0 - self.delta = delta - self.zeta = zeta - - self.cn = np.array([0.0,0.0,0.0]) - self.sn = np.array([0.0,np.arcsin(self.delta),-self.zeta]) - - self._from_mxh(thetas = thetas) - - def _from_mxh(self, thetas = None): - - self.R, self.Z = from_mxh_to_RZ(self.R0, self.a, self.kappa, self.Z0, self.cn, self.sn, thetas = thetas) - - def reconstruct_from_RZ(self, R, Z): - - self.R = R - self.Z = Z - - if len(self.R.shape) == 1: - self.R = self.R[np.newaxis,:] - self.Z = self.Z[np.newaxis,:] - - self._to_miller() - - def _to_mxh(self, n_coeff=6): - - self.cn = np.zeros((self.R.shape[0],n_coeff)) - self.sn = np.zeros((self.R.shape[0],n_coeff)) - self.gn = np.zeros((self.R.shape[0],4)) - for i in range(self.R.shape[0]): - self.cn[i,:], self.sn[i,:], self.gn[i,:] = from_RZ_to_mxh(self.R[i,:], self.Z[i,:], n_coeff=n_coeff) - - [self.R0, self.a, self.Z0, self.kappa] = self.gn.T - - def _to_miller(self): - - Rmin = np.min(self.R, axis=-1) - Rmax = np.max(self.R, axis=-1) - Zmax = np.max(self.Z, axis=-1) - Zmin = np.min(self.Z, axis=-1) - - self.R0 = 0.5* (Rmax + Rmin) - self.Z0 = 0.5* (Zmax + Zmin) - - self.a = (Rmax - Rmin) / 2 - - # Elongations - - self.kappa_u = (Zmax - self.Z0) / self.a - self.kappa_l = (self.Z0 - Zmin) / self.a - self.kappa = (self.kappa_u + self.kappa_l) / 2 - - # Triangularities - - RatZmax = self.R[np.arange(self.R.shape[0]),np.argmax(self.Z,axis=-1)] - self.delta_u = (self.R0-RatZmax) / self.a - - RatZmin = self.R[np.arange(self.R.shape[0]),np.argmin(self.Z,axis=-1)] - self.delta_l = (self.R0-RatZmin) / self.a - - self.delta = (self.delta_u + self.delta_l) / 2 - - # Squareness (not parallel for the time being) - self.zeta = np.zeros(self.R0.shape) - for i in range(self.R0.shape[0]): - try: - Ri, Zi, zeta_uo = find_squareness_points(self.R[i,:], self.Z[i,:]) - except AttributeError: - zeta_uo = np.nan - self.zeta[i] = zeta_uo - - def plot(self, ax = None, color = 'r', label = None, plot_extremes=False): - - if ax is None: - fig, ax = plt.subplots() - - for i in range(self.R.shape[0]): - ax.plot(self.R[i,:], self.Z[i,:], color = color, label = label) - - if plot_extremes: - ax.plot([self.R[i,self.Z[i,:].argmax()]], [self.Z[i,:].max()], 'o', color=color, markersize=5) - ax.plot([self.R[i,self.Z[i,:].argmin()]], [self.Z[i,:].min()], 'o', color=color, markersize=5) - - ax.set_aspect('equal') - ax.set_xlabel('R [m]') - ax.set_ylabel('Z [m]') - if label is not None: - ax.legend() - GRAPHICStools.addDenseAxis(ax) - -def find_squareness_points(R, Z, debug = False): - - # Reference point (A) - A_r = R_of_maxZ = R[Z.argmax()] - A_z = Z_of_maxR = Z[R.argmax()] - - # Upper Outer Squareness point (D) - C_r = R_of_maxR = R.max() - C_z = Z_of_maxZ = Z.max() - - # Find intersection with separatrix (C) - Ri_uo, Zi_uo = find_intersection_squareness(R, Z, R_of_maxZ, Z_of_maxR, R_of_maxR, Z_of_maxZ) - C_r, C_z = Ri_uo, Zi_uo - - # Find point B - B_r = Rs_uo = R_of_maxZ + (R_of_maxR-R_of_maxZ)/2 - B_z = Zs_uo = Z_of_maxR + (Z_of_maxZ-Z_of_maxR)/2 - - # Squareness is defined as the distance BC divided by the distance AB - zeta_uo = np.sqrt((C_r-B_r)**2 + (C_z-B_z)**2) / np.sqrt((A_r-B_r)**2 + (A_z-B_z)**2) - #zeta_uo = np.sqrt((Ri_uo-Rs_uo)**2 + (Zi_uo-Zs_uo)**2) / np.sqrt((R_of_maxZ-R_of_maxR)**2 + (Z_of_maxZ-Z_of_maxR)**2) - - if debug: - plt.ion() - fig, ax = plt.subplots() - ax.plot(R, Z, 'o-', markersize=3, color='k') - ax.plot([Ri_uo], [Zi_uo], 'o', color='k', label = 'C', markersize=6) - - ax.plot([R_of_maxZ], [Z_of_maxR], 'o', color='b', label = 'A') - ax.axvline(x=R_of_maxZ, ls='--', color='b') - ax.axhline(y=Z_of_maxR, ls='--', color='b') - - ax.plot([R_of_maxR], [Z_of_maxZ], 'o', color='r', label = 'D') - ax.axhline(y=Z_of_maxZ, ls='--', color='r') - ax.axvline(x=R_of_maxR, ls='--', color='r') - - ax.plot([Rs_uo], [Zs_uo], 'o', color='g', label = 'B') - - # Connect A and D - ax.plot([R_of_maxZ, R_of_maxR], [Z_of_maxR, Z_of_maxZ], 'm--') - - # Connect maxZ with maxR - ax.plot([R_of_maxZ, R_of_maxR], [Z_of_maxZ, Z_of_maxR], 'm--') - - ax.set_aspect('equal') - ax.legend() - ax.set_xlabel('R [m]') - ax.set_ylabel('Z [m]') - - return Ri_uo, Zi_uo, zeta_uo - -def find_intersection_squareness(R, Z, Ax, Az, Dx, Dz): - - R_line = np.linspace(Ax, Dx, 100) - Z_line = np.linspace(Az, Dz, 100) - line1 = LineString(zip(R_line, Z_line)) - line2 = LineString(zip(R, Z)) - intersection = line1.intersection(line2) - - return intersection.x, intersection.y - -# ----------------------------------------------------------------------------- -# Utilities for parameterizations -# ----------------------------------------------------------------------------- - -def from_RZ_to_mxh(R, Z, n_coeff=3): - """ - Calculates MXH Coefficients for a flux surface - """ - Z = np.roll(Z, -np.argmax(R)) - R = np.roll(R, -np.argmax(R)) - if Z[1] < Z[0]: # reverses array so that theta increases - Z = np.flip(Z) - R = np.flip(R) - - # compute bounding box for each flux surface - r = 0.5*(np.max(R)-np.min(R)) - kappa = 0.5*(np.max(Z) - np.min(Z))/r - R0 = 0.5*(np.max(R)+np.min(R)) - Z0 = 0.5*(np.max(Z)+np.min(Z)) - bbox = [R0, r, Z0, kappa] - - # solve for polar angles - # need to use np.clip to avoid floating-point precision errors - theta_r = np.arccos(np.clip(((R - R0) / r), -1, 1)) - theta = np.arcsin(np.clip(((Z - Z0) / r / kappa),-1,1)) - - # Find the continuation of theta and theta_r to [0,2pi] - theta_r_cont = np.copy(theta_r) - theta_cont = np.copy(theta) - - max_theta = np.argmax(theta) - min_theta = np.argmin(theta) - max_theta_r = np.argmax(theta_r) - min_theta_r = np.argmin(theta_r) - - theta_cont[:max_theta] = theta_cont[:max_theta] - theta_cont[max_theta:max_theta_r] = np.pi-theta[max_theta:max_theta_r] - theta_cont[max_theta_r:min_theta] = np.pi-theta[max_theta_r:min_theta] - theta_cont[min_theta:] = 2*np.pi+theta[min_theta:] - - theta_r_cont[:max_theta] = theta_r_cont[:max_theta] - theta_r_cont[max_theta:max_theta_r] = theta_r[max_theta:max_theta_r] - theta_r_cont[max_theta_r:min_theta] = 2*np.pi - theta_r[max_theta_r:min_theta] - theta_r_cont[min_theta:] = 2*np.pi - theta_r[min_theta:] - - theta_r_cont = theta_r_cont - theta_cont ; theta_r_cont[-1] = theta_r_cont[0] - - # Fourier decompose to find coefficients - - c, s = np.zeros(n_coeff), np.zeros(n_coeff) - - def f_theta_r(theta): - return np.interp(theta, theta_cont, theta_r_cont) - - for i in np.arange(n_coeff): - s[i] = quad(f_theta_r,0,2*np.pi, weight="sin", wvar=i)[0]/np.pi - c[i] = quad(f_theta_r,0,2*np.pi, weight="cos", wvar=i)[0]/np.pi - - c[0] /= 2 - - return c, s, bbox - -def from_mxh_to_RZ(R0, a, kappa, Z0, cn, sn, thetas = None): - - if thetas is None: - thetas = np.linspace(0, 2 * np.pi, 100) - - # Prepare data to always have the first dimension a batch (e.g. a radius) for parallel computation - if IOtools.isfloat(R0): - R0 = [R0] - a = [a] - kappa = [kappa] - Z0 = [Z0] - cn = np.array(cn)[np.newaxis,:] - sn = np.array(sn)[np.newaxis,:] - - R0 = np.array(R0) - a = np.array(a) - kappa = np.array(kappa) - Z0 = np.array(Z0) - - R = np.zeros((R0.shape[0],len(thetas))) - Z = np.zeros((R0.shape[0],len(thetas))) - n = np.arange(1, sn.shape[1]) - for i,theta in enumerate(thetas): - theta_R = theta + cn[:,0] + np.sum( cn[:,1:]*np.cos(n*theta) + sn[:,1:]*np.sin(n*theta), axis=-1 ) - R[:,i] = R0 + a*np.cos(theta_R) - Z[:,i] = Z0 + kappa*a*np.sin(theta) - - return R, Z - -# -------------------------------------------------------------------------------------------------------------- -# Fixed boundary stuff -# -------------------------------------------------------------------------------------------------------------- - -class freegs_millerized: - - def __init__(self, R, a, kappa_sep, delta_sep, zeta_sep, z0): - - print("> Fixed-boundary equilibrium with FREEGS") - - print("\t- Initializing miller geometry") - print(f"\t\t* R={R} m, a={a} m, kappa_sep={kappa_sep}, delta_sep={delta_sep}, zeta_sep={zeta_sep}, z0={z0} m") - - self.R0 = R - self.a = a - self.kappa_sep = kappa_sep - self.delta_sep = delta_sep - self.zeta_sep = zeta_sep - self.Z0 = z0 - - thetas = np.linspace(0, 2*np.pi, 1000, endpoint=False) - - self.mitim_separatrix = mitim_flux_surfaces() - self.mitim_separatrix.reconstruct_from_miller(self.R0, self.a, self.kappa_sep, self.Z0, self.delta_sep, self.zeta_sep, thetas = thetas) - self.R_sep, self.Z_sep = self.mitim_separatrix.R[0,:], self.mitim_separatrix.Z[0,:] - - def prep(self, p0_MPa, Ip_MA, B_T, - beta_pol = None, n_coils = 10, resol_eq = 2**8+1, - parameters_profiles = {'alpha_m':2.0, 'alpha_n':2.0, 'Raxis':1.0}, - constraint_miller_squareness_point = False): - - print("\t- Initializing plasma parameters") - if beta_pol is not None: - print(f"\t\t* beta_pol={beta_pol:.5f}, Ip={Ip_MA:.5f} MA, B={B_T:.5f} T") - else: - print(f"\t\t* p0={p0_MPa:.5f} MPa, Ip={Ip_MA:.5f} MA, B={B_T:.5f} T") - - self.p0_MPa = p0_MPa - self.Ip_MA = Ip_MA - self.B_T = B_T - self.beta_pol = beta_pol - self.parameters_profiles = parameters_profiles - - print(f"\t- Preparing equilibrium with FREEGS, with a resolution of {resol_eq}x{resol_eq}") - self._define_coils(n_coils) - self._define_eq(resol = resol_eq) - self._define_gs() - - # Define xpoints - print("\t\t* Defining upper and lower x-points") - self.xpoints = [ - (self.R0 - self.a*self.delta_sep, self.Z0+self.a*self.kappa_sep), - (self.R0 - self.a*self.delta_sep, self.Z0-self.a*self.kappa_sep), - ] - - # Define isoflux - print("\t\t* Defining midplane separatrix (isoflux)") - self.isoflux = [ - (self.xpoints[0][0], self.xpoints[0][1], self.R0 + self.a, self.Z0), # Upper x-point with outer midplane - (self.xpoints[0][0], self.xpoints[0][1], self.R0 - self.a, self.Z0), # Upper x-point with inner midplane - (self.xpoints[0][0], self.xpoints[0][1], self.xpoints[1][0], self.xpoints[1][1]), # Between x-points - ] - - print("\t\t* Defining squareness isoflux point") - - # Find squareness point - if constraint_miller_squareness_point: - Rsq, Zsq, _ = find_squareness_points(self.R_sep, self.Z_sep) - - self.isoflux.append( - (self.xpoints[0][0], self.xpoints[0][1], Rsq, Zsq) # Upper x-point with squareness point - ) - self.isoflux.append( - (self.xpoints[0][0], self.xpoints[0][1], Rsq, -Zsq) # Upper x-point with squareness point - ) - - # Combine - self.constrain = freegs.control.constrain( - isoflux=self.isoflux, - xpoints=self.xpoints, - ) - - def _define_coils(self, n, rel_distance_coils = 0.5, updown_coils = True): - - print(f"\t- Defining {n} coils{' (up-down symmetric)' if updown_coils else ''} at a distance of {rel_distance_coils}*a from the separatrix") - - self.distance_coils = self.a*rel_distance_coils - self.updown_coils = updown_coils - - if self.updown_coils: - thetas = np.linspace(0, np.pi, n) - else: - thetas = np.linspace(0, 2*np.pi, n, endpoint=False) - - self.mitim_coils_surface = mitim_flux_surfaces() - self.mitim_coils_surface.reconstruct_from_miller(self.R0, (self.a+self.distance_coils), self.kappa_sep, self.Z0, self.delta_sep, self.zeta_sep, thetas = thetas) - self.Rcoils, self.Zcoils = self.mitim_coils_surface.R[0,:], self.mitim_coils_surface.Z[0,:] - - self.coils = [] - for num, (R, Z) in enumerate(zip(self.Rcoils, self.Zcoils)): - - if self.updown_coils and Z > 0: - coilU = freegs.machine.Coil( - R, - Z - ) - coilL = freegs.machine.Coil( - R, - -Z - ) - coil = freegs.machine.Circuit( [ ('U', coilU, 1.0 ), ('L', coilL, 1.0 ) ] ) - else: - - coil = freegs.machine.Coil( - R, - Z - ) - - self.coils.append( - (f"coil_{num}", coil) - ) - - def _define_eq(self, resol=2**9+1): - - print("\t- Defining equilibrium") - self.tokamak = freegs.machine.Machine(self.coils) - - a = self.a + self.distance_coils - - Rmin = (self.R0-a) - a*0.25 - Rmax = (self.R0+a) + a*0.25 - - b = a*self.kappa_sep - Zmin = (self.Z0 - b) - b*0.25 - Zmax = (self.Z0 + b) + b*0.25 - - self.eq = freegs.Equilibrium(tokamak=self.tokamak, - Rmin=Rmin, Rmax=Rmax, - Zmin=Zmin, Zmax=Zmax, - nx=resol, ny=resol, - boundary=freegs.boundary.freeBoundaryHagenow) - - def _define_gs(self): - - if self.beta_pol is None: - print("\t- Defining Grad-Shafranov equilibrium: p0, Ip and vaccum R*Bt") - - self.profiles = freegs.jtor.ConstrainPaxisIp(self.eq, - self.p0_MPa*1E6, self.Ip_MA*1E6, self.R0*self.B_T, - alpha_m=self.parameters_profiles['alpha_m'], alpha_n=self.parameters_profiles['alpha_n'], Raxis=self.parameters_profiles['Raxis']) - - else: - print("\t- Defining Grad-Shafranov equilibrium: beta_pol, Ip and vaccum R*Bt") - - self.profiles = freegs.jtor.ConstrainBetapIp(self.eq, - self.beta_pol, self.Ip_MA*1E6, self.R0*self.B_T, - alpha_m=self.parameters_profiles['alpha_m'], alpha_n=self.parameters_profiles['alpha_n'], Raxis=self.parameters_profiles['Raxis']) - - def solve(self, show = False, rtol=1e-6): - - print("\t- Solving equilibrium with FREEGS") - with IOtools.timer(): - self.x,self.y = freegs.solve(self.eq, # The equilibrium to adjust - self.profiles, # The toroidal current profile function - constrain=self.constrain, # Constraint function to set coil currents - show=show, - rtol=rtol, # Default is 1e-3 - atol=1e-10, - maxits=100, # Default is 50 - convergenceInfo=True) - print("\t\t * Done!") - - self.check() - - def check(self, warning_error = 0.01, plotYN = False): - - print("\t- Checking separatrix quality (Miller vs FREEGS)") - RZ = self.eq.separatrix() - - self.mitim_separatrix_eq = mitim_flux_surfaces() - self.mitim_separatrix_eq.reconstruct_from_RZ(RZ[:,0], RZ[:,1]) - - # -------------------------------------------------------------- - # Check errors - # -------------------------------------------------------------- - - max_error = 0.0 - for key in ['R0', 'a', 'kappa_sep', 'delta_sep', 'zeta_sep']: - miller_value = getattr(self, key) - sep_value = getattr(self.mitim_separatrix_eq, key.replace('_sep', ''))[0] - error = abs( (miller_value-sep_value)/miller_value ) - print(f"\t\t* {key}: {miller_value:.3f} vs {sep_value:.3f} ({100*error:.2f}%)") - - max_error = np.max([max_error, error]) - - if max_error > warning_error: - print(f"\t\t- Maximum error in equilibrium quantities is {100*max_error:.2f}%", typeMsg='w') - else: - print(f"\t\t- Maximum error in equilibrium quantities is {100*max_error:.2f}%") - - # -------------------------------------------------------------- - # Plotting - # -------------------------------------------------------------- - - if plotYN: - - fig = plt.figure(figsize=(12,8)) - axs = fig.subplot_mosaic( - """ - AB - AB - CB - """ - ) - - # Plot direct FreeGS output - - ax = axs['A'] - self.eq.plot(axis=ax,show=False) - self.constrain.plot(axis=ax, show=False) - - for coil in self.coils: - if isinstance(coil[1],freegs.machine.Circuit): - ax.plot([coil[1]['U'].R], [coil[1]['U'].Z], 's', c='k', markersize=2) - ax.plot([coil[1]['L'].R], [coil[1]['L'].Z], 's', c='k', markersize=2) - else: - ax.plot([coil[1].R], [coil[1].Z], 's', c='k', markersize=2) - - GRAPHICStools.addLegendApart(ax,ratio=0.9,size=10) - - ax = axs['C'] - ax.plot(self.x,'-o', markersize=3, color='b', label = '$\\psi$ max change') - ax.set_xlabel('Iteration') - ax.set_ylabel('$\\psi$ max change') - ax.set_yscale('log') - ax.legend(loc='lower left',prop={'size': 10}) - - ax = axs['C'].twinx() - ax.plot(self.y,'-o', markersize=3, color='r', label = '$\\psi$ max relative change') - ax.set_ylabel('$\\psi$ max relative change') - ax.set_yscale('log') - ax.legend(loc='upper right',prop={'size': 10}) - - # Plot comparison of equilibria - - ax = axs['B'] - - self.mitim_separatrix.plot(ax=ax, color = 'b', label = 'Miller (original)', plot_extremes=True) - self.mitim_separatrix_eq.plot(ax=ax, color = 'r', label = 'Separatrix (freegs)', plot_extremes=True) - - ax.legend(prop={'size': 10}) - - plt.show() - - def derive(self, psi_surfaces = np.linspace(0,1.0,10), psi_profiles = np.linspace(0,1.0,100)): - - # Grab surfaces - Rs, Zs = [], [] - for psi_norm in psi_surfaces: - R, Z = self.find_surface(psi_norm) - Rs.append(R) - Zs.append(Z) - Rs = np.array(Rs) - Zs = np.array(Zs) - - # Calculate surface stuff in parallel - self.surfaces = mitim_flux_surfaces() - self.surfaces.reconstruct_from_RZ(Rs, Zs) - self.surfaces.psi = psi_surfaces - - # Grab profiles - self.profile_psi_norm = psi_profiles - self.profile_pressure = self.eq.pressure(psinorm =psi_profiles)*1E-6 - self.profile_q = self.eq.q(psinorm = psi_profiles) - self.profile_RB = self.eq.fpol(psinorm = psi_profiles) - - # Grab quantities - self.profile_q95 = self.eq.q(psinorm = 0.95) - self.profile_q0 = self.eq.q(psinorm = 0.0) - self.profile_betaN = self.eq.betaN() - self.profile_Li2 = self.eq.internalInductance2() - self.profile_pave = self.eq.pressure_ave() - self.profile_beta_pol = self.eq.poloidalBeta() - self.profile_Ashaf = self.eq.shafranovShift - - def find_surface(self, psi_norm = 0.5, thetas = None): - - if psi_norm == 0.0: - psi_norm = 1E-6 - - if psi_norm == 1.0: - RZ = self.eq.separatrix(npoints= 1000 if thetas is None else len(thetas)) - R, Z = RZ[:,0], RZ[:,1] - else: - if thetas is None: - thetas = np.linspace(0, 2*np.pi, 1000, endpoint=False) - - from freegs.critical import find_psisurface, find_critical - from scipy import interpolate - - psi = self.eq.psi() - opoint, xpoint = find_critical(self.eq.R, self.eq.Z, psi) - psinorm = (psi - opoint[0][2]) / (self.eq.psi_bndry - opoint[0][2]) - psifunc = interpolate.RectBivariateSpline(self.eq.R[:, 0], self.eq.Z[0, :], psinorm) - r0, z0 = opoint[0][0:2] - - R = np.zeros(len(thetas)) - Z = np.zeros(len(thetas)) - for i,theta in enumerate(thetas): - R[i],Z[i] = find_psisurface( - self.eq, - psifunc, - r0, - z0, - r0 + 10.0 * np.sin(theta), - z0 + 10.0 * np.cos(theta), - psival=psi_norm, - n=1000, - ) - return R,Z - - # -------------------------------------------------------------- - # Plotting - # -------------------------------------------------------------- - - def plot(self, axs = None, color = 'b', label = ''): - - if axs is None: - plt.ion() - fig = plt.figure(figsize=(16,7)) - axs = fig.subplot_mosaic( - """ - A12 - A34 - """) - axs = [axs['A'], axs['1'], axs['2'], axs['3'], axs['4']] - - self.plot_flux_surfaces(ax = axs[0], color = color) - self.plot_profiles(axs = axs[1:], color = color, label = label) - - def plot_flux_surfaces(self, ax = None, color = 'b'): - - if ax is None: - plt.ion() - fig, ax = plt.subplots(figsize=(12,8)) - - for i in range(self.surfaces.R.shape[0]): - ax.plot(self.surfaces.R[i],self.surfaces.Z[i], '-', label = f'$\\psi$ = {self.surfaces.psi[i]:.2f}', color = color, markersize=3) - - ax.set_aspect('equal') - ax.set_xlabel('R [m]') - ax.set_ylabel('Z [m]') - GRAPHICStools.addDenseAxis(ax) - - def plot_profiles(self, axs = None, color = 'b', label = ''): - - if axs is None: - plt.ion() - fig, axs = plt.subplots(nrows=2,ncols=2,figsize=(8,8)) - axs = axs.flatten() - - ax = axs[0] - ax.plot(self.profile_psi_norm,self.profile_pressure,'-',color=color, label = label) - ax.set_xlabel('$\\psi$') - ax.set_xlim([0,1]) - ax.set_ylabel('Pressure (MPa)') - GRAPHICStools.addDenseAxis(ax) - - ax = axs[1] - ax.plot(self.profile_psi_norm,self.profile_q,'-',color=color) - ax.axhline(y=1, color='k', ls='--', lw=0.5) - ax.set_xlabel('$\\psi$') - ax.set_ylabel('q') - ax.set_xlim([0,1]) - GRAPHICStools.addDenseAxis(ax) - - ax = axs[2] - ax.plot(self.profile_psi_norm,self.profile_RB,'-',color=color) - ax.axhline(y=self.R0*self.B_T, color=color, ls='--', lw=0.5) - ax.set_xlabel('$\\psi$') - ax.set_ylabel('$R\\cdot B_t$ ($T\\cdot m$)') - ax.set_xlim([0,1]) - GRAPHICStools.addDenseAxis(ax) - - def plot_flux_surfaces_characteristics(self, axs = None, color = 'b', label = ''): - - if axs is None: - plt.ion() - fig, axs = plt.subplots(nrows=2,ncols=2,figsize=(8,8)) - axs = axs.flatten() - - ax = axs[0] - ax.plot(self.surfaces.psi, self.surfaces.kappa, '-o', color=color, label = label, markersize=3) - ax.set_xlabel('$\\psi$') - ax.set_ylabel('$\\kappa$') - ax.set_xlim([0,1]) - GRAPHICStools.addDenseAxis(ax) - - ax = axs[1] - ax.plot(self.surfaces.psi, self.surfaces.delta, '-o', color=color, label = label, markersize=3) - ax.set_xlabel('$\\psi$') - ax.set_ylabel('$\\delta$') - ax.set_xlim([0,1]) - GRAPHICStools.addDenseAxis(ax) - - # -------------------------------------------------------------- - # Writing - # -------------------------------------------------------------- - - def write(self, filename): - - print(f"\t- Writing equilibrium to {IOtools.clipstr(filename)}") - - with open(filename, "w") as f: - geqdsk.write(self.eq, f) - - def to_profiles(self, scratch_folder = '~/scratch/'): - - # Produce geqdsk object - scratch_folder = IOtools.expandPath(scratch_folder) - file_scratch = scratch_folder / 'mitim_freegs.geqdsk' - self.write(file_scratch) - g = MITIMgeqdsk(file_scratch) - - os.remove(file_scratch) - - # From geqdsk to profiles - return g.to_profiles() - - def to_transp(self, folder = '~/scratch/', shot = '12345', runid = 'P01', ne0_20 = 1E19, Vsurf = 0.0, Zeff = 1.5, PichT_MW = 11.0, times = [0.0,1.0]): - - # Produce geqdsk object - scratch_folder = IOtools.expandPath(folder) - scratch_folder.mkdir(parents=True, exist_ok=True) - file_scratch = scratch_folder / 'mitim_freegs.geqdsk' - self.write(file_scratch) - g = MITIMgeqdsk(file_scratch) - - return g.to_transp(folder=folder, shot=shot, runid=runid, ne0_20=ne0_20, Vsurf=Vsurf, Zeff=Zeff, PichT_MW=PichT_MW, times=times) diff --git a/src/mitim_tools/gs_tools/utils/GEQplotting_old.py b/src/mitim_tools/gs_tools/utils/GEQplotting_old.py deleted file mode 100644 index fe6531b4..00000000 --- a/src/mitim_tools/gs_tools/utils/GEQplotting_old.py +++ /dev/null @@ -1,976 +0,0 @@ -import numpy as np -import matplotlib.pyplot as plt -from mitim_tools.misc_tools import GRAPHICStools -from mitim_tools.misc_tools.LOGtools import printMsg as print -from IPython import embed - - -def compareGeqdsk(geqdsks, fn=None, extraLabel="", plotAll=True, labelsGs=None): - - if fn is None: - from mitim_tools.misc_tools.GUItools import FigureNotebook - fn = FigureNotebook("GEQDSK Notebook", geometry="1600x1000") - - if labelsGs is None: - labelsGs = [] - for i, g in enumerate(geqdsks): - labelsGs.append(f"#{i + 1}") - - # ----------------------------------------------------------------------------- - # Plot All - # ----------------------------------------------------------------------------- - if plotAll: - for i, g in enumerate(geqdsks): - _ = g.plot(fn=fn, extraLabel=f"{labelsGs[i]} - ") - - # ----------------------------------------------------------------------------- - # Compare in same plot - Surfaces - # ----------------------------------------------------------------------------- - fig = fn.add_figure(label=extraLabel + "Comp. - Surfaces") - - grid = plt.GridSpec(2, 3, hspace=0.3, wspace=0.3) - ax1 = fig.add_subplot(grid[:, 0]) - ax2 = fig.add_subplot(grid[:, 1]) - ax3 = fig.add_subplot(grid[0, 2]) - ax4 = fig.add_subplot(grid[1, 2]) - axs = [ax1, ax2, ax3, ax4] - - cols = GRAPHICStools.listColors() - - for i, g in enumerate(geqdsks): - g.plotFS(axs=axs, color=cols[i], label=f"#{i + 1}") - - ax3.legend() - - # ----------------------------------------------------------------------------- - # Compare in same plot - Surfaces - # ----------------------------------------------------------------------------- - fig = fn.add_figure(label=extraLabel + "Comp. - Plasma") - grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) - - ax_plasma = [ - fig.add_subplot(grid[0, 0]), - fig.add_subplot(grid[1, 0]), - fig.add_subplot(grid[0, 1]), - fig.add_subplot(grid[1, 1]), - fig.add_subplot(grid[0, 2]), - fig.add_subplot(grid[1, 2]), - fig.add_subplot(grid[0, 3]), - ] # , - # fig.add_subplot(grid[1,3])] - - cols = GRAPHICStools.listColors() - - for i, g in enumerate(geqdsks): - g.plotFS(axs=[ax1, ax2, ax3, ax4], color=cols[i], label=f"{labelsGs[i]} ") - g.plotPlasma( - axs=ax_plasma, - legendYN=False, - color=cols[i], - label=f"{labelsGs[i]} ", - ) - - return ax_plasma, fn - -# ----------------------------------------------------------------------------- -# Plot of GEQ class -# ----------------------------------------------------------------------------- - -def plot(self, fn=None, extraLabel=""): - if fn is None: - wasProvided = False - - from mitim_tools.misc_tools.GUItools import FigureNotebook - - self.fn = FigureNotebook("GEQDSK Notebook", geometry="1600x1000") - else: - wasProvided = True - self.fn = fn - # ----------------------------------------------------------------------------- - # OMFIT Summary - # ----------------------------------------------------------------------------- - # fig = self.fn.add_figure(label=extraLabel+'OMFIT Summ') - # self.g.plot() - - # ----------------------------------------------------------------------------- - # Flux - # ----------------------------------------------------------------------------- - 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]) - ax3 = fig.add_subplot(grid[0, 2]) - ax4 = fig.add_subplot(grid[1, 2]) - - self.plotFS(axs=[ax1, ax2, ax3, ax4]) - - # ----------------------------------------------------------------------------- - # 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]) - ax3 = fig.add_subplot(grid[2, 1]) - ax4 = fig.add_subplot(grid[:2, 1]) - ax5 = fig.add_subplot(grid[2, 2]) - ax6 = fig.add_subplot(grid[:2, 2]) - ax7 = fig.add_subplot(grid[2, 3]) - ax8 = fig.add_subplot(grid[:2, 3]) - ax9 = fig.add_subplot(grid[2, 4]) - ax10 = fig.add_subplot(grid[:2, 4]) - - plotCurrents(self, - axs=[ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10], zlims_thr=[-1, 1] - ) - - # ----------------------------------------------------------------------------- - # 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]) - ax3 = fig.add_subplot(grid[2, 1]) - ax4 = fig.add_subplot(grid[:2, 1]) - ax5 = fig.add_subplot(grid[2, 2]) - ax6 = fig.add_subplot(grid[:2, 2]) - ax7 = fig.add_subplot(grid[2, 3]) - ax8 = fig.add_subplot(grid[:2, 3]) - ax9 = fig.add_subplot(grid[2, 4]) - ax10 = fig.add_subplot(grid[:2, 4]) - - plotFields(self, - axs=[ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10], zlims_thr=[-1, 1] - ) - - # ----------------------------------------------------------------------------- - # Checks - # ----------------------------------------------------------------------------- - 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() - ax2 = fig.add_subplot(grid[1, 0]) - ax3 = fig.add_subplot(grid[:, 1]) - ax4 = fig.add_subplot(grid[:, 2]) - ax5 = fig.add_subplot(grid[:, 3]) - - plotChecks(self,axs=[ax1, ax1E, ax2, ax3, ax4, ax5]) - - # ----------------------------------------------------------------------------- - # Parameterization - # ----------------------------------------------------------------------------- - 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]) - ax3 = fig.add_subplot(grid[1, 1]) - ax4 = fig.add_subplot(grid[2, 1]) - ax5 = fig.add_subplot(grid[:, 2]) - - plotParameterization(self,axs=[ax1, ax2, ax3, ax4, ax5]) - - # ----------------------------------------------------------------------------- - # Plasma - # ----------------------------------------------------------------------------- - fig = self.fn.add_figure(label=extraLabel + "Plasma") - grid = plt.GridSpec(2, 4, hspace=0.3, wspace=0.3) - - ax_plasma = [ - fig.add_subplot(grid[0, 0]), - fig.add_subplot(grid[1, 0]), - fig.add_subplot(grid[0, 1]), - fig.add_subplot(grid[1, 1]), - fig.add_subplot(grid[0, 2]), - fig.add_subplot(grid[1, 2]), - fig.add_subplot(grid[0, 3]), - fig.add_subplot(grid[1, 3]), - ] - ax_plasma = self.plotPlasma(axs=ax_plasma, legendYN=not wasProvided) - - # ----------------------------------------------------------------------------- - # 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]) - ax3 = fig.add_subplot(grid[1, 0]) - ax4 = fig.add_subplot(grid[1, 1]) - - self.plotGeometry(axs=[ax1, ax2, ax3, ax4]) - - return ax_plasma - -def plotFS(self, axs=None, color="b", label=""): - if axs is None: - plt.ion() - fig, axs = plt.subplots(ncols=4) - - ax = axs[0] - self.plotFluxSurfaces( - ax=ax, fluxes=np.linspace(0, 1, 21), rhoPol=True, sqrt=False, color=color - ) - ax.plot(self.Rb, self.Yb, lw=1, c="r") - ax.set_title("Poloidal Flux") - ax.set_aspect("equal") - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - - ax = axs[1] - self.plotFluxSurfaces( - ax=ax, fluxes=np.linspace(0, 1, 21), rhoPol=False, sqrt=True, color=color - ) - ax.plot(self.Rb, self.Yb, lw=1, c="r") - ax.set_title("Sqrt Toroidal Flux") - ax.set_aspect("equal") - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - - ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["AuxQuantities"]["RHO"] - ax.plot(x, y, lw=2, ls="-", c=color, label=label) - ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) - - ax.set_xlabel("$\\Psi_n$ (PSI_NORM)") - ax.set_ylabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_xlim([0, 1]) - ax.set_ylim([0, 1]) - - ax = axs[3] - x = self.g["AuxQuantities"]["RHO"] - y = self.g["AuxQuantities"]["RHOp"] - ax.plot(x, y, lw=2, ls="-", c=color) - ax.plot([0, 1], [0, 1], ls="--", c="k", lw=0.5) - - ax.set_ylabel("$\\sqrt{\\Psi_n}$ (RHOp)") - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_xlim([0, 1]) - ax.set_ylim([0, 1]) - -def plotCurrents(self, axs=None, zlims_thr=[-1, 1]): - if axs is None: - plt.ion() - fig, axs = plt.subplots(ncols=10) - - ax = axs[0] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jr") * 1e-6 - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_ylabel("FSA $\\langle J\\rangle$ ($MA/m^2$)") - ax.set_xlim([0, 1]) - - zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] - zlims = GRAPHICStools.aroundZeroLims(zlims) - ax.set_ylim(zlims) - - ax = axs[1] - plot2Dquantity(self, - ax=ax, var="Jr", title="Radial Current Jr", zlims=zlims, factor=1e-6 - ) - - ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jz") * 1e-6 - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] - zlims = GRAPHICStools.aroundZeroLims(zlims) - ax.set_ylim(zlims) - - ax = axs[3] - plot2Dquantity(self, - ax=ax, var="Jz", title="Vertical Current Jz", zlims=zlims, factor=1e-6 - ) - - ax = axs[4] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jt") * 1e-6 - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] - zlims = GRAPHICStools.aroundZeroLims(zlims) - ax.set_ylim(zlims) - - ax = axs[5] - plot2Dquantity(self, - ax=ax, var="Jt", title="Toroidal Current Jt", zlims=zlims, factor=1e-6 - ) - - ax = axs[6] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jp") * 1e-6 - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] - zlims = GRAPHICStools.aroundZeroLims(zlims) - ax.set_ylim(zlims) - - ax = axs[7] - plot2Dquantity(self, - ax=ax, var="Jp", title="Poloidal Current Jp", zlims=zlims, factor=1e-6 - ) - - ax = axs[8] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Jpar") * 1e-6 - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] - zlims = GRAPHICStools.aroundZeroLims(zlims) - ax.set_ylim(zlims) - - ax = axs[9] - plot2Dquantity(self, - ax=ax, var="Jpar", title="Parallel Current Jpar", zlims=zlims, factor=1e-6 - ) - -def plotFields(self, axs=None, zlims_thr=[-1, 1]): - if axs is None: - plt.ion() - fig, axs = plt.subplots(ncols=10) - - ax = axs[0] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Br") - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_ylabel("FSA $\\langle B\\rangle$ ($T$)") - ax.set_xlim([0, 1]) - - zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] - zlims = GRAPHICStools.aroundZeroLims(zlims) - ax.set_ylim(zlims) - - ax = axs[1] - plot2Dquantity(self, - ax=ax, var="Br", title="Radial Field Br", zlims=zlims, titlebar="B ($T$)" - ) - - ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Bz") - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] - zlims = GRAPHICStools.aroundZeroLims(zlims) - ax.set_ylim(zlims) - - ax = axs[3] - plot2Dquantity(self, - ax=ax, var="Bz", title="Vertical Field Bz", zlims=zlims, titlebar="B ($T$)" - ) - - ax = axs[4] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Bt") - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - zlims = [y.min(), y.max()] - # zlims = [np.min([zlims_thr[0],y.min()]),np.max([zlims_thr[1],y.max()])] - # zlims = GRAPHICStools.aroundZeroLims(zlims) - ax.set_ylim(zlims) - - ax = axs[5] - plot2Dquantity(self, - ax=ax, var="Jt", title="Toroidal Field Bt", zlims=zlims, titlebar="B ($T$)" - ) - - ax = axs[6] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g.surfAvg("Bp") - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - zlims = [np.min([zlims_thr[0], y.min()]), np.max([zlims_thr[1], y.max()])] - zlims = GRAPHICStools.aroundZeroLims(zlims) - ax.set_ylim(zlims) - - ax = axs[7] - plot2Dquantity(self, - ax=ax, var="Bp", title="Poloidal Field Bp", zlims=zlims, titlebar="B ($T$)" - ) - - ax = axs[8] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["avg"]["Bp**2"] - ax.plot(x, y, lw=2, ls="-", c="r") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - ax.set_ylabel("$\\langle B_{\\theta}^2\\rangle$") - - ax = axs[9] - x = self.g["fluxSurfaces"]["midplane"]["R"] - y = self.g["fluxSurfaces"]["midplane"]["Bt"] - ax.plot(x, y, lw=2, ls="-", c="r", label="$B_{t}$") - y = self.g["fluxSurfaces"]["midplane"]["Bp"] - ax.plot(x, y, lw=2, ls="-", c="b", label="$B_{p}$") - y = self.g["fluxSurfaces"]["midplane"]["Bz"] - ax.plot(x, y, lw=2, ls="-", c="g", label="$B_{z}$") - y = self.g["fluxSurfaces"]["midplane"]["Br"] - ax.plot(x, y, lw=2, ls="-", c="m", label="$B_{r}$") - y = self.g["fluxSurfaces"]["geo"]["bunit"] - ax.plot(x, y, lw=2, ls="-", c="c", label="$B_{unit}$") - ax.set_xlabel("$R$ LF midplane") - ax.set_ylabel("$B$ (T)") - ax.legend() - -def plotChecks(self, axs=None): - if axs is None: - plt.ion() - fig, axs = plt.subplots(ncols=8) - - ax = axs[0] - x = self.g["AuxQuantities"]["PSI_NORM"] - y1 = self.Jt - ax.plot(x, np.abs(y1), lw=2, ls="-", c="b", label="$\\langle Jt\\rangle$") - zmax = y1.max() - zmin = y1.min() - y2 = self.Jt_fb - ax.plot(x, np.abs(y2), lw=2, ls="-", c="g", label="$\\langle Jt_{FB}\\rangle$") - - y3 = self.Jerror - ax.plot( - x, - y3, - lw=1, - ls="-", - c="r", - label="$|\\langle Jt\\rangle-\\langle Jt_{FB}\\rangle|$", - ) - - ax.set_ylabel("Current Density ($MA/m^2$)") - - axE = axs[1] - yErr = np.abs(self.Jerror / self.Jt) * 100.0 - axE.plot(x, yErr, lw=0.5, ls="--", c="b") - axE.set_ylim([0, 50]) - axE.set_ylabel("Relative Error (%)") - - ax.set_title("$|\\langle Jt\\rangle|$") - ax.set_ylim(bottom=0) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\Psi_n$") - ax.legend() - - ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y1 = self.g["FFPRIM"] - ax.plot(x, y1, lw=2, ls="-", c="r", label="$FF'$") - y2 = self.g["PPRIME"] * (4 * np.pi * 1e-7) - ax.plot(x, y2, lw=2, ls="-", c="b", label="$p'*\\mu_0$") - - ax.set_ylabel("") - ax.legend() - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - - ax = axs[3] - plot2Dquantity(self, - ax=ax, - var="Jt", - title="Toroidal Current Jt", - zlims=[zmin, zmax], - cmap="viridis", - factor=1e-6, - ) - - ax = axs[4] - plot2Dquantity(self, - ax=ax, - var="Jt_fb", - title="Toroidal Current Jt (FB)", - zlims=[zmin, zmax], - cmap="viridis", - factor=1e-6, - ) - - ax = axs[5] - z = ( - np.abs(self.g["AuxQuantities"]["Jt"] - self.g["AuxQuantities"]["Jt_fb"]) - * 1e-6 - ) - zmaxx = np.max([np.abs(zmax), np.abs(zmin)]) - plot2Dquantity(self, - ax=ax, - var=z, - title="Absolute Error", - zlims=[0, zmaxx], - cmap="viridis", - direct=True, - ) - -def plotParameterization(self, axs=None): - if axs is None: - plt.ion() - fig, axs = plt.subplots(ncols=5) - - ax = axs[0] - cs, csA = self.plotFluxSurfaces( - ax=ax, fluxes=np.linspace(0, 1, 21), rhoPol=True, sqrt=False - ) - # Boundary, axis and limiter - ax.plot(self.Rb, self.Yb, lw=1, c="r") - ax.plot(self.g["RMAXIS"], self.g["ZMAXIS"], "+", markersize=10, c="r") - ax.plot([self.Rmag], [self.Zmag], "o", markersize=5, c="m") - ax.plot([self.Rmajor], [self.Zmag], "+", markersize=10, c="k") - ax.plot(self.g["RLIM"], self.g["ZLIM"], lw=1, c="k") - - import matplotlib - - path = matplotlib.path.Path( - np.transpose(np.array([self.g["RLIM"], self.g["ZLIM"]])) - ) - patch = matplotlib.patches.PathPatch(path, facecolor="none") - ax.add_patch(patch) - # for col in cs.collections: - # col.set_clip_path(patch) - # for col in csA.collections: - # col.set_clip_path(patch) - - self.plotEnclosingBox(ax=ax) - - ax.set_aspect("equal") - ax.set_title("Poloidal Flux") - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - - ax = axs[1] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["geo"]["kap"] - ax.plot(x, y, label="$\\kappa$") - y = self.g["fluxSurfaces"]["geo"]["kapl"] - ax.plot(x, y, ls="--", label="$\\kappa_L$") - y = self.g["fluxSurfaces"]["geo"]["kapu"] - ax.plot(x, y, ls="--", label="$\\kappa_U$") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - ax.set_ylabel("Elongation $\\kappa$") - ax.legend() - - ax = axs[2] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["geo"]["delta"] - ax.plot(x, y, label="$\\delta$") - y = self.g["fluxSurfaces"]["geo"]["dell"] - ax.plot(x, y, ls="--", label="$\\delta_L$") - y = self.g["fluxSurfaces"]["geo"]["delu"] - ax.plot(x, y, ls="--", label="$\\delta_U$") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - ax.set_ylabel("Triangularity $\\delta$") - ax.legend() - - ax = axs[3] - x = self.g["AuxQuantities"]["PSI_NORM"] - y = self.g["fluxSurfaces"]["geo"]["zeta"] - ax.plot(x, y, label="$\\zeta$") - y = self.g["fluxSurfaces"]["geo"]["zetail"] - ax.plot(x, y, ls="--", label="$\\zeta_{IL}$") - y = self.g["fluxSurfaces"]["geo"]["zetaiu"] - ax.plot(x, y, ls="--", label="$\\zeta_{IU}$") - y = self.g["fluxSurfaces"]["geo"]["zetaol"] - ax.plot(x, y, ls="--", label="$\\zeta_{OL}$") - y = self.g["fluxSurfaces"]["geo"]["zetaou"] - ax.plot(x, y, ls="--", label="$\\zeta_{OU}$") - ax.set_xlabel("$\\Psi_n$") - ax.set_xlim([0, 1]) - ax.set_ylabel("Squareness $\\zeta$") - ax.legend() - - ax = axs[4] - ax.text( - 0.0, - 11.0, - "Rmajor = {0:.3f}m, Rmag = {1:.3f}m (Zmag = {2:.3f}m)".format( - self.Rmajor, self.Rmag, self.Zmag - ), - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - ) - ax.text( - 0.0, - 10.0, - f"a = {self.a:.3f}m, eps = {self.eps:.3f}", - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - ) - ax.text( - 0.0, - 9.0, - "kappa = {0:.3f} (kU = {1:.3f}, kL = {2:.3f})".format( - self.kappa, self.kappaU, self.kappaL - ), - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - ) - ax.text( - 0.0, - 8.0, - f" kappa95 = {self.kappa95:.3f}, kappa995 = {self.kappa995:.3f}", - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - ) - ax.text( - 0.0, - 7.0, - f" kappa_areal = {self.kappa_a:.3f}", - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - ) - ax.text( - 0.0, - 6.0, - "delta = {0:.3f} (dU = {1:.3f}, dL = {2:.3f})".format( - self.delta, self.deltaU, self.deltaL - ), - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - ) - ax.text( - 0.0, - 5.0, - f" delta95 = {self.delta95:.3f}, delta995 = {self.delta995:.3f}", - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - ) - ax.text( - 0.0, - 4.0, - f"zeta = {self.zeta:.3f}", - color="k", - fontsize=10, - fontweight="normal", - horizontalalignment="left", - verticalalignment="bottom", - rotation=0, - ) - - ax.set_ylim([0, 12]) - ax.set_xlim([-1, 1]) - - ax.set_axis_off() - -def plotPlasma(self, axs=None, legendYN=False, color="r", label=""): - if axs is None: - plt.ion() - fig, axs = plt.subplots(ncols=7) - - ax_plasma = axs - - ax = ax_plasma[0] - ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["PRES"] * 1e-6, - "-s", - c=color, - lw=2, - markersize=3, - label="geqdsk p", - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_ylim(bottom=0) - ax.set_ylabel("pressure (MPa)") - - ax = ax_plasma[1] - ax.plot( - self.g["AuxQuantities"]["RHO"], - -self.g["PPRIME"] * 1e-6, - c=color, - lw=2, - ls="-", - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_ylabel("pressure gradient -p' (MPa/[])") - ax.axhline(y=0.0, ls="--", lw=0.5, c="k") - - ax = ax_plasma[2] - ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FPOL"], c=color, lw=2, ls="-") - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_ylabel("$F = RB_{\\phi}$ (T*m)") - - ax = ax_plasma[3] - ax.plot(self.g["AuxQuantities"]["RHO"], self.g["FFPRIM"], c=color, lw=2, ls="-") - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_ylabel("FF' (T*m/[])") - ax.axhline(y=0.0, ls="--", lw=0.5, c="k") - - ax = ax_plasma[4] - ax.plot( - self.g["AuxQuantities"]["RHO"], - np.abs(self.g["QPSI"]), - "-s", - c=color, - lw=2, - markersize=3, - label=label + "geqdsk q", - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_ylim(bottom=0) - ax.set_ylabel("safety factor q") - ax.axhline(y=1.0, ls="--", lw=0.5, c="k") - - ax = ax_plasma[5] - ax.plot( - self.g["AuxQuantities"]["RHO"], - np.abs(self.g.surfAvg("Jt") * 1e-6), - "-s", - c=color, - lw=2, - markersize=3, - label=label + "geqdsk Jt", - ) - ax.plot( - self.g["AuxQuantities"]["RHO"], - np.abs(self.g.surfAvg("Jt_fb") * 1e-6), - "--o", - c=color, - lw=2, - markersize=3, - label=label + "geqdsk Jt(fb)", - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_ylabel("FSA toroidal current density ($MA/m^2$)") - ax.axhline(y=0.0, ls="--", lw=0.5, c="k") - - if legendYN: - ax.legend() - - ax = ax_plasma[6] - ax.plot( - self.g["fluxSurfaces"]["midplane"]["R"], - np.abs(self.g["fluxSurfaces"]["midplane"]["Bt"]), - "-s", - c=color, - lw=2, - markersize=3, - label=label + "geqdsk Bt", - ) - ax.plot( - self.g["fluxSurfaces"]["midplane"]["R"], - np.abs(self.g["fluxSurfaces"]["midplane"]["Bp"]), - "--o", - c=color, - lw=2, - markersize=3, - label=label + "geqdsk Bp", - ) - ax.set_xlabel("R (m) midplane") - ax.set_ylabel("Midplane fields (abs())") - - if legendYN: - ax.legend() - - return ax_plasma - -def plotGeometry(self, axs=None, color="r"): - if axs is None: - plt.ion() - fig, axs = plt.subplots(ncols=4) - - ax = axs[0] - ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["fluxSurfaces"]["geo"]["cxArea"], - "-", - c=color, - lw=2, - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_ylim(bottom=0) - ax.set_ylabel("CX Area ($m^2$)") - - ax = axs[1] - ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["fluxSurfaces"]["geo"]["surfArea"], - "-", - c=color, - lw=2, - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_ylim(bottom=0) - ax.set_ylabel("Surface Area ($m^2$)") - - ax = axs[2] - ax.plot( - self.g["AuxQuantities"]["RHO"], - self.g["fluxSurfaces"]["geo"]["vol"], - "-", - c=color, - lw=2, - ) - ax.set_xlim([0, 1]) - ax.set_xlabel("$\\sqrt{\\phi_n}$ (RHO)") - ax.set_ylim(bottom=0) - ax.set_ylabel("Volume ($m^3$)") - -def plotFluxSurfaces( - self, - ax=None, - fluxes=[1.0], - color="b", - alpha=1.0, - rhoPol=True, - sqrt=False, - lw=1, - lwB=2, - plot1=True, - label = '', -): - x = self.g["AuxQuantities"]["R"] - y = self.g["AuxQuantities"]["Z"] - - if rhoPol: - z = self.g["AuxQuantities"]["RHOpRZ"] - else: - z = self.g["AuxQuantities"]["RHORZ"] - - if not sqrt: - z = z**2 - - cs, csA = plotSurfaces( - x, y, z, fluxes=fluxes, ax=ax, color=color, alpha=alpha, lw=lw, lwB=lwB, plot1=plot1, label = label - ) - - return cs, csA - -def plotXpointEnvelope( - self, ax=None, color="b", alpha=1.0, rhoPol=True, sqrt=False -): - flx = 0.001 - fluxes = [1.0 - flx, 1.0 - flx / 2, 1.0, 1.0 + flx / 2, 1.0 + flx] - - self.plotFluxSurfaces( - fluxes=fluxes, ax=ax, color=color, alpha=alpha, rhoPol=rhoPol, sqrt=sqrt - ) - -def plot2Dquantity( - self, - ax=None, - var="Jr", - zlims=None, - title="Radial Current", - cmap="seismic", - direct=False, - titlebar="J ($MA/m^2$)", - factor=1.0, - includeSurfs=True, -): - if ax is None: - fig, ax = plt.subplots() - - x = self.g["AuxQuantities"]["R"] - y = self.g["AuxQuantities"]["Z"] - if not direct: - z = self.g["AuxQuantities"][var] * factor - else: - z = var - - if zlims is None: - am = np.amax(np.abs(z[:, :])) - ming = -am - maxg = am - else: - ming = zlims[0] - maxg = zlims[1] - levels = np.linspace(ming, maxg, 100) - colticks = np.linspace(ming, maxg, 5) - - cs = ax.contourf(x, y, z, levels=levels, extend="both", cmap=cmap) - - cbar = GRAPHICStools.addColorbarSubplot( - ax, - cs, - barfmt="%3.1f", - title=titlebar, - fontsize=10, - fontsizeTitle=8, - ylabel="", - ticks=colticks, - orientation="bottom", - drawedges=False, - padCB="25%", - ) - - ax.set_aspect("equal") - ax.set_title(title) - ax.set_xlabel("R (m)") - ax.set_ylabel("Z (m)") - - if includeSurfs: - self.plotFluxSurfaces( - ax=ax, - fluxes=np.linspace(0, 1, 6), - rhoPol=False, - sqrt=True, - color="k", - lw=0.5, - ) - - return cs - - -def plotSurfaces(R, Z, F, fluxes=[1.0], ax=None, color="b", alpha=1.0, lw=1, lwB=2, plot1=True, label = ''): - if ax is None: - fig, ax = plt.subplots() - - [Rg, Yg] = np.meshgrid(R, Z) - - if plot1: - csA = ax.contour( - Rg, Yg, F, 1000, levels=[1.0], colors=color, alpha=alpha, linewidths=lwB - ) - else: - csA = None - cs = ax.contour( - Rg, Yg, F, 1000, levels=fluxes, colors=color, alpha=alpha, linewidths=lw, label = label - ) - - return cs, csA - From b9990b09f5695511b4d59601ff2312616515fa3b Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 8 Oct 2025 19:12:48 +0200 Subject: [PATCH 364/385] kappa aeral must be always positive, regardless of sign of psi --- src/mitim_tools/gs_tools/GEQtools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index a4401efa..9e044f45 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -108,7 +108,7 @@ def derive(self, debug=False): # Core values vp = np.array(self.g.fluxsurfaces["Vprime"]).flatten() ir = np.array(self.g.fluxsurfaces["1/R"]).flatten() - self.cx_area = cumulative_trapezoid(vp * ir, self.g.derived["psi"], initial=0.0) + self.cx_area = abs(cumulative_trapezoid(vp * ir, self.g.derived["psi"], initial=0.0)) self.kappa_a = self.cx_area[-1] / (np.pi * self.a**2) self.kappa995 = np.interp( From 6c20cf254cf30d3e76957a6e4096b65dcba94b6d Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 3 Oct 2025 18:55:16 -0400 Subject: [PATCH 365/385] Until implementation is ready, do not require the new packages --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2a3fba09..040bfb3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,8 +44,8 @@ dependencies = [ "tensorflow", "f90nml", "pyyaml", - "megpy @ git+https://github.com/gsnoep/megpy.git@unstable", - "fibe @ git+https://github.com/aaronkho/fibe.git", + # "megpy @ git+https://github.com/gsnoep/megpy.git@unstable", + # "fibe @ git+https://github.com/aaronkho/fibe.git", "scikit-image", # Stricly not for MITIM, but good to have for pygacode #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models # "vmecpp", From c74bdac246b42b674f762d922ebae5a3d0deaf68 Mon Sep 17 00:00:00 2001 From: AudreySaltzman Date: Mon, 6 Oct 2025 17:59:10 -0400 Subject: [PATCH 366/385] Adding flag to choose whether EPEDbeat uses squareness --- src/mitim_modules/maestro/utils/EPEDbeat.py | 161 +++++++++++++----- src/mitim_tools/gacode_tools/PROFILEStools.py | 4 + 2 files changed, 120 insertions(+), 45 deletions(-) diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index f3dfeb3a..edb81b9e 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -63,6 +63,12 @@ def prepare( self.ptop_multiplier = ptop_multiplier self.TioverTe = TioverTe + # Whether EPED is going to be run with Zeta + if 'zeta_flag' in kwargs: + self.zeta_flag = kwargs['zeta_flag'] + else: + self.zeta_flag = False + self._inform() def run(self, **kwargs): @@ -89,6 +95,12 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F minimum_relative_change_in_x: minimum relative change in x to streach the core, otherwise it will keep the old core ''' + # Check to make sure using full EPED if running with squareness + if self.zeta_flag and self.nn is not None: + print('Warning: zeta_flag is not implemented for NN-based EPED, ignoring it', typeMsg='warning') + self.zeta_flag = False + + # ------------------------------------------------------- # Grab inputs from profiles_current # ------------------------------------------------------- @@ -119,32 +131,50 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F kappa995 = self.profiles_current.derived['kappa995'] delta995 = self.profiles_current.derived['delta995'] + zeta995 = self.profiles_current.derived['zeta995'] if self.zeta_flag else None BetaN = self.profiles_current.derived['BetaN_engineering'] Tesep_keV = self.profiles_current.profiles['te(keV)'][-1] nesep_20 = self.profiles_current.profiles['ne(10^19/m^3)'][-1]*0.1 - if 'kappa995' in self.__dict__ and self.kappa995 is not None: kappa995 = self.kappa995 - if 'delta995' in self.__dict__ and self.delta995 is not None: delta995 = self.delta995 - if "BetaN" in self.__dict__ and self.BetaN is not None: BetaN = self.BetaN - if "Tesep_keV" in self.__dict__ and self.Tesep_keV is not None: Tesep_keV = self.Tesep_keV - if "nesep_20" in self.__dict__ and self.nesep_20 is not None: nesep_20 = self.nesep_20 + if 'kappa995' in self.__dict__ and self.kappa995 is not None: kappa995 = self.kappa995 + if 'delta995' in self.__dict__ and self.delta995 is not None: delta995 = self.delta995 + if self.zeta_flag and 'zeta995' in self.__dict__ and self.zeta995 is not None: zeta995 = self.zeta995 + if "BetaN" in self.__dict__ and self.BetaN is not None: BetaN = self.BetaN + if "Tesep_keV" in self.__dict__ and self.Tesep_keV is not None: Tesep_keV = self.Tesep_keV + if "nesep_20" in self.__dict__ and self.nesep_20 is not None: nesep_20 = self.nesep_20 nesep_ratio = nesep_20 / neped_20 # Store evaluation - self.current_evaluation = { - 'Ip': np.abs(Ip), - 'Bt': np.abs(Bt), - 'R': np.abs(R), - 'a': np.abs(a), - 'kappa995': np.abs(kappa995), - 'delta995': np.abs(delta995), - 'neped_20': np.abs(neped_20), - 'BetaN': np.abs(BetaN), - 'zeff': np.abs(zeff), - 'Tesep_keV': np.abs(Tesep_keV), - 'nesep_ratio': np.abs(nesep_ratio), - } + if self.zeta_flag: + self.current_evaluation = { + 'Ip': np.abs(Ip), + 'Bt': np.abs(Bt), + 'R': np.abs(R), + 'a': np.abs(a), + 'kappa995': np.abs(kappa995), + 'delta995': np.abs(delta995), + 'neped_20': np.abs(neped_20), + 'BetaN': np.abs(BetaN), + 'zeff': np.abs(zeff), + 'Tesep_keV': np.abs(Tesep_keV), + 'nesep_ratio': np.abs(nesep_ratio), + 'zeta': np.abs(zeta995) + } + else: + self.current_evaluation = { + 'Ip': np.abs(Ip), + 'Bt': np.abs(Bt), + 'R': np.abs(R), + 'a': np.abs(a), + 'kappa995': np.abs(kappa995), + 'delta995': np.abs(delta995), + 'neped_20': np.abs(neped_20), + 'BetaN': np.abs(BetaN), + 'zeff': np.abs(zeff), + 'Tesep_keV': np.abs(Tesep_keV), + 'nesep_ratio': np.abs(nesep_ratio) + } # --- Sometimes we may need specific EPED inputs for key, value in self.corrections_set.items(): @@ -163,6 +193,7 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F print(f'\t\t- zeff: {self.current_evaluation["zeff"]:.2f}') print(f'\t\t- tesep: {self.current_evaluation["Tesep_keV"]:.3f} keV') print(f'\t\t- nesep_ratio: {self.current_evaluation["nesep_ratio"]:.2f}') + if self.zeta_flag: print(f'\t\t- zeta: {self.current_evaluation["zeta"]:.3f}') # ------------------------------------------------------- # Run NN @@ -174,19 +205,36 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F for i in range(loopBetaN): print(f'\t\t- BetaN: {BetaN:.2f}') - inputs_to_eped = ( - self.current_evaluation["Ip"], - self.current_evaluation["Bt"], - self.current_evaluation["R"], - self.current_evaluation["a"], - self.current_evaluation["kappa995"], - self.current_evaluation["delta995"], - self.current_evaluation["neped_20"]*10, - BetaN, - self.current_evaluation["zeff"], - self.current_evaluation["Tesep_keV"]* 1E3, - self.current_evaluation["nesep_ratio"] - ) + if self.zeta_flag: + inputs_to_eped = ( + self.current_evaluation["Ip"], + self.current_evaluation["Bt"], + self.current_evaluation["R"], + self.current_evaluation["a"], + self.current_evaluation["kappa995"], + self.current_evaluation["delta995"], + self.current_evaluation["neped_20"]*10, + BetaN, + self.current_evaluation["zeff"], + self.current_evaluation["Tesep_keV"]* 1E3, + self.current_evaluation["nesep_ratio"], + self.current_evaluation["zeta"] + ) + + else: + inputs_to_eped = ( + self.current_evaluation["Ip"], + self.current_evaluation["Bt"], + self.current_evaluation["R"], + self.current_evaluation["a"], + self.current_evaluation["kappa995"], + self.current_evaluation["delta995"], + self.current_evaluation["neped_20"]*10, + BetaN, + self.current_evaluation["zeff"], + self.current_evaluation["Tesep_keV"]* 1E3, + self.current_evaluation["nesep_ratio"] + ) # ------------------------------------------------------- # Give the option to override the ptop_kPa and wtop_psipol @@ -325,27 +373,50 @@ def _run(self, loopBetaN = 1, minimum_relative_change_in_x=0.005, store_scan = F return eped_results - def _run_full_eped(self, folder, Ip, Bt, R, a, kappa995, delta995, neped19, BetaN, zeff, Tesep_eV, nesep_ratio, nproc_per_run=64, cold_start=True): + def _run_full_eped(self, folder, Ip, Bt, R, a, kappa995, delta995, neped19, BetaN, zeff, Tesep_eV, nesep_ratio, *args, nproc_per_run=64, cold_start=True): ''' Run the full EPED code with the given inputs. Returns ptop_kPa and wtop_psipol. + If zeta is provided as an extra argument, use it; otherwise set zeta to zero. ''' + # Handle optional zeta parameter + if len(args) > 0: + zeta = args[0] + else: + zeta = 0.0 + eped = EPEDtools.EPED(folder=folder) - input_params = { - 'ip': Ip, - 'bt': Bt, - 'r': R, - 'a': a, - 'kappa': kappa995, - 'delta': delta995, - 'neped': neped19, - 'betan': BetaN, - 'zeffped': zeff, - 'nesep': nesep_ratio * neped19, - 'tesep': Tesep_eV - } + if int(zeta) != 0: + input_params = { + 'ip': Ip, + 'bt': Bt, + 'r': R, + 'a': a, + 'kappa': kappa995, + 'delta': delta995, + 'neped': neped19, + 'betan': BetaN, + 'zeffped': zeff, + 'nesep': nesep_ratio * neped19, + 'tesep': Tesep_eV, + 'zeta': zeta + } + else: + input_params = { + 'ip': Ip, + 'bt': Bt, + 'r': R, + 'a': a, + 'kappa': kappa995, + 'delta': delta995, + 'neped': neped19, + 'betan': BetaN, + 'zeffped': zeff, + 'nesep': nesep_ratio * neped19, + 'tesep': Tesep_eV + } eped.run( subfolder = 'case1', diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 862209ec..96f32867 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -258,6 +258,10 @@ def derive_geometry(self, n_theta_geo=1001, **kwargs): self.derived["delta95"] = np.interp(0.95, self.derived["psi_pol_n"], self.profiles["delta(-)"]) self.derived["delta995"] = np.interp(0.995, self.derived["psi_pol_n"], self.profiles["delta(-)"]) + + self.derived["zeta95"] = np.interp(0.95, self.derived["psi_pol_n"], self.profiles["zeta(-)"]) + + self.derived["zeta995"] = np.interp(0.995, self.derived["psi_pol_n"], self.profiles["zeta(-)"]) self.derived["kappa_a"] = self.derived["surfXS"][-1] / np.pi / self.derived["a"] ** 2 From 229ef9d8b7fed2ef0beeb573930d3f9a1a4801c8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 8 Oct 2025 20:02:55 +0200 Subject: [PATCH 367/385] Do not plot fast gradients in standard plot by default, pass a color if desired --- .../plasmastate_tools/utils/state_plotting.py | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py index 77678371..9559e42f 100644 --- a/src/mitim_tools/plasmastate_tools/utils/state_plotting.py +++ b/src/mitim_tools/plasmastate_tools/utils/state_plotting.py @@ -504,7 +504,7 @@ def plot_gradients( self, axs4, color="b", - fast_color='r', + fast_color=None, lw=1.0, label="", ls="-o", @@ -591,18 +591,19 @@ def plot_gradients( markersize=ms, alpha=alpha, ) - for i in range(len(self.Species)): - if self.Species[i]["S"] != "therm": - ax.plot( - xcoord[:ix], - self.derived["aLTi"][:ix, i], - ls, - c=fast_color, - lw=lw, - markersize=ms, - alpha=alpha, - label=self.Species[i]["N"], - ) + if fast_color is not None: + for i in range(len(self.Species)): + if self.Species[i]["S"] != "therm": + ax.plot( + xcoord[:ix], + self.derived["aLTi"][:ix, i], + ls, + c=fast_color, + lw=lw, + markersize=ms, + alpha=alpha, + label=self.Species[i]["N"], + ) ax.legend(loc="best", fontsize=7) ax = axs4[5] ax.plot( From 7c267194cafe4477e855970101b7c3c5c93db9c4 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 8 Oct 2025 20:12:00 +0200 Subject: [PATCH 368/385] Bring bach megpy dependency --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 040bfb3d..b2f5c67c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,7 @@ dependencies = [ "tensorflow", "f90nml", "pyyaml", - # "megpy @ git+https://github.com/gsnoep/megpy.git@unstable", + "megpy @ git+https://github.com/gsnoep/megpy.git@unstable", # "fibe @ git+https://github.com/aaronkho/fibe.git", "scikit-image", # Stricly not for MITIM, but good to have for pygacode #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models From 841015806fd173a03fb3162f565e6111b20c0ed8 Mon Sep 17 00:00:00 2001 From: AudreySaltzman Date: Wed, 8 Oct 2025 15:02:14 -0400 Subject: [PATCH 369/385] Bug fix for EPEDbeat with squareness --- src/mitim_modules/maestro/utils/EPEDbeat.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/utils/EPEDbeat.py b/src/mitim_modules/maestro/utils/EPEDbeat.py index edb81b9e..baefe24e 100644 --- a/src/mitim_modules/maestro/utils/EPEDbeat.py +++ b/src/mitim_modules/maestro/utils/EPEDbeat.py @@ -66,6 +66,7 @@ def prepare( # Whether EPED is going to be run with Zeta if 'zeta_flag' in kwargs: self.zeta_flag = kwargs['zeta_flag'] + print('zeta_flag set to True') else: self.zeta_flag = False @@ -383,12 +384,14 @@ def _run_full_eped(self, folder, Ip, Bt, R, a, kappa995, delta995, neped19, Beta # Handle optional zeta parameter if len(args) > 0: zeta = args[0] + print('Let of args > 0, using zeta =', zeta) else: zeta = 0.0 + print('No zeta provided, setting zeta = 0.0') eped = EPEDtools.EPED(folder=folder) - if int(zeta) != 0: + if len(args) > 0: input_params = { 'ip': Ip, 'bt': Bt, @@ -403,6 +406,7 @@ def _run_full_eped(self, folder, Ip, Bt, R, a, kappa995, delta995, neped19, Beta 'tesep': Tesep_eV, 'zeta': zeta } + print('_run_full_eped input_params with zeta:', input_params) else: input_params = { 'ip': Ip, @@ -417,6 +421,7 @@ def _run_full_eped(self, folder, Ip, Bt, R, a, kappa995, delta995, neped19, Beta 'nesep': nesep_ratio * neped19, 'tesep': Tesep_eV } + print('_run_full_eped input_params without zeta:', input_params) eped.run( subfolder = 'case1', From eaa3f85f9c26ace8df6154d20d89d9650911e338 Mon Sep 17 00:00:00 2001 From: Garud Snoep Date: Wed, 8 Oct 2025 16:43:17 -0400 Subject: [PATCH 370/385] MAESTROplot: updated plot_g_quantities() to be compatible with megpy migration --- src/mitim_modules/maestro/utils/MAESTROplot.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mitim_modules/maestro/utils/MAESTROplot.py b/src/mitim_modules/maestro/utils/MAESTROplot.py index 7a502caf..685c639b 100644 --- a/src/mitim_modules/maestro/utils/MAESTROplot.py +++ b/src/mitim_modules/maestro/utils/MAESTROplot.py @@ -365,7 +365,7 @@ def _special(ax,x): def plot_g_quantities(g, axs, color = 'b', lw = 1, ms = 0): g.plotFluxSurfaces(ax=axs[0], fluxes=np.linspace(0, 1, 21), rhoPol=False, sqrt=True, color=color,lwB=lw*3, lw = lw,label='Initial geqdsk') - axs[3].plot(g.g['RHOVN'], g.g['PRES']*1E-6, '-o', markersize=ms, lw = lw, label='Initial geqdsk', color=color) - axs[4].plot(g.g['RHOVN'], g.g['QPSI'], '-o', markersize=ms, lw = lw, label='Initial geqdsk', color=color) + axs[3].plot(g.g.derived['rho_tor'], g.g.raw['pres']*1E-6, '-o', markersize=ms, lw = lw, label='Initial geqdsk', color=color) + axs[4].plot(g.g.derived['rho_tor'], g.g.raw['qpsi'], '-o', markersize=ms, lw = lw, label='Initial geqdsk', color=color) From a6784400bd5fc858a5251069fc3456ca8cce9b27 Mon Sep 17 00:00:00 2001 From: pabloprf Date: Wed, 8 Oct 2025 19:43:30 -0400 Subject: [PATCH 371/385] Bug fix MAESTRO: make _to_mxh routine work even if flux surface is one point, at the center --- src/mitim_tools/gs_tools/GEQtools.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/mitim_tools/gs_tools/GEQtools.py b/src/mitim_tools/gs_tools/GEQtools.py index 9e044f45..a4f7479e 100644 --- a/src/mitim_tools/gs_tools/GEQtools.py +++ b/src/mitim_tools/gs_tools/GEQtools.py @@ -226,6 +226,20 @@ def get_MXH_coeff_new(self, n_coeff=7, plotYN=False): else: Rf, Zf = self.g.fluxsurfaces["R"][flux], self.g.fluxsurfaces["Z"][flux] + # To avoid the following code to fail if only one point is found + if Rf.shape[0] == 1: + + min_value = 1E-7 + + kappa.append(min_value) + rmin.append(min_value) + rmaj.append(Rf[0]) + zmag.append(Zf[0]) + + sn.append(np.ones(n_coeff)*min_value) + cn.append(np.ones(n_coeff)*min_value) + continue + # Perform the MXH decompositionusing the MITIM surface class surfaces = mitim_flux_surfaces() surfaces.reconstruct_from_RZ(Rf,Zf) @@ -538,7 +552,7 @@ def _to_miller(self): for i in range(self.R0.shape[0]): try: Ri, Zi, zeta_uo = find_squareness_points(self.R[i,:], self.Z[i,:]) - except AttributeError: + except: zeta_uo = np.nan self.zeta[i] = zeta_uo From 9a91dafcdf9e1bc5a0145f010b06be7bb9cb1ea3 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Fri, 10 Oct 2025 01:02:39 +0200 Subject: [PATCH 372/385] Bug fix: FSA of Bt^2 and Bp^2 were swapped; and their calculation missed some radial dependence. Expect betas to change by <2% --- src/mitim_tools/gacode_tools/PROFILEStools.py | 60 +++++++------------ 1 file changed, 22 insertions(+), 38 deletions(-) diff --git a/src/mitim_tools/gacode_tools/PROFILEStools.py b/src/mitim_tools/gacode_tools/PROFILEStools.py index 96f32867..90d80427 100644 --- a/src/mitim_tools/gacode_tools/PROFILEStools.py +++ b/src/mitim_tools/gacode_tools/PROFILEStools.py @@ -511,8 +511,6 @@ def plot_plasma_boundary(self, ax=None, color="b"): # Set equal aspect ratio ax.set_aspect("equal") - - def calculateGeometricFactors(profiles, n_theta=1001): # ---------------------------------------- @@ -529,6 +527,8 @@ def calculateGeometricFactors(profiles, n_theta=1001): q = profiles.profiles["q(-)"] shape_coeffs = profiles.shape_cos + profiles.shape_sin + + signb = np.sign(profiles.profiles['torfluxa(Wb/radian)'][0]) # ---------------------------------------- # Derivatives as defined in expro_util.f90 @@ -590,6 +590,7 @@ def calculateGeometricFactors(profiles, n_theta=1001): dzmag, dRmag, q, + geo_signb_in=signb, n_theta=n_theta, ) @@ -619,18 +620,15 @@ def volp_surf_geo_vectorized( geo_dzmag_in, geo_drmaj_in, geo_q_in, + geo_signb_in = 1.0, n_theta=1001): """ Completety from f2py/geo/geo.f90 """ - geo_rmin_in = geo_rmin_in.clip( - 1e-10 - ) # To avoid problems at 0 (Implemented by PRF, not sure how TGYRO deals with this) - + geo_rmin_in = geo_rmin_in.clip(1e-10) # To avoid problems at 0 (Implemented by PRF, not sure how TGYRO deals with this) geo_q_in = geo_q_in.clip(1e-2) # To avoid problems at 0 with some geqdsk files that are corrupted... - [ geo_shape_cos0_in, geo_shape_cos1_in, @@ -665,8 +663,6 @@ def volp_surf_geo_vectorized( geo_shape_s_sin6_in, ] = np.array(cos_sin_s).astype(float).T - geo_signb_in = 1.0 - geov_theta = np.zeros((n_theta,geo_rmin_in.shape[0])) geov_bigr = np.zeros((n_theta,geo_rmin_in.shape[0])) geov_bigr_r = np.zeros((n_theta,geo_rmin_in.shape[0])) @@ -691,7 +687,7 @@ def volp_surf_geo_vectorized( #! Generalized Miller-type parameterization #!----------------------------------------- - theta = -0.5 * pi_2 + (i - 1) * d_theta + theta = -0.5 * pi_2 + i * d_theta geov_theta[i] = theta @@ -709,7 +705,6 @@ def volp_surf_geo_vectorized( + geo_shape_cos4_in * np.cos(4 * theta) + geo_shape_cos5_in * np.cos(5 * theta) + geo_shape_cos6_in * np.cos(6 * theta) - + geo_shape_sin3_in * np.sin(3 * theta) + x * np.sin(theta) - geo_zeta_in * np.sin(2 * theta) + geo_shape_sin3_in * np.sin(3 * theta) @@ -791,16 +786,11 @@ def volp_surf_geo_vectorized( bigz[i] = geo_zmag_in + geo_kappa_in * geo_rmin_in * np.sin(a) bigz_r[i] = geo_dzmag_in + geo_kappa_in * (1.0 + geo_s_kappa_in) * np.sin(a) bigz_t[i] = geo_kappa_in * geo_rmin_in * np.cos(a) * a_t - bigz_tt = ( - -geo_kappa_in * geo_rmin_in * np.sin(a) * a_t**2 - + geo_kappa_in * geo_rmin_in * np.cos(a) * a_tt - ) + bigz_tt = (-geo_kappa_in * geo_rmin_in * np.sin(a) * a_t**2+ geo_kappa_in * geo_rmin_in * np.cos(a) * a_tt) g_tt = geov_bigr_t[i] ** 2 + bigz_t[i] ** 2 - geov_jac_r[i] = geov_bigr[i] * ( - geov_bigr_r[i] * bigz_t[i] - geov_bigr_t[i] * bigz_r[i] - ) + geov_jac_r[i] = geov_bigr[i] * (geov_bigr_r[i] * bigz_t[i] - geov_bigr_t[i] * bigz_r[i]) geov_grad_r[i] = geov_bigr[i] * np.sqrt(g_tt) / geov_jac_r[i] @@ -814,13 +804,11 @@ def volp_surf_geo_vectorized( geov_l_r[i] = bigz_l[i] * bigz_r[i] + bigr_l[i] * geov_bigr_r[i] - geov_nsin[i] = ( - geov_bigr_r[i] * geov_bigr_t[i] + bigz_r[i] * bigz_t[i] - ) / geov_l_t[i] + geov_nsin[i] = (geov_bigr_r[i] * geov_bigr_t[i] + bigz_r[i] * bigz_t[i]) / geov_l_t[i] c = 0.0 - for i in range(n_theta): - c = c + geov_l_t[i] / (geov_bigr[i] * geov_grad_r[i]) + for i in range(n_theta - 1): + c += geov_l_t[i] / (geov_bigr[i] * geov_grad_r[i]) f = geo_rmin_in / (c * d_theta / pi_2) @@ -836,20 +824,15 @@ def volp_surf_geo_vectorized( geo_surf = geo_surf + geov_l_t[i] * geov_bigr[i] geo_surf = pi_2 * geo_surf * d_theta - # ----- - c = 0.0 - for i in range(n_theta - 1): - c = c + geov_l_t[i] / (geov_bigr[i] * geov_grad_r[i]) - f = geo_rmin_in / (c * d_theta / pi_2) - geov_b = np.zeros((n_theta,geo_rmin_in.shape[0])) + geov_bp = np.zeros((n_theta,geo_rmin_in.shape[0])) geov_g_theta = np.zeros((n_theta,geo_rmin_in.shape[0])) geov_bt = np.zeros((n_theta,geo_rmin_in.shape[0])) for i in range(n_theta): geov_bt[i] = f / geov_bigr[i] - geov_bp = (geo_rmin_in / geo_q_in) * geov_grad_r[i] / geov_bigr[i] + geov_bp[i] = (geo_rmin_in / geo_q_in) * geov_grad_r[i] / geov_bigr[i] - geov_b[i] = geo_signb_in * (geov_bt[i] ** 2 + geov_bp**2) ** 0.5 + geov_b[i] = geo_signb_in * (geov_bt[i] ** 2 + geov_bp[i] ** 2) ** 0.5 geov_g_theta[i] = ( geov_bigr[i] * geov_b[i] @@ -866,7 +849,8 @@ def volp_surf_geo_vectorized( z = (x0 - x1) / dx if i2 == n_theta: i2 -= 1 - bt_geo0 = geov_bt[i1] + (geov_bt[i2] - geov_bt[i1]) * z + # bt_geo0 = geov_bt[i1] + (geov_bt[i2] - geov_bt[i1]) * z + bt_geo0 = geov_bt[n_theta // 2] denom = 0 for i in range(n_theta - 1): @@ -879,21 +863,21 @@ def volp_surf_geo_vectorized( + geov_grad_r[i] * geov_g_theta[i] / geov_b[i] / denom ) - geo_fluxsurfave__bp2 = 0 + geo_fluxsurfave_bp2 = 0 for i in range(n_theta - 1): - geo_fluxsurfave__bp2 = ( - geo_fluxsurfave__bp2 - + geov_bt[i] ** 2 * geov_g_theta[i] / geov_b[i] / denom + geo_fluxsurfave_bp2 = ( + geo_fluxsurfave_bp2 + + geov_bp[i] ** 2 * geov_g_theta[i] / geov_b[i] / denom ) geo_fluxsurfave_bt2 = 0 for i in range(n_theta - 1): geo_fluxsurfave_bt2 = ( geo_fluxsurfave_bt2 - + geov_bp ** 2 * geov_g_theta[i] / geov_b[i] / denom + + geov_bt[i] ** 2 * geov_g_theta[i] / geov_b[i] / denom ) - return geo_volume_prime, geo_surf, geo_fluxsurfave_grad_r, geo_fluxsurfave__bp2, geo_fluxsurfave_bt2, bt_geo0 + return geo_volume_prime, geo_surf, geo_fluxsurfave_grad_r, geo_fluxsurfave_bp2, geo_fluxsurfave_bt2, bt_geo0 def xsec_area_RZ(R,Z): # calculates the cross-sectional area of the plasma for each flux surface From fdf330f52b817c400ce6598f8290479684fd88d2 Mon Sep 17 00:00:00 2001 From: Audrey Saltzman Date: Fri, 10 Oct 2025 09:51:10 -0400 Subject: [PATCH 373/385] Make the real frequency not normalized to ky --- src/mitim_tools/gacode_tools/utils/GACODEplotting.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/mitim_tools/gacode_tools/utils/GACODEplotting.py b/src/mitim_tools/gacode_tools/utils/GACODEplotting.py index 0777a2be..78520f31 100644 --- a/src/mitim_tools/gacode_tools/utils/GACODEplotting.py +++ b/src/mitim_tools/gacode_tools/utils/GACODEplotting.py @@ -89,13 +89,15 @@ def plotTGLFspectrum( elif limity: ax.set_ylim(bottom=0) + freq_coeff = 0 # The real frequencies should not be normalized by ky + if freq is not None and type(axs) == list: plot_spec( axF, kys, freq, markers=markers, - coeff=coeff, + coeff=freq_coeff, c=c, lw=lw, label=label, @@ -105,9 +107,9 @@ def plotTGLFspectrum( ) if ylabel: - axF.set_ylabel(decorateLabel("$\\omega$", coeff)) + axF.set_ylabel(decorateLabel("$\\omega$", freq_coeff)) - if coeff == 0: + if freq_coeff == 0: axF.set_yscale("symlog", linthresh=thr_symlog) elif limity: axF.set_ylim(bottom=0) From 04ef2c6ed290dc9aeaf845a22e4e6e298a709773 Mon Sep 17 00:00:00 2001 From: aaronkho Date: Fri, 10 Oct 2025 11:01:07 -0400 Subject: [PATCH 374/385] Added FiBE equilibrium initialization method --- .../maestro/utils/MAESTRObeat.py | 66 ++++++++++++++++++- 1 file changed, 65 insertions(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/utils/MAESTRObeat.py b/src/mitim_modules/maestro/utils/MAESTRObeat.py index 7b1c281d..ae9936ce 100644 --- a/src/mitim_modules/maestro/utils/MAESTRObeat.py +++ b/src/mitim_modules/maestro/utils/MAESTRObeat.py @@ -40,6 +40,8 @@ def define_initializer(self, initializer): self.initialize = initializer_from_previous(self) elif initializer == 'freegs': self.initialize = initializer_from_freegs(self) + elif initializer == 'fibe': + self.initialize = initializer_from_fibe(self) elif initializer == 'geqdsk': self.initialize = initializer_from_geqdsk(self) elif initializer == 'profiles': @@ -269,6 +271,68 @@ def __call__(self, # Call the geqdsk initializer super().__call__(geqdsk_file = self.folder / 'freegs.geqdsk',**kwargs_geqdsk) +# -------------------------------------------------------------------------------------------- +# Initializer from FiBE: create the equilibrium, convert to geqdsk and call the geqdsk initializer +# -------------------------------------------------------------------------------------------- + +class initializer_from_fibe(initializer_from_geqdsk): + ''' + Idea is to write geqdsk and then call the geqdsk initializer + ''' + def __init__(self, beat_instance, label = 'fibe'): + super().__init__(beat_instance, label = label) + + def __call__(self, + R, + a, + kappa_sep, + delta_sep, + zeta_sep, + z0, + p0_MPa = 1.0, + Ip_MA = 1.0, + B_T = 5.4, + **kwargs_geqdsk + ): + + p0 = p0_MPa * 1.0e6 + Ip = Ip_MA * 1.0e6 + # If profiles exist, substitute the pressure and density guesses by something better (not perfect though, no ions) + if ('ne' in kwargs_geqdsk.get('profiles_insert',{})) and ('Te' in kwargs_geqdsk.get('profiles_insert',{})): + print('\t- Using ne profile instead of the ne0 guess') + ne0_20 = kwargs_geqdsk['profiles_insert']['ne'][1][0] + print('\t- Using Te profile for a better estimation of pressure, instead of the p0 guess') + Te0_keV = kwargs_geqdsk['profiles_insert']['Te'][1][0] + p0 = 2 * (Te0_keV*1E3) * 1.602176634E-19 * (ne0_20 * 1E20) + # If betaN provided, use it to estimate the pressure + elif 'BetaN' in kwargs_geqdsk: + print('\t- Using BetaN for a better estimation of pressure, instead of the p0 guess') + pvol_MPa = ( Ip_MA / (a * B_T) ) * (B_T ** 2 / (2 * 4 * np.pi * 1e-7)) / 1e6 * kwargs_geqdsk['BetaN'] * 1E-2 + p0 = pvol_MPa * 3.0 * 1.0e6 + + # Run FiBE to generate equilibrium + from fibe import FixedBoundaryEquilibrium + eq = FixedBoundaryEquilibrium() + eq.define_grid_and_boundary_from_mxh( + nr=129, + nz=129, + rgeo=R, + zgeo=z0, + rminor=a, + kappa=kappa_sep, + cos_coeffs=[0.0, 0.0, 0.0], + sin_coeffs=[0.0, np.arcsin(delta_sep), -zeta_sep) + ) + eq.initialize_profiles_with_minimal_input(p0, Ip, B_T) + eq.initialize_psi() + eq.solve_psi() + + # Convert to geqdsk and write it to initialization folder + eq.to_geqdsk(self.folder / 'fibe.geqdsk') + + # Call the geqdsk initializer + super().__call__(geqdsk_file = self.folder / 'fibe.geqdsk',**kwargs_geqdsk) + # -------------------------------------------------------------------------------------------- # [Generic] Profile creator: Insert profiles # -------------------------------------------------------------------------------------------- @@ -488,4 +552,4 @@ def _inform_save(self, eped_results = None): if eped_results is None: eped_results = np.load(beat_eped_for_save.folder_output / 'eped_results.npy', allow_pickle=True).item() - beat_eped_for_save._inform_save(eped_results) \ No newline at end of file + beat_eped_for_save._inform_save(eped_results) From 30535b09b70eca02c53cce21066f69c76774bb92 Mon Sep 17 00:00:00 2001 From: aaronkho Date: Fri, 10 Oct 2025 11:03:04 -0400 Subject: [PATCH 375/385] Re-added FiBE installation dependency --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b2f5c67c..b2f9d044 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ dependencies = [ "f90nml", "pyyaml", "megpy @ git+https://github.com/gsnoep/megpy.git@unstable", - # "fibe @ git+https://github.com/aaronkho/fibe.git", + "fibe @ git+https://github.com/aaronkho/fibe.git@geqdsk_update", "scikit-image", # Stricly not for MITIM, but good to have for pygacode #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models # "vmecpp", From 122a597bceea3e2e31cc32738fc22331ff01e423 Mon Sep 17 00:00:00 2001 From: aaronkho Date: Fri, 10 Oct 2025 11:26:22 -0400 Subject: [PATCH 376/385] Switched fibe dependency to main branch due to merge --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b2f9d044..2a3fba09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ dependencies = [ "f90nml", "pyyaml", "megpy @ git+https://github.com/gsnoep/megpy.git@unstable", - "fibe @ git+https://github.com/aaronkho/fibe.git@geqdsk_update", + "fibe @ git+https://github.com/aaronkho/fibe.git", "scikit-image", # Stricly not for MITIM, but good to have for pygacode #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models # "vmecpp", From 8d5ad5e20e6916b2d07fd0f93b60b6ae6c1fa3b1 Mon Sep 17 00:00:00 2001 From: Audrey Saltzman Date: Fri, 10 Oct 2025 17:10:20 -0400 Subject: [PATCH 377/385] Bug fix to run MAESTRObeat.py, uneven ) --- src/mitim_modules/maestro/utils/MAESTRObeat.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/mitim_modules/maestro/utils/MAESTRObeat.py b/src/mitim_modules/maestro/utils/MAESTRObeat.py index ae9936ce..7cd8d49f 100644 --- a/src/mitim_modules/maestro/utils/MAESTRObeat.py +++ b/src/mitim_modules/maestro/utils/MAESTRObeat.py @@ -321,8 +321,7 @@ def __call__(self, rminor=a, kappa=kappa_sep, cos_coeffs=[0.0, 0.0, 0.0], - sin_coeffs=[0.0, np.arcsin(delta_sep), -zeta_sep) - ) + sin_coeffs=[0.0, np.arcsin(delta_sep), -zeta_sep]) eq.initialize_profiles_with_minimal_input(p0, Ip, B_T) eq.initialize_psi() eq.solve_psi() From 476bdab3ceb532103b73ff2353663bc6ccede457 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Sat, 11 Oct 2025 11:41:24 +0200 Subject: [PATCH 378/385] Logic of keep_files had weird wording, now fixed to avoid user confusion --- src/mitim_modules/powertorch/physics_models/transport_tglf.py | 4 ++-- templates/namelist.portals.yaml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mitim_modules/powertorch/physics_models/transport_tglf.py b/src/mitim_modules/powertorch/physics_models/transport_tglf.py index 2e9ae17f..cf0c14cf 100644 --- a/src/mitim_modules/powertorch/physics_models/transport_tglf.py +++ b/src/mitim_modules/powertorch/physics_models/transport_tglf.py @@ -60,7 +60,7 @@ def _evaluate_tglf(self, pass_info = True): "minutes": 2, }, attempts_execution=2, - only_minimal_files=keep_tglf_files in ['minimal'], + only_minimal_files=keep_tglf_files in ['none'], **simulation_options["run"] ) @@ -117,7 +117,7 @@ def _evaluate_tglf(self, pass_info = True): extra_name=self.name, cores_per_tglf_instance=cores_per_tglf_instance, Qi_includes_fast=Qi_includes_fast, - only_minimal_files=keep_tglf_files in ['minimal', 'base'], + only_minimal_files=keep_tglf_files in ['none', 'base'], reuse_scan_ball_file=reuse_scan_ball_file, **simulation_options["run"] ) diff --git a/templates/namelist.portals.yaml b/templates/namelist.portals.yaml index 1fa03cd5..7a7597e6 100644 --- a/templates/namelist.portals.yaml +++ b/templates/namelist.portals.yaml @@ -136,7 +136,7 @@ transport: # [EXPERIMENTAL] If True, store previous evaluations and reuse them if they are within the delta of all inputs (to capture combinations) reuse_scan_ball: false - # Files to keep from simulation (minimal: only retrieve minimal files always, none: always minimal files, base: retrieve all files) + # Files to keep from simulation (none: all runs will only keep minimal files; base: minimal files for scans only; all: retrieve all files) keep_files: "base" # Number of cores to use per TGLF instance From 438d9b53641f336b1be7878ad45274aaa001cdf2 Mon Sep 17 00:00:00 2001 From: AudreySaltzman Date: Wed, 15 Oct 2025 11:04:26 -0400 Subject: [PATCH 379/385] Adjustments to make FIBE initializer work with MAESTRO --- src/mitim_modules/maestro/scripts/run_maestro.py | 8 ++++++++ src/mitim_modules/maestro/utils/MAESTRObeat.py | 4 ++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/src/mitim_modules/maestro/scripts/run_maestro.py b/src/mitim_modules/maestro/scripts/run_maestro.py index c96e0337..a3934e56 100644 --- a/src/mitim_modules/maestro/scripts/run_maestro.py +++ b/src/mitim_modules/maestro/scripts/run_maestro.py @@ -88,6 +88,14 @@ def parse_maestro_nml(file_path): delta_sep = maestro_namelist["machine"]["separatrix"]["parameters"]["delta_sep"] n_mxh = maestro_namelist["machine"]["separatrix"]["parameters"]["n_mxh"] geometry = {'R': R, 'a': a, 'kappa_sep': kappa_sep, 'delta_sep': delta_sep, 'zeta_sep': 0.0, 'z0': 0.0, 'coeffs_MXH' : n_mxh} + elif separatrix_type == 'fibe': + R = maestro_namelist["machine"]["separatrix"]["parameters"]["R"] + a = maestro_namelist["machine"]["separatrix"]["parameters"]["a"] + kappa_sep = maestro_namelist["machine"]["separatrix"]["parameters"]["kappa_sep"] + delta_sep = maestro_namelist["machine"]["separatrix"]["parameters"]["delta_sep"] + zeta_sep = maestro_namelist["machine"]["separatrix"]["parameters"]["zeta_sep"] + n_mxh = maestro_namelist["machine"]["separatrix"]["parameters"]["n_mxh"] + geometry = {'R': R, 'a': a, 'kappa_sep': kappa_sep, 'delta_sep': delta_sep, 'zeta_sep': 0.0, 'z0': 0.0, 'coeffs_MXH' : n_mxh} elif separatrix_type == "geqdsk": # Initialize geometry from geqdsk file geqdsk_file = maestro_namelist["machine"]["separatrix"]["parameters"]["geqdsk_file"] diff --git a/src/mitim_modules/maestro/utils/MAESTRObeat.py b/src/mitim_modules/maestro/utils/MAESTRObeat.py index 7cd8d49f..506b225c 100644 --- a/src/mitim_modules/maestro/utils/MAESTRObeat.py +++ b/src/mitim_modules/maestro/utils/MAESTRObeat.py @@ -313,7 +313,7 @@ def __call__(self, # Run FiBE to generate equilibrium from fibe import FixedBoundaryEquilibrium eq = FixedBoundaryEquilibrium() - eq.define_grid_and_boundary_from_mxh( + eq.define_grid_and_boundary_with_mxh( nr=129, nz=129, rgeo=R, @@ -327,7 +327,7 @@ def __call__(self, eq.solve_psi() # Convert to geqdsk and write it to initialization folder - eq.to_geqdsk(self.folder / 'fibe.geqdsk') + eq.to_geqdsk(str(self.folder.absolute()) + '/fibe.geqdsk') # Call the geqdsk initializer super().__call__(geqdsk_file = self.folder / 'fibe.geqdsk',**kwargs_geqdsk) From 29938a50e5468a265b860290eff1d0d95a2bfa7b Mon Sep 17 00:00:00 2001 From: AudreySaltzman Date: Wed, 15 Oct 2025 12:00:27 -0400 Subject: [PATCH 380/385] Fix for pathlib compatibility --- src/mitim_modules/maestro/utils/MAESTRObeat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mitim_modules/maestro/utils/MAESTRObeat.py b/src/mitim_modules/maestro/utils/MAESTRObeat.py index 506b225c..4991b4a1 100644 --- a/src/mitim_modules/maestro/utils/MAESTRObeat.py +++ b/src/mitim_modules/maestro/utils/MAESTRObeat.py @@ -327,7 +327,7 @@ def __call__(self, eq.solve_psi() # Convert to geqdsk and write it to initialization folder - eq.to_geqdsk(str(self.folder.absolute()) + '/fibe.geqdsk') + eq.to_geqdsk(str(self.folder / 'fibe.geqdsk')) # Call the geqdsk initializer super().__call__(geqdsk_file = self.folder / 'fibe.geqdsk',**kwargs_geqdsk) From 10949cc378302dcac231cc13d344e2f1e6f14bc4 Mon Sep 17 00:00:00 2001 From: aaronkho Date: Wed, 15 Oct 2025 16:24:03 -0400 Subject: [PATCH 381/385] Updated megpy and fibe requirements to PyPI versions --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2a3fba09..aa7981e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,8 +44,8 @@ dependencies = [ "tensorflow", "f90nml", "pyyaml", - "megpy @ git+https://github.com/gsnoep/megpy.git@unstable", - "fibe @ git+https://github.com/aaronkho/fibe.git", + "megpy>=2.0", + "fibe>=0.3", "scikit-image", # Stricly not for MITIM, but good to have for pygacode #"onnx2pytorch", # Stricly not for MITIM, but good to use ONNX models # "vmecpp", From 8fe2b5186cf8f7879ca0ac83b568a5e0a60edc72 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 15 Oct 2025 16:50:41 -0400 Subject: [PATCH 382/385] Recover VITALS (after documentation followup) --- src/mitim_modules/vitals/VITALSmain.py | 16 ++++++---------- src/mitim_tools/gacode_tools/TGLFtools.py | 4 ++++ src/mitim_tools/gacode_tools/TGYROtools.py | 4 +--- 3 files changed, 11 insertions(+), 13 deletions(-) diff --git a/src/mitim_modules/vitals/VITALSmain.py b/src/mitim_modules/vitals/VITALSmain.py index e265ce58..6983276f 100644 --- a/src/mitim_modules/vitals/VITALSmain.py +++ b/src/mitim_modules/vitals/VITALSmain.py @@ -74,6 +74,8 @@ def prep( if not classLoaded: with open(tglf_class_file, "rb") as f: tglf_read = pickle.load(f) + tglf_tmp = TGLFtools.TGLF() + tglf_read.run_specifications = tglf_tmp.run_specifications self.tglf = TGLFtools.TGLF(alreadyRun=tglf_read) self.tglf.FolderGACODE, self.tglf.rhos = self.folder, [rho] @@ -214,16 +216,10 @@ def run(self, paramsfile, resultsfile): value = tglf.results["tglf1"]["output"][0].neTeSpectrum_level dictOFs[iquant]["value"] = value - dictOFs[iquant]["error"] = np.abs( - dictOFs[iquant]["value"] * self.VITALSparameters["rel_error"] - ) + dictOFs[iquant]["error"] = np.abs(dictOFs[iquant]["value"] * self.VITALSparameters["rel_error"]) else: - dictOFs[iquant]["value"] = self.VITALSparameters["experimentalVals"][ - iquant[:-4] - ] - dictOFs[iquant]["error"] = self.VITALSparameters["std_deviation"][ - iquant[:-4] - ] + dictOFs[iquant]["value"] = self.VITALSparameters["experimentalVals"][iquant[:-4]] + dictOFs[iquant]["error"] = self.VITALSparameters["std_deviation"][iquant[:-4]] # Write stuff self.write(dictOFs, resultsfile) @@ -287,7 +283,7 @@ def runTGLF( numSim = self.folder.name - variation = TGLFtools.completeVariation(variation, tglf.inputs_files[tglf.rhos[0]]) + variation = TGLFtools.completeVariation_TGLF(variation, tglf.inputs_files[tglf.rhos[0]]) extraOptions = self.TGLFparameters["extraOptions"] multipliers = {} diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 002dc9c9..32c1daff 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -362,6 +362,8 @@ def prepare_for_save_TGLF(self): if "convolution_fun_fluct" in tglf_copy.results[label]: tglf_copy.results[label]["convolution_fun_fluct"] = None + del tglf_copy.run_specifications + return tglf_copy def save_pkl(self, file): @@ -407,6 +409,8 @@ def prep_using_tgyro( # If inputgacode is already a PROFILEStools object, just use it profiles = inputgacode + else: + profiles = None # TGYRO class. It checks existence and creates input.profiles/input.gacode diff --git a/src/mitim_tools/gacode_tools/TGYROtools.py b/src/mitim_tools/gacode_tools/TGYROtools.py index a53813d6..9dfd1676 100644 --- a/src/mitim_tools/gacode_tools/TGYROtools.py +++ b/src/mitim_tools/gacode_tools/TGYROtools.py @@ -251,9 +251,7 @@ def run( Tepred, Tipred, nepred = PredictionSet self.FolderTGYRO = IOtools.expandPath(self.FolderGACODE / subFolderTGYRO) - self.FolderTGYRO_tmp = ( - self.FolderTGYRO / "tmp_tgyro_run" - ) # Folder to run TGYRO on (or to retrieve the raw outputs from a cluster) + self.FolderTGYRO_tmp = self.FolderTGYRO / "tmp_tgyro_run" # Folder to run TGYRO on (or to retrieve the raw outputs from a cluster) inputclass_TGYRO = TGYROinput( input_profiles=self.profiles, From 77febf204179166b9df2b65abd564a8d85f4fec1 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 15 Oct 2025 16:51:07 -0400 Subject: [PATCH 383/385] Additional folder expansions for robustness (after documentation followup) --- src/mitim_tools/misc_tools/FARMINGtools.py | 4 ++-- src/mitim_tools/simulation_tools/SIMtools.py | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/mitim_tools/misc_tools/FARMINGtools.py b/src/mitim_tools/misc_tools/FARMINGtools.py index 63f517f2..93265fc6 100644 --- a/src/mitim_tools/misc_tools/FARMINGtools.py +++ b/src/mitim_tools/misc_tools/FARMINGtools.py @@ -208,8 +208,8 @@ def run( self.output_files = curateOutFiles(self.output_files) # Relative paths - self.input_files = [path.relative_to(self.folder_local) for path in self.input_files] - self.input_folders = [path.relative_to(self.folder_local) for path in self.input_folders] + self.input_files = [IOtools.expandPath(path).relative_to(self.folder_local) for path in self.input_files] + self.input_folders = [IOtools.expandPath(path).relative_to(self.folder_local) for path in self.input_folders] # Process self.full_process( diff --git a/src/mitim_tools/simulation_tools/SIMtools.py b/src/mitim_tools/simulation_tools/SIMtools.py index a22530b8..4bc6518f 100644 --- a/src/mitim_tools/simulation_tools/SIMtools.py +++ b/src/mitim_tools/simulation_tools/SIMtools.py @@ -260,6 +260,9 @@ def _run_prepare( if len(rhosEvaluate) == len(rhos): # All radii need to be evaluated IOtools.askNewFolder(Folder_sim, force=forceIfcold_start) + + # Once created, expand here + Folder_sim = IOtools.expandPath(Folder_sim) # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Change this specific run From 5573f866ae4a1c8459a8455d12ec92ff160cddbb Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 15 Oct 2025 16:51:14 -0400 Subject: [PATCH 384/385] Updated documentation --- docs/capabilities/misc_capabilities.rst | 2 +- docs/capabilities/optimization.rst | 2 +- docs/capabilities/tglf_capabilities.rst | 49 +++++++---------------- docs/capabilities/tgyro_capabilities.rst | 10 +---- docs/capabilities/transp_capabilities.rst | 21 +--------- docs/capabilities/vitals_capabilities.rst | 6 +-- docs/faq.rst | 8 +--- docs/index.rst | 6 +-- docs/installation.rst | 46 ++++++++++++--------- docs/requirements.txt | 29 -------------- 10 files changed, 53 insertions(+), 126 deletions(-) diff --git a/docs/capabilities/misc_capabilities.rst b/docs/capabilities/misc_capabilities.rst index 41e23b1e..aa8ffed2 100644 --- a/docs/capabilities/misc_capabilities.rst +++ b/docs/capabilities/misc_capabilities.rst @@ -10,7 +10,7 @@ Miscellaneous Interpret Tokamak equilibrium ----------------------------- -MITIM has a quick g-eqdsk file reader and visualizer that is based on the ``omfit-classes`` package. +MITIM has a quick g-eqdsk file reader and visualizer that is based on the ``megpy`` package. To open and plot a g-eqdsk file: diff --git a/docs/capabilities/optimization.rst b/docs/capabilities/optimization.rst index f3c14478..4b4d7782 100644 --- a/docs/capabilities/optimization.rst +++ b/docs/capabilities/optimization.rst @@ -105,7 +105,7 @@ Now we can create and launch the MITIM optimization process from the beginning ( .. code-block:: python - MITIM_BO = STRATEGYtools.MITIM_BO( opt_fun1D, cold_startYN = True ) + MITIM_BO = STRATEGYtools.MITIM_BO( opt_fun1D, cold_start = True ) MITIM_BO.run() Once finished, we can plot the results easily with: diff --git a/docs/capabilities/tglf_capabilities.rst b/docs/capabilities/tglf_capabilities.rst index 41862a74..8ac3ed7c 100644 --- a/docs/capabilities/tglf_capabilities.rst +++ b/docs/capabilities/tglf_capabilities.rst @@ -68,30 +68,31 @@ To generate the input files (input.tglf) to TGLF at each radial location, MITIM .. code-block:: python - cdf = tglf.prep(folder,inputgacode=inputgacode_file,cold_start=False ) + _ = tglf.prep(inputgacode_file,folder,cold_start=False) -.. tip:: - The ``.prep()`` method, when applied to a case that starts with an input.gacode file, launches a `TGYRO` run for a "zero" iteration to generate *input.tglf* at specific ``rho`` locations from the *input.gacode*. This method to generate input files is inspired by how the `OMFIT framework `_ works. +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 ``cold_start`` flag can be provided. + +The set of control inputs to TGLF (saturation rule, electromagnetic effects, basis functions, etc.) are provided following the following sequential logic: + +1. Each code has a set of default settings, which for TGLF are specified in ``templates/input.tglf.controls``. This is the base namelist of settings that will be used if no other specification is provided. +2. Then, a ``code_settings`` argument can be provided to the ``run()`` command. This argument refers to a specific set of settings that are specified in ``templates/input.tglf.models.yaml``, and that will overwrite the default settings in ``input.tglf.controls``. +3. Finally, an ``extraOptions`` argument can be provided to the ``run()`` command, which is a dictionary of specific settings to change from the previous two steps. -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 ``cold_start`` 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 ``code_settings`` indicates the base case to start with. -The user is referred to ``templates/input.tglf.models.yaml`` 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): .. code-block:: python tglf.run( subfolder = 'yes_em_folder', - code_settings = 5, - extraOptions = {}, + code_settings = 'SAT2', + extraOptions = {'USE_BPER':True}, cold_start = False ) tglf.read( label = 'yes_em' ) tglf.run( subfolder = 'no_em_folder', - code_settings = 5, + code_settings = 'SAT2', extraOptions = {'USE_BPER':False}, cold_start = False ) @@ -139,7 +140,7 @@ Similarly as in the previous section, you need to run the ``prep()`` command, bu .. code-block:: python - cdf = tglf.prep(folder,cold_start=False) + cdf = tglf.prep_using_tgyro(folder,cold_start=False) .. note:: @@ -215,26 +216,4 @@ Run 1D scans of TGLF input parameter TGLF aliases ------------ -MITIM provides a few useful aliases, including for the TGLF tools: - -- To plot results that exist in a folder ``run1/``, with or without a suffix and with or without an input.gacode file (for normalizations): - - .. code-block:: bash - - mitim_plot_tglf run1/ - mitim_plot_tglf run1/ --suffix _0.55 --gacode input.gacode - - -- To run TGLF in a folder ``run1/`` using input file ``input.tglf``, with or without an input.gacode file (for normalizations): - - .. code-block:: bash - - mitim_run_tglf --folder run1/ --tglf input.tglf - mitim_run_tglf --folder run1/ --tglf input.tglf --gacode input.gacode - -- To run a parameter scan in a folder ``scan1/`` using input file ``input.tglf``, with or without an input.gacode file (for normalizations): - - .. code-block:: bash - - mitim_run_tglf --folder scan1/ --tglf input.tglf --gacode input.gacode --scan RLTS_2 - +MITIM provides a few useful aliases, including for the TGLF tools: :ref:`Shell Scripts` diff --git a/docs/capabilities/tgyro_capabilities.rst b/docs/capabilities/tgyro_capabilities.rst index e1a35c1b..0efec858 100644 --- a/docs/capabilities/tgyro_capabilities.rst +++ b/docs/capabilities/tgyro_capabilities.rst @@ -83,7 +83,7 @@ Now TGYRO can be run: PredictionSet = PredictionSet, TGLFsettings = TGLFsettings, TGYRO_solver_options = solver, - Physics_options = physics_options) + TGYRO_physics_options = physics_options) Read: @@ -133,11 +133,5 @@ Plot results: TGYRO aliases ------------- -MITIM provides a few useful aliases, including for the TGYRO tools: - -- To plot results that exist in a folder ``run1/``: - - .. code-block:: bash - - mitim_plot_tgyro run1/ +MITIM provides a few useful aliases, including for the TGYRO tools: :ref:`Shell Scripts` diff --git a/docs/capabilities/transp_capabilities.rst b/docs/capabilities/transp_capabilities.rst index 8d194bb9..8fb3bfae 100644 --- a/docs/capabilities/transp_capabilities.rst +++ b/docs/capabilities/transp_capabilities.rst @@ -123,24 +123,5 @@ If TRANSP has already been run and the .CDF results file already exists (``cdf_f TRANSP aliases -------------- -MITIM provides a few useful aliases, including for the TRANSP tools: - -- To read TRANSP results in CDF files (which stores the results in the ``cdfs`` list. First run can be plotted with ``cdfs[0].plot``): - - .. code-block:: bash - - mitim_read_transp 12345A01.CDF 12345A02.CDF - -- To interact with the TRANSP globus grid: - - .. code-block:: bash - - # To check status of runs under username pablorf - mitim_trcheck pablorf - - # To remove from the grid CMOD run numbers 88664P01, 88664P03 from user pablorf - mitim_trclean 88664P CMOD --numbers 1,3 - - # To get results file (intermediate or final) from CMOD run 152895P01 from user pablorf - mitim_trlook 152895P01 CMOD +MITIM provides a few useful aliases, including for the TRANSP tools: :ref:`Shell Scripts` diff --git a/docs/capabilities/vitals_capabilities.rst b/docs/capabilities/vitals_capabilities.rst index 62f08239..8e9e5da6 100644 --- a/docs/capabilities/vitals_capabilities.rst +++ b/docs/capabilities/vitals_capabilities.rst @@ -38,8 +38,8 @@ As a starting point of VITALS, you need to prepare and run TGLF for the base cas rho = 0.5 tglf = TGLFtools.TGLF( rhos = [ rho ] ) - cdf = tglf.prep( folder, inputgacode = inputgacode_file) - tglf.run( subfolder = 'run_base', code_settings = 5) + cdf = tglf.prep( inputgacode_file, folder ) + tglf.run( subfolder = 'run_base', code_settings = 'SAT3') tglf.read( label = 'run_base' ) @@ -120,7 +120,7 @@ Once the VITALS object has been created, parameters such as the TGLF control inp .. code-block:: python - vitals_fun.TGLFparameters['code_settings'] = 5 + vitals_fun.TGLFparameters['code_settings'] = 'SAT3' vitals_fun.TGLFparameters['extraOptions'] = {} .. note:: diff --git a/docs/faq.rst b/docs/faq.rst index 9205ad98..bc59e530 100644 --- a/docs/faq.rst +++ b/docs/faq.rst @@ -30,7 +30,7 @@ Issues during MITIM installation .. code-block:: console - pip3 install -e MITIM-fusion\[pyqt\] --no-cache + pip3 install -e MITIM-fusion\[pyqt\] --no-cache-dir Issues during MITIM tests ------------------------- @@ -61,9 +61,3 @@ Issues during MITIM tests Make sure you that, if you have keys, you have added them to authorized_keys in both server and tunnel machines. - - -Issues during PORTALS simulations ---------------------------------- - -Nothing here yet. diff --git a/docs/index.rst b/docs/index.rst index e95508de..f4ab8834 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,8 +1,8 @@ MITIM: a toolbox for modeling tasks in plasma physics and fusion energy ======================================================================= -The **MITIM** (MIT Integrated Modeling) is a versatile and user-friendly Python library designed for *plasma physics* and *fusion energy* researchers, distributed as the `MITIM-fusion `_ GitHub repository. -Developed in 2018 by `Pablo Rodriguez-Fernandez `_ at the MIT Plasma Science and Fusion Center, this light-weight, command-line, +The **MITIM** (MIT Integrated Modeling) is a versatile and user-friendly Python library designed for plasma physics and fusion energy researchers, distributed as the `MITIM-fusion `_ GitHub repository. +Spearheaded by `Pablo Rodriguez-Fernandez `_ at the MIT Plasma Science and Fusion Center, this light-weight, command-line, object-oriented toolbox streamlines the execution and interpretation of physics models and simplifies complex optimization tasks. MITIM stands out for its modular nature, making it particularly useful for integrating models with optimization workflows. @@ -28,7 +28,7 @@ Overview -------- Developed at the MIT Plasma Science and Fusion Center, MITIM emerged in 2023 as a progression from the PORTALS project (*Performance Optimization of Reactors via Training of Active Learning Surrogates*). -This evolution marks a significant enhancement in our approach to transport and optimization in plasma physics research. +This evolution marked a significant enhancement in our approach to transport and optimization in plasma physics research. MITIM's core functionality revolves around the standalone execution of codes and the nuanced interpretation of results through object-oriented Python scripts. This enables researchers to seamlessly integrate these scripts into custom surrogate-based optimization frameworks, diff --git a/docs/installation.rst b/docs/installation.rst index 7558ecf2..5446bcf0 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -36,13 +36,6 @@ Use ``pip`` to install all the required MITIM requirements: The optional argument ``[pyqt]`` added in the intallation command above must only be used if the machine allows for graphic interfaces. If running in a computing cluster, remove that flag. The ``pyqt`` package is used to create condensed figures into a single notebook when interpreting and plotting simulation results. - - If you wish to install all capabilities (including compatibility with `OMFIT `_), it is recommended that ``pip`` is run as follows: - - .. code-block:: console - - pip3 install -e MITIM-fusion[pyqt,omfit] - If you were unsuccessful in the installation, check out our :ref:`Frequently Asked Questions` section. @@ -53,11 +46,11 @@ User configuration In ``MITIM-fusion/templates/``, there is a ``config_user_example.json`` with specifications of where to run certain codes and what the login requirements are. There are also options to specify the default verbose level and the default DPI for the figures in notebooks. Users need to specify their own configurations in a file that follows the same structure. -There are different options to handle this config file. +There are different options to handle this config file: 1. Create a new file named ``config_user.json`` **in the same folder** ``MITIM-fusion/templates/``. MITIM will automatically look for this file when running the code. 2. Create a new file anywhere in your machine. Then, **set the environment variable** ``MITIM_CONFIG`` to the path of this file. MITIM will automatically look for this file when running the code. -3. Create a new file anywhere in your machine. **Do this at the beginning of your script**: +3. Create a new file anywhere in your machine. Then, **add these lines at the beginning of your script**: .. code-block:: python @@ -87,16 +80,29 @@ In this example, the ``identity`` option is only required if you are running in { "preferences": { "tglf": "engaging", + "neo": "local", "tgyro": "perlmutter", "verbose_level": "5", "dpi_notebook": "80" }, + "local": { + "machine": "local", + "username": "YOUR_USERNAME", + "scratch": "/Users/YOUR_USERNAME/scratch/", + "modules": "", + "cores_per_node": 8, + "gpus_per_node": 0 + }, "engaging": { "machine": "eofe7.mit.edu", "username": "YOUR_USERNAME", "scratch": "/pool001/YOUR_USERNAME/scratch/", + "modules": "", + "cores_per_node": 64, + "gpus_per_node": 0, "slurm": { "partition": "sched_mit_psfc", + "exclusive": false, "exclude": "node584" } }, @@ -104,12 +110,16 @@ In this example, the ``identity`` option is only required if you are running in "machine": "perlmutter.nersc.gov", "username": "YOUR_USERNAME", "scratch": "/pscratch/sd/p/YOUR_USERNAME/scratch/", + "modules": "", "identity": "/Users/YOUR_USERNAME/.ssh/id_rsa_nersc", + "cores_per_node": 32, + "gpus_per_node": 4, "slurm": { "account": "YOUR_ACCOUNT", "partition": "YOUR_PARTITION", "constraint": "gpu", - "mem": "4GB" + "mem": "4GB", + "email": "optional@email" } } } @@ -119,20 +129,15 @@ MITIM will attempt to create SSH and SFTP connections to that machine, and will .. 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. + 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. Source at shell initialization (recommended) +.. dropdown:: 1. Send specific commands per code (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:: 2. Send specific commands per code - - Finally, you can populate the ``modules`` option per machine in your ``config_user.json`` file. For example: + You can populate the ``modules`` option per machine in your ``config_user.json`` file. For example: .. code-block:: console @@ -142,10 +147,13 @@ There are several ways to make sure that the shell understands the command: ... } - 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. +.. dropdown:: 2. Source at shell initialization + + If 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. License and contributions diff --git a/docs/requirements.txt b/docs/requirements.txt index 9d91ba68..ae1a84d8 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -5,32 +5,3 @@ sphinx-copybutton sphinx-togglebutton sphinx-panels sphinxcontrib-images -h5py -matplotlib -importlib-metadata -IPython -psutil -pip -numpy -matplotlib -argparse -h5py -netCDF4 -xarray -pandas -xlsxwriter -statsmodels -dill -IPython -pyDOE -multiprocessing_on_dill -deap -paramiko -tqdm -shapely -freegs -botorch -scikit-image -psutil -onnx2pytorch -tensorflow \ No newline at end of file From d602b66aabc6c2ba54ed9c8197e5c063dc0552a8 Mon Sep 17 00:00:00 2001 From: Pablo RF Date: Wed, 15 Oct 2025 17:02:57 -0400 Subject: [PATCH 385/385] Removed deprecated TGLF instructions --- src/mitim_tools/gacode_tools/TGLFtools.py | 84 ----------------------- 1 file changed, 84 deletions(-) diff --git a/src/mitim_tools/gacode_tools/TGLFtools.py b/src/mitim_tools/gacode_tools/TGLFtools.py index 32c1daff..66bd75f6 100644 --- a/src/mitim_tools/gacode_tools/TGLFtools.py +++ b/src/mitim_tools/gacode_tools/TGLFtools.py @@ -38,90 +38,6 @@ def __init__( avTime=0.0, # Averaging window to extract CDF file alreadyRun=None, # Option2: Do more stuff with a class that has already been created and store ): - """ - TGLF class that manages the run and the results. - - The philosophy of this is that a single 'tglf' class will handle the tglf simulation and results - at one time slice but with possibility of several radii at once. - - It can also handle different TGLF settings, running them one by one, storing results in folders and then - grabbing them. - - Scans can also be run. At several radii at once if wanted. - - *Note* - The 'run' command does not require label. When performing a 'read', the results extracted from - the specified folder will be saved with the label indicated in 'read', in the "results" or "scans" - dictionaries. Plotting then can happen with more than one label of the same category. - - *Note* - The 'run' command uses input.tglf from the specified folder, but one can change the Settings presets, - extraOptions and multipliers. The modified inputs is not rewritten in the actual folder, it is only written - in the tmp folder on which the simulation takes place. - - *Note* - After a 'prep' command, the class can be detached from the file system, as it stores the input tglf file - to run later with different options. It also stores the Normalizations, since the runs are expected - to only change dimensionless parameteres. - - ************************************** - ***** Example use for standalone ***** - ************************************** - - # Initialize class, by specifying where the inputs to TGLF come from (TRANSP cdf) - tglf = TGLF(cdf='~/testTGLF/12345B12.CDF',time=1.45,avTime=0.1,rhos=[0.4,0.6]) - - # Prepare TGLF (this will create input.tglf in the specified folder) - cdf = tglf.prep_using_tgyro('~/testTGLF/') - - # Run standalone TGLF (this will find the input.tglf in the previous folder, - # and then copy to this specify TGLF run, and run it there) - tglf.run(subfolder='tglf1/',Settings=1,extraOptions={'NS':3}) - - # Read results - tglf.read(label='run1',folder='~/testTGLF/tglf1/') - - # Plot - plt.ion(); tglf.plot(labels=['run1']) - - ********************************* - ***** Example use for scans ***** - ********************************* - - # Initialize class, by specifying where the inputs to TGLF come from (TRANSP cdf) - tglf = TGLF(cdf='~/testTGLF/12345B12.CDF',time=1.45,avTime=0.1,rhos=[0.4,0.6]) - - # Prepare TGLF (this will create input.tglf in the specified folder) - cdf = tglf.prep_using_tgyro('~/testTGLF/') - - # Run - tglf.run_scan('scan1/',Settings=1,varUpDown=np.linspace(0.5,2.0,20),variable='RLTS_2') - - # Read scan - tglf.read_scan(label='scan1',variable='RLTS_2') - - # Plot - plt.ion(); tglf.plot_scan(labels=['scan1'],variableLabel='RLTS_2') - - **************************** - ***** Special analysis ***** - **************************** - - Following the prep phase, we can run "runAnalysis()" and select among the different options: - - Chi_inc - - D and V for trace impurity - Then, plotAnalysis() with the right option for different labels too - - **************************** - ***** Do more stuff with a class that has already been created and store - **************************** - - tglf = TGLF(alreadyRun=previousClass) - tglf.FolderGACODE = '~/testTGLF/' - - ** Modify the class as wish, and do run,read, etc ** - ** Because normalizations are stored in the prep phase, that's all ready ** - """ super().__init__(rhos=rhos)