Add error spread score (#10)

frazane · web-flow · commit 02afe44cbc09 · 2023-12-22T11:02:11.000+01:00
* added ess

* docstrings

* updated docstrings

* wip

* update docs
diff --git a/docs/api/error_spread.md b/docs/api/error_spread.md
@@ -0,0 +1,3 @@
+# Error Spread Score
+
+::: scoringrules.error_spread_score
diff --git a/docs/api/logarithmic.md b/docs/api/logarithmic.md
@@ -0,0 +1,3 @@
+# Logarithmic Score
+
+::: scoringrules.logs_normal
diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css
@@ -43,3 +43,9 @@ h6.heading {
   border-radius: 2pt;
   padding: 0pt 5pt 2pt 5pt;
 }
+
+
+/* Maximum space for text block */
+.md-grid {
+  max-width: 90%; /* or 100%, if you want to stretch to full-width */
+}
diff --git a/mkdocs.yaml b/mkdocs.yaml
@@ -13,6 +13,7 @@ nav:
       - Analytical formulations: api/crps/analytical.md
       - Weighted versions: api/crps/weighted.md
     - Logarithmic Score: api/logarithmic.md
+    - Error Spread Score: api/error_spread.md
     - Energy Score: 
       - api/energy/index.md
       - Ensemble-based estimation: api/energy/ensemble.md
diff --git a/scoringrules/__init__.py b/scoringrules/__init__.py
@@ -16,6 +16,7 @@
     twenergy_score,
     vrenergy_score,
 )
+from scoringrules._error_spread import error_spread_score
 from scoringrules._logs import logs_normal
 from scoringrules._variogram import (
     owvariogram_score,
@@ -40,6 +41,7 @@
     "vrcrps_ensemble",
     "logs_normal",
     "brier_score",
+    "error_spread_score",
     "energy_score",
     "owenergy_score",
     "twenergy_score",
diff --git a/scoringrules/_crps.py b/scoringrules/_crps.py
@@ -100,8 +100,6 @@ def twcrps_ensemble(
         Chaining function used to emphasise particular outcomes. For example, a function that
         only considers values above a certain threshold $t$ by projecting forecasts and observations
         to $\[t, \inf)$.
-    v_funcargs: tuple
-        Additional arguments to the chaining function.
     axis: int
         The axis corresponding to the ensemble. Default is the last axis.
     backend: str
@@ -158,8 +156,6 @@ def owcrps_ensemble(
         The observed values.
     w_func: tp.Callable
         Weight function used to emphasise particular outcomes.
-    w_funcargs: tuple
-        Additional arguments to the weight function.
     axis: int
         The axis corresponding to the ensemble. Default is the last axis.
     backend: str
@@ -235,8 +231,6 @@ def vrcrps_ensemble(
         The observed values.
     w_func: tp.Callable
         Weight function used to emphasise particular outcomes.
-    w_funcargs: tuple
-        Additional arguments to the weight function.
     axis: int
         The axis corresponding to the ensemble. Default is the last axis.
     backend: str
diff --git a/scoringrules/_energy.py b/scoringrules/_energy.py
@@ -143,8 +143,6 @@ def owenergy_score(
         The observed values, where the variables dimension is by default the last axis.
     w_func: tp.Callable
         Weight function used to emphasise particular outcomes.
-    w_funcargs: tuple
-        Additional arguments to the weight function.
     m_axis: int
         The axis corresponding to the ensemble dimension. Defaults to -2.
     v_axis: int or tuple(int)
@@ -211,8 +209,6 @@ def vrenergy_score(
         The observed values, where the variables dimension is by default the last axis.
     w_func: tp.Callable
         Weight function used to emphasise particular outcomes.
-    w_funcargs: tuple
-        Additional arguments to the weight function.
     m_axis: int
         The axis corresponding to the ensemble dimension. Defaults to -2.
     v_axis: int or tuple(int)
diff --git a/scoringrules/_error_spread.py b/scoringrules/_error_spread.py
@@ -0,0 +1,62 @@
+import typing as tp
+
+from scoringrules.backend import backends
+from scoringrules.core import error_spread
+
+if tp.TYPE_CHECKING:
+    from scoringrules.core.typing import Array, ArrayLike, Backend
+
+
+def error_spread_score(
+    forecasts: "Array",
+    observations: "ArrayLike",
+    /,
+    axis: int = -1,
+    *,
+    backend: "Backend" = None,
+) -> "Array":
+    r"""Compute the error-spread score (ESS) for a finite ensemble.
+
+    The error spread score [(Christensen et al., 2015)](https://doi.org/10.1002/qj.2375) is given by:
+    
+    $$ESS = \left(s^2 - e^2 - e \cdot s \cdot g\right)^2$$
+
+    where the mean $m$, variance $s^2$, and skewness $g$ of the ensemble forecast of size $F$ are computed as follows:
+    
+    $$m = \frac{1}{F} \sum_{f=1}^{F} X_f, \quad s^2 = \frac{1}{F-1} \sum_{f=1}^{F} (X_f - m)^2, \quad g = \frac{F}{(F-1)(F-2)} \sum_{f=1}^{F} \left(\frac{X_f - m}{s}\right)^3$$
+
+    The error in the ensemble mean $e$ is calculated as $e = m - y$, where $y$ is the observed value.
+
+    Parameters
+    ----------
+    forecasts: Array
+        The predicted forecast ensemble, where the ensemble dimension is by default
+        represented by the last axis.
+    observations: ArrayLike
+        The observed values.
+    axis: int
+        The axis corresponding to the ensemble. Default is the last axis.
+    backend: str
+        The name of the backend used for computations. Defaults to 'numba' if available, else 'numpy'.
+
+    Returns
+    -------
+    - Array
+        An array of error spread scores for each ensemble forecast, which should be averaged to get meaningful values.
+    """
+    B = backends.active if backend is None else backends[backend]
+    forecasts, observations = map(B.asarray, (forecasts, observations))
+
+    if axis != -1:
+        forecasts = B.moveaxis(forecasts, axis, -1)
+
+    if B.name == "numba":
+        return error_spread._ess_gufunc(forecasts, observations)
+
+    return error_spread.ess(forecasts, observations, backend=backend)
+
+
+
+    # \[
+    #     ESS = \left(s^2 - e^2 - e \cdot s \cdot g\right)^2
+    # \]
diff --git a/scoringrules/_variogram.py b/scoringrules/_variogram.py
@@ -94,8 +94,6 @@ def twvariogram_score(
         The order of the Variogram Score. Typical values are 0.5, 1.0 or 2.0. Defaults to 1.0.
     v_func: tp.Callable
         Chaining function used to emphasise particular outcomes.
-    v_funcargs: tuple
-        Additional arguments to the chaining function.
     m_axis: int
         The axis corresponding to the ensemble dimension. Defaults to -2.
     v_axis: int
@@ -152,8 +150,6 @@ def owvariogram_score(
         The order of the Variogram Score. Typical values are 0.5, 1.0 or 2.0. Defaults to 1.0.
     w_func: tp.Callable
         Weight function used to emphasise particular outcomes.
-    w_funcargs: tuple
-        Additional arguments to the weight function.
     m_axis: int
         The axis corresponding to the ensemble dimension. Defaults to -2.
     v_axis: int
@@ -223,8 +219,6 @@ def vrvariogram_score(
         The order of the Variogram Score. Typical values are 0.5, 1.0 or 2.0. Defaults to 1.0.
     w_func: tp.Callable
         Weight function used to emphasise particular outcomes.
-    w_funcargs: tuple
-        Additional arguments to the weight function.
     m_axis: int
         The axis corresponding to the ensemble dimension. Defaults to -2.
     v_axis: int
diff --git a/scoringrules/backend/registry.py b/scoringrules/backend/registry.py
@@ -1,4 +1,3 @@
-import sys
 import typing as tp
 from importlib.util import find_spec
 
@@ -20,7 +19,7 @@
 }
 
 try:
-    import numba  # type: ignore
+    import numba  # noqa: F401
 
     _NUMBA_IMPORTED = True
 except ImportError:
@@ -31,8 +30,9 @@ class BackendsRegistry(dict[str, ArrayBackend]):
     """A dict-like container of registered backends."""
 
     def __init__(self):
-        for backend in self.available_backends:
-            self.register_backend(backend)
+        self.register_backend("numpy")
+        if _NUMBA_IMPORTED:
+            self.register_backend("numba")
 
         self._active = "numba" if _NUMBA_IMPORTED else "numpy"
 
diff --git a/scoringrules/backend/torch.py b/scoringrules/backend/torch.py
@@ -84,7 +84,7 @@ def expand_dims(self, x: "Tensor", /, axis: int | tuple[int] = 0) -> "Tensor":
     def squeeze(
         self, x: "Tensor", /, *, axis: int | tuple[int, ...] | None = None
     ) -> "Tensor":
-        return torch.squeeze(x, dim=axis)
+        return torch.squeeze(x, dim=() if axis is None else axis)
 
     def stack(
         self, arrays: tuple["Tensor", ...] | list["Tensor"], /, *, axis: int = 0
diff --git a/scoringrules/core/error_spread/__init__.py b/scoringrules/core/error_spread/__init__.py
@@ -0,0 +1,4 @@
+from ._score import error_spread_score as ess
+from ._gufunc import _error_spread_score_gufunc as _ess_gufunc
+
+__all__ = ["ess"]
diff --git a/scoringrules/core/error_spread/_gufunc.py b/scoringrules/core/error_spread/_gufunc.py
@@ -0,0 +1,44 @@
+import math
+
+import numpy as np
+from numba import guvectorize
+
+INV_SQRT_PI = 1 / np.sqrt(np.pi)
+EPSILON = 1e-6
+
+
+@guvectorize(
+    [
+        "void(float32[:], float32[:], float32[:])",
+        "void(float64[:], float64[:], float64[:])",
+    ],
+    "(n),()->()",
+)
+def _error_spread_score_gufunc(
+    forecasts: np.ndarray, observation: np.ndarray, out: np.ndarray
+):
+    """Error spread score."""
+    M = forecasts.shape[0]
+    m = 0.0
+    s = 0.0
+    g = 0.0
+    e = 0.0
+
+    for i in range(M):
+        m += forecasts[i]
+    m /= M
+
+    for i in range(M):
+        s += (forecasts[i] - m) ** 2
+    s /= (M - 1)
+    s = np.sqrt(s)
+
+    for i in range(M):
+        g += ((forecasts[i] - m) / s) ** 3
+    g /= ((M - 1) * (M - 2))
+
+    for i in range(M):
+        e += forecasts[i] - observation[i]
+
+    out[0] = (s**2 - e**2 - e * s * g) ** 2
+
diff --git a/scoringrules/core/error_spread/_score.py b/scoringrules/core/error_spread/_score.py
@@ -0,0 +1,21 @@
+import typing as tp
+
+from scoringrules.backend import backends
+
+if tp.TYPE_CHECKING:
+    from scoringrules.core.typing import Array, Backend
+
+
+def error_spread_score(
+    fcts: "Array", obs: "Array", backend: "Backend" = None
+) -> "Array":
+    """Error spread score."""
+    B = backends.active if backend is None else backends[backend]
+    M: int = fcts.shape[-1]
+    m = B.mean(fcts, axis=-1)
+    s = B.sqrt(B.sum((fcts - m[...,None]) ** 2, axis=-1) / (M - 1))
+    g = B.sum(((fcts - m[...,None]) / s[...,None]) ** 3, axis=-1) / ((M - 1) * (M - 2))
+    e = m - obs
+    return B.squeeze((s**2 - e**2 - e * s * g) ** 2)
+
+
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,3 +1,7 @@
 from scoringrules.backend import backends
 
-BACKENDS = backends.available_backends
+RUN_TESTS = ["numba","jax","torch"]
+BACKENDS = [b for b in backends.available_backends if b in RUN_TESTS]
+
+for backend in RUN_TESTS:
+    backends.register_backend(backend)
diff --git a/tests/test_error_spread.py b/tests/test_error_spread.py
@@ -0,0 +1,21 @@
+import numpy as np
+import pytest
+from scoringrules._error_spread import error_spread_score
+
+from .conftest import BACKENDS
+
+ENSEMBLE_SIZE = 11
+N = 5
+
+
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_error_spread_score(backend):
+    obs = np.random.randn(N,)
+    fcts = obs[:,None] + np.random.randn(N, ENSEMBLE_SIZE)
+    res = error_spread_score(fcts, obs, backend=backend)
+
+    # approx zero when perfect forecast
+    perfect_fcts = obs[..., None] + np.random.randn(N, ENSEMBLE_SIZE) * 1e-5
+    res = error_spread_score(perfect_fcts, obs, backend=backend)
+    res = np.asarray(res)
+    assert not np.any(res - 0.0 > 0.0001)

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# Error Spread Score`
	`2`	`+`
	`3`	`+::: scoringrules.error_spread_score`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# Logarithmic Score`
	`2`	`+`
	`3`	`+::: scoringrules.logs_normal`