"""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 numpy as np
from xgc_analysis.data_access.catalog import resolve_catalog_interval_steps
from xgc_analysis.data_access.products.heatdiag import HeatDiag
from xgc_analysis.simulation import Simulation
from xgc_analysis.analysis.geometry import (
finite_pair_count as _finite_pair_count,
interp_radius_at_psi as _interp_radius_at_psi,
interp_radius_map as _interp_radius_map,
midplane_radius_maps as _midplane_radius_maps,
)
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.
Attributes
----------
label, step, time_s, lambda_q_mm, s_mm
Target name, logical step, physics time in seconds, and fitted widths in millimeters.
"""
label: str
step: int
time_s: float
lambda_q_mm: float
s_mm: float
[docs]
@dataclass
class DivertorWidthScalings:
"""Static Eich #14 and Chang Eq. (6) widths from Simulation inputs.
Attributes
----------
eich_mm, chang_mm
Eich and Chang upstream width estimates in millimeters.
bpol_midplane_t, minor_radius_m, ion_temperature_ev, ion_name
Equilibrium and main-ion inputs used in the scaling.
"""
eich_mm: float
chang_mm: float
bpol_midplane_t: float
minor_radius_m: float
ion_temperature_ev: float
ion_name: str
[docs]
@dataclass
class DivertorTargetTotalsPoint:
"""One integrated particle and power total for a divertor target and time.
Attributes
----------
label, step, time_s
Target name, logical step, and physics time in seconds.
particle_i_rate_per_s, particle_e_rate_per_s, power_i_w, power_e_w
Integrated species particle rates and powers in SI units.
"""
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.
Attributes
----------
target_label, channel, component, view
Selected target, diagnostic channel, species component, and map orientation.
x_mm, y_values, values
Strike-point distance in millimeters, second map coordinate, and mapped rates.
"""
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.
Attributes
----------
check_cancel, phase_callback, log_callback
Optional callable hooks for cancellation, progress phase, and log messages.
"""
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)
local_index_by_step = _heatdiag_local_index_by_logical_step(heatdiag)
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)
local_index_by_step = _heatdiag_local_index_by_logical_step(heatdiag)
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)
local_index_by_step = _heatdiag_local_index_by_logical_step(heatdiag)
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 estimate_divertor_width_scalings(simulation: Simulation) -> DivertorWidthScalings:
"""Evaluate Eich (2013) #14 and Chang et al. (2021) Eq. (6).
Both scalings predict upstream outer-midplane width in millimeters. Chang
uses the initial main-ion temperature at psi_N=0.99, so its estimate is
constant over a heatdiag time series. The ion poloidal gyroradius uses
``sqrt(2*m_i*T_i)/(|q_i|*B_pol)`` with T_i converted from eV to joules.
Parameters
----------
simulation
Initialized Simulation with equilibrium, mesh, and species.
Returns
-------
DivertorWidthScalings: static predictions and input parameters.
"""
plane = simulation.mesh.get_plane(0)
psi_mid, r_mid_out, _ = _midplane_radius_maps(plane)
r_sep_out = _interp_radius_at_psi(psi_mid, r_mid_out, 1.0, "outer")
br, bz, _ = simulation.magnetic_field.compute_background_field(r_sep_out, 0.0)
bpol = float(np.hypot(br, bz))
# surf_map contains the full midplane-crossing surfaces, excluding the
# separatrix legs that have similar psi_N but a much smaller rmin.
surfaces = np.asarray(plane.surf_map, dtype=int)
psi_surf = np.asarray(plane.psi_surf)[surfaces]/float(plane.x_psi)
minor_radius = _interp_radius_at_psi(psi_surf, np.asarray(plane.rmin)[surfaces], 1.0, "minor-radius")
ions = [species for species in simulation.species if species.charge_C > 0]
if not ions:
raise ValueError("Chang heat-load scaling requires a positive-ion Species.")
ion = ions[0]
profile = ion.initial_profiles.get("temperature")
if profile is None:
raise ValueError("Chang heat-load scaling requires the main ion's initial temperature profile.")
psi_temp = np.asarray(profile.psi_n, dtype=float)
temp_ev = np.asarray(profile.value, dtype=float)
if psi_temp.size < 2 or not np.isfinite(psi_temp).all() or not np.isfinite(temp_ev).all():
raise ValueError("The main ion's initial temperature profile is incomplete.")
order = np.argsort(psi_temp)
if not (psi_temp[order[0]] <= 0.99 <= psi_temp[order[-1]]):
raise ValueError("The main ion's initial temperature profile does not bracket psi_N=0.99.")
ti_ev = float(np.interp(0.99, psi_temp[order], temp_ev[order]))
if not all(np.isfinite(value) and value > 0 for value in (bpol, minor_radius, ti_ev, ion.mass_kg)):
raise ValueError("Eich/Chang heat-load scaling inputs must be finite and positive.")
rho_i_pol = math.sqrt(2.0*ion.mass_kg*ti_ev*E_CHARGE_C)/(abs(ion.charge_C)*bpol)
eich_mm = 0.63*bpol**(-1.19)
chang_mm = eich_mm*(1.0 + 1.08e-10*(bpol*minor_radius/rho_i_pol)**4)
return DivertorWidthScalings(
eich_mm=eich_mm,
chang_mm=chang_mm,
bpol_midplane_t=bpol,
minor_radius_m=minor_radius,
ion_temperature_ev=ti_ev,
ion_name=ion.name,
)
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 _heatdiag_local_index_by_logical_step(heatdiag) -> dict[int, int]:
"""Map catalog logical steps to positions in the loaded HeatDiag time array."""
local_steps = heatdiag.list_step_indices("time")
mapping = {}
for position, local_step in enumerate(local_steps):
info = heatdiag.get_step_info(local_step)
mapping[int(info["logical_step"])] = position
return mapping
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:
"""Sum the component values selected by a target wall mask."""
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]:
"""Extend a target mask across contiguous valid wall points."""
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 _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.
Parameters
----------
x_m, q0, q_bg, lambda_q, s, s0
Upstream coordinate, width, spreading, and offset in meters; amplitude
and background use the same units as the fitted profile.
Returns
-------
ndarray
modeled heat-load profile at x_m.
"""
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.
Parameters
----------
x_m, y, fit_window_mm, smoothing_sigma_mm
Profile coordinates in meters, measured values, optional fit interval in millimeters, and Gaussian smoothing width in millimeters.
Returns
-------
tuple
fitted parameters and a sampled fit curve.
"""
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),
x_scale="jac",
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:
"""Interpolate interval rates onto the selected target coordinate."""
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,
)