neural_lam.metrics#

Evaluation metrics shared across training and validation routines.

Attributes#

DEFINED_METRICS

Functions#

`crps_gauss`(→ torch.Tensor)	Continuous Ranked Probability Score (CRPS) for a Gaussian predictive
`get_metric`(→ collections.abc.Callable[Ellipsis, ...)	Get a metric function by name.
`mae`(→ torch.Tensor)	(Unweighted) Mean Absolute Error.
`mask_and_reduce_metric`(→ torch.Tensor)	Apply a boolean mask and optionally reduce a per-entry metric tensor.
`mse`(→ torch.Tensor)	(Unweighted) Mean Squared Error.
`nll`(→ torch.Tensor)	Negative Log Likelihood loss for an isotropic Gaussian likelihood.
`wmae`(→ torch.Tensor)	Weighted Mean Absolute Error.
`wmse`(→ torch.Tensor)	Weighted Mean Squared Error.

Module Contents#

neural_lam.metrics.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[source]#

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:

Reduced metric values. Shape is one of (...,), (..., num_variables), (..., N), or (..., N, num_variables) depending on average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.get_metric(metric_name: str) → collections.abc.Callable[Ellipsis, torch.Tensor][source]#

Get a metric function by name.

Parameters:: metric_name (str) – Name of the metric. Must be a key in DEFINED_METRICS.
Returns:: Function implementing the requested metric.
Return type:: callable
Raises:: AssertionError – If metric_name (case-insensitive) is not a key in DEFINED_METRICS.

neural_lam.metrics.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[source]#

(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:

Reduced metric values. Shape is one of (...,), (..., num_variables), (..., N), or (..., N, num_variables) depending on average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.mask_and_reduce_metric(metric_entry_vals: torch.Tensor, mask: torch.Tensor | None, average_grid: bool, sum_vars: bool) → torch.Tensor[source]#

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:

Reduced metric values. Shape is one of (...,), (..., num_variables), (..., N), or (..., N, num_variables) depending on average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.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[source]#

(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:

Reduced metric values. Shape is one of (...,), (..., num_variables), (..., N), or (..., N, num_variables) depending on average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.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[source]#

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:

Reduced metric values. Shape is one of (...,), (..., num_variables), (..., N), or (..., N, num_variables) depending on average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.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[source]#

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:

Reduced metric values. Shape is one of (...,), (..., num_variables), (..., N), or (..., N, num_variables) depending on average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.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[source]#

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:

Reduced metric values. Shape is one of (...,), (..., num_variables), (..., N), or (..., N, num_variables) depending on average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.DEFINED_METRICS#