crps ensemble with dask

Author: frazane

scoringrules/backend/gufuncs.py

@@ -106,7 +106,7 @@ def _crps_ensemble_qd_gufunc(
 ):
     """CRPS estimator based on the quantile decomposition form."""
     obs = observation[0]
-    M = forecasts.shape[0]
+    M = forecasts.shape[-1]
 
     if np.isnan(obs):
         out[0] = np.nan
@@ -194,7 +194,7 @@ def _crps_ensemble_pwm_gufunc(
 ):
     """CRPS estimator based on the probability weighted moment (PWM) form."""
     obs = observation[0]
-    M = forecasts.shape[0]
+    M = forecasts.shape[-1]
 
     if np.isnan(obs):
         out[0] = np.nan

scoringrules/backend/numba.py

@@ -48,6 +48,11 @@ def crps_ensemble(
             "akr_circperm",
             "fair",
         ]:
+            if hasattr(forecasts, "__dask_graph__"):
+                raise NotImplementedError(
+                    "Estimators that need sorting of the ensemble do not yet work with dask. "
+                    "Sort the ensemble beforehand or try another estimator (nrg, akr, akr_circperm or fair)"
+                )
             forecasts = np.sort(forecasts, axis=-1)
 
         if estimator == "int":