fix bugs in variogram scores

Author: sallen12

scoringrules/_variogram.py

@@ -15,7 +15,7 @@ def vs_ensemble(
     m_axis: int = -2,
     v_axis: int = -1,
     *,
-    p: float = 1.0,
+    p: float = 0.5,
     estimator: str = "nrg",
     backend: "Backend" = None,
 ) -> "Array":

scoringrules/core/variogram/_score.py

@@ -23,7 +23,7 @@ def vs_ensemble(
     obs_diff = B.abs(B.expand_dims(obs, -2) - B.expand_dims(obs, -1)) ** p  # (... D D)
 
     if estimator == "nrg":
-        vfct = B.sum(B.abs(fct_diff) ** p, axis=-3) / M  # (... D D)
+        vfct = B.sum(fct_diff, axis=-3) / M  # (... D D)
         out = B.sum((obs_diff - vfct) ** 2, axis=(-2, -1))  # (...)
 
     elif estimator == "fair":
@@ -52,29 +52,20 @@ def owvs_ensemble(
 ) -> "Array":
     """Compute the Outcome-Weighted Variogram Score for a multivariate finite ensemble."""
     B = backends.active if backend is None else backends[backend]
-    M: int = fct.shape[-2]
-    wbar = B.mean(fw, axis=-1)
-
-    fct_diff = B.expand_dims(fct, -2) - B.expand_dims(fct, -1)  # (... M D D)
-    fct_diff = B.abs(fct_diff) ** p  # (... M D D)
+    M = fct.shape[-2]
+    wbar = B.sum(fw, -1) / M
 
-    obs_diff = B.expand_dims(obs, -2) - B.expand_dims(obs, -1)  # (... D D)
-    obs_diff = B.abs(obs_diff) ** p  # (... D D)
-    del obs, fct
-
-    E_1 = (fct_diff - B.expand_dims(obs_diff, -3)) ** 2  # (... M D D)
-    E_1 = B.sum(E_1, axis=(-2, -1))  # (... M)
-    E_1 = B.sum(E_1 * fw * B.expand_dims(ow, -1), axis=-1) / (M * wbar)  # (...)
+    fct_diff = (
+        B.abs(B.expand_dims(fct, -2) - B.expand_dims(fct, -1)) ** p
+    )  # (... M D D)
+    obs_diff = B.abs(B.expand_dims(obs, -2) - B.expand_dims(obs, -1)) ** p  # (... D D)
 
-    fct_diff_spread = B.expand_dims(fct_diff, -3) - B.expand_dims(
-        fct_diff, -4
-    )  # (... M M D D)
-    fw_prod = B.expand_dims(fw, -2) * B.expand_dims(fw, -1)  # (... M M)
-    E_2 = B.sum(fct_diff_spread**2, axis=(-2, -1))  # (... M M)
-    E_2 *= fw_prod * B.expand_dims(ow, (-2, -1))  # (... M M)
-    E_2 = B.sum(E_2, axis=(-2, -1)) / (M**2 * wbar**2)  # (...)
+    vfct = B.sum(fct_diff * B.expand_dims(fw, (-2, -1)), axis=-3) / (
+        M * B.expand_dims(wbar, (-2, -1))
+    )  # (... D D)
+    out = B.sum(((obs_diff - vfct) ** 2), axis=(-2, -1)) * ow  # (...)
 
-    return E_1 - 0.5 * E_2
+    return out
 
 
 def vrvs_ensemble(

tests/test_variogram.py

@@ -49,4 +49,4 @@ def test_variogram_score_correctness(backend):
     np.testing.assert_allclose(res, 0.05083489, rtol=1e-5)
 
     res = sr.vs_ensemble(obs, fct.T, p=1.0, backend=backend)
-    np.testing.assert_allclose(res, 0.04856365, rtol=1e-5)
+    np.testing.assert_allclose(res, 0.04856366, rtol=1e-5)

tests/test_wvariogram.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+
 import scoringrules as sr
 from scoringrules.backend import backends
 
@@ -23,9 +24,7 @@ def test_owvs_vs_vs(backend):
         lambda x: backends[backend].mean(x) * 0.0 + 1.0,
         backend=backend,
     )
-    np.testing.assert_allclose(
-        res, resw, rtol=1e-3
-    )  # TODO: not sure why tolerance must be so high
+    np.testing.assert_allclose(res, resw, rtol=1e-3)
 
 
 @pytest.mark.parametrize("backend", BACKENDS)