externalize adios2py and the xarray adios2 reader

.pre-commit-config.yaml

@@ -63,7 +63,7 @@ repos:
         additional_dependencies:
           - pytest
           - xarray
-          - adios2
+          - adios2py
 
   - repo: https://github.com/codespell-project/codespell
     rev: "v2.3.0"

pyproject.toml

@@ -6,7 +6,7 @@ authors = [
 description = "Python utilities for reading PSC data"
 readme = "README.md"
 license.file = "LICENSE"
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 classifiers = [
   "Development Status :: 1 - Planning",
   "Intended Audience :: Science/Research",
@@ -23,7 +23,7 @@ classifiers = [
   "Typing :: Typed",
 ]
 dynamic = ["version"]
-dependencies = ["xarray", "adios2", "typing-extensions"]
+dependencies = ["xarray", "xarray-adios2", "typing-extensions"]
 
 [project.optional-dependencies]
 test = ["pytest >=6", "pytest-cov >=3"]
@@ -42,9 +42,6 @@ Homepage = "https://github.com/psc-code/pscpy"
 Discussions = "https://github.com/psc-code/pscpy/discussions"
 Changelog = "https://github.com/psc-code/pscpy/releases"
 
-[project.entry-points."xarray.backends"]
-pscadios2_engine = "pscpy.pscadios2:PscAdios2BackendEntrypoint"
-
 [build-system]
 build-backend = "setuptools.build_meta"
 requires = ["setuptools>=42", "setuptools-scm>=7"]
@@ -81,7 +78,7 @@ report.exclude_also = ['\.\.\.', 'if typing.TYPE_CHECKING:']
 
 [tool.mypy]
 files = ["src", "tests"]
-python_version = "3.8"
+python_version = "3.10"
 warn_unused_configs = true
 strict = true
 enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
@@ -147,3 +144,6 @@ messages_control.disable = [
   "missing-function-docstring",
   "wrong-import-position",
 ]
+
+[tool.uv]
+reinstall-package = ["pscpy"]

src/pscpy/__init__.py

@@ -8,11 +8,14 @@
 
 import pathlib
 
-from pscpy import pscadios2  # noqa: F401
-
 from ._version import version as __version__
+from .psc import decode_psc
 
 sample_dir = pathlib.Path(__file__).parent / "sample"
 
 
-__all__ = ["__version__"]
+__all__ = [
+    "__version__",
+    "decode_psc",
+    "sample_dir",
+]

src/pscpy/psc.py

@@ -1,23 +1,12 @@
 from __future__ import annotations
 
-from typing import Any, Iterable
+from collections.abc import Iterable
+from typing import Any
 
 import numpy as np
+import xarray as xr
 from numpy.typing import ArrayLike, NDArray
 
-from .adios2py import File
-
-
-def _get_array_attribute(
-    file: File, attribute_name: str, default: ArrayLike | None
-) -> NDArray[np.floating[Any]]:
-    if attribute_name in file.attribute_names:
-        return np.asarray(file.get_attribute(attribute_name))
-    if default is not None:
-        return np.asarray(default)
-    error_messsage = f"Missing attribute '{attribute_name}' with no default specified."
-    raise KeyError(error_messsage)
-
 
 class RunInfo:
     """Global information about the PSC run
@@ -28,23 +17,22 @@ class RunInfo:
 
     def __init__(
         self,
-        file: File,
+        ds: xr.Dataset,
         length: ArrayLike | None = None,
         corner: ArrayLike | None = None,
     ) -> None:
-        assert len(file.variable_names) > 0
-        var = next(iter(file.variable_names))
-        self.gdims = np.asarray(file.get_variable(var).shape)[0:3]
+        first_var = ds[next(iter(ds))]
+        self.gdims = np.asarray(first_var.shape)[::-1][:3]
 
-        self.length = _get_array_attribute(file, "length", length)
-        self.corner = _get_array_attribute(file, "corner", corner)
+        self.length = ds.attrs.get("length", length)
+        self.corner = ds.attrs.get("corner", corner)
 
         self.x = self._get_coord(0)
         self.y = self._get_coord(1)
         self.z = self._get_coord(2)
 
-    def _get_coord(self, coord_idx: int) -> NDArray[np.floating[Any]]:
-        return np.linspace(
+    def _get_coord(self, coord_idx: int) -> NDArray[Any]:
+        return np.linspace(  # type: ignore[no-any-return]
             start=self.corner[coord_idx],
             stop=self.corner[coord_idx] + self.length[coord_idx],
             num=self.gdims[coord_idx],
@@ -101,3 +89,44 @@ def get_field_to_component(species_names: Iterable[str]) -> dict[str, dict[str,
             )
 
     return field_to_component
+
+
+def decode_psc(
+    ds: xr.Dataset,
+    species_names: Iterable[str],
+    length: ArrayLike | None = None,
+    corner: ArrayLike | None = None,
+) -> xr.Dataset:
+    da = ds[next(iter(ds))]  # first dataset
+    if da.dims[0] == "dim_0_1":
+        # for compatibility, if dimensions weren't saved as attribute in the .bp file,
+        # fix them up here
+        ds = ds.rename_dims(
+            {
+                da.dims[0]: "step",
+                da.dims[1]: f"comp_{da.name}",
+                da.dims[2]: "z",
+                da.dims[3]: "y",
+                da.dims[4]: "x",
+            }
+        )
+    ds = ds.squeeze("step")
+    field_to_component = get_field_to_component(species_names)
+
+    data_vars = {}
+    for var_name in ds:
+        if var_name in field_to_component:
+            for field, component in field_to_component[var_name].items():  # type: ignore[index]
+                data_vars[field] = ds[var_name].isel({f"comp_{var_name}": component})
+    ds = ds.assign(data_vars)
+
+    if length is not None:
+        run_info = RunInfo(ds, length=length, corner=corner)
+        coords = {
+            "x": ("x", run_info.x),
+            "y": ("y", run_info.y),
+            "z": ("z", run_info.z),
+        }
+        ds = ds.assign_coords(coords)
+
+    return ds