Support different sources for trapped particle fraction

docs/configuration.rst

@@ -1183,6 +1183,25 @@ Geometry dicts for all geometry types can contain the following additional keys.
   higher resolution mesh with ``nrho_hires = nrho * hi_res_fac``, used for
   ``j`` to ``psi`` conversions.
 
+``trapped_fraction_source`` (str [default = 'SAUTER'])
+  Selects how the effective trapped particle fraction, used by the
+  neoclassical ``bootstrap_current`` and ``conductivity`` models, is
+  computed. Computed once at geometry construction time. Options are:
+
+  * ``'SAUTER'`` (default)
+    Uses the analytical approximation from
+    `O. Sauter, Fusion Eng. Des. 112, 633 (2016) <https://doi.org/10.1016/j.fusengdes.2016.04.033>`_. Supported
+    by every geometry source.
+
+  * ``'FILE'``
+    Reads the value precomputed by the input equilibrium code directly from
+    the geometry file. Only supported for CHEASE and IMAS
+    geometries.
+
+  * ``'EXACT'``
+    Computes the exact bounce-averaged integral directly from the traced 2D
+    equilibrium. Only supported for EQDSK and IMAS geometries.
+
 
 Geometry dicts for all non-circular geometry types can contain the following
 additional keys.

docs/links.rst

@@ -16,6 +16,7 @@
 .. _flax_link: https://github.com/google/flax
 .. _qualikiz-pythontools_link: https://gitlab.com/qualikiz-group/QuaLiKiz-pythontools
 .. _sauter_link: https://doi.org/10.1063/1.873240
+.. _sauter2016_link: https://doi.org/10.1016/j.fusengdes.2016.04.033
 .. _bosch-hale_link: https://doi.org/10.1088/0029-5515/32/4/I07
 .. _lin-liu_link: https://doi.org/10.1063/1.1610472
 .. _albajar2001_link: https://doi.org/10.1088/0029-5515/41/6/301
@@ -49,6 +50,7 @@
 .. |flax| replace:: `Flax <flax_link_>`_
 .. |qualikiz-pythontools| replace:: `QuaLiKiz Pythontools <qualikiz-pythontools_link_>`_
 .. |sauter99| replace:: `[Sauter PoP 1999] <sauter_link_>`_
+.. |sauter16| replace:: `[Sauter, Fusion Eng. Des. 2016] <sauter2016_link_>`_
 .. |bosch-hale| replace:: `[H.-S. Bosch and G.M. Hale NF 1992] <bosch-hale_link_>`_
 .. |lin-liu| replace:: `[Lin-Liu, Chan, Prater, PoP 2003] <lin-liu_link_>`_
 .. |albajar2001| replace:: `Albajar NF 2001 <albajar2001_link_>`_

docs/physics_models.rst

@@ -321,6 +321,26 @@ used in the current diffusion equation. The Sauter model is a widely-used
 analytical formulation that provides a relatively fast and differentiable
 approximation for these neoclassical quantities.
 
+These formulations, as well as the Redl bootstrap current model, all depend
+on the effective trapped particle fraction, :math:`f_t`. By default this is
+calculated with the analytical approximation of |sauter16| (Eqs. 33-34),
+which only requires the local inverse aspect ratio and triangularity of each
+flux surface.
+
+Where available, TORAX can instead use the full bounce-averaged trapped
+particle fraction integral,
+
+.. math::
+
+  f_t = 1 - \frac{3}{4} \langle B^2 \rangle
+  \int_0^{1/B_\mathrm{max}} \frac{\lambda \, d\lambda}{\langle \sqrt{1 -
+  \lambda B} \rangle},
+
+evaluated directly from the poloidal variation of :math:`B` on each flux
+surface, where :math:`\langle \cdot \rangle` denotes a flux surface average.
+This is most impactful at low aspect ratio (e.g. spherical tokamaks), where
+the analytical approximation is least accurate.
+
 Future work can incorporate more recent neoclassical physics parameterizations,
 and also set neoclassical transport coefficients themselves. This can be of
 importance for ion heat transport in the inner core. When extending TORAX to

torax/_src/geometry/base.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Base class for geometry configuration."""
-from typing import Annotated, Any
+import enum
+from typing import Annotated, Any, ClassVar
 
 import numpy as np
 import pydantic
@@ -21,6 +22,30 @@
 import typing_extensions
 
 
+@enum.unique
+class TrappedFractionSource(enum.StrEnum):
+  """Selects how the effective trapped particle fraction is computed.
+
+  Not every option is supported by every geometry source; see
+  `BaseGeometryConfig._supported_trapped_fraction_sources`.
+
+  Attributes:
+    SAUTER: Uses the analytic approximation from [1]. Supported by all
+      geometry sources.
+    FILE: Reads the value precomputed by the input equilibrium/geometry code
+      directly from the geometry file. Only supported for CHEASE and IMAS
+      sources.
+    EXACT: Computes the full bounce-averaged integral directly from the 2D
+      equilibrium. Only supported for EQDSK and IMAS sources.
+
+  [1] O. Sauter, Fusion Engineering and Design 112 (2016) 633-645, Eqs 33+34.
+  """
+
+  SAUTER = 'SAUTER'
+  FILE = 'FILE'
+  EXACT = 'EXACT'
+
+
 class BaseGeometryConfig(torax_pydantic.BaseModelFrozen):
   """Base class for all geometry configuration classes.
 
@@ -33,13 +58,25 @@ class BaseGeometryConfig(torax_pydantic.BaseModelFrozen):
     hires_factor: Only used when the initial condition ``psi`` is from plasma
       current. Sets up a higher resolution mesh with ``nrho_hires = nrho *
       hi_res_fac``, used for ``j`` to ``psi`` conversions.
+    trapped_fraction_source: Selects how the effective trapped particle
+      fraction is computed. See `TrappedFractionSource`.
+    _supported_trapped_fraction_sources: Overridden per subclass to restrict
+      which `TrappedFractionSource` options that geometry source actually
+      supports.
   """
 
+  _supported_trapped_fraction_sources: ClassVar[
+      frozenset[TrappedFractionSource]
+  ] = frozenset(TrappedFractionSource)
+
   n_rho: Annotated[int | None, torax_pydantic.TIME_INVARIANT] = None
   face_centers: Annotated[
       torax_pydantic.NumpyArray1DSorted | None, torax_pydantic.TIME_INVARIANT
   ] = None
   hires_factor: pydantic.PositiveInt = 4
+  trapped_fraction_source: Annotated[
+      TrappedFractionSource, torax_pydantic.TIME_INVARIANT
+  ] = TrappedFractionSource.SAUTER
 
   @pydantic.model_validator(mode='before')
   @classmethod
@@ -76,6 +113,23 @@ def _validate_n_rho_or_face_centers(self) -> typing_extensions.Self:
 
     return self
 
+  @pydantic.model_validator(mode='after')
+  def _validate_trapped_fraction_source(self) -> typing_extensions.Self:
+    """Validates that trapped_fraction_source is supported by this geometry."""
+    if (
+        self.trapped_fraction_source
+        not in self._supported_trapped_fraction_sources
+    ):
+      allowed = ', '.join(
+          sorted(s.value for s in self._supported_trapped_fraction_sources)
+      )
+      raise ValueError(
+          f'trapped_fraction_source={self.trapped_fraction_source.value} is'
+          f' not supported for {type(self).__name__}. Supported options:'
+          f' {allowed}.'
+      )
+    return self
+
   def get_face_centers(self) -> np.ndarray:
     """Returns face_centers, computing from n_rho if needed."""
     if self.face_centers is not None:

torax/_src/geometry/chease.py

@@ -12,14 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Functions for loading and representing a CHEASE geometry."""
-from typing import Annotated, Literal
+from typing import Annotated, ClassVar, Literal
 import numpy as np
 import pydantic
 from torax._src import constants
 from torax._src.geometry import base
 from torax._src.geometry import geometry
 from torax._src.geometry import geometry_loader
 from torax._src.geometry import standard_geometry
+from torax._src.neoclassical.formulas import formulas
 from torax._src.torax_pydantic import torax_pydantic
 import typing_extensions
 
@@ -39,6 +40,13 @@ class CheaseConfig(base.BaseGeometryConfig):
     B_0: Vacuum toroidal magnetic field at `R_major` [T].
   """
 
+  _supported_trapped_fraction_sources: ClassVar[
+      frozenset[base.TrappedFractionSource]
+  ] = frozenset({
+      base.TrappedFractionSource.SAUTER,
+      base.TrappedFractionSource.FILE,
+  })
+
   geometry_type: Annotated[Literal['chease'], torax_pydantic.TIME_INVARIANT] = (
       'chease'
   )
@@ -67,6 +75,7 @@ def build_geometry(self) -> standard_geometry.StandardGeometry:
         a_minor=self.a_minor,
         B_0=self.B_0,
         hires_factor=self.hires_factor,
+        trapped_fraction_source=self.trapped_fraction_source,
     )
 
     return standard_geometry.build_standard_geometry(intermediates)
@@ -84,6 +93,7 @@ def _construct_intermediates_from_chease(
     a_minor: float,
     B_0: float,
     hires_factor: int,
+    trapped_fraction_source: base.TrappedFractionSource,
 ) -> standard_geometry.StandardGeometryIntermediates:
   """Constructs a StandardGeometryIntermediates from a CHEASE file.
 
@@ -103,6 +113,8 @@ def _construct_intermediates_from_chease(
     B_0: Vacuum toroidal magnetic field at `R_major` [T].
     hires_factor: Grid refinement factor for poloidal flux <--> plasma current
       calculations.
+    trapped_fraction_source: Selects how the effective trapped particle
+      fraction is computed; see `base.TrappedFractionSource`.
 
   Returns:
     A StandardGeometry instance based on the input file. This can then be
@@ -149,6 +161,19 @@ def _construct_intermediates_from_chease(
   )
   flux_surf_avg_B2 = chease_data['<B**2>'] * B_0**2
   flux_surf_avg_1_over_B2 = chease_data['<1/B**2>'] / B_0**2
+  match trapped_fraction_source:
+    case base.TrappedFractionSource.FILE:
+      trapped_fraction = chease_data['FTRAP']
+    case base.TrappedFractionSource.SAUTER:
+      trapped_fraction = formulas.calculate_sauter_trapped_fraction(
+          epsilon=(R_out_chease - R_in_chease) / (R_out_chease + R_in_chease),
+          delta=0.5
+          * (chease_data['delta_upper'] + chease_data['delta_bottom']),
+      )
+    case _:
+      raise ValueError(
+          f'Unknown trapped_fraction_source: {trapped_fraction_source}'
+      )
 
   rhon = np.sqrt(Phi / Phi[-1])
   vpr = 4 * np.pi * Phi[-1] * rhon / (F * flux_surf_avg_1_over_R2)
@@ -173,6 +198,7 @@ def _construct_intermediates_from_chease(
       flux_surf_avg_grad_psi2=flux_surf_avg_grad_psi2,
       flux_surf_avg_B2=flux_surf_avg_B2,
       flux_surf_avg_1_over_B2=flux_surf_avg_1_over_B2,
+      trapped_fraction=trapped_fraction,
       delta_upper_face=chease_data['delta_upper'],
       delta_lower_face=chease_data['delta_bottom'],
       elongation=chease_data['elongation'],

torax/_src/geometry/circular_geometry.py

@@ -18,6 +18,7 @@
 import pydantic
 from torax._src.geometry import base
 from torax._src.geometry import geometry
+from torax._src.neoclassical.formulas import formulas
 from torax._src.torax_pydantic import torax_pydantic
 import typing_extensions
 
@@ -217,6 +218,9 @@ def _build_circular_geometry(
   # Analytical expressions for  <1/B^2> (gm4) and <B^2> (gm5)
   epsilon = (R_out - R_in) / (R_out + R_in)
   epsilon_face = (R_out_face - R_in_face) / (R_out_face + R_in_face)
+  trapped_fraction_face = formulas.calculate_sauter_trapped_fraction(
+      epsilon=epsilon_face, delta=delta_face
+  )
   gm4 = B_0**-2 * (1.0 + 1.5 * epsilon**2)
   gm4_face = B_0**-2 * (1.0 + 1.5 * epsilon_face**2)
   gm5 = B_0**2 / np.sqrt(1.0 - epsilon**2)
@@ -244,6 +248,7 @@ def _build_circular_geometry(
       spr=spr,
       spr_face=spr_face,
       delta_face=delta_face,
+      trapped_fraction_face=trapped_fraction_face,
       g0=g0,
       g0_face=g0_face,
       g1=g1,