Add option to have ensemble weights

scoringrules/_crps.py

@@ -4,6 +4,7 @@
 from scoringrules.core import crps, stats
 from scoringrules.core.utils import (
     univariate_array_check,
+    univariate_weight_check,
     estimator_check,
     uv_weighted_score_weights,
     uv_weighted_score_chain,
@@ -19,8 +20,9 @@ def crps_ensemble(
     fct: "Array",
     m_axis: int = -1,
     *,
-    sorted_ensemble: bool = False,
+    ens_w: "Array" = None,
     estimator: str = "qd",
+    sorted_ensemble: bool = False,
     backend: "Backend" = None,
 ) -> "Array":
     r"""Estimate the Continuous Ranked Probability Score (CRPS) for a finite ensemble.
@@ -56,11 +58,15 @@ def crps_ensemble(
         represented by the last axis.
     m_axis : int
         The axis corresponding to the ensemble. Default is the last axis.
+    ens_w : array, shape (..., m)
+        Weights assigned to the ensemble members. Array with the same shape as fct.
+        Default is equal weighting. Weights are normalised so that they sum to one
+        across the ensemble members.
+    estimator : str
+        Indicates the CRPS estimator to be used.
     sorted_ensemble : bool
         Boolean indicating whether the ensemble members are already in ascending order.
         Default is False.
-    estimator : str
-        Indicates the CRPS estimator to be used.
     backend : str, optional
         The name of the backend used for computations. Defaults to ``numba`` if available, else ``numpy``.
 
@@ -103,12 +109,21 @@ def crps_ensemble(
     array([0.69605316, 0.32865417, 0.39048665])
     """
     obs, fct = univariate_array_check(obs, fct, m_axis, backend=backend)
-    fct = univariate_sort_ens(fct, estimator, sorted_ensemble, backend=backend)
-    if backend == "numba":
-        estimator_check(estimator, crps.estimator_gufuncs)
-        return crps.estimator_gufuncs[estimator](obs, fct)
+    fct, ens_w = univariate_sort_ens(
+        fct, ens_w, m_axis, estimator, sorted_ensemble, backend=backend
+    )
+    if ens_w is None:
+        if backend == "numba":
+            estimator_check(estimator, crps.estimator_gufuncs)
+            return crps.estimator_gufuncs[estimator](obs, fct)
+        else:
+            return crps.ensemble(obs, fct, estimator, backend=backend)
     else:
-        return crps.ensemble(obs, fct, estimator, backend=backend)
+        if backend == "numba":
+            estimator_check(estimator, crps.estimator_gufuncs_w)
+            return crps.estimator_gufuncs_w[estimator](obs, fct, ens_w)
+        else:
+            return crps.ensemble_w(obs, fct, ens_w, estimator, backend=backend)
 
 
 def twcrps_ensemble(
@@ -118,6 +133,7 @@ def twcrps_ensemble(
     b: float = float("inf"),
     m_axis: int = -1,
     *,
+    ens_w: "Array" = None,
     v_func: tp.Callable[["ArrayLike"], "ArrayLike"] = None,
     estimator: str = "qd",
     sorted_ensemble: bool = False,
@@ -151,10 +167,19 @@ def twcrps_ensemble(
         to values in the range [a, b].
     m_axis : int
         The axis corresponding to the ensemble. Default is the last axis.
+    ens_w : array, shape (..., m)
+        Weights assigned to the ensemble members. Array with the same shape as fct.
+        Default is equal weighting. Weights are normalised so that they sum to one
+        across the ensemble members.
     v_func : callable, array_like -> array_like
         Chaining function used to emphasise particular outcomes. For example, a function that
         only considers values above a certain threshold :math:`t` by projecting forecasts and observations
         to :math:`[t, \inf)`.
+    estimator : str
+        Indicates the CRPS estimator to be used.
+    sorted_ensemble : bool
+        Boolean indicating whether the ensemble members are already in ascending order.
+        Default is False.
     backend : str, optional
         The name of the backend used for computations. Defaults to ``numba`` if available, else ``numpy``.
 
@@ -200,8 +225,9 @@ def twcrps_ensemble(
         obs,
         fct,
         m_axis=m_axis,
-        sorted_ensemble=sorted_ensemble,
+        ens_w=ens_w,
         estimator=estimator,
+        sorted_ensemble=sorted_ensemble,
         backend=backend,
     )
 
@@ -213,6 +239,7 @@ def owcrps_ensemble(
     b: float = float("inf"),
     m_axis: int = -1,
     *,
+    ens_w: "Array" = None,
     w_func: tp.Callable[["ArrayLike"], "ArrayLike"] = None,
     backend: "Backend" = None,
 ) -> "Array":
@@ -250,6 +277,10 @@ def owcrps_ensemble(
         to values in the range [a, b].
     m_axis : int
         The axis corresponding to the ensemble. Default is the last axis.
+    ens_w : array, shape (..., m)
+        Weights assigned to the ensemble members. Array with the same shape as fct.
+        Default is equal weighting. Weights are normalised so that they sum to one
+        across the ensemble members.
     w_func : callable, array_like -> array_like
         Weight function used to emphasise particular outcomes.
     backend : str, optional
@@ -288,10 +319,17 @@ def owcrps_ensemble(
     """
     obs, fct = univariate_array_check(obs, fct, m_axis, backend=backend)
     obs_w, fct_w = uv_weighted_score_weights(obs, fct, a, b, w_func, backend=backend)
-    if backend == "numba":
-        return crps.estimator_gufuncs["ownrg"](obs, fct, obs_w, fct_w)
+    if ens_w is None:
+        if backend == "numba":
+            return crps.estimator_gufuncs["ownrg"](obs, fct, obs_w, fct_w)
+        else:
+            return crps.ow_ensemble(obs, fct, obs_w, fct_w, backend=backend)
     else:
-        return crps.ow_ensemble(obs, fct, obs_w, fct_w, backend=backend)
+        ens_w = univariate_weight_check(ens_w, fct, m_axis, backend=backend)
+        if backend == "numba":
+            return crps.estimator_gufuncs_w["ownrg"](obs, fct, obs_w, fct_w, ens_w)
+        else:
+            return crps.ow_ensemble_w(obs, fct, obs_w, fct_w, ens_w, backend=backend)
 
 
 def vrcrps_ensemble(
@@ -301,6 +339,7 @@ def vrcrps_ensemble(
     b: float = float("inf"),
     m_axis: int = -1,
     *,
+    ens_w: "Array" = None,
     w_func: tp.Callable[["ArrayLike"], "ArrayLike"] = None,
     backend: "Backend" = None,
 ) -> "Array":
@@ -336,6 +375,10 @@ def vrcrps_ensemble(
         to values in the range [a, b].
     m_axis : int
         The axis corresponding to the ensemble. Default is the last axis.
+    ens_w : array, shape (..., m)
+        Weights assigned to the ensemble members. Array with the same shape as fct.
+        Default is equal weighting. Weights are normalised so that they sum to one
+        across the ensemble members.
     w_func : callable, array_like -> array_like
         Weight function used to emphasise particular outcomes.
     backend : str, optional
@@ -374,10 +417,17 @@ def vrcrps_ensemble(
     """
     obs, fct = univariate_array_check(obs, fct, m_axis, backend=backend)
     obs_w, fct_w = uv_weighted_score_weights(obs, fct, a, b, w_func, backend=backend)
-    if backend == "numba":
-        return crps.estimator_gufuncs["vrnrg"](obs, fct, obs_w, fct_w)
+    if ens_w is None:
+        if backend == "numba":
+            return crps.estimator_gufuncs["vrnrg"](obs, fct, obs_w, fct_w)
+        else:
+            return crps.vr_ensemble(obs, fct, obs_w, fct_w, backend=backend)
     else:
-        return crps.vr_ensemble(obs, fct, obs_w, fct_w, backend=backend)
+        ens_w = univariate_weight_check(ens_w, fct, m_axis, backend=backend)
+        if backend == "numba":
+            return crps.estimator_gufuncs_w["vrnrg"](obs, fct, obs_w, fct_w, ens_w)
+        else:
+            return crps.vr_ensemble_w(obs, fct, obs_w, fct_w, ens_w, backend=backend)
 
 
 def crps_quantile(

scoringrules/_energy.py

@@ -3,6 +3,7 @@
 from scoringrules.core import energy
 from scoringrules.core.utils import (
     multivariate_array_check,
+    multivariate_weight_check,
     estimator_check,
     mv_weighted_score_chain,
     mv_weighted_score_weights,
@@ -18,6 +19,7 @@ def es_ensemble(
     m_axis: int = -2,
     v_axis: int = -1,
     *,
+    ens_w: "Array" = None,
     estimator: str = "nrg",
     backend: "Backend" = None,
 ) -> "Array":
@@ -43,6 +45,9 @@ def es_ensemble(
     v_axis : int
         The axis corresponding to the variables dimension on the forecasts array (or the observations
         array with an extra dimension on `m_axis`). Defaults to -1.
+    ens_w : array_like
+        Weights assigned to the ensemble members. Array with one less dimension than fct (without the v_axis dimension).
+        Default is equal weighting. Weights are normalised so that they sum to one across the ensemble members.
     estimator : str
         The energy score estimator to be used.
     backend : str
@@ -66,11 +71,19 @@ def es_ensemble(
         Some theoretical background on scoring rules for multivariate forecasts.
     """
     obs, fct = multivariate_array_check(obs, fct, m_axis, v_axis, backend=backend)
-    if backend == "numba":
-        estimator_check(estimator, energy.estimator_gufuncs)
-        return energy.estimator_gufuncs[estimator](obs, fct)
+    if ens_w is None:
+        if backend == "numba":
+            estimator_check(estimator, energy.estimator_gufuncs)
+            return energy.estimator_gufuncs[estimator](obs, fct)
+        else:
+            return energy.es(obs, fct, estimator=estimator, backend=backend)
     else:
-        return energy.es(obs, fct, estimator=estimator, backend=backend)
+        ens_w = multivariate_weight_check(ens_w, fct, m_axis, backend=backend)
+        if backend == "numba":
+            estimator_check(estimator, energy.estimator_gufuncs_w)
+            return energy.estimator_gufuncs_w[estimator](obs, fct, ens_w)
+        else:
+            return energy.es_w(obs, fct, ens_w, estimator=estimator, backend=backend)
 
 
 def twes_ensemble(
@@ -80,6 +93,7 @@ def twes_ensemble(
     m_axis: int = -2,
     v_axis: int = -1,
     *,
+    ens_w: "Array" = None,
     estimator: str = "nrg",
     backend: "Backend" = None,
 ) -> "Array":
@@ -110,6 +124,9 @@ def twes_ensemble(
         The axis corresponding to the ensemble dimension. Defaults to -2.
     v_axis : int or tuple of int
         The axis corresponding to the variables dimension. Defaults to -1.
+    ens_w : array_like
+        Weights assigned to the ensemble members. Array with one less dimension than fct (without the v_axis dimension).
+        Default is equal weighting. Weights are normalised so that they sum to one across the ensemble members.
     estimator : str
         The energy score estimator to be used.
     backend : str
@@ -122,7 +139,13 @@ def twes_ensemble(
     """
     obs, fct = mv_weighted_score_chain(obs, fct, v_func)
     return es_ensemble(
-        obs, fct, m_axis=m_axis, v_axis=v_axis, estimator=estimator, backend=backend
+        obs,
+        fct,
+        m_axis=m_axis,
+        v_axis=v_axis,
+        ens_w=ens_w,
+        estimator=estimator,
+        backend=backend,
     )
 
 
@@ -133,6 +156,7 @@ def owes_ensemble(
     m_axis: int = -2,
     v_axis: int = -1,
     *,
+    ens_w: "Array" = None,
     backend: "Backend" = None,
 ) -> "Array":
     r"""Compute the Outcome-Weighted Energy Score (owES) for a finite multivariate ensemble.
@@ -165,6 +189,9 @@ def owes_ensemble(
         The axis corresponding to the ensemble dimension. Defaults to -2.
     v_axis : int or tuple of ints
         The axis corresponding to the variables dimension. Defaults to -1.
+    ens_w : array_like
+        Weights assigned to the ensemble members. Array with one less dimension than fct (without the v_axis dimension).
+        Default is equal weighting. Weights are normalised so that they sum to one across the ensemble members.
     backend : str
         The name of the backend used for computations. Defaults to 'numba' if available, else 'numpy'.
 
@@ -175,17 +202,25 @@ def owes_ensemble(
     """
     obs, fct = multivariate_array_check(obs, fct, m_axis, v_axis, backend=backend)
     obs_w, fct_w = mv_weighted_score_weights(obs, fct, w_func=w_func, backend=backend)
-    if backend == "numba":
-        return energy.estimator_gufuncs["ownrg"](obs, fct, obs_w, fct_w)
+    if ens_w is None:
+        if backend == "numba":
+            return energy.estimator_gufuncs["ownrg"](obs, fct, obs_w, fct_w)
+        else:
+            return energy.owes(obs, fct, obs_w, fct_w, backend=backend)
     else:
-        return energy.owes(obs, fct, obs_w, fct_w, backend=backend)
+        ens_w = multivariate_weight_check(ens_w, fct, m_axis, backend=backend)
+        if backend == "numba":
+            return energy.estimator_gufuncs_w["ownrg"](obs, fct, obs_w, fct_w, ens_w)
+        else:
+            return energy.owes_w(obs, fct, obs_w, fct_w, ens_w=ens_w, backend=backend)
 
 
 def vres_ensemble(
     obs: "Array",
     fct: "Array",
     w_func: tp.Callable[["ArrayLike"], "ArrayLike"],
     *,
+    ens_w: "Array" = None,
     m_axis: int = -2,
     v_axis: int = -1,
     backend: "Backend" = None,
@@ -221,6 +256,9 @@ def vres_ensemble(
         The axis corresponding to the ensemble dimension. Defaults to -2.
     v_axis : int or tuple of int
         The axis corresponding to the variables dimension. Defaults to -1.
+    ens_w : array_like
+        Weights assigned to the ensemble members. Array with one less dimension than fct (without the v_axis dimension).
+        Default is equal weighting. Weights are normalised so that they sum to one across the ensemble members.
     backend : str
         The name of the backend used for computations. Defaults to 'numba' if available, else 'numpy'.
 
@@ -231,7 +269,14 @@ def vres_ensemble(
     """
     obs, fct = multivariate_array_check(obs, fct, m_axis, v_axis, backend=backend)
     obs_w, fct_w = mv_weighted_score_weights(obs, fct, w_func=w_func, backend=backend)
-    if backend == "numba":
-        return energy.estimator_gufuncs["vrnrg"](obs, fct, obs_w, fct_w)
+    if ens_w is None:
+        if backend == "numba":
+            return energy.estimator_gufuncs["vrnrg"](obs, fct, obs_w, fct_w)
+        else:
+            return energy.vres(obs, fct, obs_w, fct_w, backend=backend)
     else:
-        return energy.vres(obs, fct, obs_w, fct_w, backend=backend)
+        ens_w = multivariate_weight_check(ens_w, fct, m_axis, backend=backend)
+        if backend == "numba":
+            return energy.estimator_gufuncs_w["vrnrg"](obs, fct, obs_w, fct_w, ens_w)
+        else:
+            return energy.vres_w(obs, fct, obs_w, fct_w, ens_w=ens_w, backend=backend)