Implement array submodule (#6)

* Implement array submodule

---------

Co-authored-by: Olivier Iffrig <[email protected]>

docs/conf.py

@@ -55,6 +55,7 @@
 autoapi_root = "_api"
 autoapi_member_order = "alphabetical"
 autoapi_add_toctree_entry = False
+autoapi_own_page_level = "function"
 
 # napoleon configuration
 napoleon_google_docstring = False
@@ -86,3 +87,9 @@
 html_css_files = ["style.css"]
 
 html_logo = "_static/earthkit-meteo.png"
+
+
+def setup(app):
+    from skip_api_rules import _skip_api_items
+
+    app.connect("autoapi-skip-member", _skip_api_items)

docs/skip_api_rules.py

@@ -0,0 +1,28 @@
+# (C) Copyright 2021 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+
+# define skip rules for autoapi
+def _skip_api_items(app, what, name, obj, skip, options):
+    # print(f"{what=} {name=}")
+
+    if (
+        what == "module"
+        and ".array" not in name
+        and name not in ["earthkit.meteo.solar", "earthkit.meteo.solar.array"]
+    ):
+        skip = True
+    elif what == "package" and ".array" not in name and len(name.split(".")) > 2:
+        skip = True
+    elif what == "function" and ".array" not in name:
+        skip = True
+
+    # if not skip:
+    #     print(f"{what} {name}")
+    return skip

earthkit/meteo/__init__.py

@@ -1,11 +1,11 @@
-# (C) Copyright 2023 ECMWF.
+# (C) Copyright 2021 ECMWF.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-
+#
 
 try:
     # NOTE: the `version.py` file must not be present in the git repository

earthkit/meteo/constants/__init__.py

@@ -1 +1,10 @@
+# (C) Copyright 2021 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
 from .constants import *  # noqa

earthkit/meteo/constants/constants.py

@@ -1,10 +1,11 @@
-# (C) Copyright 2021- ECMWF.
-
+# (C) Copyright 2021 ECMWF.
+#
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
-
-# In applying this licence, ECMWF does not waive the privileges and immunities granted to it by
-# virtue of its status as an intergovernmental organisation nor does it submit to any jurisdiction.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
 
 import numpy as np
 

earthkit/meteo/extreme/__init__.py

@@ -1,3 +1,20 @@
+# (C) Copyright 2021 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+"""
+Extreme index functions.
+
+The API is split into two levels. The low level functions are in the ``array`` submodule and they
+can be used to operate on numpy arrays. The high level functions are still to be developed and
+planned to work with objects like *earthkit.data FieldLists* or *xarray DataSets*.
+"""
+
 from .cpf import *  # noqa
 from .efi import *  # noqa
 from .sot import *  # noqa

earthkit/meteo/extreme/array/__init__.py

@@ -0,0 +1,16 @@
+# (C) Copyright 2021 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+"""
+Extreme index functions operating on numpy arrays.
+"""
+
+from .cpf import *  # noqa
+from .efi import *  # noqa
+from .sot import *  # noqa

earthkit/meteo/extreme/array/cpf.py

@@ -0,0 +1,99 @@
+# (C) Copyright 2021 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+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
+    """
+    nclim, npoints = clim.shape
+    nens, npoints_ens = ens.shape
+    assert npoints == npoints_ens
+
+    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)
+        for iq in range(nens):
+            # quantile level of forecast
+            tau_f = (iq + 1.0) / (nens + 1.0)
+            if tau_f >= tau_c:
+                # quantile values of forecast and climatology
+                qv_f = ens[iq, :]
+                qv_c = clim[icl, :]
+
+                # lowest climate quantile: interpolate between 2 consecutive quantiles
+                if iq < 2:
+                    # quantile value and quantile level of climatology at previous
+                    qv_c_2 = clim[icl - 1, :]
+                    tau_c_2 = (icl - 1) / (nclim - 1)
+
+                    # condition of crossing situtaion:
+                    idx = (qv_f < qv_c) & (qv_c_2 < qv_c)
+
+                    # intersection between two lines
+                    tau_i = (
+                        tau_c * (qv_c_2[idx] - qv_f[idx])
+                        + tau_c_2 * (qv_f[idx] - qv_c[idx])
+                    ) / (qv_c_2[idx] - qv_c[idx])
+
+                    # populate matrix, no values below 0
+                    cpf[idx] = np.maximum(tau_i, 0)
+                    mask[idx] = True
+
+                # check crossing cases
+                idx = (qv_f < qv_c) & (~mask)
+                cpf[idx] = tau_f
+                mask[idx] = True
+
+                # largest climate quantile: interpolate
+                if iq == nens - 1:
+                    qv_c_2 = clim[nclim - 1, :]
+                    tau_c_2 = 1.0
+
+                    idx = (qv_f > qv_c) & (qv_c_2 > qv_c) & (~mask)
+
+                    tau_i = (
+                        tau_c * (qv_c_2[idx] - qv_f[idx])
+                        + tau_c_2 * (qv_f[idx] - qv_c[idx])
+                    ) / (qv_c_2[idx] - qv_c[idx])
+
+                    # populate matrix, no values above 1
+                    cpf[idx] = np.minimum(tau_i, 1)
+
+                # speed up process
+                break
+
+    return cpf

earthkit/meteo/extreme/array/efi.py

@@ -0,0 +1,148 @@
+# (C) Copyright 2021 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+import numpy as np
+
+# import numba
+# from numba import float64, float32
+
+
+def efi(clim, ens, eps=-0.1):
+    """Compute Extreme Forecast Index (EFI)
+
+    Parameters
+    ----------
+    clim: numpy array (nclim, npoints)
+        Sorted per-point climatology
+    ens: numpy array (nens, npoints)
+        Ensemble forecast
+    eps: (float)
+        Epsilon factor for zero values
+
+    Returns
+    -------
+    numpy array (npoints)
+        EFI values
+    """
+    # locate missing values
+    missing_mask = np.logical_or(
+        np.sum(np.isnan(clim), axis=0), np.sum(np.isnan(ens), axis=0)
+    )
+
+    # Compute fraction of the forecast below climatology
+    nclim, npoints = clim.shape
+    nens, npoints_ens = ens.shape
+    assert npoints == npoints_ens
+    frac = np.zeros_like(clim)
+    ##################################
+    for icl in range(nclim):
+        frac[icl, :] = np.sum(ens[:, :] <= clim[icl, np.newaxis, :], axis=0)
+    ##################################
+    frac /= nens
+
+    # Compute formula coefficients
+    p = np.linspace(0.0, 1.0, nclim)
+    dp = 1 / (nclim - 1)
+    dFdp = np.diff(frac, axis=0) / dp
+
+    acosdiff = np.diff(np.arccos(np.sqrt(p)))
+    proddiff = np.diff(np.sqrt(p * (1.0 - p)))
+
+    acoef = (1.0 - 2.0 * p[:-1]) * acosdiff + proddiff
+
+    # compute EFI from coefficients
+    efi = np.zeros(npoints)
+    ##################################
+    if eps > 0:
+        efimax = np.zeros(npoints)
+        for icl in range(nclim - 1):
+            mask = clim[icl + 1, :] > eps
+            dEFI = np.where(
+                mask,
+                (2.0 * frac[icl, :] - 1.0) * acosdiff[icl]
+                + acoef[icl] * dFdp[icl, :]
+                - proddiff[icl],
+                0.0,
+            )
+            defimax = np.where(mask, -acosdiff[icl] - proddiff[icl], 0.0)
+            efi += dEFI
+            efimax += defimax
+        efimax = np.fmax(efimax, eps)
+        efi /= efimax
+    else:
+        for icl in range(nclim - 1):
+            dEFI = (
+                (2.0 * frac[icl, :] - 1.0) * acosdiff[icl]
+                + acoef[icl] * dFdp[icl, :]
+                - proddiff[icl]
+            )
+            efi += dEFI
+        efi *= 2.0 / np.pi
+    ##################################
+
+    # apply missing values
+    efi[missing_mask] = np.nan
+
+    return efi
+
+
+# @numba.jit(float64[:](float64[:,:], float64[:,:]), fastmath=False, nopython=True, nogil=True, cache=True)
+# @numba.jit(nopython=True)
+# def efi_numba(clim, ens):
+#     """Compute EFI
+
+#     Parameters
+#     ----------
+#     clim: numpy array (nclim, npoints)
+#         Sorted per-point climatology
+#     ens: numpy array (nens, npoints)
+#         Ensemble forecast
+
+#     Returns
+#     -------
+#     numpy array (npoints)
+#         EFI values
+#     """
+
+#     # Compute fraction of the forecast below climatology
+#     nclim, npoints = clim.shape
+#     nens, npoints_ens = ens.shape
+#     assert npoints == npoints_ens
+#     frac = np.zeros_like(clim)
+#     ##################################
+#     for ifo in numba.prange(nens):
+#         for icl in range(nclim):
+#             for i in range(npoints):
+#                if ens[ifo, i] <= clim[icl, i]:
+#                    frac[icl, i] += 1
+#     ##################################
+#     frac /= nens
+
+#     # Compute formula coefficients
+#     p = np.linspace(0., 1., nclim)
+#     dp = 1 / (nclim - 1)  #np.diff(p)
+
+#     acosdiff = np.diff(np.arccos(np.sqrt(p)))
+#     proddiff = np.diff(np.sqrt(p * (1. - p)))
+
+#     acoef = (1. - 2. * p[:-1]) * acosdiff + proddiff
+
+#     # TODO: handle epsilon
+#     efi = np.zeros(npoints)
+#     ##################################
+#     for icl in numba.prange(nclim-1):
+#         for i in range(npoints):
+#             dFdp = (frac[icl+1, i] - frac[icl, i]) / dp
+#              # XXX: why proddiff here?!
+#             dEFI = (2. * frac[icl, i] - 1.) * acosdiff[icl] + acoef[icl] * dFdp - proddiff[icl]
+#             efi[i] += dEFI
+#     efi *= 2. / np.pi
+#     ##################################
+
+#     return efi