Consolidate rotation calculations and compute ExB drift

PiperOrigin-RevId: 837611623

Author: hamelphi

torax/_src/neoclassical/formulas.py

@@ -13,14 +13,16 @@
 # limitations under the License.
 """Common formulas used in neoclassical models."""
 
+import jax
 import jax.numpy as jnp
 from torax._src import array_typing
 from torax._src import constants
+from torax._src.fvm import cell_variable
 from torax._src.geometry import geometry as geometry_lib
-
-# pylint: disable=invalid-name
+from torax._src.physics import collisions
 
 
+# pylint: disable=invalid-name
 def calculate_f_trap(
     geo: geometry_lib.Geometry,
 ) -> array_typing.FloatVectorFace:
@@ -180,3 +182,127 @@ def calculate_nu_i_star(
           * (geo.epsilon_face + constants.CONSTANTS.eps) ** 1.5
       )
   )
+
+
+# Functions to calculate the neoclassical poloidal velocity.
+def _calculate_neoclassical_k_neo(
+    nu_star: array_typing.FloatScalar, epsilon: array_typing.FloatScalar
+):
+  """Calculates the neoclassical coefficient k_neo.
+
+  Equation (6.135) from
+  Hinton, F. L., & Hazeltine, R. D.,
+  "Theory of plasma transport in toroidal confinement systems"
+  Rev. Mod. Phys. 48(2), 239–308. (1976)
+  https://doi.org/10.1103/RevModPhys.48.239
+
+  Limits:
+    - Banana regime (nu_star -> 0): ~1.17
+    - Pfirsch-Schluter regime (nu_star -> inf): ~ -2.1
+
+  Args:
+    nu_star : The normalized ion collisionality.
+    epsilon : The inverse aspect ratio.
+
+  Returns:
+    k_neo : The neoclassical coefficient.
+  """
+  # Calculate the first term (Banana-Plateau transition)
+  # (1.17 - 0.35 * sqrt(nu)) / (1 + 0.7 * sqrt(nu))
+  sqrt_nu = jnp.sqrt(nu_star)
+  term1 = (1.17 - 0.35 * sqrt_nu) / (1.0 + 0.7 * sqrt_nu)
+
+  # Calculate the second term (Pfirsch-Schluter driver)
+  # 2.1 * nu^2 * epsilon^3
+  ps_factor = (nu_star**2) * (epsilon**3)
+  term2 = 2.1 * ps_factor
+
+  # Calculate the final denominator (Switching function)
+  # 1 + nu^2 * epsilon^3
+  denominator = 1.0 + ps_factor
+
+  return (term1 - term2) / denominator
+
+# TODO(b/381199010): Implement alternative Sauter-based k_neo calculation.
+# See Sauter (1999) Eq. 17a-17b
+
+
+@jax.jit
+def calculate_poloidal_velocity(
+    T_i: cell_variable.CellVariable,
+    n_i: array_typing.FloatVectorFace,
+    q: array_typing.FloatVectorFace,
+    Z_eff: array_typing.FloatVectorFace,
+    Z_i: array_typing.FloatVectorFace,
+    B_tor: array_typing.FloatVectorFace,
+    B_total_squared: array_typing.FloatVectorFace,
+    geo: geometry_lib.Geometry,
+    rotation_multiplier: array_typing.FloatScalar = 1.0,
+) -> cell_variable.CellVariable:
+  """Computes the neoclassical ion poloidal velocity profile.
+
+  Implementing eq.33 from
+  Y. B. Kim , P. H. Diamond , R. J. Groebner.
+  "Neoclassical poloidal and toroidal rotation in tokamaks"
+  Phys. Fluids B 3, 2050–2060 (1991)
+  https://doi.org/10.1063/1.859671
+
+  Eq. 33 can be simplified to the following form in SI units:
+  v_pol = k_neo * (dT/dr) * (B_tor / <B^2>) / (Z * e)
+
+  Args:
+    T_i: Ion temperature as a cell variable [keV].
+    n_i: Ion density on the face grid [m^-3].
+    q: Safety factor on the face grid.
+    Z_eff: Effective charge on the face grid.
+    Z_i: Main ion charge on the face grid.
+    B_tor: Toroidal magnetic field on the face grid [T].
+    B_total_squared: Total magnetic field (toroidal + poloidal) on the face grid
+      [T].
+    geo : Geometry
+    rotation_multiplier: A multiplier to apply to the poloidal velocity.
+  Returns:
+    v_pol : Poloidal velocity profile [m/s].
+  """
+  # Note: all computations are performed on the face grid.
+
+  T_i_face = T_i.face_value()
+  epsilon = geo.epsilon_face
+
+  # Calculate Neoclassical Coefficient k_i
+  log_lambda_ii = collisions.calculate_log_lambda_ii(
+      T_i_face,
+      n_i,
+      Z_eff,
+  )
+  nu_i_star = calculate_nu_i_star(
+      q=q,
+      geo=geo,
+      n_i=n_i,
+      T_i=T_i_face,
+      Z_eff=Z_eff,
+      log_lambda_ii=log_lambda_ii,
+  )
+  k_neo = _calculate_neoclassical_k_neo(nu_i_star, epsilon)
+
+  # Calculate Radial Temperature Gradient (dT/dr)
+  grad_Ti = T_i.face_grad(geo.r_mid) * constants.CONSTANTS.keV_to_J  # [J/m]
+
+  # Calculate Poloidal Velocity
+  # v_pol = k_i * (dT/dr) * (B_tor / <B^2>) / (Z * e)
+  B_total_squared_safe = jnp.maximum(B_total_squared, constants.CONSTANTS.eps)
+  v_pol = (
+      k_neo
+      * grad_Ti
+      * (B_tor / B_total_squared_safe)
+      / (constants.CONSTANTS.q_e * Z_i)
+  )
+
+  v_pol = rotation_multiplier * v_pol
+
+  return cell_variable.CellVariable(
+      value=geometry_lib.face_to_cell(v_pol),
+      dr=geo.drho_norm,
+      right_face_constraint=v_pol[-1],
+      right_face_grad_constraint=None,
+  )

torax/_src/neoclassical/tests/formulas_test.py

@@ -64,7 +64,7 @@ def setUp(self):
             torax_config
         )
     )
-    runtime_params, geo = (
+    runtime_params, self.geo = (
         build_runtime_params.get_consistent_runtime_params_and_geometry(
             t=torax_config.numerics.t_initial,
             runtime_params_provider=params_provider,
@@ -75,7 +75,7 @@ def setUp(self):
     neoclassical_models = torax_config.neoclassical.build_models()
     self.core_profiles = initialization.initial_core_profiles(
         runtime_params,
-        geo,
+        self.geo,
         source_models=source_models,
         neoclassical_models=neoclassical_models,
     )
@@ -85,14 +85,14 @@ def setUp(self):
     )
     self.nu_e_star = formulas.calculate_nu_e_star(
         q=self.core_profiles.q_face,
-        geo=geo,
+        geo=self.geo,
         n_e=self.core_profiles.n_e.face_value(),
         T_e=self.core_profiles.T_e.face_value(),
         Z_eff=self.core_profiles.Z_eff_face,
         log_lambda_ei=log_lambda_ei,
     )
 
-    self.f_trap = formulas.calculate_f_trap(geo)
+    self.f_trap = formulas.calculate_f_trap(self.geo)
 
   def test_calculate_f_trap_positive_triangularity(self):
     geo = mock.create_autospec(
@@ -128,10 +128,30 @@ def test_L32_values_are_correct(self):
     )
     np.testing.assert_allclose(L32, _L32_EXPECTED, atol=_A_TOL, rtol=_R_TOL)
 
+  def test_calculate_poloidal_velocity_values_are_correct(self):
+    poloidal_velocity = formulas.calculate_poloidal_velocity(
+        T_i=self.core_profiles.T_i,
+        n_i=self.core_profiles.n_i.face_value(),
+        q=self.core_profiles.q_face,
+        Z_eff=self.core_profiles.Z_eff_face,
+        Z_i=self.core_profiles.Z_i_face,
+        B_tor=np.ones_like(self.geo.rho_face_norm),
+        B_total_squared=np.ones_like(self.geo.rho_face_norm),
+        geo=self.geo,
+    )
+    np.testing.assert_allclose(
+        _POLOIDAL_VELOCITY_EXPECTED,
+        poloidal_velocity.face_value(),
+        atol=_A_TOL,
+        rtol=_R_TOL,
+    )
+
 
 # Reference values from running test code in a notebook.
 # The test thus does not directly test the implementation, but rather
 # guards against unexpected modifications.
+# If a change is expected to theese reference values, the new values can b
+# copied/pasted from the logs of a failing test.
 _L31_EXPECTED = np.array([
     0.0,
     0.25942749,
@@ -158,6 +178,19 @@ def test_L32_values_are_correct(self):
     0.08557197,
     0.16296924,
 ])
+_POLOIDAL_VELOCITY_EXPECTED = np.array([
+    -1485.871716,
+    -2507.496827,
+    -3933.755809,
+    -4537.621566,
+    -4854.858931,
+    -5031.592012,
+    -5073.608117,
+    -4858.248803,
+    -3559.941551,
+    3011.279478,
+    17562.499243,
+])
 
 if __name__ == '__main__':
   absltest.main()

torax/_src/physics/formulas.py

@@ -38,6 +38,7 @@
 from torax._src.geometry import geometry
 from torax._src.physics import psi_calculations
 
+
 # pylint: disable=invalid-name
 
 

torax/_src/physics/rotation.py

@@ -0,0 +1,148 @@
+# Copyright 2025 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.
+"""Calculations related to the rotation of the plasma."""
+
+from jax import numpy as jnp
+from torax._src import array_typing
+from torax._src import constants
+from torax._src.fvm import cell_variable
+from torax._src.geometry import geometry
+from torax._src.neoclassical import formulas as neoclassical_formulas
+from torax._src.physics import psi_calculations
+
+
+# pylint: disable=invalid-name
+def _calculate_radial_electric_field(
+    pressure_thermal_i: cell_variable.CellVariable,
+    toroidal_velocity: cell_variable.CellVariable,
+    poloidal_velocity: cell_variable.CellVariable,
+    n_i: cell_variable.CellVariable,
+    Z_i_face: array_typing.FloatVector,
+    B_pol_face: array_typing.FloatVector,
+    B_tor_face: array_typing.FloatVector,
+    geo: geometry.Geometry,
+) -> cell_variable.CellVariable:
+  """Calculates the radial electric field Er.
+
+  Er = (1 / (Zi * e * ni)) * dpi/dr - v_phi * B_theta + v_theta * B_phi
+
+  Args:
+    pressure_thermal_i: Pressure profile as a cell variable.
+    toroidal_velocity: Toroidal velocity profile as a cell variable.
+    poloidal_velocity: Poloidal velocity profile as a cell variable.
+    n_i: Main ion density profile as a cell variable.
+    Z_i_face: Main ion charge on the face grid.
+    B_pol_face: Flux-surface-averaged poloidal magnetic field on the face grid.
+    B_tor_face: Flux-surface-averaged toroidal magnetic field on the face grid.
+    geo: Geometry object.
+
+  Returns:
+    Er: Radial electric field [V/m] on the cell grid.
+  """
+  # Calculate dpi/dr with respect to a midplane-averaged radial coordinate.
+  dpi_dr = pressure_thermal_i.face_grad(geo.r_mid)
+
+  # Calculate Er
+  denominator = Z_i_face * constants.CONSTANTS.q_e * n_i.face_value()
+  Er = (
+      (1.0 / denominator) * dpi_dr
+      - toroidal_velocity.face_value() * B_pol_face
+      + poloidal_velocity.face_value() * B_tor_face
+  )
+  return cell_variable.CellVariable(
+      value=geometry.face_to_cell(Er),
+      dr=geo.drho_norm,
+      right_face_constraint=Er[-1],
+      right_face_grad_constraint=None,
+  )
+
+
+def _calculate_v_ExB(
+    Er_face: array_typing.FloatVectorFace,
+    B_total_face: array_typing.FloatVectorFace,
+) -> array_typing.FloatVectorFace:
+  """Calculates the ExB velocity, on the face grid."""
+  B_total_face = jnp.maximum(B_total_face, constants.CONSTANTS.eps)
+  return jnp.where(B_total_face > 0, Er_face / B_total_face, 0.0)
+
+
+def calculate_rotation(
+    T_i: cell_variable.CellVariable,
+    psi: cell_variable.CellVariable,
+    n_i: cell_variable.CellVariable,
+    q_face: array_typing.FloatVectorFace,
+    Z_eff_face: array_typing.FloatVectorFace,
+    Z_i_face: array_typing.FloatVector,
+    toroidal_velocity: cell_variable.CellVariable,
+    pressure_thermal_i: cell_variable.CellVariable,
+    geo: geometry.Geometry,
+    rotation_multiplier: float = 1.0,
+):
+  """Calculates quantities related to the rotation of the plasma.
+
+  Args:
+    T_i: Ion temperature profile as a cell variable.
+    psi: Poloidal flux profile as a cell variable.
+    n_i: Main ion density profile as a cell variable.
+    q_face: Safety factor on the face grid.
+    Z_eff_face: Effective charge on the face grid.
+    Z_i_face: Main ion charge on the face grid.
+    toroidal_velocity: Toroidal velocity profile as a cell variable.
+    pressure_thermal_i: Pressure profile as a cell variable.
+    geo: Geometry object.
+    rotation_multiplier: A multiplier to apply to the poloidal 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].
+  """
+
+  # Flux surface average of `B_phi = F/R`.
+  B_tor_face = geo.F_face / geo.R_major_profile_face  # Tesla
+
+  # flux-surface-averaged B_theta.
+  B_pol_squared_face = psi_calculations.calc_bpol_squared(
+      geo, psi
+  )  # On the face grid.
+  B_pol_face = jnp.sqrt(B_pol_squared_face)  # Tesla
+  B_total_squared_face = B_pol_squared_face + B_tor_face**2
+  B_total_face = jnp.sqrt(B_total_squared_face)
+
+  poloidal_velocity = neoclassical_formulas.calculate_poloidal_velocity(
+      T_i=T_i,
+      n_i=n_i.face_value(),
+      q=q_face,
+      Z_eff=Z_eff_face,
+      Z_i=Z_i_face,
+      B_tor=B_tor_face,
+      B_total_squared=B_total_squared_face,
+      geo=geo,
+      rotation_multiplier=rotation_multiplier,
+  )
+
+  Er = _calculate_radial_electric_field(
+      pressure_thermal_i=pressure_thermal_i,
+      toroidal_velocity=toroidal_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)
+
+  return v_ExB, Er, poloidal_velocity