Merge branch 'main' into theo-brown/add-prad

docs/code_structure.rst

@@ -232,7 +232,7 @@ inside |torax.sim.run_simulation()|_.
 .. |torax.sources| replace:: ``torax.sources``
 .. _torax.sources: https://github.com/google-deepmind/torax/tree/main/torax/sources
 .. |QLKNN| replace:: ``QLKNN``
-.. _QLKNN: https://github.com/google-deepmind/torax/blob/main/torax/transport_model/qlknn_wrapper.py
+.. _QLKNN: https://github.com/google-deepmind/torax/blob/main/torax/transport_model/qlknn_transport_model.py
 .. |torax.transport_model| replace:: ``torax.transport_model``
 .. _torax.transport_model: https://github.com/google-deepmind/torax/blob/main/torax/transport_model
 .. |torax.time_step_calculator| replace:: ``torax.time_step_calculator``

docs/configuration.rst

@@ -281,27 +281,10 @@ Configures boundary conditions, initial conditions, and prescribed time-dependen
   Toggles units of ``ne_bound_right``.
 
 ``set_pedestal`` (bool = True), **time-varying-scalar**
-  Set internal boundary conditions if True. Do not set internal boundary conditions if False.
+  If True use the configured pedestal model to set internal boundary conditions. Do not set internal boundary conditions if False.
   Internal boundary conditions are set using an adaptive localized source term. While a common use-case is to mock up a pedestal, this feature
   can also be used for L-mode modeling with a desired internal boundary condition below :math:`\hat{\rho}=1`.
 
-``Tiped`` (float = 5.0), **time-varying-scalar**
-  Internal boundary condition for ion temperature at :math:`\hat{\rho}` = ``Ped_top``, in units of keV.
-
-``Teped`` (float = 5.0), **time-varying-scalar**
-  Internal boundary condition for electron temperature at :math:`\hat{\rho}` = ``Ped_top``, in units of keV.
-
-``neped`` (float = 0.7), **time-varying-scalar**
-  Internal boundary condition for electron density at  :math:`\hat{\rho}` = ``Ped_top``, in units of keV.
-  In units of reference density if ``neped_is_fGW==False``. In units of Greenwald fraction if ``neped_is_fGW==True``.
-
-``neped_is_fGW`` (bool = False)
-  Toggles units of ``neped``.
-
-``Ped_top`` (float = 0.91), **time-varying-scalar**
-  Location of internal boundary condition, in units of :math:`\hat{\rho}`. In practice, the closest cell
-  gridpoint to ``Ped_top`` will be used.
-
 ``nu`` (float = 3.0)
   Peaking coefficient of initial current profile: :math:`j = j_0(1 - \hat{\rho}^2)^\nu`. :math:`j_0` is calculated
   to be consistent with a desired total current. Only used if ``initial_psi_from_j==True``, otherwise the ``psi`` profile from the geometry file is used.
@@ -388,6 +371,43 @@ output_dir
 
 .. _time_step_calculator:
 
+pedestal
+--------
+If ``set_pedestal`` is set to True in the ``profile_conditions`` section, then
+a pedestal model config is required.
+
+In TORAX we aim to support different models for computing the pedestal width,
+and electron density, ion temperature and electron temperature at the pedestal
+top. These models will only be used if the ``set_pedestal`` option in the
+``profile_conditions`` section is set to True.
+
+The model can be configured by setting the ``pedestal_model`` key in the
+``pedestal`` section of the configuration. If this field is not set, then
+the default model is ``'set_tped_nped'``.
+
+The following models are currently supported:
+
+set_tped_nped
+^^^^^^^^^^^^^
+Directly specify the pedestal width, electron density, ion temperature and
+electron temperature.
+
+``neped`` (float = 0.7) **time-varying-scalar**
+  Electron density at the pedestal top.
+  In units of reference density if ``neped_is_fGW==False``. In units of Greenwald fraction if ``neped_is_fGW==True``.
+
+``neped_is_fGW`` (bool = False) **time-varying-scalar**
+  Toggles units of ``neped``.
+
+``Tiped`` (float = 5.0) **time-varying-scalar**
+  Ion temperature at the pedestal top in units of keV.
+
+``Teped`` (float = 5.0) **time-varying-scalar**
+  Electron temperature at the pedestal top in units of keV.
+
+``rho_norm_ped`` (float = 0.91) **time-varying-scalar**
+  Location of pedestal top, in units of :math:`\hat{\rho}`.
+
 geometry
 --------
 

docs/output.rst

@@ -320,6 +320,15 @@ analysis and inspection.
 ``P_icrh_tot`` (time) [W]:
   Total ion cyclotron resonance heating power.
 
+``P_LH_hi_dens`` (time) [W]: H-mode transition power for high density branch,
+  according to Eq 3 from Martin 2008.
+
+``P_LH_low_dens`` (time) [W]: H-mode transition power for low density branch,
+  according to Eq 4 from Ryter 2014.
+
+``ne_min_P_LH`` (time) [nref]:  Density corresponding to the minimum P_LH,
+  according to Eq 3 from Ryter 2014.
+
 ``E_cumulative_fusion`` (time) [J]:
   Total cumulative fusion energy.
 

pyproject.toml

@@ -30,8 +30,7 @@ dependencies = [
     "chex>=0.1.85",
     "equinox>=0.11.3",
     "PyYAML>=6.0.1",
-    # DataTree is only working with development version of xarray for now.
-    "xarray @ git+https://github.com/pydata/xarray",
+    "xarray>=2024.11.0",
     "netcdf4>=1.6.5,<1.7.1",
     "h5netcdf>=1.3.0",
     "scipy>=1.12.0",

run_simulation_main.py

@@ -33,7 +33,7 @@
 from torax.config import build_sim
 from torax.config import config_loader
 from torax.plotting import plotruns_lib
-from torax.transport_model import qlknn_wrapper
+from torax.transport_model import qlknn_transport_model
 
 
 # String used when prompting the user to make a choice of command
@@ -104,9 +104,9 @@
     'Path to the qlknn model network parameters (if using a QLKNN transport'
     ' model). If not set, then it will use the value from the config in the'
     ' "model_path" field in the qlknn_params. If that is not set, it will look'
-    f' for the "{qlknn_wrapper.MODEL_PATH_ENV_VAR}" env variable.'
+    f' for the "{qlknn_transport_model.MODEL_PATH_ENV_VAR}" env variable.'
     ' Finally, if this is also not set, it uses a hardcoded default path'
-    f' "{qlknn_wrapper.DEFAULT_MODEL_PATH}".',
+    f' "{qlknn_transport_model.DEFAULT_MODEL_PATH}".',
 )
 
 _OUTPUT_DIR = flags.DEFINE_string(
@@ -274,6 +274,11 @@ def change_config(
     new_stepper_builder = build_sim.build_stepper_builder_from_config(
         sim_config['stepper']
     )
+    new_pedestal_model_builder = (
+        build_sim.build_pedestal_model_builder_from_config(
+            sim_config['pedestal'] if 'pedestal' in sim_config else {}
+        )
+    )
   else:
     # Assume the config module has several methods to define the individual Sim
     # attributes (the "advanced", more Python-forward configuration method).
@@ -282,6 +287,7 @@ def change_config(
     new_transport_model_builder = config_module.get_transport_model_builder()
     source_models_builder = config_module.get_sources_builder()
     new_stepper_builder = config_module.get_stepper_builder()
+    new_pedestal_model_builder = config_module.get_pedestal_model_builder()
   new_source_params = {
       name: runtime_params
       for name, runtime_params in source_models_builder.runtime_params.items()
@@ -301,6 +307,7 @@ def change_config(
       transport_runtime_params=new_transport_model_builder.runtime_params,
       source_runtime_params=new_source_params,
       stepper_runtime_params=new_stepper_builder.runtime_params,
+      pedestal_runtime_params=new_pedestal_model_builder.runtime_params,
   )
   return sim, new_runtime_params
 

torax/__init__.py

@@ -99,6 +99,7 @@
     'source_profile',
     'explicit_source_profiles',
     'source_models',
+    'pedestal_model',
     'time_step_calculator',
     'coeffs_callback',
     'evolving_names',

torax/calc_coeffs.py

@@ -28,15 +28,17 @@
 from torax import state
 from torax.config import runtime_params_slice
 from torax.fvm import block_1d_coeffs
+from torax.pedestal_model import pedestal_model as pedestal_model_lib
 from torax.sources import source_models as source_models_lib
 from torax.sources import source_profiles as source_profiles_lib
 from torax.transport_model import transport_model as transport_model_lib
 
 
-def calculate_pereverzev_flux(
+def _calculate_pereverzev_flux(
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
     geo: geometry.Geometry,
     core_profiles: state.CoreProfiles,
+    pedestal_model_output: pedestal_model_lib.PedestalModelOutput,
 ) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]:
   """Adds Pereverzev-Corrigan flux to diffusion terms."""
 
@@ -80,7 +82,7 @@ def calculate_pereverzev_flux(
       jnp.logical_and(
           dynamic_runtime_params_slice.profile_conditions.set_pedestal,
           geo.rho_face_norm
-          > dynamic_runtime_params_slice.profile_conditions.Ped_top,
+          > pedestal_model_output.rho_norm_ped_top,
       ),
       0.0,
       chi_face_per_ion,
@@ -89,7 +91,7 @@ def calculate_pereverzev_flux(
       jnp.logical_and(
           dynamic_runtime_params_slice.profile_conditions.set_pedestal,
           geo.rho_face_norm
-          > dynamic_runtime_params_slice.profile_conditions.Ped_top,
+          > pedestal_model_output.rho_norm_ped_top,
       ),
       0.0,
       chi_face_per_el,
@@ -110,7 +112,7 @@ def calculate_pereverzev_flux(
       jnp.logical_and(
           dynamic_runtime_params_slice.profile_conditions.set_pedestal,
           geo.rho_face_norm
-          > dynamic_runtime_params_slice.profile_conditions.Ped_top,
+          > pedestal_model_output.rho_norm_ped_top,
       ),
       0.0,
       d_face_per_el * geo.g1_over_vpr_face,
@@ -120,7 +122,7 @@ def calculate_pereverzev_flux(
       jnp.logical_and(
           dynamic_runtime_params_slice.profile_conditions.set_pedestal,
           geo.rho_face_norm
-          > dynamic_runtime_params_slice.profile_conditions.Ped_top,
+          > pedestal_model_output.rho_norm_ped_top,
       ),
       0.0,
       v_face_per_el * geo.g0_face,
@@ -147,6 +149,7 @@ def calc_coeffs(
     transport_model: transport_model_lib.TransportModel,
     explicit_source_profiles: source_profiles_lib.SourceProfiles,
     source_models: source_models_lib.SourceModels,
+    pedestal_model: pedestal_model_lib.PedestalModel,
     evolving_names: tuple[str, ...],
     use_pereverzev: bool = False,
     explicit_call: bool = False,
@@ -173,6 +176,7 @@ def calc_coeffs(
     source_models: All TORAX source/sink functions that generate the explicit
       and implicit source profiles used as terms for the core profiles
       equations.
+    pedestal_model: A PedestalModel subclass, calculates pedestal values.
     evolving_names: The names of the evolving variables in the order that their
       coefficients should be written to `coeffs`.
     use_pereverzev: Toggle whether to calculate Pereverzev terms
@@ -202,6 +206,7 @@ def calc_coeffs(
         transport_model,
         explicit_source_profiles,
         source_models,
+        pedestal_model,
         evolving_names,
         use_pereverzev,
     )
@@ -212,6 +217,7 @@ def calc_coeffs(
     static_argnames=[
         'static_runtime_params_slice',
         'transport_model',
+        'pedestal_model',
         'source_models',
         'evolving_names',
     ],
@@ -224,6 +230,7 @@ def _calc_coeffs_full(
     transport_model: transport_model_lib.TransportModel,
     explicit_source_profiles: source_profiles_lib.SourceProfiles,
     source_models: source_models_lib.SourceModels,
+    pedestal_model: pedestal_model_lib.PedestalModel,
     evolving_names: tuple[str, ...],
     use_pereverzev: bool = False,
 ) -> block_1d_coeffs.Block1DCoeffs:
@@ -249,6 +256,7 @@ def _calc_coeffs_full(
     source_models: All TORAX source/sink functions that generate the explicit
       and implicit source profiles used as terms for the core profiles
       equations.
+    pedestal_model: A PedestalModel subclass, calculates pedestal values.
     evolving_names: The names of the evolving variables in the order that their
       coefficients should be written to `coeffs`.
     use_pereverzev: Toggle whether to calculate Pereverzev terms
@@ -259,11 +267,25 @@ def _calc_coeffs_full(
 
   consts = constants.CONSTANTS
 
+  pedestal_model_output: pedestal_model_lib.PedestalModelOutput = jax.lax.cond(
+      dynamic_runtime_params_slice.profile_conditions.set_pedestal,
+      lambda: pedestal_model(
+          dynamic_runtime_params_slice, geo, core_profiles
+      ),
+      # TODO(b/380271610): Refactor to avoid needing dummy output.
+      lambda: pedestal_model_lib.PedestalModelOutput(
+          neped=0.0,
+          Tiped=0.0,
+          Teped=0.0,
+          rho_norm_ped_top=0.0,
+      ),
+  )
+
   # Boolean mask for enforcing internal temperature boundary conditions to
   # model the pedestal.
   mask = physics.internal_boundary(
       geo,
-      dynamic_runtime_params_slice.profile_conditions.Ped_top,
+      pedestal_model_output.rho_norm_ped_top,
       dynamic_runtime_params_slice.profile_conditions.set_pedestal,
   )
 
@@ -403,7 +425,7 @@ def _calc_coeffs_full(
 
   # Diffusion term coefficients
   transport_coeffs = transport_model(
-      dynamic_runtime_params_slice, geo, core_profiles
+      dynamic_runtime_params_slice, geo, core_profiles, pedestal_model_output
   )
   chi_face_ion = transport_coeffs.chi_face_ion
   chi_face_el = transport_coeffs.chi_face_el
@@ -562,24 +584,11 @@ def _calc_coeffs_full(
       source_models,
   )
 
-  # calculate neped
-  # pylint: disable=invalid-name
-  nGW = (
-      dynamic_runtime_params_slice.profile_conditions.Ip_tot
-      / (jnp.pi * geo.Rmin**2)
-      * 1e20
-      / dynamic_runtime_params_slice.numerics.nref
-  )
-  # pylint: enable=invalid-name
-  neped_unnorm = jnp.where(
-      dynamic_runtime_params_slice.profile_conditions.neped_is_fGW,
-      dynamic_runtime_params_slice.profile_conditions.neped * nGW,
-      dynamic_runtime_params_slice.profile_conditions.neped,
-  )
-
   source_ne += jnp.where(
       dynamic_runtime_params_slice.profile_conditions.set_pedestal,
-      mask * dynamic_runtime_params_slice.numerics.largeValue_n * neped_unnorm,
+      mask
+      * dynamic_runtime_params_slice.numerics.largeValue_n
+      * pedestal_model_output.neped,
       0.0,
   )
   source_mat_nn += jnp.where(
@@ -602,10 +611,11 @@ def _calc_coeffs_full(
       v_face_per_el,
   ) = jax.lax.cond(
       use_pereverzev,
-      lambda: calculate_pereverzev_flux(
+      lambda: _calculate_pereverzev_flux(
           dynamic_runtime_params_slice,
           geo,
           core_profiles,
+          pedestal_model_output,
       ),
       lambda: tuple([jnp.zeros_like(geo.rho_face)] * 6),
   )
@@ -720,14 +730,14 @@ def _calc_coeffs_full(
       dynamic_runtime_params_slice.profile_conditions.set_pedestal,
       mask
       * dynamic_runtime_params_slice.numerics.largeValue_T
-      * dynamic_runtime_params_slice.profile_conditions.Tiped,
+      * pedestal_model_output.Tiped,
       0.0,
   )
   source_e += jnp.where(
       dynamic_runtime_params_slice.profile_conditions.set_pedestal,
       mask
       * dynamic_runtime_params_slice.numerics.largeValue_T
-      * dynamic_runtime_params_slice.profile_conditions.Teped,
+      * pedestal_model_output.Teped,
       0.0,
   )