def test_boundary_kwarg_same_as_grid_constructor_kwarg(): ds = datasets["2d_left"] grid1 = Grid(ds, periodic=False) grid2 = Grid(ds, periodic=False, boundary={"X": "fill", "Y": "fill"}) actual1 = grid1.interp(ds.data_g, ("X", "Y"), boundary={"X": "fill", "Y": "fill"}) actual2 = grid2.interp(ds.data_g, ("X", "Y")) xr.testing.assert_identical(actual1, actual2)
def test_interp_conservative(grid_type, variable, axis, metric_weighted): ds, coords, metrics = datasets_grid_metric(grid_type) grid = Grid(ds, coords=coords, metrics=metrics) metric = grid.get_metric(ds[variable], metric_weighted) expected_raw = grid.interp(ds[variable] * metric, axis) metric_new = grid.get_metric(expected_raw, metric_weighted) expected = expected_raw / metric_new new = grid.interp(ds[variable], axis, metric_weighted=metric_weighted) assert new.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_vector_connected_grid_x_to_y(ds, ds_face_connections_x_to_y, boundary): # one face connection, rotated grid = Grid( ds, face_connections=ds_face_connections_x_to_y, boundary=boundary, fill_value=1, periodic=False, ) # TODO: Remove the periodic once that is deprecated. # ! Set boundary on grid, so it is applied to all axes. # TODO: modify the non velocity tests too (after release) # modify the values of the dataset, so we know what to expect from the output # TODO: Maybe change this in the dataset definition? u_modifier = xr.DataArray([-2, -1], dims="face") v_modifier = xr.DataArray([1, 1], dims="face") u = ds.u * 0 + u_modifier v = ds.v * 0 + v_modifier # no need to check for diff vs interp. They all go through the same dispatch # v is the interesting variable here because it involves a sign change for this # connection (see https://github.com/xgcm/xgcm/issues/410#issue-1098348557) v_out = grid.interp({"Y": v}, "X", other_component={"X": u}) # the test case is set up in a way that all interpolated values for u should be 1 # if the face connection is done properly np.testing.assert_allclose(v_out.data, 1)
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_get_metric_with_conditions_02a(periodic): # Condition 2, case a: interpolate metric with matching axis to desired dimensions ds, coords, _ = datasets_grid_metric("C") grid = Grid(ds, coords=coords, periodic=periodic, boundary="extend") grid.set_metrics(("X", "Y"), "area_e") get_metric = grid.get_metric(ds.v, ("X", "Y")) expected_metric = grid.interp(ds["area_e"], ("X", "Y")) xr.testing.assert_allclose(get_metric, expected_metric)
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_get_metric_with_conditions_04a(): # Condition 4, case a: 1 metric on the wrong position (must interpolate before multiplying) ds, coords, _ = datasets_grid_metric("C") grid = Grid(ds, coords=coords) grid.set_metrics(("X"), "dx_t") grid.set_metrics(("Y"), "dy_n") get_metric = grid.get_metric(ds.v, ("X", "Y")) interp_metric = grid.interp(ds.dx_t, "Y") expected_metric = (interp_metric * ds.dy_n).reset_coords(drop=True) xr.testing.assert_allclose(get_metric, expected_metric)
def test_get_metric_with_conditions_02b(): # Condition 2, case b: get_metric should select for the metric with matching axes and interpolate from there, # even if other metrics in the desired positions are available ds, coords, _ = datasets_grid_metric("C") grid = Grid(ds, coords=coords) grid.set_metrics(("X", "Y"), "area_e") grid.set_metrics(("X"), "dx_n") grid.set_metrics(("Y"), "dx_n") get_metric = grid.get_metric(ds.v, ("X", "Y")) expected_metric = grid.interp(ds["area_e"], ("X", "Y")) xr.testing.assert_allclose(get_metric, expected_metric)
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_interp_like( metric_axes, metric_name, periodic, boundary, boundary_expected, fill_value ): ds, coords, _ = datasets_grid_metric("C") grid = Grid(ds, coords=coords, periodic=periodic) grid.set_metrics(metric_axes, metric_name) metric_available = grid._metrics.get(frozenset(metric_axes), None) metric_available = metric_available[0] interp_metric = grid.interp_like( metric_available, ds.u, boundary=boundary, fill_value=fill_value ) expected_metric = grid.interp( ds[metric_name], metric_axes, boundary=boundary_expected, fill_value=fill_value ) xr.testing.assert_allclose(interp_metric, expected_metric)