Source code for xbout.tests.test_geometries
import numpy as np
from xarray import Dataset, DataArray
from xarray.testing import assert_equal
import pytest
from xbout.geometries import (
apply_geometry,
register_geometry,
UnregisteredGeometryError,
REGISTERED_GEOMETRIES,
)
[docs]
class TestGeometryRegistration:
[docs]
def test_unregistered_geometry(self):
with pytest.raises(
UnregisteredGeometryError, match="tesseract is not a registered geometry"
):
apply_geometry(ds=Dataset(), geometry_name="tesseract")
[docs]
def test_register_new_geometry(self):
@register_geometry(name="Schwarzschild")
def add_schwarzschild_coords(ds, coordinates=None):
ds["event_horizon"] = 4.0
return ds
assert "Schwarzschild" in REGISTERED_GEOMETRIES.keys()
original = Dataset()
original["dy"] = DataArray(np.ones((3, 4)), dims=("x", "y"))
metadata = {
"bout_tdim": "t",
"bout_xdim": "x",
"bout_ydim": "y",
"bout_zdim": "z",
"keep_xboundaries": True,
"keep_yboundaries": True,
}
original.attrs["metadata"] = metadata
original["dy"].attrs["metadata"] = metadata
updated = apply_geometry(ds=original, geometry_name="Schwarzschild")
assert_equal(updated["event_horizon"], DataArray(4.0))
# clean up
del REGISTERED_GEOMETRIES["Schwarzschild"]