Combine Waltz and Victor rules for E×B shear suppression

Refactors the rotation/shear suppression in the QLKNN transport model to simultaneously apply both the Waltz rule and the Victor rule (Van de Plassche 2020) to different physical components of the E×B shear, rather than treating them as mutually exclusive options.

The E×B shearing rate is decomposed into:

Poloidal/pressure contributions (from ∇Pᵢ and v_θB_φ): Waltz rule applied, giving pure suppression f_waltz = -α.
Toroidal contribution (from -v_φB_θ): Victor rule applied, a fitted model that can produce both suppression and enhancement depending on local plasma parameters.
The combined scaling factor becomes:

f_rot = clip(1 + f_waltz * |γ_pp| / γ_max + f_victor * |γ_tor| / γ_max, 0)
Changes:

rotation.py: Introduce RotationOutput dataclass. Split _calculate_radial_electric_field to return separate Er components for poloidal/pressure and toroidal contributions.
qualikiz_based_transport_model.py: Propagate separate v_ExB components through to QualikizInputs with individual gamma_E_GB_poloidal_pressure and gamma_E_GB_toroidal shearing rates.
qlknn_transport_model.py: Apply both rules simultaneously. Remove ShearSuppressionModel enum.
pydantic_model.py: Remove shear_suppression_model config field. Retain shear_suppression_alpha for Waltz rule strength.
tglf_based_transport_model.py, post_processing.py: Update to new RotationOutput API.
Docs: Update physics_models.rst to describe combined model; add Waltz 1998 reference to links.rst.
Tests: Remove separate Waltz suppression sim test. Update unit tests and regenerate reference .nc data.
PiperOrigin-RevId: 868764762

Author: hamelphi

docs/configuration.rst

@@ -1365,19 +1365,9 @@ 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'``.
+  Alpha parameter for the Waltz rule shear suppression applied to
+  poloidal/pressure contributions.
 
 ``output_mode_contributions`` (bool [default = False])
   If ``True``, output individual ITG/TEM/ETG contributions to transport

docs/links.rst

@@ -26,6 +26,7 @@
 .. _mavrin2017_link: https://doi.org/10.1007/s10894-017-0136-z
 .. _kim1991_link: https://doi.org/10.1063/1.859671
 .. _hinton1976_link: https://doi.org/10.1103/RevModPhys.48.239
+.. _waltz1998_link: https://doi.org/10.1063/1.872847
 
 .. Define substitutions using link targets
 .. |qlknn10d| replace:: `[van de Plassche et al, Phys. Plasmas 2020] <qlknn10d_link_>`_
@@ -55,3 +56,4 @@
 .. |mavrin2017| replace:: `[Mavrin, J Fusion Energ 2017] <mavrin2017_link_>`_
 .. |kim1991| replace:: `[Kim et al, Phys. Fluids B 1991] <kim1991_link_>`_
 .. |hinton1976| replace:: `[Hinton & Hazeltine, Rev. Mod. Phys. 1976] <hinton1976_link_>`_
+.. |waltz1998| replace:: `[Waltz et al, Phys. Plasmas 1998] <waltz1998_link_>`_

docs/physics_models.rst

@@ -386,42 +386,39 @@ be enabled through the transport model configuration.
     `True` in the transport model configuration.
 
 *   **QLKNN:** The QLKNN transport model
-    incorporates a "rotation rule" that reduces turbulent fluxes (specifically
-    for ITG and TEM modes), see |qlknn10d|. This reduction is based on the
-    ratio of the :math:`E \times B` shearing rate (:math:`\gamma_{E \times B}`)
-    to the maximum linear growth rate (:math:`\gamma_{max}`). The scaling
-    factor is calculated as:
+    incorporates a "rotation rule" that modifies turbulent fluxes (specifically
+    for ITG and TEM modes) based on E×B shear. The modification combines two
+    physical effects:
+
+    * **Waltz rule** applied to poloidal/pressure contributions (from
+      :math:`\nabla P_i` and :math:`v_\theta B_\phi`): Simple suppression
+      :math:`f_{waltz} = -\alpha`, where :math:`\alpha` is set by the
+      ``shear_suppression_alpha`` parameter (default 1.0). See |waltz1998|.
+
+    * **Victor rule** applied to toroidal contributions (from
+      :math:`-v_\phi B_\theta`): A fitted model that can produce both
+      suppression and enhancement depending on local plasma parameters,
+      due to competing stabilizing (E×B shear) and destabilizing (parallel
+      velocity shear) effects. See |qlknn10d|.
+
+    The combined scaling factor is:
 
     .. math::
-      f_{rot} = 1 + f_{rule} \frac{\gamma_{E \times B}}{\gamma_{max}}
-
-    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.
+      f_{rot} = \text{clip}\left(1 + f_{waltz}\frac{|\gamma_{pp}|}{\gamma_{max}}
+      + f_{victor}\frac{|\gamma_{tor}|}{\gamma_{max}}, 0\right)
+
+    where :math:`\gamma_{pp}` and :math:`\gamma_{tor}` are the E×B shearing
+    rates from poloidal/pressure and toroidal contributions respectively.
 
     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.
+    * ``off``: No rotation correction is applied.
 
-  * ``half_radius``: The rotation correction is only applied to the outer
-    half of the radius (:math:`\hat{\rho} > 0.5`).
+    * ``half_radius``: The rotation correction is only applied to the outer
+      half of the radius (:math:`\hat{\rho} > 0.5`).
 
-  * ``full_radius``: The rotation correction is applied everywhere.
+    * ``full_radius``: The rotation correction is applied everywhere.
 
 **Tuning Rotation Effects:**
 

torax/_src/output_tools/post_processing.py

@@ -956,7 +956,7 @@ def cumulative_values():
       sim_state.core_profiles, sim_state.geometry
   )
 
-  _, radial_electric_field, poloidal_velocity = rotation.calculate_rotation(
+  rotation_output = rotation.calculate_rotation(
       T_i=sim_state.core_profiles.T_i,
       psi=sim_state.core_profiles.psi,
       n_i=sim_state.core_profiles.n_i,
@@ -1039,8 +1039,8 @@ def cumulative_values():
       beta_pol=beta_pol,
       beta_N=beta_N,
       impurity_species=impurity_radiation_outputs,
-      poloidal_velocity=poloidal_velocity.face_value(),
-      radial_electric_field=radial_electric_field.face_value(),
+      poloidal_velocity=rotation_output.poloidal_velocity.face_value(),
+      radial_electric_field=rotation_output.Er.face_value(),
       first_step=jnp.array(False),
   )
 

torax/_src/physics/rotation.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Calculations related to the rotation of the plasma."""
+import dataclasses
 
 from jax import numpy as jnp
 from torax._src import array_typing
@@ -24,6 +25,27 @@
 
 
 # pylint: disable=invalid-name
+@dataclasses.dataclass(frozen=True)
+class RotationOutput:
+  """Structured output from rotation calculations.
+
+  Attributes:
+    v_ExB: Total ExB velocity profile on the face grid [m/s].
+    v_ExB_poloidal_pressure: v_ExB from poloidal rotation + pressure gradient
+      contributions: (dpi/dr / (Zi * e * ni) + v_theta * B_phi) / B_total.
+    v_ExB_toroidal: v_ExB from toroidal velocity contribution: -v_phi * B_theta
+      / B_total.
+    Er: Total radial electric field as a cell variable [V/m].
+    poloidal_velocity: Poloidal velocity as a cell variable [m/s].
+  """
+
+  v_ExB: array_typing.FloatVectorFace
+  v_ExB_poloidal_pressure: array_typing.FloatVectorFace
+  v_ExB_toroidal: array_typing.FloatVectorFace
+  Er: cell_variable.CellVariable
+  poloidal_velocity: cell_variable.CellVariable
+
+
 def _calculate_radial_electric_field(
     pressure_thermal_i: cell_variable.CellVariable,
     toroidal_angular_velocity: cell_variable.CellVariable,
@@ -33,11 +55,21 @@ def _calculate_radial_electric_field(
     B_pol_face: array_typing.FloatVector,
     B_tor_face: array_typing.FloatVector,
     geo: geometry.Geometry,
-) -> cell_variable.CellVariable:
-  """Calculates the radial electric field Er.
+) -> tuple[
+    cell_variable.CellVariable,
+    array_typing.FloatVectorFace,
+    array_typing.FloatVectorFace,
+]:
+  """Calculates the radial electric field Er with separate components.
 
   Er = (1 / (Zi * e * ni)) * dpi/dr - v_phi * B_theta + v_theta * B_phi
 
+  We split Er into:
+  - Er_poloidal_pressure: (1 / (Zi * e * ni)) * dpi/dr + v_theta * B_phi
+    This contribution is from the pressure gradient and poloidal rotation.
+  - Er_toroidal: -v_phi * B_theta
+    This contribution is from toroidal velocity.
+
   Args:
     pressure_thermal_i: Pressure profile as a cell variable.
     toroidal_angular_velocity: Toroidal velocity profile as a cell variable.
@@ -49,7 +81,10 @@ def _calculate_radial_electric_field(
     geo: Geometry object.
 
   Returns:
-    Er: Radial electric field [V/m] on the cell grid.
+    Er: Radial electric field [V/m] as a CellVariable.
+    Er_poloidal_pressure_face: Er contribution from pressure gradient and
+      poloidal rotation [V/m] on face grid.
+    Er_toroidal_face: Er contribution from toroidal velocity [V/m] on face grid.
   """
   # Calculate dpi/dr with respect to a midplane-averaged radial coordinate.
   dpi_dr = pressure_thermal_i.face_grad(
@@ -58,19 +93,27 @@ def _calculate_radial_electric_field(
 
   # Calculate Er
   denominator = Z_i_face * constants.CONSTANTS.q_e * n_i.face_value()
-  Er = (
+
+  Er_poloidal_pressure_face = (
       math_utils.safe_divide(jnp.array(1.0), denominator) * dpi_dr
-      - toroidal_angular_velocity.face_value()
+      + poloidal_velocity.face_value() * B_tor_face
+  )
+
+  Er_toroidal_face = (
+      -toroidal_angular_velocity.face_value()
       * geo.R_major_profile_face
       * B_pol_face
-      + poloidal_velocity.face_value() * B_tor_face
   )
-  return cell_variable.CellVariable(
-      value=geometry.face_to_cell(Er),
+
+  Er_face = Er_poloidal_pressure_face + Er_toroidal_face
+
+  Er = cell_variable.CellVariable(
+      value=geometry.face_to_cell(Er_face),
       face_centers=geo.rho_face_norm,
-      right_face_constraint=Er[-1],
+      right_face_constraint=Er_face[-1],
       right_face_grad_constraint=None,
   )
+  return Er, Er_poloidal_pressure_face, Er_toroidal_face
 
 
 def _calculate_v_ExB(
@@ -93,11 +136,7 @@ def calculate_rotation(
     pressure_thermal_i: cell_variable.CellVariable,
     geo: geometry.Geometry,
     poloidal_velocity_multiplier: array_typing.FloatScalar = 1.0,
-) -> tuple[
-    array_typing.FloatVectorFace,
-    cell_variable.CellVariable,
-    cell_variable.CellVariable,
-]:
+) -> RotationOutput:
   """Calculates quantities related to the rotation of the plasma.
 
   Args:
@@ -114,9 +153,8 @@ def calculate_rotation(
       velocity.
 
   Returns:
-    v_ExB: ExB velocity profile on the face grid [m/s].
-    Er: Radial electric field as a cell variable [V/m] .
-    poloidal_velocity: Poloidal velocity as a cell variable [m/s].
+    RotationOutput with v_ExB, separated v_ExB components, Er, and
+    poloidal_velocity.
   """
 
   # Flux surface average of `B_phi = F/R`.
@@ -142,17 +180,29 @@ def calculate_rotation(
       poloidal_velocity_multiplier=poloidal_velocity_multiplier,
   )
 
-  Er = _calculate_radial_electric_field(
-      pressure_thermal_i=pressure_thermal_i,
-      toroidal_angular_velocity=toroidal_angular_velocity,
-      poloidal_velocity=poloidal_velocity,
-      n_i=n_i,
-      Z_i_face=Z_i_face,
-      B_pol_face=B_pol_face,
-      B_tor_face=B_tor_face,
-      geo=geo,
+  Er, Er_poloidal_pressure_face, Er_toroidal_face = (
+      _calculate_radial_electric_field(
+          pressure_thermal_i=pressure_thermal_i,
+          toroidal_angular_velocity=toroidal_angular_velocity,
+          poloidal_velocity=poloidal_velocity,
+          n_i=n_i,
+          Z_i_face=Z_i_face,
+          B_pol_face=B_pol_face,
+          B_tor_face=B_tor_face,
+          geo=geo,
+      )
   )
 
   v_ExB = _calculate_v_ExB(Er.face_value(), B_total_face)
+  v_ExB_poloidal_pressure = _calculate_v_ExB(
+      Er_poloidal_pressure_face, B_total_face
+  )
+  v_ExB_toroidal = _calculate_v_ExB(Er_toroidal_face, B_total_face)
 
-  return v_ExB, Er, poloidal_velocity
+  return RotationOutput(
+      v_ExB=v_ExB,
+      v_ExB_poloidal_pressure=v_ExB_poloidal_pressure,
+      v_ExB_toroidal=v_ExB_toroidal,
+      Er=Er,
+      poloidal_velocity=poloidal_velocity,
+  )

torax/_src/physics/tests/rotation_test.py

@@ -31,7 +31,7 @@ def setUp(self):
     ).build_geometry()
 
   def test_calculate_rotation_shapes_are_correct(self):
-    v_ExB, Er, poloidal_velocity = rotation.calculate_rotation(
+    rotation_output = rotation.calculate_rotation(
         T_i=core_profile_helpers.make_constant_core_profile(
             geo=self.geo, value=1.0
         ),
@@ -50,16 +50,19 @@ def test_calculate_rotation_shapes_are_correct(self):
         ),
         geo=self.geo,
     )
-    self.assertEqual(v_ExB.shape, self.geo.rho_face_norm.shape)
-    self.assertEqual(Er.face_value().shape, self.geo.rho_face_norm.shape)
+    self.assertEqual(rotation_output.v_ExB.shape, self.geo.rho_face_norm.shape)
     self.assertEqual(
-        poloidal_velocity.face_value().shape, self.geo.rho_face_norm.shape
+        rotation_output.Er.face_value().shape, self.geo.rho_face_norm.shape
+    )
+    self.assertEqual(
+        rotation_output.poloidal_velocity.face_value().shape,
+        self.geo.rho_face_norm.shape,
     )
 
   def test_electric_field_is_zero_for_zero_velocities_and_constant_pressure(
       self,
   ):
-    E_r = rotation._calculate_radial_electric_field(
+    E_r, _, _ = rotation._calculate_radial_electric_field(
         pressure_thermal_i=core_profile_helpers.make_constant_core_profile(
             geo=self.geo, value=1.0
         ),
@@ -75,13 +78,12 @@ def test_electric_field_is_zero_for_zero_velocities_and_constant_pressure(
         B_tor_face=np.ones_like(self.geo.rho_face_norm),
         geo=self.geo,
     )
-    # For constant profiles and zero velocities, E_r should be zero.
     np.testing.assert_allclose(
         E_r.value, np.zeros_like(self.geo.rho_norm), atol=1e-12
     )
 
   def test_electric_field_is_not_zero_for_toroidal_velocity(self):
-    E_r = rotation._calculate_radial_electric_field(
+    E_r, _, _ = rotation._calculate_radial_electric_field(
         pressure_thermal_i=core_profile_helpers.make_constant_core_profile(
             geo=self.geo, value=1.0
         ),
@@ -97,12 +99,11 @@ def test_electric_field_is_not_zero_for_toroidal_velocity(self):
         B_tor_face=np.ones_like(self.geo.rho_face_norm),
         geo=self.geo,
     )
-    # E_r should not be all zeros when toroidal velocity is non-zero.
     self.assertTrue(np.any(np.abs(E_r.value) > 1e-12))
 
   def test_electric_field_is_not_zero_for_poloidal_velocity(self):
     """Test that radial electric field is not zero for non-zero poloidal velocity."""
-    E_r = rotation._calculate_radial_electric_field(
+    E_r, _, _ = rotation._calculate_radial_electric_field(
         pressure_thermal_i=core_profile_helpers.make_constant_core_profile(
             geo=self.geo, value=1.0
         ),
@@ -118,12 +119,11 @@ def test_electric_field_is_not_zero_for_poloidal_velocity(self):
         B_tor_face=np.ones_like(self.geo.rho_face_norm),
         geo=self.geo,
     )
-    # E_r should not be all zeros when poloidal velocity is non-zero.
     self.assertTrue(np.any(np.abs(E_r.value) > 1e-12))
 
   def test_electric_field_is_not_zero_for_non_constant_pressure(self):
     """Test that radial electric field is not zero for non-zero poloidal velocity."""
-    E_r = rotation._calculate_radial_electric_field(
+    E_r, _, _ = rotation._calculate_radial_electric_field(
         pressure_thermal_i=cell_variable.CellVariable(
             value=np.linspace(1.0, 2.0, self.geo.rho_norm.size),
             face_centers=self.geo.rho_face_norm,
@@ -140,7 +140,6 @@ def test_electric_field_is_not_zero_for_non_constant_pressure(self):
         B_tor_face=np.ones_like(self.geo.rho_face_norm),
         geo=self.geo,
     )
-    # E_r should not be all zeros when poloidal velocity is non-zero.
     self.assertTrue(np.any(np.abs(E_r.value) > 1e-12))
 
   def test_v_ExB_is_zero_for_zero_electric_field(self):