From 4c5565a1618834ea343d0acf470c2f857a8106e5 Mon Sep 17 00:00:00 2001 From: James McClung Date: Tue, 25 Mar 2025 11:16:13 -0400 Subject: [PATCH 1/3] test_xarray_adios2: use fixtures for ds --- tests/test_xarray_adios2.py | 103 ++++++++++++++++-------------------- 1 file changed, 47 insertions(+), 56 deletions(-) diff --git a/tests/test_xarray_adios2.py b/tests/test_xarray_adios2.py index c1ac651..46b47ab 100644 --- a/tests/test_xarray_adios2.py +++ b/tests/test_xarray_adios2.py @@ -1,8 +1,5 @@ from __future__ import annotations -import os -from typing import Any - import adios2py import numpy as np import pytest @@ -70,11 +67,14 @@ def test_filename_4(tmp_path): return filename -def _open_dataset(filename: os.PathLike[Any], *, decode: bool = False) -> xr.Dataset: - ds = xr.open_dataset(filename) - if decode: - ds = _decode_dataset(ds) - return ds +@pytest.fixture +def ds_pfd_raw() -> xr.Dataset: + return xr.open_dataset(pscpy.sample_dir / "pfd.000000400.bp") + + +@pytest.fixture +def ds_pfd_moments_raw() -> xr.Dataset: + return xr.open_dataset(pscpy.sample_dir / "pfd_moments.000000400.bp") def _decode_dataset(ds: xr.Dataset) -> xr.Dataset: @@ -86,77 +86,70 @@ def _decode_dataset(ds: xr.Dataset) -> xr.Dataset: ) -def test_open_dataset(): - ds_decoded = _open_dataset(pscpy.sample_dir / "pfd.000000400.bp", decode=True) - assert "jx_ec" in ds_decoded - assert ds_decoded.coords.keys() == set({"x", "y", "z"}) - assert ds_decoded.jx_ec.sizes == dict(x=1, y=128, z=512) # noqa: C408 - assert np.allclose( - ds_decoded.jx_ec.z.data, np.linspace(-25.6, 25.6, 512, endpoint=False).data - ) +@pytest.fixture +def ds_pfd_decoded(ds_pfd_raw) -> xr.Dataset: + return _decode_dataset(ds_pfd_raw) -def test_component(): - ds_raw = _open_dataset(pscpy.sample_dir / "pfd.000000400.bp") - ds_decoded = _decode_dataset(ds_raw) - assert np.all(ds_raw.jeh.isel(dim_1_9=0).data == ds_decoded.jx_ec.data) +@pytest.fixture +def ds_pfd_moments_decoded(ds_pfd_moments_raw) -> xr.Dataset: + return _decode_dataset(ds_pfd_moments_raw) -def test_selection(): - ds_raw = _open_dataset(pscpy.sample_dir / "pfd.000000400.bp") - ds_decoded = _decode_dataset(ds_raw) - assert np.all( - ds_raw.jeh.isel(dim_1_9=0, dim_3_128=slice(0, 10), dim_2_512=slice(0, 40)).data - == ds_decoded.jx_ec.isel(y=slice(0, 10), z=slice(0, 40)).data - ) +def test_open_dataset(ds_pfd_decoded): + assert "jx_ec" in ds_pfd_decoded + assert ds_pfd_decoded.coords.keys() == set({"x", "y", "z"}) + assert ds_pfd_decoded.jx_ec.sizes == dict(x=1, y=128, z=512) # noqa: C408 + assert np.allclose(ds_pfd_decoded.jx_ec.z.data, np.linspace(-25.6, 25.6, 512, endpoint=False).data) + + +def test_component(ds_pfd_raw, ds_pfd_decoded): + assert np.all(ds_pfd_raw.jeh.isel(dim_1_9=0).data == ds_pfd_decoded.jx_ec.data) + + +def test_selection(ds_pfd_raw, ds_pfd_decoded): + data_raw = ds_pfd_raw.jeh.isel(dim_1_9=0, dim_3_128=slice(0, 10), dim_2_512=slice(0, 40)).data + data_decoded = ds_pfd_decoded.jx_ec.isel(y=slice(0, 10), z=slice(0, 40)).data + assert np.all(data_raw == data_decoded) -def test_nbytes(): - ds_raw = _open_dataset(pscpy.sample_dir / "pfd.000000400.bp") - ds_decoded = _decode_dataset(ds_raw) - assert ds_decoded.nbytes == ds_decoded.nbytes +def test_nbytes(ds_pfd_decoded): + # FIXME this passes, go figure + assert ds_pfd_decoded.nbytes == ds_pfd_decoded.nbytes -def test_missing_length(): - ds_raw = _open_dataset(pscpy.sample_dir / "pfd.000000400.bp") +def test_missing_length(ds_pfd_raw): with pytest.raises(ValueError, match=r".*length.*"): pscpy.decode_psc( - ds_raw, + ds_pfd_raw, species_names=["e", "i"], corner=[0, -6.4, -25.6], ) -def test_missing_corner(): - ds_raw = _open_dataset(pscpy.sample_dir / "pfd.000000400.bp") +def test_missing_corner(ds_pfd_raw): with pytest.raises(ValueError, match=r".*corner.*"): pscpy.decode_psc( - ds_raw, + ds_pfd_raw, species_names=["e", "i"], length=[1, 12.8, 51.2], ) -def test_computed(): - ds_raw = _open_dataset(pscpy.sample_dir / "pfd.000000400.bp") - ds_decoded = _decode_dataset(ds_raw) - ds_raw = ds_raw.assign(jx=ds_raw.jeh.isel(dim_1_9=0)) - assert np.all(ds_raw.jx.data == ds_decoded.jx_ec.data) +def test_computed(ds_pfd_raw, ds_pfd_decoded): + ds_pfd_raw = ds_pfd_raw.assign(jx=ds_pfd_raw.jeh.isel(dim_1_9=0)) + assert np.all(ds_pfd_raw.jx.data == ds_pfd_decoded.jx_ec.data) -def test_computed_via_lambda(): - ds_raw = _open_dataset(pscpy.sample_dir / "pfd.000000400.bp") - ds_decoded = _decode_dataset(ds_raw) - ds_raw = ds_raw.assign(jx=lambda ds: ds.jeh.isel(dim_1_9=0)) - assert np.all(ds_raw.jx.data == ds_decoded.jx_ec.data) +def test_computed_via_lambda(ds_pfd_raw, ds_pfd_decoded): + ds_pfd_raw = ds_pfd_raw.assign(jx=lambda ds: ds.jeh.isel(dim_1_9=0)) + assert np.all(ds_pfd_raw.jx.data == ds_pfd_decoded.jx_ec.data) -def test_pfd_moments(): - ds_raw = _open_dataset(pscpy.sample_dir / "pfd_moments.000000400.bp") - ds_decoded = _decode_dataset(ds_raw) - assert "all_1st" in ds_raw - assert "rho_i" in ds_decoded - assert np.all(ds_decoded.rho_i.data == ds_raw.all_1st.isel(dim_1_26=13).data) +def test_pfd_moments(ds_pfd_moments_raw, ds_pfd_moments_decoded): + assert "all_1st" in ds_pfd_moments_raw + assert "rho_i" in ds_pfd_moments_decoded + assert np.all(ds_pfd_moments_decoded.rho_i.data == ds_pfd_moments_raw.all_1st.isel(dim_1_26=13).data) def test_open_dataset_steps(test_filename): @@ -204,9 +197,7 @@ def test_open_dataset_3_step(test_filename_3, mode): with adios2py.File(test_filename_3, mode=mode) as file: for n, step in enumerate(file.steps): ds = xr.open_dataset(Adios2Store(step)) - assert ds.time == np.datetime64("2020-01-01T00:01:40") + np.timedelta64( - 10 * n, "s" - ) + assert ds.time == np.datetime64("2020-01-01T00:01:40") + np.timedelta64(10 * n, "s") def test_open_dataset_4(test_filename_4): From 9383a253beca183cad6fb7adec1b26472ee4238c Mon Sep 17 00:00:00 2001 From: James McClung Date: Tue, 25 Mar 2025 11:21:30 -0400 Subject: [PATCH 2/3] test_xarray_adios2: fix nbytes test --- tests/test_xarray_adios2.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_xarray_adios2.py b/tests/test_xarray_adios2.py index 46b47ab..a537df3 100644 --- a/tests/test_xarray_adios2.py +++ b/tests/test_xarray_adios2.py @@ -113,9 +113,12 @@ def test_selection(ds_pfd_raw, ds_pfd_decoded): assert np.all(data_raw == data_decoded) -def test_nbytes(ds_pfd_decoded): - # FIXME this passes, go figure - assert ds_pfd_decoded.nbytes == ds_pfd_decoded.nbytes +def _get_nbytes(ds: xr.Dataset) -> int: + return sum(arr.nbytes for arr in ds.data_vars.values()) + + +def test_nbytes(ds_pfd_raw, ds_pfd_decoded): + assert _get_nbytes(ds_pfd_raw) == _get_nbytes(ds_pfd_decoded) def test_missing_length(ds_pfd_raw): From fa182caf490ff26542a4353d1a4af53f2bb8994b Mon Sep 17 00:00:00 2001 From: James McClung Date: Fri, 6 Jun 2025 12:05:15 -0400 Subject: [PATCH 3/3] test_xarray_adios2: format --- tests/test_xarray_adios2.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/test_xarray_adios2.py b/tests/test_xarray_adios2.py index a537df3..aafbe89 100644 --- a/tests/test_xarray_adios2.py +++ b/tests/test_xarray_adios2.py @@ -100,7 +100,9 @@ def test_open_dataset(ds_pfd_decoded): assert "jx_ec" in ds_pfd_decoded assert ds_pfd_decoded.coords.keys() == set({"x", "y", "z"}) assert ds_pfd_decoded.jx_ec.sizes == dict(x=1, y=128, z=512) # noqa: C408 - assert np.allclose(ds_pfd_decoded.jx_ec.z.data, np.linspace(-25.6, 25.6, 512, endpoint=False).data) + assert np.allclose( + ds_pfd_decoded.jx_ec.z.data, np.linspace(-25.6, 25.6, 512, endpoint=False).data + ) def test_component(ds_pfd_raw, ds_pfd_decoded): @@ -108,7 +110,9 @@ def test_component(ds_pfd_raw, ds_pfd_decoded): def test_selection(ds_pfd_raw, ds_pfd_decoded): - data_raw = ds_pfd_raw.jeh.isel(dim_1_9=0, dim_3_128=slice(0, 10), dim_2_512=slice(0, 40)).data + data_raw = ds_pfd_raw.jeh.isel( + dim_1_9=0, dim_3_128=slice(0, 10), dim_2_512=slice(0, 40) + ).data data_decoded = ds_pfd_decoded.jx_ec.isel(y=slice(0, 10), z=slice(0, 40)).data assert np.all(data_raw == data_decoded) @@ -152,7 +156,10 @@ def test_computed_via_lambda(ds_pfd_raw, ds_pfd_decoded): def test_pfd_moments(ds_pfd_moments_raw, ds_pfd_moments_decoded): assert "all_1st" in ds_pfd_moments_raw assert "rho_i" in ds_pfd_moments_decoded - assert np.all(ds_pfd_moments_decoded.rho_i.data == ds_pfd_moments_raw.all_1st.isel(dim_1_26=13).data) + assert np.all( + ds_pfd_moments_decoded.rho_i.data + == ds_pfd_moments_raw.all_1st.isel(dim_1_26=13).data + ) def test_open_dataset_steps(test_filename): @@ -200,7 +207,9 @@ def test_open_dataset_3_step(test_filename_3, mode): with adios2py.File(test_filename_3, mode=mode) as file: for n, step in enumerate(file.steps): ds = xr.open_dataset(Adios2Store(step)) - assert ds.time == np.datetime64("2020-01-01T00:01:40") + np.timedelta64(10 * n, "s") + assert ds.time == np.datetime64("2020-01-01T00:01:40") + np.timedelta64( + 10 * n, "s" + ) def test_open_dataset_4(test_filename_4):