Skip to content

Commit 8f977f9

Browse files
hamelphiTorax team
authored andcommitted
Enable fast ions by default and update sim test references
PiperOrigin-RevId: 868861998
1 parent f004b76 commit 8f977f9

75 files changed

Lines changed: 698 additions & 78 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

torax/_src/config/numerics.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ class RuntimeParams:
5353
evolve_density: bool = dataclasses.field(metadata={'static': True})
5454
exact_t_final: bool = dataclasses.field(metadata={'static': True})
5555
adaptive_dt: bool = dataclasses.field(metadata={'static': True})
56+
enable_fast_ions: bool = dataclasses.field(metadata={'static': True})
5657

5758
@functools.cached_property
5859
def evolving_names(self) -> tuple[str, ...]:
@@ -128,6 +129,7 @@ class Numerics(torax_pydantic.BaseModelFrozen):
128129
evolve_electron_heat: Annotated[bool, torax_pydantic.JAX_STATIC] = True
129130
evolve_current: Annotated[bool, torax_pydantic.JAX_STATIC] = False
130131
evolve_density: Annotated[bool, torax_pydantic.JAX_STATIC] = False
132+
enable_fast_ions: Annotated[bool, torax_pydantic.JAX_STATIC] = True
131133
resistivity_multiplier: torax_pydantic.TimeVaryingScalar = (
132134
torax_pydantic.ValidatedDefault(1.0)
133135
)
@@ -188,4 +190,5 @@ def build_runtime_params(self, t: chex.Numeric) -> RuntimeParams:
188190
evolve_density=self.evolve_density,
189191
exact_t_final=self.exact_t_final,
190192
adaptive_dt=self.adaptive_dt,
193+
enable_fast_ions=self.enable_fast_ions,
191194
)

torax/_src/core_profiles/initialization.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
from torax._src.neoclassical import neoclassical_models as neoclassical_models_lib
3434
from torax._src.neoclassical.bootstrap_current import base as bootstrap_current_base
3535
from torax._src.physics import psi_calculations
36+
from torax._src.sources import source as source_lib
3637
from torax._src.sources import source_models as source_models_lib
3738
from torax._src.sources import source_profile_builders
3839
from torax._src.sources import source_profiles as source_profiles_lib
@@ -97,6 +98,11 @@ def initial_core_profiles(
9798
face_centers=geo.rho_face_norm,
9899
)
99100

101+
fast_ions_list = []
102+
if runtime_params.numerics.enable_fast_ions:
103+
for s in source_models.standard_sources.values():
104+
fast_ions_list.extend(s.zero_fast_ions(geo))
105+
100106
core_profiles = state.CoreProfiles(
101107
T_i=T_i,
102108
T_e=T_e,
@@ -114,6 +120,7 @@ def initial_core_profiles(
114120
A_impurity_face=ions.A_impurity_face,
115121
Z_eff=ions.Z_eff,
116122
Z_eff_face=ions.Z_eff_face,
123+
fast_ions=tuple(fast_ions_list),
117124
psi=psi,
118125
psidot=psidot,
119126
q_face=jnp.zeros_like(geo.rho_face, dtype=jax_utils.get_dtype()),

torax/_src/core_profiles/tests/convertors_test.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ def setUp(self):
8787
toroidal_angular_velocity=mock.ANY,
8888
charge_state_info=mock.ANY,
8989
charge_state_info_face=mock.ANY,
90+
fast_ions=mock.ANY,
9091
)
9192

9293
def test_core_profiles_to_solver_x_tuple(self):

torax/_src/core_profiles/updaters.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,7 @@ def update_core_and_source_profiles_after_step(
200200
toroidal_angular_velocity=updated_core_profiles_t_plus_dt.toroidal_angular_velocity,
201201
charge_state_info=ions.charge_state_info,
202202
charge_state_info_face=ions.charge_state_info_face,
203+
fast_ions=core_profiles_t_plus_dt.fast_ions,
203204
)
204205

205206
conductivity = neoclassical_models.conductivity.calculate_conductivity(
@@ -224,6 +225,13 @@ def update_core_and_source_profiles_after_step(
224225
conductivity=conductivity,
225226
)
226227

228+
if runtime_params_t_plus_dt.numerics.enable_fast_ions:
229+
intermediate_core_profiles = dataclasses.replace(
230+
intermediate_core_profiles,
231+
# Flatten the tuple of tuples.
232+
fast_ions=sum(total_source_profiles.fast_ions.values(), ()),
233+
)
234+
227235
if (
228236
not runtime_params_t_plus_dt.numerics.evolve_current
229237
and runtime_params_t_plus_dt.profile_conditions.psidot is not None

torax/_src/imas_tools/output/equilibrium.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -120,9 +120,7 @@ def torax_state_to_imas_equilibrium(
120120
eq.profiles_1d.gm2 = gm2
121121

122122
# Quantities useful for coupling with equilibrium code
123-
eq.profiles_1d.pressure = (
124-
sim_state.core_profiles.pressure_thermal_total.face_value()
125-
)
123+
eq.profiles_1d.pressure = sim_state.core_profiles.pressure_total.face_value()
126124
eq.profiles_1d.dpressure_dpsi = post_processed_outputs.pprime
127125

128126
# <j.B>/B_0, could be useful to calculate and use instead of FF'

torax/_src/output_tools/output.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -568,11 +568,12 @@ def _save_core_profiles(
568568
if attr_name == "impurity_fractions":
569569
continue
570570

571-
# Skip charge_state_info since it is not needed in the output.
571+
# Skip attributes that are not needed in the output.
572572
if attr_name in (
573573
"charge_state_info",
574574
"charge_state_info_face",
575575
"impurity_density_scaling",
576+
"fast_ions",
576577
):
577578
continue
578579

@@ -647,6 +648,22 @@ def _save_core_profiles(
647648
coords={MAIN_ION: main_ions, TIME: self.times},
648649
name="main_ion_fractions",
649650
)
651+
# Handle fast ions
652+
first_fast_ions = self.core_profiles[0].fast_ions
653+
for i, first_fi in enumerate(first_fast_ions):
654+
source_key = f"{first_fi.source}_{first_fi.species}"
655+
n_data = np.stack([
656+
cp.fast_ions[i].n.cell_plus_boundaries() for cp in self.core_profiles
657+
])
658+
T_data = np.stack([
659+
cp.fast_ions[i].T.cell_plus_boundaries() for cp in self.core_profiles
660+
])
661+
xr_dict[f"n_fast_ion_{source_key}"] = self._pack_into_data_array(
662+
f"n_fast_ion_{source_key}", n_data
663+
)
664+
xr_dict[f"T_fast_ion_{source_key}"] = self._pack_into_data_array(
665+
f"T_fast_ion_{source_key}", T_data
666+
)
650667

651668
return xr_dict
652669

torax/_src/output_tools/post_processing.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -964,7 +964,7 @@ def cumulative_values():
964964
Z_eff_face=sim_state.core_profiles.Z_eff_face,
965965
Z_i_face=sim_state.core_profiles.Z_i_face,
966966
toroidal_angular_velocity=sim_state.core_profiles.toroidal_angular_velocity,
967-
pressure_thermal_i=sim_state.core_profiles.pressure_thermal_i,
967+
pressure_total_i=sim_state.core_profiles.pressure_total_i,
968968
geo=sim_state.geometry,
969969
poloidal_velocity_multiplier=runtime_params.neoclassical.poloidal_velocity_multiplier,
970970
)

torax/_src/physics/fast_ions.py

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
# Copyright 2024 DeepMind Technologies Limited
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Fast ion physics classes."""
16+
17+
import dataclasses
18+
19+
import jax
20+
from jax import numpy as jnp
21+
from torax._src import constants
22+
from torax._src import math_utils
23+
from torax._src.fvm import cell_variable
24+
25+
26+
# pylint: disable=invalid-name
27+
@jax.tree_util.register_dataclass
28+
@dataclasses.dataclass(frozen=True)
29+
class FastIon:
30+
"""State of a fast ion species.
31+
32+
Attributes:
33+
species: Species name (e.g. 'He3').
34+
source: Source name (e.g. 'ICRH').
35+
n: Density [m^-3].
36+
T: Temperature [keV].
37+
"""
38+
39+
species: str = dataclasses.field(metadata={'static': True})
40+
source: str = dataclasses.field(metadata={'static': True})
41+
n: cell_variable.CellVariable
42+
T: cell_variable.CellVariable
43+
44+
45+
def bimaxwellian_split(
46+
power_deposition: jax.Array,
47+
T_e: jax.Array,
48+
n_e: jax.Array,
49+
T_i: jax.Array,
50+
minority_concentration: jax.Array | float,
51+
P_total_W: float,
52+
charge_number: int,
53+
mass_number: float,
54+
) -> tuple[jax.Array, jax.Array]:
55+
"""Returns (n_tail, T_tail) using the Power Balance Closure.
56+
57+
Splits a minority species density into a bulk thermal component and a
58+
high-energy tail component based on Stix theory power balance.
59+
60+
Args:
61+
power_deposition: Power deposition profile [MW/m^3 / MW_in]. Normalized per
62+
MW of input power.
63+
T_e: Electron temperature profile [keV].
64+
n_e: Electron density profile [m^-3].
65+
T_i: Ion temperature profile [keV].
66+
minority_concentration: Minority species fractional concentration
67+
(n_minority/n_e).
68+
P_total_W: Total absolute power absorbed [W].
69+
charge_number: Charge number of the minority species (e.g. 2 for He3).
70+
mass_number: Mass number of the minority species (e.g. 3.016 for He3).
71+
72+
Returns:
73+
Tuple containing:
74+
n_tail: Density of the fast tail component [m^-3].
75+
T_tail: Temperature of the fast tail component [keV].
76+
"""
77+
keV_to_J = constants.CONSTANTS.keV_to_J
78+
79+
n_total = n_e * minority_concentration
80+
81+
# Calculate T_tail (The high-energy slope from Stix)
82+
# power_deposition is normalized (per MW input).
83+
# Calculate absolute power density in MW/m^3.
84+
# P_total_W is in Watts, convert to MW: P_total_W / 1e6
85+
# Result p_dens_mw is in MW/m^3.
86+
p_dens_mw = power_deposition * (P_total_W / 1.0e6)
87+
88+
n_e20 = n_e / 1.0e20
89+
90+
# Stix parameter xi [Stix, Nucl. Fusion 15, 737 (1975)].
91+
# xi = (0.24 * sqrt(T_e) * A_f * p_dens) / (n_e20^2 * Z_f^2 * c_fast)
92+
xi = (0.24 * jnp.sqrt(T_e) * mass_number * p_dens_mw) / (
93+
n_e20**2 * charge_number**2 * minority_concentration
94+
)
95+
96+
T_tail = T_e * (1.0 + xi)
97+
T_bulk = T_i
98+
99+
# Spitzer slowing-down time on electrons:
100+
# tau_s = 6.27e8 * A_f * T_e[eV]^1.5 / (Z_f^2 * n_e[cm^-3] * ln_lambda)
101+
# Converting to T_e[keV] and n_e20 [1e20 m^-3] with ln_lambda ~ 15:
102+
# 6.27e8 * 1000^1.5 / 1e14 / 15 ~ 0.013.
103+
tau_s = (0.013 * mass_number * T_e**1.5) / (charge_number**2 * n_e20)
104+
105+
# Solve for n_tail using Power Balance:
106+
# P_abs (W/m^3) = 1.5 * n_tail * (T_tail - T_bulk) * e / tau_s
107+
p_abs_w = p_dens_mw * 1.0e6 # MW/m^3 -> W/m^3
108+
109+
dT_joules = (T_tail - T_bulk) * keV_to_J
110+
111+
n_tail = math_utils.safe_divide(p_abs_w * tau_s, 1.5 * dT_joules)
112+
113+
# Constraints and Particle Conservation
114+
# Tail density cannot exceed 99% of total He3
115+
n_tail = jnp.clip(n_tail, 0.0, n_total * 0.99)
116+
117+
return n_tail, T_tail

torax/_src/physics/formulas.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ def calc_pprime(
6565
with respect to the normalized toroidal flux coordinate, on the face grid.
6666
"""
6767

68-
p_total_face = core_profiles.pressure_thermal_total.face_value()
68+
p_total_face = core_profiles.pressure_total.face_value()
6969
psi = core_profiles.psi.face_value()
7070
n_e = core_profiles.n_e.face_value()
7171
n_i = core_profiles.n_i.face_value()
@@ -87,6 +87,11 @@ def calc_pprime(
8787
+ dni_drhon * T_i
8888
+ dnimp_drhon * T_i
8989
)
90+
for fi in core_profiles.fast_ions:
91+
dptot_drhon += constants.CONSTANTS.keV_to_J * (
92+
fi.n.face_value() * fi.T.face_grad()
93+
+ fi.n.face_grad() * fi.T.face_value()
94+
)
9095

9196
# Calculate on-axis value with L'Hôpital's rule using 2nd order forward
9297
# difference approximation for second derivative at edge.
@@ -232,7 +237,7 @@ def calculate_betas(
232237
Tuple of beta_tor, beta_pol, and beta_N
233238
"""
234239
p_total_volume_avg = math_utils.volume_average(
235-
core_profiles.pressure_thermal_total.value, geo
240+
core_profiles.pressure_total.value, geo
236241
)
237242

238243
magnetic_pressure_on_axis = geo.B_0**2 / (2 * constants.CONSTANTS.mu_0)

torax/_src/physics/rotation.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525

2626
# pylint: disable=invalid-name
2727
def _calculate_radial_electric_field(
28-
pressure_thermal_i: cell_variable.CellVariable,
28+
pressure_total_i: cell_variable.CellVariable,
2929
toroidal_angular_velocity: cell_variable.CellVariable,
3030
poloidal_velocity: cell_variable.CellVariable,
3131
n_i: cell_variable.CellVariable,
@@ -39,7 +39,7 @@ def _calculate_radial_electric_field(
3939
Er = (1 / (Zi * e * ni)) * dpi/dr - v_phi * B_theta + v_theta * B_phi
4040
4141
Args:
42-
pressure_thermal_i: Pressure profile as a cell variable.
42+
pressure_total_i: Ion pressure profile (thermal + fast) as a cell variable.
4343
toroidal_angular_velocity: Toroidal velocity profile as a cell variable.
4444
poloidal_velocity: Poloidal velocity profile as a cell variable.
4545
n_i: Main ion density profile as a cell variable.
@@ -52,7 +52,7 @@ def _calculate_radial_electric_field(
5252
Er: Radial electric field [V/m] on the cell grid.
5353
"""
5454
# Calculate dpi/dr with respect to a midplane-averaged radial coordinate.
55-
dpi_dr = pressure_thermal_i.face_grad(
55+
dpi_dr = pressure_total_i.face_grad(
5656
x=geo.r_mid, x_left=geo.r_mid_face[0], x_right=geo.r_mid_face[-1]
5757
)
5858

@@ -90,7 +90,7 @@ def calculate_rotation(
9090
Z_eff_face: array_typing.FloatVectorFace,
9191
Z_i_face: array_typing.FloatVector,
9292
toroidal_angular_velocity: cell_variable.CellVariable,
93-
pressure_thermal_i: cell_variable.CellVariable,
93+
pressure_total_i: cell_variable.CellVariable,
9494
geo: geometry.Geometry,
9595
poloidal_velocity_multiplier: array_typing.FloatScalar = 1.0,
9696
) -> tuple[
@@ -108,7 +108,7 @@ def calculate_rotation(
108108
Z_eff_face: Effective charge on the face grid.
109109
Z_i_face: Main ion charge on the face grid.
110110
toroidal_angular_velocity: Toroidal velocity profile as a cell variable.
111-
pressure_thermal_i: Pressure profile as a cell variable.
111+
pressure_total_i: Total ion pressure (thermal + fast) as a cell variable.
112112
geo: Geometry object.
113113
poloidal_velocity_multiplier: A multiplier to apply to the poloidal
114114
velocity.
@@ -143,7 +143,7 @@ def calculate_rotation(
143143
)
144144

145145
Er = _calculate_radial_electric_field(
146-
pressure_thermal_i=pressure_thermal_i,
146+
pressure_total_i=pressure_total_i,
147147
toroidal_angular_velocity=toroidal_angular_velocity,
148148
poloidal_velocity=poloidal_velocity,
149149
n_i=n_i,

0 commit comments

Comments
 (0)