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_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_padding_fill(self, boundary_width, fill_value): ds, coords, _ = datasets_grid_metric("C") grid = Grid(ds, coords=coords) data = ds.tracer # iterate over all axes expected = data.copy(deep=True) for ax, widths in boundary_width.items(): dim = grid._get_dims_from_axis( data, ax)[0] # ? not entirely sure why this is a list expected = expected.pad({dim: widths}, "constant", constant_values=fill_value) # we intentionally strip all coords from padded arrays expected = _strip_all_coords(expected) result = pad( data, grid, boundary="fill", boundary_width=boundary_width, fill_value=fill_value, other_component=None, ) xr.testing.assert_allclose(expected, result)
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_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_padding_mixed(self, boundary_width): ds, coords, _ = datasets_grid_metric("C") grid = Grid(ds, coords=coords) data = ds.tracer axis_padding_mapping = {"X": "periodic", "Y": "extend"} method_mapping = { "periodic": "wrap", "extend": "edge", } # iterate over all axes expected = data.copy(deep=True) for ax, widths in boundary_width.items(): dim = grid._get_dims_from_axis( data, ax)[0] # ? not entirely sure why this is a list expected = expected.pad({dim: widths}, method_mapping[axis_padding_mapping[ax]]) # we intentionally strip all coords from padded arrays expected = _strip_all_coords(expected) result = pad( data, grid, boundary=axis_padding_mapping, boundary_width=boundary_width, fill_value=None, other_component=None, ) xr.testing.assert_allclose(expected, result)
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_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_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_set_metric_key_errors(metric_axes, add_metric): ds, coords, metrics = datasets_grid_metric("C") grid = Grid(ds, coords=coords, metrics=metrics) if len(metric_axes) == 1: with pytest.raises(KeyError, match="not found in dataset."): grid.set_metrics(metric_axes, add_metric) else: with pytest.raises(KeyError, match="not compatible with grid axes"): grid.set_metrics(metric_axes, add_metric)
def test_set_metric_value_errors(metric_axes, overwrite_metric, add_metric): ds, coords, metrics = datasets_grid_metric("C") if add_metric is not None: ds = ds.assign_coords({overwrite_metric: ds[add_metric] * 10}) grid = Grid(ds, coords=coords, metrics=metrics) with pytest.raises(ValueError, match="setting overwrite=True."): grid.set_metrics(metric_axes, overwrite_metric)
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_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_get_metric_orig(axes, data_var, drop_vars, metric_expected_list): ds, coords, metrics = datasets_grid_metric("C") # drop metrics according to drop_vars input, and remove from metrics input if drop_vars: ds = ds.drop_vars(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=coords, metrics=metrics) 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_cgrid(self, funcname, boundary, periodic): ds, coords, metrics = datasets_grid_metric("C") grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic) if funcname == "cumint": # cumint needs a boundary kwargs = dict(boundary=boundary) else: # integrate and average don't use the boundary input kwargs = dict() func = getattr(grid, funcname) # test tracer position for axis, metric_name, dim in zip( ["X", "Y", "Z", ["X", "Y"], ["X", "Y", "Z"]], ["dx_t", "dy_t", "dz_t", "area_t", "volume_t"], ["xt", "yt", "zt", ["xt", "yt"], ["xt", "yt", "zt"]], ): new = func(ds.tracer, axis, **kwargs) expected = _expected_result( ds.tracer, ds[metric_name], grid, dim, axis, funcname, **kwargs ) assert_allclose(new, expected) # test with tuple input if list is provided if isinstance(axis, list): new = func(ds.tracer, tuple(axis), **kwargs) assert_allclose(new, expected) # test u positon for axis, metric_name, dim in zip( ["X", "Y", ["X", "Y"]], ["dx_e", "dy_e", "area_e"], # need more metrics? ["xu", "yt", ["xu", "yt"]], ): new = func(ds.u, axis, **kwargs) expected = _expected_result( ds.u, ds[metric_name], grid, dim, axis, funcname, **kwargs ) assert_allclose(new, expected) # test v positon for axis, metric_name, dim in zip( ["X", "Y", ["X", "Y"]], ["dx_n", "dy_n", "area_n"], # need more metrics? ["xt", "yu", ["xt", "yu"]], ): new = func(ds.v, axis, **kwargs) expected = _expected_result( ds.v, ds[metric_name], grid, dim, axis, funcname, **kwargs ) assert_allclose(new, expected)
def test_weighted_metric( self, funcname, grid_type, variable, axis, metric_weighted, periodic, boundary ): """tests the correct execution of weighted ops along a single axis""" # metric_weighted allows the interpolation of e.g. a surface flux to be conservative # It multiplies the values with a metric like the area, then performs interpolation # and divides by the same metric (area) for the new grid position ds, coords, metrics = datasets_grid_metric(grid_type) grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic) func = getattr(grid, funcname) metric = grid.get_metric(ds[variable], metric_weighted) expected_raw = func(ds[variable] * metric, axis, boundary=boundary) metric_new = grid.get_metric(expected_raw, metric_weighted) expected = expected_raw / metric_new new = func( ds[variable], axis, metric_weighted=metric_weighted, boundary=boundary ) assert new.equals(expected) @pytest.mark.parametrize( "multi_axis", ["X", ["X"], ("Y"), ["X", "Y"], ("Y", "X")] ) 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_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_03b(): # Condition 3: use provided metrics with matching dimensions to calculate for required metric ds, coords, metrics = datasets_grid_metric("C") grid = Grid(ds, coords=coords) grid.set_metrics(("X", "Y"), "area_t") grid.set_metrics(("Z"), "dz_t") get_metric = grid.get_metric(ds.tracer, ("X", "Y", "Z")) metric_var_1 = ds.area_t metric_var_2 = ds.dz_t expected_metric = (metric_var_1 * metric_var_2).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_bgrid(funcname): ds, coords, metrics = datasets_grid_metric("B") grid = Grid(ds, coords=coords, metrics=metrics) if funcname == "cumint": # cumint needs a boundary... kwargs = dict(boundary="fill") else: kwargs = dict() func = getattr(grid, funcname) # test tracer position for axis, metric_name, dim in zip( ["X", "Y", "Z", ["X", "Y"], ["X", "Y", "Z"]], ["dx_t", "dy_t", "dz_t", "area_t", "volume_t"], ["xt", "yt", "zt", ["xt", "yt"], ["xt", "yt", "zt"]], ): new = func(ds.tracer, axis, **kwargs) expected = _expected_result(ds.tracer, ds[metric_name], grid, dim, axis, funcname) assert new.equals(expected) # test with tuple input if list is provided if isinstance(axis, list): new = func(ds.tracer, tuple(axis), **kwargs) assert new.equals(expected) # test u position for axis, metric_name, dim in zip( ["X", "Y", ["X", "Y"]], ["dx_ne", "dy_ne", "area_ne"], # need more metrics? ["xu", "yu", ["xu", "yu"]], ): new = func(ds.u, axis, **kwargs) expected = _expected_result(ds.u, ds[metric_name], grid, dim, axis, funcname) assert new.equals(expected) # test v position for axis, metric_name, dim in zip( ["X", "Y", ["X", "Y"]], ["dx_ne", "dy_ne", "area_ne"], # need more metrics? ["xu", "yu", ["xu", "yu"]], ): new = func(ds.v, axis, **kwargs) expected = _expected_result(ds.v, ds[metric_name], grid, dim, axis, funcname) assert new.equals(expected)
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_metric_axes_missing_from_array(self, funcname, periodic, boundary): ds, coords, metrics = datasets_grid_metric("C") grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic) if funcname == "cumint": # cumint needs a boundary kwargs = dict(boundary="fill") else: kwargs = dict() func = getattr(grid, funcname) with pytest.raises(ValueError, match="Did not find single matching dimension"): func(ds.tracer.mean("xt"), "X", **kwargs) with pytest.raises(ValueError, match="Did not find single matching dimension"): func(ds.tracer.mean("xt"), ["X", "Y", "Z"], **kwargs)
def test_weighted_metric(self, funcname, grid_type, variable, axis, metric_weighted, periodic, boundary): """tests the correct execution of weighted ops along a single axis""" # metric_weighted allows the interpolation of e.g. a surface flux to be conservative # It multiplies the values with a metric like the area, then performs interpolation # and divides by the same metric (area) for the new grid position ds, coords, metrics = datasets_grid_metric(grid_type) grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic) func = getattr(grid, funcname) metric = grid.get_metric(ds[variable], metric_weighted) expected_raw = func(ds[variable] * metric, axis, boundary=boundary) metric_new = grid.get_metric(expected_raw, metric_weighted) expected = expected_raw / metric_new new = func(ds[variable], axis, metric_weighted=metric_weighted, boundary=boundary) assert new.equals(expected)
def test_missingdim(funcname): ds, coords, metrics = datasets_grid_metric("C") 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.*" with pytest.raises(KeyError, match=match_message): func(ds.tracer.mean("xt"), "X", **kwargs) match_message = ( "Unable to find any combinations of metrics for array dims.*X.*Y.*Z.*") with pytest.raises(KeyError, match=match_message): func(ds.tracer.mean("xt"), ["X", "Y", "Z"])
def test_assign_metric(key, metric_vars): ds, coords, _ = datasets_grid_metric("C") _ = Grid(ds, coords=coords, metrics={key: metric_vars})