def test_viscosity_filter(vector_grid_type_and_input_ds, filter_args):
"""Test all viscosity-based filters: filters that use a vector Laplacian."""
grid_type, da_u, da_v, grid_vars, geolat_u = vector_grid_type_and_input_ds
filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **filter_args)
filtered_u, filtered_v = filter.apply_to_vector(da_u,
da_v,
dims=["y", "x"])
# check conservation under solid body rotation: u = cos(lat), v=0;
data_u = np.cos(geolat_u / 360 * 2 * np.pi)
data_v = np.zeros_like(data_u)
da_u = xr.DataArray(data_u, dims=["y", "x"])
da_v = xr.DataArray(data_v, dims=["y", "x"])
filtered_u, filtered_v = filter.apply_to_vector(da_u,
da_v,
dims=["y", "x"])
xr.testing.assert_allclose(filtered_u, da_u, atol=1e-12)
xr.testing.assert_allclose(filtered_v, da_v, atol=1e-12)
# check that we get an error if we pass vector Laplacian to .apply, where
# the latter method is for scalar Laplacians only
with pytest.raises(ValueError, match=r"Provided Laplacian *"):
filtered_u = filter.apply(da_u, dims=["y", "x"])
# check that we get an error if we leave out any required grid_vars
for gv in grid_vars:
grid_vars_missing = {k: v for k, v in grid_vars.items() if k != gv}
with pytest.raises(ValueError, match=r"Provided `grid_vars` .*"):
filter = Filter(grid_type=grid_type,
grid_vars=grid_vars_missing,
**filter_args)
def test_iterated_viscosity_filter(vector_grid_type_and_input_ds, filter_args,
n_iterations):
"""Test error in the iterated Gaussian filter for vectors"""
grid_type, da_u, da_v, grid_vars, _ = vector_grid_type_and_input_ds
filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **filter_args)
filtered_u, filtered_v = filter.apply_to_vector(da_u,
da_v,
dims=["y", "x"])
iterated_filter_args = filter_args.copy()
iterated_filter_args["n_iterations"] = n_iterations
iterated_filter = Filter(grid_type=grid_type,
grid_vars=grid_vars,
**iterated_filter_args)
iteratively_filtered_u, iteratively_filtered_v = iterated_filter.apply_to_vector(
da_u, da_v, dims=["y", "x"])
area = 1
for k, v in grid_vars.items():
if "area" in k:
area = v
break
# The following tests whether a relative error bound in L^2 holds.
# See the "Factoring the Gaussian Filter" section of the docs for details.
difference = (filtered_u - iteratively_filtered_u)**2 + (
filtered_v - iteratively_filtered_v)**2
unfiltered = da_u**2 + da_v**2
assert (difference * area).sum() < (
(0.01 * (1 + n_iterations))**2) * (unfiltered * area).sum()
def test_diffusion_filter(grid_type_and_input_ds, filter_args):
"""Test all diffusion-based filters: filters that use a scalar Laplacian."""
grid_type, da, grid_vars = grid_type_and_input_ds
filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **filter_args)
filter.plot_shape()
filtered = filter.apply(da, dims=["y", "x"])
# check conservation
area = 1
for k, v in grid_vars.items():
if "area" in k:
area = v
break
da_sum = (da * area).sum()
filtered_sum = (filtered * area).sum()
xr.testing.assert_allclose(da_sum, filtered_sum)
# check that we get an error if we pass scalar Laplacian to .apply_to vector,
# where the latter method is for vector Laplacians only
with pytest.raises(ValueError, match=r"Provided Laplacian *"):
filtered_u, filtered_v = filter.apply_to_vector(da, da, dims=["y", "x"])
# check variance reduction
assert (filtered**2).sum() < (da**2).sum()
# check that we get an error if we leave out any required grid_vars
for gv in grid_vars:
grid_vars_missing = {k: v for k, v in grid_vars.items() if k != gv}
with pytest.raises(ValueError, match=r"Provided `grid_vars` .*"):
filter = Filter(
grid_type=grid_type, grid_vars=grid_vars_missing, **filter_args
)
bad_filter_args = copy.deepcopy(filter_args)
# check that we get an error when transition_width <= 1
bad_filter_args["transition_width"] = 1
with pytest.raises(ValueError, match=r"Transition width .*"):
filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **bad_filter_args)
bad_filter_args["transition_width"] = np.pi
# check that we get an error if ndim > 2 and n_steps = 0
bad_filter_args["ndim"] = 3
bad_filter_args["n_steps"] = 0
with pytest.raises(ValueError, match=r"When ndim > 2, you .*"):
filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **bad_filter_args)
# check that we get a warning if n_steps < n_steps_default
bad_filter_args["ndim"] = 2
bad_filter_args["n_steps"] = 3
with pytest.warns(UserWarning, match=r"You have set n_steps .*"):
filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **bad_filter_args)
# check that we get an error if we pass dx_min != 1 to a regular scalar Laplacian
if grid_type in area_weighted_regular_grids:
bad_filter_args["filter_scale"] = 3 # restore good value for filter scale
bad_filter_args["dx_min"] = 3
with pytest.raises(ValueError, match=r"Provided Laplacian .*"):
filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **bad_filter_args)
def test_diffusion_filter(grid_type_and_input_ds, filter_args):
"""Test all diffusion-based filters: filters that use a scalar Laplacian."""
grid_type, da, grid_vars = grid_type_and_input_ds
filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **filter_args)
filter.plot_shape()
filtered = filter.apply(da, dims=["y", "x"])
# check conservation
# this would need to be replaced by a proper area-weighted integral
da_sum = da.sum()
filtered_sum = filtered.sum()
xr.testing.assert_allclose(da_sum, filtered_sum)
# check that we get an error if we pass scalar Laplacian to .apply_to vector,
# where the latter method is for vector Laplacians only
with pytest.raises(ValueError, match=r"Provided Laplacian *"):
filtered_u, filtered_v = filter.apply_to_vector(da,
da,
dims=["y", "x"])
# check variance reduction
assert (filtered**2).sum() < (da**2).sum()
# check that we get an error if we leave out any required grid_vars
for gv in grid_vars:
grid_vars_missing = {k: v for k, v in grid_vars.items() if k != gv}
with pytest.raises(ValueError, match=r"Provided `grid_vars` .*"):
filter = Filter(grid_type=grid_type,
grid_vars=grid_vars_missing,
**filter_args)
bad_filter_args = copy.deepcopy(filter_args)
# check that we get an error if ndim > 2 and n_steps = 0
bad_filter_args["ndim"] = 3
bad_filter_args["n_steps"] = 0
with pytest.raises(ValueError, match=r"When ndim > 2, you .*"):
filter = Filter(grid_type=grid_type,
grid_vars=grid_vars,
**bad_filter_args)
# check that we get a warning if n_steps < n_steps_default
bad_filter_args["ndim"] = 2
bad_filter_args["n_steps"] = 3
with pytest.warns(UserWarning, match=r"Warning: You have set n_steps .*"):
filter = Filter(grid_type=grid_type,
grid_vars=grid_vars,
**bad_filter_args)
# check that we get a warning if numerical instability possible
bad_filter_args["n_steps"] = 0
bad_filter_args["filter_scale"] = 1000
with pytest.warns(UserWarning,
match=r"Warning: Filter scale much larger .*"):
filter = Filter(grid_type=grid_type,
grid_vars=grid_vars,
**bad_filter_args)
def _get_results(grid_data_and_input_ds, filter_args):
grid_type, data, grid_vars = grid_data_and_input_ds
filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **filter_args)
if len(data) == 2: # vector grid
(filtered_u, filtered_v) = filter.apply_to_vector(*data,
dims=["y", "x"])
filtered = np.stack([filtered_u.data, filtered_v.data])
else:
filtered = filter.apply(data, dims=["y", "x"])
filtered = filtered.data
filtered = filtered.astype("f4") # use single precision to save space
return filtered
