neural_lam.metrics ================== .. py:module:: neural_lam.metrics .. autoapi-nested-parse:: Evaluation metrics shared across training and validation routines. Attributes ---------- .. autoapisummary:: neural_lam.metrics.DEFINED_METRICS Functions --------- .. autoapisummary:: neural_lam.metrics.crps_gauss neural_lam.metrics.get_metric neural_lam.metrics.mae neural_lam.metrics.mask_and_reduce_metric neural_lam.metrics.mse neural_lam.metrics.nll neural_lam.metrics.wmae neural_lam.metrics.wmse Module Contents --------------- .. py:function:: 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). :param pred: 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. :type pred: torch.Tensor :param target: Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. :type target: torch.Tensor :param pred_std: Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation of the Gaussian. :type pred_std: torch.Tensor :param mask: Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. :type mask: torch.Tensor or None, optional :param average_grid: If True, average over the grid dimension (default True). :type average_grid: bool, optional :param sum_vars: If True, sum over the variable dimension (default True). :type sum_vars: bool, optional :returns: Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. :rtype: torch.Tensor .. py:function:: get_metric(metric_name: str) -> collections.abc.Callable[Ellipsis, torch.Tensor] Get a metric function by name. :param metric_name: Name of the metric. Must be a key in ``DEFINED_METRICS``. :type metric_name: str :returns: Function implementing the requested metric. :rtype: callable :raises AssertionError: If ``metric_name`` (case-insensitive) is not a key in ``DEFINED_METRICS``. .. py:function:: 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. :param pred: 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. :type pred: torch.Tensor :param target: Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. :type target: torch.Tensor :param pred_std: Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation (unused; ``pred_std`` is replaced by ones internally). :type pred_std: torch.Tensor :param mask: Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. :type mask: torch.Tensor or None, optional :param average_grid: If True, average over the grid dimension (default True). :type average_grid: bool, optional :param sum_vars: If True, sum over the variable dimension (default True). :type sum_vars: bool, optional :returns: Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. :rtype: torch.Tensor .. py:function:: 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. :param metric_entry_vals: 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). :type metric_entry_vals: torch.Tensor :param mask: Shape ``(N,)``. Boolean mask selecting which grid nodes to include. ``None`` means all nodes are used. :type mask: torch.Tensor or None :param average_grid: If True, average over the grid dimension ``N``. :type average_grid: bool :param sum_vars: If True, sum over the variable dimension ``num_variables``. :type sum_vars: bool :returns: Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. :rtype: torch.Tensor .. py:function:: 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. :param pred: 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. :type pred: torch.Tensor :param target: Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. :type target: torch.Tensor :param pred_std: Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation (unused; ``pred_std`` is replaced by ones internally). :type pred_std: torch.Tensor :param mask: Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. :type mask: torch.Tensor or None, optional :param average_grid: If True, average over the grid dimension (default True). :type average_grid: bool, optional :param sum_vars: If True, sum over the variable dimension (default True). :type sum_vars: bool, optional :returns: Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. :rtype: torch.Tensor .. py:function:: 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. :param pred: 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. :type pred: torch.Tensor :param target: Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. :type target: torch.Tensor :param pred_std: Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation of the Gaussian. :type pred_std: torch.Tensor :param mask: Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. :type mask: torch.Tensor or None, optional :param average_grid: If True, average over the grid dimension (default True). :type average_grid: bool, optional :param sum_vars: If True, sum over the variable dimension (default True). :type sum_vars: bool, optional :returns: Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. :rtype: torch.Tensor .. py:function:: 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. :param pred: 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. :type pred: torch.Tensor :param target: Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. :type target: torch.Tensor :param pred_std: Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation used as per-entry weight. :type pred_std: torch.Tensor :param mask: Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. :type mask: torch.Tensor or None, optional :param average_grid: If True, average over the grid dimension (default True). :type average_grid: bool, optional :param sum_vars: If True, sum over the variable dimension (default True). :type sum_vars: bool, optional :returns: Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. :rtype: torch.Tensor .. py:function:: 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. :param pred: 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. :type pred: torch.Tensor :param target: Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as ``pred``. :type target: torch.Tensor :param pred_std: Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted standard deviation used as per-entry weight. :type pred_std: torch.Tensor :param mask: Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all nodes. :type mask: torch.Tensor or None, optional :param average_grid: If True, average over the grid dimension (default True). :type average_grid: bool, optional :param sum_vars: If True, sum over the variable dimension (default True). :type sum_vars: bool, optional :returns: Reduced metric values. Shape is one of ``(...,)``, ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)`` depending on ``average_grid`` and ``sum_vars``. :rtype: torch.Tensor .. py:data:: DEFINED_METRICS