Source code for xbout.tests.test_grid

from xarray import DataArray, open_dataset, merge
from xarray.testing import assert_equal
import pytest
import numpy as np

from xbout.load import open_boutdataset
from xbout.geometries import register_geometry, REGISTERED_GEOMETRIES


[docs] @pytest.fixture def create_example_grid_file(tmp_path_factory): """ Mocks up a set of BOUT-like netCDF files, and return the temporary test directory containing them. Deletes the temporary directory once that test is done. """ # Create grid dataset arr = np.arange(15).reshape(5, 3) grid = DataArray(data=arr, name="arr", dims=["x", "y"]).to_dataset() grid["dy"] = DataArray(np.ones((5, 3)), dims=["x", "y"]) grid = grid.set_coords(["dy"]) # Create temporary directory save_dir = tmp_path_factory.mktemp("griddata") # Save filepath = save_dir.joinpath("grid.nc") grid.to_netcdf(filepath, engine="netcdf4") return filepath
[docs] class TestOpenGrid:
[docs] def test_open_grid(self, create_example_grid_file): example_grid = create_example_grid_file with pytest.warns(UserWarning): result = open_boutdataset(datapath=example_grid) result = result.drop_vars(["x", "y"]) assert_equal(result, open_dataset(example_grid)) result.close()
[docs] def test_open_grid_extra_dims(self, create_example_grid_file, tmp_path_factory): example_grid = open_dataset(create_example_grid_file) new_var = DataArray( name="new", data=[[1, 2], [8, 9], [16, 17], [27, 28], [37, 38]], dims=["x", "w"], ) dodgy_grid_directory = tmp_path_factory.mktemp("dodgy_grid") dodgy_grid_path = dodgy_grid_directory.joinpath("dodgy_grid.nc") merge([example_grid, new_var]).to_netcdf(dodgy_grid_path, engine="netcdf4") with pytest.warns( UserWarning, match="drop all variables containing " "the dimensions 'w'" ): result = open_boutdataset(datapath=dodgy_grid_path) result = result.drop_vars(["x", "y"]) assert_equal(result, example_grid) result.close()
[docs] def test_open_grid_apply_geometry(self, create_example_grid_file): @register_geometry(name="Schwarzschild") def add_schwarzschild_coords(ds, coordinates=None): ds["event_horizon"] = 4.0 ds["event_horizon"].attrs = ds.attrs.copy() return ds example_grid = create_example_grid_file result = result = open_boutdataset( datapath=example_grid, geometry="Schwarzschild" ) assert_equal(result["event_horizon"], DataArray(4.0)) # clean up del REGISTERED_GEOMETRIES["Schwarzschild"] result.close()
[docs] def test_open_grid_chunks(self, create_example_grid_file): example_grid = create_example_grid_file with pytest.warns(UserWarning): result = open_boutdataset(datapath=example_grid, chunks={"x": 4, "y": 5}) result = result.drop_vars(["x", "y"]) assert_equal(result, open_dataset(example_grid)) result.close()
[docs] def test_open_grid_chunks_not_in_grid(self, create_example_grid_file): example_grid = create_example_grid_file with pytest.warns(UserWarning): result = open_boutdataset( datapath=example_grid, chunks={"anonexistantdimension": 5} ) result = result.drop_vars(["x", "y"]) assert_equal(result, open_dataset(example_grid)) result.close()