Source code for xgc_analysis.oneddiag

"""Reader and utilities for XGC 1D diagnostics (``xgc.oneddiag.bp``).

This module refactors the 1D diagnostic reader to follow the common
``XGC-Analysis`` data layout used by other readers:

    ``self.data[var_name][step_index] = scalar | np.ndarray``

Key conventions
---------------
- Species-resolved variables are normalized to standardized dotted keys such as
  ``"e.gc_density_df_1d"`` and ``"i2.parallel_flow_df_1d"``.
- Static arrays that are effectively constant in time (for example ``psi``) are
  stored in ``self.static_data`` and exposed via compatibility aliases
  ``self.psi``, ``self.psi00``, and ``self.psi_mks``.
- Derived quantities are stored separately in ``self.derived_data`` /
  ``self.derived_static_data`` and accessed through ``od.derived``.

The module also provides lightweight species views (``od.e``, ``od.i``, ...)
for backward-compatible attribute access in notebooks.
"""

from __future__ import annotations

import os
from typing import Dict, Optional

import matplotlib.pyplot as plt
import numpy as np

from .accessor_mixin import ArrayAccessorMixin
from .bp_reader_mixin import BPReaderMixin
from .constants import echarge, m_p, mu_0
from .time_step_utils import build_last_occurrence_step_mask


class _OneDSpeciesView:
    """Lightweight attribute view onto one species namespace in ``OneDDiag``.

    Examples
    --------
    ``od.e.gc_density_df_1d`` returns the stacked time-series array for the
    standardized variable ``"e.gc_density_df_1d"``. The special attributes
    ``mass`` and ``species`` expose metadata from ``OneDDiag.mass_by_prefix`` and
    ``OneDDiag.species_by_prefix``.
    """

    def __init__(self, owner: "OneDDiag", prefix: str, *, prefer_derived: bool = False):
        """
        Store the owner and species namespace represented by this view.

        Parameters
        ----------
        owner : OneDDiag
            Reader instance that owns the standardized and derived variables.
        prefix : str
            Species prefix such as ``"e"`` or ``"i"``.
        prefer_derived : bool, optional
            If True, resolve attributes only against derived species variables.
        """
        self._owner = owner
        self.prefix = prefix
        self.spname = prefix
        self._prefer_derived = prefer_derived

    def __getattr__(self, name):
        """
        Resolve one species attribute to a stored diagnostic array or metadata.

        Parameters
        ----------
        name : str
            Attribute requested from the species view.
        """
        if name in ("mass", "species"):
            if name == "mass":
                return self._owner.mass_by_prefix.get(self.prefix)
            return self._owner.species_by_prefix.get(self.prefix)

        key = f"{self.prefix}.{name}"

        if self._prefer_derived:
            if self._owner.has_derived_var(key):
                return self._owner.get_derived_array(key)
            raise AttributeError(f"Derived species variable '{key}' not found.")

        if self._owner.has_var(key):
            return self._owner.get_array(key)
        if self._owner.has_derived_var(key):
            return self._owner.get_derived_array(key)

        raise AttributeError(f"Species variable '{key}' not found.")

    def __repr__(self):
        """Return a compact debugging representation of this species view."""
        return f"_OneDSpeciesView(prefix={self.prefix!r})"


class _OneDDerivedView:
    """Attribute view onto derived variables and per-species derived subviews.

    This enables access patterns such as ``od.derived.shear_r`` and
    ``od.derived.e.T``.
    """

    def __init__(self, owner: "OneDDiag"):
        """
        Build derived-variable subviews for all known species prefixes.

        Parameters
        ----------
        owner : OneDDiag
            Reader instance that owns the derived variables.
        """
        self._owner = owner
        for prefix in owner.SPECIES_PREFIXES:
            setattr(self, prefix, _OneDSpeciesView(owner, prefix, prefer_derived=True))

    def __getattr__(self, name):
        """
        Resolve a top-level derived diagnostic variable by attribute name.

        Parameters
        ----------
        name : str
            Derived variable name requested from ``oneddiag.derived``.
        """
        if self._owner.has_derived_var(name):
            return self._owner.get_derived_array(name)
        raise AttributeError(f"Derived variable '{name}' not found.")


[docs] class OneDDiag(BPReaderMixin, ArrayAccessorMixin): """Reader for XGC 1D diagnostics with standard nested-dict storage layout. Parameters ---------- path : str, default "./" Directory containing the oneddiag BP file. filename : str, default "xgc.oneddiag.bp" Diagnostic filename. simulation : Simulation or None, optional If provided, ``simulation.species`` is used to attach ``Species`` objects and species masses to standardized prefixes detected from the file. catalog : SimulationCatalog or None, 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. steps : iterable[int] or None, optional Logical XGC steps to read from ``catalog``. If omitted, all available oneddiag steps are read. missing : {"raise", "skip", "zero"}, optional Missing-step policy for catalog read planning. source_reader : callable or None, optional Optional read-plan backend hook. Attributes ---------- data : dict[str, dict[int, object]] Time-dependent raw variables in the common reader layout. static_data : dict[str, np.ndarray] Time-independent arrays promoted out of ``data`` (for example ``psi``). derived_data : dict[str, dict[int, object]] Time-dependent derived quantities (e.g. ``"e.T"``, ``"shear_r"``). species_by_prefix : dict[str, Species | None] Mapping from standardized prefixes (``e``, ``i``, ``i2``, ...) to ``Species`` objects when available. mass_by_prefix : dict[str, float | None] Species masses in atomic mass units, from ``Species`` metadata or internal fallback defaults. """ SPECIES_PREFIXES = ["e", "i", "i2", "i3", "i4", "i5", "i6", "i7", "i8", "i9"] _SPECIES_PARSE_ORDER = sorted(SPECIES_PREFIXES, key=len, reverse=True) _IGNORED_VARS = {"samples", "gsamples"} _STATIC_CANDIDATES = {"psi", "psi00", "psi_mks"} echarge = echarge protmass = m_p mu0 = mu_0 def __init__( self, path: str = "./", filename: str = "xgc.oneddiag.bp", simulation=None, catalog=None, steps=None, missing: str = "raise", source_reader=None, ): """ Initialize and load one-dimensional diagnostic data. Parameters ---------- path : str, optional Dataset directory containing the oneddiag product. filename : str, optional Catalog product key for the diagnostic output. simulation : Simulation or None, optional Simulation object used to inherit species metadata and, when available, a catalog. catalog : SimulationCatalog or None, optional Catalog used to resolve logical steps into ADIOS source fragments. steps : iterable[int] or None, optional Explicit logical XGC steps to read. If omitted, all catalog steps are 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.path = path self.filename = filename self.simulation = simulation 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=None, read_all_steps=True) self.data: Dict[str, Dict[int, object]] = {} self.static_data: Dict[str, np.ndarray] = {} self.derived_data: Dict[str, Dict[int, object]] = {} self.derived_static_data: Dict[str, np.ndarray] = {} self.available_steps = [] self.active_species_prefixes = [] self.electron_on = False self.species_by_prefix: Dict[str, object] = {p: None for p in self.SPECIES_PREFIXES} self.mass_by_prefix: Dict[str, Optional[float]] = {p: None for p in self.SPECIES_PREFIXES} for prefix in self.SPECIES_PREFIXES: setattr(self, prefix, _OneDSpeciesView(self, prefix)) self.derived = _OneDDerivedView(self) if self.catalog is None: raise RuntimeError("OneDDiag requires a catalog; direct xgc.oneddiag.bp reads are disabled.") self.load_data_from_catalog() self._initialize_species_metadata(simulation=simulation) self.post_process() @classmethod def _split_species_var(cls, var_name: str): """Split a raw BP variable name into ``(species_prefix, remainder)``. Matching is performed longest-prefix-first so names like ``i2_*`` are not incorrectly parsed as ``i_*``. """ for prefix in cls._SPECIES_PARSE_ORDER: head = prefix + "_" if var_name.startswith(head): return prefix, var_name[len(head):] return None, var_name @classmethod def _standardize_var_name(cls, var_name: str) -> str: """Convert a raw BP variable name to the standardized reader key format.""" prefix, rest = cls._split_species_var(var_name) return f"{prefix}.{rest}" if prefix is not None else rest @staticmethod def _normalize_step_value(value): """Normalize one BP-read value to either a scalar or a squeezed ndarray.""" arr = np.asarray(value) if arr.ndim == 0: return arr.item() return np.squeeze(arr) def _data_file_path(self) -> str: """Return the full path to the oneddiag BP file.""" return os.path.join(self.path, self.filename) def _promote_static_arrays(self): """Move known constant-in-time arrays from ``self.data`` to ``self.static_data``. The promoted arrays are also exposed as compatibility attributes ``self.psi``, ``self.psi00``, and ``self.psi_mks``. """ for key in list(self._STATIC_CANDIDATES): if key not in self.data or not self.data[key]: continue first_step = sorted(self.data[key].keys())[0] self.static_data[key] = np.asarray(self.data[key][first_step]) del self.data[key] self.psi = self.static_data.get("psi") self.psi00 = self.static_data.get("psi00") self.psi_mks = self.static_data.get("psi_mks") def _refresh_cached_time_series(self): """Cache common time-series arrays and the list of available step indices.""" self.available_steps = self.list_step_indices() if self.has_var("time"): self.time = self.get_array("time") if self.has_var("gstep"): self.gstep = self.get_array("gstep") self.step = self.gstep elif self.has_var("step"): self.step = self.get_array("step") self.gstep = self.step def _detect_active_species(self): """Detect which standardized species prefixes are present in loaded data. This method only inspects variable names and does not require ``Species`` metadata. It also populates legacy convenience attributes: ``electron_on`` and ``sps``. """ self.active_species_prefixes = [ p for p in self.SPECIES_PREFIXES if any(k.startswith(p + ".") for k in self.data.keys()) ] self.electron_on = "e" in self.active_species_prefixes self.sps = [getattr(self, p) for p in self.active_species_prefixes] def _initialize_species_metadata(self, simulation=None): """ Attach Species objects and masses to standardized species prefixes. Species-prefix presence is detected from the oneddiag file. If a ``Simulation`` object is provided, its ``simulation.species`` sequence is mapped onto the active prefixes in file order. """ if simulation is not None: sim_species = getattr(simulation, "species", None) if sim_species is not None: for prefix, sp in zip(self.active_species_prefixes, sim_species): self.species_by_prefix[prefix] = sp # Fallback masses used when no Species metadata is available. default_mass_list = [5.45e-4, 2.0, 12.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] for i, prefix in enumerate(self.SPECIES_PREFIXES): sp_obj = self.species_by_prefix.get(prefix) if sp_obj is not None and hasattr(sp_obj, "mass_au"): self.mass_by_prefix[prefix] = float(sp_obj.mass_au) elif i < len(default_mass_list): self.mass_by_prefix[prefix] = float(default_mass_list[i])
[docs] def load_data(self): """Disabled legacy direct-file reader. One-dimensional diagnostic data must be read through catalog read plans so the same code path works for directory and campaign backends. """ raise RuntimeError("OneDDiag direct xgc.oneddiag.bp reads are disabled; use a catalog.")
[docs] def load_data_from_catalog(self): """ Read the oneddiag product through catalog read plans. All variables advertised by the catalog product are read unless they are ignored by this reader. The resulting raw step-major data are then normalized by the same path used for direct BP reads. """ product = self.catalog.get_product(self.filename) variables = sorted(name for name in product.variables if name not in self._IGNORED_VARS) step_variables = self._read_catalog_product( self.catalog, self.filename, variables, steps=self.catalog_steps, read_all_steps=self.catalog_steps is None, missing=self.missing, require_all_variables=True, source_reader=self.source_reader, ) self._load_raw_step_data(step_variables)
def _load_raw_step_data(self, raw): """ Normalize step-major raw oneddiag data into ``self.data``. Parameters ---------- raw : dict[int, dict[str, object]] Step-major variable dictionary returned by the catalog read-plan path. """ for step_idx, variables in raw.items(): for raw_name, raw_value in variables.items(): if raw_name in self._IGNORED_VARS: continue std_name = self._standardize_var_name(raw_name) self.data.setdefault(std_name, {})[step_idx] = self._normalize_step_value(raw_value) self._promote_static_arrays() self._refresh_cached_time_series() self._detect_active_species()
[docs] def has_derived_var(self, var_name): """Return ``True`` if a derived variable exists in dynamic or static storage.""" return var_name in self.derived_data or var_name in self.derived_static_data
[docs] def get_derived_item(self, var_name, step_index): """Return one derived item for a specific ``step_index``.""" if var_name not in self.derived_data: raise KeyError(f"Derived variable '{var_name}' not found in derived_data.") if step_index not in self.derived_data[var_name]: raise KeyError(f"Step {step_index} not found for derived variable '{var_name}'.") return self.derived_data[var_name][step_index]
[docs] def get_derived_array(self, var_name): """Return a derived variable stacked over steps, or a derived static array.""" if var_name in self.derived_static_data: return self.derived_static_data[var_name] if var_name not in self.derived_data: raise KeyError(f"Derived variable '{var_name}' not found.") step_dict = self.derived_data[var_name] return np.stack([step_dict[k] for k in sorted(step_dict.keys())])
def _store_derived_series(self, var_name, values): """Store a derived time-series array into ``self.derived_data``. The first dimension of ``values`` must match the number of loaded steps. """ arr = np.asarray(values) if arr.shape[0] != len(self.available_steps): raise ValueError( f"Derived series '{var_name}' first dimension {arr.shape[0]} does not match number of available steps {len(self.available_steps)}." ) for i, step_idx in enumerate(self.available_steps): self.derived_data.setdefault(var_name, {})[step_idx] = arr[i]
[docs] def get_static_array(self, var_name): """Return a static array from ``self.static_data`` (for example ``psi``).""" if var_name not in self.static_data: raise KeyError(f"Static variable '{var_name}' not found.") return self.static_data[var_name]
[docs] def get_profile(self, var_name, step_index=0): """Return a time-dependent oneddiag variable as an ndarray for one step.""" return self.get_as(var_name, step_index, np.ndarray)
[docs] def get_scalar(self, var_name, step_index=0): """Return a scalar time-dependent oneddiag variable for one step.""" return self.get_as(var_name, step_index, (int, float, np.integer, np.floating))
[docs] def list_species(self): """Return active species prefixes detected in the file.""" return list(self.active_species_prefixes)
[docs] def get_species_view(self, prefix): """Return the lightweight species view object for a standardized prefix.""" if prefix not in self.SPECIES_PREFIXES: raise KeyError(f"Unknown species prefix '{prefix}'.") return getattr(self, prefix)
[docs] def get_species_by_prefix(self, prefix): """Return the associated ``Species`` object (if available) for ``prefix``.""" if prefix not in self.species_by_prefix: raise KeyError(f"Unknown species prefix '{prefix}'.") return self.species_by_prefix[prefix]
[docs] def d_dpsi(self, var, psi): """Compute ``d(var)/d(psi)`` on a non-uniform 1D grid. Parameters ---------- var : np.ndarray Array of shape ``(n_step, n_psi)``. psi : np.ndarray 1D non-uniform coordinate array of length ``n_psi``. Returns ------- np.ndarray Numerical derivative with the same shape as ``var``. """ n = len(psi) dvar_dpsi = np.zeros_like(var) h0 = psi[1] - psi[0] h1 = psi[2] - psi[1] dvar_dpsi[:, 0] = (-var[:, 2] * h0 + var[:, 1] * (h0 + h1) - var[:, 0] * h1) / (h0 * h1 * (h0 + h1)) for i in range(1, n - 1): h0 = psi[i] - psi[i - 1] h1 = psi[i + 1] - psi[i] term1 = -h1**2 * var[:, i - 1] term2 = (h1**2 - h0**2) * var[:, i] term3 = h0**2 * var[:, i + 1] dvar_dpsi[:, i] = (term1 + term2 + term3) / (h0 * h1 * (h0 + h1)) h0 = psi[-2] - psi[-3] h1 = psi[-1] - psi[-2] dvar_dpsi[:, -1] = (var[:, -3] * h1 - var[:, -2] * (h0 + h1) + var[:, -1] * h0) / (h0 * h1 * (h0 + h1)) return dvar_dpsi
[docs] def post_process(self): """Compute commonly used derived quantities and store them in ``od.derived``. Derived quantities currently include: - species temperatures ``<prefix>.T`` - species gradient scale lengths ``<prefix>.Ln`` and ``<prefix>.Lt`` - reference ``density`` and ``Ln`` - ``grad_psi_sqr`` and ``shear_r`` Notes ----- The reference species for ``density`` / ``shear_r`` is electrons if present, otherwise the main ion prefix ``i`` when available. """ if not self.active_species_prefixes or self.psi_mks is None: return for prefix in self.active_species_prefixes: sp = getattr(self, prefix) try: Teperp = sp.perp_temperature_df_1d Tepara = sp.parallel_mean_en_df_1d - 0.5 * sp.mass * self.protmass * sp.parallel_flow_df_1d**2 / self.echarge self._store_derived_series(f"{prefix}.T", (Teperp + Tepara) / 3.0 * 2.0) except AttributeError: continue shear_src = "e" if self.electron_on else ("i" if "i" in self.active_species_prefixes else None) if shear_src is not None: try: src = getattr(self, shear_src) shear = self.d_dpsi(src.poloidal_ExB_flow_1d, self.psi_mks) grad_psi_sqr = src.grad_psi_sqr_1d self._store_derived_series("grad_psi_sqr", grad_psi_sqr) self._store_derived_series("shear_r", shear * np.sqrt(grad_psi_sqr)) self.grad_psi_sqr = self.get_derived_array("grad_psi_sqr") self.shear_r = self.get_derived_array("shear_r") except AttributeError: pass dens_prefix = "e" if self.electron_on else ("i" if "i" in self.active_species_prefixes else None) if dens_prefix is not None and self.has_derived_var("grad_psi_sqr"): try: density = getattr(self, dens_prefix).gc_density_df_1d self._store_derived_series("density", density) Ln = density / self.d_dpsi(density, self.psi_mks) / np.sqrt(self.get_derived_array("grad_psi_sqr")) self._store_derived_series("Ln", Ln) self.density = self.get_derived_array("density") self.Ln = self.get_derived_array("Ln") except AttributeError: pass if self.has_derived_var("grad_psi_sqr"): grad_psi_sqr = self.get_derived_array("grad_psi_sqr") for prefix in self.active_species_prefixes: sp = getattr(self, prefix) try: sp_Ln = sp.gc_density_df_1d / self.d_dpsi(sp.gc_density_df_1d, self.psi_mks) / np.sqrt(grad_psi_sqr) self._store_derived_series(f"{prefix}.Ln", sp_Ln) except AttributeError: pass try: sp_T = self.get_derived_array(f"{prefix}.T") sp_Lt = sp_T / self.d_dpsi(sp_T, self.psi_mks) / np.sqrt(grad_psi_sqr) self._store_derived_series(f"{prefix}.Lt", sp_Lt) except (AttributeError, KeyError): pass
[docs] def get_time_mask(self): """Build a mask selecting the last occurrence of each diagnostic step. This is a legacy helper kept for compatibility with older notebook workflows that expect ``self.tmask`` to index a monotonic subset of the stored time history. Overlapping diagnostic segments are resolved by keeping the last occurrence for each repeated step value. """ if hasattr(self, "gstep"): step_values = self.gstep elif hasattr(self, "step"): step_values = self.step else: raise KeyError("OneDDiag does not contain a 'gstep' or legacy 'step' time series.") self.tmask = build_last_occurrence_step_mask(step_values) return self.tmask
# Plotting helpers kept for now (candidate for future move to plotting.py).
[docs] def plot1d_if(self, var, time=None, varstr=None, psi=None, xlim=None, initial=True): """Plot first/last profiles from a time-series array (legacy helper). Parameters ---------- var : np.ndarray Array of shape ``(n_step, n_psi)``. time : np.ndarray or None, optional Time array aligned with ``var`` for labels. If omitted, labels are ``Initial`` / ``Final``. varstr : str or None, optional Label/title string. psi : np.ndarray or None, optional X-axis coordinate. Defaults to ``self.psi``. xlim : tuple[float, float] or None, optional Restrict plotting to a psi interval. initial : bool, default True Whether to also plot the first time slice. """ if psi is None: psi = self.psi if varstr is None: varstr = '' tunit = 1E3 if time is None: tstr0 = 'Initial' tstr1 = 'Final' else: tstr0 = 't=%3.3f ms' % (time[0] * tunit) tstr1 = 't=%3.3f ms' % (time[-1] * tunit) lbl = [varstr + ' ' + tstr0, varstr + ' ' + tstr1] fig, ax = plt.subplots() if xlim is None: if initial: ax.plot(psi, var[0, :], label=lbl[0]) ax.plot(psi, var[-1, :], label=lbl[1]) title_string = varstr else: msk = (psi >= xlim[0]) & (psi <= xlim[1]) if initial: ax.plot(psi[msk], var[0, msk], label=lbl[0]) ax.plot(psi[msk], var[-1, msk], label=lbl[1]) title_string = varstr + ' near edge' ax.legend() ax.set_xlabel('Normalized Pol. Flux') ax.set_ylabel(varstr) ax.set_title(title_string)
[docs] def report_profiles(self, sp_names=None, init_idx=0, end_idx=-1, edge_lim=[0.85, 1.05], show_edge=True): """Generate a small set of legacy profile plots for quick inspection. This method is retained for backward compatibility and is a candidate for migration into ``plotting.py`` in a future cleanup. """ if sp_names is None: sp_names = [sp.spname for sp in self.sps] linestyles = ['-', '-', '--', '--', '--', '--', '--', '--', '--', '--'] tunit = 1E3 fig, ax = plt.subplots() for i, sp1d in enumerate(self.sps): plt.plot(self.psi, sp1d.T[0, :] / tunit, label=sp_names[i], linestyle=linestyles[i]) plt.legend() plt.xlabel('Normalized Pol. Flux') plt.ylabel('Temperature (keV)') plt.title('Initial Temperature') dunit = 1E19 fig, ax = plt.subplots() for i, sp1d in enumerate(self.sps): plt.plot(self.psi, sp1d.gc_density_df_1d[0, :] / dunit, label=sp_names[i], linestyle=linestyles[i]) plt.legend() plt.xlabel('Normalized Pol. Flux') plt.ylabel('Density ($10^{19} m^{-3}$)') plt.title('Initial Density') i0 = init_idx i1 = end_idx for i, sp1d in enumerate(self.sps): density_str = 'density (m^-3)' if sp_names[i] == 'e' else 'g.c. density (m^-3)' self.plot1d_if(sp1d.gc_density_df_1d[i0:i1, :], time=self.time[i0:i1], varstr=sp_names[i] + ' ' + density_str) if show_edge: self.plot1d_if(sp1d.gc_density_df_1d[i0:i1, :], time=self.time[i0:i1], varstr=sp_names[i] + ' ' + density_str, xlim=edge_lim) for i, sp1d in enumerate(self.sps): self.plot1d_if(sp1d.T[i0:i1, :], time=self.time[i0:i1], varstr=sp_names[i] + ' Temperature (keV)') if show_edge: self.plot1d_if(sp1d.T[i0:i1, :], time=self.time[i0:i1], varstr=sp_names[i] + ' Temperature (keV)', xlim=edge_lim) for i, sp1d in enumerate(self.sps): self.plot1d_if(sp1d.parallel_flow_df_1d[i0:i1, :], time=self.time[i0:i1], varstr=sp_names[i] + ' parallel flow FSA (m/s)') if show_edge: self.plot1d_if(sp1d.parallel_flow_df_1d[i0:i1, :], time=self.time[i0:i1], varstr=sp_names[i] + ' parallel flow FSA (m/s)', xlim=edge_lim)