diff --git a/docs/configuration.rst b/docs/configuration.rst index 13a993094..ff072433f 100644 --- a/docs/configuration.rst +++ b/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) `_. 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. diff --git a/docs/links.rst b/docs/links.rst index e8dd94782..97b4db82c 100644 --- a/docs/links.rst +++ b/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 `_ .. |qualikiz-pythontools| replace:: `QuaLiKiz Pythontools `_ .. |sauter99| replace:: `[Sauter PoP 1999] `_ +.. |sauter16| replace:: `[Sauter, Fusion Eng. Des. 2016] `_ .. |bosch-hale| replace:: `[H.-S. Bosch and G.M. Hale NF 1992] `_ .. |lin-liu| replace:: `[Lin-Liu, Chan, Prater, PoP 2003] `_ .. |albajar2001| replace:: `Albajar NF 2001 `_ diff --git a/docs/physics_models.rst b/docs/physics_models.rst index 858d2367c..82cf16d2b 100644 --- a/docs/physics_models.rst +++ b/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 diff --git a/torax/_src/geometry/base.py b/torax/_src/geometry/base.py index c05b1c9a4..32875bfdd 100644 --- a/torax/_src/geometry/base.py +++ b/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: diff --git a/torax/_src/geometry/chease.py b/torax/_src/geometry/chease.py index 2fd24e529..c2eb979c3 100644 --- a/torax/_src/geometry/chease.py +++ b/torax/_src/geometry/chease.py @@ -12,7 +12,7 @@ # 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 @@ -20,6 +20,7 @@ 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_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'], diff --git a/torax/_src/geometry/circular_geometry.py b/torax/_src/geometry/circular_geometry.py index a1949f1c3..735e3a558 100644 --- a/torax/_src/geometry/circular_geometry.py +++ b/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 (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, diff --git a/torax/_src/geometry/eqdsk.py b/torax/_src/geometry/eqdsk.py index 23d4dcb41..8c24de176 100644 --- a/torax/_src/geometry/eqdsk.py +++ b/torax/_src/geometry/eqdsk.py @@ -15,7 +15,7 @@ import json import logging -from typing import Annotated, Any, Literal +from typing import Annotated, Any, ClassVar, Literal import contourpy import eqdsk @@ -30,6 +30,7 @@ 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 @@ -65,6 +66,13 @@ class EQDSKConfig(base.BaseGeometryConfig): grid. Needed to avoid divergent integrations in diverted geometries. """ + _supported_trapped_fraction_sources: ClassVar[ + frozenset[base.TrappedFractionSource] + ] = frozenset({ + base.TrappedFractionSource.SAUTER, + base.TrappedFractionSource.EXACT, + }) + cocos: torax_pydantic.COCOSInt = ... # pyrefly: ignore[bad-assignment] geometry_file: str | None = None eqdsk_object: eqdsk.EQDSKInterface | None = None @@ -127,6 +135,7 @@ def build_geometry(self) -> standard_geometry.StandardGeometry: cocos=self.cocos, n_surfaces=self.n_surfaces, last_surface_factor=self.last_surface_factor, + trapped_fraction_source=self.trapped_fraction_source, ) return standard_geometry.build_standard_geometry(intermediates) @@ -141,6 +150,7 @@ def _construct_intermediates_from_eqdsk( n_surfaces: int, last_surface_factor: float, cocos: int, + trapped_fraction_source: base.TrappedFractionSource, ) -> standard_geometry.StandardGeometryIntermediates: """Constructs a StandardGeometryIntermediates from EQDSK. @@ -167,6 +177,8 @@ def _construct_intermediates_from_eqdsk( grid. Needed to avoid divergent integrations in diverted geometries. cocos: COCOS convention of the EQDSK file, specified as an integer between 1-8 or 11-18 inclusive. + trapped_fraction_source: Selects how the effective trapped particle + fraction is computed; see `base.TrappedFractionSource`. Returns: A StandardGeometryIntermediates instance based on the input file or object. @@ -321,6 +333,12 @@ def calculate_area(x, z): flux_surf_avg_grad_psi2 = np.empty(len(surfaces) + 1) # <|grad(psi)|**2> flux_surf_avg_B2 = np.empty(len(surfaces) + 1) # flux_surf_avg_1_over_B2 = np.empty(len(surfaces) + 1) # <1/B**2> + compute_exact_trapped_fraction = ( + trapped_fraction_source == base.TrappedFractionSource.EXACT + ) + exact_trapped_fraction = ( + np.empty(len(surfaces) + 1) if compute_exact_trapped_fraction else None + ) # Effective trapped fraction int_dl_over_Bp = np.empty(len(surfaces) + 1) # int(Rdl / | grad(psi) |) Ip = np.empty(len(surfaces) + 1) # Toroidal plasma current delta_upper_face = np.empty(len(surfaces) + 1) # Upper face delta @@ -389,6 +407,16 @@ def calculate_area(x, z): / surface_int_dl_over_bpol ) + if compute_exact_trapped_fraction: + surface_B = np.sqrt(surface_B2) + surface_trapped_fraction = ( + formulas.calculate_bounce_averaged_trapped_fraction( + B=surface_B, + dl_over_Bp=surface_dl / surface_Bpol, + flux_surf_avg_B2=surface_FSA_B2, + ) + ) + # Volumes and areas area = calculate_area(x_surface, z_surface) volume = area * 2 * np.pi * R_major @@ -425,6 +453,8 @@ def calculate_area(x, z): flux_surf_avg_grad_psi2_over_R2[n + 1] = surface_FSA_abs_grad_psi2_over_R2 flux_surf_avg_B2[n + 1] = surface_FSA_B2 flux_surf_avg_1_over_B2[n + 1] = surface_FSA_1_over_B2 + if compute_exact_trapped_fraction: + exact_trapped_fraction[n + 1] = surface_trapped_fraction Ip[n + 1] = surface_int_bpol_dl / constants.CONSTANTS.mu_0 delta_upper_face[n + 1] = surface_delta_upper_face delta_lower_face[n + 1] = surface_delta_lower_face @@ -445,6 +475,9 @@ def calculate_area(x, z): flux_surf_avg_grad_psi2_over_R2[0] = 0 flux_surf_avg_B2[0] = Btor_axis**2 flux_surf_avg_1_over_B2[0] = 1 / Btor_axis**2 + if compute_exact_trapped_fraction: + # No trapped particles on the magnetic axis, where B is uniform. + exact_trapped_fraction[0] = 0.0 Ip[0] = 0 delta_upper_face[0] = delta_upper_face[1] delta_lower_face[0] = delta_lower_face[1] @@ -459,6 +492,31 @@ def calculate_area(x, z): rhon = np.sqrt(Phi / Phi[-1]) vpr = 4 * np.pi * Phi[-1] * rhon / (F * flux_surf_avg_1_over_R2) + sauter_trapped_fraction = formulas.calculate_sauter_trapped_fraction( + epsilon=(R_outboard - R_inboard) / (R_outboard + R_inboard), + delta=0.5 * (delta_upper_face + delta_lower_face), + ) + + match trapped_fraction_source: + case base.TrappedFractionSource.EXACT: + # Fill any unreliable values (NaN, or outside the physically valid + # [0, 1] range, e.g. surfaces too close to the magnetic axis for the + # integral to resolve well) with the Sauter approximation. + exact_is_unreliable = ( + np.isnan(exact_trapped_fraction) + | (exact_trapped_fraction < 0.0) + | (exact_trapped_fraction > 1.0) + ) + trapped_fraction = np.where( + exact_is_unreliable, sauter_trapped_fraction, exact_trapped_fraction + ) + case base.TrappedFractionSource.SAUTER: + trapped_fraction = sauter_trapped_fraction + case _: + raise ValueError( + f'Unknown trapped_fraction_source: {trapped_fraction_source}' + ) + # ------------------------------------ # # ---- 6. Sense-check the results ---- # # ------------------------------------ # @@ -496,6 +554,7 @@ def calculate_area(x, z): flux_surf_avg_grad_psi2_over_R2=flux_surf_avg_grad_psi2_over_R2, 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=delta_upper_face, delta_lower_face=delta_lower_face, elongation=elongation, diff --git a/torax/_src/geometry/fbt.py b/torax/_src/geometry/fbt.py index 5b838b0ad..665115b45 100644 --- a/torax/_src/geometry/fbt.py +++ b/torax/_src/geometry/fbt.py @@ -18,6 +18,7 @@ import logging from typing import Annotated from typing import Any +from typing import ClassVar from typing import Literal, TypeAlias import jax @@ -30,6 +31,7 @@ from torax._src.geometry import geometry_loader from torax._src.geometry import geometry_provider 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 @@ -81,6 +83,10 @@ class FBTConfig(base.BaseGeometryConfig): edge quantities when diverted. """ + _supported_trapped_fraction_sources: ClassVar[ + frozenset[base.TrappedFractionSource] + ] = frozenset({base.TrappedFractionSource.SAUTER}) + geometry_type: Annotated[Literal['fbt'], torax_pydantic.TIME_INVARIANT] = ( 'fbt' ) @@ -443,6 +449,10 @@ def _from_fbt( num=B_0**2, denom=np.sqrt(1.0 - LY['epsilon'] ** 2), eps=1e-7 ) flux_surf_avg_1_over_B2 = B_0**-2 * (1.0 + 1.5 * LY['epsilon'] ** 2) + trapped_fraction = formulas.calculate_sauter_trapped_fraction( + epsilon=LY['epsilon'], + delta=0.5 * (LY['deltau'] + LY['deltal']), + ) # Edge/Divertor geometry # These parameters are optional as older FBT files may not contain them. @@ -497,6 +507,7 @@ def _from_fbt( flux_surf_avg_grad_psi2=LY['Q4Q'], # pyrefly: ignore[bad-argument-type] flux_surf_avg_B2=flux_surf_avg_B2, # pyrefly: ignore[bad-argument-type] flux_surf_avg_1_over_B2=flux_surf_avg_1_over_B2, + trapped_fraction=trapped_fraction, # pyrefly: ignore[bad-argument-type] delta_upper_face=LY['deltau'], # pyrefly: ignore[bad-argument-type] delta_lower_face=LY['deltal'], # pyrefly: ignore[bad-argument-type] elongation=LY['kappa'], # pyrefly: ignore[bad-argument-type] diff --git a/torax/_src/geometry/geometry.py b/torax/_src/geometry/geometry.py index 8307e4360..77ea8ed5d 100644 --- a/torax/_src/geometry/geometry.py +++ b/torax/_src/geometry/geometry.py @@ -172,6 +172,8 @@ class Geometry: location of the upper extent of the flux surface. Lower triangularity is defined as (R_major_local - R_lower) / a_minor_local, where R_lower is the radial location of the lower extent of the flux surface. + trapped_fraction_face: Effective trapped particle fraction on the face + grid [dimensionless], computed at geometry construction time. elongation: Plasma elongation profile on cell grid [dimensionless]. Elongation is defined as (Z_upper - Z_lower) / (2.0 * a_minor_local), where Z_upper and Z_lower are the Z coordinates of the upper and lower @@ -201,6 +203,7 @@ class Geometry: spr: array_typing.Array spr_face: array_typing.Array delta_face: array_typing.Array + trapped_fraction_face: array_typing.Array elongation: array_typing.Array elongation_face: array_typing.Array g0: array_typing.Array diff --git a/torax/_src/geometry/geometry_provider.py b/torax/_src/geometry/geometry_provider.py index 94b583d03..7b5224724 100644 --- a/torax/_src/geometry/geometry_provider.py +++ b/torax/_src/geometry/geometry_provider.py @@ -130,6 +130,7 @@ class TimeDependentGeometryProvider: spr: interpolated_param.InterpolatedVarSingleAxis spr_face: interpolated_param.InterpolatedVarSingleAxis delta_face: interpolated_param.InterpolatedVarSingleAxis + trapped_fraction_face: interpolated_param.InterpolatedVarSingleAxis elongation: interpolated_param.InterpolatedVarSingleAxis elongation_face: interpolated_param.InterpolatedVarSingleAxis g0: interpolated_param.InterpolatedVarSingleAxis diff --git a/torax/_src/geometry/imas.py b/torax/_src/geometry/imas.py index d0296f260..b748613bc 100644 --- a/torax/_src/geometry/imas.py +++ b/torax/_src/geometry/imas.py @@ -149,6 +149,7 @@ def build_geometry(self) -> standard_geometry.StandardGeometry: explicit_convert=self.explicit_convert, slice_index=self.slice_index, slice_time=self.slice_time, + trapped_fraction_source=self.trapped_fraction_source, ) intermediates = standard_geometry.StandardGeometryIntermediates( geometry_type=geometry.GeometryType.IMAS, **inputs @@ -179,6 +180,7 @@ def build_geometry_provider( face_centers=self.get_face_centers(), hires_factor=self.hires_factor, explicit_convert=self.explicit_convert, + trapped_fraction_source=self.trapped_fraction_source, ) geometries = {} for t, inputs in all_inputs.items(): diff --git a/torax/_src/geometry/standard_geometry.py b/torax/_src/geometry/standard_geometry.py index 9d1946300..ca85e40ce 100644 --- a/torax/_src/geometry/standard_geometry.py +++ b/torax/_src/geometry/standard_geometry.py @@ -197,6 +197,8 @@ class StandardGeometryIntermediates: [:math:`\mathrm{T}^2`]. flux_surf_avg_1_over_B2: Flux surface average of :math:`1/B^2` [:math:`\mathrm{T}^{-2}`]. + trapped_fraction: Effective trapped particle fraction [dimensionless]; + see `torax._src.geometry.base.TrappedFractionSource`. delta_upper_face: Upper triangularity [dimensionless]. See `Geometry` docstring for definition. delta_lower_face: Lower triangularity [dimensionless]. See `Geometry` @@ -257,6 +259,7 @@ class StandardGeometryIntermediates: R_OMP: array_typing.FloatScalar | None R_target: array_typing.FloatScalar | None B_pol_OMP: array_typing.FloatScalar | None + trapped_fraction: array_typing.Array def __post_init__(self): """Enforces sign conventions, extrapolates edge, and smooths near-axis. @@ -492,6 +495,10 @@ def build_standard_geometry( # average triangularity delta_face = 0.5 * (delta_upper_face + delta_lower_face) + trapped_fraction_face = rhon_interpolation_func( + rho_face_norm, intermediates.trapped_fraction + ) + # elongation elongation = rhon_interpolation_func(rho_norm, intermediates.elongation) elongation_face = rhon_interpolation_func( @@ -570,6 +577,7 @@ def build_standard_geometry( spr=spr_cell, spr_face=spr_face, delta_face=delta_face, + trapped_fraction_face=trapped_fraction_face, g0=g0, g0_face=g0_face, g1=g1, diff --git a/torax/_src/geometry/tests/chease_test.py b/torax/_src/geometry/tests/chease_test.py index b6604c0b0..8b4227e61 100644 --- a/torax/_src/geometry/tests/chease_test.py +++ b/torax/_src/geometry/tests/chease_test.py @@ -13,6 +13,8 @@ # limitations under the License. from absl.testing import absltest from absl.testing import parameterized +import numpy as np +from torax._src.geometry import base from torax._src.geometry import chease # pylint: disable=invalid-name @@ -28,6 +30,44 @@ def test_access_z_magnetic_axis_raises_error_for_chease_geometry(self): with self.assertRaisesRegex(ValueError, 'does not have a z magnetic axis'): geo.z_magnetic_axis() + def test_trapped_fraction_is_physically_sensible(self): + """Tests that the exact trapped particle fraction (FTRAP) is sensible.""" + geo = chease.CheaseConfig( + trapped_fraction_source=base.TrappedFractionSource.FILE, + ).build_geometry() + trapped_fraction = geo.trapped_fraction_face + self.assertTrue(np.all(trapped_fraction >= 0.0)) + self.assertTrue(np.all(trapped_fraction <= 1.0)) + self.assertGreater( + np.mean(np.diff(trapped_fraction) >= -1e-6), + 0.8, + ) + + def test_trapped_fraction_source_exact_not_supported(self): + """Tests that EXACT is rejected for CHEASE (no full 2D equilibrium).""" + with self.assertRaisesRegex(ValueError, 'not supported for CheaseConfig'): + chease.CheaseConfig( + trapped_fraction_source=base.TrappedFractionSource.EXACT, + ) + + def test_trapped_fraction_geometry_consistent_with_sauter(self): + """Tests that the exact and Sauter trapped fractions roughly agree.""" + geo_sauter = chease.CheaseConfig( + trapped_fraction_source=base.TrappedFractionSource.SAUTER, + ).build_geometry() + geo_geometry = chease.CheaseConfig( + trapped_fraction_source=base.TrappedFractionSource.FILE, + ).build_geometry() + # Moderately coarse tolerance: Sauter is only an analytic approximation, + # so it need not match the exact integral closely, but a large deviation + # would indicate a bug rather than the expected model discrepancy. + np.testing.assert_allclose( + geo_geometry.trapped_fraction_face, + geo_sauter.trapped_fraction_face, + atol=0.05, + rtol=0.15, + ) + if __name__ == '__main__': absltest.main() diff --git a/torax/_src/geometry/tests/eqdsk_test.py b/torax/_src/geometry/tests/eqdsk_test.py index e12d302f1..a2b9f8e3b 100644 --- a/torax/_src/geometry/tests/eqdsk_test.py +++ b/torax/_src/geometry/tests/eqdsk_test.py @@ -18,6 +18,7 @@ import eqdsk as eqdsk_lib import numpy as np from torax._src import array_typing +from torax._src.geometry import base from torax._src.geometry import eqdsk from torax._src.geometry import geometry_loader @@ -84,6 +85,59 @@ def test_build_geometry_from_eqdsk_object(self): else: self.assertEqual(val1, val2, msg=f'Field "{name}" mismatch.') + def test_trapped_fraction_is_physically_sensible(self): + """Tests that the exact trapped particle fraction is well-behaved.""" + geo = eqdsk.EQDSKConfig( + geometry_file='iterhybrid_cocos11.eqdsk', + cocos=11, + trapped_fraction_source=base.TrappedFractionSource.EXACT, + ).build_geometry() + trapped_fraction = geo.trapped_fraction_face + self.assertIsNotNone(trapped_fraction) + # No trapped particles on the magnetic axis, where B is uniform. + self.assertAlmostEqual(float(trapped_fraction[0]), 0.0) + # The trapped particle fraction is a fraction, so must lie in [0, 1]. + self.assertTrue(np.all(trapped_fraction >= 0.0)) + self.assertTrue(np.all(trapped_fraction <= 1.0)) + # Trapped fraction increases with normalized radius over most of the + # profile (small deviations from strict monotonicity are possible near + # the edge for diverted geometries, due to the X-point). + self.assertGreater( + np.mean(np.diff(trapped_fraction) >= -1e-6), + 0.8, + ) + + def test_trapped_fraction_source_file_not_supported(self): + """Tests that FILE is rejected for EQDSK (no precomputed value).""" + with self.assertRaisesRegex(ValueError, 'not supported for EQDSKConfig'): + eqdsk.EQDSKConfig( + geometry_file='iterhybrid_cocos11.eqdsk', + cocos=11, + trapped_fraction_source=base.TrappedFractionSource.FILE, + ) + + def test_trapped_fraction_geometry_consistent_with_sauter(self): + """Tests that the exact and Sauter trapped fractions roughly agree.""" + geo_sauter = eqdsk.EQDSKConfig( + geometry_file='iterhybrid_cocos11.eqdsk', + cocos=11, + trapped_fraction_source=base.TrappedFractionSource.SAUTER, + ).build_geometry() + geo_geometry = eqdsk.EQDSKConfig( + geometry_file='iterhybrid_cocos11.eqdsk', + cocos=11, + trapped_fraction_source=base.TrappedFractionSource.EXACT, + ).build_geometry() + # Moderately coarse tolerance: Sauter is only an analytic approximation, + # so it need not match the exact integral closely, but a large deviation + # would indicate a bug rather than the expected model discrepancy. + np.testing.assert_allclose( + geo_geometry.trapped_fraction_face, + geo_sauter.trapped_fraction_face, + atol=0.05, + rtol=0.15, + ) + def test_eqdsk_serialization_round_trip(self): """Test that EQDSKConfig with eqdsk_object can be serialized and deserialized.""" geo_dir = geometry_loader.get_geometry_dir() diff --git a/torax/_src/geometry/tests/fbt_test.py b/torax/_src/geometry/tests/fbt_test.py index b7cd08f90..6a4d8a61d 100644 --- a/torax/_src/geometry/tests/fbt_test.py +++ b/torax/_src/geometry/tests/fbt_test.py @@ -17,6 +17,7 @@ from absl.testing import absltest from absl.testing import parameterized import numpy as np +from torax._src.geometry import base from torax._src.geometry import fbt from torax._src.geometry import geometry from torax._src.geometry import geometry_loader @@ -32,6 +33,17 @@ class FBTGeometryTest(parameterized.TestCase): + @parameterized.parameters([ + base.TrappedFractionSource.FILE, + base.TrappedFractionSource.EXACT, + ]) + def test_trapped_fraction_source_not_supported( + self, trapped_fraction_source: base.TrappedFractionSource + ): + """Tests that FBT only supports SAUTER (no file or 2D equilibrium data).""" + with self.assertRaisesRegex(ValueError, 'not supported for FBTConfig'): + fbt.FBTConfig(trapped_fraction_source=trapped_fraction_source) + def test_edge_geometry_params_are_propagated(self): """Tests that edge geometry parameters are propagated to StandardGeometry.""" intermediate = standard_geometry.StandardGeometryIntermediates( @@ -68,6 +80,7 @@ def test_edge_geometry_params_are_propagated(self): R_OMP=np.array(8.2), R_target=np.array(7.0), B_pol_OMP=np.array(0.5), + trapped_fraction=np.arange(0, 1.0, 0.01), ) geo = standard_geometry.build_standard_geometry(intermediate) self.assertTrue(geo.diverted) diff --git a/torax/_src/geometry/tests/imas_test.py b/torax/_src/geometry/tests/imas_test.py index ec00e2dc2..dfe8c44bf 100644 --- a/torax/_src/geometry/tests/imas_test.py +++ b/torax/_src/geometry/tests/imas_test.py @@ -14,10 +14,12 @@ from absl.testing import absltest from absl.testing import parameterized import numpy as np +from torax._src.geometry import base from torax._src.geometry import chease from torax._src.geometry import eqdsk from torax._src.geometry import imas from torax._src.geometry import standard_geometry +from torax._src.imas_tools.input import loader # pylint: disable=invalid-name @@ -77,6 +79,99 @@ def test_gm4_gm5_terms(self): np.testing.assert_allclose(eqdsk_geo.gm5, chease_geo.gm5, rtol=0.02) np.testing.assert_allclose(imas_geo.gm5, chease_geo.gm5, rtol=0.01) + def test_trapped_fraction_is_loaded_from_file(self): + """Tests that the exact trapped particle fraction is loaded from IMAS.""" + geo = imas.IMASConfig( + imas_filepath='ITERhybrid_COCOS17_IDS_ddv4.nc', + trapped_fraction_source=base.TrappedFractionSource.FILE, + ).build_geometry() + trapped_fraction = geo.trapped_fraction_face + self.assertIsNotNone(trapped_fraction) + self.assertTrue(np.all(trapped_fraction >= 0.0)) + self.assertTrue(np.all(trapped_fraction <= 1.0)) + # Trapped fraction increases with normalized radius over most of the + # profile (small deviations from strict monotonicity are possible near + # the edge for diverted geometries). + self.assertGreater( + np.mean(np.diff(trapped_fraction) >= -1e-6), + 0.8, + ) + + def test_trapped_fraction_file_requires_provided_value(self): + """Tests that FILE raises when the equilibrium doesn't provide it.""" + equilibrium_object = loader.load_imas_data( + 'ITERhybrid_COCOS17_IDS_ddv4.nc', 'equilibrium' + ) + equilibrium_object.time_slice[0].profiles_1d.trapped_fraction = [] + + with self.assertRaisesRegex(ValueError, 'trapped_fraction_source=FILE'): + imas.IMASConfig( + equilibrium_object=equilibrium_object, + trapped_fraction_source=base.TrappedFractionSource.FILE, + ).build_geometry() + + @parameterized.named_parameters( + dict(testcase_name='not_provided_by_equilibrium_code', strip=True), + dict(testcase_name='provided_by_equilibrium_code', strip=False), + ) + def test_trapped_fraction_exact_is_computed(self, strip: bool): + """Tests the exact bounce-averaged integral, computed from the 2D grid. + + `EXACT` always computes the integral directly from the full 2D + equilibrium, regardless of whether the equilibrium code separately + provides `profiles_1d.trapped_fraction` (use `FILE` for that instead). + """ + equilibrium_object = loader.load_imas_data( + 'ITERhybrid_COCOS17_IDS_ddv4.nc', 'equilibrium' + ) + if strip: + equilibrium_object.time_slice[0].profiles_1d.trapped_fraction = [] + + geo = imas.IMASConfig( + equilibrium_object=equilibrium_object, + trapped_fraction_source=base.TrappedFractionSource.EXACT, + ).build_geometry() + trapped_fraction = geo.trapped_fraction_face + self.assertTrue(np.all(trapped_fraction >= 0.0)) + self.assertTrue(np.all(trapped_fraction <= 1.0)) + self.assertGreater( + np.mean(np.diff(trapped_fraction) >= -1e-6), + 0.8, + ) + + @parameterized.named_parameters( + dict( + testcase_name='file', + trapped_fraction_source=base.TrappedFractionSource.FILE, + ), + dict( + testcase_name='exact', + trapped_fraction_source=base.TrappedFractionSource.EXACT, + ), + ) + def test_trapped_fraction_geometry_consistent_with_sauter( + self, trapped_fraction_source: base.TrappedFractionSource + ): + """Tests that the exact and Sauter trapped fractions roughly agree.""" + geo_sauter = imas.IMASConfig( + imas_filepath='ITERhybrid_COCOS17_IDS_ddv4.nc', + trapped_fraction_source=base.TrappedFractionSource.SAUTER, + ).build_geometry() + geo_geometry = imas.IMASConfig( + imas_filepath='ITERhybrid_COCOS17_IDS_ddv4.nc', + trapped_fraction_source=trapped_fraction_source, + ).build_geometry() + + # Moderately coarse tolerance: Sauter is only an analytic approximation, + # so it need not match the exact integral closely, but a large deviation + # would indicate a bug rather than the expected model discrepancy. + np.testing.assert_allclose( + geo_geometry.trapped_fraction_face, + geo_sauter.trapped_fraction_face, + atol=0.05, + rtol=0.15, + ) + if __name__ == '__main__': absltest.main() diff --git a/torax/_src/geometry/tests/standard_geometry_test.py b/torax/_src/geometry/tests/standard_geometry_test.py index 5b22aeee7..cc29d8aac 100644 --- a/torax/_src/geometry/tests/standard_geometry_test.py +++ b/torax/_src/geometry/tests/standard_geometry_test.py @@ -72,6 +72,7 @@ def foo(geo: geometry.Geometry): R_OMP=None, R_target=None, B_pol_OMP=None, + trapped_fraction=np.arange(0, 1.0, 0.01), ) geo = standard_geometry.build_standard_geometry(intermediate) foo(geo) @@ -185,6 +186,7 @@ def _make_intermediates(self, **overrides): R_OMP=None, R_target=None, B_pol_OMP=None, + trapped_fraction=np.linspace(0.0, 0.5, 100), ) defaults.update(overrides) return standard_geometry.StandardGeometryIntermediates(**defaults) # pyrefly: ignore[bad-argument-type] diff --git a/torax/_src/imas_tools/input/equilibrium.py b/torax/_src/imas_tools/input/equilibrium.py index aec5d2c2e..03bfe2cba 100644 --- a/torax/_src/imas_tools/input/equilibrium.py +++ b/torax/_src/imas_tools/input/equilibrium.py @@ -17,10 +17,13 @@ import logging from typing import Any +import contourpy from imas import ids_toplevel import numpy as np import scipy +from torax._src.geometry import base from torax._src.imas_tools.input import loader +from torax._src.neoclassical.formulas import formulas # TODO(b/379832500) - Modify for consistency when we have a fixed TORAX COCOS. @@ -77,12 +80,115 @@ def _load_equilibrium( return equilibrium +# Below this many contour vertices, the poloidal (R, Z) equilibrium grid does +# not resolve the flux surface well enough for an accurate line integral +# (e.g. flux surfaces very close to the magnetic axis, which can be much +# smaller than a single grid cell). Surfaces below this threshold fall back +# to the Sauter approximation instead of the exact integral. +_MIN_CONTOUR_POINTS_FOR_EXACT_INTEGRAL = 20 + +# IMAS DD `equilibrium_profiles_2d_grid_type` identifier index for a +# rectangular (R, Z) grid, the only grid type currently supported for the +# exact bounce-averaged trapped fraction calculation below. +_IMAS_RECTANGULAR_GRID_TYPE = 1 + + +def _calculate_exact_trapped_fraction( + IMAS_data: Any, + flux_surf_avg_B2: np.ndarray, +) -> np.ndarray | None: + """Computes the trapped fraction from the full 2D equilibrium, if possible. + + Used to implement `TrappedFractionSource.EXACT`. Builds flux surface + contours from `profiles_2d` (mirroring the approach used for EQDSK + geometries) at each of the `profiles_1d.psi` grid points, and applies + `formulas.calculate_bounce_averaged_trapped_fraction` to each. + + Args: + IMAS_data: A single equilibrium IDS time slice. + flux_surf_avg_B2: Flux surface average of B^2 on the `profiles_1d.psi` + grid (i.e. `profiles_1d.gm5`). + + Returns: + The trapped fraction on the `profiles_1d.psi` grid, with NaN at any + surface too close to the magnetic axis for the 2D grid to resolve + reliably (the caller should fill these gaps with the Sauter + approximation), or None if no exact data is available at all (e.g. no + `profiles_2d`, or not on a rectangular grid), in which case the caller + should fall back to the Sauter approximation entirely. + """ + if not IMAS_data.profiles_2d or not IMAS_data.profiles_2d[0].psi: + return None + profiles_2d = IMAS_data.profiles_2d[0] + if profiles_2d.grid_type.index != _IMAS_RECTANGULAR_GRID_TYPE: + return None + + psi_1d = np.asarray(IMAS_data.profiles_1d.psi) + F_1d = np.asarray(IMAS_data.profiles_1d.f) + R = np.asarray(profiles_2d.r) + Z = np.asarray(profiles_2d.z) + psi_2d = np.asarray(profiles_2d.psi) + R_1D = R[:, 0] + Z_1D = Z[0, :] + + boundary_r = np.asarray(IMAS_data.boundary.outline.r) + boundary_z = np.asarray(IMAS_data.boundary.outline.z) + offset = 0.01 + mask = ( + (R > boundary_r.min() - offset) + & (R < boundary_r.max() + offset) + & (Z > boundary_z.min() - offset) + & (Z < boundary_z.max() + offset) + ) + masked_psi_2d = np.ma.masked_where(~mask, psi_2d) + + psi_2d_interpolator = scipy.interpolate.RectBivariateSpline( + R_1D, Z_1D, psi_2d, kx=3, ky=3, s=0 + ) + psi_contour_generator = contourpy.contour_generator(R, Z, masked_psi_2d) + + # No trapped particles on the magnetic axis (n=0), where B is uniform; no + # contour is defined there either way. + trapped_fraction = np.full(len(psi_1d), np.nan) + trapped_fraction[0] = 0.0 + for n in range(1, len(psi_1d)): + vertices = psi_contour_generator.create_contour(psi_1d[n]) + if ( + not vertices + or len(vertices[0]) < _MIN_CONTOUR_POINTS_FOR_EXACT_INTEGRAL + ): + # Contour generation failed, or the flux surface is too small for the + # grid to resolve well (typically only an issue very close to the + # magnetic axis). Leave as NaN; the caller falls back to Sauter. + continue + x_surface, z_surface = vertices[0].T[0], vertices[0].T[1] + surface_dl = np.sqrt( + np.gradient(x_surface) ** 2 + np.gradient(z_surface) ** 2 + ) + surface_dpsi_x = psi_2d_interpolator.ev(x_surface, z_surface, dx=1) + surface_dpsi_z = psi_2d_interpolator.ev(x_surface, z_surface, dy=1) + surface_Bpol = np.sqrt(surface_dpsi_x**2 + surface_dpsi_z**2) / ( + 2 * np.pi * x_surface + ) + surface_Btor = F_1d[n] / x_surface + surface_B = np.sqrt(surface_Bpol**2 + surface_Btor**2) + trapped_fraction[n] = formulas.calculate_bounce_averaged_trapped_fraction( + B=surface_B, + dl_over_Bp=surface_dl / surface_Bpol, + flux_surf_avg_B2=flux_surf_avg_B2[n], + ) + return trapped_fraction + + def _geometry_from_single_slice( equilibrium: ids_toplevel.IDSToplevel, face_centers: np.ndarray, Ip_from_parameters: bool = False, hires_factor: int = 4, slice_index: int = 0, + trapped_fraction_source: base.TrappedFractionSource = ( + base.TrappedFractionSource.SAUTER + ), ) -> dict[str, Any]: """Extracts geometry data from a single time slice of an equilibrium IDS. @@ -94,6 +200,8 @@ def _geometry_from_single_slice( hires_factor: Grid refinement factor for poloidal flux <--> plasma current calculations. slice_index: Index of the time slice to process. + trapped_fraction_source: Selects how the effective trapped particle + fraction is computed; see `base.TrappedFractionSource`. Returns: A dict of intermediate geometry values for building a StandardGeometry. @@ -205,6 +313,57 @@ def _geometry_from_single_slice( z_magnetic_axis = np.asarray(IMAS_data.global_quantities.magnetic_axis.z) + sauter_trapped_fraction = formulas.calculate_sauter_trapped_fraction( + epsilon=(R_out - R_in) / (R_out + R_in), + delta=0.5 + * ( + IMAS_data.profiles_1d.triangularity_upper + + IMAS_data.profiles_1d.triangularity_lower + ), + ) + + match trapped_fraction_source: + case base.TrappedFractionSource.SAUTER: + trapped_fraction = sauter_trapped_fraction + case base.TrappedFractionSource.FILE: + if not IMAS_data.profiles_1d.trapped_fraction: + raise ValueError( + "trapped_fraction_source=FILE requires the equilibrium IDS to" + " populate profiles_1d.trapped_fraction, but this IDS does" + " not. Use trapped_fraction_source=EXACT to compute it directly" + " from the 2D equilibrium instead, or SAUTER for the analytic" + " approximation." + ) + trapped_fraction = np.asarray(IMAS_data.profiles_1d.trapped_fraction) + case base.TrappedFractionSource.EXACT: + exact_trapped_fraction = _calculate_exact_trapped_fraction( + IMAS_data, np.asarray(IMAS_data.profiles_1d.gm5) + ) + if exact_trapped_fraction is None: + raise ValueError( + "trapped_fraction_source=EXACT requires a rectangular" + " profiles_2d psi grid to compute the bounce-averaged integral," + " but this equilibrium IDS does not provide one. Use" + " trapped_fraction_source=FILE to read a value precomputed by" + " the equilibrium code instead (if available), or SAUTER for" + " the analytic approximation." + ) + # Fill any unreliable values (NaN, or outside the physically valid + # [0, 1] range, e.g. surfaces too close to the magnetic axis for the + # grid to resolve well) with the Sauter approximation. + exact_is_unreliable = ( + np.isnan(exact_trapped_fraction) + | (exact_trapped_fraction < 0.0) + | (exact_trapped_fraction > 1.0) + ) + trapped_fraction = np.where( + exact_is_unreliable, sauter_trapped_fraction, exact_trapped_fraction + ) + case _: + raise ValueError( + f"Unknown trapped_fraction_source: {trapped_fraction_source}" + ) + # TODO(b/446608829): Add support for edge geometries from IMAS. return { @@ -226,6 +385,7 @@ def _geometry_from_single_slice( "flux_surf_avg_grad_psi2_over_R2": flux_surf_avg_grad_psi2_over_R2, "flux_surf_avg_B2": IMAS_data.profiles_1d.gm5, "flux_surf_avg_1_over_B2": IMAS_data.profiles_1d.gm4, + "trapped_fraction": trapped_fraction, "delta_upper_face": IMAS_data.profiles_1d.triangularity_upper, "delta_lower_face": IMAS_data.profiles_1d.triangularity_lower, "elongation": IMAS_data.profiles_1d.elongation, @@ -255,6 +415,9 @@ def geometry_from_IMAS( imas_uri: str | None = None, imas_filepath: str | None = None, explicit_convert: bool = False, + trapped_fraction_source: base.TrappedFractionSource = ( + base.TrappedFractionSource.SAUTER + ), ) -> Mapping[float, dict[str, Any]]: """Constructs geometry intermediates for all time slices in an IMAS IDS. @@ -276,6 +439,8 @@ def geometry_from_IMAS( version. If True, an explicit conversion will be attempted. Explicit conversion is recommended when converting between major DD versions. https://imas-python.readthedocs.io/en/latest/multi-dd.html#conversion-of-idss-between-dd-versions + trapped_fraction_source: Selects how the effective trapped particle + fraction is computed; see `base.TrappedFractionSource`. Returns: A mapping from times to dicts of intermediate geometry values, one per @@ -301,6 +466,7 @@ def geometry_from_IMAS( face_centers=face_centers, Ip_from_parameters=Ip_from_parameters, hires_factor=hires_factor, + trapped_fraction_source=trapped_fraction_source, ) return intermediates @@ -317,6 +483,9 @@ def geometry_from_single_IMAS_slice( explicit_convert: bool = False, slice_index: int = 0, slice_time: float | None = None, + trapped_fraction_source: base.TrappedFractionSource = ( + base.TrappedFractionSource.SAUTER + ), ) -> dict[str, Any]: """Constructs geometry intermediates for a single time slice in an IMAS IDS. @@ -356,6 +525,8 @@ def geometry_from_single_IMAS_slice( slice_time: Time (in seconds) of the IDS time slice to load. The slice whose time is closest to this value is selected. When provided, takes precedence over ``slice_index``. + trapped_fraction_source: Selects how the effective trapped particle + fraction is computed; see `base.TrappedFractionSource`. Returns: A dict of intermediate geometry values for building a StandardGeometry, @@ -391,4 +562,5 @@ def geometry_from_single_IMAS_slice( face_centers=face_centers, Ip_from_parameters=Ip_from_parameters, hires_factor=hires_factor, + trapped_fraction_source=trapped_fraction_source, ) diff --git a/torax/_src/neoclassical/bootstrap_current/redl.py b/torax/_src/neoclassical/bootstrap_current/redl.py index 7b7104a9f..c7c029cd5 100644 --- a/torax/_src/neoclassical/bootstrap_current/redl.py +++ b/torax/_src/neoclassical/bootstrap_current/redl.py @@ -117,7 +117,7 @@ def _calculate_bootstrap_current( # collisionality and for multi-species plasmas. # Effective trapped particle fraction - f_trap = formulas.calculate_f_trap(geo) + f_trap = geo.trapped_fraction_face # Collision frequencies log_lambda_ei = collisions.calculate_log_lambda_ei( diff --git a/torax/_src/neoclassical/bootstrap_current/sauter.py b/torax/_src/neoclassical/bootstrap_current/sauter.py index 4cac40b1c..49afa87b5 100644 --- a/torax/_src/neoclassical/bootstrap_current/sauter.py +++ b/torax/_src/neoclassical/bootstrap_current/sauter.py @@ -108,7 +108,7 @@ def _calculate_bootstrap_current( # corrections. # Effective trapped particle fraction - f_trap = formulas.calculate_f_trap(geo) + f_trap = geo.trapped_fraction_face # Spitzer conductivity log_lambda_ei = collisions.calculate_log_lambda_ei( diff --git a/torax/_src/neoclassical/conductivity/sauter.py b/torax/_src/neoclassical/conductivity/sauter.py index 1c50e114e..db3743581 100644 --- a/torax/_src/neoclassical/conductivity/sauter.py +++ b/torax/_src/neoclassical/conductivity/sauter.py @@ -91,7 +91,7 @@ def _calculate_conductivity( # Formulas from Sauter PoP 1999. # Effective trapped particle fraction - f_trap = formulas.calculate_f_trap(geo) + f_trap = geo.trapped_fraction_face # Spitzer conductivity NZ = 0.58 + 0.74 / (0.76 + Z_eff_face) diff --git a/torax/_src/neoclassical/formulas/formulas.py b/torax/_src/neoclassical/formulas/formulas.py index f865d2876..f551cdceb 100644 --- a/torax/_src/neoclassical/formulas/formulas.py +++ b/torax/_src/neoclassical/formulas/formulas.py @@ -15,6 +15,7 @@ import jax import jax.numpy as jnp +import numpy as np from torax._src import array_typing from torax._src import constants from torax._src.fvm import cell_variable @@ -26,29 +27,67 @@ # pylint: disable=invalid-name -def calculate_f_trap( - geo: geometry_lib.Geometry, -) -> array_typing.FloatVectorFace: - """Calculates the effective trapped particle fraction. +def calculate_sauter_trapped_fraction( + epsilon: array_typing.Array, delta: array_typing.Array +) -> array_typing.Array: + """Analytic approximation for the effective trapped particle fraction. - From O. Sauter, Fusion Engineering and Design 112 (2016) 633-645. Eqs 33+34. + From O. Sauter, Fusion Engineering and Design 112 (2016) 633-645, Eqs 33+34. Args: - geo: The magnetic geometry. + epsilon: Local midplane inverse aspect ratio of each flux surface. + delta: Average triangularity of each flux surface. Returns: - The effective trapped particle fraction. + The effective trapped particle fraction of each flux surface. """ + epsilon_effective = 0.67 * (1.0 - 1.4 * np.abs(delta) * delta) * epsilon + aa = (1.0 - epsilon) / (1.0 + epsilon) + return 1.0 - np.sqrt(aa) * (1.0 - epsilon_effective) / ( + 1.0 + 2.0 * np.sqrt(epsilon_effective) + ) + + +def calculate_bounce_averaged_trapped_fraction( + B: array_typing.Array, + dl_over_Bp: array_typing.Array, + flux_surf_avg_B2: array_typing.Array, +) -> array_typing.Array: + r"""Effective trapped particle fraction of one flux surface, exactly. + + Computed from the full bounce-averaged integral, as opposed to the + `calculate_sauter_trapped_fraction` analytic approximation: - epsilon_effective = ( - 0.67 - * (1.0 - 1.4 * jnp.abs(geo.delta_face) * geo.delta_face) - * geo.epsilon_face + .. math:: + f_t = 1 - \frac{3}{4} \langle B^2 \rangle + \int_0^{1/B_{max}} \frac{\lambda \, d\lambda}{\langle \sqrt{1 - + \lambda B} \rangle} + + where :math:`\langle . \rangle` is the flux surface average, using the same + :math:`dl/B_p` weighting as other flux surface averages. This requires the + full poloidal variation of :math:`|B|` on the flux surface. + + Args: + B: :math:`|B|` at samples of a poloidal contour around one flux surface + [:math:`\mathrm{T}`]. + dl_over_Bp: Poloidal line-element weights :math:`dl / B_p` at the same + contour samples [:math:`\mathrm{m/T}`]. + flux_surf_avg_B2: Flux surface average of :math:`B^2` for this flux + surface [:math:`\mathrm{T}^2`]. + + Returns: + The effective trapped particle fraction of this flux surface. + """ + B_max = B.max() + lam = np.linspace(0.0, 1.0, 101) / B_max + sqrt_term = np.sqrt( + np.clip(1.0 - lam[:, np.newaxis] * B[np.newaxis, :], 0.0, None) ) - 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) + h_lambda = np.sum(sqrt_term * dl_over_Bp[np.newaxis, :], axis=1) / np.sum( + dl_over_Bp ) + bounce_integral = np.trapezoid(lam / np.maximum(h_lambda, 1e-10), lam) + return 1.0 - 0.75 * flux_surf_avg_B2 * bounce_integral # TODO(b/428166775): currently we have two very similar implementations for diff --git a/torax/_src/neoclassical/formulas/tests/formulas_test.py b/torax/_src/neoclassical/formulas/tests/formulas_test.py index fe7f26c0f..8bf20ed70 100644 --- a/torax/_src/neoclassical/formulas/tests/formulas_test.py +++ b/torax/_src/neoclassical/formulas/tests/formulas_test.py @@ -12,14 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest import mock - from absl.testing import absltest from absl.testing import parameterized import numpy as np from torax._src.config import build_runtime_params from torax._src.core_profiles import initialization -from torax._src.geometry import geometry from torax._src.neoclassical.formulas import formulas from torax._src.physics import collisions from torax._src.torax_pydantic import model_config @@ -91,27 +88,17 @@ def setUp(self): log_lambda_ei=log_lambda_ei, ) - self.f_trap = formulas.calculate_f_trap(self.geo) - - def test_calculate_f_trap_positive_triangularity(self): - geo = mock.create_autospec( - geometry.Geometry, - instance=True, - delta_face=np.array(0.2), - epsilon_face=np.array(0.1), + def test_calculate_sauter_trapped_fraction_positive_triangularity(self): + result = formulas.calculate_sauter_trapped_fraction( + epsilon=np.array(0.1), delta=np.array(0.2) ) - result = formulas.calculate_f_trap(geo) expected = 0.4362384616678634 np.testing.assert_allclose(result, expected) - def test_calculate_f_trap_negative_triangularity(self): - geo = mock.create_autospec( - geometry.Geometry, - instance=True, - delta_face=np.array(-0.2), - epsilon_face=np.array(0.1), + def test_calculate_sauter_trapped_fraction_negative_triangularity(self): + result = formulas.calculate_sauter_trapped_fraction( + epsilon=np.array(0.1), delta=np.array(-0.2) ) - result = formulas.calculate_f_trap(geo) expected = 0.45134158459680895 np.testing.assert_allclose(result, expected) diff --git a/torax/_src/neoclassical/formulas/tests/redl_test.py b/torax/_src/neoclassical/formulas/tests/redl_test.py index 338206ed7..2de976f81 100644 --- a/torax/_src/neoclassical/formulas/tests/redl_test.py +++ b/torax/_src/neoclassical/formulas/tests/redl_test.py @@ -88,7 +88,7 @@ def setUp(self): log_lambda_ei=log_lambda_ei, ) - self.f_trap = formulas.calculate_f_trap(self.geo) + self.f_trap = self.geo.trapped_fraction_face def test_L31_values_are_correct(self): L31 = redl_formulas.calculate_L31( @@ -105,28 +105,28 @@ def test_L32_values_are_correct(self): _L31_EXPECTED = np.array([ 0.0, - 0.24302010813886246, - 0.36213029680638326, - 0.44486357141119376, - 0.5036505543784661, - 0.5456217675849331, - 0.5743686943725517, - 0.5884757302502829, - 0.5788428389680405, - 0.5133804621522303, + 0.24300339539816412, + 0.36212967141791325, + 0.44486186876271433, + 0.5036494714316562, + 0.545621750277896, + 0.5743683188240756, + 0.5884751857797844, + 0.5788407842140618, + 0.5133748206351552, 0.28646594502191125, ]) _L32_EXPECTED = np.array([ 0.0, - -0.040995932527566226, - -0.08550417010706743, - -0.10122043852379645, - -0.10048876254207006, - -0.09165040157223098, - -0.07639140154722468, - -0.05120454722561363, - -0.006618464913488831, - 0.07582955103486738, + -0.04099565836762635, + -0.08550416711714443, + -0.10122050749428985, + -0.10048885740690477, + -0.09165040387502599, + -0.07639146897224214, + -0.05120467415215507, + -0.006619092698929574, + 0.07582707623300719, 0.18210763398539093, ]) diff --git a/torax/_src/neoclassical/formulas/tests/sauter_test.py b/torax/_src/neoclassical/formulas/tests/sauter_test.py index dd0816f5b..3f58306d7 100644 --- a/torax/_src/neoclassical/formulas/tests/sauter_test.py +++ b/torax/_src/neoclassical/formulas/tests/sauter_test.py @@ -91,7 +91,7 @@ def setUp(self): log_lambda_ei=log_lambda_ei, ) - self.f_trap = formulas.calculate_f_trap(self.geo) + self.f_trap = self.geo.trapped_fraction_face def test_L31_values_are_correct(self): L31 = sauter_formulas.calculate_L31( @@ -107,29 +107,29 @@ def test_L32_values_are_correct(self): _L31_EXPECTED = np.array([ 0.0, - 0.25942749, - 0.39198664, - 0.48032915, - 0.53634519, - 0.57082292, - 0.5894148, - 0.59111759, - 0.56839259, - 0.5001917, - 0.33682819, + 0.2594107633186986, + 0.39198603870801796, + 0.48032755509905906, + 0.5363442162030589, + 0.570822902946326, + 0.5894144834221119, + 0.5911171365474494, + 0.5683909023668218, + 0.5001868536054601, + 0.33682819436163186, ]) _L32_EXPECTED = np.array([ 0.0, - -0.03501634, - -0.07505174, - -0.09268982, - -0.09172278, - -0.08055448, - -0.06213122, - -0.03385067, - 0.01149523, - 0.08557197, - 0.16296924, + -0.03501691867331952, + -0.07505176373149916, + -0.0926899359270229, + -0.0917228945357012, + -0.08055448284343791, + -0.06213128742807339, + -0.03385079089994658, + 0.011494643548927597, + 0.08556965525319629, + 0.16296924403319402, ]) if __name__ == '__main__': diff --git a/torax/_src/orchestration/tests/sim_state_test.py b/torax/_src/orchestration/tests/sim_state_test.py index 93d3a6004..246601a90 100644 --- a/torax/_src/orchestration/tests/sim_state_test.py +++ b/torax/_src/orchestration/tests/sim_state_test.py @@ -52,6 +52,7 @@ def _make_geometry(self, **overrides) -> standard_geometry.StandardGeometry: flux_surf_avg_grad_psi2_over_R2=np.linspace(0.01, 1.0, 10), flux_surf_avg_B2=np.linspace(25.0, 30.0, 10), flux_surf_avg_1_over_B2=np.linspace(0.03, 0.04, 10), + trapped_fraction=np.linspace(0.0, 0.5, 10), delta_upper_face=np.linspace(0.0, 0.3, 10), delta_lower_face=np.linspace(0.0, 0.3, 10), elongation=np.linspace(1.0, 1.7, 10), diff --git a/torax/_src/test_utils/references.json b/torax/_src/test_utils/references.json index 5e94011db..5e2759791 100644 --- a/torax/_src/test_utils/references.json +++ b/torax/_src/test_utils/references.json @@ -168,191 +168,191 @@ }, "chease_references_Ip_from_chease": { "psi": [ - 0.027987856747679257, + 0.027987856747679264, 0.25802053150777976, 0.7512547932373561, - 1.528708510114555, + 1.5287085101145557, 2.615697515298346, 4.109279464112691, 6.083132046719322, 8.548002981591813, - 11.417797255915486, + 11.417797255915483, 14.544591130509547, 17.801364061495992, - 21.109232630763035, + 21.109232630763028, 24.425647141348136, 27.716706270506247, 30.953531107440565, 34.11189477326518, 37.17235977851071, - 40.119267600702344, + 40.11926760070234, 42.937960721543085, 45.61729884452383, 48.143819194434435, 50.49230227280328, 52.66771218657442, - 54.76107429655608, + 54.761074296556075, 56.79416882696004 ], "psi_face_grad": [ 0.0, - 5.750816869002513, - 12.330856543239406, - 19.43634292192997, + 5.750816869002514, + 12.330856543239404, + 19.43634292192998, 27.174725129594773, - 37.33954872035861, - 49.34631456516578, - 61.621773371812246, - 71.74485685809199, - 78.16984686485144, - 81.4193232746611, - 82.69671423167624, - 82.91036276462751, - 82.27647822895273, - 80.92062092335802, - 78.95909164561543, - 76.51162513113833, - 73.67269555479096, - 70.4673280210185, - 66.98345307451862, - 63.16300874776515, - 58.712076959221235, - 54.38524784427858, - 52.3340527495416, - 50.82736326009918, + 37.33954872035862, + 49.34631456516575, + 61.62177337181221, + 71.74485685809188, + 78.16984686485154, + 81.41932327466104, + 82.69671423167604, + 82.91036276462764, + 82.2764782289527, + 80.92062092335786, + 78.95909164561546, + 76.51162513113825, + 73.67269555479079, + 70.46732802101866, + 66.98345307451856, + 63.16300874776506, + 58.7120769592212, + 54.38524784427846, + 52.33405274954147, + 50.82736326009906, 50.41748508435972 ], "psidot": [ - 0.018996072254309097, - 0.02500202031296845, - 0.031338773378838745, - 0.03868328280898349, - 0.05299467477740712, - 0.07198210090577462, - 0.11974110051938826, - 0.33574868167253263, - 0.975780784130611, - 1.7948330791789693, - 1.839253779269071, - 1.0422908200755818, - 0.35501066028069256, - 0.11457588869572854, - 0.07324057006414057, - 0.07045893350049669, - 0.07154945891442197, - 0.07366953948700813, - 0.07805304209657025, - 0.08252448382248863, - 0.08435187644333315, - 0.11926590254333072, - 0.21052605840812091, - 0.3845679042885645, - 0.6046308208774739 + 0.01899598263975594, + 0.025001584097881297, + 0.0313381988081886, + 0.038682878210207536, + 0.05299432864450665, + 0.07198196066892859, + 0.11974071599065529, + 0.3357474355268937, + 0.97577754082544, + 1.7948286373425513, + 1.8392483073085446, + 1.0422884204963254, + 0.3550102413323787, + 0.11457560227412626, + 0.07324040977481404, + 0.07045874144091978, + 0.07154914766840711, + 0.07366915921386581, + 0.07805245221292942, + 0.08252390850459809, + 0.08435108710425139, + 0.1192647270027188, + 0.21052354270094503, + 0.3845600792209157, + 0.6046242949507349 ], "j_total": [ - 813160.9908847061, + 813160.9908847059, 884353.6628799367, - 940038.6859902174, - 1004231.0472487116, - 1197720.4210286674, - 1368504.2452566356, - 1437133.9811527752, - 1328921.3312816792, - 1097679.8401292658, - 879657.1851242019, - 731828.2648608742, - 638410.3378917036, - 560602.3284088922, - 492467.49330409843, - 432072.10054247826, - 383610.8071653922, - 341163.4830060164, - 304636.17759270134, - 276649.77860184625, - 247762.80897923285, - 211372.2943842105, - 244447.8942875597, - 350771.5245092025, - 501684.3148104644, - 522719.6370215425 + 940038.6859902188, + 1004231.0472487104, + 1197720.421028668, + 1368504.2452566326, + 1437133.981152774, + 1328921.3312816736, + 1097679.8401292812, + 879657.1851241911, + 731828.2648608637, + 638410.3378917262, + 560602.3284088803, + 492467.49330409017, + 432072.1005424921, + 383610.80716537987, + 341163.4830060095, + 304636.1775927303, + 276649.7786018271, + 247762.80897922878, + 211372.29438421573, + 244447.89428754957, + 350771.52450920007, + 501684.3148104631, + 522719.637021559 ], "q": [ - 1.7477328615081547, - 1.7477328615081547, - 1.6302016955961902, - 1.551355364974149, + 1.7477328615081544, + 1.7477328615081544, + 1.6302016955961904, + 1.551355364974148, 1.4794470338947647, - 1.3458774900767738, - 1.2220841671000315, - 1.141743210354947, - 1.12073724167867, - 1.1571984368682962, - 1.234460226175709, - 1.3369310845585383, - 1.4547119979688514, - 1.588079532628141, - 1.7388952421394317, - 1.9093858755838253, - 2.101827868430554, - 2.3192467208551566, - 2.5673748996204018, - 2.850956946255101, - 3.1825246522401462, - 3.5949796873732214, - 4.065801379217709, - 4.4172101179162135, - 4.745896372882916, + 1.3458774900767732, + 1.2220841671000322, + 1.1417432103549476, + 1.1207372416786718, + 1.1571984368682946, + 1.2344602261757098, + 1.3369310845585414, + 1.4547119979688492, + 1.5880795326281416, + 1.7388952421394352, + 1.9093858755838247, + 2.101827868430556, + 2.319246720855162, + 2.5673748996203964, + 2.850956946255103, + 3.1825246522401507, + 3.5949796873732236, + 4.065801379217719, + 4.417210117916224, + 4.745896372882927, 4.983832298285873 ], "s": [ -0.0, - -0.03604804430932131, - -0.11807194454358688, - -0.1454563011087649, - -0.29119080774593703, - -0.47895077990130097, - -0.48703968044756485, - -0.2956937457296233, - 0.052440025411032924, - 0.42804329193023366, - 0.7170628430731953, - 0.9018555843838894, - 1.032470140237754, - 1.1596678921942585, - 1.2898072214480516, - 1.4220354239044362, - 1.5562626967774909, - 1.7007598754941313, - 1.8581005105745956, - 2.0390789929569144, - 2.3103133450727102, - 2.5767043592462664, - 2.3157298338849213, - 1.7904041117839449, - 1.4658231947836782, - 1.2533771637033337 + -0.03604804430932115, + -0.11807194454358734, + -0.14545630110876517, + -0.2911908077459367, + -0.4789507799012993, + -0.487039680447562, + -0.2956937457296198, + 0.05244002541102533, + 0.42804329193022933, + 0.7170628430732142, + 0.9018555843838799, + 1.0324701402377394, + 1.1596678921942811, + 1.2898072214480487, + 1.4220354239044262, + 1.5562626967775144, + 1.7007598754941087, + 1.8581005105745758, + 2.039078992956951, + 2.310313345072707, + 2.5767043592462735, + 2.3157298338849426, + 1.7904041117839475, + 1.4658231947836504, + 1.2533771637032702 ] }, "chease_references_Ip_from_runtime_params": { "psi": [ - 0.03566926771324071, + 0.035669267713240714, 0.32883559097917703, 0.9574405279554149, - 1.9482704086409488, + 1.9482704086409497, 3.3335891265690307, 5.237092308831824, 7.752678865894843, 10.894046282742263, - 14.55147032830675, + 14.551470328306747, 18.53642883379377, 22.687039814973087, - 26.90276989466861, + 26.902769894668598, 31.129391393149195, 35.323698595603325, 39.44888662591045, 43.47407941390097, 47.37450475133448, - 51.13020655385683, + 51.13020655385682, 54.72250447202719, 58.13720069776861, 61.35713752810074, @@ -363,141 +363,141 @@ ], "psi_face_grad": [ 0.0, - 7.329158081648409, - 15.715123424405943, - 24.77074701713834, - 34.63296794820205, - 47.5875795565698, - 62.88966392657549, - 78.53418542118546, - 91.4356011391124, - 99.62396263717541, - 103.7652745294829, - 105.39325199238833, - 105.66553746201453, - 104.85768006135318, - 103.12970075767808, + 7.32915808164841, + 15.715123424405942, + 24.770747017138362, + 34.632967948202044, + 47.58757955656982, + 62.88966392657547, + 78.53418542118541, + 91.43560113911226, + 99.62396263717551, + 103.7652745294828, + 105.39325199238796, + 105.66553746201484, + 104.85768006135316, + 103.12970075767794, 100.62981969976317, - 97.51063343583778, - 93.89254506305899, - 89.80744795425883, - 85.3674056435356, - 80.49842075830345, - 74.8259078906606, - 69.31155831249474, - 66.69740219379408, - 64.77719403908814, + 97.51063343583768, + 93.89254506305882, + 89.80744795425909, + 85.3674056435355, + 80.49842075830318, + 74.82590789066062, + 69.31155831249455, + 66.69740219379405, + 64.77719403908779, 64.25482269382653 ], "psidot": [ - 0.024048393241628024, - 0.03165563353561173, - 0.03967773134427998, - 0.04897536898024741, - 0.0670929243280754, - 0.09112949365674046, - 0.1515893796791297, - 0.4250430136726126, - 1.2352821337794926, - 2.2721414681545835, - 2.3283692504582554, - 1.3194685227924063, - 0.4494199182053732, - 0.1450460602295121, - 0.09271870483546152, - 0.08919818479127689, - 0.09057998420967996, - 0.09326565531704084, - 0.09881750801711751, - 0.10448164788765779, - 0.10679938417799156, - 0.15101174664116138, - 0.26657619547629646, - 0.4869751410210283, - 0.7656480809369889 + 0.02404827965283221, + 0.03165508069099443, + 0.03967700335291807, + 0.04897485645875391, + 0.06709248590599773, + 0.09112931604689754, + 0.15158889269899634, + 0.42504143556582025, + 1.2352780265530574, + 2.2721358431283325, + 2.328362320789424, + 1.3194654839395648, + 0.4494193876205266, + 0.14504569747241036, + 0.09271850182002422, + 0.08919794151675323, + 0.09057958994517452, + 0.09326517356711879, + 0.09881676064679595, + 0.10448091888803607, + 0.10679838380153639, + 0.15101025660117245, + 0.2665730058036467, + 0.48696521783944424, + 0.7656398040134632 ], "j_total": [ 1036337.1993547073, 1127069.0656605966, - 1198037.1292334995, - 1279847.413583486, - 1526440.9393130587, - 1744097.260852705, - 1831562.7800897774, - 1693650.612930658, - 1398943.7074078445, - 1121083.6154745142, - 932682.2891361101, - 813625.3325309503, - 714462.5154029084, - 627627.7250200276, - 550656.5066633691, - 488894.7625333805, - 434797.55234676687, - 388245.14044268907, - 352577.73057510896, - 315762.58384266647, - 269384.5056991217, - 311537.87379607663, - 447042.57016526576, - 639374.1505052428, - 666182.7248860064 + 1198037.1292335014, + 1279847.4135834838, + 1526440.9393130608, + 1744097.2608527034, + 1831562.780089775, + 1693650.61293065, + 1398943.7074078599, + 1121083.6154745014, + 932682.2891360901, + 813625.3325310008, + 714462.5154028847, + 627627.7250200156, + 550656.5066633822, + 488894.76253337227, + 434797.55234675755, + 388245.1404427291, + 352577.7305750764, + 315762.58384264604, + 269384.5056991505, + 311537.87379605253, + 447042.57016528543, + 639374.1505051983, + 666182.7248860548 ], "q": [ - 1.3713569158288057, - 1.3713569158288057, - 1.279136199065655, - 1.2172695012609593, - 1.1608466981522922, - 1.056041483526507, - 0.9589071712202829, - 0.8958677164596723, - 0.8793854032570487, - 0.9079946451405767, - 0.9686180341203997, - 1.0490216950053617, - 1.1414383759038902, - 1.2460850842367805, - 1.3644224862556837, + 1.3713569158288055, + 1.3713569158288055, + 1.2791361990656551, + 1.2172695012609582, + 1.1608466981522925, + 1.0560414835265066, + 0.9589071712202831, + 0.8958677164596728, + 0.8793854032570502, + 0.9079946451405758, + 0.9686180341204006, + 1.0490216950053652, + 1.1414383759038869, + 1.2460850842367808, + 1.3644224862556857, 1.4981977984943338, - 1.649197223863326, - 1.819794718177623, - 2.014488255992146, - 2.23700063727323, - 2.4971649202035744, - 2.820797368477212, - 3.190227158036696, - 3.465959688282214, - 3.723863043430813, + 1.6491972238633277, + 1.8197947181776264, + 2.0144882559921404, + 2.2370006372732325, + 2.497164920203583, + 2.8207973684772116, + 3.1902271580367043, + 3.4659596882822155, + 3.7238630434308333, 3.9105592394066315 ], "s": [ -0.0, - -0.03604804430932119, - -0.11807194454358667, - -0.1454563011087658, - -0.2911908077459369, - -0.4789507799012987, - -0.48703968044756635, - -0.29569374572962437, - 0.0524400254110341, - 0.4280432919302306, - 0.71706284307319, - 0.9018555843838995, - 1.032470140237763, - 1.1596678921942518, - 1.2898072214480547, - 1.4220354239044277, - 1.5562626967774966, - 1.7007598754941442, - 1.858100510574589, - 2.0390789929568833, - 2.310313345072706, - 2.5767043592463095, - 2.315729833884929, - 1.7904041117838982, - 1.4658231947836486, - 1.253377163703465 + -0.036048044309321084, + -0.11807194454358734, + -0.14545630110876548, + -0.29119080774593636, + -0.47895077990129825, + -0.48703968044756324, + -0.2956937457296193, + 0.05244002541102766, + 0.42804329193022717, + 0.7170628430732121, + 0.901855584383881, + 1.032470140237736, + 1.1596678921942811, + 1.2898072214480565, + 1.422035423904421, + 1.5562626967775117, + 1.7007598754941164, + 1.8581005105745745, + 2.039078992956945, + 2.3103133450726996, + 2.5767043592462935, + 2.3157298338849417, + 1.7904041117839227, + 1.4658231947836518, + 1.2533771637033213 ] } } \ No newline at end of file diff --git a/torax/examples/step_flattop_bgb.py b/torax/examples/step_flattop_bgb.py index 3815caec6..c4e5d3d02 100644 --- a/torax/examples/step_flattop_bgb.py +++ b/torax/examples/step_flattop_bgb.py @@ -97,6 +97,11 @@ "geometry_type": "IMAS", "imas_filepath": "STEP_SPP_001_ECHD_ftop.nc", "n_rho": 100, + # STEP is a low-aspect-ratio (spherical tokamak) scenario, where the + # Sauter analytic approximation for the trapped particle fraction is + # least accurate. Compute the exact value directly from the 2D + # equilibrium instead. + "trapped_fraction_source": "EXACT", }, "pedestal": { "model_name": "set_T_ped_n_ped", diff --git a/torax/tests/test_data/test_all_transport_fusion_qlknn.nc b/torax/tests/test_data/test_all_transport_fusion_qlknn.nc index d2264b90b..39fb8ebcd 100644 Binary files a/torax/tests/test_data/test_all_transport_fusion_qlknn.nc and b/torax/tests/test_data/test_all_transport_fusion_qlknn.nc differ diff --git a/torax/tests/test_data/test_bohmgyrobohm_all.nc b/torax/tests/test_data/test_bohmgyrobohm_all.nc index 7da3106d2..e483036f1 100644 Binary files a/torax/tests/test_data/test_bohmgyrobohm_all.nc and b/torax/tests/test_data/test_bohmgyrobohm_all.nc differ diff --git a/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc b/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc index 084333079..8d0d4b36a 100644 Binary files a/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc and b/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc differ diff --git a/torax/tests/test_data/test_changing_config_after.nc b/torax/tests/test_data/test_changing_config_after.nc index e4558e2c2..6fb12a9eb 100644 Binary files a/torax/tests/test_data/test_changing_config_after.nc and b/torax/tests/test_data/test_changing_config_after.nc differ diff --git a/torax/tests/test_data/test_changing_config_before.nc b/torax/tests/test_data/test_changing_config_before.nc index a935e2b20..a9594ba07 100644 Binary files a/torax/tests/test_data/test_changing_config_before.nc and b/torax/tests/test_data/test_changing_config_before.nc differ diff --git a/torax/tests/test_data/test_chease.nc b/torax/tests/test_data/test_chease.nc index 9fcc72286..5d44ecdb5 100644 Binary files a/torax/tests/test_data/test_chease.nc and b/torax/tests/test_data/test_chease.nc differ diff --git a/torax/tests/test_data/test_combined_transport.nc b/torax/tests/test_data/test_combined_transport.nc index bc808b02c..db6b60e3f 100644 Binary files a/torax/tests/test_data/test_combined_transport.nc and b/torax/tests/test_data/test_combined_transport.nc differ diff --git a/torax/tests/test_data/test_fixed_dt.nc b/torax/tests/test_data/test_fixed_dt.nc index 17c692a78..6f48b79be 100644 Binary files a/torax/tests/test_data/test_fixed_dt.nc and b/torax/tests/test_data/test_fixed_dt.nc differ diff --git a/torax/tests/test_data/test_imas_profiles_and_geo.nc b/torax/tests/test_data/test_imas_profiles_and_geo.nc index 18e45216d..aaa9baeb2 100644 Binary files a/torax/tests/test_data/test_imas_profiles_and_geo.nc and b/torax/tests/test_data/test_imas_profiles_and_geo.nc differ diff --git a/torax/tests/test_data/test_implicit.nc b/torax/tests/test_data/test_implicit.nc index 7dbf71eaf..2aded02cb 100644 Binary files a/torax/tests/test_data/test_implicit.nc and b/torax/tests/test_data/test_implicit.nc differ diff --git a/torax/tests/test_data/test_implicit_short_optimizer.nc b/torax/tests/test_data/test_implicit_short_optimizer.nc index 795165d3d..72dca065d 100644 Binary files a/torax/tests/test_data/test_implicit_short_optimizer.nc and b/torax/tests/test_data/test_implicit_short_optimizer.nc differ diff --git a/torax/tests/test_data/test_iterbaseline_mockup.nc b/torax/tests/test_data/test_iterbaseline_mockup.nc index 56eed20cb..097e53291 100644 Binary files a/torax/tests/test_data/test_iterbaseline_mockup.nc and b/torax/tests/test_data/test_iterbaseline_mockup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_lh_transition.nc b/torax/tests/test_data/test_iterhybrid_lh_transition.nc index bcb9ccda8..174ccd6a0 100644 Binary files a/torax/tests/test_data/test_iterhybrid_lh_transition.nc and b/torax/tests/test_data/test_iterhybrid_lh_transition.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_lh_transition_internal_boundary_condition.nc b/torax/tests/test_data/test_iterhybrid_lh_transition_internal_boundary_condition.nc index 30ecd390b..6a9ebfa12 100644 Binary files a/torax/tests/test_data/test_iterhybrid_lh_transition_internal_boundary_condition.nc and b/torax/tests/test_data/test_iterhybrid_lh_transition_internal_boundary_condition.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_mockup.nc b/torax/tests/test_data/test_iterhybrid_mockup.nc index 9b98bacff..60b9d5150 100644 Binary files a/torax/tests/test_data/test_iterhybrid_mockup.nc and b/torax/tests/test_data/test_iterhybrid_mockup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc index 5e273298b..e0f72f566 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc index 91651b7fb..418e88781 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc index 9a27104ba..057c135fa 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc index 63a6257ac..87cd4d015 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc index 8189af44e..fcc050fa9 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc index e360435f4..dcff6ae24 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc index 508075a18..aad786af3 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc index d353ea593..b35a2f34a 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc index f26b98c12..cb8ff68fe 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc index 2200a5008..cf29d16af 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc index 3b5787e30..a55ad9b71 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc index 7e46601df..5f58ec0c5 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc index db547207c..27aca7aba 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc index 70834ac85..bcac50f4a 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc index 385ba61c8..76e8f8e01 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc index 31a798fc6..123572080 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc index 64770b145..2b0814173 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc index 2ec473d95..40e2e8df6 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc index 9ecadbe53..8d15c2bf6 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc index 691fab3b9..3173a2abf 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_radiation_collapse.nc b/torax/tests/test_data/test_iterhybrid_radiation_collapse.nc index d4e087fb2..66842d0eb 100644 Binary files a/torax/tests/test_data/test_iterhybrid_radiation_collapse.nc and b/torax/tests/test_data/test_iterhybrid_radiation_collapse.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_rampup.nc b/torax/tests/test_data/test_iterhybrid_rampup.nc index cb9d2e54b..9b0cb4488 100644 Binary files a/torax/tests/test_data/test_iterhybrid_rampup.nc and b/torax/tests/test_data/test_iterhybrid_rampup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc b/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc index 8277b79e3..a3a639c8b 100644 Binary files a/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc and b/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc differ diff --git a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc index 9d59c6f84..5bc83da36 100644 Binary files a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc and b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc differ diff --git a/torax/tests/test_data/test_ne_qlknn_defromchie.nc b/torax/tests/test_data/test_ne_qlknn_defromchie.nc index 492d5607a..d67bec146 100644 Binary files a/torax/tests/test_data/test_ne_qlknn_defromchie.nc and b/torax/tests/test_data/test_ne_qlknn_defromchie.nc differ diff --git a/torax/tests/test_data/test_particle_sources_cgm.nc b/torax/tests/test_data/test_particle_sources_cgm.nc index 07e7df48e..6ae0b1d6f 100644 Binary files a/torax/tests/test_data/test_particle_sources_cgm.nc and b/torax/tests/test_data/test_particle_sources_cgm.nc differ diff --git a/torax/tests/test_data/test_prescribed_generic_current_source.nc b/torax/tests/test_data/test_prescribed_generic_current_source.nc index 004a46c09..d64273822 100644 Binary files a/torax/tests/test_data/test_prescribed_generic_current_source.nc and b/torax/tests/test_data/test_prescribed_generic_current_source.nc differ diff --git a/torax/tests/test_data/test_prescribed_timedependent_ne.nc b/torax/tests/test_data/test_prescribed_timedependent_ne.nc index 8679b8154..9d42bd21d 100644 Binary files a/torax/tests/test_data/test_prescribed_timedependent_ne.nc and b/torax/tests/test_data/test_prescribed_timedependent_ne.nc differ diff --git a/torax/tests/test_data/test_prescribed_transport.nc b/torax/tests/test_data/test_prescribed_transport.nc index 2fe8a36e3..e34f1bb92 100644 Binary files a/torax/tests/test_data/test_prescribed_transport.nc and b/torax/tests/test_data/test_prescribed_transport.nc differ diff --git a/torax/tests/test_data/test_psi_and_heat.nc b/torax/tests/test_data/test_psi_and_heat.nc index ed9b3ed79..bd2264e0c 100644 Binary files a/torax/tests/test_data/test_psi_and_heat.nc and b/torax/tests/test_data/test_psi_and_heat.nc differ diff --git a/torax/tests/test_data/test_psi_heat_dens.nc b/torax/tests/test_data/test_psi_heat_dens.nc index 9ed481a73..8f5af5487 100644 Binary files a/torax/tests/test_data/test_psi_heat_dens.nc and b/torax/tests/test_data/test_psi_heat_dens.nc differ diff --git a/torax/tests/test_data/test_psichease_ip_chease_vloop.nc b/torax/tests/test_data/test_psichease_ip_chease_vloop.nc index 70f1eba31..48ca3c4b2 100644 Binary files a/torax/tests/test_data/test_psichease_ip_chease_vloop.nc and b/torax/tests/test_data/test_psichease_ip_chease_vloop.nc differ diff --git a/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.nc b/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.nc index a5953d7f7..b5971d3c9 100644 Binary files a/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.nc and b/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.nc differ diff --git a/torax/tests/test_data/test_psichease_prescribed_johm.nc b/torax/tests/test_data/test_psichease_prescribed_johm.nc index f1ee269ea..84f0bac58 100644 Binary files a/torax/tests/test_data/test_psichease_prescribed_johm.nc and b/torax/tests/test_data/test_psichease_prescribed_johm.nc differ diff --git a/torax/tests/test_data/test_psichease_prescribed_jtot.nc b/torax/tests/test_data/test_psichease_prescribed_jtot.nc index f6bdadfd2..30a15fbec 100644 Binary files a/torax/tests/test_data/test_psichease_prescribed_jtot.nc and b/torax/tests/test_data/test_psichease_prescribed_jtot.nc differ diff --git a/torax/tests/test_data/test_psichease_prescribed_jtot_vloop.nc b/torax/tests/test_data/test_psichease_prescribed_jtot_vloop.nc index 91c8b0587..e1e9e5e95 100644 Binary files a/torax/tests/test_data/test_psichease_prescribed_jtot_vloop.nc and b/torax/tests/test_data/test_psichease_prescribed_jtot_vloop.nc differ diff --git a/torax/tests/test_data/test_semiimplicit_convection.nc b/torax/tests/test_data/test_semiimplicit_convection.nc index 2032dc654..621a5d7b2 100644 Binary files a/torax/tests/test_data/test_semiimplicit_convection.nc and b/torax/tests/test_data/test_semiimplicit_convection.nc differ diff --git a/torax/tests/test_data/test_step_flattop_bgb.nc b/torax/tests/test_data/test_step_flattop_bgb.nc index e03592b78..99e99f3bd 100644 Binary files a/torax/tests/test_data/test_step_flattop_bgb.nc and b/torax/tests/test_data/test_step_flattop_bgb.nc differ diff --git a/torax/tests/test_data/test_timedependence.nc b/torax/tests/test_data/test_timedependence.nc index 60eed33a9..7aa4853b5 100644 Binary files a/torax/tests/test_data/test_timedependence.nc and b/torax/tests/test_data/test_timedependence.nc differ