Refactor neoclassical models.

docs/configuration.rst

@@ -1691,6 +1691,14 @@ conductivity
   The name of the Sauter model to use. If not provided, the default is to use
   the Sauter model with default values.
 
+transport
+^^^^^^^^^
+``model_name`` (str [default = 'zeros'])
+  The name of the neoclassical transport model. Currently, only the `zeros`
+  model is supported, which sets all neoclassical transport coefficients to
+  zero. This is a placeholder for future implementations of neoclassical
+  transport models.
+
 
 Additional Notes
 ================

docs/interfacing_with_surrogates.rst

@@ -21,14 +21,14 @@ tensors for the neural network.
           dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
           geo: geometry.Geometry,
           core_profiles: state.CoreProfiles,
-      ) -> state.CoreTransport:
+      ) -> TurbulentTransport:
         input_tensor = self._prepare_input(dynamic_runtime_params_slice, geo, core_profiles)
 
         output_tensor = self._call_surrogate_model(input_tensor)
 
         chi_i, chi_e, d_e, v_e = self._parse_output(output_tensor)
 
-        return state.CoreTransport(
+        return TurbulentTransport(
             chi_face_ion=chi_i,
             chi_face_electron=chi_e,
             d_e=d_e,

docs/output.rst

@@ -137,23 +137,33 @@ These are called out in the list of profiles below, and generate relate to:
   Poloidal cross-sectional area of each flux surface [:math:`m^2`].
 
 ``chi_bohm_e`` (time, rho_face_norm) [:math:`m^2/s`]
-  Bohm component of electron heat conductivity. Only output if active.
+  Bohm component of electron heat turbulent conductivity. Only output if active.
 
 ``chi_bohm_i`` (time, rho_face_norm) [:math:`m^2/s`]
-  Bohm component of ion heat conductivity. Only output if active.
+  Bohm component of ion heat turbulent conductivity. Only output if active.
 
 ``chi_gyrobohm_e`` (time, rho_face_norm) [:math:`m^2/s`]
-  Gyro-Bohm component of electron heat conductivity. Only output if active.
+  Gyro-Bohm component of electron heat turbulent conductivity. Only output if
+  active.
 
 ``chi_gyrobohm_i`` (time, rho_face_norm) [:math:`m^2/s`]
-  Gyro-Bohm component of ion heat conductivity. Only output if active.
+  Gyro-Bohm component of ion heat turbulent conductivity. Only output if active.
+
+``chi_neo_e`` (time, rho_face_norm)
+  Neoclassical electron heat conductivity [:math:`m^2/s`].
+
+``chi_neo_i`` (time, rho_face_norm)
+  Neoclassical ion heat conductivity [:math:`m^2/s`].
 
 ``chi_turb_e`` (time, rho_face_norm)
   Total turbulent electron heat conductivity [:math:`m^2/s`].
 
 ``chi_turb_i`` (time, rho_face_norm)
   Total turbulent ion heat conductivity [:math:`m^2/s`].
 
+``D_neo_e`` (time, rho_face_norm)
+  Neoclassical electron particle diffusivity [:math:`m^2/s`].
+
 ``D_turb_e`` (time, rho_face_norm)
   Total turbulent electron particle diffusivity [:math:`m^2/s`].
 
@@ -164,6 +174,9 @@ These are called out in the list of profiles below, and generate relate to:
 ``elongation`` (time, rho_norm)
   Elongation of each flux surface [dimensionless].
 
+``epsilon`` (time, rho_norm)
+  Local inverse aspect ratio at each flux surface [dimensionless].
+
 ``F`` (time, rho_norm)
   Flux function :math:`F=B_{tor}R` , constant on any given flux surface
   [:math:`T m`].
@@ -348,16 +361,24 @@ These are called out in the list of profiles below, and generate relate to:
   Loop voltage profile :math:`V_{loop}=\frac{\partial\psi}{\partial t}`
   [:math:`V`].
 
-``V_turb_e`` (time, rho_face_norm)
-  Turbulent electron particle convection [:math:`m/s`].
-
 ``volume`` (time, rho_norm)
   Plasma volume enclosed by each flux surface [:math:`m^3`].
 
 ``vpr`` (time, rho_norm)
   Derivative of plasma volume enclosed by each flux surface with respect to the
   normalized toroidal flux coordinate rho_norm [:math:`m^3`].
 
+``V_neo_e`` (time, rho_face_norm)
+  Neoclassical electron particle convection velocity [:math:`m/s`]. Contains
+  all components apart from the Ware pinch, which is output separately.
+
+``V_turb_e`` (time, rho_face_norm)
+  Turbulent electron particle convection [:math:`m/s`].
+
+``V_ware_e`` (time, rho_face_norm)
+  Ware pinch electron particle convection velocity [:math:`m/s`], i.e. the
+  component of neoclassical convection arising from the parallel electric field.
+
 ``Z_eff`` (time, rho_norm)
   Effective charge profile defined as
   :math:`(Z_i^2n_i + Z_{impurity}^2n_{impurity})/n_e` [dimensionless].

torax/_src/fvm/calc_coeffs.py

@@ -273,7 +273,7 @@ def calc_coeffs(
       and implicit source profiles used as terms for the core profiles
       equations.
     neoclassical_models: Neoclassical models for the time step. These include
-      the conductivity model, bootstrap current, etc.
+      the conductivity model, bootstrap current, and neoclassical transport.
     pedestal_model: A PedestalModel subclass, calculates pedestal values.
     evolving_names: The names of the evolving variables in the order that their
       coefficients should be written to `coeffs`.
@@ -358,7 +358,7 @@ def _calc_coeffs_full(
       and implicit source profiles used as terms for the core profiles
       equations.
     neoclassical_models: Neoclassical models for the time step. These include
-      the conductivity model, bootstrap current, etc.
+      the conductivity model, bootstrap current, and neoclassical transport.
     pedestal_model: A PedestalModel subclass, calculates pedestal values.
     evolving_names: The names of the evolving variables in the order that their
       coefficients should be written to `coeffs`.
@@ -429,40 +429,50 @@ def _calc_coeffs_full(
   tic_dens_el = geo.vpr
 
   # Diffusion term coefficients
-  transport_coeffs = transport_model(
+  turbulent_transport = transport_model(
       dynamic_runtime_params_slice, geo, core_profiles, pedestal_model_output
   )
-  chi_face_ion = transport_coeffs.chi_face_ion
-  chi_face_el = transport_coeffs.chi_face_el
-  d_face_el = transport_coeffs.d_face_el
-  v_face_el = transport_coeffs.v_face_el
+  neoclassical_transport = (
+      neoclassical_models.transport.calculate_neoclassical_transport(
+          dynamic_runtime_params_slice, geo, core_profiles
+      )
+  )
+
+  chi_face_ion_total = (
+      turbulent_transport.chi_face_ion + neoclassical_transport.chi_neo_i
+  )
+  chi_face_el_total = (
+      turbulent_transport.chi_face_el + neoclassical_transport.chi_neo_e
+  )
+  d_face_el_total = (
+      turbulent_transport.d_face_el + neoclassical_transport.D_neo_e
+  )
+  v_face_el_total = (
+      turbulent_transport.v_face_el
+      + neoclassical_transport.V_neo_e
+      + neoclassical_transport.V_ware_e
+  )
   d_face_psi = geo.g2g3_over_rhon_face
   # No poloidal flux convection term
   v_face_psi = jnp.zeros_like(d_face_psi)
 
-  if static_runtime_params_slice.evolve_density:
-    if d_face_el is None or v_face_el is None:
-      raise NotImplementedError(
-          f'{type(transport_model)} does not support the density equation.'
-      )
-
   # entire coefficient preceding dT/dr in heat transport equations
   full_chi_face_ion = (
       geo.g1_over_vpr_face
       * core_profiles.n_i.face_value()
       * consts.keV2J
-      * chi_face_ion
+      * chi_face_ion_total
   )
   full_chi_face_el = (
       geo.g1_over_vpr_face
       * core_profiles.n_e.face_value()
       * consts.keV2J
-      * chi_face_el
+      * chi_face_el_total
   )
 
   # entire coefficient preceding dne/dr in particle equation
-  full_d_face_el = geo.g1_over_vpr_face * d_face_el
-  full_v_face_el = geo.g0_face * v_face_el
+  full_d_face_el = geo.g1_over_vpr_face * d_face_el_total
+  full_v_face_el = geo.g0_face * v_face_el_total
 
   # density source terms. Initialize source matrix to zero
   source_mat_nn = jnp.zeros_like(geo.rho)
@@ -695,7 +705,7 @@ def _calc_coeffs_full(
       auxiliary_outputs=(
           merged_source_profiles,
           conductivity,
-          transport_coeffs,
+          state.CoreTransport(**turbulent_transport, **neoclassical_transport),
       ),
   )
 

torax/_src/fvm/newton_raphson_solve_block.py

@@ -172,7 +172,7 @@ def newton_raphson_solve_block(
     source_models: Collection of source callables to generate source PDE
       coefficients.
     neoclassical_models: Collection of neoclassical models for calculating
-      conductivity, bootstrap current, etc.
+      conductivity, bootstrap current, and neoclassical transport.
     pedestal_model: Model of the pedestal's behavior.
     coeffs_callback: Calculates diffusion, convection etc. coefficients given a
       core_profiles. Repeatedly called by the iterative optimizer.

torax/_src/fvm/residual_and_loss.py

@@ -243,7 +243,7 @@ def theta_method_block_residual(
     source_models: Collection of source callables to generate source PDE
       coefficients.
     neoclassical_models: Collection of neoclassical models for calculating
-      conductivity, bootstrap current, etc.
+      conductivity, bootstrap current, and neoclassical_transport.
     coeffs_old: The coefficients calculated at x_old.
     evolving_names: The names of variables within the core profiles that should
       evolve.
@@ -374,7 +374,7 @@ def theta_method_block_loss(
     source_models: Collection of source callables to generate source PDE
       coefficients.
     neoclassical_models: Collection of neoclassical models for calculating
-      conductivity, bootstrap current, etc.
+      conductivity, bootstrap current, and neoclassical_transport.
     coeffs_old: The coefficients calculated at x_old.
     evolving_names: The names of variables within the core profiles that should
       evolve.
@@ -461,7 +461,7 @@ def jaxopt_solver(
     source_models: Collection of source callables to generate source PDE
       coefficients.
     neoclassical_models: Collection of neoclassical models for calculating
-      conductivity, bootstrap current, etc.
+      conductivity, bootstrap current, and neoclassical_transport.
     coeffs_old: The coefficients calculated at x_old.
     evolving_names: The names of variables within the core profiles that should
       evolve.

torax/_src/geometry/geometry.py

@@ -257,6 +257,18 @@ def r_mid_face(self) -> chex.Array:
     """Midplane radius of the plasma on the face grid [m]."""
     return (self.R_out_face - self.R_in_face) / 2
 
+  @property
+  def epsilon(self) -> chex.Array:
+    """Local midplane inverse aspect ratio [dimensionless]."""
+    return (self.R_out - self.R_in) / (self.R_out + self.R_in)
+
+  @property
+  def epsilon_face(self) -> chex.Array:
+    """Local midplane inverse aspect ratio on the face grid [dimensionless]."""
+    return (self.R_out_face - self.R_in_face) / (
+        self.R_out_face + self.R_in_face
+    )
+
   @property
   def drho(self) -> chex.Array:
     """Grid size for rho [m]."""

torax/_src/neoclassical/bootstrap_current/sauter.py

@@ -17,14 +17,15 @@
 
 import chex
 import jax.numpy as jnp
-from torax._src import constants
 from torax._src import jax_utils
 from torax._src import state
 from torax._src.config import runtime_params_slice
 from torax._src.fvm import cell_variable
 from torax._src.geometry import geometry as geometry_lib
+from torax._src.neoclassical import formulas
 from torax._src.neoclassical.bootstrap_current import base
 from torax._src.neoclassical.bootstrap_current import runtime_params
+from torax._src.physics import collisions
 
 
 @chex.dataclass(frozen=True)
@@ -109,48 +110,29 @@ def _calculate_bootstrap_current(
   # Formulas from Sauter PoP 1999. Future work can include Redl PoP 2021
   # corrections.
 
-  # local r/R0 on face grid
-  epsilon = (geo.R_out_face - geo.R_in_face) / (geo.R_out_face + geo.R_in_face)
-  epseff = (
-      0.67 * (1.0 - 1.4 * jnp.abs(geo.delta_face) * geo.delta_face) * epsilon
-  )
-  aa = (1.0 - epsilon) / (1.0 + epsilon)
-  ftrap = 1.0 - jnp.sqrt(aa) * (1.0 - epseff) / (1.0 + 2.0 * jnp.sqrt(epseff))
+  # Effective trapped particle fraction
+  ftrap = formulas.calculate_ftrap(geo)
 
   # Spitzer conductivity
-  lnLame = (
-      31.3 - 0.5 * jnp.log(n_e.face_value()) + jnp.log(T_e.face_value() * 1e3)
-  )
-  lnLami = (
-      30
-      - 3 * jnp.log(Z_i_face)
-      - 0.5 * jnp.log(n_i.face_value())
-      + 1.5 * jnp.log(T_i.face_value() * 1e3)
+  lambda_ei = collisions.calculate_lambda_ei(T_e.face_value(), n_e.face_value())
+  lambda_ii = collisions.calculate_lambda_ii(
+      T_i.face_value(), n_i.face_value(), Z_i_face
   )
-
-  nu_e_star = (
-      6.921e-18
-      * q_face
-      * geo.R_major
-      * n_e.face_value()
-      * Z_eff_face
-      * lnLame
-      / (
-          ((T_e.face_value() * 1e3) ** 2)
-          * (epsilon + constants.CONSTANTS.eps) ** 1.5
-      )
+  nu_e_star = formulas.calculate_nu_e_star(
+      q=q_face,
+      geo=geo,
+      n_e=n_e.face_value(),
+      T_e=T_e.face_value(),
+      Z_eff=Z_eff_face,
+      lambda_ei=lambda_ei,
   )
-  nu_i_star = (
-      4.9e-18
-      * q_face
-      * geo.R_major
-      * n_i.face_value()
-      * Z_eff_face**4
-      * lnLami
-      / (
-          ((T_i.face_value() * 1e3) ** 2)
-          * (epsilon + constants.CONSTANTS.eps) ** 1.5
-      )
+  nu_i_star = formulas.calculate_nu_i_star(
+      q=q_face,
+      geo=geo,
+      n_i=n_i.face_value(),
+      T_i=T_i.face_value(),
+      Z_eff=Z_eff_face,
+      lambda_ii=lambda_ii,
   )
 
   # Calculate terms needed for bootstrap current