"""Evaluation metrics shared across training and validation routines."""
# Standard library
from collections.abc import Callable
# Third-party
import torch
[docs]
def get_metric(metric_name: str) -> Callable[..., torch.Tensor]:
"""
Get a metric function by name.
Parameters
----------
metric_name : str
Name of the metric. Must be a key in ``DEFINED_METRICS``.
Returns
-------
callable
Function implementing the requested metric.
Raises
------
AssertionError
If ``metric_name`` (case-insensitive) is not a key in
``DEFINED_METRICS``.
"""
metric_name_lower = metric_name.lower()
assert metric_name_lower in DEFINED_METRICS, (
f"Unknown metric: {metric_name}"
)
return DEFINED_METRICS[metric_name_lower]
[docs]
def mask_and_reduce_metric(
metric_entry_vals: torch.Tensor,
mask: torch.Tensor | None,
average_grid: bool,
sum_vars: bool,
) -> torch.Tensor:
"""
Apply a boolean mask and optionally reduce a per-entry metric tensor.
Parameters
----------
metric_entry_vals : torch.Tensor
Shape ``(..., N, num_variables)``. Per-entry metric values. ``(...)``
denotes any number of broadcastable batch dimensions, ``N`` is
the number of grid nodes, and ``num_variables`` is the number of
variables in the gridded representation (e.g. state features).
mask : torch.Tensor or None
Shape ``(N,)``. Boolean mask selecting which grid nodes to
include. ``None`` means all nodes are used.
average_grid : bool
If True, average over the grid dimension ``N``.
sum_vars : bool
If True, sum over the variable dimension ``num_variables``.
Returns
-------
torch.Tensor
Reduced metric values. Shape is one of ``(...,)``,
``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
depending on ``average_grid`` and ``sum_vars``.
"""
# Only keep grid nodes in mask
if mask is not None:
metric_entry_vals = metric_entry_vals[
..., mask, :
] # (..., num_selected_nodes, num_variables)
# Optionally reduce last two dimensions
if average_grid: # Reduce grid first
metric_entry_vals = torch.mean(
metric_entry_vals, dim=-2
) # (..., num_variables)
if sum_vars: # Reduce vars second
metric_entry_vals = torch.sum(
metric_entry_vals, dim=-1
) # (..., num_grid_nodes) or (...,)
return metric_entry_vals
[docs]
def wmse(
pred: torch.Tensor,
target: torch.Tensor,
pred_std: torch.Tensor,
mask: torch.Tensor | None = None,
average_grid: bool = True,
sum_vars: bool = True,
) -> torch.Tensor:
"""
Weighted Mean Squared Error.
Parameters
----------
pred : torch.Tensor
Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes
any number of broadcastable batch dimensions, ``N`` is the number
of grid nodes, and ``num_variables`` is the number of state variables.
target : torch.Tensor
Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
``pred``.
pred_std : torch.Tensor
Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
standard deviation used as per-entry weight.
mask : torch.Tensor or None, optional
Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
nodes.
average_grid : bool, optional
If True, average over the grid dimension (default True).
sum_vars : bool, optional
If True, sum over the variable dimension (default True).
Returns
-------
torch.Tensor
Reduced metric values. Shape is one of ``(...,)``,
``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
depending on ``average_grid`` and ``sum_vars``.
"""
entry_mse = torch.nn.functional.mse_loss(
pred, target, reduction="none"
) # (..., num_grid_nodes, num_variables)
entry_mse_weighted = entry_mse / (
pred_std**2
) # (..., num_grid_nodes, num_variables)
return mask_and_reduce_metric(
entry_mse_weighted,
mask=mask,
average_grid=average_grid,
sum_vars=sum_vars,
)
[docs]
def mse(
pred: torch.Tensor,
target: torch.Tensor,
pred_std: torch.Tensor,
mask: torch.Tensor | None = None,
average_grid: bool = True,
sum_vars: bool = True,
) -> torch.Tensor:
"""
(Unweighted) Mean Squared Error.
Parameters
----------
pred : torch.Tensor
Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes
any number of broadcastable batch dimensions, ``N`` is the number
of grid nodes, and ``num_variables`` is the number of state variables.
target : torch.Tensor
Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
``pred``.
pred_std : torch.Tensor
Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
standard deviation (unused; ``pred_std`` is replaced by ones
internally).
mask : torch.Tensor or None, optional
Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
nodes.
average_grid : bool, optional
If True, average over the grid dimension (default True).
sum_vars : bool, optional
If True, sum over the variable dimension (default True).
Returns
-------
torch.Tensor
Reduced metric values. Shape is one of ``(...,)``,
``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
depending on ``average_grid`` and ``sum_vars``.
"""
# Replace pred_std with constant ones
return wmse(
pred, target, torch.ones_like(pred_std), mask, average_grid, sum_vars
)
[docs]
def wmae(
pred: torch.Tensor,
target: torch.Tensor,
pred_std: torch.Tensor,
mask: torch.Tensor | None = None,
average_grid: bool = True,
sum_vars: bool = True,
) -> torch.Tensor:
"""
Weighted Mean Absolute Error.
Parameters
----------
pred : torch.Tensor
Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes
any number of broadcastable batch dimensions, ``N`` is the number
of grid nodes, and ``num_variables`` is the number of state variables.
target : torch.Tensor
Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
``pred``.
pred_std : torch.Tensor
Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
standard deviation used as per-entry weight.
mask : torch.Tensor or None, optional
Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
nodes.
average_grid : bool, optional
If True, average over the grid dimension (default True).
sum_vars : bool, optional
If True, sum over the variable dimension (default True).
Returns
-------
torch.Tensor
Reduced metric values. Shape is one of ``(...,)``,
``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
depending on ``average_grid`` and ``sum_vars``.
"""
entry_mae = torch.nn.functional.l1_loss(
pred, target, reduction="none"
) # (..., num_grid_nodes, num_variables)
entry_mae_weighted = (
entry_mae / pred_std
) # (..., num_grid_nodes, num_variables)
return mask_and_reduce_metric(
entry_mae_weighted,
mask=mask,
average_grid=average_grid,
sum_vars=sum_vars,
)
[docs]
def mae(
pred: torch.Tensor,
target: torch.Tensor,
pred_std: torch.Tensor,
mask: torch.Tensor | None = None,
average_grid: bool = True,
sum_vars: bool = True,
) -> torch.Tensor:
"""
(Unweighted) Mean Absolute Error.
Parameters
----------
pred : torch.Tensor
Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes
any number of broadcastable batch dimensions, ``N`` is the number
of grid nodes, and ``num_variables`` is the number of state variables.
target : torch.Tensor
Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
``pred``.
pred_std : torch.Tensor
Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
standard deviation (unused; ``pred_std`` is replaced by ones
internally).
mask : torch.Tensor or None, optional
Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
nodes.
average_grid : bool, optional
If True, average over the grid dimension (default True).
sum_vars : bool, optional
If True, sum over the variable dimension (default True).
Returns
-------
torch.Tensor
Reduced metric values. Shape is one of ``(...,)``,
``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
depending on ``average_grid`` and ``sum_vars``.
"""
# Replace pred_std with constant ones
return wmae(
pred, target, torch.ones_like(pred_std), mask, average_grid, sum_vars
)
[docs]
def nll(
pred: torch.Tensor,
target: torch.Tensor,
pred_std: torch.Tensor,
mask: torch.Tensor | None = None,
average_grid: bool = True,
sum_vars: bool = True,
) -> torch.Tensor:
"""
Negative Log Likelihood loss for an isotropic Gaussian likelihood.
Parameters
----------
pred : torch.Tensor
Shape ``(..., N, num_variables)``. Predicted mean. ``(...)`` denotes
any number of broadcastable batch dimensions, ``N`` is the number
of grid nodes, and ``num_variables`` is the number of state variables.
target : torch.Tensor
Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
``pred``.
pred_std : torch.Tensor
Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
standard deviation of the Gaussian.
mask : torch.Tensor or None, optional
Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
nodes.
average_grid : bool, optional
If True, average over the grid dimension (default True).
sum_vars : bool, optional
If True, sum over the variable dimension (default True).
Returns
-------
torch.Tensor
Reduced metric values. Shape is one of ``(...,)``,
``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
depending on ``average_grid`` and ``sum_vars``.
"""
# Broadcast pred_std if shaped (num_variables,) via distribution internals
dist = torch.distributions.Normal(
pred, pred_std
) # (..., num_grid_nodes, num_variables)
entry_nll = -dist.log_prob(target) # (..., num_grid_nodes, num_variables)
return mask_and_reduce_metric(
entry_nll, mask=mask, average_grid=average_grid, sum_vars=sum_vars
)
[docs]
def crps_gauss(
pred: torch.Tensor,
target: torch.Tensor,
pred_std: torch.Tensor,
mask: torch.Tensor | None = None,
average_grid: bool = True,
sum_vars: bool = True,
) -> torch.Tensor:
"""
Continuous Ranked Probability Score (CRPS) for a Gaussian predictive
distribution (closed-form expression, negated for minimisation).
Parameters
----------
pred : torch.Tensor
Shape ``(..., N, num_variables)``. Predicted mean. ``(...)`` denotes
any number of broadcastable batch dimensions, ``N`` is the number
of grid nodes, and ``num_variables`` is the number of state variables.
target : torch.Tensor
Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
``pred``.
pred_std : torch.Tensor
Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
standard deviation of the Gaussian.
mask : torch.Tensor or None, optional
Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
nodes.
average_grid : bool, optional
If True, average over the grid dimension (default True).
sum_vars : bool, optional
If True, sum over the variable dimension (default True).
Returns
-------
torch.Tensor
Reduced metric values. Shape is one of ``(...,)``,
``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
depending on ``average_grid`` and ``sum_vars``.
"""
std_normal = torch.distributions.Normal(
torch.zeros((), device=pred.device), torch.ones((), device=pred.device)
)
target_standard = (
target - pred
) / pred_std # (..., num_grid_nodes, num_variables)
entry_crps = -pred_std * (
torch.pi ** (-0.5)
- 2 * torch.exp(std_normal.log_prob(target_standard))
- target_standard * (2 * std_normal.cdf(target_standard) - 1)
) # (..., num_grid_nodes, num_variables)
return mask_and_reduce_metric(
entry_crps, mask=mask, average_grid=average_grid, sum_vars=sum_vars
)
DEFINED_METRICS = {
"mse": mse,
"mae": mae,
"wmse": wmse,
"wmae": wmae,
"nll": nll,
"crps_gauss": crps_gauss,
}
