def test_rps_category_edges_None_works(o, f_prob, input_distributions):
"""Test that rps expects inputs to have category_edges dim if category_edges is None."""
o = o.rename({"time": "category"})
f_prob = f_prob.rename({"time": "category"}).mean("member")
rps(o,
f_prob,
category_edges=None,
dim=[],
input_distributions=input_distributions)
def test_rps_wilks_example():
"""Test with values from Wilks, D. S. (2006). Statistical methods in the
atmospheric sciences (2nd ed, Vol. 91). Amsterdam ; Boston: Academic Press. p.301
"""
category_edges = np.array([-0.01, 0.01, 0.24, 10])
Obs = xr.DataArray([0.0001])
F1 = xr.DataArray([0] * 2 + [0.1] * 5 + [0.3] * 3, dims="member")
F2 = xr.DataArray([0] * 2 + [0.1] * 3 + [0.3] * 5, dims="member")
np.testing.assert_allclose(rps(Obs, F2, category_edges), 0.89)
np.testing.assert_allclose(rps(Obs, F1, category_edges), 0.73)
def test_rps_fair_category_edges_None(o, f_prob):
"""Test that RPS without category_edges works for fair==True if forecast[member]
set."""
rps(
o.rename({"time": "category"}),
f_prob.mean("member").rename({
"time": "category"
}).assign_coords(member=f_prob.member.size),
category_edges=None,
dim=None,
fair=True,
input_distributions="p",
)
def test_rps_wilks_example():
"""Test with values from Wilks, D. S. (2006). Statistical methods in the
atmospheric sciences (2nd ed, Vol. 91). Amsterdam ; Boston: Academic Press. p.301.
"""
category_edges = np.array([0.0, 0.01, 0.24, 1.0])
# first example
# xhistogram way with np.array category_edges
Obs = xr.DataArray([0.0001]) # .expand_dims('time') # no precip
F1 = xr.DataArray([0] * 2 + [0.1] * 5 + [0.3] * 3,
dims="member") # .expand_dims('time')
F2 = xr.DataArray([0] * 2 + [0.1] * 3 + [0.3] * 5,
dims="member") # .expand_dims('time')
np.testing.assert_allclose(rps_xhist(Obs, F1, category_edges), 0.73)
np.testing.assert_allclose(rps_xhist(Obs, F2, category_edges), 0.89)
# xr way with xr.DataArray category_edges
xr_category_edges = xr.DataArray(category_edges,
dims="category_edge",
coords={"category_edge": category_edges})
assert_allclose(rps(Obs, F1, category_edges),
rps(Obs, F1, xr_category_edges))
assert_allclose(rps(Obs, F2, category_edges),
rps(Obs, F2, xr_category_edges))
# second example
Obs = xr.DataArray([0.3]) # larger than 0.25
np.testing.assert_allclose(rps_xhist(Obs, F1, category_edges), 0.53)
np.testing.assert_allclose(rps_xhist(Obs, F2, category_edges), 0.29)
# xr way with xr.DataArray category_edges
assert_allclose(rps(Obs, F1, category_edges),
rps(Obs, F1, xr_category_edges))
assert_allclose(rps(Obs, F2, category_edges),
rps(Obs, F2, xr_category_edges))
def test_2_category_rps_equals_brier_score(o, f_prob):
"""Test that RPS for two categories equals the Brier Score."""
category_edges = np.array([0.0, 0.5, 1.0])
assert_allclose(
rps(o, f_prob, category_edges=category_edges, dim=None),
brier_score(o > 0.5, (f_prob > 0.5).mean("member"), dim=None),
)
def test_rps_last_edge_included(o, f_prob):
"""Test that last edges is included."""
category_edges_np = np.array([0, 0.2, 0.4, 0.6, 0.8, 1.0])
o = xr.ones_like(o)
f_prob = xr.ones_like(f_prob)
res_actual = rps(o, f_prob, dim="time", category_edges=category_edges_np)
assert (res_actual == 0).all()
def test_rps_reduce_dim(o, f_prob, category_edges, dim, fair_bool):
"""Test that rps reduced dim and works for (chunked) ds and da"""
actual = rps(o,
f_prob,
category_edges=category_edges,
dim=dim,
fair=fair_bool)
assert_only_dim_reduced(dim, actual, o)
def test_2_category_rps_equals_brier_score(o, f_prob, fair_bool):
"""Test that RPS for two categories equals the Brier Score."""
category_edges = np.array([0.0, 0.5, 1.0])
assert_allclose(
rps(o, f_prob, category_edges=category_edges, dim=None,
fair=fair_bool).drop(
["forecasts_category_edge", "observations_category_edge"]),
brier_score(o > 0.5, (f_prob > 0.5), dim=None, fair=fair_bool),
)
def test_rps_wilks_example_pdf():
"""Test xs.rps(category_edges=None, input_distributions='p') with values from
Wilks, D. S. (2006). Statistical methods in the atmospheric sciences (2nd ed,
Vol. 91). Amsterdam ; Boston: Academic Press. p.301.
"""
Obs = xr.DataArray([1.0, 0.0, 0.0], dims="category") # no precip
F1 = xr.DataArray([0.2, 0.5, 0.3], dims="category")
F2 = xr.DataArray([0.2, 0.3, 0.5], dims="category")
np.testing.assert_allclose(
rps(Obs, F1, category_edges=None, input_distributions="p"), 0.73)
np.testing.assert_allclose(
rps(Obs, F2, category_edges=None, input_distributions="p"), 0.89)
# second example
Obs = xr.DataArray([0.0, 0.0, 1.0], dims="category") # larger than 0.25
np.testing.assert_allclose(
rps(Obs, F1, category_edges=None, input_distributions="p"), 0.53)
np.testing.assert_allclose(
rps(Obs, F2, category_edges=None, input_distributions="p"), 0.29)
def test_rps_accessor(o, f_prob, outer_bool):
category_edges = np.linspace(0, 1, 6)
actual = rps(o, f_prob, category_edges=category_edges)
ds = xr.Dataset()
ds["o"] = o
ds["f_prob"] = f_prob
if outer_bool:
ds = ds.drop_vars("f_prob")
expected = ds.xs.rps("o", f_prob, category_edges=category_edges)
else:
expected = ds.xs.rps("o", "f_prob", category_edges=category_edges)
assert_allclose(actual, expected)
def test_rps_category_edges_None(o, f_prob, fair_bool):
"""Test rps with category_edges as None expecting o and f_prob are already CDFs."""
e = [0, 0.2, 0.4, 0.6, 0.8, 1.0]
bin_dim = "category_edge"
edges = xr.DataArray(e, dims=bin_dim, coords={bin_dim: e})
o_c = o < edges # CDF
f_prob_c = f_prob < edges # CDF
actual = rps(o_c,
f_prob_c,
dim="time",
fair=fair_bool,
category_edges=None)
assert set(["lon", "lat"]) == set(actual.dims)
assert "quantile" not in actual.dims
def test_rps_keeps_masked(o, f_prob, fair_bool, category_edges):
"""Test rps keeps NaNs."""
o = o.where(o.lat > 1)
f_prob = f_prob.where(f_prob.lat > 1)
actual = rps(o, f_prob, dim="time", category_edges=category_edges)
assert set(["lon", "lat"]) == set(actual.dims)
assert actual.isel(lat=[0, 1]).isnull().all()
assert actual.isel(lat=slice(2, None)).notnull().all()
# test forecasts_category_edge no repeats
assert ("[-np.inf, 0.2), [0.2, 0.4), [0.4, 0.6), [0.6, 0.8), [0.8, np.inf]"
in actual.coords["forecasts_category_edge"].values)
# one more category internally used than category_edges provided
assert len(category_edges) + 1 == str(
actual.coords["forecasts_category_edge"].values).count("[")
def test_rps_category_edges_xrDataArray(o, f_prob, fair_bool):
"""Test rps with category_edges as xrDataArray for forecast and observations edges."""
category_edges = f_prob.quantile(q=[0.2, 0.4, 0.6, 0.8],
dim=["time", "member"]).rename(
{"quantile": "category_edge"})
actual = rps(
o,
f_prob,
dim="time",
fair=fair_bool,
category_edges=category_edges,
)
assert set(["lon", "lat"]) == set(actual.dims)
assert "category_edge" not in actual.dims
def test_rps_new_identical_old_xhistogram(o, f_prob, fair_bool):
"""Test that new rps algorithm is identical to old algorithm with xhistogram.
Makes a difference whether full range of f_prob is covered or not."""
category_edges_np = np.array([0, 0.2, 0.4, 0.6, 0.8, 1.0])
category_edges_xr = xr.DataArray(
category_edges_np,
dims="category_edge",
coords={"category_edge": category_edges_np},
)
dim = "time"
actual = rps(o, f_prob, dim=dim, category_edges=category_edges_xr)
expected = rps_xhist(o, f_prob, dim=dim, category_edges=category_edges_np)
drop = ["observations_category_edge", "forecasts_category_edge"]
assert_allclose(
actual.rename("histogram_category_edge").drop(drop),
expected.drop(drop))
def test_rps_api_and_inputs(o, f_prob, category_edges, keep_attrs, input_type,
chunk_bool):
"""Test that rps keeps attributes, chunking, input types."""
o, f_prob = modify_inputs(o, f_prob, input_type, chunk_bool)
category_edges = xr.DataArray(category_edges, dims="category_edge")
if "Dataset" in input_type:
category_edges = category_edges.to_dataset(name="var")
if "multidim" in input_type:
category_edges["var2"] = category_edges["var"]
actual = rps(o, f_prob, category_edges, keep_attrs=keep_attrs)
# test that returns chunks
assert_chunk(actual, chunk_bool)
# test that attributes are kept
assert_keep_attrs(actual, o, keep_attrs)
# test that input types equal output types
assign_type_input_output(actual, o)
def test_rps_category_edges_None(o, f_prob, input_distributions):
"""Test rps with category_edges as None expecting o and f_prob are already CDFs."""
e = [0, 0.2, 0.4, 0.6, 0.8, 1.0]
bin_dim = "category"
edges = xr.DataArray(e, dims=bin_dim, coords={bin_dim: e})
o_c = o < edges # CDF
f_c = (f_prob < edges).mean("member") # CDF
actual = rps(
o_c,
f_c,
dim="time",
fair=False,
category_edges=None,
input_distributions=input_distributions,
)
assert set(["lon", "lat"]) == set(actual.dims)
assert "quantile" not in actual.dims
assert "member" not in actual.dims
def test_rps_vs_fair_rps(o, f_prob, category_edges, dim):
"""Test that fair rps is smaller (e.g. better) or equal than rps due to ensemble-
size adjustment."""
frps = rps(o, f_prob, dim=dim, fair=True, category_edges=category_edges)
ufrps = rps(o, f_prob, dim=dim, fair=False, category_edges=category_edges)
assert (frps <= ufrps).all()
def test_rps_dask(o_dask, f_prob_dask, category_edges, fair_bool):
"""Test that rps returns dask array if provided dask array"""
assert (rps(o_dask,
f_prob_dask,
category_edges=category_edges,
fair=fair_bool).chunks is not None)
def test_rps_perfect_values(o, category_edges, fair_bool):
"""Test values for perfect forecast"""
f = xr.concat(10 * [o], dim="member")
res = rps(o, f, category_edges=category_edges, fair=fair_bool)
assert (res == 0).all()
def test_rps_category_edges_None_fails(o, f_prob):
"""Test that rps expects inputs to have category_edges dim if category_edges is None."""
with pytest.raises(ValueError, match="Expected dimension"):
rps(o, f_prob, category_edges=None, dim=[], input_distributions="c")
