Add utility functions for recentering data

src/pscpy/__init__.py

@@ -9,13 +9,16 @@
 import pathlib
 
 from ._version import version as __version__
+from .postprocessing import auto_recenter, get_recentered
 from .psc import decode_psc
 
 sample_dir = pathlib.Path(__file__).parent / "sample"
 
 
 __all__ = [
     "__version__",
+    "auto_recenter",
     "decode_psc",
+    "get_recentered",
     "sample_dir",
 ]

src/pscpy/postprocessing.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from typing import Literal, TypeAlias
+
+import xarray as xr
+
+BoundaryInterpMethod: TypeAlias = Literal["periodic", "pad", "zero"]
+
+
+def get_recentered(
+    da: xr.DataArray,
+    dim: str,
+    interp_dir: Literal[-1, 1],
+    *,
+    boundary: BoundaryInterpMethod = "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.
+    """
+
+    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)
+
+
+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]
+
+
+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.
+
+    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`.
+    """
+
+    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):
+            continue
+
+        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
+                )
+                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")
+            ):
+                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
+
+        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)

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],

tests/test_postprocessing.py

@@ -0,0 +1,155 @@
+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.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.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.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.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.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.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]])
+
+
+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]])
+
+
+@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)
+    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]])
+
+
+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]])
+
+
+@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]]
+    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])
+
+
+@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])