Merge 08bed70 into 659568e

xskillscore/core/accessor.py

@@ -110,36 +110,66 @@ def smape(self, a, b, dim, weights=None, skipna=False):
         b = self._in_ds(b)
         return smape(a, b, dim, weights=weights, skipna=skipna)
 
-    def crps_gaussian(self, observations, mu, sig):
+    def crps_gaussian(self, observations, mu, sig, dim=None, weights=None):
         observations = self._in_ds(observations)
         mu = self._in_ds(mu)
         sig = self._in_ds(sig)
-        return crps_gaussian(observations, mu, sig)
+        return crps_gaussian(observations, mu, sig, dim=dim, weights=weights)
 
     def crps_ensemble(
-        self, observations, forecasts, weights=None, issorted=False, dim='member'
+        self,
+        observations,
+        forecasts,
+        member_weights=None,
+        issorted=False,
+        dim=None,
+        member_dim='member',
+        weights=None,
     ):
         observations = self._in_ds(observations)
         forecasts = self._in_ds(forecasts)
         return crps_ensemble(
-            observations, forecasts, weights=weights, issorted=issorted, dim='member'
+            observations,
+            forecasts,
+            member_weights=member_weights,
+            issorted=issorted,
+            member_dim=member_dim,
+            dim=dim,
+            weights=weights,
         )
 
-    def crps_quadrature(self, x, cdf_or_dist, xmin=None, xmax=None, tol=1e-6):
+    def crps_quadrature(
+        self, x, cdf_or_dist, xmin=None, xmax=None, tol=1e-6, dim=None, weights=None
+    ):
         x = self._in_ds(x)
         cdf_or_dist = self._in_ds(cdf_or_dist)
-        return crps_quadrature(x, cdf_or_dist, xmin=xmin, xmax=xmax, tol=1e-6)
+        return crps_quadrature(
+            x, cdf_or_dist, xmin=xmin, xmax=xmax, tol=1e-6, dim=dim, weights=weights
+        )
 
     def threshold_brier_score(
-        self, observations, forecasts, threshold, issorted=False, dim='member'
+        self,
+        observations,
+        forecasts,
+        threshold,
+        issorted=False,
+        dim=None,
+        member_dim='member',
+        weights=None,
     ):
         observations = self._in_ds(observations)
         forecasts = self._in_ds(forecasts)
         return threshold_brier_score(
-            observations, forecasts, threshold, issorted=issorted, dim='member'
+            observations,
+            forecasts,
+            threshold,
+            issorted=issorted,
+            dim=dim,
+            member_dim=member_dim,
+            weights=weights,
         )
 
-    def brier_score(self, observations, forecasts):
+    def brier_score(self, observations, forecasts, dim=None, weights=None):
         observations = self._in_ds(observations)
         forecasts = self._in_ds(forecasts)
-        return brier_score(observations, forecasts)
+        return brier_score(observations, forecasts, dim=dim, weights=weights)

xskillscore/core/probabilistic.py

@@ -17,10 +17,11 @@
 ]
 
 
-def xr_crps_gaussian(observations, mu, sig):
+def xr_crps_gaussian(observations, mu, sig, dim=None, weights=None):
     """
     xarray version of properscoring.crps_gaussian: Continuous Ranked
      Probability Score with a Gaussian distribution.
+
     Parameters
     ----------
     observations : xarray.Dataset or xarray.DataArray
@@ -29,13 +30,21 @@ def xr_crps_gaussian(observations, mu, sig):
         The mean of the forecast normal distribution.
     sig : xarray.Dataset or xarray.DataArray
         The standard deviation of the forecast distribution.
+    dim : str or list of str, optional
+        Dimension to mean over after calculating crps_gaussian.
+        Defaults to None implying averaging.
+    weights : xr.DataArray with dimensions from dim, optional
+        Weigths for `weighted.mean(dim)`. Defaults to None implying unweighted mean.
+
     Returns
     -------
-    xarray.Dataset or xarray.DataArray
+    xarray.Dataset or xarray.DataArray reduced by dimension dim
+
     See Also
     --------
     properscoring.crps_gaussian
     xarray.apply_ufunc
+
     """
     # check if same dimensions
     if isinstance(mu, (int, float)):
@@ -46,7 +55,7 @@ def xr_crps_gaussian(observations, mu, sig):
         observations, mu = xr.broadcast(observations, mu)
     if sig.dims != observations.dims:
         observations, sig = xr.broadcast(observations, sig)
-    return xr.apply_ufunc(
+    res = xr.apply_ufunc(
         crps_gaussian,
         observations,
         mu,
@@ -55,28 +64,47 @@ def xr_crps_gaussian(observations, mu, sig):
         dask='parallelized',
         output_dtypes=[float],
     )
+    if dim is None:
+        return res
+    else:
+        if weights is not None:
+            return res.weighted(weights).mean(dim)
+        else:
+            return res.mean(dim)
 
 
-def xr_crps_quadrature(x, cdf_or_dist, xmin=None, xmax=None, tol=1e-6):
+def xr_crps_quadrature(
+    x, cdf_or_dist, xmin=None, xmax=None, tol=1e-6, dim=None, weights=None
+):
     """
     xarray version of properscoring.crps_quadrature: Continuous Ranked
      Probability Score with numerical integration of the normal distribution
+
     Parameters
     ----------
     x : xarray.Dataset or xarray.DataArray
         Observations associated with the forecast distribution ``cdf_or_dist``.
     cdf_or_dist : callable or scipy.stats.distribution
         Function which returns the cumulative density of the forecast
         distribution at value x.
+    xmin, xmax, tol: see properscoring.crps_quadrature
+    dim : str or list of str, optional
+        Dimension to mean over after calculating crps_gaussian.
+        Defaults to None implying averaging.
+    weights : xr.DataArray with dimensions from dim, optional
+        Weigths for `weighted.mean(dim)`. Defaults to None implying unweighted mean.
+
     Returns
     -------
     xarray.Dataset or xarray.DataArray
+
     See Also
     --------
     properscoring.crps_quadrature
     xarray.apply_ufunc
+
     """
-    return xr.apply_ufunc(
+    res = xr.apply_ufunc(
         crps_quadrature,
         x,
         cdf_or_dist,
@@ -87,90 +115,144 @@ def xr_crps_quadrature(x, cdf_or_dist, xmin=None, xmax=None, tol=1e-6):
         dask='parallelized',
         output_dtypes=[float],
     )
+    if dim is None:
+        return res
+    else:
+        if weights is not None:
+            return res.weighted(weights).mean(dim)
+        else:
+            return res.mean(dim)
 
 
 def xr_crps_ensemble(
-    observations, forecasts, weights=None, issorted=False, dim='member'
+    observations,
+    forecasts,
+    member_weights=None,
+    issorted=False,
+    member_dim='member',
+    dim=None,
+    weights=None,
 ):
     """
     xarray version of properscoring.crps_ensemble: Continuous Ranked
      Probability Score with the ensemble distribution
+
     Parameters
     ----------
     observations : xarray.Dataset or xarray.DataArray
         The observations or set of observations.
     forecasts : xarray.Dataset or xarray.DataArray
         Forecast with required member dimension ``dim``.
-    weights : xarray.Dataset or xarray.DataArray
+    member_weights : xarray.Dataset or xarray.DataArray
         If provided, the CRPS is calculated exactly with the assigned
         probability weights to each forecast. Weights should be positive,
         but do not need to be normalized. By default, each forecast is
         weighted equally.
     issorted : bool, optional
         Optimization flag to indicate that the elements of `ensemble` are
         already sorted along `axis`.
-    dim : str, optional
+    member_dim : str, optional
         Name of ensemble member dimension. By default, 'member'.
+    dim : str or list of str, optional
+        Dimension to mean over after calculating crps_gaussian.
+        Defaults to None implying averaging.
+    weights : xr.DataArray with dimensions from dim, optional
+        Weigths for `weighted.mean(dim)`. Defaults to None implying unweighted mean.
+
     Returns
     -------
     xarray.Dataset or xarray.DataArray
+
     See Also
     --------
     properscoring.crps_ensemble
     xarray.apply_ufunc
+
     """
-    return xr.apply_ufunc(
+    res = xr.apply_ufunc(
         crps_ensemble,
         observations,
         forecasts,
-        input_core_dims=[[], [dim]],
-        kwargs={'axis': -1, 'issorted': issorted, 'weights': weights},
+        input_core_dims=[[], [member_dim]],
+        kwargs={'axis': -1, 'issorted': issorted, 'weights': member_weights},
         dask='parallelized',
         output_dtypes=[float],
     )
+    if dim is None:
+        return res
+    else:
+        if weights is not None:
+            return res.weighted(weights).mean(dim)
+        else:
+            return res.mean(dim)
 
 
-def xr_brier_score(observations, forecasts):
+def xr_brier_score(observations, forecasts, dim=None, weights=None):
     """
     xarray version of properscoring.brier_score: Calculate Brier score (BS).
+
     ..math:
-        BS(p, k) = (p_1 - k)^2,
+        BS(p, k) = (p_1 - k)^2
+
     Parameters
     ----------
     observations : xarray.Dataset or xarray.DataArray
         The observations or set of observations.
     forecasts : xarray.Dataset or xarray.DataArray
         The forecasts associated with the observations.
+    dim : str or list of str, optional
+        Dimension to mean over after calculating crps_gaussian.
+        Defaults to None implying averaging.
+    weights : xr.DataArray with dimensions from dim, optional
+        Weigths for `weighted.mean(dim)`. Defaults to None implying unweighted mean.
+
     Returns
     -------
     xarray.Dataset or xarray.DataArray
+
     References
     ----------
     Gneiting, Tilmann, and Adrian E Raftery. “Strictly Proper Scoring Rules,
       Prediction, and Estimation.” Journal of the American Statistical
       Association 102, no. 477 (March 1, 2007): 359–78.
       https://doi.org/10/c6758w.
+
     See Also
     --------
     properscoring.brier_score
     xarray.apply_ufunc
+
     """
-    return xr.apply_ufunc(
+    res = xr.apply_ufunc(
         brier_score,
         observations,
         forecasts,
         input_core_dims=[[], []],
         dask='parallelized',
         output_dtypes=[float],
     )
+    if dim is None:
+        return res
+    else:
+        if weights is not None:
+            return res.weighted(weights).mean(dim)
+        else:
+            return res.mean(dim)
 
 
 def xr_threshold_brier_score(
-    observations, forecasts, threshold, issorted=False, dim='member'
+    observations,
+    forecasts,
+    threshold,
+    issorted=False,
+    member_dim='member',
+    dim=None,
+    weights=None,
 ):
     """
     xarray version of properscoring.threshold_brier_score: Calculate the Brier
      scores of an ensemble for exceeding given thresholds.
+
     Parameters
     ----------
     observations : xarray.Dataset or xarray.DataArray
@@ -182,23 +264,32 @@ def xr_threshold_brier_score(
     issorted : bool, optional
         Optimization flag to indicate that the elements of `ensemble` are
         already sorted along `axis`.
-    dim : str, optional
+    member_dim : str, optional
         Name of ensemble member dimension. By default, 'member'.
+    dim : str or list of str, optional
+        Dimension to mean over after calculating crps_gaussian.
+        Defaults to None implying averaging.
+    weights : xr.DataArray with dimensions from dim, optional
+        Weigths for `weighted.mean(dim)`. Defaults to None implying unweighted mean.
+
     Returns
     -------
     xarray.Dataset or xarray.DataArray
         (If ``threshold`` is a scalar, the result will have the same shape as
         observations. Otherwise, it will have an additional final dimension
         corresponding to the threshold levels. Not implemented yet.)
+
     References
     ----------
     Gneiting, T. and Ranjan, R. Comparing density forecasts using threshold-
        and quantile-weighted scoring rules. J. Bus. Econ. Stat. 29, 411-422
        (2011). http://www.stat.washington.edu/research/reports/2008/tr533.pdf
+
     See Also
     --------
     properscoring.threshold_brier_score
     xarray.apply_ufunc
+
     """
     if isinstance(threshold, list):
         threshold.sort()
@@ -210,18 +301,18 @@ def xr_threshold_brier_score(
             raise ValueError(
                 'please provide threshold with threshold dim, found', threshold.dims,
             )
-        input_core_dims = [[], [dim], ['threshold']]
+        input_core_dims = [[], [member_dim], ['threshold']]
         output_core_dims = [['threshold']]
     elif isinstance(threshold, (int, float)):
-        input_core_dims = [[], [dim], []]
+        input_core_dims = [[], [member_dim], []]
         output_core_dims = [[]]
     else:
         raise ValueError(
             'Please provide threshold as list, int, float \
             or xr.object with threshold dimension; found',
             type(threshold),
         )
-    return xr.apply_ufunc(
+    res = xr.apply_ufunc(
         threshold_brier_score,
         observations,
         forecasts,
@@ -232,3 +323,10 @@ def xr_threshold_brier_score(
         dask='parallelized',
         output_dtypes=[float],
     )
+    if dim is None:
+        return res
+    else:
+        if weights is not None:
+            return res.weighted(weights).mean(dim)
+        else:
+            return res.mean(dim)