Merge pull request #20 from ecmwf/fix/no_danu_dep

refactor: port crps logic across from danu to hyve

Author: Oisin-M

src/hyve/tools/crps_stations.py

@@ -10,7 +10,86 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
-from danu import stats, utils
+
+
+def crps_minmax(x, y):
+    """
+    Computes CRPS from x using y as reference,
+    first x dimension must be ensembles, next dimensions can be arbitrary
+    x: ensemble data (n_ens, n_points)
+    y: observation/analysis data (n_points)
+    returns: crps (n_points)
+    REFERENCE
+      Hersbach, 2000: Decomposition of the Continuous Ranked Probability Score for Ensemble Prediction Systems.
+      Weather and Forecasting 15: 559-570.
+    """
+
+    # first sort ensemble
+    x.sort(axis=0)
+
+    # construct alpha and beta, size nens+1
+    n_ens = x.shape[0]
+    shape = (n_ens + 1,) + x.shape[1:]
+    alpha = np.zeros(shape)
+    beta = np.zeros(shape)
+
+    # x[i+1]-x[i] and x[i]-y[i] arrays
+    diffxy = x - y.reshape(1, *(y.shape))
+    diffxx = x[1:] - x[:-1]  # x[i+1]-x[i], size ens-1
+
+    # if i == 0
+    alpha[0] = 0
+    beta[0] = np.fmax(diffxy[0], 0)  # x(0)-y
+    # if i == n_ens
+    alpha[-1] = np.fmax(-diffxy[-1], 0)  # y-x(n)
+    beta[-1] = 0
+    # else
+    alpha[1:-1] = np.fmin(
+        diffxx, np.fmax(-diffxy[:-1], 0)
+    )  # x(i+1)-x(i) or y-x(i) or 0
+    beta[1:-1] = np.fmin(diffxx, np.fmax(diffxy[1:], 0))  # 0 or x(i+1)-y or x(i+1)-x(i)
+
+    # compute crps
+    p_exp = (np.arange(n_ens + 1) / float(n_ens)).reshape(n_ens + 1, *([1] * y.ndim))
+    crps = np.sum(alpha * (p_exp**2) + beta * ((1 - p_exp) ** 2), axis=0)
+    #
+    # p = np.arange(n_ens+1)/float(n_ens)
+    # alpha_mean = alpha.mean(axis=1)
+    # beat_mean = beta.mean(axis=1)
+    # crps = alpha_mean*(p**2) + beat_mean*((1-p)**2)
+    # crps_mean = crps2.sum()
+    #
+    # p_exp = np.expand_dims(np.arange(n_ens+1)/float(n_ens), axis=1)
+    # crps = np.nansum(alpha*(p_exp**2) + beta*((1-p_exp)**2), axis=0)
+    # crps_mean = crps.mean()
+    return crps
+
+
+def crps_masked(x, y):
+    n_steps = x.shape[0]
+
+    mask = np.logical_not(np.isnan(x[0, 0]))
+    crps = np.ones(y.shape) * np.nan
+    for i in range(n_steps):
+        xi = x[i]
+        yi = y[i]
+        crps_masked = crps_minmax(xi[:, mask], yi[mask])
+        crps[i][mask] = crps_masked
+
+    return crps
+
+
+def forecast_crps(x, y, core_dims=["lat", "lon"]):
+    crps = xr.apply_ufunc(
+        crps_masked,
+        x,
+        y,
+        input_core_dims=[["ensemble", *core_dims], core_dims],
+        dask="parallelized",
+        output_core_dims=[core_dims],
+        output_dtypes=[np.float32],
+    )
+    return crps
 
 
 @dask.delayed
@@ -58,7 +137,6 @@ def coord_dmh(dates):
     return days_months
 
 
-@utils.timer
 def compute_score(
     out_dir, reforecast_dir, ds_reanalysis, ds_clim, core_dims, with_init=False
 ):
@@ -113,14 +191,12 @@ def compute_score(
         persistence = ds_reanalysis.sel(time=date_persistence)
 
         crps_pers = persistence_crps(reanalysis, persistence)
-        crps_refo = stats.forecast_crps(reforecast, reanalysis, core_dims=core_dims)
+        crps_refo = forecast_crps(reforecast, reanalysis, core_dims=core_dims)
         if ds_clim is not None:
             log.debug(coord_dmh(date_range))
             climatology = ds_clim.sel(time=coord_dmh(date_range))
             climatology.coords["time"] = date_range
-            crps_clim = stats.forecast_crps(
-                climatology, reanalysis, core_dims=core_dims
-            )
+            crps_clim = forecast_crps(climatology, reanalysis, core_dims=core_dims)
             crps_refo, crps_pers, crps_clim = dask.compute(
                 crps_refo, crps_pers, crps_clim
             )