"""Sheath diagnostics reader integrated with Simulation mesh data."""
from typing import Dict
import os
import numpy as np
from .accessor_mixin import ArrayAccessorMixin
from .bp_reader_mixin import BPReaderMixin
[docs]
class SheathData(BPReaderMixin, ArrayAccessorMixin):
"""
Class to handle / construct sheath data utilizing xgc.mesh.bp and xgc.sheathdiag.bp files.
Mesh data are accessed via the Simulation instance (mesh/plane).
xgc.sheathdiag.bp
- nwall : Integer representing the number of wall nodes.
- sheath_pot : Sheath Potential
- sheath_ilost : Sheath Ion Loss
- sheath_lost : Sheath Electron Loss
- sheath_nphi : Sheath Electron Density
Attributes
----------
simulation : Simulation
Simulation instance providing mesh and magnetic field data.
data : dict[str, dict[int, np.ndarray]]
Nested dictionary storing time-dependent sheath diagnostics in the
common XGC-Analysis layout: ``self.data[var_name][step_idx] = ndarray``.
wall_data : dict[str, np.ndarray]
Static wall-related geometry/derived arrays:
``wall_rz``, ``wall_phi``, ``wall_psi``.
Returns:
SheathData: An instance of the SheathData class.
"""
TIME_VARS = ["sheath_pot", "sheath_ilost", "sheath_lost", "sheath_nphi"]
def __init__(
self,
steps,
simulation,
sheath_file="xgc.sheathdiag.bp",
data_dir="./",
catalog=None,
missing="raise",
source_reader=None,
):
"""
Initializes the SheathData instance.
Args:
------
steps (int, tuple):
Single integer timestep or
tuple of integers (start_step_inclusive, end_step_exclusive) to read timestep data.
sheath_file (str):
Path to the sheath diagnostic file (default: "xgc.sheathdiag.bp").
data_dir (str):
Directory containing the mesh and sheath files (default: ".").
simulation (Simulation):
xgc_analysis Simulation instance used to source mesh data.
catalog (SimulationCatalog, optional):
Optional catalog used to resolve logical steps into BP sources.
If omitted, ``simulation.catalog`` is used when present. Direct
local filename reads are disabled when no catalog is available.
missing (str):
Missing-step policy for catalog read planning.
source_reader (callable, optional):
Optional read-plan backend hook.
"""
# Construct full paths for sheath files
self.requested_steps = steps
self.simulation = simulation
if data_dir == "./":
data_dir = self.simulation.data_directory
self.sheath_file_path = os.path.join(data_dir, sheath_file)
self.catalog = catalog if catalog is not None else getattr(simulation, "catalog", None)
self.missing = missing
self.source_reader = source_reader
self._init_bp_reader_state(variables=self.TIME_VARS, read_all_steps=False)
# Time-dependent sheath diagnostics in common nested-dict layout.
self.data: Dict[str, Dict[int, np.ndarray]] = {}
# Static wall geometry/derived arrays.
self.wall_data: Dict[str, np.ndarray] = {}
if self.catalog is None:
raise RuntimeError("SheathData requires a catalog; direct xgc.sheathdiag.bp reads are disabled.")
self._load_sheath_data_from_catalog()
self._build_wall_data()
def _requested_step_list(self):
"""
Return requested sheath logical steps as a list.
Returns
-------
list[int]
Step values requested by the constructor.
"""
if isinstance(self.requested_steps, int):
return [int(self.requested_steps)]
if isinstance(self.requested_steps, tuple) and len(self.requested_steps) == 2:
return list(range(int(self.requested_steps[0]), int(self.requested_steps[1])))
raise TypeError("Steps must be an integer or a tuple of two integers (start, end).")
def _requested_step_range_and_list(self):
"""
Return the legacy ADIOS step range and expected step list.
Returns
-------
tuple
``(step_range, expected_steps)`` retained only for compatibility
with older private helper logic. Current reads use catalog logical
steps from :meth:`_requested_step_list`.
"""
expected_steps = self._requested_step_list()
if isinstance(self.requested_steps, int):
step_range = (int(self.requested_steps), int(self.requested_steps) + 1)
else:
step_range = (int(self.requested_steps[0]), int(self.requested_steps[1]))
return step_range, expected_steps
def _load_sheath_data(self):
"""
Disabled legacy direct-file reader.
Sheath diagnostic data must be read through catalog read plans so the
same code path works for directory and campaign backends.
"""
raise RuntimeError("SheathData direct xgc.sheathdiag.bp reads are disabled; use a catalog.")
def _load_sheath_data_from_catalog(self):
"""
Load sheath diagnostic data through catalog read plans.
Returns
-------
None
Populates ``self.data`` and ``self.step_index_info`` in place.
"""
print("Loading sheath diagnostic data... ")
step_variables = self._read_catalog_product(
self.catalog,
"xgc.sheathdiag.bp",
self.TIME_VARS,
steps=self._requested_step_list(),
missing=self.missing,
require_all_variables=True,
source_reader=self.source_reader,
)
for step_index in sorted(step_variables):
for var_name, value in step_variables[step_index].items():
self.data.setdefault(var_name, {})[step_index] = np.asarray(value)
def _build_wall_data(self):
"""Build static wall-related arrays and store them in ``self.wall_data``."""
self.wall_data["wall_rz"] = self.simulation.mesh.wall_rz
self.wall_data["wall_phi"] = self.simulation.mesh.wall_phi
self.wall_data["wall_psi"] = self._compute_wall_psi()
def _compute_wall_psi(self) -> np.ndarray:
"""Compute poloidal flux (psi) at wall nodes on plane 0."""
plane = self.simulation.mesh.get_plane(0)
wall_nodes = getattr(plane, "wall_nodes", None)
if wall_nodes is None:
raise AttributeError("Simulation mesh is missing wall_nodes.")
psi = self.simulation.magnetic_field.psi_pd.data
return psi[wall_nodes]
# ====================================================================
# Assembling Spatial Data
# ====================================================================
# Assemble wall spatial data using mesh wall nodes
@property
def wall_rz(self):
"""Return wall node (R, Z) coordinates for each mesh plane."""
return self.wall_data["wall_rz"]
@property
def wall_phi(self) -> np.array:
"""Return toroidal angle(s) corresponding to wall planes."""
return self.wall_data["wall_phi"]
@property
def wall_psi(self) -> np.array:
"""
Returns poloidal magnetic flux (psi) at wall nodes for plane 0.
"""
return self.wall_data["wall_psi"]
[docs]
def get_wall_array(self, var_name: str) -> np.ndarray:
"""Return a static wall-related array from ``self.wall_data``."""
if var_name not in self.wall_data:
raise KeyError(f"Wall variable '{var_name}' not found in SheathData.")
return self.wall_data[var_name]
# --- Typed accessors ---
[docs]
def get_sheath_array(self, var_name: str, step_index: int):
"""Return time-dependent sheath diagnostic ``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))