diff --git a/docs/configuration.rst b/docs/configuration.rst
index 6b29f10b6..b6cc54312 100644
--- a/docs/configuration.rst
+++ b/docs/configuration.rst
@@ -1469,58 +1469,6 @@ transport model.
Width of HWHM Gaussian smoothing kernel operating on transport model outputs.
If using the ``QLKNN_7_11`` transport model, the default is set to 0.1.
-``smooth_everywhere`` (bool [default = False])
- Smooth across entire radial domain regardless of inner and outer patches.
-
-``apply_inner_patch`` (**time-varying-scalar** [default = False])
- If ``True``, set a patch for inner core transport coefficients below
- ``rho_inner``. Typically used as an ad-hoc measure for MHD (e.g. sawteeth) or
- EM (e.g. KBM) transport in the inner-core. If using a
- `CombinedTransportModel`, ensure that the inner patch is only set on the
- global model rather than its component models to avoid conflicts.
-
-``D_e_inner`` (**time-varying-scalar** [default = 0.2])
- Particle diffusivity value for inner transport patch.
-
-``V_e_inner`` (**time-varying-scalar** [default = 0.0])
- Particle convection value for inner transport patch.
-
-``chi_i_inner`` (**time-varying-scalar** [default = 1.0])
- Ion heat conduction value for inner transport patch.
-
-``chi_e_inner`` (**time-varying-scalar** [default = 1.0])
- Electron heat conduction value for inner transport patch.
-
-``rho_inner`` (**time-varying-scalar** [default = 0.3])
- :math:`\hat{\rho}` below which inner patch is applied.
- Note that ``rho_inner`` and ``rho_outer`` must have the same interpolation
- mode to simplify the validation test ``rho_inner < rho_outer`` at all times.
-
-``apply_outer_patch`` (**time-varying-scalar** [default = False])
- If ``True``, set a patch for outer core transport coefficients above
- ``rho_outer``. Useful for the L-mode near-edge region where models like
- QLKNN10D are not applicable. Only used if ``set_pedestal==False``.
- If using a `CombinedTransportModel`, ensure that the outer patch is
- only set on the global model rather than its component models to avoid
- conflicts.
-
-``D_e_outer`` (**time-varying-scalar** [default = 0.2])
- Particle diffusivity value for outer transport patch.
-
-``V_e_outer`` (**time-varying-scalar** [default = 0.0])
- Particle convection value for outer transport patch.
-
-``chi_i_outer`` (**time-varying-scalar** [default = 1.0])
- Ion heat conduction value for outer transport patch.
-
-``chi_e_outer`` (**time-varying-scalar** [default = 1.0])
- Electron heat conduction value for outer transport patch.
-
-``rho_outer`` (**time-varying-scalar** [default = 0.9])
- :math:`\hat{\rho}` above which outer patch is applied.
- Note that ``rho_inner`` and ``rho_outer`` must have the same interpolation
- mode to simplify the validation test ``rho_inner < rho_outer`` at all times.
-
``fast_ion_stabilization`` (**time-varying-scalar** [default = False])
If ``True``, apply a fast ion stabilization correction to the :math:`R/L_{Ti}`
input of quasilinear transport models (QLKNN, TGLFNN, QuaLiKiz). The fast ion
@@ -3003,18 +2951,6 @@ CHEASE geometry), is shown below. The configuration file is also available in
},
'transport': {
'model_name': 'qlknn',
- 'apply_inner_patch': True,
- 'D_e_inner': 0.25,
- 'V_e_inner': 0.0,
- 'chi_i_inner': 1.5,
- 'chi_e_inner': 1.5,
- 'rho_inner': 0.3,
- 'apply_outer_patch': True,
- 'D_e_outer': 0.1,
- 'V_e_outer': 0.0,
- 'chi_i_outer': 2.0,
- 'chi_e_outer': 2.0,
- 'rho_outer': 0.9,
'chi_min': 0.05,
'chi_max': 100,
'D_e_min': 0.05,
diff --git a/torax/_src/config/runtime_params.py b/torax/_src/config/runtime_params.py
index 72aae9548..b7a289109 100644
--- a/torax/_src/config/runtime_params.py
+++ b/torax/_src/config/runtime_params.py
@@ -86,7 +86,7 @@ class RuntimeParams:
profile_conditions: profile_conditions.RuntimeParams
solver: solver_params.RuntimeParams
sources: Mapping[str, sources_params.RuntimeParams]
- transport: transport_model_params.RuntimeParams
+ transport: transport_model_params.CombinedRuntimeParams
time_step_calculator: time_step_calculator_runtime_params.RuntimeParams
diff --git a/torax/_src/transport_model/combined.py b/torax/_src/transport_model/combined.py
index f45dd448d..a8e0c059b 100644
--- a/torax/_src/transport_model/combined.py
+++ b/torax/_src/transport_model/combined.py
@@ -17,12 +17,12 @@
A class for combining transport models.
"""
+# pylint: disable=invalid-name
+
import dataclasses
-from typing import Callable, Sequence, Tuple
-import chex
+from typing import Callable, Sequence
import jax
import jax.numpy as jnp
-from torax._src import array_typing
from torax._src import constants
from torax._src import jax_utils
from torax._src import state
@@ -37,23 +37,8 @@
MIN_SMOOTHING_WIDTH = 1e-5
-@chex.dataclass
-class SmoothingZoneParams:
- rho_min: array_typing.FloatScalar
- rho_max: array_typing.FloatScalar
- smoothing_width: array_typing.FloatScalar
-
-
-@jax.tree_util.register_dataclass
-@dataclasses.dataclass(frozen=True)
-class RuntimeParams(transport_runtime_params_lib.RuntimeParams):
- """Runtime parameters for the CombinedTransportModel."""
-
- transport_model_params: Tuple[transport_runtime_params_lib.RuntimeParams, ...]
- pedestal_transport_model_params: Tuple[
- transport_runtime_params_lib.RuntimeParams, ...
- ]
- smoothing_zones: Tuple[SmoothingZoneParams, ...]
+SmoothingZoneParams = transport_runtime_params_lib.SmoothingZoneParams
+RuntimeParams = transport_runtime_params_lib.CombinedRuntimeParams
@dataclasses.dataclass(frozen=True, eq=False)
@@ -82,23 +67,12 @@ def __call__(
core_profiles,
pedestal_model_output,
)
-
- # In contrast to the base TransportModel, we do not apply domain restriction
- # or output masking (enabled/disabled channels) as these are handled at the
- # component model level in call_implementation here.
-
# Apply min/max clipping
transport_coeffs = self._apply_clipping(
transport_runtime_params,
transport_coeffs,
)
- # In contrast to the base TransportModel, we do not apply patches, as these
- # should be handled by instantiating constant component models instead.
- # However, the rho_inner and rho_outer arguments are currently required
- # in the case where the inner/outer region are to be excluded from
- # smoothing.
-
transport_coeffs = self._smooth_coeffs(
runtime_params,
geo,
@@ -108,19 +82,22 @@ def __call__(
return transport_coeffs
+ # pytype: disable=signature-mismatch
def call_implementation(
self,
- transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
+ transport_runtime_params: transport_runtime_params_lib.CombinedRuntimeParams,
runtime_params: runtime_params_lib.RuntimeParams,
geo: geometry.Geometry,
core_profiles: state.CoreProfiles,
pedestal_model_output: pedestal_model_output_lib.PedestalModelOutput,
) -> transport_model_lib.TurbulentTransport:
+ # pytype: enable=signature-mismatch
r"""Calculates transport coefficients using the Combined model.
Args:
transport_runtime_params: Input runtime parameters for this transport
- model. Can change without triggering a JAX recompilation.
+ model (expected to be an instance of combined.RuntimeParams at runtime).
+ Can change without triggering a JAX recompilation.
runtime_params: Runtime parameters for the simulation at the current time.
geo: Geometry of the torus at the current time.
core_profiles: Core plasma profiles.
@@ -255,6 +232,42 @@ def _combine(
return transport_model_lib.TurbulentTransport(**accumulators)
+ def _apply_clipping(
+ self,
+ transport_runtime_params: transport_runtime_params_lib.CombinedRuntimeParams,
+ transport_coeffs: transport_model_lib.TurbulentTransport,
+ ) -> transport_model_lib.TurbulentTransport:
+ """Applies min/max clipping to transport coefficients for PDE stability."""
+ assert isinstance(transport_runtime_params, RuntimeParams)
+ chi_face_ion = jnp.clip(
+ transport_coeffs.chi_face_ion,
+ transport_runtime_params.chi_min,
+ transport_runtime_params.chi_max,
+ )
+ chi_face_el = jnp.clip(
+ transport_coeffs.chi_face_el,
+ transport_runtime_params.chi_min,
+ transport_runtime_params.chi_max,
+ )
+ d_face_el = jnp.clip(
+ transport_coeffs.d_face_el,
+ transport_runtime_params.D_e_min,
+ transport_runtime_params.D_e_max,
+ )
+ v_face_el = jnp.clip(
+ transport_coeffs.v_face_el,
+ transport_runtime_params.V_e_min,
+ transport_runtime_params.V_e_max,
+ )
+
+ return dataclasses.replace(
+ transport_coeffs,
+ chi_face_ion=chi_face_ion,
+ chi_face_el=chi_face_el,
+ d_face_el=d_face_el,
+ v_face_el=v_face_el,
+ )
+
def _smooth_coeffs(
self,
runtime_params: runtime_params_lib.RuntimeParams,
diff --git a/torax/_src/transport_model/pydantic_model.py b/torax/_src/transport_model/pydantic_model.py
index 631232e12..2a833ae73 100644
--- a/torax/_src/transport_model/pydantic_model.py
+++ b/torax/_src/transport_model/pydantic_model.py
@@ -16,7 +16,6 @@
import copy
import dataclasses
-import itertools
from typing import Annotated, Any, Literal, Sequence
from absl import logging
import chex
@@ -170,9 +169,6 @@ def _conform_data(cls, data: Any) -> Any:
# The QLK version this specific QLKNN was trained on tends to
# underpredict ITG electron heat flux in shaped, high-beta scenarios.
data['ITG_flux_ratio_correction'] = 2.0
- else:
- if 'smoothing_width' not in data:
- data['smoothing_width'] = 0.1
return data
def build_transport_model(self) -> qlknn_transport_model.QLKNNTransportModel:
@@ -457,14 +453,19 @@ class SmoothingZone(torax_pydantic.BaseModelFrozen):
smoothing_width: pydantic.NonNegativeFloat
-class CombinedTransportModel(pydantic_model_base.TransportBase):
+class CombinedTransportModel(torax_pydantic.BaseModelFrozen):
"""Model for the Combined transport model.
- Note: smoothing and patches should be applied on the combined model, not the
- individual component models.
-
Attributes:
model_name: The transport model to use. Hardcoded to 'combined'.
+ chi_min: Lower bound on heat conductivity.
+ chi_max: Upper bound on heat conductivity (can be helpful for stability).
+ D_e_min: minimum electron density diffusivity.
+ D_e_max: maximum electron density diffusivity.
+ V_e_min: minimum electron density convection.
+ V_e_max: maximum electron density convection.
+ smoothing_width: Width of HWHM Gaussian smoothing kernel operating on
+ transport model outputs.
transport_models: A sequence of transport models, whose outputs will be
summed to give the combined core transport coefficients.
pedestal_transport_models: A sequence of models that will be combined for
@@ -475,6 +476,13 @@ class CombinedTransportModel(pydantic_model_base.TransportBase):
means that zone will not be used for the smoothing of other zones.
"""
+ chi_min: torax_pydantic.MeterSquaredPerSecond = 0.05
+ chi_max: torax_pydantic.MeterSquaredPerSecond = 100.0
+ D_e_min: torax_pydantic.MeterSquaredPerSecond = 0.05
+ D_e_max: torax_pydantic.MeterSquaredPerSecond = 100.0
+ V_e_min: torax_pydantic.MeterPerSecond = -50.0
+ V_e_max: torax_pydantic.MeterPerSecond = 50.0
+ smoothing_width: pydantic.NonNegativeFloat = 0.0
# TODO(b/434175938) V2: rename `transport_models` to `core_transport_models`
transport_models: Sequence[CombinedCompatibleTransportModel] = pydantic.Field(
default_factory=list
@@ -506,7 +514,6 @@ def build_transport_model(self) -> combined.CombinedTransportModel:
)
def build_runtime_params(self, t: chex.Numeric) -> combined.RuntimeParams:
- base_kwargs = dataclasses.asdict(super().build_runtime_params(t))
transport_model_params = tuple(
model.build_runtime_params(t) for model in self.transport_models
)
@@ -525,28 +532,18 @@ def build_runtime_params(self, t: chex.Numeric) -> combined.RuntimeParams:
)
)
return combined.RuntimeParams(
+ chi_min=self.chi_min,
+ chi_max=self.chi_max,
+ D_e_min=self.D_e_min,
+ D_e_max=self.D_e_max,
+ V_e_min=self.V_e_min,
+ V_e_max=self.V_e_max,
+ smoothing_width=self.smoothing_width,
transport_model_params=transport_model_params,
pedestal_transport_model_params=pedestal_transport_model_params,
smoothing_zones=tuple(smoothing_zones),
- **base_kwargs,
)
- @pydantic.model_validator(mode='after')
- def _check_no_smoothing_in_components(self) -> typing_extensions.Self:
- for model_list in ['transport_models', 'pedestal_transport_models']:
- for i, model in enumerate(getattr(self, model_list)):
- if model.smoothing_width > 0.0:
- logging.warning(
- 'smoothing_width > 0.0 is not supported for component models of'
- ' CombinedTransportModel; instead, smoothing_width should be set'
- ' on the CombinedTransportModel itself. Smoothing width set on %s'
- ' component %i (%s) will be ignored.',
- model_list,
- i,
- model.model_name,
- )
- return self
-
@pydantic.model_validator(mode='after')
def _check_smoothing_width_minimum(self) -> typing_extensions.Self:
smoothing_widths = [
@@ -566,29 +563,24 @@ def _check_smoothing_width_minimum(self) -> typing_extensions.Self:
@pydantic.model_validator(mode='after')
def _check_fields(self) -> typing_extensions.Self:
- super()._check_fields() # pyrefly: ignore[not-callable]
- if (
- any([
- np.any(model.apply_inner_patch.value)
- or np.any(model.apply_outer_patch.value)
- # Use itertools.chain to iterate over both lists of models without
- # needing to make a new list.
- for model in itertools.chain(
- self.transport_models, self.pedestal_transport_models
- )
- ])
- or np.any(self.apply_inner_patch.value)
- or np.any(self.apply_outer_patch.value)
- ):
- raise ValueError(
- 'apply_inner_patch and apply_outer_patch not supported for'
- ' CombinedTransportModel or its component models.'
- )
- if np.any(self.rho_min.value != 0.0) or np.any(self.rho_max.value != 1.0):
- raise ValueError(
- 'rho_min and rho_max should not be set for CombinedTransportModel, as'
- ' it should be applied across the whole rho domain.'
+ if not self.chi_min < self.chi_max:
+ raise ValueError('chi_min must be less than chi_max.')
+ if not self.D_e_min < self.D_e_max:
+ raise ValueError('D_e_min must be less than D_e_max.')
+ if not self.V_e_min < self.V_e_max:
+ raise ValueError('V_e_min must be less than V_e_max.')
+ if self.smoothing_width == 0.0:
+ has_qlknn = any(
+ isinstance(m, qlknn_transport_model.QLKNNTransportModel)
+ for m in list(self.transport_models)
+ + list(self.pedestal_transport_models)
)
+ if has_qlknn:
+ logging.warning(
+ 'QLKNN transport model is configured in CombinedTransportModel'
+ ' with smoothing_width=0. Stiff QLKNN transport coefficients'
+ ' without spatial smoothing may degrade solver convergence.'
+ )
if any([
np.any(model.rho_min.value != 0.0) or np.any(model.rho_max.value != 1.0)
for model in self.pedestal_transport_models
diff --git a/torax/_src/transport_model/pydantic_model_base.py b/torax/_src/transport_model/pydantic_model_base.py
index a5da19131..20a1ad7aa 100644
--- a/torax/_src/transport_model/pydantic_model_base.py
+++ b/torax/_src/transport_model/pydantic_model_base.py
@@ -32,96 +32,28 @@ class TransportBase(torax_pydantic.BaseModelFrozen, abc.ABC):
"""Base model holding parameters common to all transport models.
Attributes:
- chi_min: Lower bound on heat conductivity.
- chi_max: Upper bound on heat conductivity (can be helpful for stability).
- D_e_min: minimum electron density diffusivity.
- D_e_max: maximum electron density diffusivity.
- V_e_min: minimum electron density convection.
- V_e_max: minimum electron density convection.
rho_min: normalized radius above which this model is applied.
rho_max: normalized radius below which this model is applied.
- apply_inner_patch: set inner core transport coefficients (ad-hoc MHD/EM
- transport).
- D_e_inner: inner core electron density diffusivity.
- V_e_inner: inner core electron density convection.
- chi_i_inner: inner core ion heat equation diffusion term.
- chi_e_inner: inner core electron heat equation diffusion term.
- rho_inner: normalized radius below which inner patch is applied.
- apply_outer_patch: set outer core transport coefficients (ad-hoc MHD/EM
- transport). Only used when pedestal.set_pedestal = False Useful for L-mode
- near-edge region where QLKNN10D is not applicable.
- D_e_outer: outer core electron density diffusivity.
- V_e_outer: outer core electron density convection.
- chi_i_outer: outer core ion heat equation diffusion term.
- chi_e_outer: outer core electron heat equation diffusion term.
- rho_outer: normalized radius above which outer patch is applied.
- smoothing_width: Width of HWHM Gaussian smoothing kernel operating on
- transport model outputs.
- smooth_everywhere: Smooth across entire radial domain regardless of inner
- and outer patches.
disable_chi_i: If True, sets the ion heat conductivity output to zero.
disable_chi_e: If True, sets the electron heat conductivity output to zero.
disable_D_e: If True, sets the electron diffusivity output to zero.
disable_V_e: If True, sets the electron convection output to zero.
- merge_mode: Defines how this model is combined with previous models in a
+ fast_ion_stabilization: If True, applies fast ion stabilization.
+ fast_ion_stabilization_model: Fast ion stabilization model config.
+ fast_ion_stabilization_multiplier: Fast ion stabilization multiplier.
+ merge_mode: Defines how transport coefficients are combined within a
CombinedTransportModel. 'add' (default) adds to the accumulated value.
'overwrite' overwrites the previous value in this model's valid domain
and prevents subsequent 'add' models in the sequence from modifying this
region.
"""
- chi_min: torax_pydantic.MeterSquaredPerSecond = 0.05
- chi_max: torax_pydantic.MeterSquaredPerSecond = 100.0
- D_e_min: torax_pydantic.MeterSquaredPerSecond = 0.05
- D_e_max: torax_pydantic.MeterSquaredPerSecond = 100.0
- V_e_min: torax_pydantic.MeterPerSecond = -50.0
- V_e_max: torax_pydantic.MeterPerSecond = 50.0
rho_min: torax_pydantic.UnitIntervalTimeVaryingScalar = (
torax_pydantic.ValidatedDefault(0.0)
)
rho_max: torax_pydantic.UnitIntervalTimeVaryingScalar = (
torax_pydantic.ValidatedDefault(1.0)
)
- # TODO(b/434175938): Remove patch mechanism in V2 due to duplication with
- # combined transport model framework.
- apply_inner_patch: interpolated_param_1d.TimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(False)
- )
- D_e_inner: torax_pydantic.PositiveTimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(0.2)
- )
- V_e_inner: interpolated_param_1d.TimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(0.0)
- )
- chi_i_inner: torax_pydantic.PositiveTimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(1.0)
- )
- chi_e_inner: torax_pydantic.PositiveTimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(1.0)
- )
- rho_inner: torax_pydantic.UnitIntervalTimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(0.3)
- )
- apply_outer_patch: interpolated_param_1d.TimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(False)
- )
- D_e_outer: interpolated_param_1d.PositiveTimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(0.2)
- )
- V_e_outer: interpolated_param_1d.TimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(0.0)
- )
- chi_i_outer: interpolated_param_1d.PositiveTimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(1.0)
- )
- chi_e_outer: interpolated_param_1d.PositiveTimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(1.0)
- )
- rho_outer: torax_pydantic.UnitIntervalTimeVaryingScalar = (
- torax_pydantic.ValidatedDefault(0.9)
- )
- smoothing_width: pydantic.NonNegativeFloat = 0.0
- smooth_everywhere: bool = False
disable_chi_i: interpolated_param_1d.TimeVaryingScalar = (
torax_pydantic.ValidatedDefault(False)
)
@@ -148,34 +80,17 @@ class TransportBase(torax_pydantic.BaseModelFrozen, abc.ABC):
@pydantic.model_validator(mode='after')
def _check_fields(self) -> typing_extensions.Self:
- if not self.chi_max > self.chi_min:
- raise ValueError('chi_min must be less than chi_max.')
- if not self.D_e_min < self.D_e_max:
- raise ValueError('D_e_min must be less than D_e_max.')
- if not self.V_e_min < self.V_e_max:
- raise ValueError('V_e_min must be less than V_e_max.')
- # For the time-varying parameter pairs (rho_min, rho_max),
- # (rho_inner, rho_outer), we have relative magnitude constraints. These need
- # to be held at all times. We validate this by checking the inequality at
- # the combined time points (knots) of each pair. This is sufficient
- # both for STEP and PIECEWISE_LINEAR interpolation modes. However, if the
- # interpolation modes are mixed, the constraint check becomes more
- # complicated. For now, we only support the same interpolation mode for
- # these pairs.
- if self.rho_outer.interpolation_mode != self.rho_inner.interpolation_mode:
- raise ValueError(
- 'rho_outer and rho_inner must have the same interpolation mode.'
- )
+ # For the time-varying parameter pair (rho_min, rho_max), we have relative
+ # magnitude constraints that must hold at all times. We validate this by
+ # checking the inequality at the combined time points (knots) of the pair.
+ # This check is sufficient to guarantee the inequality at all times
+ # provided that both parameters use the same interpolation mode (either step
+ # or linear interpolation). We therefore require their interpolation modes
+ # to match.
if self.rho_max.interpolation_mode != self.rho_min.interpolation_mode:
raise ValueError(
'rho_max and rho_min must have the same interpolation mode.'
)
- all_times_inner_outer = np.union1d(self.rho_inner.time, self.rho_outer.time)
- if not np.all(
- self.rho_outer.get_value(all_times_inner_outer)
- > self.rho_inner.get_value(all_times_inner_outer)
- ):
- raise ValueError('rho_outer must be greater than rho_inner for all time.')
all_times_min_max = np.union1d(self.rho_min.time, self.rho_max.time)
if not np.all(
self.rho_max.get_value(all_times_min_max)
@@ -195,28 +110,8 @@ def build_runtime_params(
)
),
fast_ion_stabilization_multiplier=self.fast_ion_stabilization_multiplier,
- chi_min=self.chi_min,
- chi_max=self.chi_max,
- D_e_min=self.D_e_min,
- D_e_max=self.D_e_max,
- V_e_min=self.V_e_min,
- V_e_max=self.V_e_max,
rho_min=self.rho_min.get_value(t),
rho_max=self.rho_max.get_value(t),
- apply_inner_patch=self.apply_inner_patch.get_value(t),
- D_e_inner=self.D_e_inner.get_value(t),
- V_e_inner=self.V_e_inner.get_value(t),
- chi_i_inner=self.chi_i_inner.get_value(t),
- chi_e_inner=self.chi_e_inner.get_value(t),
- rho_inner=self.rho_inner.get_value(t),
- apply_outer_patch=self.apply_outer_patch.get_value(t),
- D_e_outer=self.D_e_outer.get_value(t),
- V_e_outer=self.V_e_outer.get_value(t),
- chi_i_outer=self.chi_i_outer.get_value(t),
- chi_e_outer=self.chi_e_outer.get_value(t),
- rho_outer=self.rho_outer.get_value(t),
- smoothing_width=self.smoothing_width,
- smooth_everywhere=self.smooth_everywhere,
disable_chi_i=self.disable_chi_i.get_value(t),
disable_chi_e=self.disable_chi_e.get_value(t),
disable_D_e=self.disable_D_e.get_value(t),
diff --git a/torax/_src/transport_model/runtime_params.py b/torax/_src/transport_model/runtime_params.py
index 381e516f0..662672a4f 100644
--- a/torax/_src/transport_model/runtime_params.py
+++ b/torax/_src/transport_model/runtime_params.py
@@ -30,28 +30,8 @@
class RuntimeParams:
"""Input params for the transport model which can be used as compiled args."""
- chi_min: float
- chi_max: float
- D_e_min: float
- D_e_max: float
- V_e_min: float
- V_e_max: float
rho_min: array_typing.FloatScalar
rho_max: array_typing.FloatScalar
- apply_inner_patch: array_typing.BoolScalar
- D_e_inner: array_typing.FloatScalar
- V_e_inner: array_typing.FloatScalar
- chi_i_inner: array_typing.FloatScalar
- chi_e_inner: array_typing.FloatScalar
- rho_inner: array_typing.FloatScalar
- apply_outer_patch: array_typing.BoolScalar
- D_e_outer: array_typing.FloatScalar
- V_e_outer: array_typing.FloatScalar
- chi_i_outer: array_typing.FloatScalar
- chi_e_outer: array_typing.FloatScalar
- rho_outer: array_typing.FloatScalar
- smoothing_width: float
- smooth_everywhere: bool
disable_chi_i: array_typing.BoolScalar
disable_chi_e: array_typing.BoolScalar
disable_D_e: array_typing.BoolScalar
@@ -62,3 +42,30 @@ class RuntimeParams:
)
fast_ion_stabilization_multiplier: float
merge_mode: enums.MergeMode = dataclasses.field(metadata={'static': True})
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class SmoothingZoneParams:
+ """Runtime parameters for a radial smoothing zone."""
+
+ rho_min: array_typing.FloatScalar
+ rho_max: array_typing.FloatScalar
+ smoothing_width: float
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class CombinedRuntimeParams:
+ """Runtime parameters for the CombinedTransportModel."""
+
+ chi_min: float
+ chi_max: float
+ D_e_min: float
+ D_e_max: float
+ V_e_min: float
+ V_e_max: float
+ smoothing_width: float
+ transport_model_params: tuple[RuntimeParams, ...]
+ pedestal_transport_model_params: tuple[RuntimeParams, ...]
+ smoothing_zones: tuple[SmoothingZoneParams, ...]
diff --git a/torax/_src/transport_model/tests/combined_test.py b/torax/_src/transport_model/tests/combined_test.py
index a7aeeff76..4341b3aab 100644
--- a/torax/_src/transport_model/tests/combined_test.py
+++ b/torax/_src/transport_model/tests/combined_test.py
@@ -379,37 +379,6 @@ def test_smoothing_width_shortcut(self):
)
)
- def test_error_if_patches_set_on_children(self):
- config = default_configs.get_default_config_dict()
- config['transport'] = {
- 'model_name': 'combined',
- 'transport_models': [
- {'model_name': 'constant', 'apply_inner_patch': True},
- {'model_name': 'constant'},
- ],
- 'pedestal_transport_models': [{'model_name': 'constant'}],
- }
- with self.assertRaisesRegex(
- ValueError, '(?=.*patch)(?=.*CombinedTransportModel)'
- ):
- model_config.ToraxConfig.from_dict(config)
-
- def test_error_if_patches_set_on_self(self):
- config = default_configs.get_default_config_dict()
- config['transport'] = {
- 'model_name': 'combined',
- 'transport_models': [
- {'model_name': 'constant'},
- {'model_name': 'constant'},
- ],
- 'pedestal_transport_models': [{'model_name': 'constant'}],
- 'apply_inner_patch': True,
- }
- with self.assertRaisesRegex(
- ValueError, '(?=.*patch)(?=.*CombinedTransportModel)'
- ):
- model_config.ToraxConfig.from_dict(config)
-
def test_error_if_pedestal_model_defines_rho_min(self):
config = default_configs.get_default_config_dict()
config['transport'] = {
diff --git a/torax/_src/transport_model/tests/pydantic_model_test.py b/torax/_src/transport_model/tests/pydantic_model_test.py
index cdaf9890c..27aeab0dd 100644
--- a/torax/_src/transport_model/tests/pydantic_model_test.py
+++ b/torax/_src/transport_model/tests/pydantic_model_test.py
@@ -97,7 +97,6 @@ def test_qlknn_defaults(self):
self.assertIsInstance(
transport, transport_pydantic_model.QLKNNTransportModel
)
- self.assertEqual(transport.smoothing_width, 0.1)
self.assertEqual(transport.ETG_correction_factor, 1.0 / 3.0)
@parameterized.parameters(
@@ -150,39 +149,18 @@ def test_qlknn_model_errors(self, model_name, model_path):
)
@parameterized.named_parameters(
- (
- 'mixed_modes_inner_outer_fails',
- {'rho_inner': 0.3, 'rho_outer': ({0: 0.9}, 'step')},
- 'rho_outer and rho_inner must have the same interpolation mode.',
- True,
- ),
(
'mixed_modes_min_max_fails',
{'rho_min': 0.0, 'rho_max': ({0: 1.0}, 'step')},
'rho_max and rho_min must have the same interpolation mode.',
True,
),
- (
- 'inner_greater_than_outer_fails',
- {'rho_inner': 0.9, 'rho_outer': 0.3},
- 'rho_outer must be greater than rho_inner for all time.',
- True,
- ),
(
'min_greater_than_max_fails',
{'rho_min': 0.9, 'rho_max': 0.3},
'rho_max must be greater than rho_min for all time.',
True,
),
- (
- 'time_varying_inner_gt_outer_fails',
- {
- 'rho_inner': {0: 0.3, 1: 0.95},
- 'rho_outer': {0: 0.9, 1: 0.9},
- },
- 'rho_outer must be greater than rho_inner for all time.',
- True,
- ),
(
'time_varying_min_gt_max_fails',
{
@@ -195,8 +173,6 @@ def test_qlknn_model_errors(self, model_name, model_path):
(
'time_varying_linear_succeeds',
{
- 'rho_inner': {0: 0.2, 1: 0.3},
- 'rho_outer': {0: 0.8, 1: 0.9},
'rho_min': {0: 0.0, 1: 0.1},
'rho_max': {0: 1.0, 1: 0.95},
},
@@ -206,8 +182,6 @@ def test_qlknn_model_errors(self, model_name, model_path):
(
'time_varying_step_succeeds',
{
- 'rho_inner': ({0: 0.2, 1: 0.3}, 'step'),
- 'rho_outer': ({0: 0.8, 1: 0.9}, 'step'),
'rho_min': ({0: 0.0, 1: 0.1}, 'step'),
'rho_max': ({0: 1.0, 1: 0.95}, 'step'),
},
@@ -465,38 +439,6 @@ def test_invalid_overlapping_overwrite_part_time(self):
):
transport_pydantic_model.CombinedTransportModel(transport_models=[m1, m2])
- def test_smoothing_in_combined_core_component_logs_warning(self):
- component_model = transport_pydantic_model.ConstantTransportModel(
- smoothing_width=0.1
- )
- with self.assertLogs(level='WARNING') as log_watcher:
- transport_pydantic_model.CombinedTransportModel(
- transport_models=[component_model]
- )
- self.assertLen(log_watcher.output, 1)
- self.assertIn(
- 'smoothing_width > 0.0 is not supported for component models of'
- ' CombinedTransportModel',
- log_watcher.output[0],
- )
-
- def test_smoothing_in_combined_pedestal_component_logs_warning(
- self,
- ):
- component_model = transport_pydantic_model.ConstantTransportModel(
- smoothing_width=0.1
- )
- with self.assertLogs(level='WARNING') as log_watcher:
- transport_pydantic_model.CombinedTransportModel(
- pedestal_transport_models=[component_model]
- )
- self.assertLen(log_watcher.output, 1)
- self.assertIn(
- 'smoothing_width > 0.0 is not supported for component models of'
- ' CombinedTransportModel',
- log_watcher.output[0],
- )
-
if __name__ == '__main__':
absltest.main()
diff --git a/torax/_src/transport_model/transport_coefficients_builder.py b/torax/_src/transport_model/transport_coefficients_builder.py
index 569cfef40..c8060861a 100644
--- a/torax/_src/transport_model/transport_coefficients_builder.py
+++ b/torax/_src/transport_model/transport_coefficients_builder.py
@@ -24,8 +24,8 @@
from torax._src.neoclassical import neoclassical_models as neoclassical_models_lib
from torax._src.pedestal_model import pedestal_transition_state as pedestal_transition_state_lib
from torax._src.pedestal_model import runtime_params as pedestal_runtime_params_lib
+from torax._src.transport_model import combined
from torax._src.transport_model import pereverzev as pereverzev_lib
-from torax._src.transport_model import transport_model as transport_model_lib
# pylint: disable=invalid-name
@@ -37,7 +37,7 @@
)
)
def calculate_all_transport_coeffs(
- transport_model: transport_model_lib.TransportModel,
+ transport_model: combined.CombinedTransportModel,
neoclassical_models: neoclassical_models_lib.NeoclassicalModels,
runtime_params: runtime_params_lib.RuntimeParams,
geo: geometry.Geometry,
diff --git a/torax/_src/transport_model/transport_model.py b/torax/_src/transport_model/transport_model.py
index ccb63feb8..9fa87409a 100644
--- a/torax/_src/transport_model/transport_model.py
+++ b/torax/_src/transport_model/transport_model.py
@@ -25,7 +25,6 @@
import immutabledict
import jax
from jax import numpy as jnp
-from torax._src import constants
from torax._src import state
from torax._src import static_dataclass
from torax._src.config import runtime_params as runtime_params_lib
@@ -122,61 +121,6 @@ class TurbulentTransport:
class TransportModel(static_dataclass.StaticDataclass, abc.ABC):
"""Calculates various coefficients related to heat and particle transport."""
- def __call__(
- self,
- runtime_params: runtime_params_lib.RuntimeParams,
- geo: geometry.Geometry,
- core_profiles: state.CoreProfiles,
- pedestal_model_output: pedestal_model_output_lib.PedestalModelOutput,
- ) -> TurbulentTransport:
- transport_runtime_params = runtime_params.transport
-
- # Calculate the transport coefficients
- transport_coeffs = self.call_implementation(
- transport_runtime_params,
- runtime_params,
- geo,
- core_profiles,
- pedestal_model_output,
- )
-
- # Apply masking to selectively enable/disable specific channels
- transport_coeffs = self.zero_out_disabled_channels(
- transport_runtime_params, transport_coeffs
- )
-
- # Restrict the model to operating in its permissible rho domain
- transport_coeffs = self._apply_domain_restriction(
- transport_runtime_params,
- runtime_params,
- geo,
- transport_coeffs,
- pedestal_model_output,
- )
-
- # Apply min/max clipping
- transport_coeffs = self._apply_clipping(
- transport_runtime_params,
- transport_coeffs,
- )
-
- # Apply inner and outer transport patch
- transport_coeffs = self._apply_transport_patches(
- transport_runtime_params,
- runtime_params,
- geo,
- transport_coeffs,
- )
-
- transport_coeffs = self._smooth_coeffs(
- runtime_params,
- geo,
- transport_coeffs,
- pedestal_model_output,
- )
-
- return transport_coeffs
-
@abc.abstractmethod
def call_implementation(
self,
@@ -212,304 +156,6 @@ def zero_out_disabled_channels(
return dataclasses.replace(transport_coeffs, **to_replace)
- def _apply_domain_restriction(
- self,
- transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
- runtime_params: runtime_params_lib.RuntimeParams,
- geo: geometry.Geometry,
- transport_coeffs: TurbulentTransport,
- pedestal_model_output: pedestal_model_output_lib.PedestalModelOutput,
- ) -> TurbulentTransport:
- """Sets transport coefficients to zero outside the model's domain."""
- active_mask = compute_core_domain_mask(
- transport_runtime_params, runtime_params, geo, pedestal_model_output
- )
-
- coeffs_dict = dataclasses.asdict(transport_coeffs)
- to_replace = {}
-
- for channel_name, config in CHANNEL_CONFIG_STRUCT.items():
- # Mask main channel
- val = coeffs_dict[channel_name]
- to_replace[channel_name] = jnp.where(active_mask, val, 0.0) # pyrefly: ignore[bad-argument-type]
-
- # Mask sub-channels
- for sub_channel in config['sub_channels']:
- sub_val = coeffs_dict[sub_channel]
- if sub_val is not None:
- to_replace[sub_channel] = jnp.where(active_mask, sub_val, 0.0)
-
- return dataclasses.replace(transport_coeffs, **to_replace)
-
- def _apply_clipping(
- self,
- transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
- transport_coeffs: TurbulentTransport,
- ) -> TurbulentTransport:
- """Applies min/max clipping to transport coefficients for PDE stability."""
- chi_face_ion = jnp.clip(
- transport_coeffs.chi_face_ion,
- transport_runtime_params.chi_min,
- transport_runtime_params.chi_max,
- )
- chi_face_el = jnp.clip(
- transport_coeffs.chi_face_el,
- transport_runtime_params.chi_min,
- transport_runtime_params.chi_max,
- )
- d_face_el = jnp.clip(
- transport_coeffs.d_face_el,
- transport_runtime_params.D_e_min,
- transport_runtime_params.D_e_max,
- )
- v_face_el = jnp.clip(
- transport_coeffs.v_face_el,
- transport_runtime_params.V_e_min,
- transport_runtime_params.V_e_max,
- )
-
- return dataclasses.replace(
- transport_coeffs,
- chi_face_ion=chi_face_ion,
- chi_face_el=chi_face_el,
- d_face_el=d_face_el,
- v_face_el=v_face_el,
- )
-
- def _apply_transport_patches(
- self,
- transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
- runtime_params: runtime_params_lib.RuntimeParams,
- geo: geometry.Geometry,
- transport_coeffs: TurbulentTransport,
- ) -> TurbulentTransport:
- """Applies inner and outer transport patches to transport coefficients."""
- consts = constants.CONSTANTS
-
- # Apply inner and outer patch constant transport coefficients. rho_inner and
- # rho_outer are shifted by consts.eps (1e-7) to avoid ambiguities if their
- # values are close to and geo.rho_face_norm values.
- chi_face_ion = jnp.where(
- jnp.logical_and(
- transport_runtime_params.apply_inner_patch,
- geo.rho_face_norm < transport_runtime_params.rho_inner + consts.eps,
- ),
- transport_runtime_params.chi_i_inner,
- transport_coeffs.chi_face_ion,
- )
- chi_face_el = jnp.where(
- jnp.logical_and(
- transport_runtime_params.apply_inner_patch,
- geo.rho_face_norm < transport_runtime_params.rho_inner + consts.eps,
- ),
- transport_runtime_params.chi_e_inner,
- transport_coeffs.chi_face_el,
- )
- d_face_el = jnp.where(
- jnp.logical_and(
- transport_runtime_params.apply_inner_patch,
- geo.rho_face_norm < transport_runtime_params.rho_inner + consts.eps,
- ),
- transport_runtime_params.D_e_inner,
- transport_coeffs.d_face_el,
- )
- v_face_el = jnp.where(
- jnp.logical_and(
- transport_runtime_params.apply_inner_patch,
- geo.rho_face_norm < transport_runtime_params.rho_inner + consts.eps,
- ),
- transport_runtime_params.V_e_inner,
- transport_coeffs.v_face_el,
- )
-
- # Apply outer patch constant transport coefficients.
- # Due to Pereverzev-Corrigan convection, it is required
- # for the convection modes to be 'ghost' to avoid numerical instability
- chi_face_ion = jnp.where(
- jnp.logical_and(
- jnp.logical_and(
- transport_runtime_params.apply_outer_patch,
- jnp.logical_not(runtime_params.pedestal.set_pedestal),
- ),
- geo.rho_face_norm > transport_runtime_params.rho_outer - consts.eps,
- ),
- transport_runtime_params.chi_i_outer,
- chi_face_ion,
- )
- chi_face_el = jnp.where(
- jnp.logical_and(
- jnp.logical_and(
- transport_runtime_params.apply_outer_patch,
- jnp.logical_not(runtime_params.pedestal.set_pedestal),
- ),
- geo.rho_face_norm > transport_runtime_params.rho_outer - consts.eps,
- ),
- transport_runtime_params.chi_e_outer,
- chi_face_el,
- )
- d_face_el = jnp.where(
- jnp.logical_and(
- jnp.logical_and(
- transport_runtime_params.apply_outer_patch,
- jnp.logical_not(runtime_params.pedestal.set_pedestal),
- ),
- geo.rho_face_norm > transport_runtime_params.rho_outer - consts.eps,
- ),
- transport_runtime_params.D_e_outer,
- d_face_el,
- )
- v_face_el = jnp.where(
- jnp.logical_and(
- jnp.logical_and(
- transport_runtime_params.apply_outer_patch,
- jnp.logical_not(runtime_params.pedestal.set_pedestal),
- ),
- geo.rho_face_norm > transport_runtime_params.rho_outer - consts.eps,
- ),
- transport_runtime_params.V_e_outer,
- v_face_el,
- )
-
- return dataclasses.replace(
- transport_coeffs,
- chi_face_ion=chi_face_ion,
- chi_face_el=chi_face_el,
- d_face_el=d_face_el,
- v_face_el=v_face_el,
- )
-
- def _smooth_coeffs(
- self,
- runtime_params: runtime_params_lib.RuntimeParams,
- geo: geometry.Geometry,
- transport_coeffs: TurbulentTransport,
- pedestal_model_output: pedestal_model_output_lib.PedestalModelOutput,
- ) -> TurbulentTransport:
- """Gaussian smoothing of turbulent transport coefficients."""
- smoothing_matrix = _build_smoothing_matrix(
- runtime_params,
- geo,
- pedestal_model_output,
- )
-
- # Iterate over fields of the CoreTransport dataclass.
- # Ignore optional fields that are made all zero in post_init.
- def smooth_single_coeff(coeff):
- return jax.lax.cond(
- jnp.all(coeff == 0.0),
- lambda: coeff,
- lambda: jnp.dot(smoothing_matrix, coeff),
- )
-
- return jax.tree_util.tree_map(smooth_single_coeff, transport_coeffs)
-
-
-def _build_smoothing_matrix(
- runtime_params: runtime_params_lib.RuntimeParams,
- geo: geometry.Geometry,
- pedestal_model_output: pedestal_model_output_lib.PedestalModelOutput,
-) -> jax.Array:
- """Builds a smoothing matrix for the turbulent transport model.
-
- Uses a Gaussian kernel of HWHM defined in the transport config.
-
- Args:
- runtime_params: Input runtime parameters of the simulation.
- geo: Geometry of the torus.
- pedestal_model_output: Output of the pedestal model.
-
- Returns:
- kernel: A smoothing matrix for convolution with the transport outputs.
- """
-
- # To reduce the range of the convolution, weights under lower_cutoff are
- # clipped to zero
- lower_cutoff = 0.01
-
- # used for eps, small number to avoid divisions by zero for sigma = 0
- consts = constants.CONSTANTS
-
- # 1. Kernel matrix
- kernel = jnp.exp(
- -jnp.log(2)
- * (geo.rho_face_norm[:, jnp.newaxis] - geo.rho_face_norm) ** 2
- / (runtime_params.transport.smoothing_width**2 + consts.eps)
- )
-
- # 2. Masking: we do not want transport coefficients calculated in pedestal
- # region or in inner and outer transport patch regions to impact
- # transport_model calculated coefficients
- if (
- runtime_params.pedestal.mode
- == pedestal_runtime_params_lib.Mode.INTERNAL_BOUNDARY_CONDITION
- ):
- # If in INTERNAL_BOUNDARY_CONDITION mode: if set_pedestal is True, mask
- # according to the pedestal top. Otherwise, mask according to the outer
- # patch, if set.
- mask_outer_edge = jnp.where(
- runtime_params.pedestal.set_pedestal,
- pedestal_model_output.rho_norm_ped_top - consts.eps,
- jnp.where(
- runtime_params.transport.apply_outer_patch,
- runtime_params.transport.rho_outer - consts.eps,
- jnp.inf,
- ),
- )
- else:
- # If in ADAPTIVE_TRANSPORT mode, only mask according to the outer patch.
- mask_outer_edge = jnp.where(
- runtime_params.transport.apply_outer_patch,
- runtime_params.transport.rho_outer - consts.eps,
- jnp.inf,
- )
-
- mask_inner_edge = jax.lax.cond(
- runtime_params.transport.apply_inner_patch,
- lambda: runtime_params.transport.rho_inner + consts.eps,
- lambda: -consts.eps,
- )
-
- mask = jnp.where(
- jnp.logical_or(
- runtime_params.transport.smooth_everywhere,
- jnp.logical_and(
- geo.rho_face_norm > mask_inner_edge,
- geo.rho_face_norm < mask_outer_edge,
- ),
- ),
- 1.0,
- 0.0,
- )
-
- # remove impact of smoothing on inner and outer patch, or pedestal zone
-
- # first zero out all rows corresponding to grid points not to be impacted
- diag_mask = jnp.diag(mask)
- kernel = jnp.dot(diag_mask, kernel)
- # now zero out all columns corresponding to grid points not to be impacted,
- # such that they don't impact the smoothing of the other grid points
- num_rows = len(mask)
- mask_mat = jnp.tile(mask, (num_rows, 1))
- kernel *= mask_mat
- # now restore identity to the zero rows, such that smoothing is a no-op for
- # on the grid points where it shouldn't impact
- zero_row_mask = jnp.all(kernel == 0, axis=1)
- kernel = jnp.where(
- zero_row_mask[:, jnp.newaxis], jnp.eye(kernel.shape[0]), kernel
- )
-
- # 3. Normalization
- row_sums = jnp.sum(kernel, axis=1)
- kernel /= row_sums[:, jnp.newaxis]
-
- # 4. Remove small numbers
- kernel = jnp.where(kernel < lower_cutoff, 0.0, kernel)
-
- # 5. Final Normalization following removal of small numbers
- row_sums = jnp.sum(kernel, axis=1)
- kernel /= row_sums[:, jnp.newaxis]
- return kernel
-
def compute_core_domain_mask(
transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
diff --git a/torax/benchmarks/tokagrad_benchmark.py b/torax/benchmarks/tokagrad_benchmark.py
index f1b4cf6a2..d446226bf 100644
--- a/torax/benchmarks/tokagrad_benchmark.py
+++ b/torax/benchmarks/tokagrad_benchmark.py
@@ -237,26 +237,30 @@
"pellet": {"S_total": 0.0},
},
"transport": {
- "model_name": "bohm-gyrobohm",
- "chi_e_bohm_multiplier": 1.0,
- "chi_i_bohm_multiplier": 1.0,
- "chi_e_gyrobohm_multiplier": 1.0,
- "chi_i_gyrobohm_multiplier": 1.0,
- "chi_e_bohm_coeff": 8e-5,
- "chi_e_gyrobohm_coeff": 5e-6,
- "chi_i_bohm_coeff": 8e-5,
- "chi_i_gyrobohm_coeff": 5e-6,
- "D_face_c1": 1.0,
- "D_face_c2": 0.3,
- "V_face_coeff": 0.0,
+ "model_name": "combined",
"chi_min": 0.05,
"chi_max": 100.0,
"D_e_min": 0.05,
"D_e_max": 50.0,
"V_e_min": -10.0,
"V_e_max": 10.0,
- "smooth_everywhere": False,
"smoothing_width": 0.05,
+ "transport_models": [
+ {
+ "model_name": "bohm-gyrobohm",
+ "chi_e_bohm_multiplier": 1.0,
+ "chi_i_bohm_multiplier": 1.0,
+ "chi_e_gyrobohm_multiplier": 1.0,
+ "chi_i_gyrobohm_multiplier": 1.0,
+ "chi_e_bohm_coeff": 8e-5,
+ "chi_e_gyrobohm_coeff": 5e-6,
+ "chi_i_bohm_coeff": 8e-5,
+ "chi_i_gyrobohm_coeff": 5e-6,
+ "D_face_c1": 1.0,
+ "D_face_c2": 0.3,
+ "V_face_coeff": 0.0,
+ }
+ ],
},
"solver": {
"solver_type": "linear",
diff --git a/torax/tests/test_data/test_combined_transport.nc b/torax/tests/test_data/test_combined_transport.nc
index 2a59e77a2..592b402db 100644
Binary files a/torax/tests/test_data/test_combined_transport.nc and b/torax/tests/test_data/test_combined_transport.nc differ
diff --git a/torax/tutorials/torax_tutorial_exercises.ipynb b/torax/tutorials/torax_tutorial_exercises.ipynb
index dbe368bfb..49a54d438 100644
--- a/torax/tutorials/torax_tutorial_exercises.ipynb
+++ b/torax/tutorials/torax_tutorial_exercises.ipynb
@@ -1,19 +1,16 @@
{
"cells": [
{
- "metadata": {
- "id": "wH4_0fmxGQL8"
- },
"cell_type": "markdown",
"source": [
"# TORAX Tutorial\n",
"![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\")
\n"
- ]
+ ],
+ "metadata": {
+ "id": "wH4_0fmxGQL8"
+ }
},
{
- "metadata": {
- "id": "m5geWYHBKDd6"
- },
"cell_type": "code",
"source": [
"#@title LICENSE\n",
@@ -31,14 +28,13 @@
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
],
- "outputs": [],
- "execution_count": 153
- },
- {
"metadata": {
- "id": "9UYOakVOJ_zM",
- "cellView": "form"
+ "id": "m5geWYHBKDd6"
},
+ "execution_count": 153,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Install TORAX\n",
@@ -47,14 +43,14 @@
"# To install the current version on GitHub, use `%pip install git+https://github.com/google-deepmind/torax.git`.\n",
"%pip install torax"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "I4TZNOJJGBR_",
- "cellView": "form"
+ "cellView": "form",
+ "id": "9UYOakVOJ_zM"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Imports\n",
@@ -68,14 +64,14 @@
"\n",
"jax.config.update('jax_enable_x64', True)"
],
- "outputs": [],
- "execution_count": 154
- },
- {
"metadata": {
- "id": "A7MMV9poKDd6",
- "cellView": "form"
+ "cellView": "form",
+ "id": "I4TZNOJJGBR_"
},
+ "execution_count": 154,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Base config\n",
@@ -188,34 +184,42 @@
" 'rho_norm_ped_top': 0.95, # set ped top location in normalized radius\n",
" },\n",
" 'transport': {\n",
- " 'model_name': 'qlknn', # Using QLKNN_7_11 default\n",
- " # set inner core transport coefficients (ad-hoc MHD/EM transport)\n",
- " 'apply_inner_patch': True,\n",
- " 'D_e_inner': 0.15,\n",
- " 'V_e_inner': 0.0,\n",
- " 'chi_i_inner': 0.3,\n",
- " 'chi_e_inner': 0.3,\n",
- " 'rho_inner': 0.1, # radius below which patch transport is applied\n",
- " # set outer core transport coefficients (L-mode near edge region)\n",
- " 'apply_outer_patch': True,\n",
- " 'D_e_outer': 0.1,\n",
- " 'V_e_outer': 0.0,\n",
- " 'chi_i_outer': 2.0,\n",
- " 'chi_e_outer': 2.0,\n",
- " 'rho_outer': 0.95, # radius above which patch transport is applied\n",
- " # allowed chi and diffusivity bounds\n",
- " 'chi_min': 0.05, # minimum chi\n",
- " 'chi_max': 100, # maximum chi (can be helpful for stability)\n",
- " 'D_e_min': 0.05, # minimum electron diffusivity\n",
- " 'D_e_max': 50, # maximum electron diffusivity\n",
- " 'V_e_min': -10, # minimum electron convection\n",
- " 'V_e_max': 10, # minimum electron convection\n",
+ " 'model_name': 'combined',\n",
+ " 'chi_min': 0.05,\n",
+ " 'chi_max': 100,\n",
+ " 'D_e_min': 0.05,\n",
+ " 'D_e_max': 50,\n",
+ " 'V_e_min': -10,\n",
+ " 'V_e_max': 10,\n",
" 'smoothing_width': 0.1,\n",
- " 'DV_effective': True,\n",
- " 'include_ITG': True, # to toggle ITG modes on or off\n",
- " 'include_TEM': True, # to toggle TEM modes on or off\n",
- " 'include_ETG': True, # to toggle ETG modes on or off\n",
- " 'avoid_big_negative_s': False,\n",
+ " 'transport_models': [\n",
+ " {\n",
+ " 'model_name': 'constant',\n",
+ " 'rho_max': 0.1,\n",
+ " 'chi_i': 0.3,\n",
+ " 'chi_e': 0.3,\n",
+ " 'D_e': 0.15,\n",
+ " 'V_e': 0.0,\n",
+ " },\n",
+ " {\n",
+ " 'model_name': 'qlknn',\n",
+ " 'rho_min': 0.1,\n",
+ " 'rho_max': 0.95,\n",
+ " 'DV_effective': True,\n",
+ " 'include_ITG': True,\n",
+ " 'include_TEM': True,\n",
+ " 'include_ETG': True,\n",
+ " 'avoid_big_negative_s': False,\n",
+ " },\n",
+ " {\n",
+ " 'model_name': 'constant',\n",
+ " 'rho_min': 0.95,\n",
+ " 'chi_i': 2.0,\n",
+ " 'chi_e': 2.0,\n",
+ " 'D_e': 0.1,\n",
+ " 'V_e': 0.0,\n",
+ " },\n",
+ " ],\n",
" },\n",
" 'solver': {\n",
" 'solver_type': 'linear', # linear solver with picard iteration\n",
@@ -229,16 +233,16 @@
" 'time_step_calculator': {\n",
" 'calculator_type': 'fixed',\n",
" },\n",
- "}\n"
+ "}"
],
- "outputs": [],
- "execution_count": 155
- },
- {
"metadata": {
- "id": "xSzvOYW_KDd6",
- "cellView": "form"
+ "cellView": "form",
+ "id": "A7MMV9poKDd6"
},
+ "execution_count": 155,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Detailed summary plot function for a single simulation\n",
@@ -449,14 +453,14 @@
"\n",
" plt.show()\n"
],
- "outputs": [],
- "execution_count": 176
- },
- {
"metadata": {
- "id": "qmq9COuoKDd6",
- "cellView": "form"
+ "cellView": "form",
+ "id": "xSzvOYW_KDd6"
},
+ "execution_count": 176,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Comparison timetrace summary plot for multiple simulations\n",
@@ -578,14 +582,14 @@
" plt.tight_layout()\n",
" plt.show()"
],
- "outputs": [],
- "execution_count": 157
- },
- {
"metadata": {
- "id": "FL-_GqQaGBSA",
- "cellView": "form"
+ "cellView": "form",
+ "id": "qmq9COuoKDd6"
},
+ "execution_count": 157,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Comparison profile summary plot for multiple simulations\n",
@@ -690,14 +694,14 @@
" plt.tight_layout()\n",
" plt.show()"
],
- "outputs": [],
- "execution_count": 158
- },
- {
"metadata": {
- "id": "do-4obv1GBSA",
- "cellView": "form"
+ "cellView": "form",
+ "id": "FL-_GqQaGBSA"
},
+ "execution_count": 158,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title set_LH_transition_time_function\n",
@@ -725,14 +729,14 @@
" if not (allowed_range[0] <= input_obj <= allowed_range[1]):\n",
" raise ValueError(f\"Input value must be between {allowed_range[0]} and {allowed_range[1]}.\")\n"
],
- "outputs": [],
- "execution_count": 159
- },
- {
"metadata": {
- "id": "O4B4jZmAGBSA",
- "cellView": "form"
+ "cellView": "form",
+ "id": "do-4obv1GBSA"
},
+ "execution_count": 159,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Modify config function\n",
@@ -848,14 +852,14 @@
" if not (allowed_range[0] <= input_obj <= allowed_range[1]):\n",
" raise ValueError(f\"{variable_name} values must be between {allowed_range[0]} and {allowed_range[1]}.\")\n"
],
- "outputs": [],
- "execution_count": 160
- },
- {
"metadata": {
- "id": "HRKQQkpJGBSA",
- "cellView": "form"
+ "cellView": "form",
+ "id": "O4B4jZmAGBSA"
},
+ "execution_count": 160,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"# @title Function to build and launch TORAX sim\n",
@@ -864,13 +868,14 @@
" data_tree, _ = torax.run_simulation(torax_config, log_timestep_info=False)\n",
" return data_tree"
],
- "outputs": [],
- "execution_count": 161
- },
- {
"metadata": {
- "id": "FDC7xuNLKDd7"
+ "cellView": "form",
+ "id": "HRKQQkpJGBSA"
},
+ "execution_count": 161,
+ "outputs": []
+ },
+ {
"cell_type": "markdown",
"source": [
"# TORAX exercise\n",
@@ -970,17 +975,21 @@
"* Maintain q_min > 1\n",
"* Total input power at t=100s must be > 50MW.\n",
"* Ip cannot exceed 18MA during the current rampup. Note that the current ramp does not need to be monotonic.\n"
- ]
+ ],
+ "metadata": {
+ "id": "FDC7xuNLKDd7"
+ }
},
{
- "metadata": {
- "id": "3OPtAiy4KDd7"
- },
"cell_type": "code",
"source": [
"config0 = set_LH_transition_time(LH_transition_time = 80)\n",
"out0 = run_sim(config0)"
],
+ "metadata": {
+ "id": "3OPtAiy4KDd7"
+ },
+ "execution_count": 162,
"outputs": [
{
"name": "stderr",
@@ -989,17 +998,17 @@
"Simulating (t=150.00000): 100%|██████████| 100/100 [00:08<00:00, 11.76it/s]\n"
]
}
- ],
- "execution_count": 162
+ ]
},
{
- "metadata": {
- "id": "P1DAzLDAKDd7"
- },
"cell_type": "code",
"source": [
"detailed_plot_single_sim(out0, time = 150)"
],
+ "metadata": {
+ "id": "P1DAzLDAKDd7"
+ },
+ "execution_count": 177,
"outputs": [
{
"name": "stdout",
@@ -1026,17 +1035,16 @@
},
"output_type": "display_data"
}
- ],
- "execution_count": 177
+ ]
},
{
+ "cell_type": "code",
+ "source": [],
"metadata": {
"id": "ksx5Hd5SKDd7"
},
- "cell_type": "code",
- "source": [],
- "outputs": [],
- "execution_count": 163
+ "execution_count": 163,
+ "outputs": []
}
],
"metadata": {
@@ -1065,6 +1073,6 @@
"version": "3.12.8"
}
},
- "nbformat": 4,
- "nbformat_minor": 0
+ "nbformat_minor": 0,
+ "nbformat": 4
}
diff --git a/torax/tutorials/torax_tutorial_exercises_with_solutions.ipynb b/torax/tutorials/torax_tutorial_exercises_with_solutions.ipynb
index a79d4be61..fa62f8752 100644
--- a/torax/tutorials/torax_tutorial_exercises_with_solutions.ipynb
+++ b/torax/tutorials/torax_tutorial_exercises_with_solutions.ipynb
@@ -1,19 +1,16 @@
{
"cells": [
{
- "metadata": {
- "id": "wAoJBvlGGJte"
- },
"cell_type": "markdown",
"source": [
"# TORAX Tutorial with Solutions\n",
"\n"
- ]
+ ],
+ "metadata": {
+ "id": "wAoJBvlGGJte"
+ }
},
{
- "metadata": {
- "id": "m5geWYHBKDd6"
- },
"cell_type": "code",
"source": [
"#@title LICENSE\n",
@@ -31,13 +28,13 @@
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "vflmKg8bKNTl"
+ "id": "m5geWYHBKDd6"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Install TORAX\n",
@@ -45,13 +42,13 @@
"# To install the current version on GitHub, use `%pip install git+https://github.com/google-deepmind/torax.git`.\n",
"%pip install torax"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "cufUHoEGKDd6"
+ "id": "vflmKg8bKNTl"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Imports\n",
@@ -65,13 +62,13 @@
"\n",
"jax.config.update('jax_enable_x64', True)"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "A7MMV9poKDd6"
+ "id": "cufUHoEGKDd6"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Base config\n",
@@ -184,34 +181,42 @@
" 'rho_norm_ped_top': 0.95, # set ped top location in normalized radius\n",
" },\n",
" 'transport': {\n",
- " 'model_name': 'qlknn', # Using QLKNN_7_11 default\n",
- " # set inner core transport coefficients (ad-hoc MHD/EM transport)\n",
- " 'apply_inner_patch': True,\n",
- " 'D_e_inner': 0.15,\n",
- " 'V_e_inner': 0.0,\n",
- " 'chi_i_inner': 0.3,\n",
- " 'chi_e_inner': 0.3,\n",
- " 'rho_inner': 0.1, # radius below which patch transport is applied\n",
- " # set outer core transport coefficients (L-mode near edge region)\n",
- " 'apply_outer_patch': True,\n",
- " 'D_e_outer': 0.1,\n",
- " 'V_e_outer': 0.0,\n",
- " 'chi_i_outer': 2.0,\n",
- " 'chi_e_outer': 2.0,\n",
- " 'rho_outer': 0.95, # radius above which patch transport is applied\n",
- " # allowed chi and diffusivity bounds\n",
- " 'chi_min': 0.05, # minimum chi\n",
- " 'chi_max': 100, # maximum chi (can be helpful for stability)\n",
- " 'D_e_min': 0.05, # minimum electron diffusivity\n",
- " 'D_e_max': 50, # maximum electron diffusivity\n",
- " 'V_e_min': -10, # minimum electron convection\n",
- " 'V_e_max': 10, # minimum electron convection\n",
+ " 'model_name': 'combined',\n",
+ " 'chi_min': 0.05,\n",
+ " 'chi_max': 100,\n",
+ " 'D_e_min': 0.05,\n",
+ " 'D_e_max': 50,\n",
+ " 'V_e_min': -10,\n",
+ " 'V_e_max': 10,\n",
" 'smoothing_width': 0.1,\n",
- " 'DV_effective': True,\n",
- " 'include_ITG': True, # to toggle ITG modes on or off\n",
- " 'include_TEM': True, # to toggle TEM modes on or off\n",
- " 'include_ETG': True, # to toggle ETG modes on or off\n",
- " 'avoid_big_negative_s': False,\n",
+ " 'transport_models': [\n",
+ " {\n",
+ " 'model_name': 'constant',\n",
+ " 'rho_max': 0.1,\n",
+ " 'chi_i': 0.3,\n",
+ " 'chi_e': 0.3,\n",
+ " 'D_e': 0.15,\n",
+ " 'V_e': 0.0,\n",
+ " },\n",
+ " {\n",
+ " 'model_name': 'qlknn',\n",
+ " 'rho_min': 0.1,\n",
+ " 'rho_max': 0.95,\n",
+ " 'DV_effective': True,\n",
+ " 'include_ITG': True,\n",
+ " 'include_TEM': True,\n",
+ " 'include_ETG': True,\n",
+ " 'avoid_big_negative_s': False,\n",
+ " },\n",
+ " {\n",
+ " 'model_name': 'constant',\n",
+ " 'rho_min': 0.95,\n",
+ " 'chi_i': 2.0,\n",
+ " 'chi_e': 2.0,\n",
+ " 'D_e': 0.1,\n",
+ " 'V_e': 0.0,\n",
+ " },\n",
+ " ],\n",
" },\n",
" 'solver': {\n",
" 'solver_type': 'linear', # linear solver with picard iteration\n",
@@ -225,15 +230,15 @@
" 'time_step_calculator': {\n",
" 'calculator_type': 'fixed',\n",
" },\n",
- "}\n"
+ "}"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "xSzvOYW_KDd6"
+ "id": "A7MMV9poKDd6"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Detailed summary plot function for a single simulation\n",
@@ -444,13 +449,13 @@
"\n",
" plt.show()\n"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "32fJVuAbQhc1"
+ "id": "xSzvOYW_KDd6"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Comparison timetrace summary plot for multiple simulations\n",
@@ -572,13 +577,13 @@
" plt.tight_layout()\n",
" plt.show()"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "1_eekusPKDd6"
+ "id": "32fJVuAbQhc1"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Comparison profile summary plot for multiple simulations\n",
@@ -683,13 +688,13 @@
" plt.tight_layout()\n",
" plt.show()"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "Xz6YJLQAKDd6"
+ "id": "1_eekusPKDd6"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title set_LH_transition_time_function\n",
@@ -717,13 +722,13 @@
" if not (allowed_range[0] <= input_obj <= allowed_range[1]):\n",
" raise ValueError(f\"Input value must be between {allowed_range[0]} and {allowed_range[1]}.\")\n"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "4-S6_hbNKDd7"
+ "id": "Xz6YJLQAKDd6"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"#@title Modify config function\n",
@@ -839,13 +844,13 @@
" if not (allowed_range[0] <= input_obj <= allowed_range[1]):\n",
" raise ValueError(f\"{variable_name} values must be between {allowed_range[0]} and {allowed_range[1]}.\")\n"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "gPm-tYA5KDd7"
+ "id": "4-S6_hbNKDd7"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"# @title Function to build and launch TORAX sim\n",
@@ -854,13 +859,13 @@
" data_tree, _ = torax.run_simulation(torax_config, log_timestep_info=False)\n",
" return data_tree"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "w15Nr5w1CgzC"
+ "id": "gPm-tYA5KDd7"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "markdown",
"source": [
"# TORAX exercise\n",
@@ -960,12 +965,12 @@
"* Maintain q_min > 1\n",
"* Total input power at t=100s must be > 50MW.\n",
"* Ip cannot exceed 18MA during the current rampup. Note that the current ramp does not need to be monotonic.\n"
- ]
+ ],
+ "metadata": {
+ "id": "w15Nr5w1CgzC"
+ }
},
{
- "metadata": {
- "id": "xdwIAbOBVer2"
- },
"cell_type": "code",
"source": [
"# @title SOLUTION: question 1\n",
@@ -997,26 +1002,26 @@
"# 6. On the other hand, q_min > 1 for the faster rampup rate, whereas q_min<1 for the slower rampup rate since a lot more current has diffused inwards. In some scenarios, q_min > 1 is a constraint to avoid sawteeth instabilities.\n",
"\n"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "OuDaKHE1Ver2"
+ "id": "xdwIAbOBVer2"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"# Optional, can investigate the more detailed plots\n",
"# detailed_plot_single_sim(out0, time = 60)\n",
"# detailed_plot_single_sim(out3, time = 120)"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "_5BQpHCSVer2"
+ "id": "OuDaKHE1Ver2"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"# @title SOLUTION: question 2\n",
@@ -1048,13 +1053,13 @@
"# This leads to a higher q-profile when there is heating during ramp-up. The NBI current also contributes here, since it displaces some of the Ohmic current\n",
"# (which in stationary state peaks on axis), with a broader current profile.\n"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "tpcEG586Ver2"
+ "id": "_5BQpHCSVer2"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"# @title SOLUTION: question 3\n",
@@ -1131,13 +1136,13 @@
"\n",
"# To achieve qmin>1.5, it is necessary to reduce the total current while ensuring that the ECCD deposition is sufficiently off-axis and not too strong"
],
- "outputs": [],
- "execution_count": null
- },
- {
"metadata": {
- "id": "HwWO7cJ9Ver2"
+ "id": "tpcEG586Ver2"
},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
"cell_type": "code",
"source": [
"# @title SOLUTION: question 4\n",
@@ -1168,8 +1173,11 @@
"# 2. Increasing EC power during the current ramp maintains a higher current diffusion time and avoids q<1 on-axis\n",
"# 3. The location and power of the ECCD makes sure that q>1 in the deposition region"
],
- "outputs": [],
- "execution_count": null
+ "metadata": {
+ "id": "HwWO7cJ9Ver2"
+ },
+ "execution_count": null,
+ "outputs": []
}
],
"metadata": {
@@ -1199,6 +1207,6 @@
"version": "3.12.8"
}
},
- "nbformat": 4,
- "nbformat_minor": 0
+ "nbformat_minor": 0,
+ "nbformat": 4
}