"""Compute SOL wall-loss source rates on the 1D flux-surface grid."""
from __future__ import annotations
from dataclasses import dataclass
import warnings
import numpy as np
from .sol_wall_mapping import SOLWallVolumeMap
[docs]
@dataclass
class SOLWallLossRatesResult:
"""Wall-loss source rates mapped to adjacent SOL surface-pair volumes."""
lower_surface: np.ndarray
upper_surface: np.ndarray
psi_lower: np.ndarray
psi_upper: np.ndarray
psi_center: np.ndarray
psi_center_norm: np.ndarray
volume_shell: np.ndarray
segment_to_bin: np.ndarray
time_left: np.ndarray
time_right: np.ndarray
dt: np.ndarray
interval_mask_used: np.ndarray
particle_rate_e: np.ndarray
particle_rate_i: np.ndarray
energy_rate_e: np.ndarray
energy_rate_i: np.ndarray
particle_rate_e_avg: np.ndarray
particle_rate_i_avg: np.ndarray
energy_rate_e_avg: np.ndarray
energy_rate_i_avg: np.ndarray
def _in_interval(s: float, s0: float, s1: float, wraps: bool) -> bool:
if not wraps:
return (s >= s0) and (s <= s1)
return (s >= s0) or (s <= s1)
def _remap_interval_to_target_length(
interval: tuple[float, float, bool],
source_length: float,
target_length: float,
) -> tuple[float, float, bool]:
s0, s1, wraps = interval
if source_length <= 0.0 or target_length <= 0.0:
raise ValueError("Wall-curve total length must be positive.")
fac = target_length / source_length
return float((s0 * fac) % target_length), float((s1 * fac) % target_length), bool(wraps)
def _build_heatdiag_segment_to_bin_map(
hd_s: np.ndarray,
mesh_intervals: list[tuple[tuple[float, float, bool], tuple[float, float, bool]]],
mesh_length: float,
hd_length: float,
) -> np.ndarray:
"""Map each heatdiag wall segment midpoint to one SOL bin index."""
nseg = hd_s.shape[0]
nbins = len(mesh_intervals)
seg_to_bin = np.full(nseg, -1, dtype=int)
hd_intervals: list[tuple[tuple[float, float, bool], tuple[float, float, bool]]] = []
for int0_mesh, int1_mesh in mesh_intervals:
int0_hd = _remap_interval_to_target_length(int0_mesh, mesh_length, hd_length)
int1_hd = _remap_interval_to_target_length(int1_mesh, mesh_length, hd_length)
hd_intervals.append((int0_hd, int1_hd))
for iseg, sval in enumerate(hd_s):
found = []
for ibin in range(nbins):
(s0a, s1a, wa), (s0b, s1b, wb) = hd_intervals[ibin]
if _in_interval(float(sval), s0a, s1a, wa) or _in_interval(float(sval), s0b, s1b, wb):
found.append(ibin)
if len(found) == 1:
seg_to_bin[iseg] = found[0]
elif len(found) > 1:
seg_to_bin[iseg] = found[0]
return seg_to_bin
def _weighted_time_average(rate: np.ndarray, dt: np.ndarray, mask: np.ndarray | None = None) -> np.ndarray:
if mask is None:
mask_use = dt > 0.0
else:
mask_use = np.asarray(mask, dtype=bool) & (dt > 0.0)
if not np.any(mask_use):
return np.zeros(rate.shape[1], dtype=float)
dt_use = dt[mask_use]
rate_use = rate[mask_use]
wsum = np.sum(dt_use)
if wsum <= 0.0:
return np.zeros(rate.shape[1], dtype=float)
return np.sum(rate_use * dt_use[:, None], axis=0) / wsum
[docs]
def compute_sol_wall_loss_rates(
plane,
heatdiag,
*,
phi_index: int = 0,
psi_norm_min: float = 1.0,
psi_norm_max: float | None = None,
interval_sample: str = "right",
time_window: tuple[float, float] | None = None,
) -> SOLWallLossRatesResult:
"""
Compute SOL wall-loss source rates on adjacent-surface shell volumes.
Parameters
----------
plane : Plane
Plane object with ``surf_map``, ``psi_surf``, ``x_psi``, ``vol_1d`` and wall data.
heatdiag : HeatDiag
HeatDiag reader with wall polygon and time-dependent wall loads.
phi_index : int
Toroidal index for selecting one heatdiag wall curve.
psi_norm_min : float
Lower normalized-psi cutoff for SOL pairing (default: 1.0).
psi_norm_max : float | None
Optional upper normalized-psi cutoff for SOL pairing.
interval_sample : str
Which time sample to use for interval-accumulated wall loads:
``"right"`` uses sample k+1 for [k,k+1], ``"left"`` uses sample k.
time_window : tuple[float, float] | None
Optional averaging window [t0, t1] in seconds. If provided, only
heatdiag intervals overlapping this window contribute to time averages.
Returns
-------
SOLWallLossRatesResult
Per-interval volumetric particle/energy loss rates [m^-3 s^-1, W m^-3]
for electrons and ions, and their dt-weighted averages.
"""
if interval_sample not in ("right", "left"):
raise ValueError("interval_sample must be 'right' or 'left'.")
mesh_wc = plane.get_wall_curve(set_inboard_origin=True)
hd_wc = heatdiag.get_wall_curve(
phi_index=phi_index,
set_inboard_origin=True,
r_axis=float(plane.axis_r),
)
mapper = SOLWallVolumeMap(plane, mesh_wc)
bounds = mapper.build_from_surf_map(psi_norm_min=psi_norm_min, psi_norm_max=psi_norm_max)
if len(bounds) == 0:
raise ValueError("No SOL adjacent-surface wall bounds were found.")
surf_map = np.asarray(plane.surf_map, dtype=int)
vol_1d = np.asarray(plane.vol_1d, dtype=float)
if surf_map.shape[0] != vol_1d.shape[0]:
raise ValueError("plane.surf_map and plane.vol_1d must have equal length.")
surf_to_map = {int(surf): i for i, surf in enumerate(surf_map)}
lower_surface = np.array([b.lower_surface for b in bounds], dtype=int)
upper_surface = np.array([b.upper_surface for b in bounds], dtype=int)
psi_lower = np.array([b.psi_lower for b in bounds], dtype=float)
psi_upper = np.array([b.psi_upper for b in bounds], dtype=float)
psi_center = 0.5 * (psi_lower + psi_upper)
psi_center_norm = psi_center / float(plane.x_psi)
volume_shell = np.zeros(len(bounds), dtype=float)
intervals = []
for ibin, b in enumerate(bounds):
if b.lower_surface not in surf_to_map or b.upper_surface not in surf_to_map:
raise ValueError("Surface in SOL bounds not found in plane.surf_map.")
ilo = surf_to_map[b.lower_surface]
ihi = surf_to_map[b.upper_surface]
volume_shell[ibin] = 0.5 * (vol_1d[ilo] + vol_1d[ihi])
intervals.append(b.intervals)
if np.any(volume_shell <= 0.0):
raise ValueError("Encountered non-positive SOL shell volume.")
hd_s = np.asarray(hd_wc.s_vertex, dtype=float)
seg_to_bin = _build_heatdiag_segment_to_bin_map(
hd_s,
intervals,
mesh_wc.total_length,
hd_wc.total_length,
)
tmask = np.asarray(heatdiag.get_time_mask(), dtype=int)
time_all = np.asarray(heatdiag.get_array("time")).reshape(-1)
times = time_all[tmask]
if times.shape[0] < 2:
warnings.warn(
"Need at least two unique heatdiag time samples; returning zero wall-loss rates.",
RuntimeWarning,
)
nbins = len(bounds)
zavg = np.zeros(nbins, dtype=float)
zstep = np.zeros((0, nbins), dtype=float)
return SOLWallLossRatesResult(
lower_surface=lower_surface,
upper_surface=upper_surface,
psi_lower=psi_lower,
psi_upper=psi_upper,
psi_center=psi_center,
psi_center_norm=psi_center_norm,
volume_shell=volume_shell,
segment_to_bin=seg_to_bin,
time_left=np.zeros(0, dtype=float),
time_right=np.zeros(0, dtype=float),
dt=np.zeros(0, dtype=float),
interval_mask_used=np.zeros(0, dtype=bool),
particle_rate_e=zstep.copy(),
particle_rate_i=zstep.copy(),
energy_rate_e=zstep.copy(),
energy_rate_i=zstep.copy(),
particle_rate_e_avg=zavg.copy(),
particle_rate_i_avg=zavg.copy(),
energy_rate_e_avg=zavg.copy(),
energy_rate_i_avg=zavg.copy(),
)
e_number = np.asarray(heatdiag.get_array("e_number"))
i_number = np.asarray(heatdiag.get_array("i_number"))
e_energy = np.asarray(heatdiag.get_array("e_para_energy")) + np.asarray(heatdiag.get_array("e_perp_energy"))
i_energy = np.asarray(heatdiag.get_array("i_para_energy")) + np.asarray(heatdiag.get_array("i_perp_energy"))
# Use a single toroidal index for now.
e_number = e_number[:, phi_index, :]
i_number = i_number[:, phi_index, :]
e_energy = e_energy[:, phi_index, :]
i_energy = i_energy[:, phi_index, :]
if e_number.shape[1] != hd_s.shape[0]:
raise ValueError(
"HeatDiag wall segment count does not match HeatDiag wall curve vertex count "
f"({e_number.shape[1]} vs {hd_s.shape[0]})."
)
nint = times.shape[0] - 1
nbins = len(bounds)
time_left = times[:-1].copy()
time_right = times[1:].copy()
dt = time_right - time_left
particle_rate_e = np.zeros((nint, nbins), dtype=float)
particle_rate_i = np.zeros((nint, nbins), dtype=float)
energy_rate_e = np.zeros((nint, nbins), dtype=float)
energy_rate_i = np.zeros((nint, nbins), dtype=float)
for k in range(nint):
if dt[k] <= 0.0:
warnings.warn(
f"Skipping non-positive heatdiag interval dt={dt[k]} at k={k}.",
RuntimeWarning,
)
continue
idx_use = tmask[k + 1] if interval_sample == "right" else tmask[k]
en = np.asarray(e_number[idx_use, :], dtype=float)
inum = np.asarray(i_number[idx_use, :], dtype=float)
ee = np.asarray(e_energy[idx_use, :], dtype=float)
ie = np.asarray(i_energy[idx_use, :], dtype=float)
for ibin in range(nbins):
mask = seg_to_bin == ibin
if not np.any(mask):
continue
inv = 1.0 / (dt[k] * volume_shell[ibin])
particle_rate_e[k, ibin] = np.sum(en[mask]) * inv
particle_rate_i[k, ibin] = np.sum(inum[mask]) * inv
energy_rate_e[k, ibin] = np.sum(ee[mask]) * inv
energy_rate_i[k, ibin] = np.sum(ie[mask]) * inv
# Time-window selection for averaging (if requested).
if time_window is not None:
t0_req, t1_req = float(time_window[0]), float(time_window[1])
if t1_req < t0_req:
t0_req, t1_req = t1_req, t0_req
# Clip requested window to available heatdiag time coverage.
t_hd_min = float(times[0])
t_hd_max = float(times[-1])
t0 = max(t0_req, t_hd_min)
t1 = min(t1_req, t_hd_max)
if t1 > t0:
interval_mask_used = (time_right > t0) & (time_left < t1) & (dt > 0.0)
else:
interval_mask_used = np.zeros(nint, dtype=bool)
if not np.any(interval_mask_used):
mids = 0.5 * (time_left + time_right)
k_near = int(np.argmin(np.abs(mids - 0.5 * (t0_req + t1_req))))
if nint > 0 and dt[k_near] > 0.0:
interval_mask_used = np.zeros(nint, dtype=bool)
interval_mask_used[k_near] = True
else:
interval_mask_used = dt > 0.0
else:
interval_mask_used = dt > 0.0
return SOLWallLossRatesResult(
lower_surface=lower_surface,
upper_surface=upper_surface,
psi_lower=psi_lower,
psi_upper=psi_upper,
psi_center=psi_center,
psi_center_norm=psi_center_norm,
volume_shell=volume_shell,
segment_to_bin=seg_to_bin,
time_left=time_left,
time_right=time_right,
dt=dt,
interval_mask_used=interval_mask_used,
particle_rate_e=particle_rate_e,
particle_rate_i=particle_rate_i,
energy_rate_e=energy_rate_e,
energy_rate_i=energy_rate_i,
particle_rate_e_avg=_weighted_time_average(particle_rate_e, dt, interval_mask_used),
particle_rate_i_avg=_weighted_time_average(particle_rate_i, dt, interval_mask_used),
energy_rate_e_avg=_weighted_time_average(energy_rate_e, dt, interval_mask_used),
energy_rate_i_avg=_weighted_time_average(energy_rate_i, dt, interval_mask_used),
)
[docs]
def interpolate_wall_loss_rates_to_surf_map(
plane,
result: SOLWallLossRatesResult,
*,
use_time_average: bool = True,
) -> dict[str, np.ndarray]:
"""
Interpolate SOL wall-loss rates from bin centers to ``plane.surf_map`` psi grid.
Core surfaces (psi/psi_x < 1) are set to zero.
Returns
-------
dict[str, np.ndarray]
Keys ``particle_e``, ``particle_i``, ``energy_e``, ``energy_i``.
If ``use_time_average=True``, each value has shape ``(npsi,)``.
Otherwise each value has shape ``(n_time_intervals, npsi)``.
"""
surf_map = np.asarray(plane.surf_map, dtype=int)
psi_surf = np.asarray(plane.psi_surf, dtype=float)
psi_norm = psi_surf[surf_map] / float(plane.x_psi)
src_psi = np.asarray(result.psi_center_norm, dtype=float)
def _interp_1d(y_src: np.ndarray) -> np.ndarray:
out = np.zeros_like(psi_norm, dtype=float)
mask = psi_norm >= 1.0 #+1.0e-5
if np.any(mask) and src_psi.size > 0:
out[mask] = np.interp(psi_norm[mask], src_psi, y_src, left=y_src[0], right=y_src[-1])
return out
def _interp_2d(y_src: np.ndarray) -> np.ndarray:
out = np.zeros((y_src.shape[0], psi_norm.shape[0]), dtype=float)
mask = psi_norm >= 1.0 #+1.0e-5
if np.any(mask) and src_psi.size > 0:
for it in range(y_src.shape[0]):
out[it, mask] = np.interp(psi_norm[mask], src_psi, y_src[it], left=y_src[it, 0], right=y_src[it, -1])
return out
if use_time_average:
return {
"particle_e": _interp_1d(np.asarray(result.particle_rate_e_avg, dtype=float)),
"particle_i": _interp_1d(np.asarray(result.particle_rate_i_avg, dtype=float)),
"energy_e": _interp_1d(np.asarray(result.energy_rate_e_avg, dtype=float)),
"energy_i": _interp_1d(np.asarray(result.energy_rate_i_avg, dtype=float)),
}
return {
"particle_e": _interp_2d(np.asarray(result.particle_rate_e, dtype=float)),
"particle_i": _interp_2d(np.asarray(result.particle_rate_i, dtype=float)),
"energy_e": _interp_2d(np.asarray(result.energy_rate_e, dtype=float)),
"energy_i": _interp_2d(np.asarray(result.energy_rate_i, dtype=float)),
}