diff --git a/src/pscpy/__init__.py b/src/pscpy/__init__.py index 684402b..87c2ef7 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 auto_recenter, get_recentered from .psc import decode_psc sample_dir = pathlib.Path(__file__).parent / "sample" @@ -16,6 +17,8 @@ __all__ = [ "__version__", + "auto_recenter", "decode_psc", + "get_recentered", "sample_dir", ] diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py new file mode 100644 index 0000000..c9af8fa --- /dev/null +++ b/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) 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], diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py new file mode 100644 index 0000000..fdf11e1 --- /dev/null +++ b/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])