def test_keep_coords_deprecation(): ds, coords, metrics = datasets_grid_metric("B") ds = ds.assign_coords(yt_bis=ds["yt"], xt_bis=ds["xt"]) grid = Grid(ds, coords=coords, metrics=metrics) for axis_name in grid.axes.keys(): with pytest.warns(DeprecationWarning): grid.diff(ds.tracer, axis_name, keep_coords=False)
def test_wrong_input_type_vector(self): ds, _, _ = datasets_grid_metric("C") grid = Grid(ds) msg = "Dictionary inputs must have a DataArray as value. Got .*?" with pytest.raises( TypeError, match=msg, ): grid.diff({"X": "not_a_dataarray"}, "X")
def test_wrong_axis_vector_input_axis(self): ds, _, _ = datasets_grid_metric("C") grid = Grid(ds) msg = "Vector component with unknown axis provided. Grid has axes .*?" with pytest.raises( ValueError, match=msg, ): grid.diff({"wrong": xr.DataArray()}, "X")
def test_wrong_input_type_scalar(self): ds, _, _ = datasets_grid_metric("C") grid = Grid(ds) msg = "All data arguments must be either a DataArray or Dictionary .*?" with pytest.raises( TypeError, match=msg, ): grid.diff("not_a_dataarray", "X")
def test_multiple_keys_vector_input(self): ds, _, _ = datasets_grid_metric("C") grid = Grid(ds) msg = "Vector components provided as dictionaries should contain exactly one key/value pair. .*?" with pytest.raises( ValueError, match=msg, ): grid.diff({"X": xr.DataArray(), "Y": xr.DataArray()}, "X")
def test_vector_missing_other_component(self, ds, ds_face_connections_x_to_y): grid = Grid(ds, face_connections=ds_face_connections_x_to_y) msg = "Padding vector components requires `other_component` input" with pytest.raises(ValueError, match=msg): grid.diff( {"X": ds.u}, "X", other_component=None, )
def test_diff_interp_cubed_sphere(cs, cubed_sphere_connections): grid = Grid(cs, face_connections=cubed_sphere_connections) face, _ = xr.broadcast(cs.face, cs.data_c) face_diff_x = grid.diff(face, "X") np.testing.assert_allclose(face_diff_x[:, 0, 0], [-3, 1, 1, 1, 1, 2]) np.testing.assert_allclose(face_diff_x[:, -1, 0], [-3, 1, 1, 1, 1, 2]) face_diff_y = grid.diff(face, "Y") np.testing.assert_allclose(face_diff_y[:, 0, 0], [-4, -3, -2, -1, 2, 5]) np.testing.assert_allclose(face_diff_y[:, 0, -1], [-4, -3, -2, -1, 2, 5])
def test_derivative_uniform_grid(self): # this is a uniform grid # a non-uniform grid would provide a more rigorous test dx = 10.0 dy = 10.0 arr = [ [1.0, 2.0, 4.0, 3.0], [4.0, 7.0, 1.0, 2.0], [3.0, 1.0, 0.0, 9.0], [8.0, 5.0, 2.0, 1.0], ] ds = xr.Dataset( {"foo": (("XC", "YC"), arr)}, coords={ "XC": (("XC", ), [0.5, 1.5, 2.5, 3.5]), "XG": (("XG", ), [0, 1.0, 2.0, 3.0]), "dXC": (("XC", ), [dx, dx, dx, dx]), "dXG": (("XG", ), [dx, dx, dx, dx]), "YC": (("YC", ), [0.5, 1.5, 2.5, 3.5]), "YG": (("YG", ), [0, 1.0, 2.0, 3.0]), "dYC": (("YC", ), [dy, dy, dy, dy]), "dYG": (("YG", ), [dy, dy, dy, dy]), }, ) grid = Grid( ds, coords={ "X": { "center": "XC", "left": "XG" }, "Y": { "center": "YC", "left": "YG" }, }, metrics={ ("X", ): ["dXC", "dXG"], ("Y", ): ["dYC", "dYG"] }, periodic=True, ) # Test x direction dfoo_dx = grid.derivative(ds.foo, "X") expected = grid.diff(ds.foo, "X") / dx assert dfoo_dx.equals(expected) # Test y direction dfoo_dy = grid.derivative(ds.foo, "Y") expected = grid.diff(ds.foo, "Y") / dy assert dfoo_dy.equals(expected)
def test_grid_ops(all_datasets): """ Check that we get the same answer using Axis or Grid objects """ ds, periodic, _ = all_datasets grid = Grid(ds, periodic=periodic) for axis_name in grid.axes.keys(): try: ax_periodic = axis_name in periodic except TypeError: ax_periodic = periodic axis = Axis(ds, axis_name, periodic=ax_periodic) bcs = [None] if ax_periodic else ["fill", "extend"] for varname in ["data_c", "data_g"]: for boundary in bcs: da_interp = grid.interp(ds[varname], axis_name, boundary=boundary) da_interp_ax = axis.interp(ds[varname], boundary=boundary) assert da_interp.equals(da_interp_ax) da_diff = grid.diff(ds[varname], axis_name, boundary=boundary) da_diff_ax = axis.diff(ds[varname], boundary=boundary) assert da_diff.equals(da_diff_ax) if boundary is not None: da_cumsum = grid.cumsum(ds[varname], axis_name, boundary=boundary) da_cumsum_ax = axis.cumsum(ds[varname], boundary=boundary) assert da_cumsum.equals(da_cumsum_ax)
def test_grid_ops(all_datasets): """ Check that we get the same answer using Axis or Grid objects """ ds, periodic, expected = all_datasets grid = Grid(ds, periodic=periodic) for axis_name in grid.axes.keys(): try: ax_periodic = axis_name in periodic except TypeError: ax_periodic = periodic axis = Axis(ds, axis_name, periodic=ax_periodic) bcs = [None] if ax_periodic else ['fill', 'extend'] for varname in ['data_c', 'data_g']: for boundary in bcs: da_interp = grid.interp(ds[varname], axis_name, boundary=boundary) da_interp_ax = axis.interp(ds[varname], boundary=boundary) assert da_interp.equals(da_interp_ax) da_diff = grid.diff(ds[varname], axis_name, boundary=boundary) da_diff_ax = axis.diff(ds[varname], boundary=boundary) assert da_diff.equals(da_diff_ax) if boundary is not None: da_cumsum = grid.cumsum(ds[varname], axis_name, boundary=boundary) da_cumsum_ax = axis.cumsum(ds[varname], boundary=boundary) assert da_cumsum.equals(da_cumsum_ax)
def test_diff_interp_connected_grid_x_to_y(ds, ds_face_connections_x_to_y): # one face connection, rotated grid = Grid(ds, face_connections=ds_face_connections_x_to_y) diff_x = grid.diff(ds.data_c, 'X', boundary='fill') interp_x = grid.interp(ds.data_c, 'X', boundary='fill') diff_y = grid.diff(ds.data_c, 'Y', boundary='fill') interp_y = grid.interp(ds.data_c, 'Y', boundary='fill') # make sure the face connection got applied correctly # non-same axis connections require rotation # ravel everything to avoid dealing with broadcasting np.testing.assert_allclose( diff_y.data[1, 0, :].ravel(), ds.data_c.data[1, 0, :].ravel() - ds.data_c.data[0, ::-1, -1].ravel()) np.testing.assert_allclose( interp_y.data[1, 0, :].ravel(), 0.5 * (ds.data_c.data[1, 0, :].ravel() + ds.data_c.data[0, ::-1, -1].ravel()))
def test_grid_no_coords(periodic_1d): """Ensure that you can use xgcm with Xarray datasets that don't have dimension coordinates.""" ds, periodic, expected = periodic_1d ds_nocoords = ds.drop_vars(list(ds.dims.keys())) coords = expected["axes"] grid = Grid(ds_nocoords, periodic=periodic, coords=coords) diff = grid.diff(ds["data_c"], "X") assert len(diff.coords) == 0 interp = grid.interp(ds["data_c"], "X") assert len(interp.coords) == 0
def test_grid_no_coords(periodic_1d): ds, periodic, expected = periodic_1d ds_nocoords = ds.drop_dims(list(ds.dims.keys())) coords = expected["axes"] grid = Grid(ds_nocoords, periodic=periodic, coords=coords) diff = grid.diff(ds["data_c"], "X") assert len(diff.coords) == 0 interp = grid.interp(ds["data_c"], "X") assert len(interp.coords) == 0
def test_grid_no_coords(periodic_1d): ds, periodic, expected = periodic_1d grid_expected = Grid(ds, periodic=periodic) ds_nocoords = ds.drop(list(ds.dims.keys())) coords = expected['axes'] grid = Grid(ds_nocoords, periodic=periodic, coords=coords) diff = grid.diff(ds['data_c'], 'X') assert len(diff.coords) == 0 interp = grid.interp(ds['data_c'], 'X') assert len(interp.coords) == 0
def test_diff_interp_connected_grid_x_to_x(ds, ds_face_connections_x_to_x): # simplest scenario with one face connection grid = Grid(ds, face_connections=ds_face_connections_x_to_x) diff_x = grid.diff(ds.data_c, "X", boundary="fill") interp_x = grid.interp(ds.data_c, "X", boundary="fill") # make sure the face connection got applied correctly np.testing.assert_allclose(diff_x[1, :, 0], ds.data_c[1, :, 0] - ds.data_c[0, :, -1]) np.testing.assert_allclose( interp_x[1, :, 0], 0.5 * (ds.data_c[1, :, 0] + ds.data_c[0, :, -1])) # make sure the left boundary got applied correctly np.testing.assert_allclose(diff_x[0, :, 0], ds.data_c[0, :, 0] - 0.0) np.testing.assert_allclose(interp_x[0, :, 0], 0.5 * (ds.data_c[0, :, 0] + 0.0))
def test_derivative_uniform_grid(): # this is a uniform grid # a non-uniform grid would provide a more rigorous test dx = 10.0 ds = xr.Dataset( {"foo": (("XC",), [1.0, 2.0, 4.0, 3.0])}, coords={ "XC": (("XC",), [0.5, 1.5, 2.5, 3.5]), "XG": (("XG",), [0, 1.0, 2.0, 3.0]), "dXC": (("XC",), [dx, dx, dx, dx]), "dXG": (("XG",), [dx, dx, dx, dx]), }, ) grid = Grid( ds, coords={"X": {"center": "XC", "left": "XG"}}, metrics={("X",): ["dXC", "dXG"]}, periodic=True, ) dfoo_dx = grid.derivative(ds.foo, "X") expected = grid.diff(ds.foo, "X") / dx assert dfoo_dx.equals(expected)