Implement Waltz rotation rule.

PiperOrigin-RevId: 866495099

Author: hamelphi

docs/configuration.rst

@@ -1365,6 +1365,24 @@ It is recommended to not set ``qlknn_model_name``,  or
 
   * ``full_radius``: The rotation correction is applied everywhere.
 
+``shear_suppression_model`` (str [default = 'waltz_rule'])
+  Selects the shear suppression model used for rotation effects on transport.
+  Options are:
+
+  * ``waltz_rule``: Simple model from Waltz et al., PoP 1998
+    (https://doi.org/10.1063/1.872847): :math:`f_{rot} = -\alpha`.
+
+  * ``vandeplassche2020``: Fitted rotation rule from Van de Plassche et al.,
+    PoP 2020 (https://doi.org/10.1063/1.5134126).
+
+``shear_suppression_alpha`` (float [default = 1.0])
+  Alpha parameter for the Waltz rule shear suppression model. Only used when
+  ``shear_suppression_model`` is ``'waltz_rule'``.
+
+``output_mode_contributions`` (bool [default = False])
+  If ``True``, output individual ITG/TEM/ETG contributions to transport
+  coefficients in addition to the total values.
+
 tglfnn-ukaea
 ^^^^^^^^^^^^
 

docs/physics_models.rst

@@ -395,11 +395,26 @@ be enabled through the transport model configuration.
     .. math::
       f_{rot} = 1 + f_{rule} \frac{\gamma_{E \times B}}{\gamma_{max}}
 
-    Here, :math:`f_{rule}` is a factor derived from experimental observations,
-    depending on the safety factor, magnetic shear, and inverse aspect ratio.
-    This rule effectively suppresses turbulent transport when :math:`E \times B`
-    shear is strong. The application of this rule is controlled by the
-    `rotation_mode` configuration parameter. Options for `rotation_mode` are:
+    The factor :math:`f_{rule}` determines the strength of shear suppression
+    and can be computed using one of two models, controlled by the
+    ``shear_suppression_model`` configuration parameter:
+
+    * ``waltz_rule``: Simple model from Waltz et al., PoP 1998:
+      :math:`f_{rule} = -\alpha`, where :math:`\alpha` is set by the
+      ``shear_suppression_alpha`` parameter (default 1.0). This model provides
+      pure suppression of turbulent transport.
+
+    * ``vandeplassche2020``: Fitted rotation rule from Van de Plassche et al.,
+      PoP 2020. This model calculates :math:`f_{rule}` based on the safety
+      factor, magnetic shear, and inverse aspect ratio. Note that this model
+      can produce both suppression and enhancement of transport depending on
+      local plasma parameters. This is because it was fitted for purely
+      toroidal rotation, which has both stabilizing (ExB shear) and
+      destabilizing (parallel velocity shear) effects, with a geometry
+      dependence setting the relative impact.
+
+    The application of the rotation rule is controlled by the
+    ``rotation_mode`` configuration parameter. Options for ``rotation_mode`` are:
 
   * ``off``: No rotation correction is applied.
 
@@ -413,9 +428,9 @@ be enabled through the transport model configuration.
 Two parameters are available to fine-tune the impact of rotation, both of them
 default to 1.0:
 
-*   `rotation_multiplier`: Located in the transport model configs, this
+*   ``rotation_multiplier``: Located in the transport model configs, this
     parameter scales the :math:`E \times B` shear term.
-*   `poloidal_velocity_multiplier`: Found under the `neoclassical`
+*   ``poloidal_velocity_multiplier``: Found under the ``neoclassical``
     configuration, this parameter directly scales the poloidal velocity term.
 
 Edge models

torax/_src/transport_model/pydantic_model.py

@@ -108,6 +108,13 @@ class QLKNNTransportModel(pydantic_model_base.TransportBase):
       D.
     rotation_multiplier: Multiplier for rotation.
     rotation_mode: Mode for rotation, either HALF_RADIUS, FULL_RADIUS or OFF.
+    shear_suppression_model: Shear suppression model for rotation effects.
+      Either WALTZ_RULE [Waltz et al., PoP
+      1998](https://doi.org/10.1063/1.872847) or VANDEPLASSCHE2020 [Van de
+      Plassche et al. PoP 2020](https://doi.org/10.1063/1.5134126).
+    shear_suppression_alpha: Alpha parameter for Waltz rule. Larger values
+      increase the suppression effect. Only used when shear_suppression_model =
+      WALTZ_RULE.
   """
 
   model_name: Annotated[Literal['qlknn'], torax_pydantic.JAX_STATIC] = 'qlknn'
@@ -130,6 +137,10 @@ class QLKNNTransportModel(pydantic_model_base.TransportBase):
   rotation_mode: Annotated[
       qualikiz_based_transport_model.RotationMode, torax_pydantic.JAX_STATIC
   ] = qualikiz_based_transport_model.RotationMode.OFF
+  shear_suppression_model: Annotated[
+      qlknn_transport_model.ShearSuppressionModel, torax_pydantic.JAX_STATIC
+  ] = qlknn_transport_model.ShearSuppressionModel.WALTZ_RULE
+  shear_suppression_alpha: pydantic.NonNegativeFloat = 1.0
   output_mode_contributions: Annotated[bool, torax_pydantic.JAX_STATIC] = False
 
   @pydantic.model_validator(mode='before')
@@ -184,6 +195,8 @@ def build_runtime_params(
         An_min=self.An_min,
         rotation_multiplier=self.rotation_multiplier,
         rotation_mode=self.rotation_mode,
+        shear_suppression_model=self.shear_suppression_model,
+        shear_suppression_alpha=self.shear_suppression_alpha,
         output_mode_contributions=self.output_mode_contributions,
         **base_kwargs,
     )

torax/_src/transport_model/qlknn_transport_model.py

@@ -14,6 +14,7 @@
 
 """A transport model that uses a QLKNN model."""
 import dataclasses
+import enum
 import functools
 import logging
 import os
@@ -37,16 +38,37 @@
 
 
 # pylint: disable=invalid-name
+class ShearSuppressionModel(enum.StrEnum):
+  """Shear suppression model for rotation effects on transport.
+
+  WALTZ_RULE: Simple model from Waltz et al., PoP 1998
+  (https://doi.org/10.1063/1.872847): f_rot_rule = -alpha.
+  VANDEPLASSCHE2020: Fitted rotation rule from Van de Plassche et al., PoP 2020
+  (https://doi.org/10.1063/1.5134126). It was fitted for purely toroidal
+  transport, so has both stabilizing (ExB shear) and destabilizing (parallel
+  velocity shear) effects, with a geometry dependence setting the relative
+  impact.
+  """
+
+  WALTZ_RULE = 'waltz_rule'
+  VANDEPLASSCHE2020 = 'vandeplassche2020'
+
+
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class RuntimeParams(qualikiz_based_transport_model.RuntimeParams):
+  """Runtime parameters for QLKNN transport model."""
   include_ITG: bool
   include_TEM: bool
   include_ETG: bool
   ITG_flux_ratio_correction: float
   ETG_correction_factor: float
   clip_inputs: bool
   clip_margin: float
+  shear_suppression_model: ShearSuppressionModel = dataclasses.field(
+      metadata={'static': True}
+  )
+  shear_suppression_alpha: float
   output_mode_contributions: bool = dataclasses.field(metadata={'static': True})
 
 
@@ -174,27 +196,27 @@ def clip_inputs(
   return feature_scan
 
 
-def _calculate_rotation_rule_factor(
-    qualikiz_inputs: qualikiz_based_transport_model.QualikizInputs,
-) -> jax.Array:
-  """Calculate the rotation scaling factor."""
-  # TODO(b/456456279): The rotation rule has been calibrated for QLKNN10D.
-  # We should validate if this formula makes sense for QLKNN_7_11.
-  return (
-      _C1 * qualikiz_inputs.q
-      + _C2 * qualikiz_inputs.smag
-      + _C3 / (_EPSILON_NN * qualikiz_inputs.x)
-      - _C4
-  )
-
-
 def _maybe_apply_rotation_rule(
     model_output: base_qlknn_model.ModelOutput,
     qualikiz_inputs: qualikiz_based_transport_model.QualikizInputs,
     rotation_mode: qualikiz_based_transport_model.RotationMode,
+    shear_suppression_model: ShearSuppressionModel,
+    shear_suppression_alpha: float,
     geo: geometry.Geometry,
 ) -> base_qlknn_model.ModelOutput:
-  """Apply the rotation scaling factor to the model output (Victor rule)."""
+  """Apply the rotation scaling factor to the model output.
+
+  Args:
+    model_output: The raw model output from QLKNN.
+    qualikiz_inputs: Prepared inputs for the quasilinear model.
+    rotation_mode: Controls where the rotation rule is applied.
+    shear_suppression_model: Which shear suppression model to use.
+    shear_suppression_alpha: Alpha parameter for Waltz rule.
+    geo: Geometry of the torus.
+
+  Returns:
+    Model output with rotation scaling applied to ITG and TEM modes.
+  """
   if rotation_mode == qualikiz_based_transport_model.RotationMode.OFF:
     # Rotation is disabled. Do not apply the rotation rule.
     return model_output
@@ -212,7 +234,21 @@ def _maybe_apply_rotation_rule(
     )
     gamma_E_GB = gamma_E_GB * scaling
 
-  f_rot_rule = _calculate_rotation_rule_factor(qualikiz_inputs)
+  if shear_suppression_model == ShearSuppressionModel.WALTZ_RULE:
+    c1, c2, c3, c4 = 0.0, 0.0, 0.0, shear_suppression_alpha
+  elif shear_suppression_model == ShearSuppressionModel.VANDEPLASSCHE2020:
+    c1, c2, c3, c4 = _C1, _C2, _C3, _C4
+  else:
+    raise ValueError(
+        f'Unknown shear suppression model: {shear_suppression_model}'
+    )
+
+  f_rot_rule = (
+      c1 * qualikiz_inputs.q
+      + c2 * qualikiz_inputs.smag
+      + c3 / (_EPSILON_NN * qualikiz_inputs.x)
+      - c4
+  )
 
   lower_bound = 1e-4
   gamma_max = jnp.maximum(model_output['gamma_max'].squeeze(), lower_bound)
@@ -347,6 +383,8 @@ def _combined(
         model_output,
         qualikiz_inputs,
         runtime_config_inputs.transport.rotation_mode,
+        runtime_config_inputs.transport.shear_suppression_model,
+        runtime_config_inputs.transport.shear_suppression_alpha,
         geo,
     )