Add a "scaled_profile" ICRH model to TORAX.

This new model allows users to provide reference ion and electron heating profiles. These profiles are then radially shifted based on the ratio of the current on-axis toroidal magnetic field (B0) to a reference B0, reflecting the change in the ICRH resonance location. The profiles are also rescaled to match a specified total absorbed power. This model is useful for approximating the adaptation of pre-computed RF heating profiles from higher-fidelity source calculations, to different magnetic field scenarios and powers. The model does not include a fast-ion calculation which are expected to be prescribed separately.

PiperOrigin-RevId: 906940622

Author: jcitrin

docs/configuration.rst

@@ -2053,46 +2053,91 @@ implemented in TORAX.
 
 icrh
 ^^^^
-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.
+Ion cyclotron resonance heating (ICRH) source. Two model implementations are
+available, selected via the ``model_name`` discriminator.
 
 ``mode`` (str [default = 'model'])
 
-``model_path`` (str | None [default = None])
-  Path to the JSON file containing the weights and configuration for the
-  surrogate model. If None, the default path
-  ``'~/toric_surrogate/TORIC_MLP_v1/toricnn.json'`` is used.
-
-``wall_inner`` (float [default = 1.24])
-  Inner radial location of first wall at plasma midplane level [m].
-
-``wall_outer`` (float [default = 2.43])
-  Outer radial location of first wall at plasma midplane level [m].
-
-``frequency`` (**time-varying-scalar** [default = 120e6])
-  ICRF wave frequency in Hz.
-
-``minority_concentration`` (**time-varying-scalar** [default = 0.03])
-  Helium-3 minority fractional concentration relative to the electron density.
-  This is a deprecated legacy option, used only when ``minority_species`` is
-  not specified. It is recommended to instead set ``minority_species`` to
-  'He3' and use the plasma composition to set the minority concentration.
-
-``minority_species`` (str | None [default = None])
-  Symbol of the minority species (e.g. 'He3'). If specified, the minority
-  concentration is automatically extracted from the plasma composition.
-  Presently, only 'He3' is supported.
-
 ``P_total`` (**time-varying-scalar** [default = 10e6])
   Total injected source power in W.
 
+``absorption_fraction`` (**time-varying-scalar** [default = 1.0])
+  Fraction of ``P_total`` absorbed by the plasma.
+
+``model_name`` (str)
+  Selects the ICRH model. Available options:
+
+* ``'toric_nn'`` (default)
+    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.
+
+    ``model_path`` (str | None [default = None])
+      Path to the JSON file containing the weights and configuration for the
+      surrogate model. If None, the default path
+      ``'~/toric_surrogate/TORIC_MLP_v1/toricnn.json'`` is used.
+
+    ``wall_inner`` (float [default = 1.24])
+      Inner radial location of first wall at plasma midplane level [m].
+
+    ``wall_outer`` (float [default = 2.43])
+      Outer radial location of first wall at plasma midplane level [m].
+
+    ``frequency`` (**time-varying-scalar** [default = 120e6])
+      ICRF wave frequency in Hz.
+
+    ``minority_concentration`` (**time-varying-scalar** [default = 0.03])
+      Helium-3 minority fractional concentration relative to the electron
+      density. This is a deprecated legacy option, used only when
+      ``minority_species`` is not specified. It is recommended to instead set
+      ``minority_species`` to 'He3' and use the plasma composition to set the
+      minority concentration.
+
+    ``minority_species`` (str | None [default = None])
+      Symbol of the minority species (e.g. 'He3'). If specified, the minority
+      concentration is automatically extracted from the plasma composition.
+      Presently, only 'He3' is supported.
+
+* ``'scaled_profile'``
+    Prescribed-profile model with magnetic-field-dependent resonance shifting.
+    Adapts reference heating profiles to different magnetic field operating points
+    and normalises them to the target power. See :ref:`physics_models` for
+    detailed physics description.
+
+    ``heat_profile_ion`` (**time-varying-array** [default = {0: {0: 0, 1: 0}}])
+      Reference ion heating power density shape [W/m³], provided on the
+      normalised radius grid. The shape is rescaled during normalisation; only
+      the relative profile matters.
+
+    ``heat_profile_electron`` (**time-varying-array** [default = {0: {0: 0, 1: 0}}])
+      Reference electron heating power density shape [W/m³], provided on the
+      normalised radius grid.
+
+    ``reference_B0`` (**time-varying-scalar** [default = 12.2])
+      Vacuum toroidal magnetic field at which the reference profiles were
+      computed [T].
+
+    Example:
+
+    .. code-block:: python
+
+        'sources': {
+            'icrh': {
+                'model_name': 'scaled_profile',
+                'P_total': 25e6,
+                'absorption_fraction': 0.9,
+                'heat_profile_ion': {0: {0.0: 0, 0.3: 1.0, 0.6: 0.2, 1.0: 0}},
+                'heat_profile_electron': {0: {0.0: 0, 0.35: 0.8, 0.65: 0.1, 1.0: 0}},
+                'reference_B0': 12.2,
+            },
+        }
+
 See :ref:`physics_models` for more detail.
 
 impurity_radiation

docs/physics_models.rst

@@ -712,7 +712,17 @@ The following models are available:
 Ion Cyclotron Resonance Heating
 -------------------------------
 
-Presently this source is implemented for a SPARC specific ICRH scenario.
+TORAX supports ICRH through two models, selectable via the ``model_name``
+discriminator in the ``icrh`` source configuration.
+
+**Common parameters** shared by all ICRH models are the total injected power
+:math:`P_\mathrm{total}` and the absorption fraction :math:`\alpha`. The
+absorbed power is :math:`P_\mathrm{abs} = P_\mathrm{total} \times \alpha`.
+
+ToricNN Surrogate (``toric_nn``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This is the default ICRH model, currently specific to the SPARC tokamak.
 
 A core Ion Cyclotron Range of Frequencies (ICRF) heating surrogate model trained
 on TORIC ICRH spectrum solver simulations is used to provide power profiles for
@@ -728,6 +738,42 @@ 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.
 
+Scaled Profile (``scaled_profile``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+A lightweight model that takes user-prescribed reference heating profiles
+and adapts them to different magnetic field operating points using a
+physically-motivated radial shift.
+
+The model performs two operations:
+
+1. **Resonance shift:** The ICRH resonance occurs where
+   :math:`\omega = n \omega_{ci}(R)`, and since the vacuum toroidal field scales
+   as :math:`B_t \propto 1/R`, the resonance radius scales linearly with
+   the on-axis field: :math:`R_\mathrm{res} \propto B_0`. When the actual
+   :math:`B_0` differs from the reference field :math:`B_{0,\mathrm{ref}}`, the
+   heating profile shifts in normalised radius space.
+
+   For each grid point, the shifted outboard midplane radius is:
+
+   .. math::
+
+     R'_\mathrm{out}(\rho) = R_\mathrm{out}(\rho) \cdot \frac{B_0}{B_{0,\mathrm{ref}}}
+
+   where :math:`R_\mathrm{out}(\rho) = R_\mathrm{major} + r(\rho)` is the
+   outboard midplane radius, which is monotonically increasing with
+   :math:`\rho`. The reference profiles are then evaluated at the :math:`\rho`
+   coordinate that corresponds to :math:`R'_\mathrm{out}` on the original
+   :math:`R_\mathrm{out}(\rho)` mapping.
+
+2. **Power normalisation:** The shifted ion and electron heating profiles are
+   rescaled so that the volume-integrated total heating equals
+   :math:`P_\mathrm{total} \times \alpha`.
+
+This model does not include a fast-ion calculation; all fast-ion outputs are
+zero. It is useful when reference profiles from a full-wave RF code are
+available and need to be quickly approximated for different field scenarios.
+
 Fast Ion Physics
 ================
 

torax/_src/sources/ion_cyclotron_source/scaled_profile.py

@@ -0,0 +1,186 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Scaled-profile ICRH model with magnetic-field-dependent resonance shift."""
+
+import dataclasses
+from typing import Annotated, Literal
+
+import chex
+import jax
+from jax import numpy as jnp
+from torax._src import array_typing
+from torax._src import math_utils
+from torax._src import state
+from torax._src.config import runtime_params as runtime_params_lib
+from torax._src.geometry import geometry
+from torax._src.neoclassical.conductivity import base as conductivity_base
+from torax._src.physics import fast_ion as fast_ion_lib
+from torax._src.sources import runtime_params as source_runtime_params_lib
+from torax._src.sources import source
+from torax._src.sources import source_profiles
+from torax._src.sources.ion_cyclotron_source import base
+from torax._src.torax_pydantic import torax_pydantic
+
+# pylint: disable=invalid-name
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class RuntimeParams(source_runtime_params_lib.RuntimeParams):
+  """Runtime parameters for the scaled-profile ICRH model."""
+
+  P_total: array_typing.FloatScalar
+  absorption_fraction: array_typing.FloatScalar
+  heat_profile_ion: array_typing.FloatVector
+  heat_profile_electron: array_typing.FloatVector
+  reference_B0: array_typing.FloatScalar
+
+
+def scaled_profile_model_func(
+    runtime_params: runtime_params_lib.RuntimeParams,
+    geo: geometry.Geometry,
+    source_name: str,
+    core_profiles: state.CoreProfiles,
+    unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
+    unused_conductivity: conductivity_base.Conductivity | None,
+) -> tuple[
+    array_typing.FloatVectorCell,
+    array_typing.FloatVectorCell,
+    tuple[fast_ion_lib.FastIon, ...],
+]:
+  """Compute ICRH heating from prescribed profiles with B-field shift.
+
+  The model performs two operations on the reference profiles:
+  1. **Radial shift**: The ICRH resonance location in major radius scales as
+     R_res ∝ B₀. When B₀ differs from the reference field, the resonance
+     moves, and the heating profile shifts accordingly in normalised radius.
+  2. **Power normalisation**: The shifted profiles are rescaled so the
+     volume-integrated total heating equals ``P_total * absorption_fraction``.
+
+  Args:
+    runtime_params: Full simulation runtime parameters.
+    geo: Magnetic geometry.
+    source_name: Name of this source (used to look up params).
+    core_profiles: Core plasma profiles (unused by this model).
+    unused_calculated_source_profiles: Not used.
+    unused_conductivity: Not used.
+
+  Returns:
+    Tuple of (ion_heating, electron_heating, fast_ions) where fast_ions are
+    all zeros (no fast-ion model in this mode).
+  """
+  del core_profiles  # Unused.
+  source_params = runtime_params.sources[source_name]
+  assert isinstance(source_params, RuntimeParams)
+
+  ref_ion = source_params.heat_profile_ion
+  ref_el = source_params.heat_profile_electron
+
+  # --- 1. Compute resonance shift ---
+  # The ICRH resonance occurs where ω = n·ω_ci(R) and since B_t ∝ 1/R,
+  # the resonance major radius scales linearly with B₀.
+  # B_ratio > 1 means stronger field → resonance moves outward in R.
+  B_ratio = geo.B_0 / source_params.reference_B0
+
+  # Outboard midplane major radius on the cell grid: R_out(ρ) = R_major + r(ρ).
+  # This is monotonically increasing with ρ, unlike the flux-surface-averaged
+  # R_major_profile which can be constant (e.g. circular geometry).
+  R_out = geo.R_out
+  rho = geo.torax_mesh.cell_centers
+
+  # Shifted major radius for each grid point.
+  R_shifted = R_out * B_ratio
+
+  # Map back to normalised radius: for each shifted R, find the
+  # corresponding ρ on the original R_out(ρ) curve.
+  rho_shifted = jnp.interp(R_shifted, R_out, rho)
+
+  # Evaluate reference profiles at the shifted ρ positions.
+  shifted_ion = jnp.interp(rho_shifted, rho, ref_ion)
+  shifted_el = jnp.interp(rho_shifted, rho, ref_el)
+
+  # --- 2. Normalise to target power ---
+  absorbed_power = source_params.P_total * source_params.absorption_fraction
+  total_shape = shifted_ion + shifted_el
+  integrated = math_utils.volume_integration(total_shape, geo)
+  # Guard against zero integrated power (e.g. if profiles are all zero).
+  # Use safe denominator to avoid NaN gradients in dead jnp.where branches.
+  safe_integrated = jnp.where(integrated > 0, integrated, 1.0)
+  scale = jnp.where(integrated > 0, absorbed_power / safe_integrated, 0.0)
+
+  source_ion = shifted_ion * scale
+  source_el = shifted_el * scale
+
+  # --- 3. Default zero fast ions in build_fast_ions ---
+  fast_ions = base.build_fast_ions(source_name=source_name, geo=geo)
+
+  return (source_ion, source_el, fast_ions)
+
+
+class ScaledProfileIonCyclotronSourceConfig(base.IonCyclotronSourceConfig):
+  """Configuration for ICRH with prescribed, B-field-shiftable profiles.
+
+  This model takes reference ion and electron heating profiles and:
+  1. Shifts them radially based on the ratio of the actual vacuum toroidal
+     magnetic field to a reference field (``B₀ / reference_B0``).
+  2. Rescales the amplitude so that the volume-integrated total heating
+     equals ``P_total * absorption_fraction``.
+
+  This is useful when computed reference heating profiles are available at a
+  specific magnetic field, and you need to approximate them to different
+  operating points without re-running the full RF solver.
+
+  Attributes:
+    model_name: Discriminator literal for Pydantic.
+    heat_profile_ion: Reference ion heating power density shape [W/m³],
+      provided on the normalised radius grid.
+    heat_profile_electron: Reference electron heating power density shape
+      [W/m³], provided on the normalised radius grid.
+    reference_B0: Vacuum toroidal magnetic field at which the reference
+      profiles were computed [T].
+  """
+
+  model_name: Annotated[
+      Literal['scaled_profile'], torax_pydantic.JAX_STATIC
+  ] = 'scaled_profile'
+  heat_profile_ion: torax_pydantic.TimeVaryingArray = (
+      torax_pydantic.ValidatedDefault({0: {0: 0, 1: 0}})
+  )
+  heat_profile_electron: torax_pydantic.TimeVaryingArray = (
+      torax_pydantic.ValidatedDefault({0: {0: 0, 1: 0}})
+  )
+  reference_B0: torax_pydantic.TimeVaryingScalar = (
+      torax_pydantic.ValidatedDefault(12.2)
+  )
+
+  @property
+  def model_func(self) -> source.SourceProfileFunction:
+    return scaled_profile_model_func
+
+  def build_runtime_params(
+      self,
+      t: chex.Numeric,
+  ) -> RuntimeParams:
+    return RuntimeParams(
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
+        mode=self.mode,
+        is_explicit=self.is_explicit,
+        P_total=self.P_total.get_value(t),
+        absorption_fraction=self.absorption_fraction.get_value(t),
+        heat_profile_ion=self.heat_profile_ion.get_value(t),
+        heat_profile_electron=self.heat_profile_electron.get_value(t),
+        reference_B0=self.reference_B0.get_value(t),
+    )