Backend special module and statistical functions

scoringrules/backend/base.py

@@ -174,3 +174,59 @@ def apply_along_axis(
         self, func1d: tp.Callable[["Array"], "Array"], x: "Array", axis: int
     ) -> "Array":
         """Apply a function along a given axis of the input array."""
+
+    @abc.abstractmethod
+    def floor(self, x: "Array", /) -> "Array":
+        """Calculate the integer component of each element ``x_i`` of the input array ``x``."""
+
+    @abc.abstractmethod
+    def minimum(self, x: "Array", y: "ArrayLike", /) -> "Array":
+        """Calculate the minimum of each element ``x_i`` of the input array ``x`` with the value ``y``."""
+
+    @abc.abstractmethod
+    def maximum(self, x: "Array", y: "ArrayLike", /) -> "Array":
+        """Calculate the maximum of each element ``x_i`` of the input array ``x`` with the value ``y``."""
+
+    @abc.abstractmethod
+    def beta(self, x: "Array", y: "Array", /) -> "Array":
+        """Calculate the beta function at each element ``x_i`` of the input array ``x``."""
+
+    @abc.abstractmethod
+    def betainc(self, x: "Array", y: "Array", z: "Array", /) -> "Array":
+        """Calculate the regularised incomplete beta function at each element ``x_i`` of the input array ``x``."""
+
+    @abc.abstractmethod
+    def mbessel0(self, x: "Array", /) -> "Array":
+        """Calculate the modified Bessel function of the first kind of order 0 at each element ``x_i`` of the input array ``x``."""
+
+    @abc.abstractmethod
+    def mbessel1(self, x: "Array", /) -> "Array":
+        """Calculate the modified Bessel function of the first kind of order 1 at each element ``x_i`` of the input array ``x``."""
+
+    @abc.abstractmethod
+    def gamma(self, x: "Array", /) -> "Array":
+        """Calculate the gamma function at each element ``x_i`` of the input array ``x``."""
+
+    @abc.abstractmethod
+    def gammalinc(self, x: "Array", y: "Array", /) -> "Array":
+        """Calculate the lower incomplete gamma function at each element ``x_i`` of the input array ``x``."""
+
+    @abc.abstractmethod
+    def gammauinc(self, x: "Array", y: "Array", /) -> "Array":
+        """Calculate the upper incomplete gamma function at each element ``x_i`` of the input array ``x``."""
+
+    @abc.abstractmethod
+    def factorial(self, n: "ArrayLike", /) -> "ArrayLike":
+        """Calculate the factorial of the integer ``n``."""
+
+    @abc.abstractmethod
+    def hypergeometric(self, a: "Array", b: "Array", c: "Array", z: "Array") -> "Array":
+        """Calculate the hypergeometric function at each element of the inputs ``a``, ``b``, ``c``, and ``z``."""
+
+    @abc.abstractmethod
+    def comb(self, n: "ArrayLike", k: "ArrayLike", /) -> "ArrayLike":
+        """Calculate ``n`` choose ``k``."""
+
+    @abc.abstractmethod
+    def expi(self, x: "Array", /) -> "Array":
+        """Calculate the exponential integral at each element ``x_i`` of the input array ``x``."""

scoringrules/backend/jax.py

@@ -163,6 +163,48 @@ def apply_along_axis(
         except Exception:
             return jnp.apply_along_axis(func1d, axis, x)
 
+    def floor(self, x: "Array") -> "Array":
+        return jnp.floor(x)
+
+    def minimum(self, x: "Array", y: "ArrayLike") -> "Array":
+        return jnp.minimum(x, y)
+
+    def maximum(self, x: "Array", y: "ArrayLike") -> "Array":
+        return jnp.maximum(x, y)
+
+    def beta(self, x: "Array", y: "Array") -> "Array":
+        return jsp.special.beta(x, y)
+
+    def betainc(self, x: "Array", y: "Array", z: "Array") -> "Array":
+        return jsp.special.betainc(x, y, z)
+
+    def mbessel0(self, x: "Array") -> "Array":
+        return jsp.special.jv(0, x)
+
+    def mbessel1(self, x: "Array") -> "Array":
+        return jsp.special.jv(1, x)
+
+    def gamma(self, x: "Array") -> "Array":
+        return jsp.special.gamma(x)
+
+    def gammalinc(self, x: "Array", y: "Array") -> "Array":
+        return jsp.special.gammainc(x, y) * jsp.special.gamma(x)
+
+    def gammauinc(self, x: "Array", y: "Array") -> "Array":
+        return jsp.special.gammaincc(x, y) * jsp.special.gamma(x)
+
+    def factorial(self, n: "ArrayLike") -> "ArrayLike":
+        return jsp.special.factorial(n)
+
+    def hypergeometric(self, a: "Array", b: "Array", c: "Array", z: "Array"):
+        return jsp.special.hyp2f1(a, b, c, z)
+
+    def comb(self, n: "ArrayLike", k: "ArrayLike") -> "ArrayLike":
+        return jsp.special.comb(n, k)
+
+    def expi(self, x: "Array") -> "Array":
+        return jsp.special.expi(x)
+
 
 if __name__ == "__main__":
     B = JaxBackend()

scoringrules/backend/numpy.py

@@ -1,7 +1,19 @@
 import typing as tp
 
 import numpy as np
-from scipy.special import erf
+from scipy.special import (
+    beta,
+    betainc,
+    comb,
+    erf,
+    expi,
+    factorial,
+    gamma,
+    gammainc,
+    gammaincc,
+    hyp2f1,
+    jv,
+)
 
 if tp.TYPE_CHECKING:
     from numpy.typing import NDArray
@@ -150,6 +162,50 @@ def erf(self, x):
     def apply_along_axis(self, func1d: tp.Callable, x: "NDArray", axis: int):
         return np.apply_along_axis(func1d, axis, x)
 
+    def floor(self, x: "NDArray") -> "NDArray":
+        return np.floor(x)
+
+    def minimum(self, x: "NDArray", y: "ArrayLike") -> "NDArray":
+        return np.minimum(x, y)
+
+    def maximum(self, x: "NDArray", y: "ArrayLike") -> "NDArray":
+        return np.maximum(x, y)
+
+    def beta(self, x: "NDArray", y: "NDArray") -> "NDArray":
+        return beta(x, y)
+
+    def betainc(self, x: "NDArray", y: "NDArray", z: "NDArray") -> "NDArray":
+        return betainc(x, y, z)
+
+    def mbessel0(self, x: "NDArray") -> "NDArray":
+        return jv(0, x)
+
+    def mbessel1(self, x: "NDArray") -> "NDArray":
+        return jv(1, x)
+
+    def gamma(self, x: "NDArray") -> "NDArray":
+        return gamma(x)
+
+    def gammalinc(self, x: "NDArray", y: "NDArray") -> "NDArray":
+        return gammainc(x, y) * gamma(x)
+
+    def gammauinc(self, x: "NDArray", y: "NDArray") -> "NDArray":
+        return gammaincc(x, y) * gamma(x)
+
+    def factorial(self, n: "ArrayLike") -> "ArrayLike":
+        return factorial(n)
+
+    def hypergeometric(
+        self, a: "NDArray", b: "NDArray", c: "NDArray", z: "NDArray"
+    ) -> "NDArray":
+        return hyp2f1(a, b, c, z)
+
+    def comb(self, n: "ArrayLike", k: "ArrayLike") -> "ArrayLike":
+        return comb(n, k)
+
+    def expi(self, x: "NDArray") -> "NDArray":
+        return expi(x)
+
 
 class NumbaBackend(NumpyBackend):
     """Numba backend."""

scoringrules/backend/tensorflow.py

@@ -17,7 +17,7 @@
 
 
 class TensorflowBackend(ArrayBackend):
-    """Tensorflow backend"""
+    """Tensorflow backend."""
 
     name = "tensorflow"
 
@@ -201,6 +201,52 @@ def apply_along_axis(
                 [func1d(x_i) for x_i in tf.unstack(x, axis=axis)], axis=axis
             )
 
+    def floor(self, x: "Tensor") -> "Tensor":
+        return tf.math.floor(x)
+
+    def minimum(self, x: "Tensor", y: "TensorLike") -> "Tensor":
+        return tf.math.minimum(x, y)
+
+    def maximum(self, x: "Tensor", y: "TensorLike") -> "Tensor":
+        return tf.math.maximum(x, y)
+
+    def beta(self, x: "Tensor", y: "Tensor") -> "Tensor":
+        return tf.math.exp(
+            tf.math.lgamma(x) + tf.math.lgamma(y) - tf.math.lgamma(x + y)
+        )
+
+    def betainc(self, x: "Tensor", y: "Tensor", z: "Tensor") -> "Tensor":
+        return tf.math.betainc(x, y, z)
+
+    def mbessel0(self, x: "Tensor") -> "Tensor":
+        return tf.math.bessel_i0e(x)
+
+    def mbessel1(self, x: "Tensor") -> "Tensor":
+        return tf.math.bessel_i1e(x)
+
+    def gamma(self, x: "Tensor") -> "Tensor":
+        return tf.math.exp(tf.math.lgamma(x))
+
+    def gammalinc(self, x: "Tensor", y: "Tensor") -> "Tensor":
+        return tf.math.igamma(x, y) * tf.math.exp(tf.math.lgamma(x))
+
+    def gammauinc(self, x: "Tensor", y: "Tensor") -> "Tensor":
+        return tf.math.igammac(x, y) * tf.math.exp(tf.math.lgamma(x))
+
+    def factorial(self, n: "TensorLike") -> "TensorLike":
+        return tf.math.exp(tf.math.lgamma(n + 1))
+
+    def hypergeometric(
+        self, a: "Tensor", b: "Tensor", c: "Tensor", z: "Tensor"
+    ) -> "Tensor":
+        raise NotImplementedError
+
+    def comb(self, n: "Tensor", k: "Tensor") -> "Tensor":
+        return self.factorial(n) / (self.factorial(k) * self.factorial(n - k))
+
+    def expi(self, x: "Tensor") -> "Tensor":
+        return tf.math.special.expint(x)
+
 
 if __name__ == "__main__":
     B = TensorflowBackend()

scoringrules/backend/torch.py

@@ -176,6 +176,50 @@ def apply_along_axis(
                 [func1d(x_i) for x_i in torch.unbind(x, dim=axis)], dim=axis
             )
 
+    def floor(self, x: "Tensor") -> "Tensor":
+        return torch.floor(x)
+
+    def minimum(self, x: "Tensor", y: "TensorLike") -> "Tensor":
+        return torch.minimum(x, y)
+
+    def maximum(self, x: "Tensor", y: "TensorLike") -> "Tensor":
+        return torch.maximum(x, y)
+
+    def beta(self, x: "Tensor", y: "Tensor") -> "Tensor":
+        return torch.exp(torch.lgamma(x) + torch.lgamma(y) - torch.lgamma(x + y))
+
+    def betainc(self, x: "Tensor", y: "Tensor", z: "Tensor") -> "Tensor":
+        return None
+
+    def mbessel0(self, x: "Tensor") -> "Tensor":
+        return torch.special.i0(x)
+
+    def mbessel1(self, x: "Tensor") -> "Tensor":
+        return torch.special.i1(x)
+
+    def gamma(self, x: "Tensor") -> "Tensor":
+        return torch.exp(torch.lgamma(x))
+
+    def gammalinc(self, x: "Tensor", y: "Tensor") -> "Tensor":
+        return torch.special.gammainc(x, y) * torch.exp(torch.lgamma(x))
+
+    def gammauinc(self, x: "Tensor", y: "Tensor") -> "Tensor":
+        return torch.special.gammaincc(x, y) * torch.exp(torch.lgamma(x))
+
+    def factorial(self, n: "TensorLike") -> "TensorLike":
+        return torch.exp(torch.lgamma(n + 1))
+
+    def hypergeometric(
+        self, a: "Tensor", b: "Tensor", c: "Tensor", z: "Tensor"
+    ) -> "Tensor":
+        return None
+
+    def comb(self, n: "Tensor", k: "Tensor") -> "Tensor":
+        return self.factorial(n) / (self.factorial(k) * self.factorial(n - k))
+
+    def expi(self, x: "Tensor") -> "Tensor":
+        return None
+
 
 if __name__ == "__main__":
     B = TorchBackend()

scoringrules/core/error_spread/__init__.py

@@ -1,4 +1,4 @@
-from ._score import error_spread_score as ess
 from ._gufunc import _error_spread_score_gufunc as _ess_gufunc
+from ._score import error_spread_score as ess
 
-__all__ = ["ess"]
+__all__ = ["ess", "_ess_gufunc"]

scoringrules/core/error_spread/_gufunc.py

@@ -1,5 +1,3 @@
-import math
-
 import numpy as np
 from numba import guvectorize
 
@@ -30,15 +28,14 @@ def _error_spread_score_gufunc(
 
     for i in range(M):
         s += (forecasts[i] - m) ** 2
-    s /= (M - 1)
+    s /= M - 1
     s = np.sqrt(s)
 
     for i in range(M):
         g += ((forecasts[i] - m) / s) ** 3
-    g /= ((M - 1) * (M - 2))
+    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
-

scoringrules/core/error_spread/_score.py

@@ -13,9 +13,9 @@ def 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))
+    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)
-
-