Type-check and unit-test the mpnn and rfd3na model packages

models/mpnn/src/mpnn/collate/feature_collator.py

@@ -43,7 +43,7 @@ class FeatureCollator:
     scalar values (requiring consistency checks across the batch).
     """
 
-    def __init__(self, default_padding: Dict[str, Any] = None):
+    def __init__(self, default_padding: Dict[str, Any] | None = None):
         """
         Initialize the FeatureCollator.
 
@@ -170,7 +170,7 @@ def __call__(self, pipeline_outputs: List[Dict[str, Any]]) -> Dict[str, Any]:
 
         return network_inputs
 
-    def _deep_equal(self, a, b):
+    def _deep_equal(self, a: Any, b: Any) -> bool:
         if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
             return torch.equal(a, b)
         if isinstance(a, dict) and isinstance(b, dict):
@@ -200,7 +200,9 @@ class TokenBudgetAwareFeatureCollator(FeatureCollator):
     """
 
     def __init__(
-        self, max_tokens_with_padding: int, default_padding: Dict[str, Any] = None
+        self,
+        max_tokens_with_padding: int,
+        default_padding: Dict[str, Any] | None = None,
     ):
         super().__init__(default_padding)
         self.max_tokens_with_padding = max_tokens_with_padding
@@ -237,7 +239,7 @@ def __call__(self, pipeline_outputs: List[Dict[str, Any]]) -> Dict[str, Any]:
         examples_with_L.sort(key=lambda x: x[0])
 
         # Apply token budget constraint by removing largest examples first.
-        filtered_examples = []
+        filtered_examples: List[tuple[int, Dict[str, Any]]] = []
         max_length = 0
         for L, example in examples_with_L:
             new_batch_size = len(filtered_examples) + 1

models/mpnn/src/mpnn/inference_engines/mpnn.py

@@ -109,7 +109,11 @@ def _validate_model_config(self) -> None:
             raise TypeError("checkpoint_path must be a string path.")
 
         # Check that the checkpoint path exists.
-        ckpt_path = Path(_absolute_path_or_none(self.checkpoint_path))
+        abs_checkpoint_path = _absolute_path_or_none(self.checkpoint_path)
+        # checkpoint_path is a non-empty str here (set in __init__), so the absolute
+        # form is never None.
+        assert abs_checkpoint_path is not None
+        ckpt_path = Path(abs_checkpoint_path)
         if not ckpt_path.is_file():
             raise FileNotFoundError(
                 f"checkpoint_path does not exist: {self.checkpoint_path}"
@@ -147,8 +151,11 @@ def _validate_all(self) -> None:
 
     def _post_process_engine_config(self) -> None:
         """Normalize paths into absolute paths."""
-        # Make checkpoint path absolute.
-        self.checkpoint_path = _absolute_path_or_none(self.checkpoint_path)
+        # Make checkpoint path absolute. checkpoint_path is a non-empty str (set in
+        # __init__), so the absolute form is never None and the attribute stays `str`.
+        abs_checkpoint_path = _absolute_path_or_none(self.checkpoint_path)
+        assert abs_checkpoint_path is not None
+        self.checkpoint_path = abs_checkpoint_path
 
         # Make output directory absolute.
         if self.out_directory is not None:
@@ -250,8 +257,13 @@ def run(
                     "'atom_arrays' and 'input_dicts' must have the same length."
                 )
 
-        # Determine the number of inputs.
-        num_inputs = len(input_dicts) if input_dicts is not None else len(atom_arrays)
+        # Determine the number of inputs. The guard above ensures at least one of
+        # input_dicts / atom_arrays is non-None; prefer input_dicts when present.
+        if input_dicts is not None:
+            num_inputs = len(input_dicts)
+        else:
+            assert atom_arrays is not None
+            num_inputs = len(atom_arrays)
         results: list[MPNNInferenceOutput] = []
         for input_idx in range(num_inputs):
             # Construct the per-input MPNNInferenceInput.

models/mpnn/src/mpnn/loss/nll_loss.py

@@ -1,9 +1,15 @@
+from typing import Any
+
 import torch
 import torch.nn as nn
 
 
 class LabelSmoothedNLLLoss(nn.Module):
-    def __init__(self, label_smoothing_eps=0.1, normalization_constant=6000.0):
+    def __init__(
+        self,
+        label_smoothing_eps: float = 0.1,
+        normalization_constant: float = 6000.0,
+    ) -> None:
         """
         Label smoothed negative log likelihood loss for Protein/Ligand MPNN.
 
@@ -20,7 +26,12 @@ def __init__(self, label_smoothing_eps=0.1, normalization_constant=6000.0):
         self.label_smoothing_eps = label_smoothing_eps
         self.normalization_constant = normalization_constant
 
-    def forward(self, network_input, network_output, loss_input):
+    def forward(
+        self,
+        network_input: dict[str, Any],
+        network_output: dict[str, Any],
+        loss_input: dict[str, Any],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
         """
         Given the network_input (same as input_features to the model), network
         output, and loss input, compute the loss.

models/mpnn/src/mpnn/metrics/nll.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import torch
 from atomworks.ml.transforms.base import ConvertToTorch
 from mpnn.collate.feature_collator import FeatureCollator
@@ -22,10 +24,10 @@ class NLL(Metric):
 
     def __init__(
         self,
-        return_per_example_metrics=False,
-        return_per_residue_metrics=False,
-        **kwargs,
-    ):
+        return_per_example_metrics: bool = False,
+        return_per_residue_metrics: bool = False,
+        **kwargs: Any,
+    ) -> None:
         """
         Initialize the NLL metric.
 
@@ -42,7 +44,7 @@ def __init__(
         self.return_per_residue_metrics = return_per_residue_metrics
 
     @property
-    def kwargs_to_compute_args(self):
+    def kwargs_to_compute_args(self) -> dict[str, Any]:
         """
         Map input keys to the compute method arguments.
 
@@ -56,7 +58,9 @@ def kwargs_to_compute_args(self):
             "mask_for_loss": ("network_output", "input_features", "mask_for_loss"),
         }
 
-    def get_per_residue_mask(self, mask_for_loss, **kwargs):
+    def get_per_residue_mask(
+        self, mask_for_loss: torch.Tensor, **kwargs: Any
+    ) -> torch.Tensor:
         """
         Get the per-residue mask for computing NLL.
 
@@ -73,7 +77,9 @@ def get_per_residue_mask(self, mask_for_loss, **kwargs):
         per_residue_mask = mask_for_loss
         return per_residue_mask
 
-    def compute_nll_metrics(self, S, log_probs, per_residue_mask):
+    def compute_nll_metrics(
+        self, S: torch.Tensor, log_probs: torch.Tensor, per_residue_mask: torch.Tensor
+    ) -> dict[str, torch.Tensor]:
         """
         Compute NLL and perplexity metrics using the provided per-residue mask.
         Args:
@@ -143,7 +149,15 @@ def compute_nll_metrics(self, S, log_probs, per_residue_mask):
         }
         return nll_dict
 
-    def compute(self, log_probs, S, mask_for_loss, **kwargs):
+    # MetricManager introspects the base ``Metric.compute(**kwargs)`` and dispatches by the
+    # names in ``kwargs_to_compute_args``, so declaring explicit params here is by design.
+    def compute(  # type: ignore[override]
+        self,
+        log_probs: torch.Tensor,
+        S: torch.Tensor,
+        mask_for_loss: torch.Tensor,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         """
         Compute the negative log likelihood (NLL) and perplexity, meaned
         across all residues that are included in the loss calculation.
@@ -182,7 +196,7 @@ def compute(self, log_probs, S, mask_for_loss, **kwargs):
         nll_metrics = self.compute_nll_metrics(S, log_probs, per_residue_mask)
 
         # Prepare the metric dictionary.
-        metric_dict = {
+        metric_dict: dict[str, Any] = {
             "mean_nll": nll_metrics["mean_nll"].detach().item(),
             "mean_perplexity": nll_metrics["mean_perplexity"].detach().item(),
         }
@@ -220,8 +234,8 @@ def __init__(
         interface_distance_threshold: float = 5.0,
         return_per_example_metrics: bool = False,
         return_per_residue_metrics: bool = False,
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> None:
         """
         Initialize the InterfaceNLL metric.
 
@@ -244,7 +258,7 @@ def __init__(
         self.interface_distance_threshold = interface_distance_threshold
 
     @property
-    def kwargs_to_compute_args(self):
+    def kwargs_to_compute_args(self) -> dict[str, Any]:
         """
         Map input keys to the compute method arguments.
 
@@ -257,7 +271,9 @@ def kwargs_to_compute_args(self):
         args_mapping["atom_array"] = ("network_input", "atom_array")
         return args_mapping
 
-    def get_per_residue_mask(self, mask_for_loss, **kwargs):
+    def get_per_residue_mask(
+        self, mask_for_loss: torch.Tensor, **kwargs: Any
+    ) -> torch.Tensor:
         """
         Get the per-residue mask for computing interface NLL.
 
@@ -334,7 +350,16 @@ def get_per_residue_mask(self, mask_for_loss, **kwargs):
 
         return combined_mask
 
-    def compute(self, log_probs, S, mask_for_loss, atom_array, **kwargs):
+    # MetricManager introspects the base ``Metric.compute(**kwargs)`` and dispatches by the
+    # names in ``kwargs_to_compute_args``, so declaring explicit params here is by design.
+    def compute(  # type: ignore[override]
+        self,
+        log_probs: torch.Tensor,
+        S: torch.Tensor,
+        mask_for_loss: torch.Tensor,
+        atom_array: Any,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         """
         Compute the interface negative log likelihood (NLL) and perplexity,
         averaged across interface residues only.
@@ -379,7 +404,7 @@ class SampledNLL(NLL):
     """
 
     @property
-    def kwargs_to_compute_args(self):
+    def kwargs_to_compute_args(self) -> dict[str, Any]:
         # Build the mapping explicitly: score the *sampled* sequence
         # (`S_sampled`) using the raw `logits`.
         return {
@@ -388,7 +413,13 @@ def kwargs_to_compute_args(self):
             "mask_for_loss": ("network_output", "input_features", "mask_for_loss"),
         }
 
-    def compute(self, logits, S, mask_for_loss, **kwargs):
+    def compute(  # type: ignore[override]
+        self,
+        logits: torch.Tensor,
+        S: torch.Tensor,
+        mask_for_loss: torch.Tensor,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         """Convert raw logits to log-probabilities and delegate to ``NLL``.
 
         Args:
@@ -417,7 +448,7 @@ class SampledLigandInterfaceNLL(InterfaceNLL):
     """
 
     @property
-    def kwargs_to_compute_args(self):
+    def kwargs_to_compute_args(self) -> dict[str, Any]:
         # Build the mapping explicitly: score the *sampled* sequence
         # (`S_sampled`) using the raw `logits`; `atom_array` derives the
         # polymer-ligand interface mask.
@@ -428,7 +459,14 @@ def kwargs_to_compute_args(self):
             "atom_array": ("network_input", "atom_array"),
         }
 
-    def compute(self, logits, S, mask_for_loss, atom_array, **kwargs):
+    def compute(  # type: ignore[override]
+        self,
+        logits: torch.Tensor,
+        S: torch.Tensor,
+        mask_for_loss: torch.Tensor,
+        atom_array: Any,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         """Convert raw logits to log-probabilities and delegate to ``InterfaceNLL``.
 
         Args:
@@ -463,7 +501,13 @@ class MPNNConfidence(SampledNLL):
     in ``(0, 1]``; higher means the model is more confident in the sequence.
     """
 
-    def compute(self, logits, S, mask_for_loss, **kwargs):
+    def compute(  # type: ignore[override]
+        self,
+        logits: torch.Tensor,
+        S: torch.Tensor,
+        mask_for_loss: torch.Tensor,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         """Compute the sampled-sequence NLL and derived confidence.
 
         Args:
@@ -501,7 +545,14 @@ class MPNNLigandInterfaceConfidence(SampledLigandInterfaceNLL):
     prefixed keys).
     """
 
-    def compute(self, logits, S, mask_for_loss, atom_array, **kwargs):
+    def compute(  # type: ignore[override]
+        self,
+        logits: torch.Tensor,
+        S: torch.Tensor,
+        mask_for_loss: torch.Tensor,
+        atom_array: Any,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         """Compute the interface NLL and derived interface confidence.
 
         Args: