Merge branch 'main' into tglf-surrogate

.github/workflows/pytest.yml

@@ -79,38 +79,8 @@ jobs:
     - run: pip freeze
 
     # Run tests (in parallel)
-    # TODO(b/326579500): tests should be discovered automatically
     - name: Run core tests
       run: |
         pytest \
-          torax/config/tests/runtime_params_slice.py \
-          torax/config/tests/runtime_params.py \
-          torax/fvm/tests/fvm.py \
-          torax/sources/tests/bootstrap_current_source.py \
-          torax/sources/tests/electron_density_sources.py \
-          torax/sources/tests/generic_current_source.py \
-          torax/sources/tests/formulas.py \
-          torax/sources/tests/fusion_heat_source.py \
-          torax/sources/tests/generic_ion_el_heat_source.py \
-          torax/sources/tests/qei_source.py \
-          torax/sources/tests/source_models.py \
-          torax/sources/tests/source.py \
-          torax/spectators/tests/plotting.py \
-          torax/spectators/tests/spectator.py \
-          torax/tests/boundary_conditions.py \
-          torax/tests/geometry.py \
-          torax/tests/interpolated_param.py \
-          torax/tests/jax_utils.py \
-          torax/tests/math_utils.py \
-          torax/tests/persistent_cache.py \
-          torax/tests/physics.py \
-          torax/tests/test_run_simulation_main.py \
-          torax/tests/sim_custom_sources.py \
-          torax/tests/sim_time_dependence.py \
-          torax/tests/sim.py \
-          torax/tests/sim_no_compile.py \
-          torax/tests/state.py \
-          torax/transport_model/tests/qlknn_wrapper.py \
-          torax/transport_model/tests/transport_model.py \
           -vv -n auto \
           --shard-id=$((${{ matrix.shard-id }} - 1)) --num-shards=${{ env.PYTEST_NUM_SHARDS }}

docs/configuration.rst

@@ -200,6 +200,9 @@ Defines the distribution of ion species. Currently restricted to a single main i
 ``Zimp`` (float = 10.0), **time-varying-scalar**
   Impurity charge state.
 
+``Aimp`` (float = 20.18), **time-varying-scalar**
+  Impurity mass in amu units.
+
 ``Zeff`` (float = 1.0), **time-varying-scalar**
   Plasma effective charge, defined as :math:`Z_{eff}=\sum_i Z_i^2 \hat{n}_i`, where :math:`\hat{n}_i` is
   the normalized ion density :math:`n_i/n_e`. For a given :math:`Z_{eff}` and :math:`Z_{imp}`, a consistent :math:`\hat{n}_i` is calculated,
@@ -210,8 +213,8 @@ Profile conditions
 
 Configures boundary conditions, initial conditions, and prescribed time-dependence of temperature, density, and current.
 
-``Ip`` (float = 15.0), **time-varying-scalar**
-  Plasma current in MA. Boundary condition for the :math:`\psi` equation.
+``Ip_tot`` (float = 15.0), **time-varying-scalar**
+  Total plasma current in MA. Boundary condition for the :math:`\psi` equation.
 
 ``Ti_bound_right`` (float | None [default]), **time-varying-scalar**
   Ion temperature boundary condition at :math:`\hat{\rho}=1` in units of keV.
@@ -819,7 +822,7 @@ A utility source module that allows for a time dependent Gaussian ion and electr
   Gaussian width of source profile in units of :math:`\hat{\rho}`.
 
 ``Ptot`` (float = 120e6), **time-varying-scalar**
-  Total source power in MW.
+  Total injected source power in W.
 
 ``el_heat_fraction`` (float = 0.66666), **time-varying-scalar**
   Electron heating fraction.
@@ -976,6 +979,35 @@ By default, both the manual and Gaussian profiles are zero. The manual and Gauss
     ``cd_efficiency`` **time-varying-scalar**
         Dimensionless local efficiency profile for conversion of EC power to current.
 
+ion_cyclotron_source
+^^^^^^^^^^^^^^^^^^^^
+Ion cyclotron heating using a surrogate model of the TORIC ICRH spectrum
+solver simulation https://meetings.aps.org/Meeting/DPP24/Session/NP12.106.
+This source is currently SPARC specific.
+
+Weights and configuration for the surrogate model are needed to use this source.
+By default these are expected to be found under
+``'~/toric_surrogate/TORIC_MLP_v1/toricnn.json'``. To use a different file path
+an alternative path can be provided using the ``TORIC_NN_MODEL_PATH``
+environment variable which should point to a compatible JSON file.
+
+``mode`` (str = 'model')
+
+``wall_inner`` (float = 1.24)
+  Inner radial location of first wall at plasma midplane level [m].
+
+``wall_outer`` (float = 2.43)
+  Outer radial location of first wall at plasma midplane level [m].
+
+``frequency`` (float = 120e6) **time-varying-scalar**
+  ICRF wave frequency in Hz.
+
+``minority_concentration`` (float = 3.0) **time-varying-scalar**
+  Helium-3 minority concentration relative to the electron density in %.
+
+``Ptot`` (float = 120e6), **time-varying-scalar**
+  Total injected source power in W.
+
 See :ref:`physics_models` for more detail.
 
 stepper
@@ -1156,7 +1188,7 @@ The configuration file is also available in ``torax/examples/iterhybrid_rampup.p
               'Zimp': 10,
           },
           'profile_conditions': {
-              'Ip': {0: 3, 80: 10.5},
+              'Ip_tot': {0: 3, 80: 10.5},
               # initial condition ion temperature for r=0 and r=Rmin
               'Ti': {0.0: {0.0: 6.0, 1.0: 0.1}},
               'Ti_bound_right': 0.1,  # boundary condition ion temp for r=Rmin
@@ -1287,6 +1319,10 @@ previous simulation and then run the simulation from that point in time. For
 all subsequent steps the dynamic runtime parameters will be constructed using
 the given runtime parameter configuration (from ``t=10`` onwards).
 
+If the requested restart time is not exactly available in the state file, the
+simulation will restart from the closest available time. A warning will be
+logged in this case.
+
 We envisage this feature being useful for example to:
 
 * restart a(n expensive) simulation that was healthy up till a certain time and

docs/physics_models.rst

@@ -420,3 +420,23 @@ Bremsstrahlung
 
 Uses the model from Wesson, John, and David J. Campbell. Tokamaks. Vol. 149.
 An optional correction for relativistic effects from Stott PPCF 2005 can be enabled with the flag "use_relativistic_correction".
+
+Ion Cyclotron Resonance
+^^^^^^^^^^^^^^^^^^^^^^^
+
+Presently this source is implemented for a SPARC specific ICRH scenario.
+
+A core Ion Cyclotron Range of Frequencies (ICRF) heating surrogate model trained
+on TORIC ICRH spectrum solver simulation
+https://meetings.aps.org/Meeting/DPP24/Session/NP12.106 is used to provide power
+profiles for Helium-3, Tritium (via its second harmonic) and electrons.
+
+A "Stix distribution" [Stix, Nuc. Fus. 1975] is used to model the non-thermal
+Helium-3 distribution based on an analytic solution to the Fokker-Planck
+equation to estimate the birth energy of Helium-3.
+
+TORAX partitions the Helium-3 power between ions and electrons using the
+parameterized model of Mikkelsen, as for Fusion Power.
+
+It is assumed that all tritium heating goes to ions and all electron heating
+goes to electrons.

pyproject.toml

@@ -71,3 +71,14 @@ include-package-data = true
 
 [tool.setuptools.packages.find]
 include = ["torax*"]
+
+[tool.pytest.ini_options]
+addopts = [
+    "--ignore=torax/tests/test_data",
+    "--ignore=torax/tests/scripts",
+    "--ignore-glob=*lib*",
+    # Leads to an ImportError in collection for `pytest` without qualikiz_tools.
+    "--ignore=torax/transport_model/qualikiz_wrapper.py",
+]
+testpaths = "**/tests"
+python_files = "*.py"

torax/__init__.py

@@ -109,4 +109,5 @@
     'tol',
     'delta_reduction_factor',
     'file_restart',
+    'ds',
 ]

torax/array_typing.py

@@ -23,7 +23,7 @@
 F = TypeVar("F", bound=Callable)
 ScalarFloat = jt.Float[jt.Array, ""]
 ScalarBool = jt.Bool[jt.Array, ""]
-ArrayFloat = jt.Bool[jt.Array, "rhon"]
+ArrayFloat = jt.Float[jt.Array, "rhon"]
 
 
 def typed(function: F) -> F:

torax/calc_coeffs.py

@@ -565,7 +565,7 @@ def _calc_coeffs_full(
   # calculate neped
   # pylint: disable=invalid-name
   nGW = (
-      dynamic_runtime_params_slice.profile_conditions.Ip
+      dynamic_runtime_params_slice.profile_conditions.Ip_tot
       / (jnp.pi * geo.Rmin**2)
       * 1e20
       / dynamic_runtime_params_slice.numerics.nref

torax/config/build_sim.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """Functions to build sim.Sim objects, which are used to run TORAX."""
-
+from collections.abc import MutableMapping
 import copy
 from typing import Any
 
@@ -184,7 +184,7 @@ def build_geometry_provider_from_config(
 
 
 def build_sim_from_config(
-    config: dict[str, Any],
+    config: MutableMapping[str, Any],
 ) -> sim_lib.Sim:
   """Builds a sim.Sim object from the given TORAX config.
 

torax/config/plasma_composition.py

@@ -47,6 +47,9 @@ class PlasmaComposition(
   Zimp: interpolated_param.TimeInterpolatedInput = (
       10.0  # impurity charge state assumed for dilution
   )
+  Aimp: interpolated_param.TimeInterpolatedInput = (
+      20.18  # impurity mass in amu
+  )
 
   def make_provider(
       self,
@@ -67,6 +70,7 @@ def make_provider(
             torax_mesh.face_centers,
         ),
         Zimp=config_args.get_interpolated_var_single_axis(self.Zimp),
+        Aimp=config_args.get_interpolated_var_single_axis(self.Aimp),
     )
 
   def __post_init__(self):
@@ -87,6 +91,7 @@ class PlasmaCompositionProvider(
   Zeff: interpolated_param.InterpolatedVarTimeRho
   Zeff_face: interpolated_param.InterpolatedVarTimeRho
   Zimp: interpolated_param.InterpolatedVarSingleAxis
+  Aimp: interpolated_param.InterpolatedVarSingleAxis
 
   def build_dynamic_params(self, t: chex.Numeric) -> DynamicPlasmaComposition:
     return DynamicPlasmaComposition(**self.get_dynamic_params_kwargs(t))
@@ -99,3 +104,4 @@ class DynamicPlasmaComposition:
   Zeff: array_typing.ArrayFloat
   Zeff_face: array_typing.ArrayFloat
   Zimp: array_typing.ScalarFloat
+  Aimp: array_typing.ScalarFloat

torax/config/profile_conditions.py

@@ -38,7 +38,7 @@ class ProfileConditions(
   # total plasma current in MA
   # Note that if Ip_from_parameters=False in geometry, then this Ip will be
   # overwritten by values from the geometry data
-  Ip: interpolated_param.TimeInterpolatedInput = 15.0
+  Ip_tot: interpolated_param.TimeInterpolatedInput = 15.0
 
   # Temperature boundary conditions at r=Rmin. If this is `None` the boundary
   # condition will instead be taken from `Ti` and `Te` at rhon=1.
@@ -187,7 +187,7 @@ class ProfileConditionsProvider(
   """Provider to retrieve initial and prescribed values and boundary conditions."""
 
   runtime_params_config: ProfileConditions
-  Ip: interpolated_param.InterpolatedVarSingleAxis
+  Ip_tot: interpolated_param.InterpolatedVarSingleAxis
   Ti_bound_right: (
       interpolated_param.InterpolatedVarSingleAxis
       | interpolated_param.InterpolatedVarTimeRho
@@ -224,7 +224,7 @@ def build_dynamic_params(
 class DynamicProfileConditions:
   """Prescribed values and boundary conditions for the core profiles."""
 
-  Ip: array_typing.ScalarFloat
+  Ip_tot: array_typing.ScalarFloat
   Ti_bound_right: array_typing.ScalarFloat
   Te_bound_right: array_typing.ScalarFloat
   # Temperature profiles defined on the cell grid.

torax/config/runtime_params_slice.py

@@ -229,6 +229,12 @@ def __init__(
     self._stepper = stepper
     self._construct_providers()
 
+  @property
+  def runtime_params_provider(
+      self,
+  ) -> general_runtime_params_lib.GeneralRuntimeParamsProvider:
+    return self._runtime_params_provider
+
   def _construct_providers(self):
     self._runtime_params_provider = (
         self._runtime_params.make_provider(
@@ -275,11 +281,11 @@ def make_ip_consistent(
       # If Ip is from parameters, renormalise psi etc to match the Ip in the
       # parameters.
       # pylint: disable=invalid-name
-      param_Ip = dynamic_runtime_params_slice.profile_conditions.Ip
-      Ip_scale_factor = param_Ip * 1e6 / geo.Ip
+      param_Ip = dynamic_runtime_params_slice.profile_conditions.Ip_tot
+      Ip_scale_factor = param_Ip * 1e6 / geo.Ip_profile_face[-1]
       geo = dataclasses.replace(
           geo,
-          Ip=geo.Ip * Ip_scale_factor,
+          Ip_profile_face=geo.Ip_profile_face * Ip_scale_factor,
           psi_from_Ip=geo.psi_from_Ip * Ip_scale_factor,
           jtot=geo.jtot * Ip_scale_factor,
           jtot_face=geo.jtot_face * Ip_scale_factor,
@@ -291,7 +297,7 @@ def make_ip_consistent(
           dynamic_runtime_params_slice,
           profile_conditions=dataclasses.replace(
               dynamic_runtime_params_slice.profile_conditions,
-              Ip=geo.Ip / 1e6,
+              Ip_tot=geo.Ip_profile_face[-1] / 1e6,
           ),
       )
   return dynamic_runtime_params_slice, geo

torax/config/tests/build_sim.py

@@ -151,7 +151,7 @@ def test_general_runtime_params_with_time_dependent_args(self):
         },
         'profile_conditions': {
             'ne_is_fGW': False,  # scalar fields.
-            'Ip': {0: 0.2, 1: 0.4, 2: 0.6},  # time-dependent.
+            'Ip_tot': {0: 0.2, 1: 0.4, 2: 0.6},  # time-dependent.
         },
         'numerics': {
             'q_correction_factor': 0.2,  # scalar fields.
@@ -177,7 +177,7 @@ def test_general_runtime_params_with_time_dependent_args(self):
         dynamic_runtime_params_slice.plasma_composition.Zeff, 0.25
     )
     np.testing.assert_allclose(
-        dynamic_runtime_params_slice.profile_conditions.Ip, 0.5
+        dynamic_runtime_params_slice.profile_conditions.Ip_tot, 0.5
     )
     np.testing.assert_allclose(
         dynamic_runtime_params_slice.numerics.resistivity_mult, 0.6
@@ -256,10 +256,12 @@ def test_build_time_dependent_geometry_from_chease(self):
   def test_chease_geometry_updates_Ip(self):
     """Tests that the Ip is updated when using chease geometry."""
     runtime_params = runtime_params_lib.GeneralRuntimeParams()
-    original_Ip = runtime_params.profile_conditions.Ip
+    original_Ip_tot = runtime_params.profile_conditions.Ip_tot
     geo_provider = build_sim.build_geometry_provider_from_config({
         'geometry_type': 'chease',
-        'Ip_from_parameters': False,  # this will force update runtime_params.Ip
+        'Ip_from_parameters': (
+            False
+        ),  # this will force update runtime_params.Ip_tot
     })
     runtime_params_provider = (
         runtime_params_slice.DynamicRuntimeParamsSliceProvider(
@@ -279,7 +281,9 @@ def test_chease_geometry_updates_Ip(self):
     )
     self.assertIsInstance(geo, geometry.StandardGeometry)
     self.assertIsNotNone(dynamic_slice)
-    self.assertNotEqual(dynamic_slice.profile_conditions.Ip, original_Ip)
+    self.assertNotEqual(
+        dynamic_slice.profile_conditions.Ip_tot, original_Ip_tot
+    )
     # pylint: enable=invalid-name
 
   def test_empty_source_config_only_has_defaults_turned_off(self):