Add type hints to all functions in metrics.py #446

neural_lam/metrics.py

@@ -1,8 +1,11 @@
+# Standard library
+from typing import Optional
+
 # Third-party
 import torch
 
 
-def get_metric(metric_name):
+def get_metric(metric_name: str):
     """
     Get a defined metric with given name
 
@@ -18,7 +21,12 @@ def get_metric(metric_name):
     return DEFINED_METRICS[metric_name_lower]
 
 
-def mask_and_reduce_metric(metric_entry_vals, mask, average_grid, sum_vars):
+def mask_and_reduce_metric(
+    metric_entry_vals: torch.Tensor,
+    mask: Optional[torch.Tensor],
+    average_grid: bool,
+    sum_vars: bool,
+) -> torch.Tensor:
     """
     Masks and (optionally) reduces entry-wise metric values
 
@@ -53,7 +61,14 @@ def mask_and_reduce_metric(metric_entry_vals, mask, average_grid, sum_vars):
     return metric_entry_vals
 
 
-def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
+def wmse(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    pred_std: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    average_grid: bool = True,
+    sum_vars: bool = True,
+) -> torch.Tensor:
     """
     Weighted Mean Squared Error
 
@@ -84,7 +99,14 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     )
 
 
-def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
+def mse(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    pred_std: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    average_grid: bool = True,
+    sum_vars: bool = True,
+) -> torch.Tensor:
     """
     (Unweighted) Mean Squared Error
 
@@ -108,7 +130,14 @@ def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     )
 
 
-def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
+def wmae(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    pred_std: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    average_grid: bool = True,
+    sum_vars: bool = True,
+) -> torch.Tensor:
     """
     Weighted Mean Absolute Error
 
@@ -139,7 +168,14 @@ def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     )
 
 
-def mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
+def mae(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    pred_std: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    average_grid: bool = True,
+    sum_vars: bool = True,
+) -> torch.Tensor:
     """
     (Unweighted) Mean Absolute Error
 
@@ -163,7 +199,14 @@ def mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     )
 
 
-def nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
+def nll(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    pred_std: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    average_grid: bool = True,
+    sum_vars: bool = True,
+) -> torch.Tensor:
     """
     Negative Log Likelihood loss, for isotropic Gaussian likelihood
 
@@ -191,8 +234,13 @@ def nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
 
 
 def crps_gauss(
-    pred, target, pred_std, mask=None, average_grid=True, sum_vars=True
-):
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    pred_std: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    average_grid: bool = True,
+    sum_vars: bool = True,
+) -> torch.Tensor:
     """
     (Negative) Continuous Ranked Probability Score (CRPS)
     Closed-form expression based on Gaussian predictive distribution