adapt RPS

CHANGELOG.rst

@@ -22,6 +22,8 @@ Features
   without replacement. (:issue:`215`, :pr:`225`) `Aaron Spring`_
 - Added receiver operating characteristic (ROC) :py:func:`~xskillscore.roc`.
   (:issue:`114`, :issue:`256`, :pr:`236`, :pr:`259`) `Aaron Spring`_
+- Added many options for ``category_edges`` in :py:func:`~xskillscore.rps`, which
+  allows multi-dimensional edges. (:issue:`275`, :pr:`277`) `Aaron Spring`_
 
 Breaking changes
 ~~~~~~~~~~~~~~~~
@@ -40,6 +42,8 @@ Bug Fixes
   (:issue:`255`, :pr:`211`) `Aaron Spring`_
 - Passing weights no longer triggers eager computation.
   (:issue:`218`, :pr:`224`). `Andrew Huang`_
+- :py:func:`~xskillscore.rps` not restricted to ``[0, 1]``.
+  (:issue:`266`, :pr:`277`) `Aaron Spring`_
 
 Internal Changes
 ~~~~~~~~~~~~~~~~

xskillscore/core/probabilistic.py

@@ -488,26 +488,44 @@ def rps(
     """Calculate Ranked Probability Score.
 
      .. math::
-        RPS(p, k) = 1/M \\sum_{m=1}^{M}
+        RPS(p, k) = \\sum_{m=1}^{M}
         [(\\sum_{k=1}^{m} p_k) - (\\sum_{k=1}^{m} o_k)]^{2}
 
+    where ``p`` and ``o`` are forecast and observation probabilities in ``M`` categories.
+
     Parameters
     ----------
     observations : xarray.Dataset or xarray.DataArray
         The observations or set of observations of the event.
+        Further requirements are specified based on ``category_edges``.
         Data should be boolean or logical \
         (True or 1 for event occurance, False or 0 for non-occurance).
     forecasts : xarray.Dataset or xarray.DataArray
         The forecast likelihoods of the event.
-        If ``fair==False``, forecasts should be between 0 and 1 without a dimension
-        ``member_dim`` or should be boolean (True,False) or binary (0, 1) containing a
-        member dimension (probabilities will be internally calculated by
-        ``.mean(member_dim))``. If ``fair==True``, forecasts must be boolean
-        (True,False) or binary (0, 1) containing dimension ``member_dim``.
-    category_edges : array_like
-        Category bin edges used to compute the CDFs. Similar to np.histogram, \
-        all but the last (righthand-most) bin include the left edge and exclude \
-        the right edge. The last bin includes both edges.
+        Further requirements are specified based on ``category_edges``.
+    category_edges : array_like, xr.Dataset, xr.DataArray, None
+
+        - array_like: Category bin edges used to compute the CDFs based on boolean or
+          logical (True or 1 for event occurance, False or 0 for non-occurance)
+          observations. If ``fair==False``, forecasts should be between 0 and 1 without
+          a dimension ``member_dim`` or should be boolean (True,False) or binary (0, 1)
+          containing a member dimension (probabilities will be internally calculated by
+          ``.mean(member_dim))``. If ``fair==True``, forecasts must be boolean
+          (True,False) or binary (0, 1) containing dimension ``member_dim``.
+          Similar to np.histogram, all but the last (righthand-most) bin include the
+          left edge and exclude the right edge. The last bin includes both edges.
+
+        - xr.Dataset/xr.DataArray: edges of the categories in absolute units provided
+          as dimension ``category_dim``, ``threshold`` or ``quantile``. Forecasts and
+          Observations are expected in absolute units.
+
+        - tuple of xr.Dataset/xr.DataArray: same as xr.Dataset/xr.DataArray where the
+          first item is taken as ``category_edges`` for observations and the second item
+          for ``category_edges`` for forecasts.
+
+        - None: expect than observations and forecasts are already CDFs containing
+          ``category_dim`` dimension.
+
     dim : str or list of str, optional
         Dimension over which to compute mean after computing ``rps``.
         Defaults to None implying averaging over all dimensions.
@@ -550,48 +568,122 @@ def rps(
     * C. A. T. Ferro. Fair scores for ensemble forecasts. Q.R.J. Meteorol. Soc., 140:
       1917–1923, 2013. doi: 10.1002/qj.2270.
     * https://www-miklip.dkrz.de/about/problems/
+
     """
     bin_names = ["category"]
-    M = forecasts[member_dim].size
     bin_dim = f"{bin_names[0]}_bin"
-    # histogram(dim=[]) not allowed therefore add fake member dim
-    # to apply over when multi-dim observations
-    if len(observations.dims) == 1:
-        observations = histogram(
-            observations, bins=[category_edges], bin_names=bin_names, dim=None
-        )
-    else:
-        observations = histogram(
-            observations.expand_dims(member_dim),
+    M = forecasts[member_dim].size
+
+    def _bool_to_int(ds):
+        """convert xr.object of dtype bool to int to evade:
+        TypeError: numpy boolean subtract, the `-` operator, is not supported"""
+
+        def _helper_bool_to_int(da):
+            if da.dtype == "bool":
+                da = da.astype("int")
+            return da
+
+        if isinstance(ds, xr.Dataset):
+            ds = ds.map(_helper_bool_to_int)
+        else:
+            ds = _helper_bool_to_int(ds)
+        return ds
+
+    forecasts = _bool_to_int(forecasts)
+
+    def _check_identical_xr_types(a, b):
+        if type(a) != type(b):
+            raise ValueError(
+                f"a and b must be same type, found {type(a)} and {type(b)}"
+            )
+        for d in [a, b]:
+            if not isinstance(d, (xr.Dataset, xr.DataArray)):
+                raise ValueError("inputs must be xr.DataArray or xr.Dataset")
+
+    def _check_bin_dim(ds, bin_dim):
+        """Assert that bin_dim is in ds. Try to guess and rename edges dimension."""
+        for d in ["quantile", "threshold", "edge"]:
+            if d in ds.dims and bin_dim not in ds.dims:
+                ds = ds.rename({d: bin_dim})
+        if bin_dim not in ds.dims:
+            raise ValueError(f"require {bin_dim} dimension, found {ds.dims}")
+        return ds
+
+    _check_identical_xr_types(observations, forecasts)
+
+    # different ways of calculating RPS based on category_edges
+    if isinstance(category_edges, (xr.Dataset, xr.DataArray)) or isinstance(
+        category_edges, tuple
+    ):
+        if isinstance(
+            category_edges, tuple
+        ):  # edges tuple of two: use for obs and forecast edges separately
+            observations_edges, forecast_edges = category_edges
+            _check_identical_xr_types(forecast_edges, forecasts)
+            _check_identical_xr_types(observations_edges, forecasts)
+        else:  # edges only given once, so use for both obs and forecasts
+            _check_identical_xr_types(category_edges, forecasts)
+            observations_edges, forecast_edges = category_edges, category_edges
+
+        # cumulative probs
+        Fc = (forecasts < forecast_edges).mean(member_dim)
+        Oc = observations < observations_edges
+        # todo: mask land
+
+    elif isinstance(category_edges, np.ndarray):
+        # histogram(dim=[]) not allowed therefore add fake member dim
+        # to apply over when multi-dim observations
+        if len(observations.dims) == 1:
+            observations = histogram(
+                observations, bins=[category_edges], bin_names=bin_names, dim=None
+            )
+        else:
+            observations = histogram(
+                observations.expand_dims(member_dim),
+                bins=[category_edges],
+                bin_names=bin_names,
+                dim=[member_dim],
+            )
+
+        forecasts = histogram(
+            forecasts,
             bins=[category_edges],
             bin_names=bin_names,
             dim=[member_dim],
         )
+        # if fair:
+        #    e = forecasts
 
-    forecasts = histogram(
-        forecasts,
-        bins=[category_edges],
-        bin_names=bin_names,
-        dim=[member_dim],
-    )
-    if fair:
-        e = forecasts
+        # normalize f.sum()=1 to make cdf
+        forecasts = forecasts / forecasts.sum(bin_dim)
+        observations = observations / observations.sum(bin_dim)
+
+        Fc = forecasts.cumsum(bin_dim)
+        Oc = observations.cumsum(bin_dim)
 
-    # normalize f.sum()=1
-    forecasts = forecasts / forecasts.sum(bin_dim)
-    observations = observations / observations.sum(bin_dim)
+    elif category_edges is None:  # expect cdfs already as inputs
+        if member_dim in forecasts.dims:
+            forecasts = forecasts.mean(member_dim)
+        Fc = forecasts
+        Oc = observations
+    else:
+        raise ValueError(
+            f"category_edges must be xr.DataArray, xr.Dataset, tuple of xr.objects, None or array-like, found {type(category_edges)}"
+        )
 
-    Fc = forecasts.cumsum(bin_dim)
-    Oc = observations.cumsum(bin_dim)
+    # check and maybe rename edges dim
+    Fc = _check_bin_dim(Fc, bin_dim)
+    Oc = _check_bin_dim(Oc, bin_dim)
 
+    # RPS formulas
     if fair:
-        Ec = e.cumsum(bin_dim)
-        res = (((Ec / M) - Oc) ** 2 - Ec * (M - Ec) / (M ** 2 * (M - 1))).sum(bin_dim)
+        Ec = Fc * M
+        res = ((Ec / M - Oc) ** 2 - Ec * (M - Ec) / (M ** 2 * (M - 1))).sum(bin_dim)
     else:
         res = ((Fc - Oc) ** 2).sum(bin_dim)
+
     if weights is not None:
         res = res.weighted(weights)
-    res = xr.apply_ufunc(np.clip, res, 0, 1, dask="allowed")  # dirty fix
     return res.mean(dim, keep_attrs=keep_attrs)
 
 

xskillscore/tests/test_probabilistic.py

@@ -484,12 +484,13 @@ def test_rps_wilks_example():
     np.testing.assert_allclose(rps(Obs, F2, category_edges), 0.29)
 
 
-def test_2_category_rps_equals_brier_score(o, f_prob):
+@pytest.mark.parametrize("fair_bool", [True, False])
+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),
-        brier_score(o > 0.5, (f_prob > 0.5).mean("member"), dim=None),
+        rps(o, f_prob, category_edges=category_edges, dim=None, fair=fair_bool),
+        brier_score(o > 0.5, (f_prob > 0.5), dim=None, fair=fair_bool),
     )
 
 
@@ -498,6 +499,7 @@ 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)
+    print(res)
     assert (res == 0).all()
 
 
@@ -512,20 +514,79 @@ def test_rps_dask(o_dask, f_prob_dask, category_edges, fair_bool):
 
 @pytest.mark.parametrize("dim", DIMS)
 def test_rps_vs_fair_rps(o, f_prob, category_edges, dim):
-    """Test that fair rps is smaller or equal than rps due to ensemble-size
-    adjustment."""
+    """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(), print("fairrps", frps, "\nufrps", ufrps)
+    # assert (frps <= ufrps).mean() >.9
+    assert (frps <= ufrps).all(), print(
+        "fairrps",
+        frps,
+        "\nufrps",
+        ufrps,
+        "\n diff: ufrps - frps, should be positive:\n",
+        ufrps - frps,
+    )
+
+
+@pytest.mark.parametrize("fair_bool", [True, False])
+def test_rps_category_edges_xrDataArray(o, f_prob, fair_bool):
+    """Test rps with category_edges as xrDataArray for forecast and observations edges."""
+    actual = rps(
+        o,
+        f_prob,
+        dim="time",
+        fair=fair_bool,
+        category_edges=f_prob.quantile(q=[0.3, 0.5, 0.7], dim=["time", "member"]),
+    )
+    assert set(["lon", "lat"]) == set(actual.dims)
+    assert "quantile" not in actual.dims
+
+
+@pytest.mark.parametrize("fair_bool", [True, False])
+def test_rps_category_edges_xrDataset(o, f_prob, fair_bool):
+    """Test rps with category_edges as xrDataArray for forecast and observations edges."""
+    o = o.to_dataset(name="var")
+    o["var2"] = o["var"] ** 2
+    f_prob = f_prob.to_dataset(name="var")
+    f_prob["var2"] = f_prob["var"] ** 2
+    actual = rps(
+        o,
+        f_prob,
+        dim="time",
+        fair=fair_bool,
+        category_edges=f_prob.quantile(q=[0.3, 0.5, 0.7], dim=["time", "member"]),
+    )
+    assert set(["lon", "lat"]) == set(actual.dims)
+    assert "quantile" not in actual.dims
+
+
+@pytest.mark.parametrize("fair_bool", [True, False])
+def test_rps_category_edges_tuple(o, f_prob, fair_bool):
+    """Test rps with category_edges as tuple of xrDataArray for forecast and observations edges separately."""
+    actual = rps(
+        o,
+        f_prob,
+        dim="time",
+        fair=fair_bool,
+        category_edges=(
+            f_prob.quantile(q=[0.3, 0.5, 0.7], dim=["time", "member"]),
+            o.quantile(q=[0.3, 0.5, 0.7], dim="time"),
+        ),
+    )
+    assert set(["lon", "lat"]) == set(actual.dims)
+    assert "quantile" not in actual.dims
 
 
-@pytest.mark.parametrize("dim", DIMS)
 @pytest.mark.parametrize("fair_bool", [True, False])
-def test_rps_limits(o, f_prob, category_edges, fair_bool, dim):
-    """Test rps between 0 and 1. Note: this only works because np.clip(rps,0,1)"""
-    res = rps(o, f_prob, dim=dim, fair=fair_bool, category_edges=category_edges)
-    assert (res <= 1.0).all(), print(res.max())
-    assert (res >= 0).all(), print(res.min())
+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."""
+    edges = xr.DataArray([0.2, 0.4, 0.6, 0.8], dims="quantile")
+    o_c = o > edges  # CDF
+    f_prob_c = f_prob > edges
+    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
 
 
 @pytest.mark.parametrize(