API Reference¶

def crps_abstract(accuracy: torch.Tensor, dispersion: torch.Tensor) -> torch.Tensor:
    """High-level function to compute the CRPS from the accuracy and dispersion terms.

    Args:
        accuracy: The accuracy term A, independent of the methods used to compute it, of shape (*batch_shape,).
        dispersion: The dispersion term D, independent of the methods used to compute it, of shape (*batch_shape,).

    Returns:
        The CRPS value for each forecast in the batch, of shape (*batch_shape,).
    """
    return accuracy - 0.5 * dispersion

def scrps_abstract(accuracy: torch.Tensor, dispersion: torch.Tensor) -> torch.Tensor:
    """High-level function to compute the SCRPS from the accuracy and dispersion terms.

    Args:
        accuracy: The accuracy term A, independent of the methods used to compute it, of shape (*batch_shape,).
        dispersion: The dispersion term D, independent of the methods used to compute it, of shape (*batch_shape,).

    Returns:
        The SCRPS value for each forecast in the batch, of shape (*batch_shape,).
    """
    return accuracy / dispersion + 0.5 * torch.log(dispersion)

def crps_analytical(
    q: Distribution,
    y: torch.Tensor,
) -> torch.Tensor:
    """Compute the (negatively-oriented, i.e., lower is better) CRPS in closed-form.

    Note:
        The input distribution must be either `torch.distributions.Normal` or `torch.distributions.StudentT`.
        There exists analytical solutions for other distributions, but they are not implemented, yet.
        Feel free to create an issue or pull request.

    Args:
        q: A PyTorch distribution object, typically a model's output distribution.
        y: Observed values, of shape (num_samples,).

    Returns:
        CRPS values for each observation, of shape (num_samples,).
    """
    if isinstance(q, Normal):
        return crps_analytical_normal(q, y)
    elif isinstance(q, StudentT):
        return crps_analytical_studentt(q, y)
    else:
        raise NotImplementedError(
            f"Detected distribution of type {type(q)}, but there are only analytical solutions for "
            "`torch.distributions.Normal` or `torch.distributions.StudentT`. Either use an alternative method, e.g. "
            "`torch_crps.crps_integral` or `torch_crps.crps_ensemble`, or create an issue for the method you need."
        )

def scrps_analytical(
    q: Distribution,
    y: torch.Tensor,
) -> torch.Tensor:
    """Compute the (negatively-oriented, i.e., lower is better) Scaled CRPS (SCRPS) in closed-form.

    Note:
        The input distribution must be either `torch.distributions.Normal` or `torch.distributions.StudentT`.
        There exists analytical solutions for other distributions, but they are not implemented, yet.
        Feel free to create an issue or pull request.

    Args:
        q: A PyTorch distribution object, typically a model's output distribution.
        y: Observed values, of shape (num_samples,).

    Returns:
        SCRPS values for each observation, of shape (num_samples,).
    """
    if isinstance(q, Normal):
        return scrps_analytical_normal(q, y)
    elif isinstance(q, StudentT):
        return scrps_analytical_studentt(q, y)
    else:
        raise NotImplementedError(
            f"Detected distribution of type {type(q)}, but there are only analytical solutions for "
            "`torch.distributions.Normal` or `torch.distributions.StudentT`. Either use an alternative method, e.g. "
            "`torch_crps.scrps_integral` or `torch_crps.scrps_ensemble`, or create an issue for the method you need."
        )

def crps_analytical_normal(
    q: Normal,
    y: torch.Tensor,
) -> torch.Tensor:
    """Compute the (negatively-oriented) CRPS in closed-form assuming a normal distribution.

    See Also:
        Gneiting & Raftery; "Strictly Proper Scoring Rules, Prediction, and Estimation"; 2007.
        Equation (5) for the analytical formula for CRPS of Normal distribution.

    Args:
        q: A PyTorch Normal distribution object, typically a model's output distribution.
        y: Observed values, of shape (num_samples,).

    Returns:
        CRPS values for each observation, of shape (num_samples,).
    """
    accuracy = _accuracy_normal(q, y)
    dispersion = _dispersion_normal(q)

    return crps_abstract(accuracy, dispersion)

def scrps_analytical_normal(
    q: Normal,
    y: torch.Tensor,
) -> torch.Tensor:
    r"""Compute the (negatively-oriented) Scaled CRPS (SCRPS) in closed-form assuming a normal distribution.

    $$
    \text{SCRPS}(F, y) = -\frac{E[|X - y|]}{E[|X - X'|]} - 0.5 \log \left( E[|X - X'|] \right)
                       = \frac{A}{D} + 0.5 \log(D)
    $$

    where $X$ and $X'$ are independent random variables drawn from the ensemble distribution, and $F(X)$ is the CDF
    of the ensemble distribution evaluated at $X$, and $y$ are the ground truth observations.

    Note:
        In contrast to the (negatively-oriented) CRPS, the SCRPS can have negative values.

    See Also:
        Bolin & Wallin; "Local scale invariance and robustness of proper scoring rules"; 2019.
        Equation (3) for the definition of the SCRPS.
        Appendix A.1 for the component formulas (Accuracy and Dispersion) for the Normal distribution

    Args:
        q: A PyTorch Normal distribution object, typically a model's output distribution.
        y: Observed values, of shape (num_samples,).

    Returns:
        SCRPS values for each observation, of shape (num_samples,).
    """
    accuracy = _accuracy_normal(q, y)
    dispersion = _dispersion_normal(q)

    return scrps_abstract(accuracy, dispersion)

def crps_analytical_studentt(
    q: StudentT,
    y: torch.Tensor,
) -> torch.Tensor:
    r"""Compute the (negatively-oriented) CRPS in closed-form assuming a StudentT distribution.

    This implements the closed-form formula from Jordan et al. (2019), see Appendix A.2.

    For the standardized StudentT distribution:

    $$
    \text{CRPS}(F_\nu, z) = z(2F_\nu(z) - 1) + 2f_\nu(z)\frac{\nu + z^2}{\nu - 1}
        - \frac{2\sqrt{\nu}}{\nu - 1} \frac{B(\frac{1}{2}, \nu - \frac{1}{2})}{B(\frac{1}{2}, \frac{\nu}{2})^2}
    $$

    where $z$ is the standardized value, $F_\nu$ is the CDF, $f_\nu$ is the PDF of the standard StudentT
    distribution, $\nu$ is the degrees of freedom, and $B$ is the beta function.

    For the location-scale transformed distribution:

    $$
    \text{CRPS}(F_{\nu,\mu,\sigma}, y) = \sigma \cdot \text{CRPS}\left(F_\nu, \frac{y-\mu}{\sigma}\right)
    $$

    where $\mu$ is the location parameter, $\sigma$ is the scale parameter, and $y$ is the observation.

    Note:
        This formula is only valid for degrees of freedom $\nu > 1$.

    See Also:
        Jordan et al.; "Evaluating Probabilistic Forecasts with scoringRules"; 2019.

    Args:
        q: A PyTorch StudentT distribution object, typically a model's output distribution.
        y: Observed values, of shape (num_samples,).

    Returns:
        CRPS values for each observation, of shape (num_samples,).
    """
    if torch.any(q.df <= 1):
        raise ValueError("StudentT SCRPS requires degrees of freedom > 1")

    accuracy = _accuracy_studentt(q, y)
    dispersion = _dispersion_studentt(q)

    return crps_abstract(accuracy, dispersion)

def scrps_analytical_studentt(
    q: StudentT,
    y: torch.Tensor,
) -> torch.Tensor:
    r"""Compute the (negatively-oriented) Scaled CRPS (SCRPS) in closed-form assuming a Student-T distribution.

    $$
    \text{SCRPS}(F, y) = -\frac{E[|X - y|]}{E[|X - X'|]} - 0.5 \log \left( E[|X - X'|] \right)
                       = \frac{A}{D} + 0.5 \log(D)
    $$

    where:

    - $F_{\nu, \mu, \sigma^2}$ is the cumulative Student-T distribution, and $F_{\nu}$ is the standardized version.
    - $A = E_F[|X - y|]$ is the accuracy term.
    - $A = \sigma [ z(2 F_{\nu}(z) - 1) +  2(\nu + z²) / (\nu*B(\nu/2, 1/2)) * F_{\nu+1}(z * \sqrt{(\nu+1)/(\nu+z²)}) ]$
    - $D = E_F[|X - X'|]$ is the dispersion term.
    - $D = \frac{ 4\sigma }{ \nu-1 } * ( \frac{ \Gamma( \nu/2 ) }{ \Gamma( (\nu-1)/2) } )^2$

    Note:
        This formula is only valid for degrees of freedom $\nu > 1$.

    See Also:
        Bolin & Wallin; "Local scale invariance and robustness of proper scoring rules"; 2019.

    Args:
        q: A PyTorch StudentT distribution object, typically a model's output distribution.
        y: Observed values, of shape (num_samples,).

    Returns:
        SCRPS values for each observation, of shape (num_samples,).
    """
    if torch.any(q.df <= 1):
        raise ValueError("StudentT SCRPS requires degrees of freedom > 1")

    accuracy = _accuracy_studentt(q, y)
    dispersion = _dispersion_studentt(q)

    return scrps_abstract(accuracy, dispersion)

def standardized_studentt_cdf_via_scipy(
    z: torch.Tensor,
    nu: torch.Tensor | float,
) -> torch.Tensor:
    """Since the `torch.distributions.StudentT` class does not have a `cdf()` method, we resort to scipy which has
    a stable implementation.

    Note:
        - The inputs `z` must be standardized.
        - This breaks differentiability and requires to move tensors to the CPU.

    Args:
        z: Standardized values at which to evaluate the CDF.
        nu: Degrees of freedom of the StudentT distribution.

    Returns:
        CDF values of the standardized StudentT distribution at `z`.
    """
    try:
        from scipy.stats import t as scipy_student_t
    except ImportError as e:
        raise ImportError(
            "scipy is required for the analytical solution for the StudentT distribution. "
            "Install `torch-crps` with the 'studentt' dependency group, e.g. `pip install torch-crps[studentt]`."
        ) from e

    z_np = z.detach().float().cpu().numpy()  # float() handles bfloat16
    nu_np = nu.detach().float().cpu().numpy() if isinstance(nu, torch.Tensor) else nu  # float() handles bfloat16

    cdf_z_np = scipy_student_t.cdf(x=z_np, df=nu_np)

    return torch.from_numpy(cdf_z_np).to(device=z.device, dtype=z.dtype)

def crps_ensemble(x: torch.Tensor, y: torch.Tensor, biased: bool = False) -> torch.Tensor:
    r"""Computes the Continuous Ranked Probability Score (CRPS) for an ensemble forecast.

    This function implements

    $$
    \text{CRPS}(F, y) = E[|X - y|] - 0.5 E[|X - X'|] = E[|X - y|] + E[X] - 2 E[X F(X)]
    $$

    where $X$ and $X'$ are independent random variables drawn from the ensemble distribution, and $F(X)$ is the CDF
    of the ensemble distribution evaluated at $X$.

    It is designed to be fully vectorized and handle any number of leading batch dimensions in the input tensors,
    as long as they are equal for `x` and `y`.

    See Also:
        Zamo & Naveau; "Estimation of the Continuous Ranked Probability Score with Limited Information and Applications
        to Ensemble Weather Forecasts"; 2017

    Note:
        - This implementation uses an efficient algorithm to compute the dispersion term E[|X - X'|] in O(m log(m))
        time, where m is the number of ensemble members. This is achieved by sorting the ensemble predictions and using
        a mathematical identity to compute the mean absolute difference. You can also see this trick
        [here][https://docs.nvidia.com/physicsnemo/25.11/_modules/physicsnemo/metrics/general/crps.html]

        - This implementation exactly matches the energy formula, see (NRG) and (eNRG), in Zamo & Naveau (2017) while
        using the compuational trick which can be read from (ePWM) in the same paper. The factors &\beta_0$ and
        $\beta_1$ in (ePWM) together equal the second term, i.e., the half mean dispersion, here. In (ePWM) they pulled
        the mean out. The energy formula and the probability weighted moment formula are equivalent.

    Args:
        x: The ensemble predictions, of shape (*batch_shape, dim_ensemble).
        y: The ground truth observations, of shape (*batch_shape).
        biased: If True, uses the biased estimator for the dispersion term $D$, i.e., divides by m². If False, uses the
            unbiased estimator which instead divides by m * (m - 1).

    Returns:
        The CRPS value for each forecast in the batch, of shape (*batch_shape).
    """
    if x.shape[:-1] != y.shape:
        raise ValueError(f"The batch dimension(s) of x {x.shape[:-1]} and y {y.shape} must be equal!")

    # Accuracy term A := E[|X - y|]
    accuracy = _accuracy_ensemble(x, y)

    # Dispersion term D := E[|X - X'|]
    dispersion = _dispersion_ensemble(x, biased)

    # CRPS value := A - 0.5 * D
    return crps_abstract(accuracy, dispersion)

def crps_ensemble_naive(x: torch.Tensor, y: torch.Tensor, biased: bool = False) -> torch.Tensor:
    """Computes the Continuous Ranked Probability Score (CRPS) for an ensemble forecast.

    This implementation uses the equality

    $$ CRPS(X, y) = E[|X - y|] - 0.5 E[|X - X'|] $$

    It is designed to be fully vectorized and handle any number of leading batch dimensions in the input tensors,
    as long as they are equal for `x` and `y`.

    See Also:
        Zamo & Naveau; "Estimation of the Continuous Ranked Probability Score with Limited Information and Applications
        to Ensemble Weather Forecasts"; 2017

    Note:
        - This implementation uses an inefficient algorithm to compute the term E[|X - X'|] in O(m²) where m is
        the number of ensemble members. This is done for clarity and educational purposes.
        - This implementation exactly matches the energy formula, see (NRG) and (eNRG), in Zamo & Naveau (2017).

    Args:
        x: The ensemble predictions, of shape (*batch_shape, dim_ensemble).
        y: The ground truth observations, of shape (*batch_shape).
        biased: If True, uses the biased estimator for $D$, i.e., divides by m². If False, uses the unbiased estimator.
            The unbiased estimator divides by m * (m - 1).

    Returns:
        The CRPS value for each forecast in the batch, of shape (*batch_shape).
    """
    if x.shape[:-1] != y.shape:
        raise ValueError(f"The batch dimension(s) of x {x.shape[:-1]} and y {y.shape} must be equal!")

    # Accuracy term A := E[|X - y|]
    accuracy = _accuracy_ensemble(x, y)

    # Dispersion term D := E[|X - X'|]
    dispersion = _dispersion_ensemble_naive(x, biased)

    # CRPS value := A - 0.5 * D
    return crps_abstract(accuracy, dispersion)

def scrps_ensemble(x: torch.Tensor, y: torch.Tensor, biased: bool = False) -> torch.Tensor:
    r"""Computes the Scaled Continuous Ranked Probability Score (SCRPS) for an ensemble forecast.

    $$
    \text{SCRPS}(F, y) = -\frac{E[|X - y|]}{E[|X - X'|]} - 0.5 \log \left( E[|X - X'|] \right)
                       = \frac{A}{D} + 0.5 \log(D)
    $$

    where $X$ and $X'$ are independent random variables drawn from the ensemble distribution, and $F(X)$ is the CDF
    of the ensemble distribution evaluated at $X$, and $y$ are the ground truth observations.

    It is designed to be fully vectorized and handle any number of leading batch dimensions in the input tensors,
    as long as they are equal for `x` and `y`.

    See Also:
        Bolin & Wallin; "Local scale invariance and robustness of proper scoring rules"; 2019.

    Note:
        This implementation uses an efficient algorithm to compute the dispersion term E[|X - X'|] in O(m log(m))
        time, where m is the number of ensemble members. This is achieved by sorting the ensemble predictions and using
        a mathematical identity to compute the mean absolute difference. You can also see this trick
        [here][https://docs.nvidia.com/physicsnemo/25.11/_modules/physicsnemo/metrics/general/crps.html]

    Args:
        x: The ensemble predictions, of shape (*batch_shape, dim_ensemble).
        y: The ground truth observations, of shape (*batch_shape).
        biased: If True, uses the biased estimator for the dispersion term $D$, i.e., divides by m². If False, uses the
            unbiased estimator which instead divides by m * (m - 1).

    Returns:
        The SCRPS value for each forecast in the batch, of shape (*batch_shape).
    """
    if x.shape[:-1] != y.shape:
        raise ValueError(f"The batch dimension(s) of x {x.shape[:-1]} and y {y.shape} must be equal!")

    # Accuracy term A := E[|X - y|]
    accuracy = _accuracy_ensemble(x, y)

    # Dispersion term D := E[|X - X'|]
    dispersion = _dispersion_ensemble(x, biased)

    # SCRPS value := A/D + 0.5 * log(D)
    return scrps_abstract(accuracy, dispersion)

def crps_integral(
    q: Distribution,
    y: torch.Tensor,
    x_min: float = -1e2,
    x_max: float = 1e2,
    x_steps: int = 5001,
) -> torch.Tensor:
    """Compute the Continuous Ranked Probability Score (CRPS) using the a (somewhat naive) integral approach.

    Note:
        This function is not differentiable with respect to `y` due to the indicator function.

    Args:
        q: A PyTorch distribution object, typically a model's output distribution. This object class must have a `cdf`
            method implemented.
        y: Observed values, of shape (num_samples,).
        x_min: Lower limit for integration for the probability space.
        x_max: Upper limit for integration for the probability space.
        x_steps: Number of steps for numerical integration.

    Returns:
        CRPS values for each observation, of shape (num_samples,).
    """

    def integrand(x: torch.Tensor) -> torch.Tensor:
        """Compute the integrand $F(x) - 1(y <= x))^2$ to be used by the torch integration function."""
        if not isinstance(q, StudentT):
            # Default case, try to access the distribution's CDF method.
            cdf_value = q.cdf(x)
        else:
            # Special case for torch's StudentT distributions which do not have a cdf method implemented.
            z = (x - q.loc) / q.scale
            cdf_value = standardized_studentt_cdf_via_scipy(z, q.df)
        indicator = (y_expanded <= x).float()
        return (cdf_value - indicator) ** 2

    # Set integration limits.
    x_values = torch.linspace(
        start=torch.tensor(x_min, dtype=y.dtype, device=y.device),
        end=torch.tensor(x_max, dtype=y.dtype, device=y.device),
        steps=x_steps,
        dtype=y.dtype,
        device=y.device,
    )

    # Reshape for proper broadcasting.
    x_values = x_values.unsqueeze(-1)  # shape: (x_steps, 1)
    y_expanded = y.unsqueeze(0)  # shape: (1, num_samples)

    # Compute the integral using the trapezoidal rule.
    integral_values = integrand(x_values)
    crps = torch.trapezoid(integral_values, x_values.squeeze(-1), dim=0)

    return crps

def normalize_by_observation(crps_fcn: Callable) -> Callable:
    """A decorator that normalizes the output of a CRPS function by the absolute maximum of the observations `y`.

    Note:
        - The resulting value is not guaranteed to be <= 1, because the (original) CRPS value can be larger than the
        normalization factor computed from the observations `y`.
        - If the observations `y` are all close to zero, then the normalization is done by 1, so the CRPS can be > 1.

    Args:
        crps_fcn: CRPS-calculating function to be wrapped. The function must accept an argument called y which is
            at the 2nd position.

    Returns:
        CRPS-calculating function which is wrapped such that the outputs are normalized by the magnitude of the
            observations.
    """

    @functools.wraps(crps_fcn)
    def wrapper(*args: WRAPPED_INPUT_TYPE, **kwargs: WRAPPED_INPUT_TYPE) -> torch.Tensor:
        """The function returned by the decorator that normalizes and forwards to the CRPS function."""
        # Find the observation 'y' from the arguments.
        if "y" in kwargs:
            y = kwargs["y"]
        elif len(args) < 2:
            raise TypeError("The observation `y` was not found in the function arguments as there is only one.")
        elif args:
            y = args[1]
        else:
            raise TypeError("The observation `y` was not found in the function arguments.")

        # Validate that y is a tenor.
        if not isinstance(y, torch.Tensor):
            raise TypeError("The observation `y` was found in the function arguments, but is not of type torch.Tensor!")

        # Calculate the normalization factor.
        abs_max_y = y.abs().max()
        if torch.isclose(abs_max_y, torch.zeros(1, device=abs_max_y.device, dtype=abs_max_y.dtype), atol=1e-6):
            # Avoid division by values close to zero.
            abs_max_y = torch.ones(1, device=abs_max_y.device, dtype=abs_max_y.dtype)

        # Call the original CRPS function.
        crps = crps_fcn(*args, **kwargs)

        # Normalize the result.
        return crps / abs_max_y

    return wrapper