Convert to psc output format v1

.pre-commit-config.yaml

@@ -64,6 +64,8 @@ repos:
           - pytest
           - xarray
           - adios2py
+          # numpy 2.5 stubs use PEP 695 syntax, rejected with python_version 3.10
+          - numpy<2.5
 
   - repo: https://github.com/codespell-project/codespell
     rev: "v2.3.0"

src/pscpy/convert/__init__.py

@@ -0,0 +1,8 @@
+"""Converters between psc output format versions.
+
+There is one module per format update, each converting from one version to
+the next (e.g. ``legacy_to_v1``, later ``v1_to_v2``). To bring old output up
+to date, run the converters in sequence::
+
+    python -m pscpy.convert.legacy_to_v1 -o OUTDIR SRC.bp ...
+"""

src/pscpy/convert/legacy_to_v1.py

@@ -0,0 +1,210 @@
+"""Convert legacy psc .bp output to psc_output_version 1.0.0.
+
+Legacy (pre-1.0.0) psc output stores all components of a field in a single
+4-d variable (e.g. "jeh" with shape (9, z, y, x)). Starting with 1.0.0, psc
+writes each component as its own (z, y, x) variable, together with
+cell-centered x/y/z coordinates and a scalar "time" variable.
+
+Usage::
+
+    python -m pscpy.convert.legacy_to_v1 -o OUTDIR [--species e,i] [--length LX LY LZ]
+        [--corner CX CY CZ] SRC.bp [SRC.bp ...]
+
+Each SRC.bp is written to OUTDIR with the same file name.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import pathlib
+from collections.abc import Sequence
+from typing import Any
+
+import adios2py
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+OUTPUT_VERSION = "1.0.0"
+"""The psc_output_version this converter writes (not necessarily the latest)."""
+
+JEH_COMPONENTS = ["jx_ec", "jy_ec", "jz_ec", "ex_ec", "ey_ec", "ez_ec", "hx_fc", "hy_fc", "hz_fc"]  # fmt: skip
+FIELD_COMPONENTS = {
+    "jeh": JEH_COMPONENTS,
+    "dive": ["dive"],
+    "rho": ["rho"],
+    "d_rho": ["d_rho"],
+    "dt_divj": ["dt_divj"],
+}
+MOMENT_FIELDS = ["all_1st", "all_1st_cc"]
+MOMENTS = ["rho", "jx", "jy", "jz", "px", "py", "pz", "txx", "tyy", "tzz", "txy", "tyz", "tzx"]  # fmt: skip
+
+
+def legacy_component_names(
+    field: str, n_components: int, species_names: Sequence[str] | None = None
+) -> list[str]:
+    """Return the names of the components stored in legacy variable `field`."""
+    if field in FIELD_COMPONENTS:
+        names = FIELD_COMPONENTS[field]
+    elif field in MOMENT_FIELDS:
+        if species_names is None:
+            message = f"species_names are required to convert moments ({field!r})"
+            raise ValueError(message)
+        names = [
+            f"{moment}_{species}" for species in species_names for moment in MOMENTS
+        ]
+    else:
+        message = f"Don't know how to convert legacy variable {field!r}"
+        raise ValueError(message)
+
+    if len(names) != n_components:
+        message = (
+            f"Legacy variable {field!r} has {n_components} components, "
+            f"expected {len(names)} ({', '.join(names)})"
+        )
+        raise ValueError(message)
+    return names
+
+
+def _unwrap_scalar(value: Any) -> Any:
+    """Legacy files sometimes store scalar attributes as 1-element arrays."""
+    arr = np.asarray(value)
+    if arr.ndim == 0:
+        return arr[()]
+    return arr[0]
+
+
+def _cell_centers(corner: float, length: float, n: int) -> NDArray[np.float64]:
+    return corner + (np.arange(n, dtype=np.float64) + 0.5) * (length / n)
+
+
+def _get_domain_attr(
+    attrs: Any, name: str, override: ArrayLike | None
+) -> NDArray[np.float64]:
+    if override is not None:
+        value = override
+    elif name in attrs:
+        value = attrs[name]
+    else:
+        message = f"{name} is missing from the file and must be provided."
+        raise ValueError(message)
+    arr = np.asarray(value, dtype=np.float64)
+    if arr.shape != (3,):
+        message = f"{name} must have 3 entries, got {arr.shape}"
+        raise ValueError(message)
+    return arr
+
+
+def convert_file(
+    src: os.PathLike[Any] | str,
+    dst: os.PathLike[Any] | str,
+    *,
+    species_names: Sequence[str] | None = None,
+    length: ArrayLike | None = None,
+    corner: ArrayLike | None = None,
+) -> None:
+    """Convert the legacy psc output file `src` to psc_output_version 1.0.0 at `dst`.
+
+    `length` and `corner` override the values stored in `src` (older files
+    don't have them). `species_names` is required for moment output.
+    """
+    if pathlib.Path(dst).exists():
+        message = f"{dst} already exists"
+        raise FileExistsError(message)
+
+    with adios2py.File(src, mode="rra") as file:
+        attrs = file.attrs
+        if "psc_output_version" in attrs:
+            message = f"{src} is not in the legacy format (psc_output_version={attrs['psc_output_version']!r})"
+            raise ValueError(message)
+        if len(file.steps) != 1:
+            message = f"{src} has {len(file.steps)} steps, expected 1"
+            raise ValueError(message)
+
+        length_arr = _get_domain_attr(attrs, "length", length)
+        corner_arr = _get_domain_attr(attrs, "corner", corner)
+        step_num = np.int32(_unwrap_scalar(attrs["step"]))
+        time = np.float64(_unwrap_scalar(attrs["time"]))
+
+        step_in = file.steps[0]
+        fields: dict[str, tuple[NDArray[Any], dict[str, Any]]] = {}
+        for field in step_in:
+            data = step_in[field][...]
+            if data.ndim != 4:
+                message = f"Legacy variable {field!r} has shape {data.shape}, expected (component, z, y, x)"
+                raise ValueError(message)
+            var_attrs = {
+                name: np.asarray(attrs[key], dtype=np.int32)
+                for name in ["ib", "im"]
+                for key in [f"{field}::{name}", name]
+                if key in attrs
+            }
+            names = legacy_component_names(field, data.shape[0], species_names)
+            for idx, name in enumerate(names):
+                fields[name] = (data[idx], var_attrs)
+
+    gdims_zyx = next(iter(fields.values()))[0].shape
+    gdims = gdims_zyx[::-1]
+
+    with adios2py.File(dst, mode="w") as out, out.steps.next() as step:
+        out.attrs["psc_output_version"] = OUTPUT_VERSION
+        out.attrs["step"] = step_num
+        out.attrs["time"] = time
+        out.attrs["length"] = length_arr
+        out.attrs["corner"] = corner_arr
+        out.attrs["step_dimension"] = "time"
+
+        step["time"] = time
+        for d, crd_name in enumerate("xyz"):
+            step[crd_name] = _cell_centers(corner_arr[d], length_arr[d], gdims[d])
+            step[crd_name].attrs["dimensions"] = crd_name
+
+        for name, (data, var_attrs) in fields.items():
+            if data.shape != gdims_zyx:
+                message = (
+                    f"Component {name!r} has shape {data.shape}, expected {gdims_zyx}"
+                )
+                raise ValueError(message)
+            step[name] = data
+            step[name].attrs["dimensions"] = "z y x"
+            for attr_name, value in var_attrs.items():
+                step[name].attrs[attr_name] = value
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(
+        prog="python -m pscpy.convert.legacy_to_v1",
+        description=f"Convert legacy psc .bp output to psc_output_version {OUTPUT_VERSION}.",
+    )
+    parser.add_argument("src", nargs="+", type=pathlib.Path, help="legacy .bp files")
+    parser.add_argument(
+        "-o", "--outdir", required=True, type=pathlib.Path, help="output directory"
+    )
+    parser.add_argument(
+        "--species",
+        type=lambda arg: arg.split(","),
+        help="comma-separated species names, required for moments (e.g. e,i)",
+    )
+    parser.add_argument(
+        "--length", nargs=3, type=float, help="domain length, if missing from the file"
+    )
+    parser.add_argument(
+        "--corner", nargs=3, type=float, help="domain corner, if missing from the file"
+    )
+    args = parser.parse_args(argv)
+
+    args.outdir.mkdir(parents=True, exist_ok=True)
+    for src in args.src:
+        dst = args.outdir / src.name
+        convert_file(
+            src,
+            dst,
+            species_names=args.species,
+            length=args.length,
+            corner=args.corner,
+        )
+        print(f"{src} -> {dst}")  # noqa: T201
+
+
+if __name__ == "__main__":
+    main()

src/pscpy/psc.py

@@ -1,108 +1,49 @@
 from __future__ import annotations
 
-from collections.abc import Generator, Hashable, Iterable
-from typing import Any
-
-import numpy as np
 import xarray as xr
-from numpy.typing import ArrayLike, NDArray
-
-
-class RunInfo:
-    """Global information about the PSC run
-
-    Currently stores domain info.
-    TODO: Should also know about timestep, species, whatever...
-    """
-
-    def __init__(
-        self,
-        ds: xr.Dataset,
-        length: ArrayLike | None = None,
-        corner: ArrayLike | None = None,
-    ) -> None:
-        first_var = ds[next(iter(ds))]
-        self.gdims = np.asarray(first_var.shape)[::-1][:3]
-
-        self.length = ds.attrs.get("length", length)
-        self.corner = ds.attrs.get("corner", corner)
+from typing_extensions import deprecated
 
-        if self.length is None:
-            message = "Dataset is missing length. A value must be manually provided."
-            raise ValueError(message)
-        if self.corner is None:
-            message = "Dataset is missing corner. A value must be manually provided."
-            raise ValueError(message)
+SUPPORTED_MAJOR_VERSION = 1
 
-        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[Any]:
-        return np.linspace(
-            start=self.corner[coord_idx],
-            stop=self.corner[coord_idx] + self.length[coord_idx],
-            num=self.gdims[coord_idx],
-            endpoint=False,
+def check_psc_output_version(attrs: dict[str, object]) -> None:
+    """Raise ValueError unless `attrs` has a supported "psc_output_version"."""
+    version = attrs.get("psc_output_version")
+    if version is None:
+        message = (
+            "Dataset has no psc_output_version attribute, so it was written in the "
+            "legacy (pre-1.0.0) psc output format. Convert it first with "
+            "`python -m pscpy.convert.legacy_to_v1`."
         )
+        raise ValueError(message)
 
-    def __repr__(self) -> str:
-        return f"Psc(gdims={self.gdims}, length={self.length}, corner={self.corner})"
-
-
-def iter_components(field: Hashable, species_names: Iterable[str]) -> Generator[str]:
-    if field == "jeh":
-        yield from ["jx_ec", "jy_ec", "jz_ec", "ex_ec", "ey_ec", "ez_ec", "hx_fc", "hy_fc", "hz_fc"]  # fmt: off
-    elif field in ["dive", "rho", "d_rho", "dt_divj"]:
-        yield str(field)
-    elif field in ["all_1st", "all_1st_cc"]:
-        moments = ["rho", "jx", "jy", "jz", "px", "py", "pz", "txx", "tyy", "tzz", "txy", "tyz", "tzx"]  # fmt: off
-        for species_name in species_names:
-            for moment in moments:
-                yield f"{moment}_{species_name}"
-
-
-def unwrap_float(arr: np.ndarray) -> float:
-    if arr.ndim == 0:
-        return float(arr)
-    return float(arr[0])
+    major = str(version).split(".", maxsplit=1)[0]
+    if major != str(SUPPORTED_MAJOR_VERSION):
+        message = (
+            f"Unsupported psc_output_version {version!r}; "
+            f"only {SUPPORTED_MAJOR_VERSION}.x.y is supported."
+        )
+        if major.isdigit() and int(major) < SUPPORTED_MAJOR_VERSION:
+            message += " Convert older output with the modules in pscpy.convert."
+        raise ValueError(message)
 
 
-def decode_psc(
-    ds: xr.Dataset,
-    species_names: Iterable[str],
-    length: ArrayLike | None = None,
-    corner: ArrayLike | None = None,
-) -> xr.Dataset:
-    dims = list(ds.dims)
-    if "dim_0_1" in dims:
-        # for compatibility, if dimensions weren't saved as attribute in the .bp file,
-        # fix them up here
-        ds = ds.rename_dims(
-            {
-                dims[0]: "step",
-                dims[1]: "component",
-                dims[2]: "z",
-                dims[3]: "y",
-                dims[4]: "x",
-            }
-        )
-    ds = ds.squeeze("step")
+@deprecated(
+    "As of PSC output version v1, `decode_psc`'s only transformative action is to squeeze `time` and rename it to `t` to preserve backwards compatibility. From now on, the officially-supported way to get usable PSC data is to call `xr.open_dataset()` or `xr.open_mfdataset()` with default options."
+)
+def decode_psc(ds: xr.Dataset) -> xr.Dataset:
+    """Decode a dataset written by psc (psc_output_version 1.x.y).
 
-    for var_name in ds:
-        components = list(iter_components(var_name, species_names))
-        for component_idx, component in enumerate(components):
-            ds = ds.assign({component: ds[var_name].isel(component=component_idx)})
-        if var_name not in components:
-            ds = ds.drop_vars([var_name])
+    Each field component is already its own variable with (z, y, x) dims and
+    cell-centered coordinates, so this only validates the format version,
+    renames "time" to "t" and drops the "t" dimension if the dataset contains
+    a single step.
+    """
+    check_psc_output_version(ds.attrs)
 
-    run_info = RunInfo(ds, length=length, corner=corner)
-    coords = {
-        "x": ("x", run_info.x),
-        "y": ("y", run_info.y),
-        "z": ("z", run_info.z),
-        "t": unwrap_float(ds.attrs["time"]),
-    }
-    ds = ds.assign_coords(coords)
+    if "time" in ds.variables or "time" in ds.dims:
+        ds = ds.rename(time="t")
+    if ds.sizes.get("t") == 1:
+        ds = ds.squeeze("t")
 
     return ds

src/pscpy/sample/legacy/continuity.000000001.bp/profiling.json

@@ -0,0 +1,3 @@
+[
+{ "rank":0, "start":"Fri_Sep_25_11:57:37_2026","InitAgg-tls_mus": 19, "InitAgg-tls":{"mus":19, "nCalls":1},"ES_AWD_mus": 87, "ES_AWD":{"mus":87, "nCalls":1},"ES_write_metadata_mus": 29, "ES_write_metadata":{"mus":29, "nCalls":1},"MetaInfoBcast_mus": 0, "MetaInfoBcast":{"mus":0, "nCalls":1},"ES_gather_write_meta_mus": 94, "ES_gather_write_meta":{"mus":94, "nCalls":1},"ES_aggregate_info_mus": 5, "ES_aggregate_info":{"mus":5, "nCalls":1},"PP_mus": 0, "PP":{"mus":0, "nCalls":3},"SelectMetaInfoGather_mus": 0, "SelectMetaInfoGather":{"mus":0, "nCalls":1},"ES_GatherMetadataBlocks_mus": 0, "ES_GatherMetadataBlocks":{"mus":0, "nCalls":1},"ES_AGG1_mus": 45, "ES_AGG1":{"mus":45, "nCalls":1},"MetadataBlockWrite_mus": 3, "MetadataBlockWrite":{"mus":3, "nCalls":1},"ES_mus": 253, "ES":{"mus":253, "nCalls":1},"ES_close_mus": 64, "ES_close":{"mus":64, "nCalls":1},"FixedMetaInfoGather_mus": 0, "FixedMetaInfoGather":{"mus":0, "nCalls":1}, "databytes":0, "metadatabytes":0, "metametadatabytes":0, "transport_0":{"type":"File_POSIX", "wbytes":512, "close":{"mus":7, "nCalls":1}, "write":{"mus":19, "nCalls":1}, "open":{"mus":56, "nCalls":1}}, "transport_1":{"type":"File_POSIX", "wbytes":2176, "close":{"mus":5, "nCalls":1}, "write":{"mus":24, "nCalls":5}, "open":{"mus":57, "nCalls":1}} }
+]

src/pscpy/sample/legacy/gauss.000000001.bp/data.0

@@ -0,0 +1,3 @@
+|—mè_ÿw¿5Ù5åëƒ?ªÛ|`?=’¦÷@—k¿9ª½Ä?d÷”cËJ]¿ú€§(©0v¿5u‘$Ôp?¤9iìbGt?_öž”TŠ€¿Ã,‹1Hf?©²,¿‚et…wü?¥žÌ»°‚?àcÏïGi?EÐÝÙ>ó?:%ÿ¢ËŚ¿r*·™Œ?BÂ\0†¿[˶jd5T?ŽšX?¼;#±pK?T²ÎeR¿IP'žšÎƒ¿:ù?G¥ÉQ¿?â¸"×yc¿Þ^
+Nþƒ¿ú¶]'¿B×9wŽa?$QªôZ¥‡¿,À_É7‚?¯ :òr€?B—mè_ÿw¿`Ù5åëƒ?jªÛ|`?:’¦÷@—k¿â8ª½Ä?>ù”cËJ]¿í€§(©0v¿Su‘$Ôp?™9iìbGt?Söž”TŠ€¿|Ã,‹1Hf?à"©²,¿¸et…wü?ञ̻°‚?¯cÏïGi?iÐÝÙ>ó?8%ÿ¢ËŚ¿“*·™Œ?»Ã‚\0†¿0˶jd5T?€šX?8;#±pK?±ÎeR¿2P'žšÎƒ¿,ø?G¥ÉQ¿Ôâ¸"×yc¿é^
+Nþƒ¿ê¶]'¿÷Ö9wŽa?QªôZ¥‡¿À_É7‚?« :òr€?

src/pscpy/sample/legacy/gauss.000000001.bp/profiling.json

@@ -0,0 +1,3 @@
+[
+{ "rank":0, "start":"Fri_Sep_25_11:57:37_2026","InitAgg-tls_mus": 19, "InitAgg-tls":{"mus":19, "nCalls":1},"ES_AWD_mus": 86, "ES_AWD":{"mus":86, "nCalls":1},"ES_write_metadata_mus": 28, "ES_write_metadata":{"mus":28, "nCalls":1},"MetaInfoBcast_mus": 0, "MetaInfoBcast":{"mus":0, "nCalls":1},"ES_gather_write_meta_mus": 91, "ES_gather_write_meta":{"mus":91, "nCalls":1},"ES_aggregate_info_mus": 5, "ES_aggregate_info":{"mus":5, "nCalls":1},"PP_mus": 0, "PP":{"mus":0, "nCalls":3},"SelectMetaInfoGather_mus": 0, "SelectMetaInfoGather":{"mus":0, "nCalls":1},"ES_GatherMetadataBlocks_mus": 0, "ES_GatherMetadataBlocks":{"mus":0, "nCalls":1},"ES_AGG1_mus": 41, "ES_AGG1":{"mus":41, "nCalls":1},"MetadataBlockWrite_mus": 2, "MetadataBlockWrite":{"mus":2, "nCalls":1},"ES_mus": 250, "ES":{"mus":250, "nCalls":1},"ES_close_mus": 65, "ES_close":{"mus":65, "nCalls":1},"FixedMetaInfoGather_mus": 0, "FixedMetaInfoGather":{"mus":0, "nCalls":1}, "databytes":0, "metadatabytes":0, "metametadatabytes":0, "transport_0":{"type":"File_POSIX", "wbytes":512, "close":{"mus":14, "nCalls":1}, "write":{"mus":20, "nCalls":1}, "open":{"mus":56, "nCalls":1}}, "transport_1":{"type":"File_POSIX", "wbytes":2168, "close":{"mus":6, "nCalls":1}, "write":{"mus":24, "nCalls":5}, "open":{"mus":69, "nCalls":1}} }
+]

src/pscpy/sample/legacy/pfd.000000001.bp/profiling.json

@@ -0,0 +1,3 @@
+[
+{ "rank":0, "start":"Fri_Sep_25_11:57:37_2026","InitAgg-tls_mus": 19, "InitAgg-tls":{"mus":19, "nCalls":1},"ES_AWD_mus": 91, "ES_AWD":{"mus":91, "nCalls":1},"ES_write_metadata_mus": 29, "ES_write_metadata":{"mus":29, "nCalls":1},"MetaInfoBcast_mus": 0, "MetaInfoBcast":{"mus":0, "nCalls":1},"ES_gather_write_meta_mus": 92, "ES_gather_write_meta":{"mus":92, "nCalls":1},"ES_aggregate_info_mus": 5, "ES_aggregate_info":{"mus":5, "nCalls":1},"PP_mus": 0, "PP":{"mus":0, "nCalls":2},"SelectMetaInfoGather_mus": 0, "SelectMetaInfoGather":{"mus":0, "nCalls":1},"ES_GatherMetadataBlocks_mus": 0, "ES_GatherMetadataBlocks":{"mus":0, "nCalls":1},"ES_AGG1_mus": 37, "ES_AGG1":{"mus":37, "nCalls":1},"MetadataBlockWrite_mus": 2, "MetadataBlockWrite":{"mus":2, "nCalls":1},"ES_mus": 251, "ES":{"mus":251, "nCalls":1},"ES_close_mus": 61, "ES_close":{"mus":61, "nCalls":1},"FixedMetaInfoGather_mus": 0, "FixedMetaInfoGather":{"mus":0, "nCalls":1}, "databytes":0, "metadatabytes":0, "metametadatabytes":0, "transport_0":{"type":"File_POSIX", "wbytes":2304, "close":{"mus":8, "nCalls":1}, "write":{"mus":25, "nCalls":1}, "open":{"mus":54, "nCalls":1}}, "transport_1":{"type":"File_POSIX", "wbytes":1272, "close":{"mus":5, "nCalls":1}, "write":{"mus":25, "nCalls":5}, "open":{"mus":68, "nCalls":1}} }
+]

src/pscpy/sample/legacy/pfd_moments.000000001.bp/profiling.json

@@ -0,0 +1,3 @@
+[
+{ "rank":0, "start":"Fri_Sep_25_11:57:37_2026","InitAgg-tls_mus": 18, "InitAgg-tls":{"mus":18, "nCalls":1},"ES_AWD_mus": 225, "ES_AWD":{"mus":225, "nCalls":1},"ES_write_metadata_mus": 29, "ES_write_metadata":{"mus":29, "nCalls":1},"MetaInfoBcast_mus": 0, "MetaInfoBcast":{"mus":0, "nCalls":1},"ES_gather_write_meta_mus": 87, "ES_gather_write_meta":{"mus":87, "nCalls":1},"ES_aggregate_info_mus": 5, "ES_aggregate_info":{"mus":5, "nCalls":1},"PP_mus": 0, "PP":{"mus":0, "nCalls":2},"SelectMetaInfoGather_mus": 0, "SelectMetaInfoGather":{"mus":0, "nCalls":1},"ES_GatherMetadataBlocks_mus": 0, "ES_GatherMetadataBlocks":{"mus":0, "nCalls":1},"ES_AGG1_mus": 39, "ES_AGG1":{"mus":39, "nCalls":1},"MetadataBlockWrite_mus": 3, "MetadataBlockWrite":{"mus":3, "nCalls":1},"ES_mus": 379, "ES":{"mus":379, "nCalls":1},"ES_close_mus": 59, "ES_close":{"mus":59, "nCalls":1},"FixedMetaInfoGather_mus": 0, "FixedMetaInfoGather":{"mus":0, "nCalls":1}, "databytes":0, "metadatabytes":0, "metametadatabytes":0, "transport_0":{"type":"File_POSIX", "wbytes":6656, "close":{"mus":9, "nCalls":1}, "write":{"mus":158, "nCalls":1}, "open":{"mus":66, "nCalls":1}}, "transport_1":{"type":"File_POSIX", "wbytes":1280, "close":{"mus":5, "nCalls":1}, "write":{"mus":24, "nCalls":5}, "open":{"mus":76, "nCalls":1}} }
+]