"""Flux-surface neoclassical transport comparison for an XGC Simulation.
The analytical curves follow the active Hinton--Hazeltine, Chang--Hinton and
Sauter terms in ``IDL_stuff/{neo_flux,bootstrap_formulas}.pro``. In particular,
the Sauter bootstrap term contains pressure and temperature gradients only;
the loop-voltage-driven Ohmic term is calculated separately. The XGC curves
use the nonadiabatic magnetic-drift and ExB moment channels, as in the IDL
diagnostic. Resolved ``f3d`` profiles are toroidally averaged first; its
axisymmetric ``_n0_df`` transport channels correspond to the ``f2d`` ``_df``
channels. Turbulent ``_turb_df`` transport is excluded from this comparison.
The local formulas are evaluated on closed surfaces with ``psi_N < 1`` and
should be interpreted cautiously in a steep pedestal.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable
import numpy as np
from scipy.ndimage import gaussian_filter1d
from xgc_analysis.data_access.catalog import read_static_variables
E_CHARGE = 1.602176634e-19
CGS_CHARGE = 4.80320425e-10
[docs]
@dataclass
class NeoResult:
"""Time-averaged XGC and model profiles on the outer-midplane grid.
Attributes
----------
profiles : dict[str, ndarray]
Named radial profiles; units and descriptions are in ``PROFILE_INFO``.
gstep_range, time_range_s : tuple
Inclusive diagnostic gsteps and physics-time extent in seconds.
sample_count : int
Number of moment samples in the average.
source : str
F2D or F3D product used for the measured moments.
smoothing_sigma_surfaces, denominator_smoothing_sigma_surfaces : float
Gaussian widths in surface-index units.
current_drive : bool
Whether loop-voltage contributions were requested.
fsa : str
Surface-averaging convention, ``volume`` or ``line``.
"""
profiles: dict[str, np.ndarray]
gstep_range: tuple[int, int]
time_range_s: tuple[float, float]
sample_count: int
source: str
smoothing_sigma_surfaces: float = 0.0
denominator_smoothing_sigma_surfaces: float = 0.0
current_drive: bool = False
fsa: str = "volume"
# Every written profile has its own XGC-style description and units attributes.
PROFILE_INFO = {
"psi_pol_norm": ("Poloidal flux divided by its value at the magnetic X point.", "1"),
"n_e": ("Flux-surface-averaged electron density, smoothed after time averaging.", "m^-3"),
"T_i": ("Flux-surface-averaged ion scalar temperature (2 T_perp + T_parallel)/3.", "eV"),
"T_e": ("Flux-surface-averaged electron scalar temperature (2 T_perp + T_parallel)/3.", "eV"),
"jpar_plasma": ("XGC total parallel plasma current, (j_parallel_i+j_parallel_e) B/B_axis, flux-surface averaged.", "A m^-2"),
"jpar_plasma_i": ("Ion contribution to XGC parallel plasma current in B/B_axis normalization.", "A m^-2"),
"jpar_plasma_e": ("Electron contribution to XGC parallel plasma current in B/B_axis normalization.", "A m^-2"),
"jpar_efit": ("MHD-equilibrium parallel current from xgc.current_drive.bp, or legacy xgc.bfield.bp, multiplied by B/B_axis before flux-surface averaging to match XGC current normalization.", "A m^-2"),
"ptl_flux_i_gc": ("Ion guiding-center particle flux: nonadiabatic magnetic-drift plus ExB moment channels, divided by <|grad psi|>.", "m^-2 s^-1"),
"ptl_flux_e": ("Electron particle flux: nonadiabatic magnetic-drift plus ExB moment channels, divided by <|grad psi|>.", "m^-2 s^-1"),
"enflux_i_gc": ("Ion kinetic-energy flux from nonadiabatic moment channels, divided by <|grad psi|>.", "W m^-2"),
"enflux_e": ("Electron kinetic-energy flux from nonadiabatic moment channels, divided by <|grad psi|>.", "W m^-2"),
"heatflux_i": ("Ion heat flux Q_i - (5/2) e T_i Gamma_i; omits electrostatic potential-energy flux as in the IDL nonadiabatic comparison.", "W m^-2"),
"heatflux_e": ("Electron heat flux Q_e - (5/2) e T_e Gamma_e; omits electrostatic potential-energy flux as in the IDL nonadiabatic comparison.", "W m^-2"),
"jboot_sauter": ("Sauter 1999 bootstrap current in the IDL B/B_axis normalization; excludes inductive current.", "A m^-2"),
"jpar_ohmic": ("Estimated inductive current: Sauter parallel conductivity times <E_parallel B/B_axis> from the loop-voltage diagnostic.", "A m^-2"),
"jpar_plasma_hh": ("Hinton-Hazeltine parallel current estimate including its inductive contribution when loop voltage is enabled.", "A m^-2"),
"jpar_ohmic_hh": ("Hinton-Hazeltine inductive contribution to parallel current.", "A m^-2"),
"ptl_flux_HH": ("Hinton-Hazeltine electron particle flux, including loop-voltage coupling when enabled.", "m^-2 s^-1"),
"heatflux_i_CH": ("Chang-Hinton 1982 ion heat flux.", "W m^-2"),
"heatflux_e_HH": ("Hinton-Hazeltine electron heat flux, including loop-voltage coupling when enabled.", "W m^-2"),
"D_eff": ("Effective electron particle diffusivity Gamma_e / (-d n_e/d r_minor), with separately smoothed denominator profile.", "m^2 s^-1"),
"chi_i": ("Effective ion heat conductivity q_i / (-e n_e dT_i/d r_minor), with separately smoothed denominator profiles.", "m^2 s^-1"),
"chi_e": ("Effective electron heat conductivity q_e / (-e n_e dT_e/d r_minor), with separately smoothed denominator profiles.", "m^2 s^-1"),
"D_HH": ("Effective diffusivity from Hinton-Hazeltine electron particle flux and the smoothed electron-density gradient.", "m^2 s^-1"),
"chi_i_ch": ("Chang-Hinton ion heat conductivity evaluated from the local analytical formula.", "m^2 s^-1"),
"chi_e_hh": ("Effective conductivity from Hinton-Hazeltine electron heat flux and the smoothed electron-temperature gradient.", "m^2 s^-1"),
"jpar_plasma_sigma": ("Temporal standard deviation of XGC total parallel plasma current before radial smoothing.", "A m^-2"),
"jpar_plasma_i_sigma": ("Temporal standard deviation of XGC ion parallel current before radial smoothing.", "A m^-2"),
"jpar_plasma_e_sigma": ("Temporal standard deviation of XGC electron parallel current before radial smoothing.", "A m^-2"),
"ptl_flux_i_gc_sigma": ("Temporal standard deviation of the ion particle flux before radial smoothing.", "m^-2 s^-1"),
"ptl_flux_e_sigma": ("Temporal standard deviation of the electron particle flux before radial smoothing.", "m^-2 s^-1"),
"enflux_i_gc_sigma": ("Temporal standard deviation of the ion energy flux before radial smoothing.", "W m^-2"),
"enflux_e_sigma": ("Temporal standard deviation of the electron energy flux before radial smoothing.", "W m^-2"),
"heatflux_i_sigma": ("Temporal standard deviation of the ion heat flux before radial smoothing.", "W m^-2"),
"heatflux_e_sigma": ("Temporal standard deviation of the electron heat flux before radial smoothing.", "W m^-2"),
"D_eff_sigma": ("Temporal standard deviation of electron diffusivity evaluated separately at each time sample.", "m^2 s^-1"),
"chi_i_sigma": ("Temporal standard deviation of ion heat conductivity evaluated separately at each time sample.", "m^2 s^-1"),
"chi_e_sigma": ("Temporal standard deviation of electron heat conductivity evaluated separately at each time sample.", "m^2 s^-1"),
"omega_tor_drift": ("Ion toroidal angular frequency from XGC drift flow.", "s^-1"),
"omega_tor_vpar": ("Ion toroidal angular frequency from parallel flow.", "s^-1"),
"omega_tor_exb": ("Ion toroidal angular frequency from radial electric-field ExB flow.", "s^-1"),
"omega_tor_dia": ("Ion toroidal angular frequency from diamagnetic flow.", "s^-1"),
"omega_pol_drift": ("Ion poloidal angular frequency from XGC drift flow.", "s^-1"),
"omega_pol_vpar": ("Ion poloidal angular frequency from parallel flow.", "s^-1"),
"omega_pol_exb": ("Ion poloidal angular frequency from radial electric-field ExB flow.", "s^-1"),
"omega_pol_dia": ("Ion poloidal angular frequency from diamagnetic flow.", "s^-1"),
}
def _smooth(values: np.ndarray, sigma: float) -> np.ndarray:
"""Apply radial Gaussian smoothing in surface-index units; zero is identity."""
return gaussian_filter1d(values, sigma, axis=0, mode="nearest") if sigma > 0 else np.asarray(values)
def _smooth_core(values: np.ndarray, core: np.ndarray, sigma: float) -> np.ndarray:
"""Smooth an analytical profile within closed surfaces and leave the SOL undefined."""
result = np.full(np.shape(values), np.nan, dtype=float)
result[core] = _smooth(np.asarray(values)[core], sigma)
return result
def _div(num: np.ndarray, den: np.ndarray) -> np.ndarray:
"""Divide profiles while marking singular or invalid denominators as NaN."""
with np.errstate(divide="ignore", invalid="ignore"):
result = np.asarray(num, dtype=float) / np.asarray(den, dtype=float)
return np.where(np.isfinite(result), result, np.nan)
def _gradient(y: np.ndarray, x: np.ndarray) -> np.ndarray:
"""Differentiate along the surface axis with the actual radial coordinate."""
return np.gradient(y, x, axis=0, edge_order=2 if len(x) > 2 else 1)
def _poloidal_field_scale(psi: np.ndarray, psi_x: float, rmin: np.ndarray, rmaj: np.ndarray, btheta_sign: float) -> np.ndarray:
"""Evaluate the IDL Bpo factor, d(physical psi)/dr_minor divided by R."""
return _gradient(psi*psi_x, rmin)/rmaj*btheta_sign
def _read_at_steps(catalog, product: str, names: list[str], steps: list[int], n_nodes: int | None = None) -> dict[str, np.ndarray]:
"""Read moments as [node,time], averaging a resolved toroidal axis."""
available = catalog.get_product(product).variables
absent = [name for name in names if name not in available]
if absent:
raise KeyError(f"{product} lacks required diagnostics: {', '.join(absent)}")
read = catalog.read_arrays(product, names, steps=steps)
result = {}
for name in names:
samples = []
for step in steps:
value = np.squeeze(np.asarray(read.arrays[name][step]))
if name != "time" and n_nodes is not None:
if value.ndim == 2 and value.shape[1] == n_nodes:
value = value.mean(axis=0)
elif value.ndim == 2 and value.shape[0] == n_nodes:
value = value.mean(axis=1)
if value.shape != (n_nodes,):
raise ValueError(f"{product}:{name} has shape {value.shape}; expected {n_nodes} nodes with an optional toroidal axis")
samples.append(value)
result[name] = np.stack(samples, axis=-1)
return result
def _step_window(catalog, product: str, gstep_window, time_window_ms, selected_frame_index):
"""Select inclusive gstep/time windows from catalog step provenance."""
infos = catalog.available_steps(product)
if not infos:
raise ValueError(f"No samples in {product}")
if gstep_window is not None:
lo, hi = gstep_window
infos = [item for item in infos if lo <= item.logical_step <= hi]
if time_window_ms is not None:
lo, hi = time_window_ms
infos = [item for item in infos if item.selected_fragment.time is not None and lo <= 1e3 * item.selected_fragment.time <= hi]
if gstep_window is None and time_window_ms is None:
if selected_frame_index is not None:
all_infos = catalog.available_steps(product)
frame = min(max(int(selected_frame_index), 0), len(all_infos) - 1)
end_step = all_infos[frame].logical_step
infos = [item for item in infos if item.logical_step <= end_step]
infos = infos[-100:]
if not infos:
raise ValueError(f"The selected averaging window has no {product} samples")
return [item.logical_step for item in infos]
def _surface_average(plane, values: np.ndarray, *, fsa: str = "volume") -> np.ndarray:
"""Flux-surface average an [node,time] array with the selected operator."""
averaged = plane.flux_avg_to_surf(values, fsa=fsa)
if np.asarray(values).ndim == 2 and np.asarray(averaged).ndim == 1:
return averaged[:, None]
return averaged
def _gradpsi_rz(simulation, plane, magnetic_field) -> np.ndarray:
"""Return the R-Z flux gradient used by the IDL diagnostic."""
velocity_grid = getattr(simulation, "velocity_grid", None)
loaded = getattr(velocity_grid, "gradpsi", None)
if loaded is not None:
gradient = np.asarray(loaded)
if gradient.shape == (plane.n_n, 2):
return gradient
if gradient.shape == (2, plane.n_n):
return gradient.T
raise ValueError(f"simulation.velocity_grid.gradpsi has unexpected shape {gradient.shape}")
catalog = simulation.catalog
if "xgc.f0.mesh.bp" in catalog.products and "gradpsi" in catalog.get_product("xgc.f0.mesh.bp").variables:
gradient = np.asarray(read_static_variables(catalog, "xgc.f0.mesh.bp", ("gradpsi",))["gradpsi"])
if gradient.shape == (plane.n_n, 2):
return gradient
if gradient.shape == (2, plane.n_n):
return gradient.T
raise ValueError(f"xgc.f0.mesh.bp:gradpsi has unexpected shape {gradient.shape}")
if plane.basis is not None and np.any(np.asarray(plane.basis) != 1):
raise ValueError("xgc.f0.mesh.bp:gradpsi is required when gradient matrices use a flux-aligned basis")
psi_nodes = np.asarray(magnetic_field.psi_pd.data).reshape(-1)
return np.column_stack((
plane.gradient_r_psi.get_csr_matrix() @ psi_nodes,
plane.gradient_z_theta.get_csr_matrix() @ psi_nodes,
))
def _time_aligned_nodes(catalog, product: str, variable: str, times: np.ndarray, steps: list[int], n_nodes: int, *, require_coverage: bool = False) -> np.ndarray:
"""Interpolate a node-centered BP diagnostic onto the selected f2d times."""
infos = catalog.available_steps(product, variable)
if not infos:
raise ValueError(f"{product} has no {variable} samples")
timed = [item for item in infos if item.selected_fragment.time is not None]
if timed:
selected = [item.logical_step for item in timed if times[0] <= item.selected_fragment.time <= times[-1]]
preceding = [item.logical_step for item in timed if item.selected_fragment.time < times[0]]
following = [item.logical_step for item in timed if item.selected_fragment.time > times[-1]]
if preceding:
selected.insert(0, preceding[-1])
if following:
selected.append(following[0])
if not selected:
selected = [min(timed, key=lambda item: abs(item.selected_fragment.time - times[-1])).logical_step]
else:
available = {item.logical_step for item in infos}
if not all(step in available for step in steps):
raise ValueError(f"{product} lacks time coordinates and matching moment gsteps")
selected = steps
read = _read_at_steps(catalog, product, [variable, "time"], selected, n_nodes)
nodes = np.asarray(read[variable], dtype=float)
sample_times = np.asarray(read["time"], dtype=float).reshape(-1)
if require_coverage and (times[0] < sample_times[0] - 1e-12 or times[-1] > sample_times[-1] + 1e-12):
raise ValueError(f"{product}:{variable} does not cover the selected moment averaging time window")
if nodes.shape[0] != n_nodes:
raise ValueError(f"{product}:{variable} has {nodes.shape[0]} nodes; mesh has {n_nodes}")
if nodes.shape[1] == 1:
return np.repeat(nodes, len(times), axis=1)
return np.stack([np.interp(times, sample_times, node_values) for node_values in nodes])
def _loop_epara(simulation, steps: list[int], times: np.ndarray, bphi: np.ndarray, b_axis: float,
*, fsa: str = "volume") -> np.ndarray:
"""Interpolate time-resolved loop voltage and calculate <E_parallel B/B_axis>."""
catalog = simulation.catalog
plane = simulation.mesh.get_plane(0)
product = next((name for name in ("xgc.loop_vol.bp", "xgc.loop_voltage.bp") if name in catalog.products), None)
if product is None:
raise FileNotFoundError("Current drive requested but xgc.loop_vol.bp / xgc.loop_voltage.bp is absent")
matched = _time_aligned_nodes(catalog, product, "loop_voltage", times, steps, plane.n_n, require_coverage=True)
factor = bphi / (2 * np.pi * plane.rz[:, 0] * b_axis)
return _surface_average(plane, matched * factor[:, None], fsa=fsa)
def _sauter(psi: np.ndarray, psi_x: float, plane, surfaces, simulation, ni, ne, ti, te, btheta_sign, fpol, epara):
"""Evaluate active Sauter 1999 bootstrap and parallel-conductivity terms."""
r0 = plane.axis_r
eps = plane.rmin[surfaces] / r0
q = np.asarray(plane.qsafety[surfaces], dtype=float).copy()
q[0] = q[1]
ft = np.asarray(plane.trapped[surfaces], dtype=float).copy()
p_i = ni * ti * E_CHARGE
p_e = ne * te * E_CHARGE
p = p_i + p_e
d_p = _gradient(p, psi) / psi_x
d_te = _gradient(te * E_CHARGE, psi) / psi_x
d_ti = _gradient(ti * E_CHARGE, psi) / psi_x
with np.errstate(invalid="ignore", divide="ignore"):
log_e = 31.3 - np.log(np.sqrt(ne) / te)
log_i = 30 - np.log(np.sqrt(ni) / ti**1.5)
nue = 6.9121e-18 * q[:, None] * r0 * ne * log_e / (te**2 * eps[:, None]**1.5)
nui = 4.90e-18 * q[:, None] * r0 * ni * log_i / (ti**2 * eps[:, None]**1.5)
nue[0] = nue[1]
nui[0] = nui[1]
f = ft[:, None]
x31 = f / (1 + (1 - .1*f)*np.sqrt(nue) + .5*(1-f)*nue)
l31 = 1.7*x31 - .95*x31**2 + .15*x31**3 + .1*x31**4
x32 = f / (1 + .26*(1-f)*np.sqrt(nue) + .18*(1-.37*f)*nue)
y32 = f / (1 + (1+.6*f)*np.sqrt(nue) + 1.7*(1-.37*f)*nue)
f32ee = (.67/1.44)*(x32-x32**4) + (1/1.22)*(x32**2-x32**4-1.2*(x32**3-x32**4)) + .8*x32**4
f32ei = -(2.49/1.44)*(y32-y32**4) + (4.95/3.48)*(y32**2-y32**4-.55*(y32**3-y32**4)) - .8*y32**4
l32 = f32ee + f32ei
x34 = f / (1 + (1-.1*f)*np.sqrt(nue) + .5*(1-.5*f)*nue)
l34 = 1.7*x34 - .95*x34**2 + .15*x34**3 + .1*x34**4
alpha0 = -1.17*(1-f)/(1-.22*f-.19*f**2)
alpha = ((alpha0+.25*(1-f**2)*np.sqrt(nui))/(1+.5*np.sqrt(nui))+.315*nui**2*f**6)/(1+.15*nui**2*f**6)
b02 = abs(np.interp(0, simulation.magnetic_field.eq_psi_grid, simulation.magnetic_field.eq_I_array)) / r0
jboot = -btheta_sign*abs(fpol) / b02 * (l31*d_p + l32*ne*d_te + l34*alpha*ne*d_ti)
sigma = np.zeros_like(jboot)
if epara is not None:
nz = .58 + .74/1.76
sigma_sp = 1.9012e4 * te**1.5 / (nz*log_e)
x = f / (1 + (.55-.1*f)*np.sqrt(nue) + .45*(1-f)*nue)
sigma = sigma_sp*(1 - 1.36*x + .59*x**2 - .23*x**3)
return jboot, sigma * epara if epara is not None else np.zeros_like(jboot), nue, nui
def _hh_ch(psi, psi_x, plane, surfaces, ne, ti, te, gpsi, bmag, b_axis, btheta_sign, epara, ion_mass, electron_mass, ion_charge, electron_charge, *, fsa="volume"):
"""Evaluate the active Hinton--Hazeltine and Chang--Hinton IDL formulas."""
rmin = plane.rmin[surfaces]
rmaj = plane.rmaj[surfaces]
eps = rmin/rmaj
eps2 = plane.epsilon[surfaces]
q = plane.qsafety[surfaces].copy()
q[0] = q[1]
gti = _gradient(ti, rmin)
gte = _gradient(te, rmin)
gn = _gradient(ne, rmin)
pe = ne*te*E_CHARGE
pi = ne*ti*E_CHARGE
gpe = (ne*gte+te*gn)*E_CHARGE
gpi = (ne*gti+ti*gn)*E_CHARGE
with np.errstate(divide="ignore", invalid="ignore"):
ni_cm = ne/1e6
mu = ion_mass/1.667e-27
log_i = 23 - np.log(2*mu/(2*mu*ti) * np.sqrt(2*ni_cm/ti))
log_e = 23.5 - np.log(np.sqrt(ni_cm)/te**1.25) - np.sqrt(1e-5+(np.log(te)-2)**2/16)
tau_i = 1 / ((4/3)*np.sqrt(np.pi)*ni_cm*(CGS_CHARGE*ion_charge/E_CHARGE)**4*log_i / (np.sqrt(ion_mass*1e3)*(E_CHARGE*ti*1e7)**1.5))
tau_e = 1 / ((4/3)*np.sqrt(2*np.pi)*ni_cm*(CGS_CHARGE*electron_charge/E_CHARGE)**4*log_e / (np.sqrt(electron_mass*1e3)*(E_CHARGE*te*1e7)**1.5))
nui = np.sqrt(ion_mass/(E_CHARGE*ti))*q[:, None]*rmaj[:, None]/(tau_i*eps2[:, None]**1.5)
nue = np.sqrt(electron_mass/(E_CHARGE*te))*q[:, None]*rmaj[:, None]/(tau_e*eps2[:, None]**1.5)
nui[0] = nui[1]
nue[0] = nue[1]
bpo = _poloidal_field_scale(psi, psi_x, rmin, rmaj, btheta_sign)
rho_i = np.sqrt(2*ti*ion_mass*E_CHARGE)/ion_charge/abs(bpo[:, None])
rho_e = np.sqrt(2*te*electron_mass*E_CHARGE)/electron_charge/abs(bpo[:, None])
b0 = bmag[0]
bmean = _surface_average(plane, bmag, fsa=fsa)[surfaces]
bmean2 = _surface_average(plane, (b0/bmag)**2, fsa=fsa)[surfaces]
bmean3 = _surface_average(plane, (bmag/b0)**2, fsa=fsa)[surfaces]
bmean4 = _surface_average(plane, bmag**2, fsa=fsa)[surfaces]
ep = eps[:, None]
def kij(k, a, b, c):
"""Evaluate one Hinton-Hazeltine transport coupling coefficient."""
return k*(1/(1+a*np.sqrt(nue)+b*nue)+ep**3*(c*c/b)*nue/(1+c*np.sqrt(ep**3)*nue))
def kin(k, a, b, c):
"""Evaluate one Hinton-Hazeltine cross-species coefficient."""
return k/((1+a*np.sqrt(nue)+b*nue)*(1+c*np.sqrt(ep**3)*nue))
k11, k12, k22 = kij(1.04,2.01,1.53,.89), kij(1.20,.76,.67,.56), kij(2.55,.45,.43,.43)
k13, k23 = kin(2.30,1.02,1.07,1.07), kin(4.19,.57,.61,.61)
k33 = (1.95-.95*np.sqrt(ep))/((1+.68*np.sqrt(nue)+.32*nue)*(1+.66*np.sqrt(ep**3)*nue))
c = (1.17-.35*np.sqrt(nui))/(1+.7*np.sqrt(nui))
kci = (c-2.1*nui**2*ep**3)/(1+nui**2*ep**3)
a1e = gpe/pe - 2.5*gte/te + ti/te*(gpi/pi-kci/(1+nue**2*ep**2)*gti/ti)
prefac = -ne*rho_e**2*np.sqrt(ep)/tau_e
particle = prefac*(k11*a1e+k12*gte/te)
if epara is not None:
particle -= k13*ne*np.sqrt(ep)*epara/bpo[:, None]
heat_e = E_CHARGE*te*prefac*(k12*a1e+k22*gte/te)-2.5*E_CHARGE*te*particle
if epara is not None:
heat_e -= k23*ne*E_CHARGE*te*np.sqrt(ep)*epara/bpo[:, None]
jboot = -ne*E_CHARGE*te*np.sqrt(ep)/bpo[:, None]*(k13*a1e+k23*gte/te)
j_ohm = np.zeros_like(jboot)
if epara is not None:
sigma = ne*electron_charge**2*tau_e/electron_mass/(.29+.46/2.08)
j_ohm = (-k33*np.sqrt(ep)+bmean[:, None]*b0/bmean4[:, None])*sigma*epara
f1 = bmean2[:, None]
f2 = 1/bmean3[:, None]
chi_i = rho_i**2*np.sqrt(ep)/tau_i*((.66+1.88*np.sqrt(ep)-1.54*ep)/(1+1.03*np.sqrt(nui)+.31*nui)*f1 + .59*nui*ep/(1+.74*nui*ep**1.5)*(f1-f2))
heat_i = -chi_i*gti*ne*E_CHARGE
return particle, heat_e, heat_i, chi_i, jboot+j_ohm, j_ohm
[docs]
def compute_neoclassical_transport(
simulation,
*,
gstep_window: tuple[int, int] | None = None,
time_window_ms: tuple[float, float] | None = None,
selected_frame_index: int | None = None,
smoothing_sigma: float = 2.0,
denominator_smoothing_sigma: float | None = None,
current_drive: bool = False,
moment_product: str = "xgc.f2d.bp",
fsa: str = "volume",
check_cancel: Callable[[], None] | None = None,
) -> NeoResult:
"""Compute XGC and analytical neoclassical profiles for a Simulation.
Windows are inclusive. With no window, the last 100 moment samples up to the
selected frame are used. ``moment_product`` may be ``xgc.f2d.bp`` or
``xgc.f3d.bp``; resolved f3d moments are averaged over toroidal planes.
Gaussian sigmas are in surface-index units; zero
disables smoothing. The ``denominator_smoothing_sigma`` controls density
and temperature profiles independently before coefficient gradients.
The current IDL coefficient set applies to one singly charged ion species
and electrons; their masses and charges are taken from ``Simulation.species``.
Analytical profiles are NaN at and beyond the separatrix.
Parameters
----------
simulation : Simulation
Initialized run with mesh, equilibrium, species, and diagnostic catalog.
gstep_window, time_window_ms : tuple, optional
Inclusive averaging limits in gsteps and milliseconds.
selected_frame_index : int, optional
Last eligible moment frame when no window is supplied.
smoothing_sigma, denominator_smoothing_sigma : float, optional
Gaussian widths in surface indices; the latter applies separately to
density and temperature before coefficient gradients are taken.
current_drive : bool, optional
Include a loop-voltage Ohmic estimate when loop-voltage data exist.
moment_product : str, optional
``xgc.f2d.bp`` or ``xgc.f3d.bp`` axisymmetric moment channels.
fsa : str, optional
``volume`` or ``line`` flux-surface averaging weights.
check_cancel : callable, optional
Callback invoked during long calculations to permit cancellation.
Returns
-------
NeoResult
XGC and analytical surface profiles with averaging provenance.
"""
if fsa not in ("volume", "line"):
raise ValueError("fsa must be 'volume' or 'line'")
if simulation.mesh is None or simulation.magnetic_field is None:
raise ValueError("A fully initialized Simulation is required")
catalog = simulation.catalog
plane = simulation.mesh.get_plane(0)
mag = simulation.magnetic_field
ions = [sp for sp in simulation.species if sp.charge_eu > 0]
electrons = [sp for sp in simulation.species if sp.charge_eu < 0]
if len(ions) != 1 or len(electrons) != 1:
raise ValueError("Neoclassical comparison requires one ion and one electron Species")
ion, electron = ions[0], electrons[0]
if not np.isclose(ion.charge_eu, 1) or not np.isclose(electron.charge_eu, -1):
raise ValueError("The IDL Hinton-Hazeltine and Sauter formulas require singly charged ions and electrons")
if moment_product not in {"xgc.f2d.bp", "xgc.f3d.bp"}:
raise ValueError(f"Unsupported moment product: {moment_product}")
steps = _step_window(catalog, moment_product, gstep_window, time_window_ms, selected_frame_index)
is_f3d = moment_product == "xgc.f3d.bp"
names = ["time", "i_den", "e_den", "i_T_para", "i_T_perp", "e_T_para", "e_T_perp", "i_u_para", "e_u_para"]
moment_aliases = {}
for species in ("i", "e"):
for channel in ("radial_flux", "radial_en_flux", "rad_mom_flux"):
for drift in ("mag", "ExB"):
canonical = f"{species}_{channel}_{drift}_df"
source_name = f"{species}_{channel}_{drift}_n0_df" if is_f3d else canonical
names.append(source_name)
moment_aliases[canonical] = source_name
for direction in ("toroidal", "poloidal"):
for component in ("df", "f0"):
canonical = f"{species}_{direction}_flow_{component}"
source_name = f"{species}_{direction}_flow_n0_{component}" if is_f3d else canonical
names.append(source_name)
moment_aliases[canonical] = source_name
if check_cancel:
check_cancel()
f2d = _read_at_steps(catalog, moment_product, names, steps, plane.n_n)
for canonical, source_name in moment_aliases.items():
if canonical != source_name:
f2d[canonical] = f2d[source_name]
times = np.asarray(f2d["time"], dtype=float).reshape(-1)
if len(times) != len(steps):
raise ValueError(f"{moment_product} time coordinate does not match selected gsteps")
if check_cancel:
check_cancel()
surf = np.asarray(plane.surf_map, dtype=int)
psi_x = plane.x_psi
psi = plane.psi_surf[surf]/psi_x
core = psi < 1
if np.count_nonzero(core) < 3:
raise ValueError("Neoclassical formulas require at least three closed surfaces with psi_N < 1")
core_surfaces = surf[core]
bvec = np.asarray(mag.bfield.data)
bmag = np.linalg.norm(bvec, axis=1)
bphi = bvec[:, 2]
btheta = np.linalg.norm(bvec[:, :2], axis=1)
gradpsi = _gradpsi_rz(simulation, plane, mag)
poloidal_direction = -bvec[:, 0]*gradpsi[:, 1] + bvec[:, 1]*gradpsi[:, 0]
btheta_sign = np.sign(poloidal_direction[1])
if btheta_sign == 0:
raise ValueError("Cannot determine poloidal magnetic-field sign")
b_axis = abs(bphi[0])
gpsi = np.sqrt(_surface_average(plane, np.sum(gradpsi**2, axis=1), fsa=fsa))[surf]
if gpsi[0] == 0:
gpsi[0] = gpsi[1]
fs = lambda values: _surface_average(plane, values, fsa=fsa)[surf]
ni = fs(f2d["i_den"])
ne = fs(f2d["e_den"])
ti = fs((2*f2d["i_T_perp"]+f2d["i_T_para"])/3)
te = fs((2*f2d["e_T_perp"]+f2d["e_T_para"])/3)
j_i = fs(ion.charge_C*f2d["i_den"]*f2d["i_u_para"]*bmag[:, None]/b_axis)
j_e = fs(electron.charge_C*f2d["e_den"]*f2d["e_u_para"]*bmag[:, None]/b_axis)
profiles_time = {"jpar_plasma_i": j_i, "jpar_plasma_e": j_e, "jpar_plasma": j_i+j_e}
for species in ("i", "e"):
for channel, output in (("radial_flux", f"ptl_flux_{species}"), ("radial_en_flux", f"enflux_{species}"), ("rad_mom_flux", f"momflux_{species}")):
values = f2d[f"{species}_{channel}_mag_df"] + f2d[f"{species}_{channel}_ExB_df"]
profiles_time[output] = fs(values)/gpsi[:, None]
profiles_time["ptl_flux_i_gc"] = profiles_time.pop("ptl_flux_i")
profiles_time["enflux_i_gc"] = profiles_time.pop("enflux_i")
profiles_time["heatflux_i"] = profiles_time["enflux_i_gc"] - 2.5*E_CHARGE*ti*profiles_time["ptl_flux_i_gc"]
profiles_time["heatflux_e"] = profiles_time["enflux_e"] - 2.5*E_CHARGE*te*profiles_time["ptl_flux_e"]
field_product = next((name for name in ("xgc.2d.bp", "xgc.3d.bp")
if name in catalog.products and catalog.available_steps(name, "epsi")), None)
if field_product is None:
raise FileNotFoundError("xgc.2d.bp:epsi or xgc.3d.bp:epsi is required for the IDL rotation profiles")
if check_cancel:
check_cancel()
radial_e = _time_aligned_nodes(catalog, field_product, "epsi", times, steps, plane.n_n)
pressure_i = E_CHARGE*f2d["i_den"]*(2*f2d["i_T_perp"]+f2d["i_T_para"])/3
pressure_grad_r = plane.gradient_r_psi.get_csr_matrix() @ pressure_i
rnode = plane.rz[:, 0, None]
b2 = bmag[:, None]**2
bt = btheta_sign*btheta[:, None]
btor = bphi[:, None]
upar = f2d["i_u_para"]
drift_tor = f2d["i_toroidal_flow_df"] + f2d["i_toroidal_flow_f0"]
drift_pol = f2d["i_poloidal_flow_df"] + f2d["i_poloidal_flow_f0"]
omega_tor = {
"drift": drift_tor/rnode,
"vpar": upar*btor/(bmag[:, None]*rnode),
"exb": -radial_e*bt/(b2*rnode),
"dia": pressure_grad_r*bt/(f2d["i_den"]*E_CHARGE*b2*rnode),
}
omega_pol = {
"drift": drift_pol,
"vpar": upar*bt/bmag[:, None],
"exb": radial_e*btor/b2,
"dia": -pressure_grad_r*btor/(f2d["i_den"]*E_CHARGE*b2),
}
for component, value in omega_tor.items():
profiles_time[f"omega_tor_{component}"] = fs(value)
for component, value in omega_pol.items():
angular = _div(fs(value), plane.rmin[surf, None])
if plane.rmin[surf[0]] == 0:
angular[0] = angular[1]
profiles_time[f"omega_pol_{component}"] = angular
epara = _loop_epara(simulation, steps, times, bphi, b_axis, fsa=fsa)[core_surfaces] if current_drive else None
fpol = np.interp(plane.psi_surf[core_surfaces], mag.eq_psi_grid, mag.eq_I_array) * np.sign(bphi[0])
ion_mass = ion.mass_kg
jboot, j_ohm, _, _ = _sauter(psi[core], psi_x, plane, core_surfaces, simulation, ni[core], ne[core], ti[core], te[core], btheta_sign, fpol[:, None], epara)
p_hh, q_e_hh, q_i_ch, chi_i_ch, j_hh, j_ohm_hh = _hh_ch(psi[core], psi_x, plane, core_surfaces, ne[core], ti[core], te[core], gpsi[core], bmag, b_axis, btheta_sign, epara, ion_mass, electron.mass_kg, abs(ion.charge_C), abs(electron.charge_C), fsa=fsa)
analytical_series = {"jboot_sauter": jboot, "jpar_ohmic": j_ohm, "ptl_flux_HH": p_hh, "heatflux_i_CH": q_i_ch, "heatflux_e_HH": q_e_hh, "chi_i_ch": chi_i_ch, "jpar_plasma_hh": j_hh, "jpar_ohmic_hh": j_ohm_hh}
for name, core_values in analytical_series.items():
series = np.full(ne.shape, np.nan)
series[core] = core_values
profiles_time[name] = series
if check_cancel:
check_cancel()
profiles = {"psi_pol_norm": psi.copy()}
for name, series in profiles_time.items():
if name in PROFILE_INFO:
if name in analytical_series:
mean_series = np.full(psi.shape, np.nan)
mean_series[core] = np.mean(series[core], axis=1)
profiles[name] = _smooth_core(mean_series, core, smoothing_sigma)
else:
profiles[name] = _smooth(np.mean(series, axis=1), smoothing_sigma)
sigma_name = name + "_sigma"
if sigma_name in PROFILE_INFO:
profiles[sigma_name] = np.std(series, axis=1, ddof=1 if len(steps) > 1 else 0)
profiles["n_e"] = _smooth(np.mean(ne, axis=1), smoothing_sigma)
profiles["T_i"] = _smooth(np.mean(ti, axis=1), smoothing_sigma)
profiles["T_e"] = _smooth(np.mean(te, axis=1), smoothing_sigma)
ef_nodes = np.asarray(mag.jpar_bg_pd.data).reshape(-1)
if ef_nodes.size != plane.n_n:
raise ValueError(f"EFIT current has {ef_nodes.size} nodes; mesh has {plane.n_n}")
profiles["jpar_efit"] = _surface_average(plane, ef_nodes*bmag/b_axis, fsa=fsa)[surf]
denominator_sigma = smoothing_sigma if denominator_smoothing_sigma is None else denominator_smoothing_sigma
density = _smooth(np.mean(ne, axis=1), denominator_sigma)
ion_temp = _smooth(np.mean(ti, axis=1), denominator_sigma)
electron_temp = _smooth(np.mean(te, axis=1), denominator_sigma)
rminor = plane.rmin[surf]
dn = _smooth(_gradient(density, rminor), denominator_sigma)
dti = _smooth(_gradient(ion_temp, rminor), denominator_sigma)
dte = _smooth(_gradient(electron_temp, rminor), denominator_sigma)
profiles["D_eff"] = _smooth(_div(profiles["ptl_flux_e"], -dn), smoothing_sigma)
profiles["chi_i"] = _smooth(_div(profiles["heatflux_i"], -E_CHARGE*density*dti), smoothing_sigma)
profiles["chi_e"] = _smooth(_div(profiles["heatflux_e"], -E_CHARGE*density*dte), smoothing_sigma)
density_time = _smooth(ne, denominator_sigma)
ion_temp_time = _smooth(ti, denominator_sigma)
electron_temp_time = _smooth(te, denominator_sigma)
dn_time = _smooth(_gradient(density_time, rminor), denominator_sigma)
dti_time = _smooth(_gradient(ion_temp_time, rminor), denominator_sigma)
dte_time = _smooth(_gradient(electron_temp_time, rminor), denominator_sigma)
coefficient_series = {
"D_eff_sigma": _smooth(_div(_smooth(profiles_time["ptl_flux_e"], smoothing_sigma), -dn_time), smoothing_sigma),
"chi_i_sigma": _smooth(_div(_smooth(profiles_time["heatflux_i"], smoothing_sigma), -E_CHARGE*density_time*dti_time), smoothing_sigma),
"chi_e_sigma": _smooth(_div(_smooth(profiles_time["heatflux_e"], smoothing_sigma), -E_CHARGE*density_time*dte_time), smoothing_sigma),
}
for name, series in coefficient_series.items():
profiles[name] = np.std(series, axis=1, ddof=1 if len(steps) > 1 else 0)
density_core = _smooth(np.mean(ne[core], axis=1), denominator_sigma)
electron_temp_core = _smooth(np.mean(te[core], axis=1), denominator_sigma)
dn_core = _smooth(_gradient(density_core, rminor[core]), denominator_sigma)
dte_core = _smooth(_gradient(electron_temp_core, rminor[core]), denominator_sigma)
profiles["D_HH"] = np.full(psi.shape, np.nan)
profiles["D_HH"][core] = _smooth(_div(profiles["ptl_flux_HH"][core], -dn_core), smoothing_sigma)
profiles["chi_e_hh"] = np.full(psi.shape, np.nan)
profiles["chi_e_hh"][core] = _smooth(_div(profiles["heatflux_e_HH"][core], -E_CHARGE*density_core*dte_core), smoothing_sigma)
return NeoResult(
profiles, (steps[0], steps[-1]), (float(times[0]), float(times[-1])),
len(steps), moment_product, smoothing_sigma, denominator_sigma, current_drive, fsa,
)
PLOT_GROUPS = {
"total_current": ("Total plasma current", "A m$^{-2}$", ("jpar_plasma", "jpar_efit", "sauter_total", "jpar_plasma_hh")),
"bootstrap_current": ("Bootstrap current", "A m$^{-2}$", ("xgc_bootstrap_estimate", "jboot_sauter", "hh_bootstrap")),
"particle_flux": ("Electron particle flux", "m$^{-2}$ s$^{-1}$", ("ptl_flux_e", "ptl_flux_HH")),
"particle_diffusivity": ("Electron particle diffusivity", "m$^2$ s$^{-1}$", ("D_eff", "D_HH")),
"heat_flux_i": ("Ion heat flux", "W m$^{-2}$", ("heatflux_i", "heatflux_i_CH")),
"heat_flux_e": ("Electron heat flux", "W m$^{-2}$", ("heatflux_e", "heatflux_e_HH")),
"energy_flux_i": ("Ion energy flux", "W m$^{-2}$", ("enflux_i_gc",)),
"energy_flux_e": ("Electron energy flux", "W m$^{-2}$", ("enflux_e", "energy_flux_hh")),
"heat_conductivity_i": ("Ion heat conductivity", "m$^2$ s$^{-1}$", ("chi_i", "chi_i_ch")),
"heat_conductivity_e": ("Electron heat conductivity", "m$^2$ s$^{-1}$", ("chi_e", "chi_e_hh")),
}
XGC_UNCERTAINTY = {
"jpar_plasma": "jpar_plasma_sigma",
"xgc_bootstrap_estimate": "jpar_plasma_sigma",
"ptl_flux_e": "ptl_flux_e_sigma",
"D_eff": "D_eff_sigma",
"heatflux_i": "heatflux_i_sigma",
"heatflux_e": "heatflux_e_sigma",
"enflux_i_gc": "enflux_i_gc_sigma",
"enflux_e": "enflux_e_sigma",
"chi_i": "chi_i_sigma",
"chi_e": "chi_e_sigma",
}