Merge pull request #28 from ecmwf/feature/cpf-dev

Enhancements to CPF

Author: oiffrig

src/earthkit/meteo/extreme/array/cpf.py

@@ -10,39 +10,13 @@
 import numpy as np
 
 
-def cpf(clim, ens, sort_clim=True, sort_ens=True):
-    """Compute Crossing Point Forecast (CPF)
-
-    WARNING: this code is experimental, use at your own risk!
-
-    Parameters
-    ----------
-    clim: numpy array (nclim, npoints)
-        Per-point climatology
-    ens: numpy array (nens, npoints)
-        Ensemble forecast
-    sort_clim: bool
-        If True, sort the climatology first
-    sort_ens: bool
-        If True, sort the ensemble first
-
-    Returns
-    -------
-    numpy array (npoints)
-        CPF values
-    """
+def _cpf(clim, ens, epsilon=None):
     nclim, npoints = clim.shape
-    nens, npoints_ens = ens.shape
-    assert npoints == npoints_ens
+    nens, _ = ens.shape
 
     cpf = np.ones(npoints, dtype=np.float32)
     mask = np.zeros(npoints, dtype=np.bool_)
 
-    if sort_clim:
-        clim = np.sort(clim, axis=0)
-    if sort_ens:
-        ens = np.sort(ens, axis=0)
-
     for icl in range(1, nclim - 1):
         # quantile level of climatology
         tau_c = icl / (nclim - 1.0)
@@ -94,4 +68,58 @@ def cpf(clim, ens, sort_clim=True, sort_ens=True):
                 # speed up process
                 break
 
+    if epsilon is not None:
+        # ens is assumed to be sorted at this point
+        mask = ens[-1, :] < epsilon
+        cpf[mask] = 0.0
+
     return cpf
+
+
+def cpf(clim, ens, sort_clim=True, sort_ens=True, epsilon=None, symmetric=False):
+    """Compute Crossing Point Forecast (CPF)
+
+    WARNING: this code is experimental, use at your own risk!
+
+    Parameters
+    ----------
+    clim: numpy array (nclim, npoints)
+        Per-point climatology
+    ens: numpy array (nens, npoints)
+        Ensemble forecast
+    sort_clim: bool
+        If True, sort the climatology first
+    sort_ens: bool
+        If True, sort the ensemble first
+    epsilon: float or None
+        If set, use this as a threshold for low-signal regions. Ignored if
+        `symmetric` is True
+    symmetric: bool
+        If True, make CPF values below 0.5 use a symmetric computation (CPF of
+        opposite values)
+
+    Returns
+    -------
+    numpy array (npoints)
+        CPF values
+    """
+    _, npoints = clim.shape
+    _, npoints_ens = ens.shape
+    assert npoints == npoints_ens
+
+    if sort_clim:
+        clim = np.sort(clim, axis=0)
+    if sort_ens:
+        ens = np.sort(ens, axis=0)
+
+    if symmetric:
+        epsilon = None
+
+    cpf_direct = _cpf(clim, ens, epsilon)
+
+    if symmetric:
+        cpf_reverse = _cpf(-clim[::-1, :], -ens[::-1, :])
+        mask = cpf_direct < 0.5
+        cpf_direct[mask] = 1 - cpf_reverse[mask]
+
+    return cpf_direct