Bring rfd3 2-error-tier modules under strict mypy

Second models/rfd3 direction-(b) slice. Clears the five modules with two
untyped-def errors each (model.RFD3, model.cfg_utils, transforms.ncaa_transforms,
transforms.ppi_transforms, utils.io) with honest annotations verified against
bodies, call sites, and sibling conventions. Tightening strip_X's f_stripped
param to dict[str, torch.Tensor] surfaced a real arg-type at its sole caller in
model.inference_sampler; resolved with a parallel 'assert f_ref is not None' in
the CFG branch (which already asserts the sibling ref_initializer_outputs).
Annotation-only otherwise; 527 strict errors across 43 modules remain.

Co-authored-by: lyskov-ai <277346777+lyskov-ai@users.noreply.github.com>

Author: lyskov

models/rfd3/src/rfd3/model/RFD3.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any
 
 import hydra
 import torch
@@ -36,7 +37,7 @@ def __init__(
         token_initializer: DictConfig | dict,
         diffusion_module: DictConfig | dict,
         inference_sampler: DictConfig | dict,
-        **_,
+        **_: Any,
     ):
         super().__init__()
         # Check for chunked P_LL mode via environment variable
@@ -74,8 +75,8 @@ def forward(
         self,
         input: dict,
         coord_atom_lvl_to_be_noised: torch.Tensor | None = None,
-        n_cycle=None,
-        **_,
+        n_cycle: int | None = None,
+        **_: Any,
     ) -> dict:
         initializer_outputs = self.token_initializer(input["f"])
 

models/rfd3/src/rfd3/model/cfg_utils.py

@@ -2,9 +2,9 @@
 
 
 def strip_f(
-    f,
-    cfg_features,
-):
+    f: dict[str, torch.Tensor],
+    cfg_features: list[str],
+) -> dict[str, torch.Tensor]:
     """
     Strips conditioning features from 'f' for classifier-free guidance.
 
@@ -67,7 +67,7 @@ def strip_f(
     return f_stripped
 
 
-def strip_X(X_L, f_stripped):
+def strip_X(X_L: torch.Tensor, f_stripped: dict[str, torch.Tensor]) -> torch.Tensor:
     """
     Strips X_L unindexed atoms from X_L
 

models/rfd3/src/rfd3/model/inference_sampler.py

@@ -278,8 +278,10 @@ def sample_diffusion_like_af3(
             if self.use_classifier_free_guidance and (
                 self.cfg_t_max is None or c_t > self.cfg_t_max
             ):
-                # CFG mode requires the reference (unconditional) initializer outputs.
+                # CFG mode requires the reference (unconditional) features and
+                # initializer outputs; RFD3.forward provides both only when CFG is on.
                 assert ref_initializer_outputs is not None
+                assert f_ref is not None
                 X_noisy_L_stripped = strip_X(X_noisy_L, f_ref)
 
                 # unconditional forward pass

models/rfd3/src/rfd3/transforms/ncaa_transforms.py

@@ -111,7 +111,7 @@ class AddIsDAminoAcidFeat(Transform):
     Adds an annotation to the atom array indicating whether each residue is a D-amino acid.
     """
 
-    def check_input(self, data) -> None:
+    def check_input(self, data: dict) -> None:
         check_contains_keys(data, ["atom_array", "feats"])
 
     def forward(self, data: dict) -> dict:
@@ -144,7 +144,7 @@ def forward(self, data: dict) -> dict:
 
 
 class StrtoBoolforIsDAminoAcidFeature(Transform):
-    def forward(self, data):
+    def forward(self, data: dict) -> dict:
         atom_array = data["atom_array"]
         convert_existing_annotations_to_bool(
             atom_array, annotations=["is_d_amino_acid"]

models/rfd3/src/rfd3/transforms/ppi_transforms.py

@@ -41,7 +41,9 @@
 # Future hotspot subsampling schemes might want to avoid giving redundant information via (say) bonded atoms
 
 
-def get_hotspot_atoms(atom_array, binder_pn_unit_iid, distance_cutoff=4.5):
+def get_hotspot_atoms(
+    atom_array: AtomArray, binder_pn_unit_iid: str, distance_cutoff: float = 4.5
+) -> np.ndarray:
     """Get hotspot atoms for a given distance cutoff.
 
     Args:
@@ -335,7 +337,7 @@ def __init__(
         self.force_crop = force_crop
         self.max_atoms_in_crop = max_atoms_in_crop
 
-    def check_input(self, data: dict):
+    def check_input(self, data: dict) -> None:
         check_contains_keys(data, ["atom_array"])
         check_is_instance(data, "atom_array", AtomArray)
         check_atom_array_annotation(data, ["pn_unit_iid", "atomize", "atom_id"])

models/rfd3/src/rfd3/utils/io.py

@@ -6,7 +6,7 @@
 import numpy as np
 import torch
 from atomworks.io.utils.io_utils import to_cif_file
-from beartype.typing import Callable, Literal, cast
+from beartype.typing import Any, Callable, Literal, cast
 from biotite.structure import AtomArray, AtomArrayStack, stack
 
 from foundry.utils.alignment import weighted_rigid_align
@@ -20,7 +20,7 @@ def dump_structures(
     base_path: PathLike,
     one_model_per_file: bool,
     extra_fields: list[str] | Literal["all"] = [],
-    **kwargs,
+    **kwargs: Any,
 ) -> None:
     """Dump structures to CIF files, given the coordinates and input AtomArray.
 
@@ -64,7 +64,7 @@ def dump_metadata(
     prediction_metadata: dict,
     base_path: PathLike,
     one_model_per_file: bool,
-):
+) -> None:
     """
     Dump JSONs of prediction metadata to disk.
 

pyproject.toml

@@ -282,6 +282,11 @@ module = [
     "rfd3.model.layers.chunked_pairwise",
     "rfd3.trainer.rfd3",
     "rfd3.utils.vizualize",
+    "rfd3.model.RFD3",
+    "rfd3.model.cfg_utils",
+    "rfd3.transforms.ncaa_transforms",
+    "rfd3.transforms.ppi_transforms",
+    "rfd3.utils.io",
 ]
 disallow_untyped_defs = true
 check_untyped_defs = true