From d6607865f03409df68a793ef28f94fc58040a91e Mon Sep 17 00:00:00 2001 From: Kai Germaschewski Date: Fri, 20 Dec 2024 17:22:48 -0500 Subject: [PATCH 1/7] pre-commit: turn off some unused checks no C++ / cmake / rst here --- .pre-commit-config.yaml | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 94b514c..cbc0c34 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -24,15 +24,15 @@ repos: - id: mixed-line-ending - id: name-tests-test args: ["--pytest-test-first"] - - id: requirements-txt-fixer + # - id: requirements-txt-fixer - id: trailing-whitespace - - repo: https://github.com/pre-commit/pygrep-hooks - rev: "v1.10.0" - hooks: - - id: rst-backticks - - id: rst-directive-colons - - id: rst-inline-touching-normal + # - repo: https://github.com/pre-commit/pygrep-hooks + # rev: "v1.10.0" + # hooks: + # - id: rst-backticks + # - id: rst-directive-colons + # - id: rst-inline-touching-normal - repo: https://github.com/pre-commit/mirrors-prettier rev: "v4.0.0-alpha.8" @@ -48,11 +48,11 @@ repos: args: ["--fix", "--show-fixes"] - id: ruff-format - - repo: https://github.com/pre-commit/mirrors-clang-format - rev: "v19.1.5" - hooks: - - id: clang-format - types_or: [c++, c, cuda] + # - repo: https://github.com/pre-commit/mirrors-clang-format + # rev: "v19.1.5" + # hooks: + # - id: clang-format + # types_or: [c++, c, cuda] - repo: https://github.com/pre-commit/mirrors-mypy rev: "v1.13.0" @@ -83,10 +83,10 @@ repos: entry: PyBind|Numpy|Cmake|CCache|Github|PyTest exclude: .pre-commit-config.yaml - - repo: https://github.com/cheshirekow/cmake-format-precommit - rev: "v0.6.13" - hooks: - - id: cmake-format + # - repo: https://github.com/cheshirekow/cmake-format-precommit + # rev: "v0.6.13" + # hooks: + # - id: cmake-format - repo: https://github.com/abravalheri/validate-pyproject rev: "v0.23" From dd787ddf3f3e52d59407ac7302c8f6919d82c814 Mon Sep 17 00:00:00 2001 From: Kai Germaschewski Date: Fri, 20 Dec 2024 17:23:26 -0500 Subject: [PATCH 2/7] adios2py: use logger for the module --- src/pscpy/adios2py/__init__.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/pscpy/adios2py/__init__.py b/src/pscpy/adios2py/__init__.py index f3b0ae3..e9f2d92 100644 --- a/src/pscpy/adios2py/__init__.py +++ b/src/pscpy/adios2py/__init__.py @@ -13,6 +13,8 @@ from numpy.typing import NDArray from typing_extensions import TypeGuard +logger = logging.getLogger(__name__) + _ad = Adios() @@ -25,10 +27,10 @@ def __init__(self, var: adios2.Variable, engine: adios2.Engine) -> None: self.name = self._name() self.shape = self._shape() self.dtype = self._dtype() - logging.debug("variable __init__ var %s engine %s", var, engine) + logger.debug("variable __init__ var %s engine %s", var, engine) def close(self) -> None: - logging.debug("adios2py.variable close") + logger.debug("adios2py.variable close") self._var = None self._engine = None @@ -126,7 +128,7 @@ class File: _state: FileState | None def __init__(self, filename: str | os.PathLike[Any], mode: str = "r") -> None: - logging.debug("adios2py: __init__ %s", filename) + logger.debug("File.__init__(%s, %s)", filename, mode) assert mode == "r" self._state = FileState(filename) self._open_vars: dict[str, Variable] = {} @@ -135,23 +137,22 @@ def __init__(self, filename: str | os.PathLike[Any], mode: str = "r") -> None: self.attribute_names: Collection[str] = self._state.io.available_attributes().keys() def __enter__(self) -> File: - logging.debug("adios2py: __enter__") + logger.debug("File.__enter__()") return self def __exit__(self, exception_type: type[BaseException] | None, exception: BaseException | None, traceback: TracebackType | None) -> None: - logging.debug("adios2py: __exit__") + logger.debug("File.__exit__()") self.close() def __del__(self) -> None: - logging.debug("adios2py: __del__") + logger.debug("File.__del__()") if FileState.is_open(self._state): self.close() def close(self) -> None: assert FileState.is_open(self._state) - logging.debug("adios2py: close") - logging.debug("open vars %s", self._open_vars) + logger.debug("File.close(): open vars %s", self._open_vars) for var in self._open_vars.values(): var.close() From d9de8dc87b8778cc2b6bfd8de188ac6427665770 Mon Sep 17 00:00:00 2001 From: Kai Germaschewski Date: Fri, 20 Dec 2024 17:38:58 -0500 Subject: [PATCH 3/7] adios2py: more extensive testing; disable mypy on tests mypy is again (or still) unhappy about not finding pscpy types :( --- .pre-commit-config.yaml | 2 +- tests/test_adios2py.py | 34 ++++++++++++++++++++++++++++++++-- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cbc0c34..a7fba72 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -58,7 +58,7 @@ repos: rev: "v1.13.0" hooks: - id: mypy - files: src|tests + files: src #|tests args: [] additional_dependencies: - pytest diff --git a/tests/test_adios2py.py b/tests/test_adios2py.py index c00e883..b48fa2b 100644 --- a/tests/test_adios2py.py +++ b/tests/test_adios2py.py @@ -1,8 +1,38 @@ from __future__ import annotations +import numpy as np + import pscpy from pscpy import adios2py -def test_adios2py_open(): - adios2py.File(pscpy.sample_dir / "pfd.000000400.bp") +def test_open_close(): + file = adios2py.File(pscpy.sample_dir / "pfd.000000400.bp") + file.close() + + +def test_with(): + with adios2py.File(pscpy.sample_dir / "pfd.000000400.bp"): + pass + + +def test_variable_names(): + with adios2py.File(pscpy.sample_dir / "pfd.000000400.bp") as file: + assert file.variable_names == set({"jeh"}) + assert file.attribute_names == set({"ib", "im", "step", "time"}) + + +def test_get_variable(): + with adios2py.File(pscpy.sample_dir / "pfd.000000400.bp") as file: + var = file.get_variable("jeh") + assert var.name == "jeh" + assert var.shape == (1, 128, 512, 9) + assert var.dtype == np.float32 + + +def test_get_attribute(): + with adios2py.File(pscpy.sample_dir / "pfd.000000400.bp") as file: + assert all(file.get_attribute("ib") == (0, 0, 0)) + assert all(file.get_attribute("im") == (1, 128, 128)) + assert np.isclose(file.get_attribute("time"), 109.38) + assert file.get_attribute("step") == 400 From 5b35bf8593b3bdd3179bb3f35d9867f662581b93 Mon Sep 17 00:00:00 2001 From: Kai Germaschewski Date: Fri, 20 Dec 2024 21:19:33 -0500 Subject: [PATCH 4/7] tests: some more testing for xarray adios2 backend --- tests/test_xarray_adios2.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_xarray_adios2.py b/tests/test_xarray_adios2.py index cf08ecd..77022a7 100644 --- a/tests/test_xarray_adios2.py +++ b/tests/test_xarray_adios2.py @@ -1,9 +1,14 @@ from __future__ import annotations +import numpy as np import xarray as xr import pscpy def test_open_dataset(): - xr.open_dataset(pscpy.sample_dir / "pfd.000000400.bp", species_names=[], length=[10, 10, 10], corner=[0, 0, 0]) + ds = xr.open_dataset(pscpy.sample_dir / "pfd.000000400.bp", species_names=[], length=[1, 12.8, 51.2], corner=[0, -6.4, -25.6]) + assert "jx_ec" in ds + assert ds.coords.keys() == set({"x", "y", "z"}) + assert ds.jx_ec.shape == (1, 128, 512) + assert np.allclose(ds.jx_ec.z, np.linspace(-25.6, 25.6, 512, endpoint=False)) From cebf0ce7e7811021b4779c9925195b4edbfc2206 Mon Sep 17 00:00:00 2001 From: Kai Germaschewski Date: Fri, 20 Dec 2024 21:24:30 -0500 Subject: [PATCH 5/7] devops: leave ruff format options at defaults So we don't get the very long lines... --- noxfile.py | 12 +++++-- pyproject.toml | 4 --- src/pscpy/adios2py/__init__.py | 31 +++++++++++++----- src/pscpy/psc.py | 19 +++++++++--- src/pscpy/pscadios2.py | 57 ++++++++++++++++++++++++++-------- tests/test_xarray_adios2.py | 7 ++++- 6 files changed, 98 insertions(+), 32 deletions(-) diff --git a/noxfile.py b/noxfile.py index e786f4c..2a9d299 100644 --- a/noxfile.py +++ b/noxfile.py @@ -19,7 +19,9 @@ def lint(session: nox.Session) -> None: Run the linter. """ session.install("pre-commit") - session.run("pre-commit", "run", "--all-files", "--show-diff-on-failure", *session.posargs) + session.run( + "pre-commit", "run", "--all-files", "--show-diff-on-failure", *session.posargs + ) @nox.session @@ -50,7 +52,9 @@ def docs(session: nox.Session) -> None: parser = argparse.ArgumentParser() parser.add_argument("--serve", action="store_true", help="Serve after building") - parser.add_argument("-b", dest="builder", default="html", help="Build target (default: html)") + parser.add_argument( + "-b", dest="builder", default="html", help="Build target (default: html)" + ) args, posargs = parser.parse_known_args(session.posargs) if args.builder != "html" and args.serve: @@ -62,7 +66,9 @@ def docs(session: nox.Session) -> None: session.chdir("docs") if args.builder == "linkcheck": - session.run("sphinx-build", "-b", "linkcheck", ".", "_build/linkcheck", *posargs) + session.run( + "sphinx-build", "-b", "linkcheck", ".", "_build/linkcheck", *posargs + ) return shared_args = ( diff --git a/pyproject.toml b/pyproject.toml index 996a2b9..46002a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,10 +95,6 @@ disallow_untyped_defs = true disallow_incomplete_defs = true -[tool.ruff] -src = ["src"] -line-length = 200 - [tool.ruff.lint] extend-select = [ "B", # flake8-bugbear diff --git a/src/pscpy/adios2py/__init__.py b/src/pscpy/adios2py/__init__.py index e9f2d92..517b644 100644 --- a/src/pscpy/adios2py/__init__.py +++ b/src/pscpy/adios2py/__init__.py @@ -39,7 +39,9 @@ def _assert_not_closed(self) -> None: error_message = "adios2py: variable is closed" raise ValueError(error_message) - def _set_selection(self, start: NDArray[np.integer[Any]], count: NDArray[np.integer[Any]]) -> None: + def _set_selection( + self, start: NDArray[np.integer[Any]], count: NDArray[np.integer[Any]] + ) -> None: self._assert_not_closed() self._var.set_selection((start[::-1], count[::-1])) @@ -59,7 +61,9 @@ def _dtype(self) -> np.dtype[Any]: return np.dtype(adios2.type_adios_to_numpy(self._var.type())) # type: ignore[no-any-return] - def __getitem__(self, args: SupportsInt | slice | tuple[SupportsInt | slice, ...]) -> NDArray[Any]: + def __getitem__( + self, args: SupportsInt | slice | tuple[SupportsInt | slice, ...] + ) -> NDArray[Any]: self._assert_not_closed() if not isinstance(args, tuple): @@ -101,7 +105,9 @@ def __getitem__(self, args: SupportsInt | slice | tuple[SupportsInt | slice, ... self._set_selection(sel_start, sel_count) - arr = np.empty(arr_shape, dtype=self.dtype, order="F") # FIXME is column-major correct? + arr = np.empty( + arr_shape, dtype=self.dtype, order="F" + ) # FIXME is column-major correct? self._engine.get(self._var, arr, adios2.bindings.Mode.Sync) return arr @@ -133,14 +139,23 @@ def __init__(self, filename: str | os.PathLike[Any], mode: str = "r") -> None: self._state = FileState(filename) self._open_vars: dict[str, Variable] = {} - self.variable_names: Collection[str] = self._state.io.available_variables().keys() - self.attribute_names: Collection[str] = self._state.io.available_attributes().keys() + self.variable_names: Collection[str] = ( + self._state.io.available_variables().keys() + ) + self.attribute_names: Collection[str] = ( + self._state.io.available_attributes().keys() + ) def __enter__(self) -> File: logger.debug("File.__enter__()") return self - def __exit__(self, exception_type: type[BaseException] | None, exception: BaseException | None, traceback: TracebackType | None) -> None: + def __exit__( + self, + exception_type: type[BaseException] | None, + exception: BaseException | None, + traceback: TracebackType | None, + ) -> None: logger.debug("File.__exit__()") self.close() @@ -163,7 +178,9 @@ def close(self) -> None: def get_variable(self, variable_name: str) -> Variable: assert FileState.is_open(self._state) - var = Variable(self._state.io.inquire_variable(variable_name), self._state.engine) + var = Variable( + self._state.io.inquire_variable(variable_name), self._state.engine + ) self._open_vars[variable_name] = var return var diff --git a/src/pscpy/psc.py b/src/pscpy/psc.py index 2db072a..8816784 100644 --- a/src/pscpy/psc.py +++ b/src/pscpy/psc.py @@ -8,7 +8,9 @@ from .adios2py import File -def _get_array_attribute(file: File, attribute_name: str, default: ArrayLike | None) -> NDArray[np.floating[Any]]: +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: @@ -24,7 +26,12 @@ class RunInfo: TODO: Should also know about timestep, species, whatever... """ - def __init__(self, file: File, length: ArrayLike | None = None, corner: ArrayLike | None = None) -> None: + def __init__( + self, + file: File, + 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] @@ -86,7 +93,11 @@ def get_field_to_component(species_names: Iterable[str]) -> dict[str, dict[str, ] for species_idx, species_name in enumerate(species_names): for moment_idx, moment in enumerate(moments): - field_to_component["all_1st"][f"{moment}_{species_name}"] = moment_idx + 13 * species_idx - field_to_component["all_1st_cc"][f"{moment}_{species_name}"] = moment_idx + 13 * species_idx + field_to_component["all_1st"][f"{moment}_{species_name}"] = ( + moment_idx + 13 * species_idx + ) + field_to_component["all_1st_cc"][f"{moment}_{species_name}"] = ( + moment_idx + 13 * species_idx + ) return field_to_component diff --git a/src/pscpy/pscadios2.py b/src/pscpy/pscadios2.py index bde3a44..aec99d8 100644 --- a/src/pscpy/pscadios2.py +++ b/src/pscpy/pscadios2.py @@ -50,7 +50,13 @@ class PscAdios2Array(BackendArray): This also takes care of slicing out the specific component of the data stored as 4-d array. """ - def __init__(self, variable_name: str, datastore: PscAdios2Store, orig_varname: str, component: int) -> None: + def __init__( + self, + variable_name: str, + datastore: PscAdios2Store, + orig_varname: str, + component: int, + ) -> None: self.variable_name = variable_name self.datastore = datastore self._orig_varname = orig_varname @@ -63,11 +69,17 @@ def get_array(self, needs_lock: bool = True) -> Variable: return self.datastore.acquire(needs_lock).get_variable(self._orig_varname) def __getitem__(self, key: indexing.ExplicitIndexer) -> Any: - return indexing.explicit_indexing_adapter(key, self.shape, indexing.IndexingSupport.BASIC, self._getitem) + return indexing.explicit_indexing_adapter( + key, self.shape, indexing.IndexingSupport.BASIC, self._getitem + ) - def _getitem(self, args: tuple[SupportsInt | slice, ...]) -> NDArray[np.floating[Any]]: + def _getitem( + self, args: tuple[SupportsInt | slice, ...] + ) -> NDArray[np.floating[Any]]: with self.datastore.lock: - return self.get_array(needs_lock=False)[(*args, self._component)] # FIXME add ... in between + return self.get_array(needs_lock=False)[ + (*args, self._component) + ] # FIXME add ... in between class PscAdios2Store(AbstractDataStore): @@ -98,13 +110,15 @@ def open( corner: ArrayLike | None = None, ) -> PscAdios2Store: if lock is None: - if mode == "r": # noqa: SIM108 + if mode == "r": lock = ADIOS2_LOCK else: lock = combine_locks([ADIOS2_LOCK, get_write_lock(filename)]) # type: ignore[no-untyped-call] manager = CachingFileManager(File, filename, mode=mode) - return PscAdios2Store(manager, species_names, mode=mode, lock=lock, length=length, corner=corner) + return PscAdios2Store( + manager, species_names, mode=mode, lock=lock, length=length, corner=corner + ) def acquire(self, needs_lock: bool = True) -> File: with self._manager.acquire_context(needs_lock) as root: @@ -125,17 +139,26 @@ def get_variables(self) -> Frozen[str, xarray.DataArray]: for field, component in field_to_component[orig_varname].items(): variables[field] = (orig_varname, component) - return FrozenDict((field, self.open_store_variable(field, *tup)) for field, tup in variables.items()) + return FrozenDict( + (field, self.open_store_variable(field, *tup)) + for field, tup in variables.items() + ) - def open_store_variable(self, field: str, orig_varname: str, component: int) -> xarray.DataArray: - data = indexing.LazilyIndexedArray(PscAdios2Array(field, self, orig_varname, component)) + def open_store_variable( + self, field: str, orig_varname: str, component: int + ) -> xarray.DataArray: + data = indexing.LazilyIndexedArray( + PscAdios2Array(field, self, orig_varname, component) + ) dims = ["x", "y", "z"] coords = {"x": self.psc.x, "y": self.psc.y, "z": self.psc.z} return xarray.DataArray(data, dims=dims, coords=coords) @override def get_attrs(self) -> Frozen[str, Any]: - return FrozenDict((name, self.ds.get_attribute(name)) for name in self.ds.attribute_names) + return FrozenDict( + (name, self.ds.get_attribute(name)) for name in self.ds.attribute_names + ) @override def get_dimensions(self) -> Never: @@ -171,7 +194,8 @@ def open_dataset( drop_variables: str | Iterable[str] | None = None, length: ArrayLike | None = None, corner: ArrayLike | None = None, - species_names: Iterable[str] | None = None, # e.g. ['e', 'i']; FIXME should be readable from file + species_names: Iterable[str] + | None = None, # e.g. ['e', 'i']; FIXME should be readable from file ) -> xarray.Dataset: if not isinstance(filename_or_obj, (str, os.PathLike)): raise NotImplementedError() @@ -188,12 +212,19 @@ def open_dataset( ) @override - def guess_can_open(self, filename_or_obj: str | os.PathLike[Any] | ReadBuffer[Any] | AbstractDataStore) -> bool: + def guess_can_open( + self, + filename_or_obj: str | os.PathLike[Any] | ReadBuffer[Any] | AbstractDataStore, + ) -> bool: if isinstance(filename_or_obj, (str, os.PathLike)): ext = pathlib.Path(filename_or_obj).suffix return ext in {".bp"} return False @override - def open_datatree(self, filename_or_obj: str | os.PathLike[Any] | ReadBuffer[Any] | AbstractDataStore, **kwargs: Any) -> DataTree: + def open_datatree( + self, + filename_or_obj: str | os.PathLike[Any] | ReadBuffer[Any] | AbstractDataStore, + **kwargs: Any, + ) -> DataTree: raise NotImplementedError() diff --git a/tests/test_xarray_adios2.py b/tests/test_xarray_adios2.py index 77022a7..36cba8e 100644 --- a/tests/test_xarray_adios2.py +++ b/tests/test_xarray_adios2.py @@ -7,7 +7,12 @@ def test_open_dataset(): - ds = xr.open_dataset(pscpy.sample_dir / "pfd.000000400.bp", species_names=[], length=[1, 12.8, 51.2], corner=[0, -6.4, -25.6]) + ds = xr.open_dataset( + pscpy.sample_dir / "pfd.000000400.bp", + species_names=[], + length=[1, 12.8, 51.2], + corner=[0, -6.4, -25.6], + ) assert "jx_ec" in ds assert ds.coords.keys() == set({"x", "y", "z"}) assert ds.jx_ec.shape == (1, 128, 512) From 755aefc3ae5e09c87a540b7ef2c6645aa91325c5 Mon Sep 17 00:00:00 2001 From: Kai Germaschewski Date: Fri, 20 Dec 2024 21:37:53 -0500 Subject: [PATCH 6/7] adios2py: enerate adios2 io name using a counter for uniqueness --- src/pscpy/adios2py/__init__.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/pscpy/adios2py/__init__.py b/src/pscpy/adios2py/__init__.py index 517b644..9c43d87 100644 --- a/src/pscpy/adios2py/__init__.py +++ b/src/pscpy/adios2py/__init__.py @@ -1,5 +1,6 @@ from __future__ import annotations +import itertools import logging import os from collections.abc import Collection @@ -118,8 +119,11 @@ def __repr__(self) -> str: class FileState: """Collects the state of a `File` to reflect the fact that they are coupled.""" + _io_count = itertools.count() + def __init__(self, filename: str | os.PathLike[Any]) -> None: - self.io_name = f"io-{filename}" + self.io_name = f"io-adios2py-{next(self._io_count)}" + logger.debug("io_name = %s", self.io_name) self.io = _ad.declare_io(self.io_name) self.engine = self.io.open(str(filename), adios2.bindings.Mode.Read) From 2a61a43e340ec670d93cd6f4c97e09fabd3c6a21 Mon Sep 17 00:00:00 2001 From: Kai Germaschewski Date: Fri, 20 Dec 2024 21:40:40 -0500 Subject: [PATCH 7/7] adios2py: make _ad a class variable, too --- src/pscpy/adios2py/__init__.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/pscpy/adios2py/__init__.py b/src/pscpy/adios2py/__init__.py index 9c43d87..1096959 100644 --- a/src/pscpy/adios2py/__init__.py +++ b/src/pscpy/adios2py/__init__.py @@ -16,8 +16,6 @@ logger = logging.getLogger(__name__) -_ad = Adios() - class Variable: """Wrapper for an `adios2.Variable` object to facilitate loading and indexing into it.""" @@ -119,14 +117,19 @@ def __repr__(self) -> str: class FileState: """Collects the state of a `File` to reflect the fact that they are coupled.""" + _ad = Adios() _io_count = itertools.count() def __init__(self, filename: str | os.PathLike[Any]) -> None: self.io_name = f"io-adios2py-{next(self._io_count)}" logger.debug("io_name = %s", self.io_name) - self.io = _ad.declare_io(self.io_name) + self.io = self._ad.declare_io(self.io_name) self.engine = self.io.open(str(filename), adios2.bindings.Mode.Read) + def close(self) -> None: + self.engine.close() + self._ad.remove_io(self.io_name) + @staticmethod def is_open(maybe_state: FileState | None) -> TypeGuard[FileState]: return maybe_state is not None @@ -175,8 +178,7 @@ def close(self) -> None: for var in self._open_vars.values(): var.close() - self._state.engine.close() - _ad.remove_io(self._state.io_name) + self._state.close() self._state = None def get_variable(self, variable_name: str) -> Variable: