"""DiffusionProfiles – Stream reader for live XGC 1‑D profile output
====================================================================
This helper wraps the **ADIOS‑2** Python bindings (⩾ v2.10) in *Stream* mode so
that client code can monitor a growing ``*.bp`` directory and always access the
**latest** diffusion‑profile snapshot without rereading the entire file.
The class blocks during construction until the requested BP directory becomes
available (configurable *timeout*). It then opens an :class:`adios2.Stream`
reader that remains *open* until :py:meth:`DiffusionProfiles.close` is called or
the object is deleted, allowing incremental calls to
:py:meth:`DiffusionProfiles.read_next_step` as new simulation output appears.
Only **one** time step is kept resident in memory—the most recent one—so the
object’s footprint stays small even for long runs.
Time‑related parameters (``timeout``, ``wait_interval`` and the ``poll_delay``
argument of :py:meth:`read_next_step`) are **always expressed in *seconds***.
Pass ``value_in_minutes * 60`` if you prefer to think in minutes.
"""
from __future__ import annotations
import os
import time
from pathlib import Path
from typing import Any, Dict, Optional
import numpy as np
import sys
from adios2 import Stream, bindings, Adios
__all__ = [
"DiffusionProfiles",
]
[docs]
class DiffusionProfiles:
"""Live reader for ``xgc.diffusion_profiles.bp``.
Reads diffusion profile data from an Adios BP file (directory) ``xgc.diffusion_profiles.bp``.
The constructor opens the BP directory named by ``bp_dir``.
Once it is found, it opens it as a persistent :class:`adios2.Stream`, drains all existing
steps while retaining only the final one. The Stream is left open until it is closed with the
member method close().
A blocking request for a new step can be posted with the member method read_next_step.
The file ``xgc.diffusion_profiles.bp`` has the following structure:
double density n_steps*{n_species, n_samples, n_surf}
double flow n_steps*{n_species, n_samples, n_surf}
double temp n_steps*{n_species, n_samples, n_surf}
int32_t n_samples n_steps*scalar
int32_t n_species n_steps*scalar
int32_t n_surf n_steps*scalar
double psi n_steps*{n_surf}
int32_t steps n_steps*{n_samples}
double time n_steps*{n_samples}
That is, there are (for example) n_steps steps recorded in the BP file. The constructor of this class
reads only the last step (i.e. the most recent one).
The property latest_step_data returns a dictionary with the following keys:
density : A 3D array (shape (n_species, n_samples, n_surf)) of the diffusion density.
flow : A 3D array (shape (n_species, n_samples, n_surf) of the parallel mean flow.
temp : A 3D array (shape (n_species, n_samples, n_surf)) of temperature (in eV).
n_samples : An integer (n_samples) representing the number of time samples.
n_species : An integer (default 2) representing the number of particle species.
n_surf : An integer (n_surf) representing the number of ψ-surfaces.
psi : A 1D array (shape (n_surf,)) containing the poloidal magnetic flux values.
steps : A 1D array (shape (n_samples,)) of time step indices.
time : A 1D array (shape (n_samples,)) of simulation times (in seconds).
Parameters
----------
bp_dir : str | os.PathLike
Path to the BP directory written by XGC. The constructor **blocks**
until the directory exists (subject to *timeout*).
wait_interval : float, optional
Polling interval *in seconds* while waiting for the directory.
Defaults to **1.0 s**.
timeout : float | None, optional
Abort after this many *seconds* if the directory still does not exist.
``None`` (default) means *wait indefinitely*.
Notes
-----
*The object keeps exactly one snapshot in memory.* Callers needing to
store a history must copy out the data returned by
:attr:`latest_step_data`.
"""
def __init__(self,
adios: adios2.Adios,
*,
bp_dir: os.PathLike | str="./xgc.diffusion_profiles.bp",
wait_interval: float = 1.0,
timeout: Optional[float] = None,
) -> None:
"""Open a live BP profile stream and retain its latest available step.
``adios`` owns the ADIOS I/O object; ``bp_dir`` names the stream.
``timeout`` limits the initial open in seconds. ``wait_interval`` is
retained for API compatibility. The caller must close this object.
"""
self.path = Path(bp_dir)
self._stream: Optional["Stream"] = None
self._step_index: int | None = None
self._data: Dict[str, np.ndarray] | None = None
#self._wait_for_file(wait_interval, timeout)
#print(f"Opening {self.path}", flush=True)
self.adios = adios
self._open_stream(timeout)
print(f"Skipping to last time step", flush=True)
self._drain_initial_steps(timeout)
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
@property
def latest_step_data(self) -> Dict[str, np.ndarray]:
"""Return a **shallow copy** of the dictionary holding the current step.
Raises
------
RuntimeError
If no step has been read yet (e.g., the writer hasn’t flushed a
step to disk).
"""
if self._data is None:
raise RuntimeError("No step has been read yet.")
# Return a shallow copy so callers cannot accidentally mutate internals
return dict(self._data)
[docs]
@classmethod
def load_last_sample_history(
cls,
adios: Adios,
*,
bp_dir: os.PathLike | str = "./xgc.diffusion_profiles.bp",
max_steps: int | None = None,
stride: int = 1,
open_timeout: float | None = 0.0,
) -> Dict[str, list]:
"""
Reconstruct per-step history using the last time sample of each profile.
Parameters
----------
adios : Adios
ADIOS handle used to create a temporary read stream.
bp_dir : str | os.PathLike
Path to ``xgc.diffusion_profiles.bp``.
max_steps : int | None
If provided and > 0, keep only the last ``max_steps`` entries.
stride : int
Keep every ``stride``-th step while scanning BP history.
open_timeout : float | None
``OpenTimeoutSecs`` passed to the temporary read stream.
Returns
-------
dict[str, list]
History dictionary with keys:
``density``, ``flow``, ``temp`` (each list of ``(n_species, n_surf)``
arrays), ``gstep`` (list[int]), and ``time`` (list[float]).
Notes
-----
This helper is independent from live-stream operation and does not change
the behavior of the default constructor.
"""
history: Dict[str, list] = {
"density": [],
"flow": [],
"temp": [],
"gstep": [],
"time": [],
}
stride = max(1, int(stride))
path = Path(bp_dir)
io_name = f"diagnosis.diffusion_profiles_history_{os.getpid()}_{int(time.time() * 1e6)}"
io = adios.declare_io(io_name)
if io.engine_type().upper() == "FILE":
io.set_engine("FileStream")
if open_timeout is not None:
io.set_parameter("OpenTimeoutSecs", f"{open_timeout}")
try:
stream = Stream(io, str(path), "r")
except Exception as exc:
print(f"[history] Failed to open {path}: {exc}", flush=True)
return history
with stream:
istep = 0
while True:
status = stream.begin_step(timeout=0.0)
if status == bindings.StepStatus.OK:
keep = (istep % stride) == 0
vars_now = stream.available_variables()
if keep and all(v in vars_now for v in ("density", "flow", "temp")):
n = np.asarray(stream.read("density"))
u = np.asarray(stream.read("flow"))
T = np.asarray(stream.read("temp"))
if n.ndim == 3 and u.ndim == 3 and T.ndim == 3:
history["density"].append(n[:, -1, :].astype(float, copy=True))
history["flow"].append(u[:, -1, :].astype(float, copy=True))
history["temp"].append(T[:, -1, :].astype(float, copy=True))
g_arr = np.asarray(stream.read("gstep")).reshape(-1)
if g_arr.size == 0:
raise ValueError("Diffusion profile scalar 'gstep' is empty.")
gstep_val = int(g_arr[0])
history["gstep"].append(gstep_val)
t_val = np.nan
if "time" in vars_now:
t_arr = np.asarray(stream.read("time")).reshape(-1)
if t_arr.size > 0:
t_val = float(t_arr[-1])
history["time"].append(t_val)
stream.end_step()
istep += 1
continue
if status in (bindings.StepStatus.EndOfStream, bindings.StepStatus.NotReady):
break
print(f"[history] Unexpected step status {status} while scanning {path}", flush=True)
break
if max_steps is not None and max_steps > 0 and len(history["gstep"]) > max_steps:
for key in ("density", "flow", "temp", "gstep", "time"):
history[key] = history[key][-max_steps:]
return history
[docs]
def read_next_step(self, *, my_timeout: float = 10.0,
poll_delay: float = 0.5) -> bool:
"""Advance the stream by **one** time step.
Examples
--------
>>> prof = DiffusionProfiles("xgc.diffusion_profiles.bp")
>>> prof.read_next_step(block=False)
False # no new data yet
>>> # ... later
>>> prof.read_next_step()
True # newest snapshot now loaded
Parameters
----------
my_timeout
Maximum wait time for a new step; if no new step is found within
my_timeout seconds, return False.
poll_delay
Sleep duration (seconds) between availability checks when
*block* is *True*.
Returns
-------
bool
*True* if a new step was successfully read, *False* otherwise.
"""
if self._stream is None:
raise RuntimeError("Stream already closed.")
found = False
print(f"Waiting for next step in xgc.diffusion_profiles.bp", flush=True)
status = self._stream.begin_step(timeout=my_timeout)
# Could raise exception on error, but for now just print to let the close() try to call
if status == bindings.StepStatus.OtherError:
print("Stream error reading profiles", file=sys.stderr)
elif status == bindings.StepStatus.NotReady:
print("Timed out waiting for profiles step", file=sys.stderror)
elif status == bindings.StepStatus.EndOfStream:
print("EndofStream found")
elif status == bindings.StepStatus.OK:
print(f"Found next diffusion profiles step", flush=True)
self._read_current_step()
self._stream.end_step()
found = True
return found
[docs]
def close(self) -> None:
"""Close the underlying stream."""
if self._stream is not None:
self._stream.close()
self._stream = None
# ------------------------------------------------------------------
# Convenience properties – keep backward compatibility
# ------------------------------------------------------------------
@property
def n_species(self) -> int | None:
"""Number of species in the **current** snapshot (or ``None``)."""
return self._data.get("n_species") if self._data else None
@property
def n_samples(self) -> int | None:
"""Number of (time) samples per species in snapshot, else ``None``."""
return self._data.get("n_samples") if self._data else None
@property
def n_surf(self) -> int | None:
"""Radial surface count (psi grid size) in snapshot, else ``None``."""
return self._data.get("n_surf") if self._data else None
# ------------------------------------------------------------------
# Representation
# ------------------------------------------------------------------
def __repr__(self) -> str:
"""Unambiguous, one‑liner representation for debugging/logging."""
if self._data is None:
return f"{self.__class__.__name__}(waiting for data)"
def _shape(name: str) -> str:
"""Format the shape of one current-step array for the representation."""
arr = self._data.get(name)
return str(arr.shape) if isinstance(arr, np.ndarray) else "N/A"
pieces = [
f"step={self._step_index}",
f"n_species={self.n_species}",
f"n_samples={self.n_samples}",
f"n_surf={self.n_surf}",
f"density.shape={_shape('density')}",
f"psi.shape={_shape('psi')}",
]
return f"{self.__class__.__name__}({', '.join(pieces)})"
# ------------------------------------------------------------------
# Internals
# ------------------------------------------------------------------
def _wait_for_file(self, interval: float, timeout: Optional[float]) -> None:
"""Block until *self.path* exists (optional timeout)."""
print(f"Waiting for file {self.path}", flush=True)
start = time.monotonic()
while not self.path.exists():
if timeout is not None and (time.monotonic() - start) > timeout:
raise FileNotFoundError(
f"Timed out after {timeout} s waiting for '{self.path}'."
)
print(f"Waiting for file {self.path}", flush=True)
time.sleep(interval)
def _open_stream(self, timeout: Optional[float]) -> None:
"""Open the BP directory in *Stream* mode (read‑only)."""
# NB: we keep a dedicated ADIOS2 Stream object – no context manager on
# purpose so that the caller controls the lifetime via `close()`.
io = self.adios.declare_io("diagnosis.diffusion_profiles")
if io.engine_type().upper() == "FILE":
io.set_engine("FileStream")
print("Setting default engine FileStream for {0}".format(self.path))
if "OpenTimeoutSecs" not in io.parameters():
io.set_parameter("OpenTimeoutSecs", f"{timeout}")
print("Setting default timeout = {0} seconds to open {1}".format(timeout, self.path))
print(f"Waiting for file {self.path}", flush=True)
self._stream = Stream(io, str(self.path), "r")
print(f"File opened", flush=True)
#self._stream = Stream(str(self.path), "r")
def _drain_initial_steps(self, timeout: float = 10.0) -> None:
"""Consume all steps currently on disk – retain only the last one."""
assert self._stream is not None
def wait_for_next_step(timeout: float = 10.0) -> bool:
"""Wait for a new step to become available."""
status = self._stream.begin_step(timeout=timeout)
if status == bindings.StepStatus.OK:
return True
elif status == bindings.StepStatus.EndOfStream:
print("EndofStream found")
return False
elif status == bindings.StepStatus.OtherError:
raise RuntimeError("Stream error")
return False
# Read the first step (if available), and loop to discard intermediates
if not wait_for_next_step(timeout=timeout):
raise RuntimeError("Timed out waiting for the first step in diffusion profiles")
while True:
print("Reading step (draining)", flush=True)
self._read_current_step()
self._stream.end_step()
# Try to get the next step (within timeout)
if not wait_for_next_step(timeout=0):
# No more steps – keep current one open
print("Last available step retained", flush=True)
return
def _read_current_step(self) -> None:
"""Read variables in *the current* step into :attr:`_data`."""
assert self._stream is not None # for mypy
variables = self._stream.available_variables()
data: Dict[str, np.ndarray] = {}
for name in variables.keys():
# ADIOS2 returns NumPy arrays directly when shape > 0.
data[name] = self._stream.read(name)
# Overwrite previous snapshot
self._data = data
# ``current_step()`` became public in v2.9 – defensive fallback
try:
self._step_index = self._stream.current_step()
except AttributeError: # pragma: no cover
self._step_index = (self._step_index or 0) + 1
# ------------------------------------------------------------------
# Context‑manager helpers (optional)
# ------------------------------------------------------------------
def __enter__(self):
"""Enter context manager → returns *self* (stream already open)."""
return self
def __exit__(self, exc_type, exc, tb):
"""Ensure the stream is closed when leaving a *with* block."""
self.close()
# ------------------------------------------------------------------------------
# Example usage:
# ------------------------------------------------------------------------------
if __name__ == "__main__":
# Create an instance of DiffusionProfiles from xgc.diffusion_profiles.bp.
diff_profiles = DiffusionProfiles("xgc.diffusion_profiles.bp")
# Print some details.
print("DiffusionProfiles read from file:")
print(diff_profiles)
# For example, check the shape of the temperature array.
print("Temperature data shape:", diff_profiles.temp.shape)
# Check simulation time values.
print("Time steps:", diff_profiles.time)