def test_grid_dict_input_boundary_fill(nonperiodic_1d): """Test axis kwarg input functionality using dict input""" ds, _, _ = nonperiodic_1d grid_direct = Grid(ds, periodic=False, boundary="fill", fill_value=5) grid_dict = Grid(ds, periodic=False, boundary={"X": "fill"}, fill_value={"X": 5}) assert grid_direct.axes["X"].fill_value == grid_dict.axes["X"].fill_value assert grid_direct.axes["X"].boundary == grid_dict.axes["X"].boundary
def test_boundary_global_input(funcname, boundary, fill_value): """Test that globally defined boundary values result in the same output as when the parameters are defined on either the grid or axis methods """ ds, coords, metrics = datasets_grid_metric("C") axis = "X" # Test results by globally specifying fill value/boundary on grid object grid_global = Grid( ds, coords=coords, metrics=metrics, periodic=False, boundary=boundary, fill_value=fill_value, ) func_global = getattr(grid_global, funcname) global_result = func_global(ds.tracer, axis) # Test results by manually specifying fill value/boundary on grid method grid_manual = Grid( ds, coords=coords, metrics=metrics, periodic=False, boundary=boundary ) func_manual = getattr(grid_manual, funcname) manual_result = func_manual( ds.tracer, axis, boundary=boundary, fill_value=fill_value ) xr.testing.assert_allclose(global_result, manual_result)
def test_set_metric(): ds, coords, metrics = datasets_grid_metric("C") expected_metrics = {k: [ds[va] for va in v] for k, v in metrics.items()} grid = Grid(ds, coords=coords, metrics=metrics) grid_manual = Grid(ds, coords=coords) for key, value in metrics.items(): grid_manual.set_metrics(key, value) assert len(grid._metrics) > 0 for k, v in expected_metrics.items(): k = frozenset(k) assert k in grid._metrics.keys() assert k in grid_manual._metrics.keys() for metric_expected, metric in zip(v, grid_manual._metrics[k]): xr.testing.assert_allclose(metric_expected.reset_coords(drop=True), metric) for metric_expected, metric in zip(v, grid._metrics[k]): xr.testing.assert_allclose(metric_expected.reset_coords(drop=True), metric)
def test_missingaxis(axis, funcname): # Error should be raised if application axes # include dimension not in datasets ds, coords, metrics = datasets_grid_metric("C") del coords[axis] del_metrics = [k for k in metrics.keys() if axis in k] for dm in del_metrics: del metrics[dm] grid = Grid(ds, coords=coords, metrics=metrics) func = getattr(grid, funcname) if funcname == "cumint": # cumint needs a boundary... kwargs = dict(boundary="fill") else: kwargs = dict() match_message = ( "Unable to find any combinations of metrics for array dims.*%s.*" % axis) with pytest.raises(KeyError, match=match_message): func(ds.tracer, ["X", "Y", "Z"]) with pytest.raises(KeyError, match=match_message): func(ds, axis, **kwargs) if axis == "Y": # test two missing axes at the same time del coords["X"] del_metrics = [k for k in metrics.keys() if "X" in k] for dm in del_metrics: del metrics[dm] grid = Grid(ds, coords=coords, metrics=metrics) func = getattr(grid, funcname) if funcname == "cumint": # cumint needs a boundary... kwargs = dict(boundary="fill") else: kwargs = dict() match_message = ( "Unable to find any combinations of metrics for array dims.*X.*Y.*Z.*" ) with pytest.raises(KeyError, match=match_message): func(ds, ["X", "Y", "Z"], **kwargs) match_message = ( "Unable to find any combinations of metrics for array dims.*X.*Y.*" ) with pytest.raises(KeyError, match=match_message): func(ds, ("X", "Y"), **kwargs)
def test_invalid_fill_value_error(): ds = datasets["1d_left"] with pytest.raises(ValueError): Axis(ds, "X", fill_value="x") with pytest.raises(ValueError): Grid(ds, fill_value="bad") with pytest.raises(ValueError): Grid(ds, fill_value={"X": "bad"})
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_dask_vs_eager(all_datasets, func, boundary): ds, coords, metrics = datasets_grid_metric("C") grid = Grid(ds, coords=coords) grid_method = getattr(grid, func) eager_result = grid_method(ds.tracer, "X", boundary=boundary) ds = ds.chunk({"xt": 1, "yt": 1, "time": 1, "zt": 1}) grid = Grid(ds, coords=coords) grid_method = getattr(grid, func) dask_result = grid_method(ds.tracer, "X", boundary=boundary).compute() xr.testing.assert_allclose(dask_result, eager_result)
def test_invalid_boundary_error(): ds = datasets["1d_left"] with pytest.raises(ValueError): Axis(ds, "X", boundary="bad") with pytest.raises(ValueError): Grid(ds, boundary="bad") with pytest.raises(ValueError): Grid(ds, boundary={"X": "bad"}) with pytest.raises(ValueError): Grid(ds, boundary={"X": 0}) with pytest.raises(ValueError): Grid(ds, boundary=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_create_grid_no_comodo(all_datasets): ds, periodic, expected = all_datasets grid_expected = Grid(ds, periodic=periodic) ds_noattr = ds.copy() for var in ds.variables: ds_noattr[var].attrs.clear() coords = expected["axes"] grid = Grid(ds_noattr, periodic=periodic, coords=coords) for axis_name_expected in grid_expected.axes: axis_expected = grid_expected.axes[axis_name_expected] axis_actual = grid.axes[axis_name_expected] _assert_axes_equal(axis_expected, axis_actual)
def test_weighted_metric_multi_axis(self, funcname, grid_type, variable, multi_axis, metric_weighted, boundary): """tests if the output for multiple axis is the same as when executing the single axis ops in serial""" ds, coords, metrics = datasets_grid_metric(grid_type) grid = Grid(ds, coords=coords, metrics=metrics) func = getattr(grid, funcname) expected = ds[variable] for ax in multi_axis: if isinstance(metric_weighted, dict): metric_weighted_axis = metric_weighted[ax] else: metric_weighted_axis = metric_weighted expected = func( expected, ax, metric_weighted=metric_weighted_axis, boundary=boundary, ) new = func( ds[variable], multi_axis, metric_weighted=metric_weighted, boundary=boundary, ) assert new.equals(expected)
def test_derivative_b_grid(self): # test derivatives with synthetic B grid data ds, coords, metrics = datasets_grid_metric("B") grid = Grid(ds, coords=coords, metrics=metrics) # tracer point var = "tracer" test_axes = ["X", "Y", "Z"] test_dx = ["dx_e", "dy_n", "dz_w"] for ax, dx in zip(test_axes, test_dx): _run_single_derivative_test(grid, ax, ds[var], ds[dx]) # zonal velocity point var = "u" test_dx = ["dx_n", "dy_e", "dz_w_ne"] for ax, dx in zip(test_axes, test_dx): _run_single_derivative_test(grid, ax, ds[var], ds[dx]) # meridional velocity point var = "v" test_dx = ["dx_n", "dy_e", "dz_w_ne"] for ax, dx in zip(test_axes, test_dx): _run_single_derivative_test(grid, ax, ds[var], ds[dx]) # vertical velocity point var = "wt" test_dx = ["dx_e", "dy_n", "dz_t"] for ax, dx in zip(test_axes, test_dx): _run_single_derivative_test(grid, ax, ds[var], ds[dx])
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_get_metric(axes, data_var, drop_vars, metric_expected_list, expected_error): ds_full = datasets() ds = ds_full["C"] metrics = ds_full["metrics"] # drop metrics according to drop_vars input, and remove from metrics input if drop_vars: print(drop_vars) ds = ds.drop(drop_vars) metrics = { k: [a for a in v if a not in drop_vars] for k, v in metrics.items() } grid = Grid(ds, coords=ds_full["coords"], metrics=metrics) if expected_error: with pytest.raises(expected_error): metric = grid.get_metric(ds[data_var], axes) else: metric = grid.get_metric(ds[data_var], axes) expected_metrics = [ ds[me].reset_coords(drop=True) for me in metric_expected_list ] expected = functools.reduce(operator.mul, expected_metrics, 1) assert metric.equals(expected)
def test_keep_coords(funcname, gridtype): ds, coords, metrics = datasets_grid_metric(gridtype) ds = ds.assign_coords(yt_bis=ds["yt"], xt_bis=ds["xt"]) grid = Grid(ds, coords=coords, metrics=metrics) func = getattr(grid, funcname) for axis_name in grid.axes.keys(): result = func(ds.tracer, axis_name) base_coords = list(result.dims) augmented_coords = [ c for c in ds.coords if set(ds[c].dims).issubset(result.dims) and c not in result.dims ] if funcname in ["integrate", "average"]: assert set(result.coords) == set(base_coords + augmented_coords) else: assert set(result.coords) == set(base_coords) # TODO: why is the behavior different for integrate and average? if funcname not in ["integrate", "average"]: result = func(ds.tracer, axis_name, keep_coords=False) assert set(result.coords) == set(base_coords) result = func(ds.tracer, axis_name, keep_coords=True) assert set(result.coords) == set(base_coords + augmented_coords)
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_set_metric_overwrite_true( metric_axes, exist_metric_varname, add_metric_varname, expected_varname ): ds, coords, metrics = datasets_grid_metric("C") ds = ds.assign_coords({add_metric_varname[0]: ds[exist_metric_varname[1]] * 10}) metrics = { k: [m for m in v if m in exist_metric_varname] for k, v in metrics.items() } grid = Grid(ds, coords=coords, metrics=metrics) for av in add_metric_varname: grid.set_metrics(metric_axes, av, overwrite=True) key = frozenset(list(metric_axes)) set_metric = grid._metrics.get(key) expected_metric = [] for ev in expected_varname: metric_var = ds[ev].reset_coords(drop=True) expected_metric.append(metric_var) assert len(set_metric) == len(expected_metric) for i in range(len(set_metric)): assert set_metric[i].equals(expected_metric[i])
def test_multiple_metrics_per_axis(): # copied from test_derivatives.py - should refactor dx = 10.0 ds = xr.Dataset( { "foo": (("XC",), [1.0, 2.0, 4.0, 3.0]), "bar": (("XG",), [10.0, 20.0, 30.0, 40.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, ) assert grid.get_metric(ds.foo, ("X",)).equals(ds.dXC.reset_coords(drop=True)) assert grid.get_metric(ds.bar, ("X",)).equals(ds.dXG.reset_coords(drop=True))
def test_grid_transform_bypass_checks(bypass_checks): """Check that the bypass checks option still delivers the right results for monotonically increasing data""" ( source, grid_kwargs, target, transform_kwargs, expected, _, ) = construct_test_source_data(cases["linear_depth_dens"]) axis = list(grid_kwargs["coords"].keys())[0] grid = Grid(source, periodic=False, **grid_kwargs) target_data = transform_kwargs.pop("target_data", None) transformed = grid.transform( source.data, axis, target, target_data=target_data, bypass_checks=bypass_checks, **transform_kwargs ) xr.testing.assert_allclose(transformed, expected.data)
def test_grid_transform_noname_targetdata(): """Check handling of a `target_data` input without name""" ( source, grid_kwargs, target, transform_kwargs, _, _, ) = construct_test_source_data(cases["linear_depth_dens"]) axis = list(grid_kwargs["coords"].keys())[0] grid = Grid(source, periodic=False, **grid_kwargs) source_da = source.data target_data = transform_kwargs.pop("target_data") target_data.name = None # the name of target_data is only used if `target` is provided as numpy array target = target.data # the high level routines should be able to deal with all cases (no error flag exception like in the mid level) with pytest.warns(UserWarning): transformed = grid.transform(source_da, axis, target, target_data=target_data, **transform_kwargs) "TRANSFORMED_DIMENSION" in transformed.dims
def test_grid_transform_multidim_other_dims_error(request, multidim_cases): # broadcast the 1d column agains some other dims and make sure that the 1d results are still valid source, grid_kwargs, target, transform_kwargs, expected, error_flag = multidim_cases na = 3 source = source * xr.DataArray(np.ones([na]), dims=["a"]) # broadcast the target, but in this case # rename one of the dimensions of the target array, which is not along the # axis of transformation (this could be the case if e.g. temperature is on a different # x grid than velocity) target_data = transform_kwargs.pop("target_data", None) if target_data is not None: target_data = target_data * xr.DataArray(np.ones([na]), dims=["a_other"]) # calculate the multidimensional result axis = list(grid_kwargs["coords"].keys())[0] grid = Grid(source, periodic=False, **grid_kwargs) with pytest.raises(ValueError): _ = grid.transform(source.data, axis, target, target_data=target_data, **transform_kwargs) else: # When target_data is none its taken as a 1D-coordinate, no checking needed pytest.skip()
def test_grid_transform_input_check(): ( source, grid_kwargs, target, transform_kwargs, _, _, ) = construct_test_source_data(cases["linear_depth_dens"]) axis = list(grid_kwargs["coords"].keys())[0] grid = Grid(source, periodic=False, **grid_kwargs) # construct output name transform_kwargs.setdefault("suffix", "") # Make sure that a sensible error is raised if xr.Dataset is provided # for either one of `source`, `target` or `target_data` input arguments. match_msg = r"needs to be a" with pytest.raises(ValueError, match=r"`da` " + match_msg): grid.transform(source, axis, target, **transform_kwargs) with pytest.raises(ValueError, match=match_msg): grid.transform(source.data, axis, target.to_dataset(name="dummy"), **transform_kwargs) transform_kwargs["target_data"] = transform_kwargs[ "target_data"].to_dataset(name="dummy") with pytest.raises(ValueError, match=match_msg): grid.transform(source.data, axis, target, **transform_kwargs)
def test_grid_transform_multidim(request, client, multidim_cases): # broadcast the 1d column agains some other dims and make sure that the 1d results are still valid source, grid_kwargs, target, transform_kwargs, expected, error_flag = multidim_cases na = 8 source = source.expand_dims(a=na).copy(deep=True) # broadcast the target_data manually target_data = transform_kwargs.pop("target_data", None) if target_data is not None: target_data = target_data.expand_dims(a=na).copy(deep=True) if client != "no_client": target_data = target_data.chunk({"a": 1}) if client != "no_client": source = source.chunk({"a": 1}) # calculate the multidimensional result axis = list(grid_kwargs["coords"].keys())[0] grid = Grid(source, periodic=False, **grid_kwargs) # the high level tests should deal with all error cases client = request.getfixturevalue(client) transformed = grid.transform(source.data, axis, target, target_data=target_data, **transform_kwargs).load() _, expected_broadcasted = xr.broadcast(transformed, expected) xr.testing.assert_allclose(transformed, expected_broadcasted.data)
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_missingaxis(self, axis, funcname, periodic, boundary): # Error should be raised if application axes include dimension not in datasets ds, coords, metrics = datasets_grid_metric("C") del coords[axis] del_metrics = [k for k in metrics.keys() if axis in k] for dm in del_metrics: del metrics[dm] grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic) func = getattr(grid, funcname) if funcname == "cumint": # cumint needs a boundary kwargs = dict(boundary=boundary) else: kwargs = dict() with pytest.raises(KeyError, match="Did not find axis"): func(ds.tracer, ["X", "Y", "Z"], **kwargs) if axis == "Y": # test two missing axes at the same time del coords["X"] del_metrics = [k for k in metrics.keys() if "X" in k] for dm in del_metrics: del metrics[dm] grid = Grid(ds, coords=coords, metrics=metrics) func = getattr(grid, funcname) if funcname == "cumint": # cumint needs a boundary kwargs = dict(boundary="fill") else: kwargs = dict() with pytest.raises(KeyError, match="Did not find axis"): func(ds.tracer, ["X", "Y", "Z"], **kwargs) with pytest.raises(KeyError, match="Did not find axis"): func(ds.tracer, ("X", "Y"), **kwargs)
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.apply_as_grid_ufunc(lambda x: x, "not_a_dataarray", "X")
def test_get_metric_with_conditions_01(): # Condition 1: metric with matching axes and dimensions exist ds, coords, metrics = datasets_grid_metric("C") grid = Grid(ds, coords=coords, metrics=metrics) get_metric = grid.get_metric(ds.v, ("X", "Y")) expected_metric = ds["area_n"].reset_coords(drop=True) xr.testing.assert_allclose(get_metric, expected_metric)
def test_transform_error_periodic(multidim_cases): source, grid_kwargs, target, transform_kwargs, expected, error_flag = multidim_cases axis = list(grid_kwargs["coords"].keys())[0] grid = Grid(source, **grid_kwargs) with pytest.raises(ValueError): _ = grid.transform(source.data, axis, target, **transform_kwargs)
def test_input_dim_notfound(): data = np.random.rand(4, 5) coord = np.random.rand(4, 5) ds = xr.DataArray(data, dims=["x", "y"], coords={ "c": (["x", "y"], coord) }).to_dataset(name="data") msg = r"Could not find dimension `other` \(for the `center` position on axis `X`\) in input dataset." with pytest.raises(ValueError, match=msg): Grid(ds, coords={"X": {"center": "other"}})
def test_input_not_dims(): data = np.random.rand(4, 5) coord = np.random.rand(4, 5) ds = xr.DataArray(data, dims=["x", "y"], coords={ "c": (["x", "y"], coord) }).to_dataset(name="data") msg = r"is not a dimension in the input dataset" with pytest.raises(ValueError, match=msg): Grid(ds, coords={"X": {"center": "c"}})