Refactor neoclassical models.

docs/output.rst

@@ -174,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`].

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,12 +17,12 @@
 
 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
@@ -110,68 +110,54 @@ 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
+  f_trap = formulas.calculate_f_trap(geo)
 
   # Spitzer conductivity
-  lambda_ei = collisions.calculate_lambda_ei(
+  log_lambda_ei = collisions.calculate_log_lambda_ei(
       T_e.face_value(), n_e.face_value()
   )
-  lambda_ii = collisions.calculate_lambda_ii(
+  log_lambda_ii = collisions.calculate_log_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
-      * lambda_ei
-      / (
-          ((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,
+      log_lambda_ei=log_lambda_ei,
   )
-  nu_i_star = (
-      4.9e-18
-      * q_face
-      * geo.R_major
-      * n_i.face_value()
-      * Z_eff_face**4
-      * lambda_ii
-      / (
-          ((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,
+      log_lambda_ii=log_lambda_ii,
   )
 
   # Calculate terms needed for bootstrap current
   denom = (
       1.0
-      + (1 - 0.1 * ftrap) * jnp.sqrt(nu_e_star)
-      + 0.5 * (1.0 - ftrap) * nu_e_star / Z_eff_face
+      + (1 - 0.1 * f_trap) * jnp.sqrt(nu_e_star)
+      + 0.5 * (1.0 - f_trap) * nu_e_star / Z_eff_face
   )
-  ft31 = ftrap / denom
-  ft32ee = ftrap / (
+  ft31 = f_trap / denom
+  ft32ee = f_trap / (
       1
-      + 0.26 * (1 - ftrap) * jnp.sqrt(nu_e_star)
-      + 0.18 * (1 - 0.37 * ftrap) * nu_e_star / jnp.sqrt(Z_eff_face)
+      + 0.26 * (1 - f_trap) * jnp.sqrt(nu_e_star)
+      + 0.18 * (1 - 0.37 * f_trap) * nu_e_star / jnp.sqrt(Z_eff_face)
   )
-  ft32ei = ftrap / (
+  ft32ei = f_trap / (
       1
-      + (1 + 0.6 * ftrap) * jnp.sqrt(nu_e_star)
-      + 0.85 * (1 - 0.37 * ftrap) * nu_e_star * (1 + Z_eff_face)
+      + (1 + 0.6 * f_trap) * jnp.sqrt(nu_e_star)
+      + 0.85 * (1 - 0.37 * f_trap) * nu_e_star * (1 + Z_eff_face)
   )
-  ft34 = ftrap / (
+  ft34 = f_trap / (
       1.0
-      + (1 - 0.1 * ftrap) * jnp.sqrt(nu_e_star)
-      + 0.5 * (1.0 - 0.5 * ftrap) * nu_e_star / Z_eff_face
+      + (1 - 0.1 * f_trap) * jnp.sqrt(nu_e_star)
+      + 0.5 * (1.0 - 0.5 * f_trap) * nu_e_star / Z_eff_face
   )
 
   F32ee = (
@@ -209,12 +195,12 @@ def _calculate_bootstrap_current(
       + 0.2 / (Z_eff_face + 1) * ft34**4
   )
 
-  alpha0 = -1.17 * (1 - ftrap) / (1 - 0.22 * ftrap - 0.19 * ftrap**2)
+  alpha0 = -1.17 * (1 - f_trap) / (1 - 0.22 * f_trap - 0.19 * f_trap**2)
   alpha = (
-      (alpha0 + 0.25 * (1 - ftrap**2) * jnp.sqrt(nu_i_star))
+      (alpha0 + 0.25 * (1 - f_trap**2) * jnp.sqrt(nu_i_star))
       / (1 + 0.5 * jnp.sqrt(nu_i_star))
-      + 0.315 * nu_i_star**2 * ftrap**6
-  ) / (1 + 0.15 * nu_i_star**2 * ftrap**6)
+      + 0.315 * nu_i_star**2 * f_trap**6
+  ) / (1 + 0.15 * nu_i_star**2 * f_trap**6)
 
   # calculate bootstrap current
   prefactor = -geo.F_face * bootstrap_multiplier * 2 * jnp.pi / geo.B_0

torax/_src/neoclassical/conductivity/sauter.py

@@ -17,11 +17,11 @@
 
 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.fvm import cell_variable
 from torax._src.geometry import geometry as geometry_lib
+from torax._src.neoclassical import formulas
 from torax._src.neoclassical.conductivity import base
 from torax._src.neoclassical.conductivity import runtime_params
 from torax._src.physics import collisions
@@ -85,43 +85,39 @@ def _calculate_conductivity(
   """Calculates sigma and sigma_face using the Sauter model."""
   # pylint: disable=invalid-name
 
-  # Formulas from Sauter PoP 1999. Future work can include Redl PoP 2021
-  # corrections.
+  # Formulas from Sauter PoP 1999.
 
-  # # 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
+  f_trap = formulas.calculate_f_trap(geo)
 
   # Spitzer conductivity
   NZ = 0.58 + 0.74 / (0.76 + Z_eff_face)
-  lambda_ei = collisions.calculate_lambda_ei(T_e.face_value(), n_e.face_value())
+  log_lambda_ei = collisions.calculate_log_lambda_ei(
+      T_e.face_value(), n_e.face_value()
+  )
 
   sigsptz = (
-      1.9012e04 * (T_e.face_value() * 1e3) ** 1.5 / Z_eff_face / NZ / lambda_ei
+      1.9012e04
+      * (T_e.face_value() * 1e3) ** 1.5
+      / Z_eff_face
+      / NZ
+      / log_lambda_ei
   )
 
-  nu_e_star = (
-      6.921e-18
-      * q_face
-      * geo.R_major
-      * n_e.face_value()
-      * Z_eff_face
-      * lambda_ei
-      / (
-          ((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,
+      log_lambda_ei=log_lambda_ei,
   )
 
   # Neoclassical correction to spitzer conductivity
-  ft33 = ftrap / (
+  ft33 = f_trap / (
       1.0
-      + (0.55 - 0.1 * ftrap) * jnp.sqrt(nu_e_star)
-      + 0.45 * (1.0 - ftrap) * nu_e_star / (Z_eff_face**1.5)
+      + (0.55 - 0.1 * f_trap) * jnp.sqrt(nu_e_star)
+      + 0.45 * (1.0 - f_trap) * nu_e_star / (Z_eff_face**1.5)
   )
   signeo_face = 1.0 - ft33 * (
       1.0

torax/_src/neoclassical/formulas.py

@@ -0,0 +1,121 @@
+# Copyright 2024 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.
+"""Common formulas used in neoclassical models."""
+
+import chex
+import jax.numpy as jnp
+from torax._src import constants
+from torax._src.geometry import geometry as geometry_lib
+
+# pylint: disable=invalid-name
+
+
+def calculate_f_trap(geo: geometry_lib.Geometry) -> chex.Array:
+  """Calculates the effective trapped particle fraction.
+
+  From O. Sauter, Fusion Engineering and Design 112 (2016) 633-645. Eqs 33+34.
+
+  Args:
+    geo: The magnetic geometry.
+  Returns:
+    The effective trapped particle fraction.
+  """
+
+  epsilon_effective = (
+      0.67
+      * (1.0 - 1.4 * jnp.abs(geo.delta_face) * geo.delta_face)
+      * geo.epsilon_face
+  )
+  aa = (1.0 - geo.epsilon_face) / (1.0 + geo.epsilon_face)
+  return 1.0 - jnp.sqrt(aa) * (1.0 - epsilon_effective) / (
+      1.0 + 2.0 * jnp.sqrt(epsilon_effective)
+  )
+
+
+# TODO(b/428166775): currently we have two very similar implementations for
+# nu_e_star. We should refactor this to have a single one in physics/collisions
+def calculate_nu_e_star(
+    q: chex.Array,
+    geo: geometry_lib.Geometry,
+    n_e: chex.Array,
+    T_e: chex.Array,
+    Z_eff: chex.Array,
+    log_lambda_ei: chex.Array,
+) -> chex.Array:
+  """Calculates the electron collisionality, nu_e_star.
+
+  This is the electron collisionality, defined as the ratio of the electron
+  collision frequency to the bounce frequency. From Sauter PoP 1999 Eq. (18b).
+
+  Args:
+    q: Safety factor.
+    geo: The geometry of the torus.
+    n_e: Electron density [m^-3].
+    T_e: Electron temperature [keV]. Converted to eV in the formula.
+    Z_eff: Effective charge.
+    log_lambda_ei: Electron-ion Coulomb logarithm.
+
+  Returns:
+    The electron collisionality.
+  """
+  return (
+      6.921e-18
+      * q
+      * geo.R_major
+      * n_e
+      * Z_eff
+      * log_lambda_ei
+      / (
+          ((T_e * 1e3) ** 2)
+          * (geo.epsilon_face + constants.CONSTANTS.eps) ** 1.5
+      )
+  )
+
+
+def calculate_nu_i_star(
+    q: chex.Array,
+    geo: geometry_lib.Geometry,
+    n_i: chex.Array,
+    T_i: chex.Array,
+    Z_eff: chex.Array,
+    log_lambda_ii: chex.Array,
+) -> chex.Array:
+  """Calculates the ion collisionality, nu_i_star.
+
+  This is the ion collisionality, defined as the ratio of the ion
+  collision frequency to the bounce frequency. From Sauter PoP 1999 Eq. (18c).
+
+  Args:
+    q: Safety factor.
+    geo: The geometry of the torus.
+    n_i: Ion density.
+    T_i: Ion temperature.
+    Z_eff: Effective charge.
+    log_lambda_ii: Ion-ion Coulomb logarithm.
+
+  Returns:
+    The ion collisionality.
+  """
+  return (
+      4.9e-18
+      * q
+      * geo.R_major
+      * n_i
+      * Z_eff**4
+      * log_lambda_ii
+      / (
+          ((T_i * 1e3) ** 2)
+          * (geo.epsilon_face + constants.CONSTANTS.eps) ** 1.5
+      )
+  )