"""Divertor heatdiag profile mapping and Eich-width fitting helpers.
This module converts XGC wall heat diagnostic data from ``xgc.heatdiag2.bp``
into upstream/midplane flux-density profiles suitable for Eich heat-load-width
fits. The main entry point is :func:`compute_divertor_eich_profiles`, which can
be called by scripts, notebooks, or thin GUI plugins.
The workflow implemented here is:
1. Create or reuse a :class:`~xgc_analysis.simulation.Simulation` and its
``Plane`` mesh.
2. Read heatdiag particle and energy totals through :class:`HeatDiag`.
3. Average one selected ADIOS frame, or a requested time window, to rates.
4. Split lower divertor wall points into inner and outer target branches using
the private-region poloidal-flux minimum.
5. Map target-wall ``psi_N`` values to inner/outer midplane radius maps built
from exact ``Plane`` flux-surface crossings with ``Z=0``. The plotted
``Delta_sep`` coordinate is therefore a midplane radial distance, not
distance along the target.
6. Estimate a uniform ``Delta_sep`` grid spacing from the heatdiag ``psi_N``
resolution mapped through the flux-surface midplane map, and remap particle
and energy loads with a locally conservative interval-overlap method.
7. Optionally smooth the remapped loads with total-load renormalization, convert
loads to toroidal-surface flux densities, and fit the total particle and
energy channels with the Eich functional form.
The helper is deliberately conservative about units and transformations:
heatdiag inputs are particle counts and energy in joules accumulated over
diagnostic intervals, rates are produced by dividing by interval duration, and
energy flux densities are displayed in ``MW m^-2`` only in the plotting layer.
"""
from __future__ import annotations
from dataclasses import dataclass
import math
from pathlib import Path
from typing import Any
import matplotlib.pyplot as plt
import numpy as np
from .heatdiag import HeatDiag
from .simulation import Simulation
E_CHARGE_C = 1.602176634e-19
[docs]
@dataclass
class DivertorEichProfile:
"""Mapped profile and optional Eich fits for one divertor target.
The profile stores a human-readable target ``label``, a mapped
``Delta_sep`` grid ``x_mm`` in millimeters, ion/electron particle and
energy flux-density arrays, and optional Eich fit results. Fit tuples
contain the fitted parameter vector ``(q0, q_bg, lambda_q, S, s0)`` in SI
units and a two-column ``(x_mm, fitted_y)`` array for plotting.
"""
label: str
x_mm: np.ndarray
particle_i: np.ndarray
particle_e: np.ndarray
energy_i: np.ndarray
energy_e: np.ndarray
fit_particle: tuple[np.ndarray, np.ndarray] | None
fit_energy: tuple[np.ndarray, np.ndarray] | None
[docs]
def compute_divertor_eich_profiles(
data_dir: str | Path,
*,
simulation: Any | None = None,
include_sheath: bool = False,
time_window: tuple[float, float] | None = None,
selected_frame_index: int | None = None,
psi_window: tuple[float, float] = (0.97, 1.12),
fit_window_mm: tuple[float, float] = (-2.0, 20.0),
smoothing_sigma_mm: float = 0.0,
show_outer: bool = True,
show_inner: bool = True,
) -> list[DivertorEichProfile]:
"""
Compute toroidally averaged divertor particle/energy flux-density profiles.
Parameters
----------
data_dir : str | Path
XGC output directory containing ``xgc.heatdiag2.bp`` and the mesh files.
simulation : Simulation | None
Optional pre-built Simulation. If omitted, a Simulation is created from
``data_dir``.
include_sheath : bool
Include ``[i,e]_potential`` in the energy channel if available.
time_window : tuple[float, float] | None
Optional heatdiag interval averaging window in seconds. If set, this
takes precedence over ``selected_frame_index``.
selected_frame_index : int | None
Optional GUI-selected heatdiag frame. When ``time_window`` is omitted,
the selected sample is converted to the diagnostic interval ending at
that frame.
psi_window : tuple[float, float]
Normalized-poloidal-flux range used to select divertor wall points.
fit_window_mm : tuple[float, float]
``Delta_sep`` window used for the Eich fit.
smoothing_sigma_mm : float
Gaussian smoothing sigma, in millimeters, applied to the displayed
particle/energy component profiles before fitting.
show_outer, show_inner : bool
Select which target profiles to compute.
Returns
-------
list[DivertorEichProfile]
One profile for each requested target.
"""
data_dir = Path(data_dir).expanduser().resolve()
sim = simulation
if sim is None:
sim = Simulation(directories=[str(data_dir)], is_stellarator=False, sim_is_axisymmetric=True)
plane = sim.mesh.get_plane(0)
heatdiag = HeatDiag(simulation=sim, data_dir=sim.data_directory)
rates = _load_rates(
heatdiag,
simulation=sim,
include_sheath=include_sheath,
time_window=time_window,
selected_frame_index=selected_frame_index,
)
r, z, psi_n = _prepare_wall_geometry(heatdiag, plane)
nseg = min(r.size, *(np.asarray(value).shape[1] for value in rates.values()))
r = r[:nseg]
z = z[:nseg]
psi_n = psi_n[:nseg]
psi_mid, r_mid_out, r_mid_in = _midplane_radius_maps(plane)
if psi_mid.size < 5:
raise ValueError("Could not build midplane radius maps from Plane surfaces.")
if show_outer and _finite_pair_count(psi_mid, r_mid_out) < 5:
raise ValueError("Could not build an outer-midplane radius map from Plane surfaces.")
if show_inner and _finite_pair_count(psi_mid, r_mid_in) < 5:
raise ValueError("Could not build an inner-midplane radius map from Plane surfaces.")
r_sep_out = _interp_radius_at_psi(psi_mid, r_mid_out, 1.0, "outer")
r_sep_in = _interp_radius_at_psi(psi_mid, r_mid_in, 1.0, "inner")
r_abs_out = _interp_radius_map(psi_mid, r_mid_out, psi_n, "outer")
r_abs_in = _interp_radius_map(psi_mid, r_mid_in, psi_n, "inner")
x_out = r_abs_out - r_sep_out
x_in = r_sep_in - r_abs_in
base = np.isfinite(psi_n) & (psi_n >= psi_window[0]) & (psi_n <= psi_window[1])
outer, inner = _target_masks_from_private_flux_minimum(r, z, psi_n, base, plane)
outer &= np.isfinite(x_out) & np.isfinite(r_abs_out)
inner &= np.isfinite(x_in) & np.isfinite(r_abs_in)
profiles = []
if show_outer:
profiles.append(
_build_target_profile(
"Outer",
outer,
x_out,
r_abs_out,
psi_n,
psi_mid,
r_mid_out,
r_sep_out,
"outer",
rates,
fit_window_mm=fit_window_mm,
smoothing_sigma_mm=smoothing_sigma_mm,
)
)
if show_inner:
profiles.append(
_build_target_profile(
"Inner",
inner,
x_in,
r_abs_in,
psi_n,
psi_mid,
r_mid_in,
r_sep_in,
"inner",
rates,
fit_window_mm=fit_window_mm,
smoothing_sigma_mm=smoothing_sigma_mm,
)
)
if not profiles:
raise ValueError("Both inner and outer target calculations are disabled.")
return profiles
[docs]
def plot_divertor_eich_profiles(
profiles: list[DivertorEichProfile],
*,
show_particle: bool = True,
show_energy: bool = True,
show_ions: bool = True,
show_electrons: bool = True,
show_total: bool = True,
xlim: tuple[float, float] | None = None,
ylim: tuple[float, float] | None = None,
):
"""
Plot divertor particle/energy profiles with optional Eich fit overlays.
Parameters
----------
profiles : list[DivertorEichProfile]
Profiles returned by :func:`compute_divertor_eich_profiles`.
show_particle, show_energy : bool
Select which physical channels are shown. At least one must be enabled.
show_ions, show_electrons, show_total : bool
Select species/component curves. Eich fits are only overlaid for total
profiles.
xlim, ylim : tuple[float, float] | None
Optional display limits. ``xlim`` is in millimeters. ``ylim`` applies to
whichever axes are displayed.
Returns
-------
matplotlib.figure.Figure
Figure owned by the caller. GUI code can save it to PNG; scripts can
display or further customize it.
"""
if not show_particle and not show_energy:
raise ValueError("At least one of particle or energy must be enabled.")
nrows = int(show_particle) + int(show_energy)
fig, axes = plt.subplots(nrows=nrows, ncols=1, figsize=(7.4, 3.6*nrows), dpi=150, squeeze=False)
axes_flat = list(axes[:, 0])
channel_axes = []
if show_particle:
channel_axes.append(("particle", axes_flat[len(channel_axes)]))
if show_energy:
channel_axes.append(("energy", axes_flat[len(channel_axes)]))
colors = {"Outer": "#0f5ea8", "Inner": "#a04400"}
for channel, ax in channel_axes:
for profile in profiles:
_plot_channel(
ax,
profile,
channel,
show_ions=show_ions,
show_electrons=show_electrons,
show_total=show_total,
color_base=colors.get(profile.label, "#333333"),
)
ax.axvline(0.0, color="black", linewidth=0.8, alpha=0.45)
ax.set_xlabel(r"$\Delta_{sep}$ [mm]")
ax.grid(alpha=0.25, linewidth=0.4)
ax.legend(loc="best", fontsize=7, ncol=2)
if xlim is not None:
ax.set_xlim(xlim)
if ylim is not None:
ax.set_ylim(ylim)
fig.suptitle("Divertor heatdiag profile mapped to midplane radius", fontsize=11)
fig.tight_layout(pad=0.35)
return fig
def _get_sorted_unique_time_indices(heatdiag) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Return de-duplicated heatdiag indices, times, and positive intervals.
Heatdiag files can contain repeated or overlapping step histories. When the
``HeatDiag`` reader exposes a ``step`` series, its ``get_time_mask`` method
selects the last occurrence of each step. If that is unavailable, this
helper falls back to the raw ADIOS step order.
"""
try:
tmask = np.asarray(heatdiag.get_time_mask(), dtype=int)
except Exception:
tmask = np.arange(len(heatdiag.available_steps), dtype=int)
time = np.asarray(heatdiag.get_array("time"), dtype=float).reshape(-1)
time = time[tmask]
if time.shape[0] < 2:
raise ValueError("Need at least two heatdiag time samples to compute rates.")
dt = np.diff(time)
if not np.any(dt > 0.0):
raise ValueError("Heatdiag time samples do not contain positive time intervals.")
return tmask, time, dt
def _weighted_average_intervals(
values: np.ndarray,
tmask: np.ndarray,
time: np.ndarray,
window,
selected_frame_index: int | None,
) -> np.ndarray:
"""Average accumulated heatdiag interval data into rates.
Parameters
----------
values : np.ndarray
Heatdiag history with shape ``(time, nphi, nseg)`` or
``(time, nseg)``. Entries are interval totals from ADIOS output.
tmask : np.ndarray
Indices into the raw heatdiag history after de-duplication.
time : np.ndarray
De-duplicated physical times, one per selected heatdiag sample.
window : tuple[float, float] | None
Optional averaging window in seconds. If set, all intervals overlapping
the window contribute with duration weights.
selected_frame_index : int | None
Raw GUI/ADIOS frame index used when ``window`` is omitted. The selected
raw index is mapped to the nearest de-duplicated interval ending at or
after that raw frame.
Returns
-------
np.ndarray
Duration-weighted rate with shape ``(nphi, nseg)``.
"""
dt = np.diff(time)
sample_idx = tmask[1:]
interval_mask = dt > 0.0
if window is not None:
t0, t1 = window
interval_mask &= (time[1:] > t0) & (time[:-1] < t1)
elif selected_frame_index is not None:
j = max(0, int(selected_frame_index))
k = int(np.searchsorted(sample_idx, j, side="left")) if sample_idx.size else 0
k = max(0, min(k, dt.size - 1))
interval_mask = np.zeros_like(dt, dtype=bool)
interval_mask[k] = dt[k] > 0.0
if not np.any(interval_mask):
mids = 0.5 * (time[:-1] + time[1:])
if window is None:
if selected_frame_index is None:
k = int(np.argmax(dt))
else:
j = max(0, int(selected_frame_index))
k = int(np.searchsorted(sample_idx, j, side="left")) if sample_idx.size else 0
k = max(0, min(k, dt.size - 1))
else:
k = int(np.argmin(np.abs(mids - 0.5 * (window[0] + window[1]))))
interval_mask = np.zeros_like(dt, dtype=bool)
interval_mask[k] = dt[k] > 0.0
arr = np.asarray(values, dtype=float)
if arr.ndim == 2:
arr = arr[:, np.newaxis, :]
if arr.ndim != 3:
raise ValueError(f"Expected heatdiag history with shape (time,nphi,nseg), got {arr.shape}")
rate = arr[sample_idx] / dt[:, None, None]
weights = dt[interval_mask]
return np.sum(rate[interval_mask] * weights[:, None, None], axis=0) / np.sum(weights)
def _species_charge_c(simulation, index: int, default: float) -> float:
"""Return a species charge in Coulombs with a controlled fallback.
The heatdiag sheath-potential contribution has to use the physical species
charge, not an implicit elementary-charge magnitude. If the Simulation
object does not expose species metadata, the caller-provided default keeps
the previous electron/proton behavior.
"""
try:
species = getattr(simulation, "species", None)
if species is None or len(species) <= index:
return float(default)
charge = float(getattr(species[index], "charge_C"))
return charge if np.isfinite(charge) else float(default)
except Exception:
return float(default)
def _load_rates(
heatdiag,
*,
simulation,
include_sheath: bool,
time_window,
selected_frame_index: int | None,
) -> dict[str, np.ndarray]:
"""Load particle and energy rates from the heatdiag reader.
Particle channels are converted from accumulated counts to ``1/s``.
Parallel and perpendicular energy channels are summed and converted to
``W``. If ``include_sheath`` is enabled, the stored potential terms are
multiplied by the signed Simulation species charges and added to the energy
channels.
"""
tmask, time, _dt = _get_sorted_unique_time_indices(heatdiag)
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"))
if include_sheath:
q_e = _species_charge_c(simulation, 0, -E_CHARGE_C)
q_i = _species_charge_c(simulation, 1, E_CHARGE_C)
e_energy = e_energy + np.asarray(heatdiag.get_array("e_potential")) * q_e
i_energy = i_energy + np.asarray(heatdiag.get_array("i_potential")) * q_i
return {
"particle_e": _weighted_average_intervals(
np.asarray(heatdiag.get_array("e_number")), tmask, time, time_window, selected_frame_index
),
"particle_i": _weighted_average_intervals(
np.asarray(heatdiag.get_array("i_number")), tmask, time, time_window, selected_frame_index
),
"energy_e": _weighted_average_intervals(e_energy, tmask, time, time_window, selected_frame_index),
"energy_i": _weighted_average_intervals(i_energy, tmask, time, time_window, selected_frame_index),
}
def _first_phi_geometry(heatdiag, name: str) -> np.ndarray:
"""Return a static wall-geometry array for the first toroidal plane.
Heatdiag wall geometry can be stored either as ``(nphi, nseg)`` or as a
single ``(nseg,)`` array. This helper normalizes both forms for target
splitting and midplane-coordinate mapping.
"""
arr = np.asarray(heatdiag.get_wall_array(name), dtype=float)
if arr.ndim == 2:
return arr[0].copy()
if arr.ndim == 1:
return arr.copy()
raise ValueError(f"Expected heatdiag wall variable '{name}' to be rank 1 or 2, got {arr.shape}")
def _prepare_wall_geometry(heatdiag, plane):
"""Load wall ``R``, ``Z``, and normalized ``psi`` arrays.
The heatdiag ``psi`` values are normalized by ``plane.x_psi`` so subsequent
selection and midplane mapping operate on ``psi_N``. All arrays are trimmed
to a common segment count.
"""
r = _first_phi_geometry(heatdiag, "r")
z = _first_phi_geometry(heatdiag, "z")
psi_n = _first_phi_geometry(heatdiag, "psi") / float(plane.x_psi)
n = min(r.size, z.size, psi_n.size)
if n < 3:
raise ValueError("Heatdiag wall geometry has too few points.")
return r[:n], z[:n], psi_n[:n]
def _wall_arclength(r: np.ndarray, z: np.ndarray) -> np.ndarray:
"""Return cumulative arclength along the ordered heatdiag wall polygon.
The arclength coordinate is used only for robustly splitting the target
branches at the private-region flux minimum. It is not used as the plotted
``Delta_sep`` coordinate.
"""
r = np.asarray(r, dtype=float)
z = np.asarray(z, dtype=float)
s = np.zeros(r.shape[0], dtype=float)
if r.shape[0] > 1:
ds = np.sqrt(np.diff(r)**2 + np.diff(z)**2)
s[1:] = np.cumsum(ds)
return s
def _target_masks_from_private_flux_minimum(
r: np.ndarray,
z: np.ndarray,
psi_n: np.ndarray,
base_mask: np.ndarray,
plane,
) -> tuple[np.ndarray, np.ndarray]:
"""Split lower divertor wall points into outer and inner target masks.
The split point is the minimum normalized poloidal flux found in the lower
private-region wall data. Points before and after that location along the
ordered wall polygon are assigned to target branches. The branch with larger
mean major radius is labeled outer. If the private-region split degenerates,
a fallback major-radius split at the X-point/axis radius is used. Each
branch is then reduced to the contiguous monotonic ``psi_N`` region around
the separatrix if folded target data are detected.
"""
z_cut = float(getattr(plane, "x_z", getattr(plane, "axis_z", 0.0)))
lower = z < z_cut
finite = np.isfinite(psi_n) & np.isfinite(r) & np.isfinite(z)
private = lower & finite & (psi_n < 1.0)
if not np.any(private):
private = lower & finite
if not np.any(private):
raise ValueError("Could not find lower wall points for inner/outer target split.")
s = _wall_arclength(r, z)
private_indices = np.flatnonzero(private)
split_index = private_indices[int(np.argmin(psi_n[private]))]
split_s = float(s[split_index])
side_after = base_mask & lower & (s >= split_s)
side_before = base_mask & lower & (s <= split_s)
if not np.any(side_after) or not np.any(side_before):
x_r = float(getattr(plane, "x_r", plane.axis_r))
outer = base_mask & lower & (r >= x_r)
inner = base_mask & lower & (r < x_r)
return (
_monotonic_region_around_separatrix(outer, psi_n, s),
_monotonic_region_around_separatrix(inner, psi_n, s),
)
mean_after = float(np.nanmean(r[side_after]))
mean_before = float(np.nanmean(r[side_before]))
if mean_after >= mean_before:
outer, inner = side_after, side_before
else:
outer, inner = side_before, side_after
return (
_monotonic_region_around_separatrix(outer, psi_n, s),
_monotonic_region_around_separatrix(inner, psi_n, s),
)
def _monotonic_region_around_separatrix(mask: np.ndarray, psi_n: np.ndarray, s: np.ndarray) -> np.ndarray:
"""Return the monotonic target subset connected to the separatrix.
Heatdiag wall curves can contain folded or repeated target regions. For
each branch, this helper inspects selected points in wall-arclength order.
If ``psi_N`` is already monotonic, the original mask is returned. Otherwise
the largest monotonic interval containing the point closest to ``psi_N=1``
is used.
"""
idx = np.flatnonzero(mask & np.isfinite(psi_n) & np.isfinite(s))
if idx.size < 3:
return mask
idx = idx[np.argsort(s[idx])]
psi = np.asarray(psi_n[idx], dtype=float)
tol = 1.0e-8
diff = np.diff(psi)
if np.all(diff >= -tol) or np.all(diff <= tol):
return mask
sep = int(np.argmin(np.abs(psi - 1.0)))
def grow(increasing: bool) -> tuple[int, int]:
left = sep
right = sep
while left > 0:
ok = psi[left - 1] <= psi[left] + tol if increasing else psi[left - 1] >= psi[left] - tol
if not ok:
break
left -= 1
while right < psi.size - 1:
ok = psi[right] <= psi[right + 1] + tol if increasing else psi[right] >= psi[right + 1] - tol
if not ok:
break
right += 1
return left, right
inc_left, inc_right = grow(True)
dec_left, dec_right = grow(False)
if inc_right - inc_left >= dec_right - dec_left:
left, right = inc_left, inc_right
else:
left, right = dec_left, dec_right
out = np.zeros_like(mask, dtype=bool)
out[idx[left:right + 1]] = True
return out
def _ordered_surface_vertices(plane, isurf: int) -> tuple[np.ndarray, bool]:
"""Return surface vertices ordered by straight-field-line poloidal angle.
The returned boolean indicates whether the surface should be closed when
constructing interpolation segments. Core surfaces are closed; open SOL or
private-flux surfaces are treated as polylines.
"""
vids = np.asarray(plane.get_surface_vertex_indices(isurf), dtype=int)
if vids.size == 0:
return vids, False
if hasattr(plane, "theta"):
theta = np.asarray(plane.theta[vids], dtype=float)
finite = np.isfinite(theta)
if np.any(finite):
vids = vids[finite][np.argsort(theta[finite])]
closed = False
try:
closed = int(np.asarray(plane.region)[int(vids[0])]) == 1
except Exception:
closed = True
return vids, closed
def _surface_midplane_crossings(plane, isurf: int, z_mid: float) -> tuple[float | None, float | None]:
"""Interpolate one flux surface to the outer and inner midplane crossings.
Each surface is traversed as a polyline in increasing ``theta``. Segment
crossings with ``Z=z_mid`` are linearly interpolated in ``(R, Z)``. The
largest crossing radius is the outboard midplane radius and the smallest
crossing radius is the inboard midplane radius.
"""
vids, closed = _ordered_surface_vertices(plane, isurf)
if vids.size < 2:
return None, None
coords = np.asarray(plane.rz[vids], dtype=float)
r = coords[:, 0]
z = coords[:, 1] - float(z_mid)
nseg = vids.size if closed else vids.size - 1
crossings = []
for i in range(nseg):
j = (i + 1) % vids.size
z0 = z[i]
z1 = z[j]
r0 = r[i]
r1 = r[j]
if not np.all(np.isfinite([z0, z1, r0, r1])):
continue
if abs(z0) <= 1.0e-12:
crossings.append(float(r0))
if z0*z1 < 0.0:
frac = -z0/(z1 - z0)
crossings.append(float(r0 + frac*(r1 - r0)))
elif abs(z1) <= 1.0e-12:
crossings.append(float(r1))
if not crossings:
return None, None
axis_r = float(plane.axis_r)
crossings = np.asarray(crossings, dtype=float)
outer = crossings[crossings >= axis_r]
inner = crossings[crossings <= axis_r]
r_outer = float(np.nanmax(outer)) if outer.size else None
r_inner = float(np.nanmin(inner)) if inner.size else None
return r_outer, r_inner
def _midplane_radius_maps(plane) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Build ``psi_N -> R_midplane`` maps for outer and inner midplanes.
For each ``Plane`` flux surface, this linearly interpolates exact
intersections with ``Z=0``. Surfaces that do not cross the outboard or
inboard midplane keep ``NaN`` for that branch. The result is sorted by
normalized flux and later used to map target-wall ``psi_N`` to
upstream/midplane radial distance from the separatrix.
"""
psi_vals = []
r_outer = []
r_inner = []
z_mid = 0.0
for isurf in range(int(plane.nsurf)):
outer, inner = _surface_midplane_crossings(plane, isurf, z_mid)
if outer is None and inner is None:
continue
psi_vals.append(float(plane.psi_surf[isurf]) / float(plane.x_psi))
r_outer.append(np.nan if outer is None else float(outer))
r_inner.append(np.nan if inner is None else float(inner))
psi = np.asarray(psi_vals, dtype=float)
order = np.argsort(psi)
psi = psi[order]
return psi, np.asarray(r_outer, dtype=float)[order], np.asarray(r_inner, dtype=float)[order]
def _finite_pair_count(x: np.ndarray, y: np.ndarray) -> int:
"""Count finite ``(x, y)`` pairs."""
return int(np.count_nonzero(np.isfinite(x) & np.isfinite(y)))
def _unique_average(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Return sorted unique ``x`` values and averaged ``y`` values."""
order = np.argsort(x)
xs = np.asarray(x[order], dtype=float)
ys = np.asarray(y[order], dtype=float)
uniq, inv = np.unique(xs, return_inverse=True)
y_acc = np.zeros(uniq.size, dtype=float)
n_acc = np.zeros(uniq.size, dtype=float)
for i, j in enumerate(inv):
y_acc[j] += ys[i]
n_acc[j] += 1.0
return uniq, y_acc/np.maximum(n_acc, 1.0)
def _interp_radius_map(psi_mid: np.ndarray, r_mid: np.ndarray, psi_query: np.ndarray, branch: str) -> np.ndarray:
"""Interpolate one branch of the midplane radius map to target ``psi_N``."""
finite = np.isfinite(psi_mid) & np.isfinite(r_mid)
if np.count_nonzero(finite) < 2:
raise ValueError(f"Not enough finite {branch}-midplane crossings for interpolation.")
psi, radius = _unique_average(psi_mid[finite], r_mid[finite])
return np.interp(psi_query, psi, radius, left=np.nan, right=np.nan)
def _interp_radius_at_psi(psi_mid: np.ndarray, r_mid: np.ndarray, psi_value: float, branch: str) -> float:
"""Interpolate one branch of the midplane radius map at a scalar ``psi_N``."""
finite = np.isfinite(psi_mid) & np.isfinite(r_mid)
psi, radius = _unique_average(psi_mid[finite], r_mid[finite])
if psi.size < 2 or psi_value < psi[0] or psi_value > psi[-1]:
raise ValueError(f"The {branch}-midplane map does not bracket psi_N={psi_value}.")
return float(np.interp(float(psi_value), psi, radius))
def _target_psi_grid_from_heatdiag(psi_values: np.ndarray) -> np.ndarray:
"""Build a target ``psi_N`` grid from heatdiag wall-segment resolution.
The spacing is the mean finite positive ``Delta psi_N`` among the selected
target segments. This keeps the final profile grid tied to the diagnostic
resolution instead of using the generally coarser flux-surface list.
Endpoints are included so the grid spans the selected target branch.
"""
psi = np.asarray(psi_values, dtype=float)
psi = np.unique(psi[np.isfinite(psi)])
if psi.size < 2:
raise ValueError("Need at least two finite heatdiag psi values to build target grid.")
psi.sort()
dpsi_values = np.diff(psi)
dpsi_values = dpsi_values[dpsi_values > 1.0e-12]
if dpsi_values.size == 0:
raise ValueError("Heatdiag psi values do not contain a positive spacing.")
dpsi = float(np.nanmean(dpsi_values))
psi_min = float(psi[0])
psi_max = float(psi[-1])
count = int(math.floor((psi_max - psi_min)/dpsi)) + 1
grid = psi_min + dpsi*np.arange(max(1, count), dtype=float)
if grid[-1] < psi_max - 1.0e-10*dpsi:
grid = np.append(grid, psi_max)
else:
grid[-1] = psi_max
return grid
def _mapped_target_grid_from_psi(
psi_values: np.ndarray,
psi_mid: np.ndarray,
r_mid: np.ndarray,
r_sep: float,
branch: str,
) -> tuple[np.ndarray, np.ndarray]:
"""Build a uniform ``Delta_sep`` grid from mapped heatdiag ``psi_N`` spacing.
A heatdiag-resolution ``psi_N`` grid is first mapped through the
flux-surface ``psi_N -> R_mid`` relation. The median positive spacing of
the resulting ``Delta_sep`` coordinates is used as the final uniform
``Delta_sep`` spacing for smoothing and fitting.
"""
psi_grid = _target_psi_grid_from_heatdiag(psi_values)
r_mapped = _interp_radius_map(psi_mid, r_mid, psi_grid, branch)
if branch == "outer":
x_mapped = r_mapped - float(r_sep)
elif branch == "inner":
x_mapped = float(r_sep) - r_mapped
else:
raise ValueError("branch must be 'outer' or 'inner'.")
finite = np.isfinite(x_mapped) & np.isfinite(r_mapped)
if np.count_nonzero(finite) < 2:
raise ValueError(f"Not enough finite {branch}-target grid points after psi-to-midplane mapping.")
x_sorted = np.unique(np.sort(x_mapped[finite]))
dx_values = np.diff(x_sorted)
dx_values = dx_values[dx_values > 1.0e-12]
if dx_values.size == 0:
raise ValueError(f"Mapped {branch}-target grid does not contain a positive Delta_sep spacing.")
dx = float(np.nanmedian(dx_values))
x_min = float(x_sorted[0])
x_max = float(x_sorted[-1])
span = x_max - x_min
if not np.isfinite(dx) or dx <= 0.0 or span <= 0.0:
raise ValueError(f"Could not build a finite uniform Delta_sep grid for the {branch} target.")
n_intervals = max(1, int(math.ceil(span/dx)))
n_points = n_intervals + 1
n_points = max(20, min(1000, n_points))
x_grid = np.linspace(x_min, x_max, n_points)
if branch == "outer":
r_grid = float(r_sep) + x_grid
else:
r_grid = float(r_sep) - x_grid
return x_grid, r_grid
def _cell_edges_from_centers(x: np.ndarray) -> np.ndarray:
"""Construct finite-volume cell edges around sorted cell centers."""
x = np.asarray(x, dtype=float)
if x.size < 2:
raise ValueError("Need at least two cell centers to construct cell edges.")
edges = np.empty(x.size + 1, dtype=float)
edges[1:-1] = 0.5*(x[:-1] + x[1:])
edges[0] = x[0] - 0.5*(x[1] - x[0])
edges[-1] = x[-1] + 0.5*(x[-1] - x[-2])
return edges
def _conservative_remap(x_src: np.ndarray, y_src: np.ndarray, x_dst: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Remap loads with local finite-volume conservation.
``y_src`` is treated as a set of loads assigned to source wall points, not
as a pointwise density. Source and destination points are interpreted as
cell centers. The source load in each cell is distributed to destination
cells by interval-overlap fraction, so load is conserved locally by mapped
``Delta_sep`` interval and globally over the target branch.
"""
finite = np.isfinite(x_src) & np.isfinite(y_src)
xs = np.asarray(x_src[finite], dtype=float)
ys = np.asarray(y_src[finite], dtype=float)
x_dst = np.asarray(x_dst, dtype=float)
if xs.size < 2:
return np.zeros_like(x_dst, dtype=float), np.gradient(x_dst)
order = np.argsort(xs)
xs = xs[order]
ys = ys[order]
uniq, inv = np.unique(xs, return_inverse=True)
if uniq.size != xs.size:
y_acc = np.zeros(uniq.size, dtype=float)
for i, j in enumerate(inv):
y_acc[j] += ys[i]
xs = uniq
ys = y_acc
if xs.size < 2:
return np.zeros_like(x_dst, dtype=float), np.gradient(x_dst)
src_edges = _cell_edges_from_centers(xs)
dst_edges = _cell_edges_from_centers(x_dst)
dst_edges[0] = min(dst_edges[0], src_edges[0])
dst_edges[-1] = max(dst_edges[-1], src_edges[-1])
widths = np.diff(dst_edges)
out = np.zeros(x_dst.size, dtype=float)
j0 = 0
for i in range(xs.size):
src_left = src_edges[i]
src_right = src_edges[i + 1]
src_width = src_right - src_left
if src_width <= 0.0:
continue
while j0 < x_dst.size and dst_edges[j0 + 1] <= src_left:
j0 += 1
j = j0
while j < x_dst.size and dst_edges[j] < src_right:
overlap = min(src_right, dst_edges[j + 1]) - max(src_left, dst_edges[j])
if overlap > 0.0:
out[j] += ys[i]*overlap/src_width
j += 1
src_total = float(np.sum(ys))
dst_total = float(np.sum(out))
if np.isfinite(src_total) and np.isfinite(dst_total) and abs(dst_total) > 0.0:
out *= src_total / dst_total
return out, widths
def _erfc_array(x: np.ndarray) -> np.ndarray:
"""Evaluate complementary error function for arrays.
SciPy is used when available. The standard-library ``math.erfc`` fallback
keeps the Eich model usable in minimal environments, albeit more slowly.
"""
try:
from scipy.special import erfc
return erfc(x)
except Exception:
return np.vectorize(math.erfc)(x)
def _gaussian_smooth_uniform(x_m: np.ndarray, y: np.ndarray, sigma_mm: float) -> np.ndarray:
"""Apply Gaussian smoothing on a nearly uniform ``x_m`` grid.
Parameters
----------
x_m : np.ndarray
Grid in meters. The grid is expected to be nearly uniform because it is
created by :func:`numpy.linspace` in :func:`_build_target_profile`.
y : np.ndarray
Profile values to smooth.
sigma_mm : float
Gaussian sigma in millimeters. Non-positive values return a copy-like
``np.asarray`` view of the input.
Returns
-------
np.ndarray
Smoothed profile. NaNs are handled by smoothing values and finite-value
weights separately, then dividing.
"""
sigma_m = max(float(sigma_mm), 0.0)*1.0e-3
if sigma_m <= 0.0 or y.size < 3:
return np.asarray(y, dtype=float)
dx = float(np.nanmedian(np.abs(np.diff(np.asarray(x_m, dtype=float)))))
if not np.isfinite(dx) or dx <= 0.0:
return np.asarray(y, dtype=float)
sigma_cells = sigma_m/dx
if not np.isfinite(sigma_cells) or sigma_cells <= 0.0:
return np.asarray(y, dtype=float)
values = np.asarray(y, dtype=float)
try:
from scipy.ndimage import gaussian_filter1d
finite = np.isfinite(values)
filled = np.where(finite, values, 0.0)
weights = gaussian_filter1d(finite.astype(float), sigma_cells, mode="nearest")
smoothed = gaussian_filter1d(filled, sigma_cells, mode="nearest")
return np.divide(smoothed, weights, out=np.full_like(values, np.nan), where=weights > 0.0)
except Exception:
radius = max(1, int(math.ceil(4.0*sigma_cells)))
offsets = np.arange(-radius, radius + 1, dtype=float)
kernel = np.exp(-0.5*(offsets/sigma_cells)**2)
kernel /= np.sum(kernel)
finite = np.isfinite(values)
filled = np.where(finite, values, 0.0)
smoothed = np.convolve(filled, kernel, mode="same")
weights = np.convolve(finite.astype(float), kernel, mode="same")
return np.divide(smoothed, weights, out=np.full_like(values, np.nan), where=weights > 0.0)
def _gaussian_smooth_conserved_load(x_m: np.ndarray, load: np.ndarray, sigma_mm: float) -> np.ndarray:
"""Smooth mapped cell loads while preserving their total.
Smoothing densities directly changes the integrated particle/energy load
when the midplane surface area varies across the grid. This helper smooths
the already conservative mapped loads and then renormalizes to the original
load sum before density conversion.
"""
values = np.asarray(load, dtype=float)
if sigma_mm <= 0.0 or values.size < 3:
return values
smoothed = _gaussian_smooth_uniform(x_m, values, sigma_mm)
src_total = float(np.nansum(values))
dst_total = float(np.nansum(smoothed))
if np.isfinite(src_total) and np.isfinite(dst_total) and abs(dst_total) > 0.0:
smoothed = smoothed * (src_total/dst_total)
return smoothed
[docs]
def eich_model(x_m: np.ndarray, q0: float, q_bg: float, lambda_q: float, s: float, s0: float) -> np.ndarray:
"""
Eich profile model from Eich et al. NF 53, 093031, Eq. 1 without ``f_x``.
The omitted target flux-expansion factor is appropriate when ``x_m`` is the
already mapped upstream/midplane radius coordinate.
Parameters are in SI units: ``x_m``, ``lambda_q``, ``S`` and ``s0`` are
meters, while ``q0`` and ``q_bg`` have the same units as the profile being
fitted.
"""
x_m = np.asarray(x_m, dtype=float)
lambda_q = max(float(lambda_q), 1.0e-9)
s = max(float(s), 1.0e-9)
arg = np.clip((s/(2.0*lambda_q))**2 - (x_m - s0)/lambda_q, -300.0, 300.0)
return 0.5*q0*np.exp(arg)*_erfc_array(s/(2.0*lambda_q) - (x_m - s0)/s) + q_bg
[docs]
def fit_eich_profile(
x_m: np.ndarray,
y: np.ndarray,
fit_window_mm,
smoothing_sigma_mm: float = 0.0,
) -> tuple[np.ndarray, np.ndarray] | None:
"""Fit ``y(x_m)`` with :func:`eich_model` and return parameters plus curve.
The fit uses positive finite points, optionally restricted to
``fit_window_mm``. The fitted ``lambda_q`` and ``S`` parameters are in
meters in the returned parameter vector; plotting text converts them to
millimeters.
"""
try:
from scipy.optimize import curve_fit
except Exception:
return None
y_fit_data = _gaussian_smooth_uniform(x_m, y, smoothing_sigma_mm)
mask = np.isfinite(x_m) & np.isfinite(y_fit_data) & (y_fit_data > 0.0)
if fit_window_mm is not None:
mask &= (x_m*1.0e3 >= fit_window_mm[0]) & (x_m*1.0e3 <= fit_window_mm[1])
if np.count_nonzero(mask) < 5:
return None
xf = x_m[mask]
yf = y_fit_data[mask]
q_bg0 = max(float(np.nanpercentile(yf, 5.0)), 0.0)
q0 = max(float(np.nanmax(yf) - q_bg0), float(np.nanmax(yf)), 1.0)
p0 = [q0, q_bg0, 1.0e-3, 1.0e-3, 0.0]
lower = [0.0, 0.0, 1.0e-5, 1.0e-5, float(np.nanmin(xf)) - 0.05]
upper = [np.inf, np.inf, 0.20, 0.20, float(np.nanmax(xf)) + 0.05]
sigma = np.maximum(np.abs(yf), 0.05*np.nanmax(np.abs(yf)))
popt, _pcov = curve_fit(
eich_model,
xf,
yf,
p0=p0,
sigma=sigma,
bounds=(lower, upper),
maxfev=10000,
)
xfit = np.linspace(float(np.nanmin(xf)), float(np.nanmax(xf)), 240)
return popt, np.column_stack([xfit*1.0e3, eich_model(xfit, *popt)])
def _build_target_profile(
label: str,
mask: np.ndarray,
x_m: np.ndarray,
r_abs_m: np.ndarray,
psi_n: np.ndarray,
psi_mid: np.ndarray,
r_mid: np.ndarray,
r_sep: float,
branch: str,
rates: dict[str, np.ndarray],
*,
fit_window_mm,
smoothing_sigma_mm: float,
) -> DivertorEichProfile:
"""Build one target profile on a heatdiag-resolution ``psi_N`` grid.
The selected heatdiag target points define a representative ``psi_N``
spacing. That spacing is mapped to midplane ``Delta_sep`` with the
flux-surface ``psi_N -> R_mid`` map, and the median mapped spacing is used
to build the final uniform ``Delta_sep`` grid. Selected loads are then
remapped conservatively onto that grid, optionally smoothed as loads with
total-load renormalization, divided by toroidal midplane surface area to
form flux densities, and fitted in total-particle and total-energy
channels.
"""
valid_mask = mask & np.isfinite(x_m) & np.isfinite(r_abs_m) & np.isfinite(psi_n)
if np.count_nonzero(valid_mask) < 5:
raise ValueError(f"Not enough heatdiag wall points for {label} target after filtering.")
x_grid, r_grid = _mapped_target_grid_from_psi(psi_n[valid_mask], psi_mid, r_mid, r_sep, branch)
grid_width = None
def interp_rate(name: str) -> np.ndarray:
nonlocal grid_width
arr = np.asarray(rates[name], dtype=float)
if arr.ndim != 2:
raise ValueError(f"Expected averaged rate '{name}' to have shape (nphi, nseg), got {arr.shape}")
n = min(arr.shape[1], valid_mask.shape[0])
arr = arr[:, :n]
local_mask = valid_mask[:n]
toroidal_sum = np.sum(arr[:, local_mask], axis=0)
load, width = _conservative_remap(x_m[:n][local_mask], toroidal_sum, x_grid)
grid_width = width if grid_width is None else grid_width
return _gaussian_smooth_conserved_load(x_grid, load, smoothing_sigma_mm)
particle_i_load = interp_rate("particle_i")
particle_e_load = interp_rate("particle_e")
energy_i_load = interp_rate("energy_i")
energy_e_load = interp_rate("energy_e")
if grid_width is None:
grid_width = np.gradient(x_grid)
area = grid_width * 2.0 * np.pi * np.maximum(np.abs(r_grid), 1.0e-6)
area = np.maximum(np.abs(area), 1.0e-20)
particle_i = particle_i_load / area
particle_e = particle_e_load / area
energy_i = energy_i_load / area
energy_e = energy_e_load / area
particle_total = particle_i + particle_e
energy_total = energy_i + energy_e
fit_particle = fit_eich_profile(x_grid, particle_total, fit_window_mm)
fit_energy = fit_eich_profile(x_grid, energy_total, fit_window_mm)
return DivertorEichProfile(
label=label,
x_mm=x_grid*1.0e3,
particle_i=particle_i,
particle_e=particle_e,
energy_i=energy_i,
energy_e=energy_e,
fit_particle=fit_particle,
fit_energy=fit_energy,
)
def _plot_channel(
ax,
profile: DivertorEichProfile,
channel: str,
*,
show_ions: bool,
show_electrons: bool,
show_total: bool,
color_base: str,
) -> None:
"""Draw one particle or energy channel for a target profile.
Species curves are drawn with dashed/dotted styles, the total is drawn as a
solid curve, and the Eich fit overlay plus fitted ``lambda``/``S`` text are
shown only when ``show_total`` is enabled and a fit is available.
"""
if channel == "particle":
yi = profile.particle_i
ye = profile.particle_e
fit = profile.fit_particle
ylabel = r"Particle flux density [m$^{-2}$ s$^{-1}$]"
else:
yi = profile.energy_i/1.0e6
ye = profile.energy_e/1.0e6
fit = profile.fit_energy
ylabel = r"Energy flux density [MW m$^{-2}$]"
if show_ions:
ax.plot(profile.x_mm, yi, linestyle="--", linewidth=1.0, label=f"{profile.label} ions")
if show_electrons:
ax.plot(profile.x_mm, ye, linestyle=":", linewidth=1.1, label=f"{profile.label} electrons")
if show_total:
total = yi + ye
ax.plot(profile.x_mm, total, linewidth=1.6, label=f"{profile.label} total")
if fit is not None:
params, fit_xy = fit
y_fit = fit_xy[:, 1]/1.0e6 if channel == "energy" else fit_xy[:, 1]
ax.plot(fit_xy[:, 0], y_fit, color=color_base, linewidth=1.2, alpha=0.85)
ax.text(
0.02,
0.92 if profile.label == "Outer" else 0.82,
f"{profile.label}: lambda={params[2]*1.0e3:.2f} mm, S={params[3]*1.0e3:.2f} mm",
transform=ax.transAxes,
fontsize=8,
color=color_base,
)
ax.set_ylabel(ylabel)