From 9750991a390969d3755dad7a704234794c97c205 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 12:16:31 -0400 Subject: [PATCH 01/22] +postprocessing: +recenter just the interface for now --- src/pscpy/__init__.py | 2 ++ src/pscpy/postprocessing.py | 21 +++++++++++++++++++++ 2 files changed, 23 insertions(+) create mode 100644 src/pscpy/postprocessing.py diff --git a/src/pscpy/__init__.py b/src/pscpy/__init__.py index 684402b..e346f3a 100644 --- a/src/pscpy/__init__.py +++ b/src/pscpy/__init__.py @@ -9,6 +9,7 @@ import pathlib from ._version import version as __version__ +from .postprocessing import recenter from .psc import decode_psc sample_dir = pathlib.Path(__file__).parent / "sample" @@ -17,5 +18,6 @@ __all__ = [ "__version__", "decode_psc", + "recenter", "sample_dir", ] diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py new file mode 100644 index 0000000..a05b50e --- /dev/null +++ b/src/pscpy/postprocessing.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from typing import Literal + +import xarray as xr + + +def recenter( + da: xr.DataArray, + dim: str, + interp_dir: Literal[-1, 1], + *, + boundary: Literal["periodic", "pad", "zero"] = "periodic", +) -> xr.DataArray: + """ + Returns a new array with values along `dim` recentered in the direction `interp_dir`. + + For example, `interp_dir=-1` interpolates node-centered values from cell-centered values, because each node center is at a lesser coordinate than the cell center of the same index. + """ + + return da From de8fca912e9d0e24fd88950044255ac25603470b Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 12:16:50 -0400 Subject: [PATCH 02/22] +test_postprocessing: +recenter tests --- tests/test_postprocessing.py | 48 ++++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 tests/test_postprocessing.py diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py new file mode 100644 index 0000000..cf2fc2d --- /dev/null +++ b/tests/test_postprocessing.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import numpy as np +import pytest +import xarray as xr + +import pscpy + + +@pytest.fixture +def test_dataarray(): + return xr.DataArray([4, 5, 6, 5], coords={"x": [0, 1, 2, 3]}) + + +def test_recenter_periodic_left(test_dataarray): + recentered = pscpy.recenter(test_dataarray, "x", -1) + assert np.array_equal(recentered.coords, test_dataarray.coords) + assert np.array_equal(recentered, [4.5, 4.5, 5.5, 5.5]) + + +def test_recenter_periodic_right(test_dataarray): + recentered = pscpy.recenter(test_dataarray, "x", 1) + assert np.array_equal(recentered.coords, test_dataarray.coords) + assert np.array_equal(recentered, [4.5, 5.5, 5.5, 4.5]) + + +def test_recenter_pad_left(test_dataarray): + recentered = pscpy.recenter(test_dataarray, "x", -1, boundary="pad") + assert np.array_equal(recentered.coords, test_dataarray.coords) + assert np.array_equal(recentered, [4, 4.5, 5.5, 5.5]) + + +def test_recenter_pad_right(test_dataarray): + recentered = pscpy.recenter(test_dataarray, "x", 1, boundary="pad") + assert np.array_equal(recentered.coords, test_dataarray.coords) + assert np.array_equal(recentered, [4.5, 5.5, 5.5, 5]) + + +def test_recenter_zero_left(test_dataarray): + recentered = pscpy.recenter(test_dataarray, "x", -1, boundary="zero") + assert np.array_equal(recentered.coords, test_dataarray.coords) + assert np.array_equal(recentered, [2, 4.5, 5.5, 5.5]) + + +def test_recenter_zero_right(test_dataarray): + recentered = pscpy.recenter(test_dataarray, "x", 1, boundary="zero") + assert np.array_equal(recentered.coords, test_dataarray.coords) + assert np.array_equal(recentered, [4.5, 5.5, 5.5, 2.5]) From 772daf98dac9479343639c090a29d6b4b34df5a8 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 12:17:04 -0400 Subject: [PATCH 03/22] postprocessing: implement recenter --- src/pscpy/postprocessing.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index a05b50e..e673e15 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -18,4 +18,14 @@ def recenter( For example, `interp_dir=-1` interpolates node-centered values from cell-centered values, because each node center is at a lesser coordinate than the cell center of the same index. """ - return da + shifted = da.roll({dim: -interp_dir}, roll_coords=False) + boundary_idx = {-1: 0, 1: -1}[interp_dir] + + if boundary == "periodic": + pass # this case is already handled by the behavior of roll() + elif boundary == "pad": + shifted[{dim: boundary_idx}] = da[{dim: boundary_idx}] + elif boundary == "zero": + shifted[{dim: boundary_idx}] = 0 + + return 0.5 * (da + shifted) From 89b6365fb7c2506435bbdfb22c5f84547d2fda08 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 12:30:55 -0400 Subject: [PATCH 04/22] postprocessing; *: rename to get_recentered --- src/pscpy/__init__.py | 4 ++-- src/pscpy/postprocessing.py | 2 +- tests/test_postprocessing.py | 12 ++++++------ 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/pscpy/__init__.py b/src/pscpy/__init__.py index e346f3a..72c1a42 100644 --- a/src/pscpy/__init__.py +++ b/src/pscpy/__init__.py @@ -9,7 +9,7 @@ import pathlib from ._version import version as __version__ -from .postprocessing import recenter +from .postprocessing import get_recentered from .psc import decode_psc sample_dir = pathlib.Path(__file__).parent / "sample" @@ -18,6 +18,6 @@ __all__ = [ "__version__", "decode_psc", - "recenter", + "get_recentered", "sample_dir", ] diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index e673e15..05552c2 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -5,7 +5,7 @@ import xarray as xr -def recenter( +def get_recentered( da: xr.DataArray, dim: str, interp_dir: Literal[-1, 1], diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index cf2fc2d..2f05df7 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -13,36 +13,36 @@ def test_dataarray(): def test_recenter_periodic_left(test_dataarray): - recentered = pscpy.recenter(test_dataarray, "x", -1) + recentered = pscpy.get_recentered(test_dataarray, "x", -1) assert np.array_equal(recentered.coords, test_dataarray.coords) assert np.array_equal(recentered, [4.5, 4.5, 5.5, 5.5]) def test_recenter_periodic_right(test_dataarray): - recentered = pscpy.recenter(test_dataarray, "x", 1) + recentered = pscpy.get_recentered(test_dataarray, "x", 1) assert np.array_equal(recentered.coords, test_dataarray.coords) assert np.array_equal(recentered, [4.5, 5.5, 5.5, 4.5]) def test_recenter_pad_left(test_dataarray): - recentered = pscpy.recenter(test_dataarray, "x", -1, boundary="pad") + recentered = pscpy.get_recentered(test_dataarray, "x", -1, boundary="pad") assert np.array_equal(recentered.coords, test_dataarray.coords) assert np.array_equal(recentered, [4, 4.5, 5.5, 5.5]) def test_recenter_pad_right(test_dataarray): - recentered = pscpy.recenter(test_dataarray, "x", 1, boundary="pad") + recentered = pscpy.get_recentered(test_dataarray, "x", 1, boundary="pad") assert np.array_equal(recentered.coords, test_dataarray.coords) assert np.array_equal(recentered, [4.5, 5.5, 5.5, 5]) def test_recenter_zero_left(test_dataarray): - recentered = pscpy.recenter(test_dataarray, "x", -1, boundary="zero") + recentered = pscpy.get_recentered(test_dataarray, "x", -1, boundary="zero") assert np.array_equal(recentered.coords, test_dataarray.coords) assert np.array_equal(recentered, [2, 4.5, 5.5, 5.5]) def test_recenter_zero_right(test_dataarray): - recentered = pscpy.recenter(test_dataarray, "x", 1, boundary="zero") + recentered = pscpy.get_recentered(test_dataarray, "x", 1, boundary="zero") assert np.array_equal(recentered.coords, test_dataarray.coords) assert np.array_equal(recentered, [4.5, 5.5, 5.5, 2.5]) From a09951878324be8ed1ed92b1e633f30410844100 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 13:45:18 -0400 Subject: [PATCH 05/22] postprocessing: +BoundaryInterpMethod --- src/pscpy/postprocessing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index 05552c2..b3ff78f 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -4,13 +4,14 @@ import xarray as xr +type BoundaryInterpMethod = Literal["periodic", "pad", "zero"] def get_recentered( da: xr.DataArray, dim: str, interp_dir: Literal[-1, 1], *, - boundary: Literal["periodic", "pad", "zero"] = "periodic", + boundary: BoundaryInterpMethod = "periodic", ) -> xr.DataArray: """ Returns a new array with values along `dim` recentered in the direction `interp_dir`. From b00b25bb373ee9cc5861aa272b09797e825f8c2f Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 13:58:35 -0400 Subject: [PATCH 06/22] postprocessing: +auto_recenter --- src/pscpy/__init__.py | 3 ++- src/pscpy/postprocessing.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/pscpy/__init__.py b/src/pscpy/__init__.py index 72c1a42..87c2ef7 100644 --- a/src/pscpy/__init__.py +++ b/src/pscpy/__init__.py @@ -9,7 +9,7 @@ import pathlib from ._version import version as __version__ -from .postprocessing import get_recentered +from .postprocessing import auto_recenter, get_recentered from .psc import decode_psc sample_dir = pathlib.Path(__file__).parent / "sample" @@ -17,6 +17,7 @@ __all__ = [ "__version__", + "auto_recenter", "decode_psc", "get_recentered", "sample_dir", diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index b3ff78f..16be49c 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -6,6 +6,7 @@ type BoundaryInterpMethod = Literal["periodic", "pad", "zero"] + def get_recentered( da: xr.DataArray, dim: str, @@ -30,3 +31,17 @@ def get_recentered( shifted[{dim: boundary_idx}] = 0 return 0.5 * (da + shifted) + + +def auto_recenter( + ds: xr.Dataset, + to_centering: Literal["nc", "cc"], + **boundaries: BoundaryInterpMethod, +): + """ + Recenters variables with names matching a particular pattern to the given centering. + + In particular, variable name ending in `"{dim}_ec"`, `"{dim}_fc"`, `"_nc"`, or `"_cc"`, where `dim` is a dimension name and a key in `boundaries`, is recentered appropriately. For example, if `to_centering="nc"` (node-centered), a variable ending in "x_ec" (i.e., the x-component of an edge-centered field) will be recentered along x, but not y or z, since it is already node-centered in those dimensions. + + Variables are also renamed appropriately. In the example above, `ex_ec` would be renamed to `ex_nc`. + """ From b2a4c5e303bb25e71049a88cddc434f20ee2cb30 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 14:45:24 -0400 Subject: [PATCH 07/22] test_postprocessing: +test_autorecenter_ec_to_nc --- tests/test_postprocessing.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index 2f05df7..c875615 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -46,3 +46,18 @@ def test_recenter_zero_right(test_dataarray): recentered = pscpy.get_recentered(test_dataarray, "x", 1, boundary="zero") assert np.array_equal(recentered.coords, test_dataarray.coords) assert np.array_equal(recentered, [4.5, 5.5, 5.5, 2.5]) + + +@pytest.fixture +def test_dataset_ec(): + coords = [[0, 1], [0, 1, 2]] + dims = ["x", "y"] + ex_ec = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) + ey_ec = xr.DataArray([[0, 2, 4], [1, 3, 5]], coords, dims) + return xr.Dataset({"ex_ec": ex_ec, "ey_ec": ey_ec}) + + +def test_autorecenter_ec_to_nc(test_dataset_ec): + pscpy.auto_recenter(test_dataset_ec, "nc", x="pad", y="pad") + assert np.array_equal(test_dataset_ec.ex_nc, [[0, 1, 2], [1.5, 2.5, 3.5]]) + assert np.array_equal(test_dataset_ec.ey_nc, [[0, 1, 3], [1, 2, 4]]) From a8d3ab33546d21f60fe36dd9f583d33d9d2d5fd3 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 14:45:38 -0400 Subject: [PATCH 08/22] postprocessing: partial impl of auto_recenter just ec to nc for now --- src/pscpy/postprocessing.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index 16be49c..0666bcf 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -45,3 +45,17 @@ def auto_recenter( Variables are also renamed appropriately. In the example above, `ex_ec` would be renamed to `ex_nc`. """ + + interp_dir = {"cc": 1, "nc": -1}[to_centering] + + for var_name in ds: + if not isinstance(var_name, str): + continue + + for dim, boundary_method in boundaries.items(): + if to_centering == "nc" and var_name.endswith(f"{dim}_ec"): + ds[var_name] = get_recentered(ds[var_name], dim, interp_dir, boundary=boundary_method) + + new_name = var_name[:-3] + "_" + to_centering + ds[new_name] = ds[var_name].rename(new_name) + del ds[var_name] From fa33c6323ffc90d89f0e303e1f2c9338850c3456 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 14:54:35 -0400 Subject: [PATCH 09/22] test_postprocessing: +test_autorecenter_ec_to_cc --- tests/test_postprocessing.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index c875615..dc0ee8b 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -61,3 +61,9 @@ def test_autorecenter_ec_to_nc(test_dataset_ec): pscpy.auto_recenter(test_dataset_ec, "nc", x="pad", y="pad") assert np.array_equal(test_dataset_ec.ex_nc, [[0, 1, 2], [1.5, 2.5, 3.5]]) assert np.array_equal(test_dataset_ec.ey_nc, [[0, 1, 3], [1, 2, 4]]) + + +def test_autorecenter_ec_to_cc(test_dataset_ec): + pscpy.auto_recenter(test_dataset_ec, "cc", x="pad", y="pad") + assert np.array_equal(test_dataset_ec.ex_cc, [[0.5, 1.5, 2], [3.5, 4.5, 5]]) + assert np.array_equal(test_dataset_ec.ey_cc, [[0.5, 2.5, 4.5], [1, 3, 5]]) From 35ae01163a6a5ca089e283749525f32cb97d2b88 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 14:54:57 -0400 Subject: [PATCH 10/22] postprocessing: pass test --- src/pscpy/postprocessing.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index 0666bcf..d00746c 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -55,6 +55,10 @@ def auto_recenter( for dim, boundary_method in boundaries.items(): if to_centering == "nc" and var_name.endswith(f"{dim}_ec"): ds[var_name] = get_recentered(ds[var_name], dim, interp_dir, boundary=boundary_method) + elif to_centering == "cc" and var_name.endswith(f"{dim}_ec"): + for other_dim, other_boundary_method in boundaries.items(): + if other_dim != dim: + ds[var_name] = get_recentered(ds[var_name], other_dim, interp_dir, boundary=other_boundary_method) new_name = var_name[:-3] + "_" + to_centering ds[new_name] = ds[var_name].rename(new_name) From 72d6307703d0c74a0c370a4d0be263b351d2a8b5 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 14:57:59 -0400 Subject: [PATCH 11/22] test_postprocessing: add "dont_touch" var --- tests/test_postprocessing.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index dc0ee8b..4d39c2b 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -54,16 +54,19 @@ def test_dataset_ec(): dims = ["x", "y"] ex_ec = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) ey_ec = xr.DataArray([[0, 2, 4], [1, 3, 5]], coords, dims) - return xr.Dataset({"ex_ec": ex_ec, "ey_ec": ey_ec}) + dont_touch = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) + return xr.Dataset({"ex_ec": ex_ec, "ey_ec": ey_ec, "dont_touch": dont_touch}) def test_autorecenter_ec_to_nc(test_dataset_ec): pscpy.auto_recenter(test_dataset_ec, "nc", x="pad", y="pad") assert np.array_equal(test_dataset_ec.ex_nc, [[0, 1, 2], [1.5, 2.5, 3.5]]) assert np.array_equal(test_dataset_ec.ey_nc, [[0, 1, 3], [1, 2, 4]]) + assert np.array_equal(test_dataset_ec.dont_touch, [[0, 1, 2], [3, 4, 5]]) def test_autorecenter_ec_to_cc(test_dataset_ec): pscpy.auto_recenter(test_dataset_ec, "cc", x="pad", y="pad") assert np.array_equal(test_dataset_ec.ex_cc, [[0.5, 1.5, 2], [3.5, 4.5, 5]]) assert np.array_equal(test_dataset_ec.ey_cc, [[0.5, 2.5, 4.5], [1, 3, 5]]) + assert np.array_equal(test_dataset_ec.dont_touch, [[0, 1, 2], [3, 4, 5]]) From 52c6a016e3660fe4c1d507fc329e8f7ca845c856 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 15:00:06 -0400 Subject: [PATCH 12/22] postprocessing: fix spurious renames --- src/pscpy/postprocessing.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index d00746c..fdab269 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -52,14 +52,19 @@ def auto_recenter( if not isinstance(var_name, str): continue + needs_rename = False + for dim, boundary_method in boundaries.items(): if to_centering == "nc" and var_name.endswith(f"{dim}_ec"): ds[var_name] = get_recentered(ds[var_name], dim, interp_dir, boundary=boundary_method) + needs_rename = True elif to_centering == "cc" and var_name.endswith(f"{dim}_ec"): for other_dim, other_boundary_method in boundaries.items(): if other_dim != dim: ds[var_name] = get_recentered(ds[var_name], other_dim, interp_dir, boundary=other_boundary_method) + needs_rename = True - new_name = var_name[:-3] + "_" + to_centering - ds[new_name] = ds[var_name].rename(new_name) - del ds[var_name] + if needs_rename: + new_name = var_name[:-3] + "_" + to_centering + ds[new_name] = ds[var_name].rename(new_name) + del ds[var_name] From c27f92d8f5bf52d730bc47e61beac72321eb131a Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 15:01:43 -0400 Subject: [PATCH 13/22] postprocessing: +_rename_var --- src/pscpy/postprocessing.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index fdab269..5c1e71c 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -33,6 +33,11 @@ def get_recentered( return 0.5 * (da + shifted) +def _rename_var(ds: xr.Dataset, old_name: str, new_name: str): + ds[new_name] = ds[old_name].rename(new_name) + del ds[old_name] + + def auto_recenter( ds: xr.Dataset, to_centering: Literal["nc", "cc"], @@ -65,6 +70,4 @@ def auto_recenter( needs_rename = True if needs_rename: - new_name = var_name[:-3] + "_" + to_centering - ds[new_name] = ds[var_name].rename(new_name) - del ds[var_name] + _rename_var(ds, var_name, var_name[:-3] + "_" + to_centering) From 58955e07fc454df1b323ca3e734f8c09095df425 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 15:15:08 -0400 Subject: [PATCH 14/22] test_postprocessing: +fc tests --- tests/test_postprocessing.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index 4d39c2b..b3ef837 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -70,3 +70,27 @@ def test_autorecenter_ec_to_cc(test_dataset_ec): assert np.array_equal(test_dataset_ec.ex_cc, [[0.5, 1.5, 2], [3.5, 4.5, 5]]) assert np.array_equal(test_dataset_ec.ey_cc, [[0.5, 2.5, 4.5], [1, 3, 5]]) assert np.array_equal(test_dataset_ec.dont_touch, [[0, 1, 2], [3, 4, 5]]) + + +@pytest.fixture +def test_dataset_fc(): + coords = [[0, 1], [0, 1, 2]] + dims = ["x", "y"] + hx_fc = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) + hy_fc = xr.DataArray([[0, 2, 4], [1, 3, 5]], coords, dims) + dont_touch = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) + return xr.Dataset({"hx_fc": hx_fc, "hy_fc": hy_fc, "dont_touch": dont_touch}) + + +def test_autorecenter_fc_to_nc(test_dataset_fc): + pscpy.auto_recenter(test_dataset_fc, "nc", x="pad", y="pad") + assert np.array_equal(test_dataset_fc.hx_nc, [[0, 0.5, 1.5], [3, 3.5, 4.5]]) + assert np.array_equal(test_dataset_fc.hy_nc, [[0, 2, 4], [0.5, 2.5, 4.5]]) + assert np.array_equal(test_dataset_fc.dont_touch, [[0, 1, 2], [3, 4, 5]]) + + +def test_autorecenter_fc_to_cc(test_dataset_fc): + pscpy.auto_recenter(test_dataset_fc, "cc", x="pad", y="pad") + assert np.array_equal(test_dataset_fc.hx_cc, [[1.5, 2.5, 3.5], [3, 4, 5]]) + assert np.array_equal(test_dataset_fc.hy_cc, [[1, 3, 4], [2, 4, 5]]) + assert np.array_equal(test_dataset_fc.dont_touch, [[0, 1, 2], [3, 4, 5]]) From 408d26978e69b5f2eaf751de55d5c3886ece2db4 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 15:15:21 -0400 Subject: [PATCH 15/22] postprocessing: pass fc tests --- src/pscpy/postprocessing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index 5c1e71c..ab80842 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -60,10 +60,10 @@ def auto_recenter( needs_rename = False for dim, boundary_method in boundaries.items(): - if to_centering == "nc" and var_name.endswith(f"{dim}_ec"): + if to_centering == "nc" and var_name.endswith(f"{dim}_ec") or to_centering == "cc" and var_name.endswith(f"{dim}_fc"): ds[var_name] = get_recentered(ds[var_name], dim, interp_dir, boundary=boundary_method) needs_rename = True - elif to_centering == "cc" and var_name.endswith(f"{dim}_ec"): + elif to_centering == "cc" and var_name.endswith(f"{dim}_ec") or to_centering == "nc" and var_name.endswith(f"{dim}_fc"): for other_dim, other_boundary_method in boundaries.items(): if other_dim != dim: ds[var_name] = get_recentered(ds[var_name], other_dim, interp_dir, boundary=other_boundary_method) From 6689864b4917600784cd89c52a3c9785b5d71a0e Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 15:21:17 -0400 Subject: [PATCH 16/22] test_postprocessing: +test_autorecenter_spurious_renames --- tests/test_postprocessing.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index b3ef837..c124c7e 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -54,22 +54,19 @@ def test_dataset_ec(): dims = ["x", "y"] ex_ec = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) ey_ec = xr.DataArray([[0, 2, 4], [1, 3, 5]], coords, dims) - dont_touch = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) - return xr.Dataset({"ex_ec": ex_ec, "ey_ec": ey_ec, "dont_touch": dont_touch}) + return xr.Dataset({"ex_ec": ex_ec, "ey_ec": ey_ec}) def test_autorecenter_ec_to_nc(test_dataset_ec): pscpy.auto_recenter(test_dataset_ec, "nc", x="pad", y="pad") assert np.array_equal(test_dataset_ec.ex_nc, [[0, 1, 2], [1.5, 2.5, 3.5]]) assert np.array_equal(test_dataset_ec.ey_nc, [[0, 1, 3], [1, 2, 4]]) - assert np.array_equal(test_dataset_ec.dont_touch, [[0, 1, 2], [3, 4, 5]]) def test_autorecenter_ec_to_cc(test_dataset_ec): pscpy.auto_recenter(test_dataset_ec, "cc", x="pad", y="pad") assert np.array_equal(test_dataset_ec.ex_cc, [[0.5, 1.5, 2], [3.5, 4.5, 5]]) assert np.array_equal(test_dataset_ec.ey_cc, [[0.5, 2.5, 4.5], [1, 3, 5]]) - assert np.array_equal(test_dataset_ec.dont_touch, [[0, 1, 2], [3, 4, 5]]) @pytest.fixture @@ -78,19 +75,30 @@ def test_dataset_fc(): dims = ["x", "y"] hx_fc = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) hy_fc = xr.DataArray([[0, 2, 4], [1, 3, 5]], coords, dims) - dont_touch = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) - return xr.Dataset({"hx_fc": hx_fc, "hy_fc": hy_fc, "dont_touch": dont_touch}) + return xr.Dataset({"hx_fc": hx_fc, "hy_fc": hy_fc}) def test_autorecenter_fc_to_nc(test_dataset_fc): pscpy.auto_recenter(test_dataset_fc, "nc", x="pad", y="pad") assert np.array_equal(test_dataset_fc.hx_nc, [[0, 0.5, 1.5], [3, 3.5, 4.5]]) assert np.array_equal(test_dataset_fc.hy_nc, [[0, 2, 4], [0.5, 2.5, 4.5]]) - assert np.array_equal(test_dataset_fc.dont_touch, [[0, 1, 2], [3, 4, 5]]) def test_autorecenter_fc_to_cc(test_dataset_fc): pscpy.auto_recenter(test_dataset_fc, "cc", x="pad", y="pad") assert np.array_equal(test_dataset_fc.hx_cc, [[1.5, 2.5, 3.5], [3, 4, 5]]) assert np.array_equal(test_dataset_fc.hy_cc, [[1, 3, 4], [2, 4, 5]]) - assert np.array_equal(test_dataset_fc.dont_touch, [[0, 1, 2], [3, 4, 5]]) + + +@pytest.fixture +def test_dataset_dont_touch(): + coords = [[0, 1, 2]] + dims = ["x"] + ex_ec = xr.DataArray([0, 1, 2], coords, dims) + dont_touch = xr.DataArray([0, 1, 2], coords, dims) + return xr.Dataset({"ex_ec": ex_ec, "dont_touch": dont_touch}) + + +def test_autorecenter_spurious_renames(test_dataset_dont_touch): + pscpy.auto_recenter(test_dataset_dont_touch, "cc", x="pad") + assert np.array_equal(test_dataset_dont_touch.dont_touch, [0, 1, 2]) From 4afed3444234fb776d6ad095c36612e786d2d060 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 15:39:27 -0400 Subject: [PATCH 17/22] test_postprocessing: +nc, cc tests --- tests/test_postprocessing.py | 38 ++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index c124c7e..20e44dd 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -90,6 +90,44 @@ def test_autorecenter_fc_to_cc(test_dataset_fc): assert np.array_equal(test_dataset_fc.hy_cc, [[1, 3, 4], [2, 4, 5]]) +@pytest.fixture +def test_dataset_nc(): + coords = [[0, 1], [0, 1, 2]] + dims = ["x", "y"] + var_nc = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) + return xr.Dataset({"var_nc": var_nc}) + + +def test_autorecenter_nc_to_nc(test_dataset_nc): + # should do nothing + pscpy.auto_recenter(test_dataset_nc, "nc", x="pad", y="pad") + assert np.array_equal(test_dataset_nc.var_nc, [[0, 1, 2], [3, 4, 5]]) + + +def test_autorecenter_nc_to_cc(test_dataset_nc): + pscpy.auto_recenter(test_dataset_nc, "cc", x="pad", y="pad") + assert np.array_equal(test_dataset_nc.var_cc, [[2, 3, 3.5], [3.5, 4.5, 5]]) + + +@pytest.fixture +def test_dataset_cc(): + coords = [[0, 1], [0, 1, 2]] + dims = ["x", "y"] + var_cc = xr.DataArray([[0, 1, 2], [3, 4, 5]], coords, dims) + return xr.Dataset({"var_cc": var_cc}) + + +def test_autorecenter_cc_to_cc(test_dataset_cc): + # should do nothing + pscpy.auto_recenter(test_dataset_cc, "cc", x="pad", y="pad") + assert np.array_equal(test_dataset_cc.var_cc, [[0, 1, 2], [3, 4, 5]]) + + +def test_autorecenter_cc_to_nc(test_dataset_cc): + pscpy.auto_recenter(test_dataset_cc, "nc", x="pad", y="pad") + assert np.array_equal(test_dataset_cc.var_nc, [[0, 0.5, 1.5], [1.5, 2, 3]]) + + @pytest.fixture def test_dataset_dont_touch(): coords = [[0, 1, 2]] From c40c75e618cb308a7da25d637958d50f20206c06 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 15:39:34 -0400 Subject: [PATCH 18/22] postprocessing: pass nc, cc tests --- src/pscpy/postprocessing.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index ab80842..71816d5 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -69,5 +69,10 @@ def auto_recenter( ds[var_name] = get_recentered(ds[var_name], other_dim, interp_dir, boundary=other_boundary_method) needs_rename = True + if to_centering == "cc" and var_name.endswith("_nc") or to_centering == "nc" and var_name.endswith("_cc"): + for dim, boundary_method in boundaries.items(): + ds[var_name] = get_recentered(ds[var_name], dim, interp_dir, boundary=boundary_method) + needs_rename = True + if needs_rename: _rename_var(ds, var_name, var_name[:-3] + "_" + to_centering) From 6a52df072a927ada081733b8c84bd5f09ebebdcb Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 15:50:39 -0400 Subject: [PATCH 19/22] postprocessing: ugly nasty reformatting this is why i can't stand line length limits --- src/pscpy/postprocessing.py | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index 71816d5..62fcf92 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -60,18 +60,33 @@ def auto_recenter( needs_rename = False for dim, boundary_method in boundaries.items(): - if to_centering == "nc" and var_name.endswith(f"{dim}_ec") or to_centering == "cc" and var_name.endswith(f"{dim}_fc"): - ds[var_name] = get_recentered(ds[var_name], dim, interp_dir, boundary=boundary_method) + if (to_centering == "nc" and var_name.endswith(f"{dim}_ec")) or ( + to_centering == "cc" and var_name.endswith(f"{dim}_fc") + ): + ds[var_name] = get_recentered( + ds[var_name], dim, interp_dir, boundary=boundary_method + ) needs_rename = True - elif to_centering == "cc" and var_name.endswith(f"{dim}_ec") or to_centering == "nc" and var_name.endswith(f"{dim}_fc"): + elif (to_centering == "cc" and var_name.endswith(f"{dim}_ec")) or ( + to_centering == "nc" and var_name.endswith(f"{dim}_fc") + ): for other_dim, other_boundary_method in boundaries.items(): if other_dim != dim: - ds[var_name] = get_recentered(ds[var_name], other_dim, interp_dir, boundary=other_boundary_method) + ds[var_name] = get_recentered( + ds[var_name], + other_dim, + interp_dir, + boundary=other_boundary_method, + ) needs_rename = True - if to_centering == "cc" and var_name.endswith("_nc") or to_centering == "nc" and var_name.endswith("_cc"): + if (to_centering == "cc" and var_name.endswith("_nc")) or ( + to_centering == "nc" and var_name.endswith("_cc") + ): for dim, boundary_method in boundaries.items(): - ds[var_name] = get_recentered(ds[var_name], dim, interp_dir, boundary=boundary_method) + ds[var_name] = get_recentered( + ds[var_name], dim, interp_dir, boundary=boundary_method + ) needs_rename = True if needs_rename: From 35f7b37293a8264a3b67c6f7781437aecd7f02a5 Mon Sep 17 00:00:00 2001 From: James McClung Date: Thu, 20 Mar 2025 16:00:56 -0400 Subject: [PATCH 20/22] postprocessing: fix typing issues --- src/pscpy/postprocessing.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index 62fcf92..c9af8fa 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -1,10 +1,10 @@ from __future__ import annotations -from typing import Literal +from typing import Literal, TypeAlias import xarray as xr -type BoundaryInterpMethod = Literal["periodic", "pad", "zero"] +BoundaryInterpMethod: TypeAlias = Literal["periodic", "pad", "zero"] def get_recentered( @@ -33,7 +33,7 @@ def get_recentered( return 0.5 * (da + shifted) -def _rename_var(ds: xr.Dataset, old_name: str, new_name: str): +def _rename_var(ds: xr.Dataset, old_name: str, new_name: str) -> None: ds[new_name] = ds[old_name].rename(new_name) del ds[old_name] @@ -42,7 +42,7 @@ def auto_recenter( ds: xr.Dataset, to_centering: Literal["nc", "cc"], **boundaries: BoundaryInterpMethod, -): +) -> None: """ Recenters variables with names matching a particular pattern to the given centering. @@ -51,7 +51,7 @@ def auto_recenter( Variables are also renamed appropriately. In the example above, `ex_ec` would be renamed to `ex_nc`. """ - interp_dir = {"cc": 1, "nc": -1}[to_centering] + interp_dir: Literal[-1, 1] = {"cc": 1, "nc": -1}[to_centering] # type: ignore[assignment] for var_name in ds: if not isinstance(var_name, str): From dae5a51b95f2cc91d023674fd97828fe83616636 Mon Sep 17 00:00:00 2001 From: Kai Germaschewski Date: Fri, 21 Mar 2025 18:34:08 +0000 Subject: [PATCH 21/22] fix mypy complaint --- src/pscpy/psc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pscpy/psc.py b/src/pscpy/psc.py index 2e3ae15..ecf5b31 100644 --- a/src/pscpy/psc.py +++ b/src/pscpy/psc.py @@ -32,7 +32,7 @@ def __init__( self.z = self._get_coord(2) def _get_coord(self, coord_idx: int) -> NDArray[Any]: - return np.linspace( # type: ignore[no-any-return] + return np.linspace( start=self.corner[coord_idx], stop=self.corner[coord_idx] + self.length[coord_idx], num=self.gdims[coord_idx], From ffcd124149cd3a205540001df814d47eaac39397 Mon Sep 17 00:00:00 2001 From: James McClung Date: Mon, 24 Mar 2025 13:33:01 -0400 Subject: [PATCH 22/22] test_postprocessing: add test for nonstr varname this seemed like the "correct" way of handling the coverage problem --- tests/test_postprocessing.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index 20e44dd..fdf11e1 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -140,3 +140,16 @@ def test_dataset_dont_touch(): def test_autorecenter_spurious_renames(test_dataset_dont_touch): pscpy.auto_recenter(test_dataset_dont_touch, "cc", x="pad") assert np.array_equal(test_dataset_dont_touch.dont_touch, [0, 1, 2]) + + +@pytest.fixture +def test_dataset_nonstr_varname(): + coords = [[0, 1, 2]] + dims = ["x"] + one = xr.DataArray([0, 1, 2], coords, dims) + return xr.Dataset({1: one}) + + +def test_nonstr_varname(test_dataset_nonstr_varname): + pscpy.auto_recenter(test_dataset_nonstr_varname, "nc", x="pad") + assert np.array_equal(test_dataset_nonstr_varname[1], [0, 1, 2])