Calculate ExB drift

PiperOrigin-RevId: 832323537

Author: hamelphi

docs/configuration.rst

@@ -339,6 +339,14 @@ time-dependence of temperature, density, and current.
   This setting is ignored for the ad-hoc circular geometry, which has no
   numerical geometry.
 
+``toroidal_velocity`` (**time-varying-array** | None [default = None])
+  Toroidal velocity profile. If not provided, the velocity will be set to zero.
+
+``toroidal_velocity_right_bc`` (**time-varying-scalar** | None [default = None])
+  Toroidal velocity boundary condition for r=a_minor. If this is ``None`` the
+  boundary condition will instead be taken from ``toroidal_velocity`` at
+  :math:`\hat{\rho}=1`.
+
 .. _numerics_dataclass:
 
 numerics

docs/output.rst

@@ -301,6 +301,9 @@ These are called out in the list of profiles below, and generate relate to:
 ``Phi`` (time, rho_norm)
   Toroidal magnetic flux at each radial grid point [:math:`Wb`].
 
+``poloidal_velocity`` (time, rho_norm)
+  Poloidal velocity [:math:`m/s`].
+
 ``pprime`` (time, rho_face_norm)
   Derivative of total pressure with respect to poloidal flux [:math:`Pa/Wb`].
 
@@ -331,6 +334,9 @@ These are called out in the list of profiles below, and generate relate to:
 ``q`` (time, rho_face_norm)
   Safety factor profile on the face grid [dimensionless].
 
+``radial_electric_field`` (time, rho_norm)
+  Radial electric field [:math:`V/m`].
+
 ``radiation_impurity_species`` (impurity_symbol, time, rho_cell_norm)
   Impurity radiation power density per species [:math:`W/m^3`]. Only output
   if the ``mavrin_fit`` model is active for ``impurity_radiation``. In this
@@ -371,6 +377,9 @@ These are called out in the list of profiles below, and generate relate to:
 ``T_i`` (time, rho_norm)
   Ion temperature [:math:`keV`].
 
+``toroidal_velocity`` (time, rho_norm)
+  Toroidal velocity [:math:`m/s`].
+
 ``v_loop`` (time, rho_norm)
   Loop voltage profile :math:`V_{loop}=\frac{\partial\psi}{\partial t}`
   [:math:`V`].

torax/_src/core_profiles/getters.py

@@ -161,6 +161,24 @@ def get_updated_electron_density(
   return n_e
 
 
+def get_updated_toroidal_velocity(
+    profile_conditions_params: profile_conditions.RuntimeParams,
+    geo: geometry.Geometry,
+) -> cell_variable.CellVariable:
+  """Gets initial and/or prescribed toroidal velocity profiles."""
+  if profile_conditions_params.toroidal_velocity is None:
+    value = jnp.zeros_like(geo.rho)
+  else:
+    value = profile_conditions_params.toroidal_velocity
+  toroidal_velocity = cell_variable.CellVariable(
+      value=value,
+      right_face_grad_constraint=None,
+      right_face_constraint=profile_conditions_params.toroidal_velocity_right_bc,
+      dr=geo.drho_norm,
+  )
+  return toroidal_velocity
+
+
 @dataclasses.dataclass(frozen=True)
 class _IonProperties:
   """Helper container for holding ion calculation outputs."""

torax/_src/core_profiles/initialization.py

@@ -32,6 +32,7 @@
 from torax._src.geometry import standard_geometry
 from torax._src.neoclassical import neoclassical_models as neoclassical_models_lib
 from torax._src.neoclassical.bootstrap_current import base as bootstrap_current_base
+from torax._src.physics import formulas
 from torax._src.physics import psi_calculations
 from torax._src.sources import source_models as source_models_lib
 from torax._src.sources import source_profile_builders
@@ -69,7 +70,9 @@ def initial_core_profiles(
       runtime_params.profile_conditions, geo
   )
   ions = getters.get_updated_ions(runtime_params, geo, n_e, T_e)
-
+  toroidal_velocity = getters.get_updated_toroidal_velocity(
+      runtime_params.profile_conditions, geo
+  )
   # Set v_loop_lcfs. Two branches:
   # 1. Set the v_loop_lcfs from profile_conditions if using the v_loop BC option
   # 2. Initialize v_loop_lcfs to 0 if using the Ip boundary condition for psi.
@@ -95,6 +98,20 @@ def initial_core_profiles(
       dr=geo.drho_norm,
   )
 
+  poloidal_velocity = cell_variable.CellVariable(
+      value=jnp.zeros_like(geo.rho),
+      dr=geo.drho_norm,
+      right_face_constraint=0.0,
+      right_face_grad_constraint=None,
+  )
+
+  radial_electric_field = cell_variable.CellVariable(
+      value=jnp.zeros_like(geo.rho),
+      dr=geo.drho_norm,
+      right_face_constraint=0.0,
+      right_face_grad_constraint=None,
+  )
+
   core_profiles = state.CoreProfiles(
       T_i=T_i,
       T_e=T_e,
@@ -121,6 +138,9 @@ def initial_core_profiles(
       j_total=jnp.zeros_like(geo.rho, dtype=jax_utils.get_dtype()),
       j_total_face=jnp.zeros_like(geo.rho_face, dtype=jax_utils.get_dtype()),
       Ip_profile_face=jnp.zeros_like(geo.rho_face, dtype=jax_utils.get_dtype()),
+      toroidal_velocity=toroidal_velocity,
+      poloidal_velocity=poloidal_velocity,
+      radial_electric_field=radial_electric_field,
   )
 
   return _init_psi_and_psi_derived(
@@ -434,6 +454,12 @@ def _calculate_all_psi_dependent_profiles(
       core_profiles,
   )
 
+  # TODO(b/456456279): Make sure poloidal_velocity has been updated.
+  # Calculate radial electric field
+  radial_electric_field = formulas.calculate_radial_electric_field(
+      core_profiles, geo
+  )
+
   # Calculate sources if they have not already been calculated.
   if not sources_are_calculated:
     source_profiles = _get_bootstrap_and_standard_source_profiles(
@@ -485,6 +511,7 @@ def _calculate_all_psi_dependent_profiles(
       psidot=psidot,
       sigma=conductivity.sigma,
       sigma_face=conductivity.sigma_face,
+      radial_electric_field=radial_electric_field,
   )
   return core_profiles
 

torax/_src/core_profiles/profile_conditions.py

@@ -58,6 +58,8 @@ class RuntimeParams:
   # If provided as array, Psi profile defined on the cell grid.
   psi: array_typing.FloatVector | None
   psidot: array_typing.FloatVector | None
+  toroidal_velocity: array_typing.FloatVector | None
+  toroidal_velocity_right_bc: array_typing.FloatScalar | None
   # Electron density profile on the cell grid.
   n_e: array_typing.FloatVector
   nbar: array_typing.FloatScalar
@@ -112,6 +114,11 @@ class ProfileConditions(torax_pydantic.BaseModelFrozen):
       `psidot` will be used instead of the internally calculated one. This is
       useful for cases where an unphysical transient `psidot` from the initial
       `psi` condition needs to be overridden.
+    toroidal_velocity: Prescribed or evolving values for toroidal velocity. If
+      None, toroidal_velocity will be initialized to zero.
+    toroidal_velocity_right_bc: Toroidal velocity boundary condition for
+      r=a_minor. If this is `None` the boundary condition will instead be taken
+      from `toroidal_velocity` at rhon=1.
     n_e: Prescribed or evolving values for electron density at different times.
     normalize_n_e_to_nbar: Whether to renormalize the density profile to have
       the desired line averaged density `nbar`.
@@ -163,6 +170,8 @@ class ProfileConditions(torax_pydantic.BaseModelFrozen):
   )
   psi: torax_pydantic.TimeVaryingArray | None = None
   psidot: torax_pydantic.TimeVaryingArray | None = None
+  toroidal_velocity: torax_pydantic.TimeVaryingArray | None = None
+  toroidal_velocity_right_bc: torax_pydantic.TimeVaryingScalar | None = None
   n_e: torax_pydantic.PositiveTimeVaryingArray = (
       torax_pydantic.ValidatedDefault({0: {0: 1.2e20, 1: 0.8e20}})
   )
@@ -398,6 +407,14 @@ def build_runtime_params(self, t: chex.Numeric) -> RuntimeParams:
           t, grid_type='face_right'
       )
 
+    if self.toroidal_velocity_right_bc is None:
+      if self.toroidal_velocity is None:
+        runtime_params['toroidal_velocity_right_bc'] = 0.0
+      else:
+        runtime_params['toroidal_velocity_right_bc'] = (
+            self.toroidal_velocity.get_value(t, grid_type='face_right')
+        )
+
     if self.n_e_right_bc is None:
       runtime_params['n_e_right_bc'] = self.n_e.get_value(
           t, grid_type='face_right'

torax/_src/core_profiles/tests/convertors_test.py

@@ -83,6 +83,9 @@ def setUp(self):
         j_total=mock.ANY,
         j_total_face=mock.ANY,
         Ip_profile_face=mock.ANY,
+        toroidal_velocity=mock.ANY,
+        poloidal_velocity=mock.ANY,
+        radial_electric_field=mock.ANY,
     )
 
   def test_core_profiles_to_solver_x_tuple(self):

torax/_src/core_profiles/tests/initialization_test.py

@@ -83,6 +83,40 @@ def test_update_psi_from_j(self):
     ).value
     np.testing.assert_allclose(psi, references.psi.value)
 
+  def test_initial_core_profiles_velocities(self):
+    config = default_configs.get_default_config_dict()
+    # Test default initialization (zeros)
+    torax_config = model_config.ToraxConfig.from_dict(config)
+    core_profiles, geo, _ = _get_initial_state(torax_config)
+    np.testing.assert_allclose(
+        core_profiles.toroidal_velocity.value, np.zeros_like(geo.rho)
+    )
+    np.testing.assert_allclose(
+        core_profiles.poloidal_velocity.value, np.zeros_like(geo.rho)
+    )
+
+  def test_initial_toroidal_velocity_from_profile_conditions(self):
+    config = default_configs.get_default_config_dict()
+    # Test initialization with toroidal_velocity from profile_conditions
+    toroidal_velocity_test = np.array([10.0, 20.0, 30.0, 40.0])
+    _, geo, _ = _get_initial_state(
+        model_config.ToraxConfig.from_dict(config)
+    )
+    config['profile_conditions']['toroidal_velocity'] = {
+        0.0: {
+            rho: value
+            for rho, value in zip(geo.rho_norm, toroidal_velocity_test)
+        }
+    }
+    torax_config = model_config.ToraxConfig.from_dict(config)
+    core_profiles, _, _ = _get_initial_state(torax_config)
+    np.testing.assert_allclose(
+        core_profiles.toroidal_velocity.value, toroidal_velocity_test
+    )
+    np.testing.assert_allclose(
+        core_profiles.poloidal_velocity.value, np.zeros_like(geo.rho)
+    )
+
   @parameterized.parameters(
       (
           {0.0: {0.0: 0.0, 1.0: 1.0}},

torax/_src/core_profiles/updaters.py

@@ -115,6 +115,9 @@ def get_prescribed_core_profile_values(
   n_i = ions.n_i.value
   n_impurity = ions.n_impurity.value
   impurity_fractions = ions.impurity_fractions
+  toroidal_velocity = getters.get_updated_toroidal_velocity(
+      runtime_params.profile_conditions, geo
+  ).value
 
   return {
       'T_i': T_i,
@@ -132,6 +135,7 @@ def get_prescribed_core_profile_values(
       'A_impurity_face': ions.A_impurity_face,
       'Z_eff': ions.Z_eff,
       'Z_eff_face': ions.Z_eff_face,
+      'toroidal_velocity': toroidal_velocity,
   }
 
 
@@ -169,7 +173,7 @@ def update_core_profiles_during_step(
       updated_core_profiles.T_e,
   )
 
-  return dataclasses.replace(
+  updated_core_profiles = dataclasses.replace(
       updated_core_profiles,
       n_i=ions.n_i,
       n_impurity=ions.n_impurity,
@@ -187,6 +191,18 @@ def update_core_profiles_during_step(
       s_face=psi_calculations.calc_s_face(geo, updated_core_profiles.psi),
   )
 
+  # TODO(b/456456279): Make sure poloidal_velocity has been updated.
+  # Computing the radial electric field requires the updated ions densities.
+  radial_electric_field = formulas.calculate_radial_electric_field(
+      updated_core_profiles, geo
+  )
+
+  updated_core_profiles = dataclasses.replace(
+      updated_core_profiles,
+      radial_electric_field=radial_electric_field,
+  )
+  return updated_core_profiles
+
 
 def update_core_and_source_profiles_after_step(
     dt: array_typing.FloatScalar,
@@ -280,16 +296,28 @@ def update_core_and_source_profiles_after_step(
       j_total=j_total,
       j_total_face=j_total_face,
       Ip_profile_face=Ip_profile_face,
+      toroidal_velocity=updated_core_profiles_t_plus_dt.toroidal_velocity,
+      poloidal_velocity=updated_core_profiles_t_plus_dt.poloidal_velocity,
+      radial_electric_field=(
+          updated_core_profiles_t_plus_dt.radial_electric_field
+      ),  # Not yet updated
   )
 
   conductivity = neoclassical_models.conductivity.calculate_conductivity(
       geo, intermediate_core_profiles
   )
 
+  # TODO(b/456456279): Make sure poloidal_velocity has been updated.
+  # Computing the radial electric field requires the updated ions densities.
+  radial_electric_field = formulas.calculate_radial_electric_field(
+      intermediate_core_profiles, geo
+  )
+
   intermediate_core_profiles = dataclasses.replace(
       intermediate_core_profiles,
       sigma=conductivity.sigma,
       sigma_face=conductivity.sigma_face,
+      radial_electric_field=radial_electric_field,
   )
 
   # build_source_profiles calculates the union with explicit + implicit

torax/_src/neoclassical/pydantic_model.py

@@ -40,6 +40,7 @@ class Neoclassical(torax_pydantic.BaseModelFrozen):
   transport: (
       transport_zeros.ZerosModelConfig | angioni_sauter.AngioniSauterModelConfig
   ) = pydantic.Field(discriminator="model_name")
+  compute_poloidal_velocity: bool = False
 
   @pydantic.model_validator(mode="before")
   @classmethod
@@ -66,6 +67,11 @@ def build_runtime_params(self) -> runtime_params.RuntimeParams:
     )
 
   def build_models(self) -> neoclassical_models.NeoclassicalModels:
+    if self.compute_poloidal_velocity:
+      # TODO(b/376326615): Implement poloidal velocity computation.
+      raise NotImplementedError(
+          "Computation of poloidal velocity is not yet implemented."
+      )
     return neoclassical_models.NeoclassicalModels(
         conductivity=self.conductivity.build_model(),
         bootstrap_current=self.bootstrap_current.build_model(),

torax/_src/neoclassical/tests/pydantic_model_test.py

@@ -81,6 +81,16 @@ def f(x: pydantic_model.Neoclassical):
         output, pydantic_model.runtime_params.RuntimeParams
     )
 
+  def test_compute_poloidal_velocity_raises_error(self):
+    model = pydantic_model.Neoclassical.from_dict(
+        {"compute_poloidal_velocity": True}
+    )
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        "Computation of poloidal velocity is not yet implemented",
+    ):
+      model.build_models()
+
 
 if __name__ == "__main__":
   absltest.main()