"""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, Callable
import matplotlib.pyplot as plt
import numpy as np
from .catalog import resolve_catalog_interval_steps
from .heatdiag import HeatDiag
from .simulation import Simulation
E_CHARGE_C = 1.602176634e-19
_HEATDIAG_INTERVAL_VARS = list(HeatDiag.MINIMAL_INTERVAL_VARS)
[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]
@dataclass
class DivertorLambdaQPoint:
"""One converged total-energy Eich width for a divertor target and time."""
label: str
step: int
time_s: float
lambda_q_mm: float
s_mm: float
[docs]
@dataclass
class DivertorTargetTotalsPoint:
"""One integrated particle and power total for a divertor target and time."""
label: str
step: int
time_s: float
particle_i_rate_per_s: float
particle_e_rate_per_s: float
power_i_w: float
power_e_w: float
@property
def particle_rate_per_s(self) -> float:
"""Return the total ion-plus-electron particle rate."""
return self.particle_i_rate_per_s + self.particle_e_rate_per_s
@property
def power_w(self) -> float:
"""Return the total ion-plus-electron power."""
return self.power_i_w + self.power_e_w
[docs]
@dataclass
class DivertorLoadMap:
"""Mapped divertor particle or energy flux-density map."""
target_label: str
channel: str
component: str
view: str
x_mm: np.ndarray
y_values: np.ndarray
values: np.ndarray
[docs]
@dataclass
class AnalysisHooks:
"""Optional callbacks for orchestration-level progress and cancellation.
The divertor helper remains usable from ordinary Python workflows without
these hooks. Callers that want progress reporting or cooperative
cancellation can provide one or more callbacks.
"""
check_cancel: Callable[[], None] | None = None
phase_callback: Callable[[str, dict[str, Any] | None], None] | None = None
log_callback: Callable[[str], None] | None = None
def _hooks_check_cancel(hooks: AnalysisHooks | None) -> None:
"""Invoke the cooperative cancellation callback when provided."""
if hooks is not None and callable(hooks.check_cancel):
hooks.check_cancel()
def _hooks_phase(
hooks: AnalysisHooks | None,
phase: str,
meta: dict[str, Any] | None = None,
) -> None:
"""Report one high-level analysis phase when provided."""
if hooks is not None and callable(hooks.phase_callback):
hooks.phase_callback(str(phase), dict(meta or {}))
def _hooks_log(hooks: AnalysisHooks | None, message: str) -> None:
"""Append one short analysis diagnostic line when provided."""
if hooks is not None and callable(hooks.log_callback):
hooks.log_callback(str(message))
[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,
hooks: AnalysisHooks | None = None,
) -> 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.
hooks : AnalysisHooks | None
Optional orchestration callbacks for cooperative cancellation and
high-level phase or log reporting. Ordinary scripts can ignore this.
Returns
-------
list[DivertorEichProfile]
One profile for each requested target.
"""
data_dir = Path(data_dir).expanduser().resolve()
sim = simulation
_hooks_phase(hooks, "simulation_setup", {"reuse": sim is not None})
_hooks_check_cancel(hooks)
if sim is None:
_hooks_log(hooks, f"Constructing Simulation from {data_dir}")
sim = Simulation(directories=[str(data_dir)], is_stellarator=False, sim_is_axisymmetric=True)
_hooks_phase(hooks, "plane_setup")
_hooks_check_cancel(hooks)
plane = sim.mesh.get_plane(0)
_hooks_phase(hooks, "heatdiag_step_selection", {"time_window": time_window is not None})
_hooks_check_cancel(hooks)
read_steps = resolve_catalog_interval_steps(
sim.catalog,
product_key="xgc.heatdiag2.bp",
coordinate_variable="time",
time_window=time_window,
selected_frame_index=selected_frame_index,
)
if read_steps:
_hooks_log(hooks, f"HeatDiag read plan uses {len(read_steps)} selected steps.")
else:
_hooks_log(hooks, "HeatDiag read plan uses the HeatDiag default step selection.")
variables = list(_HEATDIAG_INTERVAL_VARS)
if include_sheath:
variables.extend(["e_potential", "i_potential"])
_hooks_phase(hooks, "heatdiag_reader_setup", {"variable_count": len(variables)})
_hooks_check_cancel(hooks)
heatdiag = HeatDiag(
simulation=sim,
data_dir=sim.data_directory,
variables=variables,
steps=read_steps or None,
)
rates = _load_rates(
heatdiag,
simulation=sim,
include_sheath=include_sheath,
time_window=time_window,
selected_frame_index=None if read_steps else selected_frame_index,
hooks=hooks,
)
_hooks_phase(hooks, "wall_geometry")
_hooks_check_cancel(hooks)
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]
_hooks_phase(hooks, "midplane_maps")
_hooks_check_cancel(hooks)
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
_hooks_phase(hooks, "target_split", {"psi_min": psi_window[0], "psi_max": psi_window[1]})
_hooks_check_cancel(hooks)
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)
_hooks_phase(hooks, "profile_build")
_hooks_check_cancel(hooks)
profiles = []
if show_outer:
_hooks_check_cancel(hooks)
_hooks_log(hooks, "Building outer target profile.")
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:
_hooks_check_cancel(hooks)
_hooks_log(hooks, "Building inner target profile.")
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.")
_hooks_phase(hooks, "complete", {"profiles": len(profiles)})
return profiles
[docs]
def compute_divertor_lambda_q_timeseries(
data_dir: str | Path,
*,
simulation: Any | None = None,
include_sheath: bool = False,
step_range: tuple[float, float] | None = None,
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,
hooks: AnalysisHooks | None = None,
) -> list[DivertorLambdaQPoint]:
"""
Compute total-energy Eich ``lambda_q`` values over heatdiag intervals.
Parameters
----------
data_dir : str | Path
XGC output directory containing ``xgc.heatdiag2.bp`` and 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.
step_range : tuple[float, float] | None
Optional inclusive logical-step range. This takes precedence over
``time_window`` when both are provided.
time_window : tuple[float, float] | None
Optional physical-time window in seconds. Intervals overlapping the
window are fitted when ``step_range`` is omitted.
selected_frame_index : int | None
GUI-style selected heatdiag frame used when no range is provided. If no
range and no selected frame are provided, all valid heatdiag intervals
are fitted.
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 before fitting.
show_outer, show_inner : bool
Select which target branches to fit.
hooks : AnalysisHooks | None
Optional orchestration callbacks for cancellation, progress, and logs.
Returns
-------
list[DivertorLambdaQPoint]
One point for each target/time whose total-energy Eich fit converged.
"""
data_dir = Path(data_dir).expanduser().resolve()
sim = simulation
_hooks_phase(hooks, "simulation_setup", {"reuse": sim is not None})
_hooks_check_cancel(hooks)
if sim is None:
_hooks_log(hooks, f"Constructing Simulation from {data_dir}")
sim = Simulation(directories=[str(data_dir)], is_stellarator=False, sim_is_axisymmetric=True)
if not show_outer and not show_inner:
raise ValueError("Both inner and outer target calculations are disabled.")
_hooks_phase(hooks, "plane_setup")
_hooks_check_cancel(hooks)
plane = sim.mesh.get_plane(0)
_hooks_phase(hooks, "heatdiag_step_selection", {"step_range": step_range is not None, "time_window": time_window is not None})
logical_steps, catalog_time = _catalog_heatdiag_time_axis(sim.catalog)
end_indices = _select_heatdiag_interval_end_indices(
logical_steps,
catalog_time,
step_range=step_range,
time_window=None if step_range is not None else time_window,
selected_frame_index=selected_frame_index,
)
if end_indices.size == 0:
raise ValueError("The requested step/time range did not select any positive heatdiag intervals.")
read_indices = np.unique(np.concatenate([end_indices - 1, end_indices]))
read_steps = [int(logical_steps[i]) for i in read_indices]
_hooks_log(hooks, f"HeatDiag read plan uses {len(read_steps)} steps for {end_indices.size} intervals.")
variables = list(_HEATDIAG_INTERVAL_VARS)
if include_sheath:
variables.extend(["e_potential", "i_potential"])
_hooks_phase(hooks, "heatdiag_reader_setup", {"variable_count": len(variables), "step_count": len(read_steps)})
_hooks_check_cancel(hooks)
heatdiag = HeatDiag(
simulation=sim,
data_dir=sim.data_directory,
variables=variables,
steps=read_steps,
)
_hooks_phase(hooks, "wall_geometry")
_hooks_check_cancel(hooks)
r, z, psi_n = _prepare_wall_geometry(heatdiag, plane)
_hooks_phase(hooks, "midplane_maps")
_hooks_check_cancel(hooks)
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
_hooks_phase(hooks, "target_split", {"psi_min": psi_window[0], "psi_max": psi_window[1]})
_hooks_check_cancel(hooks)
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)
heatdiag_steps = np.asarray(sorted(heatdiag.list_step_indices("time")), dtype=int)
local_index_by_step = {int(step): i for i, step in enumerate(heatdiag_steps)}
time_local = np.asarray(heatdiag.get_array("time"), dtype=float).reshape(-1)
points: list[DivertorLambdaQPoint] = []
_hooks_phase(hooks, "lambda_q_fit", {"interval_count": int(end_indices.size)})
for count, end_index in enumerate(end_indices, start=1):
_hooks_check_cancel(hooks)
step = int(logical_steps[int(end_index)])
prev_step = int(logical_steps[int(end_index) - 1])
if step not in local_index_by_step or prev_step not in local_index_by_step:
_hooks_log(hooks, f"Skipping heatdiag step {step}: interval endpoints were not loaded.")
continue
local_end = local_index_by_step[step]
local_prev = local_index_by_step[prev_step]
if local_end >= time_local.size or local_prev >= time_local.size:
_hooks_log(hooks, f"Skipping heatdiag step {step}: interval time endpoint is unavailable.")
continue
dt = float(time_local[local_end] - time_local[local_prev])
if not np.isfinite(dt) or dt <= 0.0:
_hooks_log(hooks, f"Skipping heatdiag step {step}: non-positive interval duration.")
continue
_hooks_phase(hooks, "lambda_q_fit", {"index": count, "count": int(end_indices.size), "step": step})
rates = _load_interval_rates(
heatdiag,
simulation=sim,
include_sheath=include_sheath,
local_end_index=local_end,
dt=dt,
hooks=hooks,
)
nseg = min(r.size, *(np.asarray(value).shape[1] for value in rates.values()))
interval_points = []
if show_outer:
interval_points.extend(
_lambda_q_points_for_target(
"Outer",
step,
float(time_local[local_end]),
outer[:nseg],
x_out[:nseg],
r_abs_out[:nseg],
psi_n[:nseg],
psi_mid,
r_mid_out,
r_sep_out,
"outer",
rates,
fit_window_mm=fit_window_mm,
smoothing_sigma_mm=smoothing_sigma_mm,
hooks=hooks,
)
)
if show_inner:
interval_points.extend(
_lambda_q_points_for_target(
"Inner",
step,
float(time_local[local_end]),
inner[:nseg],
x_in[:nseg],
r_abs_in[:nseg],
psi_n[:nseg],
psi_mid,
r_mid_in,
r_sep_in,
"inner",
rates,
fit_window_mm=fit_window_mm,
smoothing_sigma_mm=smoothing_sigma_mm,
hooks=hooks,
)
)
if not interval_points:
_hooks_log(hooks, f"No converged total-energy Eich fit for heatdiag step {step}.")
points.extend(interval_points)
_hooks_phase(hooks, "complete", {"points": len(points)})
return points
[docs]
def compute_divertor_target_totals_timeseries(
data_dir: str | Path,
*,
simulation: Any | None = None,
include_sheath: bool = False,
step_range: tuple[float, float] | None = None,
time_window: tuple[float, float] | None = None,
selected_frame_index: int | None = None,
psi_window: tuple[float, float] = (0.97, 1.12),
show_outer: bool = True,
show_inner: bool = True,
show_total: bool = True,
show_wall: bool = True,
hooks: AnalysisHooks | None = None,
) -> list[DivertorTargetTotalsPoint]:
"""
Compute integrated divertor and wall particle/power totals over time.
Parameters
----------
data_dir : str | Path
XGC output directory containing ``xgc.heatdiag2.bp`` and 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.
step_range : tuple[float, float] | None
Optional inclusive logical-step range. This takes precedence over
``time_window`` when both are provided.
time_window : tuple[float, float] | None
Optional physical-time window in seconds. Intervals overlapping the
window are integrated when ``step_range`` is omitted.
selected_frame_index : int | None
GUI-style selected heatdiag frame used when no range is provided. If no
range and no selected frame are provided, all valid heatdiag intervals
are integrated.
psi_window : tuple[float, float]
Normalized-poloidal-flux range used to select inner/outer divertor
wall points, matching :func:`compute_divertor_eich_profiles`.
show_outer, show_inner, show_total, show_wall : bool
Select which traces to return. The ``Total`` trace sums the inner and
outer divertor masks. The wall control trace sums all heatdiag wall
segments after the reader has removed the garbage bin.
hooks : AnalysisHooks | None
Optional orchestration callbacks for cancellation, progress, and logs.
Returns
-------
list[DivertorTargetTotalsPoint]
One point for each requested trace and heatdiag interval.
"""
data_dir = Path(data_dir).expanduser().resolve()
sim = simulation
_hooks_phase(hooks, "simulation_setup", {"reuse": sim is not None})
_hooks_check_cancel(hooks)
if sim is None:
_hooks_log(hooks, f"Constructing Simulation from {data_dir}")
sim = Simulation(directories=[str(data_dir)], is_stellarator=False, sim_is_axisymmetric=True)
if not show_outer and not show_inner and not show_total and not show_wall:
raise ValueError("All divertor and wall total traces are disabled.")
_hooks_phase(hooks, "plane_setup")
_hooks_check_cancel(hooks)
plane = sim.mesh.get_plane(0)
_hooks_phase(hooks, "heatdiag_step_selection", {"step_range": step_range is not None, "time_window": time_window is not None})
logical_steps, catalog_time = _catalog_heatdiag_time_axis(sim.catalog)
end_indices = _select_heatdiag_interval_end_indices(
logical_steps,
catalog_time,
step_range=step_range,
time_window=None if step_range is not None else time_window,
selected_frame_index=selected_frame_index,
)
if end_indices.size == 0:
raise ValueError("The requested step/time range did not select any positive heatdiag intervals.")
read_indices = np.unique(np.concatenate([end_indices - 1, end_indices]))
read_steps = [int(logical_steps[i]) for i in read_indices]
_hooks_log(hooks, f"HeatDiag read plan uses {len(read_steps)} steps for {end_indices.size} intervals.")
variables = list(_HEATDIAG_INTERVAL_VARS)
if include_sheath:
variables.extend(["e_potential", "i_potential"])
_hooks_phase(hooks, "heatdiag_reader_setup", {"variable_count": len(variables), "step_count": len(read_steps)})
_hooks_check_cancel(hooks)
heatdiag = HeatDiag(
simulation=sim,
data_dir=sim.data_directory,
variables=variables,
steps=read_steps,
)
_hooks_phase(hooks, "wall_geometry")
_hooks_check_cancel(hooks)
r, z, psi_n = _prepare_wall_geometry(heatdiag, plane)
_hooks_phase(hooks, "target_split", {"psi_min": psi_window[0], "psi_max": psi_window[1]})
_hooks_check_cancel(hooks)
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)
heatdiag_steps = np.asarray(sorted(heatdiag.list_step_indices("time")), dtype=int)
local_index_by_step = {int(step): i for i, step in enumerate(heatdiag_steps)}
time_local = np.asarray(heatdiag.get_array("time"), dtype=float).reshape(-1)
points: list[DivertorTargetTotalsPoint] = []
_hooks_phase(hooks, "target_totals", {"interval_count": int(end_indices.size)})
for count, end_index in enumerate(end_indices, start=1):
_hooks_check_cancel(hooks)
step = int(logical_steps[int(end_index)])
prev_step = int(logical_steps[int(end_index) - 1])
if step not in local_index_by_step or prev_step not in local_index_by_step:
_hooks_log(hooks, f"Skipping heatdiag step {step}: interval endpoints were not loaded.")
continue
local_end = local_index_by_step[step]
local_prev = local_index_by_step[prev_step]
if local_end >= time_local.size or local_prev >= time_local.size:
_hooks_log(hooks, f"Skipping heatdiag step {step}: interval time endpoint is unavailable.")
continue
dt = float(time_local[local_end] - time_local[local_prev])
if not np.isfinite(dt) or dt <= 0.0:
_hooks_log(hooks, f"Skipping heatdiag step {step}: non-positive interval duration.")
continue
_hooks_phase(hooks, "target_totals", {"index": count, "count": int(end_indices.size), "step": step})
rates = _load_interval_rates(
heatdiag,
simulation=sim,
include_sheath=include_sheath,
local_end_index=local_end,
dt=dt,
hooks=hooks,
)
nseg = min(r.size, *(np.asarray(value).shape[1] for value in rates.values()))
time_s = float(time_local[local_end])
if show_outer:
points.append(_target_totals_point("Outer", step, time_s, outer[:nseg], rates))
if show_inner:
points.append(_target_totals_point("Inner", step, time_s, inner[:nseg], rates))
if show_total:
points.append(_target_totals_point("Total", step, time_s, (outer[:nseg] | inner[:nseg]), rates))
if show_wall:
points.append(_target_totals_point("Control", step, time_s, np.ones(nseg, dtype=bool), rates))
_hooks_phase(hooks, "complete", {"points": len(points)})
return points
[docs]
def compute_divertor_load_map(
data_dir: str | Path,
*,
simulation: Any | None = None,
include_sheath: bool = False,
view: str = "time",
channel: str = "energy",
target: str = "outer",
component: str = "total",
step_range: tuple[float, float] | None = None,
time_window: tuple[float, float] | None = None,
selected_frame_index: int | None = None,
psi_window: tuple[float, float] = (0.97, 1.12),
hooks: AnalysisHooks | None = None,
) -> DivertorLoadMap:
"""
Compute a mapped divertor particle or energy flux-density map.
Parameters
----------
data_dir : str | Path
XGC output directory containing ``xgc.heatdiag2.bp`` and 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.
view : {"time", "toroidal"}
Choose whether the second map coordinate is heatdiag time or toroidal
angle. The time view uses all selected intervals. The toroidal view
duration-averages selected intervals and requires toroidal resolution.
channel : {"particle", "energy"}
Select particle flux density or energy flux density.
target : {"outer", "inner"}
Select the divertor branch, using the same split as the Eich profile
helper.
component : {"ion", "electron", "total"}
Select species component.
step_range : tuple[float, float] | None
Optional inclusive logical-step range. This takes precedence over
``time_window`` when both are provided.
time_window : tuple[float, float] | None
Optional physical-time window in seconds. Intervals overlapping the
window are selected when ``step_range`` is omitted.
selected_frame_index : int | None
GUI-style selected heatdiag frame used when no range is provided. For
the time view, no range and no frame means all valid intervals. For the
toroidal view, no range and no frame means the latest valid interval.
psi_window : tuple[float, float]
Normalized-poloidal-flux range used to select divertor wall points.
hooks : AnalysisHooks | None
Optional orchestration callbacks for cancellation, progress, and logs.
Returns
-------
DivertorLoadMap
Mapped load values on ``(y, Delta_sep)`` coordinates. Energy values are
stored in SI units, W m^-2. Plotting converts them to MW m^-2.
"""
view = str(view or "time").strip().lower()
channel = str(channel or "energy").strip().lower()
target = str(target or "outer").strip().lower()
component = str(component or "total").strip().lower()
if view not in {"time", "toroidal"}:
raise ValueError("view must be 'time' or 'toroidal'.")
if channel not in {"particle", "energy"}:
raise ValueError("channel must be 'particle' or 'energy'.")
if target not in {"outer", "inner"}:
raise ValueError("target must be 'outer' or 'inner'.")
if component not in {"ion", "electron", "total"}:
raise ValueError("component must be 'ion', 'electron', or 'total'.")
data_dir = Path(data_dir).expanduser().resolve()
sim = simulation
_hooks_phase(hooks, "simulation_setup", {"reuse": sim is not None})
_hooks_check_cancel(hooks)
if sim is None:
_hooks_log(hooks, f"Constructing Simulation from {data_dir}")
sim = Simulation(directories=[str(data_dir)], is_stellarator=False, sim_is_axisymmetric=True)
_hooks_phase(hooks, "plane_setup")
_hooks_check_cancel(hooks)
plane = sim.mesh.get_plane(0)
_hooks_phase(hooks, "heatdiag_step_selection", {"step_range": step_range is not None, "time_window": time_window is not None})
logical_steps, catalog_time = _catalog_heatdiag_time_axis(sim.catalog)
selection_frame = selected_frame_index
if view == "toroidal" and step_range is None and time_window is None and selection_frame is None:
selection_frame = int(logical_steps.size - 1)
end_indices = _select_heatdiag_interval_end_indices(
logical_steps,
catalog_time,
step_range=step_range,
time_window=None if step_range is not None else time_window,
selected_frame_index=selection_frame,
)
if end_indices.size == 0:
raise ValueError("The requested step/time range did not select any positive heatdiag intervals.")
read_indices = np.unique(np.concatenate([end_indices - 1, end_indices]))
read_steps = [int(logical_steps[i]) for i in read_indices]
_hooks_log(hooks, f"HeatDiag read plan uses {len(read_steps)} steps for {end_indices.size} intervals.")
variables = list(_HEATDIAG_INTERVAL_VARS)
if include_sheath:
variables.extend(["e_potential", "i_potential"])
_hooks_phase(hooks, "heatdiag_reader_setup", {"variable_count": len(variables), "step_count": len(read_steps)})
_hooks_check_cancel(hooks)
heatdiag = HeatDiag(
simulation=sim,
data_dir=sim.data_directory,
variables=variables,
steps=read_steps,
)
_hooks_phase(hooks, "wall_geometry")
_hooks_check_cancel(hooks)
r, z, psi_n = _prepare_wall_geometry(heatdiag, plane)
_hooks_phase(hooks, "midplane_maps")
_hooks_check_cancel(hooks)
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.")
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
_hooks_phase(hooks, "target_split", {"psi_min": psi_window[0], "psi_max": psi_window[1]})
_hooks_check_cancel(hooks)
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)
if target == "outer":
label = "Outer"
mask = outer
x_m = x_out
psi_branch = r_mid_out
r_sep = r_sep_out
branch = "outer"
else:
label = "Inner"
mask = inner
x_m = x_in
psi_branch = r_mid_in
r_sep = r_sep_in
branch = "inner"
valid_mask = mask & np.isfinite(x_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, psi_branch, r_sep, branch)
area_full = _target_grid_toroidal_area(x_grid, r_grid)
heatdiag_steps = np.asarray(sorted(heatdiag.list_step_indices("time")), dtype=int)
local_index_by_step = {int(step): i for i, step in enumerate(heatdiag_steps)}
time_local = np.asarray(heatdiag.get_array("time"), dtype=float).reshape(-1)
rows = []
y_values = []
weights = []
_hooks_phase(hooks, "load_map", {"interval_count": int(end_indices.size), "view": view})
for count, end_index in enumerate(end_indices, start=1):
_hooks_check_cancel(hooks)
step = int(logical_steps[int(end_index)])
prev_step = int(logical_steps[int(end_index) - 1])
if step not in local_index_by_step or prev_step not in local_index_by_step:
_hooks_log(hooks, f"Skipping heatdiag step {step}: interval endpoints were not loaded.")
continue
local_end = local_index_by_step[step]
local_prev = local_index_by_step[prev_step]
if local_end >= time_local.size or local_prev >= time_local.size:
_hooks_log(hooks, f"Skipping heatdiag step {step}: interval time endpoint is unavailable.")
continue
dt = float(time_local[local_end] - time_local[local_prev])
if not np.isfinite(dt) or dt <= 0.0:
_hooks_log(hooks, f"Skipping heatdiag step {step}: non-positive interval duration.")
continue
_hooks_phase(hooks, "load_map", {"index": count, "count": int(end_indices.size), "step": step})
rates = _load_interval_rates(
heatdiag,
simulation=sim,
include_sheath=include_sheath,
local_end_index=local_end,
dt=dt,
hooks=hooks,
)
arr = _load_component_rate_array(rates, channel, component)
nseg = min(arr.shape[1], valid_mask.shape[0])
local_mask = valid_mask[:nseg]
if view == "time":
rows.append(_mapped_time_load_row(arr[:, :nseg], local_mask, x_m[:nseg], x_grid, area_full))
y_values.append(float(time_local[local_end]))
else:
rows.append(_mapped_toroidal_load_rows(arr[:, :nseg], local_mask, x_m[:nseg], x_grid, area_full))
weights.append(dt)
if not rows:
raise ValueError("No valid heatdiag intervals were available for the selected load map.")
if view == "time":
values = np.asarray(rows, dtype=float)
y_axis = np.asarray(y_values, dtype=float)
order = np.argsort(y_axis)
values = values[order]
y_axis = y_axis[order]
else:
first = np.asarray(rows[0], dtype=float)
if first.shape[0] <= 1:
raise ValueError("Toroidal-angle view requires heatdiag data with nphi > 1.")
stack = np.asarray(rows, dtype=float)
weights_arr = np.asarray(weights, dtype=float)
values = np.average(stack, axis=0, weights=weights_arr)
y_axis = np.linspace(0.0, 360.0, first.shape[0], endpoint=False)
return DivertorLoadMap(
target_label=label,
channel=channel,
component=component,
view=view,
x_mm=x_grid*1.0e3,
y_values=y_axis,
values=values,
)
[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
[docs]
def plot_divertor_lambda_q_timeseries(
points: list[DivertorLambdaQPoint],
*,
ylim: tuple[float, float] | None = None,
show_lines: bool = True,
show_symbols: bool = False,
):
"""
Plot total-energy Eich ``lambda_q`` values versus heatdiag time.
Parameters
----------
points : list[DivertorLambdaQPoint]
Points returned by :func:`compute_divertor_lambda_q_timeseries`.
ylim : tuple[float, float] | None
Optional vertical axis limits in millimeters.
show_lines, show_symbols : bool
Select line, symbol-only, or line-plus-symbol display style.
Returns
-------
matplotlib.figure.Figure
Figure owned by the caller. GUI code can save it to PNG.
"""
if not show_lines and not show_symbols:
raise ValueError("Enable at least one of lines or symbols.")
if not points:
raise ValueError("No converged total-energy Eich fits found in the selected range.")
linestyle = "-" if show_lines else "None"
marker = "o" if show_symbols else None
colors = {"Outer": "#0f5ea8", "Inner": "#a04400"}
fig, ax = plt.subplots(figsize=(7.4, 4.0), dpi=150)
plotted = False
for label in ("Outer", "Inner"):
selected = [point for point in points if point.label == label]
if not selected:
continue
order = np.argsort([point.time_s for point in selected])
time_ms = np.asarray([selected[i].time_s*1.0e3 for i in order], dtype=float)
lambda_q_mm = np.asarray([selected[i].lambda_q_mm for i in order], dtype=float)
ax.plot(
time_ms,
lambda_q_mm,
linestyle=linestyle,
marker=marker,
markersize=4.0,
linewidth=1.5,
color=colors.get(label, "#333333"),
label=label,
)
plotted = True
if not plotted:
raise ValueError("No outer or inner target lambda_q points are available.")
ax.set_title(r"Divertor heat-load width from Eich fits")
ax.set_xlabel("time [ms]")
ax.set_ylabel(r"$\lambda_q$ [mm]")
if ylim is not None:
ax.set_ylim(ylim)
ax.grid(alpha=0.25, linewidth=0.4)
ax.legend(loc="best", fontsize=8)
fig.tight_layout(pad=0.35)
return fig
[docs]
def plot_divertor_target_totals_timeseries(
points: list[DivertorTargetTotalsPoint],
*,
show_particle: bool = True,
show_power: bool = True,
show_outer: bool = True,
show_inner: bool = True,
show_total: bool = True,
show_control: bool = True,
show_ions: bool = True,
show_electrons: bool = True,
show_target_total: bool = True,
xlim_ms: tuple[float, float] | None = None,
particle_ylim: tuple[float, float] | None = None,
power_ylim_mw: tuple[float, float] | None = None,
):
"""
Plot integrated target particle flux and power versus heatdiag time.
Parameters
----------
points : list[DivertorTargetTotalsPoint]
Points returned by :func:`compute_divertor_target_totals_timeseries`.
show_particle, show_power : bool
Select particle-flux and power panels.
show_outer, show_inner, show_total, show_control : bool
Select target/control trace groups. ``Total`` is the inner-plus-outer
divertor sum, while ``Control`` is the full-wall control quantity after
excluding the heatdiag garbage bin.
show_ions, show_electrons, show_target_total : bool
Select ion, electron, and ion-plus-electron component lines for the
inner and outer target groups.
xlim_ms : tuple[float, float] | None
Optional horizontal axis range in milliseconds.
particle_ylim : tuple[float, float] | None
Optional particle-flux vertical range in particles per second.
power_ylim_mw : tuple[float, float] | None
Optional power vertical range in megawatts.
Returns
-------
matplotlib.figure.Figure
Figure owned by the caller. GUI code can save it to PNG.
"""
if not show_particle and not show_power:
raise ValueError("At least one of particle flux or power must be enabled.")
if not show_outer and not show_inner and not show_total and not show_control:
raise ValueError("At least one target/control trace must be enabled.")
if not points:
raise ValueError("No divertor target total points found in the selected range.")
nrows = int(show_particle) + int(show_power)
fig, axes = plt.subplots(nrows=nrows, ncols=1, figsize=(7.4, 3.4*nrows), dpi=150, squeeze=False)
axes_flat = list(axes[:, 0])
panel_axes = []
if show_particle:
panel_axes.append(("particle", axes_flat[len(panel_axes)]))
if show_power:
panel_axes.append(("power", axes_flat[len(panel_axes)]))
trace_specs = _target_totals_trace_specs(
show_outer=show_outer,
show_inner=show_inner,
show_total=show_total,
show_control=show_control,
show_ions=show_ions,
show_electrons=show_electrons,
show_target_total=show_target_total,
)
if not trace_specs:
raise ValueError("Enable at least one target/control or ion/electron/target-total line.")
for panel, ax in panel_axes:
plotted = False
for spec in trace_specs:
selected = [point for point in points if point.label == spec["label"]]
if not selected:
continue
order = np.argsort([point.time_s for point in selected])
time_ms = np.asarray([selected[i].time_s*1.0e3 for i in order], dtype=float)
if panel == "particle":
y = np.asarray([_target_totals_particle_value(selected[i], spec["component"]) for i in order], dtype=float)
else:
y = np.asarray([_target_totals_power_value(selected[i], spec["component"])/1.0e6 for i in order], dtype=float)
ax.plot(
time_ms,
y,
linestyle=spec["linestyle"],
linewidth=1.5,
color=spec["color"],
label=spec["plot_label"],
)
plotted = True
if not plotted:
raise ValueError("No enabled target/control traces are available.")
ax.set_xlabel("time [ms]")
if panel == "particle":
ax.set_ylabel("Particle flux [ptl/s]")
if particle_ylim is not None:
ax.set_ylim(particle_ylim)
else:
ax.set_ylabel("Power [MW]")
if power_ylim_mw is not None:
ax.set_ylim(power_ylim_mw)
if xlim_ms is not None:
ax.set_xlim(xlim_ms)
ax.grid(alpha=0.25, linewidth=0.4)
ax.legend(loc="best", fontsize=8)
fig.suptitle("Divertor heatdiag target totals", fontsize=11)
fig.tight_layout(pad=0.35)
return fig
[docs]
def plot_divertor_load_map(
load_map: DivertorLoadMap,
*,
xlim: tuple[float, float] | None = None,
ylim: tuple[float, float] | None = None,
clim: tuple[float, float] | None = None,
contour_count: int = 256,
cmap: str = "seismic",
):
"""
Plot a mapped divertor load map as filled contours without contour lines.
Parameters
----------
load_map : DivertorLoadMap
Map returned by :func:`compute_divertor_load_map`.
xlim, ylim : tuple[float, float] | None
Optional display ranges. ``xlim`` is in millimeters. ``ylim`` is in
milliseconds for the time view and degrees for the toroidal view.
clim : tuple[float, float] | None
Optional color scale in displayed units. Energy maps are displayed in
MW m^-2; particle maps are displayed in ptl m^-2 s^-1.
contour_count : int
Number of filled contour levels. Values below two are promoted to two.
cmap : str
Matplotlib colormap name. The default is ``"seismic"``.
Returns
-------
matplotlib.figure.Figure
Figure owned by the caller. GUI code can save it to PNG.
"""
contour_count = max(2, int(contour_count))
x = np.asarray(load_map.x_mm, dtype=float)
y = np.asarray(load_map.y_values, dtype=float)
z = np.asarray(load_map.values, dtype=float)
if z.ndim != 2:
raise ValueError(f"Expected a 2D load map, got shape={z.shape}")
if load_map.channel == "energy":
z_plot = z/1.0e6
cbar_label = r"Energy flux density [MW m$^{-2}$]"
else:
z_plot = z
cbar_label = r"Particle flux density [ptl m$^{-2}$ s$^{-1}$]"
if load_map.view == "time":
y_plot = y*1.0e3
y_label = "time [ms]"
else:
y_plot = y
y_label = "toroidal angle [deg]"
finite = z_plot[np.isfinite(z_plot)]
if finite.size == 0:
raise ValueError("Load map contains no finite values.")
if clim is None:
zmin = float(np.nanmin(finite))
zmax = float(np.nanmax(finite))
else:
zmin, zmax = clim
if not np.isfinite(zmin) or not np.isfinite(zmax):
raise ValueError("Color scale limits must be finite.")
if zmin == zmax:
pad = max(abs(zmin)*1.0e-6, 1.0e-12)
zmin -= pad
zmax += pad
if zmin > zmax:
zmin, zmax = zmax, zmin
levels = np.linspace(zmin, zmax, contour_count)
fig, ax = plt.subplots(figsize=(7.4, 4.4), dpi=150)
contour = ax.contourf(x, y_plot, z_plot, levels=levels, cmap=cmap, extend="both")
ax.set_xlabel(r"$\Delta_{sep}$ [mm]")
ax.set_ylabel(y_label)
if xlim is not None:
ax.set_xlim(xlim)
if ylim is not None:
ax.set_ylim(ylim)
title_channel = "energy" if load_map.channel == "energy" else "particle"
ax.set_title(f"{load_map.target_label} divertor {load_map.component} {title_channel} load")
cbar = fig.colorbar(contour, ax=ax)
cbar.set_label(cbar_label)
fig.tight_layout(pad=0.35)
return fig
def _target_totals_trace_specs(
*,
show_outer: bool,
show_inner: bool,
show_total: bool,
show_control: bool,
show_ions: bool,
show_electrons: bool,
show_target_total: bool,
) -> list[dict[str, str]]:
"""Return line specifications for target-total plots."""
specs: list[dict[str, str]] = []
def add_target(label: str, color: str) -> None:
if show_ions:
specs.append({"label": label, "component": "ion", "plot_label": f"{label} ions", "color": color, "linestyle": "--"})
if show_electrons:
specs.append({"label": label, "component": "electron", "plot_label": f"{label} electrons", "color": color, "linestyle": ":"})
if show_target_total:
specs.append({"label": label, "component": "total", "plot_label": f"{label} total", "color": color, "linestyle": "-"})
if show_outer:
add_target("Outer", "#0f5ea8")
if show_inner:
add_target("Inner", "#a04400")
if show_total:
specs.append({"label": "Total", "component": "total", "plot_label": "Total divertor", "color": "#20845a", "linestyle": "-"})
if show_control:
specs.append({"label": "Control", "component": "total", "plot_label": "Wall control", "color": "#333333", "linestyle": "--"})
return specs
def _target_totals_particle_value(point: DivertorTargetTotalsPoint, component: str) -> float:
"""Return one particle-flux component from an integrated totals point."""
if component == "ion":
return point.particle_i_rate_per_s
if component == "electron":
return point.particle_e_rate_per_s
return point.particle_rate_per_s
def _target_totals_power_value(point: DivertorTargetTotalsPoint, component: str) -> float:
"""Return one power component from an integrated totals point."""
if component == "ion":
return point.power_i_w
if component == "electron":
return point.power_e_w
return point.power_w
def _load_component_rate_array(rates: dict[str, np.ndarray], channel: str, component: str) -> np.ndarray:
"""Return one particle or energy rate array for a species component."""
if channel == "particle":
ion = np.asarray(rates["particle_i"], dtype=float)
electron = np.asarray(rates["particle_e"], dtype=float)
else:
ion = np.asarray(rates["energy_i"], dtype=float)
electron = np.asarray(rates["energy_e"], dtype=float)
if component == "ion":
return ion
if component == "electron":
return electron
return ion + electron
def _target_grid_toroidal_area(x_grid: np.ndarray, r_grid: np.ndarray) -> np.ndarray:
"""Return full-torus surface areas for mapped target grid cells."""
widths = np.diff(_cell_edges_from_centers(np.asarray(x_grid, dtype=float)))
area = widths * 2.0 * np.pi * np.maximum(np.abs(np.asarray(r_grid, dtype=float)), 1.0e-6)
return np.maximum(np.abs(area), 1.0e-20)
def _mapped_time_load_row(
arr: np.ndarray,
local_mask: np.ndarray,
x_m: np.ndarray,
x_grid: np.ndarray,
area_full: np.ndarray,
) -> np.ndarray:
"""Map one interval's toroidally summed target loads to flux density."""
if arr.ndim != 2:
raise ValueError(f"Expected rate array with shape (nphi,nseg), got {arr.shape}")
toroidal_sum = np.sum(arr[:, local_mask], axis=0)
load, _width = _conservative_remap(x_m[local_mask], toroidal_sum, x_grid)
return load / area_full
def _mapped_toroidal_load_rows(
arr: np.ndarray,
local_mask: np.ndarray,
x_m: np.ndarray,
x_grid: np.ndarray,
area_full: np.ndarray,
) -> np.ndarray:
"""Map one interval's target loads by toroidal plane to flux density."""
if arr.ndim != 2:
raise ValueError(f"Expected rate array with shape (nphi,nseg), got {arr.shape}")
nphi = int(arr.shape[0])
if nphi <= 1:
raise ValueError("Toroidal-angle view requires heatdiag data with nphi > 1.")
area_per_phi = area_full / float(nphi)
rows = []
for iphi in range(nphi):
load, _width = _conservative_remap(x_m[local_mask], arr[iphi, local_mask], x_grid)
rows.append(load / area_per_phi)
return np.asarray(rows, dtype=float)
def _catalog_heatdiag_time_axis(catalog) -> tuple[np.ndarray, np.ndarray]:
"""Return heatdiag logical steps and endpoint times from a catalog."""
records = list(catalog.available_steps("xgc.heatdiag2.bp", "time"))
if len(records) < 2:
raise ValueError("Need at least two heatdiag time samples to compute interval fits.")
logical_steps = np.asarray([int(record.logical_step) for record in records], dtype=int)
times = [_step_fragment_time(record.selected_fragment) for record in records]
if any(value is None for value in times):
result = catalog.read_arrays(
"xgc.heatdiag2.bp",
variables=["time"],
steps=[int(step) for step in logical_steps],
)
values_by_step = result.arrays.get("time", {})
times = [
float(np.asarray(values_by_step[int(step)]).reshape(-1)[0])
for step in logical_steps
]
return logical_steps, np.asarray(times, dtype=float)
def _step_fragment_time(fragment) -> float | None:
"""Return the time cached in one catalog step fragment, if present."""
value = getattr(fragment, "time", None)
if value is None:
return None
try:
return float(value)
except (TypeError, ValueError):
return None
def _select_heatdiag_interval_end_indices(
logical_steps: np.ndarray,
time_s: np.ndarray,
*,
step_range: tuple[float, float] | None,
time_window: tuple[float, float] | None,
selected_frame_index: int | None,
) -> np.ndarray:
"""Select heatdiag interval-ending sample indices for a range request."""
if logical_steps.size != time_s.size:
raise ValueError("Heatdiag step and time axes have inconsistent lengths.")
if time_s.size < 2:
raise ValueError("Need at least two heatdiag samples to compute interval fits.")
valid_interval = np.diff(time_s) > 0.0
end_indices = np.arange(1, time_s.size, dtype=int)
if step_range is not None:
lo, hi = step_range
mask = (logical_steps[1:] >= math.ceil(lo)) & (logical_steps[1:] <= math.floor(hi)) & valid_interval
return end_indices[mask]
if time_window is not None:
t0, t1 = time_window
mask = (time_s[1:] > float(t0)) & (time_s[:-1] < float(t1)) & valid_interval
return end_indices[mask]
if selected_frame_index is None:
return end_indices[valid_interval]
idx = max(1, min(int(selected_frame_index), time_s.size - 1))
if not valid_interval[idx - 1]:
valid = end_indices[valid_interval]
if valid.size == 0:
raise ValueError("Heatdiag time samples do not contain positive intervals.")
idx = int(valid[np.argmin(np.abs(valid - idx))])
return np.asarray([idx], dtype=int)
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,
hooks: AnalysisHooks | None = 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.
"""
_hooks_phase(hooks, "heatdiag_time_axis")
_hooks_check_cancel(hooks)
tmask, time, _dt = _get_sorted_unique_time_indices(heatdiag)
_hooks_log(hooks, f"HeatDiag time samples available: {time.shape[0]}")
_hooks_phase(hooks, "heatdiag_arrays", {"include_sheath": include_sheath})
_hooks_check_cancel(hooks)
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:
_hooks_phase(hooks, "heatdiag_sheath")
_hooks_check_cancel(hooks)
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
_hooks_phase(hooks, "heatdiag_rate_average")
_hooks_check_cancel(hooks)
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 _interval_rate(values: np.ndarray, local_end_index: int, dt: float, name: str) -> np.ndarray:
"""Return one accumulated heatdiag endpoint converted to an interval rate."""
arr = np.asarray(values, dtype=float)
if arr.ndim == 2:
arr = arr[:, np.newaxis, :]
if arr.ndim != 3:
raise ValueError(f"Expected heatdiag history '{name}' with shape (time,nphi,nseg), got {arr.shape}")
if local_end_index < 0 or local_end_index >= arr.shape[0]:
raise IndexError(f"local_end_index={local_end_index} out of range for heatdiag history '{name}'.")
return arr[local_end_index] / float(dt)
def _load_interval_rates(
heatdiag,
*,
simulation,
include_sheath: bool,
local_end_index: int,
dt: float,
hooks: AnalysisHooks | None = None,
) -> dict[str, np.ndarray]:
"""Load particle and energy rates for one already-read heatdiag interval."""
if not np.isfinite(dt) or dt <= 0.0:
raise ValueError("Heatdiag interval duration must be positive.")
_hooks_phase(hooks, "heatdiag_arrays", {"include_sheath": include_sheath})
_hooks_check_cancel(hooks)
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:
_hooks_phase(hooks, "heatdiag_sheath")
_hooks_check_cancel(hooks)
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
_hooks_phase(hooks, "heatdiag_rate_average")
_hooks_check_cancel(hooks)
return {
"particle_e": _interval_rate(np.asarray(heatdiag.get_array("e_number")), local_end_index, dt, "e_number"),
"particle_i": _interval_rate(np.asarray(heatdiag.get_array("i_number")), local_end_index, dt, "i_number"),
"energy_e": _interval_rate(e_energy, local_end_index, dt, "electron energy"),
"energy_i": _interval_rate(i_energy, local_end_index, dt, "ion energy"),
}
def _lambda_q_points_for_target(
label: str,
step: int,
time_s: float,
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,
hooks: AnalysisHooks | None = None,
) -> list[DivertorLambdaQPoint]:
"""Build one target profile and return its converged energy-fit width."""
try:
profile = _build_target_profile(
label,
mask,
x_m,
r_abs_m,
psi_n,
psi_mid,
r_mid,
r_sep,
branch,
rates,
fit_window_mm=fit_window_mm,
smoothing_sigma_mm=smoothing_sigma_mm,
)
except Exception as exc:
_hooks_log(hooks, f"Skipping {label.lower()} target at heatdiag step {step}: {exc}")
return []
if profile.fit_energy is None:
return []
params, _fit_xy = profile.fit_energy
lambda_q_mm = float(np.asarray(params, dtype=float)[2])*1.0e3
s_mm = float(np.asarray(params, dtype=float)[3])*1.0e3
if not np.isfinite(lambda_q_mm) or lambda_q_mm <= 0.0:
return []
return [
DivertorLambdaQPoint(
label=label,
step=int(step),
time_s=float(time_s),
lambda_q_mm=lambda_q_mm,
s_mm=s_mm,
)
]
def _target_totals_point(
label: str,
step: int,
time_s: float,
mask: np.ndarray,
rates: dict[str, np.ndarray],
) -> DivertorTargetTotalsPoint:
"""Return one integrated particle and power total for a wall mask."""
mask = np.asarray(mask, dtype=bool)
def masked_sum(name: str) -> float:
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], mask.shape[0])
if n <= 0:
return 0.0
return float(np.nansum(arr[:, :n][:, mask[:n]]))
return DivertorTargetTotalsPoint(
label=str(label),
step=int(step),
time_s=float(time_s),
particle_i_rate_per_s=masked_sum("particle_i"),
particle_e_rate_per_s=masked_sum("particle_e"),
power_i_w=masked_sum("energy_i"),
power_e_w=masked_sum("energy_e"),
)
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)