neural_lam.metrics

neural_lam.metrics#

Evaluation metrics shared across training and validation routines.

Module Contents#

neural_lam.metrics.crps_gauss(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True)#

Compute the (negative) Continuous Ranked Probability Score (CRPS).

A closed-form expression for a Gaussian predictive distribution is used.

Parameters:

  • pred (torch.Tensor) –

    Distribution mean predictions.

    • Shape: (..., N, d_state)

  • target (torch.Tensor) –

    Ground-truth values.

    • Shape: (..., N, d_state)

  • pred_std (torch.Tensor) –

    Predicted standard deviation parameter of the Gaussian.

    • Shape: (..., N, d_state) or (d_state,)

  • mask (torch.Tensor or None, optional) –

    Boolean mask selecting grid nodes. Default is None (all nodes).

    • Shape: (N,)

  • average_grid (bool, optional) – If True, average over the grid dimension. Default is True.

  • sum_vars (bool, optional) – If True, sum over the variable dimension. Default is True.

Returns:

Negative CRPS values with shape determined by average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.get_metric(metric_name)#

Retrieve a registered metric function by name.

Parameters:

metric_name (str) – Name of the metric to load (case-insensitive).

Returns:

Metric function implementing the requested metric.

Return type:

callable

Raises:

AssertionError – If metric_name is not part of DEFINED_METRICS.

neural_lam.metrics.mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True)#

Compute the unweighted Mean Absolute Error (MAE).

Parameters:

  • pred (torch.Tensor) –

    Model predictions.

    • Shape: (..., N, d_state)

  • target (torch.Tensor) –

    Ground-truth values.

    • Shape: (..., N, d_state)

  • pred_std (torch.Tensor) –

    Unused argument for compatibility with wmae().

    • Shape: (..., N, d_state) or (d_state,)

  • mask (torch.Tensor or None, optional) –

    Boolean mask selecting grid nodes. Default is None (all nodes).

    • Shape: (N,)

  • average_grid (bool, optional) – If True, average over the grid dimension. Default is True.

  • sum_vars (bool, optional) – If True, sum over the variable dimension. Default is True.

Returns:

MAE with shape determined by average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.mask_and_reduce_metric(metric_entry_vals, mask, average_grid, sum_vars)#

Apply a spatial mask and optionally reduce a per-entry metric tensor.

Parameters:

  • metric_entry_vals (torch.Tensor) –

    Entry-wise metric values.

    • Shape: (..., N, d_state) where ... are broadcastable leading dimensions.

  • mask (torch.Tensor or None) –

    Boolean mask selecting which grid nodes to include. Pass None to use all nodes.

    • Shape: (N,)

  • average_grid (bool) – If True, reduce the grid dimension N by taking the mean, producing (..., d_state).

  • sum_vars (bool) – If True, reduce the variable dimension d_state by summing, producing (..., N) or (...,) depending on average_grid.

Returns:

Reduced metric tensor.

  • Shape: one of (...,), (..., d_state), (..., N), or (..., N, d_state) depending on the reduction flags.

Return type:

torch.Tensor

neural_lam.metrics.mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True)#

Compute the unweighted Mean Squared Error (MSE).

Parameters:

  • pred (torch.Tensor) –

    Model predictions.

    • Shape: (..., N, d_state)

  • target (torch.Tensor) –

    Ground-truth values.

    • Shape: (..., N, d_state)

  • pred_std (torch.Tensor) –

    Unused argument for API parity with wmse().

    • Shape: (..., N, d_state) or (d_state,)

  • mask (torch.Tensor or None, optional) –

    Boolean mask selecting grid nodes. Default is None (all nodes).

    • Shape: (N,)

  • average_grid (bool, optional) – If True, average over the grid dimension. Default is True.

  • sum_vars (bool, optional) – If True, sum over the variable dimension. Default is True.

Returns:

MSE with shape determined by average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True)#

Compute the Negative Log Likelihood for an isotropic Gaussian likelihood.

Parameters:

  • pred (torch.Tensor) –

    Distribution mean predictions.

    • Shape: (..., N, d_state)

  • target (torch.Tensor) –

    Ground-truth values.

    • Shape: (..., N, d_state)

  • pred_std (torch.Tensor) –

    Predicted standard deviation parameter of the Gaussian.

    • Shape: (..., N, d_state) or (d_state,)

  • mask (torch.Tensor or None, optional) –

    Boolean mask selecting grid nodes. Default is None (all nodes).

    • Shape: (N,)

  • average_grid (bool, optional) – If True, average over the grid dimension. Default is True.

  • sum_vars (bool, optional) – If True, sum over the variable dimension. Default is True.

Returns:

Negative log-likelihood with shape determined by average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True)#

Compute the Weighted Mean Absolute Error (wMAE).

Parameters:

  • pred (torch.Tensor) –

    Model predictions.

    • Shape: (..., N, d_state)

  • target (torch.Tensor) –

    Ground-truth values.

    • Shape: (..., N, d_state)

  • pred_std (torch.Tensor) –

    Predicted standard deviation used as the per-entry weighting.

    • Shape: (..., N, d_state) or (d_state,)

  • mask (torch.Tensor or None, optional) –

    Boolean mask selecting grid nodes. Default is None (all nodes).

    • Shape: (N,)

  • average_grid (bool, optional) – If True, average over the grid dimension. Default is True.

  • sum_vars (bool, optional) – If True, sum over the variable dimension. Default is True.

Returns:

Weighted MAE with shape determined by average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True)#

Compute the Weighted Mean Squared Error (wMSE).

Scales the squared error at each grid node and variable by the inverse variance 1 / pred_std**2, then applies masking and reduction via mask_and_reduce_metric().

Parameters:

  • pred (torch.Tensor) –

    Model predictions.

    • Shape: (..., N, d_state)

  • target (torch.Tensor) –

    Ground-truth values.

    • Shape: (..., N, d_state)

  • pred_std (torch.Tensor) –

    Predicted standard deviation used as the per-entry weighting.

    • Shape: (..., N, d_state) or (d_state,)

  • mask (torch.Tensor or None, optional) –

    Boolean mask selecting grid nodes. Default is None (all nodes).

    • Shape: (N,)

  • average_grid (bool, optional) – If True, average over the grid dimension. Default is True.

  • sum_vars (bool, optional) – If True, sum over the variable dimension. Default is True.

Returns:

Weighted MSE after masking and reduction (see mask_and_reduce_metric()).

  • Shape: determined by average_grid and sum_vars.

Return type:

torch.Tensor

neural_lam.metrics.DEFINED_METRICS#
Contents