Skip to content

Re-numba physics and bootstrap functions #4132

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
timothy-nunn wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Re-numba physics and bootstrap functions #4132

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
26f3379
Add numba jit compiling to bootstrap current functions
timothy-nunn Mar 19, 2026
203a52d
Add numba jit compiling to physics functions
timothy-nunn Mar 19, 2026
0a791b3
Un-numba function which divides by 0 on first process iteration
timothy-nunn Mar 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions process/models/physics/bootstrap_current.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import logging
from enum import IntEnum

import numba as nb
import numpy as np
import scipy
import scipy.integrate as integrate
Expand Down Expand Up @@ -255,6 +256,7 @@ def get_bootstrap_current_fraction_value(
return model_map[model]

@staticmethod
@nb.njit(cache=True)
def bootstrap_fraction_iter89(
aspect: float,
beta: float,
Expand Down Expand Up @@ -320,6 +322,7 @@ def bootstrap_fraction_iter89(
return c_bs * (betapbs / np.sqrt(aspect)) ** 1.3

@staticmethod
@nb.njit(cache=True)
def bootstrap_fraction_wilson(
alphaj: float,
alphap: float,
Expand Down Expand Up @@ -442,6 +445,7 @@ def bootstrap_fraction_wilson(
return seps1 * betpth * (a * b).sum()

@staticmethod
@nb.njit(cache=True)
def _nevins_integral(
y: float,
nd_plasma_electrons_vol_avg: float,
Expand Down Expand Up @@ -652,6 +656,7 @@ def bootstrap_fraction_nevins(
return 1.0e6 * aibs / plasma_current

@staticmethod
@nb.njit(cache=True)
def bootstrap_fraction_sakai(
beta_poloidal: float,
q95: float,
Expand Down Expand Up @@ -714,6 +719,7 @@ def bootstrap_fraction_sakai(
)

@staticmethod
@nb.njit(cache=True)
def bootstrap_fraction_aries(
beta_poloidal: float,
ind_plasma_internal_norm: float,
Expand Down Expand Up @@ -767,6 +773,7 @@ def bootstrap_fraction_aries(
return c_bs * np.sqrt(inverse_aspect) * beta_poloidal

@staticmethod
@nb.njit(cache=True)
def bootstrap_fraction_andrade(
beta_poloidal: float,
core_pressure: float,
Expand Down Expand Up @@ -815,6 +822,7 @@ def bootstrap_fraction_andrade(
return c_bs * np.sqrt(inverse_aspect) * beta_poloidal * c_p**0.8

@staticmethod
@nb.njit(cache=True)
def bootstrap_fraction_hoang(
beta_poloidal: float,
pressure_index: float,
Expand Down Expand Up @@ -923,6 +931,7 @@ def bootstrap_fraction_wong(
return c_bs * f_peak**0.25 * beta_poloidal * np.sqrt(inverse_aspect)

@staticmethod
@nb.njit(cache=True)
def bootstrap_fraction_gi_I( # noqa: N802
beta_poloidal: float,
pressure_index: float,
Expand Down Expand Up @@ -989,6 +998,7 @@ def bootstrap_fraction_gi_I( # noqa: N802
return c_bs * np.sqrt(inverse_aspect) * beta_poloidal

@staticmethod
@nb.njit(cache=True)
def bootstrap_fraction_gi_II( # noqa: N802
beta_poloidal: float,
pressure_index: float,
Expand Down Expand Up @@ -1048,6 +1058,7 @@ def bootstrap_fraction_gi_II( # noqa: N802
return c_bs * np.sqrt(inverse_aspect) * beta_poloidal

@staticmethod
@nb.njit(cache=True)
def bootstrap_fraction_sugiyama_l_mode(
eps: float,
beta_poloidal: float,
Expand Down Expand Up @@ -1585,6 +1596,7 @@ def bootstrap_fraction_sauter(self, plasma_profile: float) -> float:
return (np.sum(da * jboot, axis=0) / physics_variables.plasma_current), jboot

@staticmethod
@nb.njit(cache=True)
def _coulomb_logarithm_sauter(
radial_elements: int, tempe: np.ndarray, ne: np.ndarray
) -> np.ndarray:
Expand Down Expand Up @@ -1704,6 +1716,7 @@ def _electron_collisionality_sauter(
)

@staticmethod
@nb.njit(cache=True)
def _ion_collisions_sauter(
radial_elements: np.ndarray,
zeff: np.ndarray,
Expand Down Expand Up @@ -2262,6 +2275,7 @@ def _calculate_l34_alpha_31_coefficient(
)

@staticmethod
@nb.njit(cache=True)
def _beta_poloidal_sauter(
radial_elements: np.ndarray,
nr: int,
Expand Down Expand Up @@ -2319,6 +2333,7 @@ def _beta_poloidal_sauter(
)

@staticmethod
@nb.njit(cache=True)
def _beta_poloidal_total_sauter(
radial_elements: np.ndarray,
nr: int,
Expand Down Expand Up @@ -2395,6 +2410,7 @@ def _beta_poloidal_total_sauter(
)

@staticmethod
@nb.njit(cache=True)
def _trapped_particle_fraction_sauter(
radial_elements: np.ndarray, triang: float, sqeps: np.ndarray, fit: int = 0
) -> np.ndarray:
Expand Down
28 changes: 28 additions & 0 deletions process/models/physics/physics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1837,6 +1837,7 @@ def plasma_composition():
# ======================================================================

@staticmethod
@nb.njit(cache=True)
def phyaux(
aspect: float,
nd_plasma_fuel_ions_vol_avg: float,
Expand Down Expand Up @@ -4142,6 +4143,7 @@ def outplas(self):
)

@staticmethod
@nb.njit(cache=True)
def calculate_plasma_masses(
m_fuel_amu: float,
m_ions_total_amu: float,
Expand Down Expand Up @@ -4669,6 +4671,7 @@ def run(self):
)

@staticmethod
@nb.njit(cache=True)
def calculate_plasma_beta(
pres_plasma: float | np.ndarray, b_field: float | np.ndarray
) -> float | np.ndarray:
Expand All @@ -4692,6 +4695,7 @@ def calculate_plasma_beta(
return 2 * constants.RMU0 * pres_plasma / (b_field**2)

@staticmethod
@nb.njit(cache=True)
def calculate_beta_norm_max_wesson(ind_plasma_internal_norm: float) -> float:
"""Calculate the Wesson normalsied beta upper limit.

Expand Down Expand Up @@ -4724,6 +4728,7 @@ def calculate_beta_norm_max_wesson(ind_plasma_internal_norm: float) -> float:
return 4 * ind_plasma_internal_norm

@staticmethod
@nb.njit(cache=True)
def calculate_beta_norm_max_original(eps: float) -> float:
"""Calculate the original scaling law normalsied beta upper limit.

Expand All @@ -4744,6 +4749,7 @@ def calculate_beta_norm_max_original(eps: float) -> float:
return 2.7 * (1.0 + 5.0 * eps**3.5)

@staticmethod
@nb.njit(cache=True)
def calculate_beta_norm_max_menard(eps: float) -> float:
"""Calculate the Menard normalsied beta upper limit.

Expand Down Expand Up @@ -4773,6 +4779,7 @@ def calculate_beta_norm_max_menard(eps: float) -> float:
return 3.12 + 3.5 * eps**1.7

@staticmethod
@nb.njit(cache=True)
def calculate_beta_norm_max_thloreus(
c_beta: float, pres_plasma_on_axis: float, pres_plasma_vol_avg: float
) -> float:
Expand Down Expand Up @@ -4816,6 +4823,7 @@ def calculate_beta_norm_max_thloreus(
)

@staticmethod
@nb.njit(cache=True)
def calculate_beta_norm_max_stambaugh(
f_c_plasma_bootstrap: float,
kappa: float,
Expand Down Expand Up @@ -4862,6 +4870,7 @@ def calculate_beta_norm_max_stambaugh(
)

@staticmethod
@nb.njit(cache=True)
def calculate_normalised_beta(
beta: float, rminor: float, c_plasma: float, b_field: float
) -> float:
Expand Down Expand Up @@ -4893,6 +4902,7 @@ def calculate_normalised_beta(
return 1.0e8 * (beta * rminor * b_field) / c_plasma

@staticmethod
@nb.njit(cache=True)
def calculate_plasma_energy_from_beta(
beta: float, b_field: float, vol_plasma: float
) -> float:
Expand All @@ -4918,6 +4928,7 @@ def calculate_plasma_energy_from_beta(
return (1.5e0 * beta * b_field**2) / (2.0e0 * constants.RMU0) * vol_plasma

@staticmethod
@nb.njit(cache=True)
def calculate_beta_limit_from_norm(
b_plasma_toroidal_on_axis: float,
beta_norm_max: float,
Expand Down Expand Up @@ -4973,6 +4984,7 @@ def calculate_beta_limit_from_norm(
)

@staticmethod
@nb.njit(cache=True)
def calculate_poloidal_beta(
b_plasma_total: float, b_plasma_poloidal_average: float, beta: float
) -> float:
Expand All @@ -4999,6 +5011,7 @@ def calculate_poloidal_beta(
return beta * (b_plasma_total / b_plasma_poloidal_average) ** 2

@staticmethod
@nb.njit(cache=True)
def fast_alpha_beta(
b_plasma_poloidal_average: float,
b_plasma_toroidal_on_axis: float,
Expand Down Expand Up @@ -5436,6 +5449,7 @@ def get_ind_internal_norm_value(self, model: IndInternalNormModel) -> float:
return model_map[model]

@staticmethod
@nb.njit(cache=True)
def calculate_volt_second_requirements(
csawth: float,
eps: float,
Expand Down Expand Up @@ -5569,6 +5583,7 @@ def calculate_volt_second_requirements(
)

@staticmethod
@nb.njit(cache=True)
def calculate_normalised_internal_inductance_iter_3(
b_plasma_poloidal_vol_avg: float,
c_plasma: float,
Expand Down Expand Up @@ -5614,6 +5629,7 @@ def calculate_normalised_internal_inductance_iter_3(
)

@staticmethod
@nb.njit(cache=True)
def calculate_internal_inductance_menard(kappa: float) -> float:
"""Calculate the Menard plasma normalized internal inductance.

Expand Down Expand Up @@ -5642,6 +5658,7 @@ def calculate_internal_inductance_menard(kappa: float) -> float:
return 3.4 - kappa

@staticmethod
@nb.njit(cache=True)
def calculate_internal_inductance_wesson(alphaj: float) -> float:
"""Calculate the Wesson plasma normalized internal inductance.

Expand Down Expand Up @@ -6101,6 +6118,7 @@ def run(self):
])

@staticmethod
@nb.njit(cache=True)
def calculate_debye_length(
temp_plasma_species_kev: float | np.ndarray,
nd_plasma_species: float | np.ndarray,
Expand All @@ -6126,6 +6144,7 @@ def calculate_debye_length(
) ** 0.5

@staticmethod
@nb.njit(cache=True)
def calculate_lorentz_factor(velocity: float | np.ndarray) -> float | np.ndarray:
"""Calculate the Lorentz factor for a given velocity.

Expand All @@ -6143,6 +6162,7 @@ def calculate_lorentz_factor(velocity: float | np.ndarray) -> float | np.ndarray
return 1 / (1 - (velocity / constants.SPEED_LIGHT) ** 2) ** 0.5

@staticmethod
@nb.njit(cache=True)
def calculate_relativistic_particle_speed(
e_kinetic: float | np.ndarray, mass: float
) -> float | np.ndarray:
Expand Down Expand Up @@ -6188,6 +6208,7 @@ def calculate_coulomb_log_from_impact(
return np.log(impact_param_max / impact_param_min)

@staticmethod
@nb.njit(cache=True)
def calculate_classical_distance_of_closest_approach(
charge1: float,
charge2: float,
Expand Down Expand Up @@ -6218,6 +6239,7 @@ def calculate_classical_distance_of_closest_approach(
)

@staticmethod
@nb.njit(cache=True)
def calculate_debroglie_wavelength(
mass: float, velocity: float | np.ndarray
) -> float | np.ndarray:
Expand All @@ -6241,6 +6263,7 @@ def calculate_debroglie_wavelength(
return (constants.PLANCK_CONSTANT / (2 * np.pi)) / (mass * velocity)

@staticmethod
@nb.njit(cache=True)
def calculate_plasma_frequency(
nd_particle: float | np.ndarray, m_particle: float, z_particle: float
) -> float | np.ndarray:
Expand Down Expand Up @@ -6270,6 +6293,7 @@ def calculate_plasma_frequency(
) / (2 * np.pi)

@staticmethod
@nb.njit(cache=True)
def calculate_larmor_frequency(
b_field: float | np.ndarray, m_particle: float, z_particle: float
) -> float | np.ndarray:
Expand All @@ -6295,6 +6319,7 @@ def calculate_larmor_frequency(
)

@staticmethod
@nb.njit(cache=True)
def calculate_upper_hybrid_frequency(
freq_plasma: float | np.ndarray, freq_larmor: float | np.ndarray
) -> float | np.ndarray:
Expand All @@ -6315,6 +6340,7 @@ def calculate_upper_hybrid_frequency(
return np.sqrt(freq_plasma**2 + freq_larmor**2)

@staticmethod
@nb.njit(cache=True)
def calculate_larmor_radius(
vel_perp: float | np.ndarray,
freq_larmor: float,
Expand All @@ -6336,6 +6362,7 @@ def calculate_larmor_radius(
return vel_perp / (freq_larmor)

@staticmethod
@nb.njit(cache=True)
def calculate_reduced_mass(mass1: float, mass2: float) -> float:
"""
Calculate the reduced mass of two particles.
Expand All @@ -6354,6 +6381,7 @@ def calculate_reduced_mass(mass1: float, mass2: float) -> float:
return (mass1 * mass2) / (mass1 + mass2)

@staticmethod
@nb.njit(cache=True)
def calculate_average_relative_velocity(
velocity_1: float | np.ndarray, velocity_2: float | np.ndarray
) -> float | np.ndarray:
Expand Down
Loading