import pytest
import numpy as np
import xarray as xr
import xarray.testing as xrt
from xbout.utils import (
_set_attrs_on_all_vars,
_update_metadata_increased_resolution,
_1d_coord_from_spacing,
)
[docs]
class TestUtils:
[docs]
def test__set_attrs_on_all_vars(self):
ds = xr.Dataset()
ds["a"] = xr.DataArray([0.0])
ds["b"] = xr.DataArray([1.0])
_set_attrs_on_all_vars(ds, "i", 42)
_set_attrs_on_all_vars(ds, "metadata", {"x": 3})
assert ds.i == 42
assert ds["a"].i == 42
assert ds["b"].i == 42
ds.metadata["x"] = 5
# no deepcopy done, so metadata on 'a' and 'b' should have changed too
assert ds["a"].metadata["x"] == 5
assert ds["b"].metadata["x"] == 5
[docs]
def test__set_attrs_on_all_vars_copy(self):
ds = xr.Dataset()
ds["a"] = xr.DataArray([0.0])
ds["b"] = xr.DataArray([1.0])
_set_attrs_on_all_vars(ds, "metadata", {"x": 3}, copy=True)
ds.metadata["x"] = 5
# deepcopy done, so metadata on 'a' and 'b' should not have changed
assert ds.metadata["x"] == 5
assert ds["a"].metadata["x"] == 3
assert ds["b"].metadata["x"] == 3
[docs]
@pytest.mark.parametrize(
"origin_at", [None, "lower", "centre", "upper", "expectfail"]
)
def test__1d_coord_from_spacing_scalar(self, origin_at):
ds = xr.Dataset()
ds["foo"] = ("testdim", np.arange(10))
if origin_at == "expectfail":
# fail with no ds for scalar argument
with pytest.raises(ValueError):
coord = _1d_coord_from_spacing(0.2, "testdim")
with pytest.raises(ValueError):
coord = _1d_coord_from_spacing(0.2, "testdim", ds, origin_at=origin_at)
return
coord = _1d_coord_from_spacing(0.2, "testdim", ds, origin_at=origin_at)
if origin_at is None or origin_at == "lower":
xrt.assert_allclose(
coord,
xr.Variable("testdim", np.linspace(0.1, 1.9, 10)),
rtol=2.0e-15,
atol=0.0,
)
elif origin_at == "centre":
xrt.assert_allclose(
coord,
xr.Variable("testdim", np.linspace(-0.9, 0.9, 10)),
rtol=2.0e-15,
atol=0.0,
)
elif origin_at == "upper":
xrt.assert_allclose(
coord,
xr.Variable("testdim", np.linspace(-1.9, -0.1, 10)),
rtol=2.0e-15,
atol=0.0,
)
else:
assert False
[docs]
@pytest.mark.parametrize(
"origin_at", [None, "lower", "centre", "upper", "expectfail"]
)
def test__1d_coord_from_spacing_da(self, origin_at):
da = xr.DataArray(0.2 * np.ones(10), dims="testdim")
da.attrs["metadata"] = {
"bout_xdim": "x",
"bout_ydim": "y",
"bout_zdim": "z",
"keep_xboundaries": True,
"keep_yboundaries": True,
}
if origin_at == "expectfail":
with pytest.raises(ValueError):
coord = _1d_coord_from_spacing(da, "testdim", origin_at=origin_at)
return
coord = _1d_coord_from_spacing(da, "testdim", origin_at=origin_at)
if origin_at is None or origin_at == "lower":
xrt.assert_allclose(
coord,
xr.Variable("testdim", np.linspace(0.1, 1.9, 10)),
rtol=2.0e-15,
atol=0.0,
)
elif origin_at == "centre":
xrt.assert_allclose(
coord,
xr.Variable("testdim", np.linspace(-0.9, 0.9, 10)),
rtol=2.0e-15,
atol=0.0,
)
elif origin_at == "upper":
xrt.assert_allclose(
coord,
xr.Variable("testdim", np.linspace(-1.9, -0.1, 10)),
rtol=2.0e-15,
atol=0.0,
)
else:
assert False