"""Reader for XGC wall heat/particle impact diagnostics (xgc.heatdiag2.bp)."""
from __future__ import annotations
import os
from typing import Dict, Iterable, Optional, Tuple
import numpy as np
from .accessor_mixin import ArrayAccessorMixin
from .bp_reader_mixin import BPReaderMixin
from .time_step_utils import build_last_occurrence_step_mask
[docs]
class HeatDiag(BPReaderMixin, ArrayAccessorMixin):
"""
Read wall heat diagnostic data from ``xgc.heatdiag2.bp``.
Time-dependent data are stored in ``self.data`` with the structure::
self.data[variable_name][step_idx] = 2D numpy array
where ``step_idx`` is the ADIOS step index in the BP file. Variables with an
extra "garbage bin" (first segment entry) are split into:
- ``<var>``: wall-segment data excluding the garbage bin
- ``<var>_garbage``: garbage-bin values with shape ``(nphi, 1)``
Static wall geometry/metadata are stored in ``self.wall_data`` as::
self.wall_data[variable_name] = 2D numpy array
The class can be constructed stand-alone using ``data_dir`` or using a
``Simulation`` instance (it will use ``simulation.data_directory`` by
default). A catalog must be supplied directly or through ``simulation``;
direct local filename reads are disabled.
"""
STATIC_VARS = {
"ds",
"psi",
"r",
"z",
"strike_angle",
"nphi",
"nseg",
"nseg1",
}
TIME_VARS = {
"time",
"gstep",
"tindex",
"e_number",
"e_para_energy",
"e_perp_energy",
"e_potential",
"i_number",
"i_para_energy",
"i_perp_energy",
"i_potential",
}
GARBAGE_BIN_VARS = {
"e_number",
"e_para_energy",
"e_perp_energy",
"e_potential",
"i_number",
"i_para_energy",
"i_perp_energy",
"i_potential",
}
DEFAULT_VARS = tuple(sorted(STATIC_VARS | TIME_VARS))
def __init__(
self,
simulation=None,
data_dir: str = "./",
filename: str = "xgc.heatdiag2.bp",
step_range: Optional[Tuple[int, int]] = None,
variables: Optional[Iterable[str]] = None,
catalog=None,
steps: Optional[Iterable[int]] = None,
missing: str = "raise",
source_reader=None,
):
"""
Initialize and load the heat diagnostic product.
Parameters
----------
simulation : Simulation or None, optional
Simulation object used for the default data directory and catalog.
data_dir : str, optional
Dataset directory used when ``simulation`` is not provided.
filename : str, optional
Catalog product key for the heat diagnostic output.
step_range : tuple[int, int] or None, optional
Legacy half-open logical step range. Ignored when ``steps`` is
provided.
variables : iterable[str] or None, optional
Variables to read. If omitted, the reader uses its standard static
and time-dependent heat diagnostic variables.
catalog : SimulationCatalog or None, optional
Catalog used to resolve logical steps into ADIOS source fragments.
If omitted, ``simulation.catalog`` is used when available.
steps : iterable[int] or None, optional
Explicit logical XGC steps to read.
missing : {"raise", "skip", "zero"}, optional
Missing-step policy for catalog read planning.
source_reader : callable or None, optional
Optional read-plan backend hook.
"""
self.simulation = simulation
if self.simulation is not None and data_dir == "./":
data_dir = self.simulation.data_directory
self.data_dir = data_dir
self.filename = filename
self.file_path = os.path.join(self.data_dir, self.filename)
self.step_range = step_range
self._variables_provided = variables is not None
self.variables = list(variables) if variables is not None else list(self.DEFAULT_VARS)
self.catalog = catalog if catalog is not None else getattr(simulation, "catalog", None)
self.catalog_steps = None if steps is None else [int(step) for step in steps]
self.missing = missing
self.source_reader = source_reader
self._init_bp_reader_state(variables=self.variables, read_all_steps=True)
self.data: Dict[str, Dict[int, np.ndarray]] = {}
self.wall_data: Dict[str, np.ndarray] = {}
self.available_steps = []
if self.catalog is None:
raise RuntimeError("HeatDiag requires a catalog; direct xgc.heatdiag2.bp reads are disabled.")
self._read_heatdiag_from_catalog()
def _read_heatdiag(self):
"""
Disabled legacy direct-file reader.
Heat diagnostic data must be read through catalog read plans so the same
code path works for directory and campaign backends.
"""
raise RuntimeError("HeatDiag direct xgc.heatdiag2.bp reads are disabled; use a catalog.")
def _read_heatdiag_from_catalog(self):
"""
Read heat diagnostic data through catalog read plans.
Returns
-------
None
Populates ``self.data``, ``self.wall_data``, and
``self.step_index_info`` in place.
"""
if self.catalog_steps is not None:
steps = self.catalog_steps
read_all_steps = False
elif self.step_range is not None:
steps = range(int(self.step_range[0]), int(self.step_range[1]))
read_all_steps = False
else:
steps = None
read_all_steps = True
step_variables = self._read_catalog_product(
self.catalog,
self.filename,
self.variables,
steps=steps,
read_all_steps=read_all_steps,
missing=self.missing,
require_all_variables=self._variables_provided,
source_reader=self.source_reader,
)
self.available_steps = sorted(step_variables)
for step_idx in self.available_steps:
self._process_step_variables(step_idx, step_variables[step_idx])
def _process_step_variables(self, step_idx: int, variables: Dict[str, object]):
"""
Normalize and store variables read for one heatdiag step.
Parameters
----------
step_idx : int
Reader-local step index.
variables : dict[str, object]
Raw variable values for one source step.
"""
for var_name, raw_value in variables.items():
arr2d = self._as_2d_array(raw_value)
if var_name in self.GARBAGE_BIN_VARS:
if arr2d.shape[1] < 1:
raise ValueError(f"Variable '{var_name}' has invalid shape {arr2d.shape}")
self._store(var_name, step_idx, arr2d[:, 1:])
self._store(f"{var_name}_garbage", step_idx, arr2d[:, :1])
continue
if var_name in self.STATIC_VARS:
if var_name not in self.wall_data:
self.wall_data[var_name] = arr2d
continue
self._store(var_name, step_idx, arr2d)
def _store(self, var_name: str, step_idx: int, value: np.ndarray):
"""
Store one time-dependent heat diagnostic array.
Parameters
----------
var_name : str
Output variable name.
step_idx : int
Reader-local step index.
value : np.ndarray
Normalized two-dimensional array for this variable and step.
"""
self.data.setdefault(var_name, {})[step_idx] = value
@staticmethod
def _as_2d_array(value) -> np.ndarray:
"""
Normalize a heat diagnostic payload to a two-dimensional array.
Scalars become ``(1, 1)`` arrays, one-dimensional arrays become
``(1, n)`` arrays, and already two-dimensional arrays are returned
unchanged. Higher-dimensional arrays are squeezed when that produces a
scalar, vector, or matrix.
Parameters
----------
value : object
Raw value returned by catalog read-plan execution.
"""
arr = np.asarray(value)
if arr.ndim == 0:
return arr.reshape(1, 1)
if arr.ndim == 1:
return arr[np.newaxis, :]
if arr.ndim == 2:
return arr
squeezed = np.squeeze(arr)
if squeezed.ndim == 0:
return squeezed.reshape(1, 1)
if squeezed.ndim == 1:
return squeezed[np.newaxis, :]
if squeezed.ndim == 2:
return squeezed
raise ValueError(f"Unsupported variable shape {arr.shape}")
[docs]
def get_wall_array(self, var_name: str) -> np.ndarray:
"""Return a static wall variable from ``self.wall_data``."""
if var_name not in self.wall_data:
raise KeyError(f"Wall variable '{var_name}' not found in HeatDiag.")
return self.wall_data[var_name]
# --- Typed accessors ---
[docs]
def get_heat_array(self, var_name: str, step_index: int):
"""Return time-dependent heatdiag variable ``var_name`` as ``np.ndarray``."""
return self.get_as(var_name, step_index, np.ndarray)
[docs]
def get_scalar(self, var_name: str, step_index: int):
"""Return ``var_name`` at ``step_index`` as a scalar numeric value."""
return self.get_as(var_name, step_index, (int, float, np.integer, np.floating))
[docs]
def get_time_mask(self) -> np.ndarray:
"""Return indices that sort and de-duplicate the diagnostic history.
This mirrors ``OneDDiag.get_time_mask()`` and is useful when heatdiag
output contains overlapping time ranges. The returned indices apply to
arrays returned by ``get_array(...)`` (stacked over ``available_steps``).
"""
if self.has_var("gstep"):
step_var = "gstep"
elif self.has_var("step"):
step_var = "step"
else:
raise KeyError("HeatDiag does not contain a 'gstep' or legacy 'step' time series.")
step_values = np.asarray(self.get_array(step_var)).reshape(-1)
self.tmask = build_last_occurrence_step_mask(step_values)
return self.tmask
[docs]
def get_wall_curve(
self,
*,
phi_index: int = 0,
verify_clockwise: bool = True,
auto_reverse: bool = False,
set_inboard_origin: bool = False,
r_axis: float | None = None,
):
"""
Build a :class:`WallCurve` from the HeatDiag wall polygon.
Parameters
----------
phi_index : int
Toroidal index for selecting one wall polygon from heatdiag arrays.
verify_clockwise : bool
Validate clockwise ordering of wall points.
auto_reverse : bool
Reverse point order automatically if not clockwise.
set_inboard_origin : bool
If True, set arclength origin at Z=0 with R<R_axis.
r_axis : float | None
Magnetic-axis R used when ``set_inboard_origin=True``. If omitted,
use ``simulation.mesh.get_plane(0).axis_r`` when available.
"""
from .wall_curve import WallCurve
r = np.asarray(self.get_wall_array("r"))
z = np.asarray(self.get_wall_array("z"))
if r.ndim != 2 or z.ndim != 2:
raise ValueError("HeatDiag wall arrays 'r' and 'z' must be 2D (nphi, nseg).")
if phi_index < 0 or phi_index >= r.shape[0]:
raise IndexError(f"phi_index={phi_index} out of range [0,{r.shape[0]-1}].")
points = np.column_stack([r[phi_index], z[phi_index]])
wc = WallCurve(points, verify_clockwise=verify_clockwise, auto_reverse=auto_reverse)
if set_inboard_origin:
if r_axis is None:
if self.simulation is None:
raise ValueError("r_axis must be provided when simulation is unavailable.")
r_axis = float(self.simulation.mesh.get_plane(0).axis_r)
wc.set_origin_at_inboard_midplane(float(r_axis))
return wc