chore(types): fix easy ty warnings shared across boltz/protenix/rf3 envs

Author: DorisMai

scripts/eval/classify_altloc_regions.py

@@ -415,7 +415,7 @@ def main(args: argparse.Namespace) -> None:
             )
         )
 
-    out_df = pd.DataFrame(all_rows, columns=OUTPUT_COLUMNS)
+    out_df = pd.DataFrame(all_rows, columns=pd.Index(OUTPUT_COLUMNS))
     args.output_file.parent.mkdir(parents=True, exist_ok=True)
     out_df.to_csv(args.output_file, index=False)
     logger.info(f"Wrote {len(out_df)} classified spans to {args.output_file}")

src/sampleworks/metrics/metric.py

@@ -26,8 +26,10 @@ def instantiate_metric_manager(
         A dictionary where keys are metric names and values are
         Hydra configurations for the metrics.
     """
-    metrics = {}
+    metrics: dict[str, Metric] = {}
     for name, cfg in metrics_cfg.items():
+        if not isinstance(name, str):
+            raise TypeError(f"metrics_cfg key must be a str, got {type(name).__name__}: {name!r}")
         metric = hydra.utils.instantiate(cfg)
         if not isinstance(metric, Metric):
             raise TypeError(f"{name} must be a Metric instance")

src/sampleworks/utils/guidance_script_utils.py

@@ -164,7 +164,7 @@ def save_losses(losses, output_dir):
 
 def get_model_and_device(
     device_str: str,
-    model_checkpoint_path: str,
+    model_checkpoint_path: str | None,
     model_type: str,
     method: str | None = None,
     model: Any = None,

src/sampleworks/utils/mmseqs2.py

@@ -297,6 +297,10 @@ def download(ID, path):
                     update_M = False
                     if M not in a3m_lines:
                         a3m_lines[M] = []
+                if M is None:
+                    # Well-formed A3M files always begin with a `>` header that sets M,
+                    # so this only triggers on malformed input. Skip rather than KeyError.
+                    continue
                 a3m_lines[M].append(line)
 
     # we run all the sequences together in one string, so this is really just a list of file dumps

tests/cli/test_guidance_cli.py

@@ -66,7 +66,7 @@ def test_model_specific_args_rf3_msa(self):
             "/data/msa.a3m",
         ] + COMMON_ARGS
         config = GuidanceConfig.from_cli(argv)
-        assert config.msa_path == "/data/msa.a3m"
+        assert config.msa_path == "/data/msa.a3m"  # ty: ignore[unresolved-attribute]
 
     def test_guidance_specific_args_fk(self):
         argv = [

tests/rewards/test_real_space_density_reward.py

@@ -476,8 +476,8 @@ def test_vmap_output_shape(self, reward_function_1vme, test_coordinates_1vme, de
 
             unique_combinations, inverse_indices = (
                 reward_function_1vme.precompute_unique_combinations(
-                    elements_batch[0, 0],  # ty: ignore[no-matching-overload, invalid-argument-type]
-                    b_factors_batch[0, 0],  # ty: ignore[no-matching-overload, invalid-argument-type]
+                    elements_batch[0, 0],
+                    b_factors_batch[0, 0],
                 )
             )
 
@@ -496,7 +496,7 @@ def test_vmap_output_shape(self, reward_function_1vme, test_coordinates_1vme, de
                 op=rf_partial,
             )
 
-            assert result.shape == torch.Size([num_particles])  # ty: ignore[unresolved-attribute]
+            assert result.shape == torch.Size([num_particles])
 
     def test_vmap_consistency(self, reward_function_1vme, test_coordinates_1vme, device):
         """Test vmap results match sequential calls."""
@@ -520,8 +520,8 @@ def test_vmap_consistency(self, reward_function_1vme, test_coordinates_1vme, dev
         occupancies_batch = einx.rearrange("n -> p e n", occupancies, p=num_particles, e=1)
 
         unique_combinations, inverse_indices = reward_function_1vme.precompute_unique_combinations(
-            elements_batch[0, 0],  # ty: ignore[no-matching-overload, invalid-argument-type]
-            b_factors_batch[0, 0],  # ty: ignore[no-matching-overload, invalid-argument-type]
+            elements_batch[0, 0],
+            b_factors_batch[0, 0],
         )
 
         rf_partial = partial(
@@ -549,7 +549,7 @@ def test_vmap_consistency(self, reward_function_1vme, test_coordinates_1vme, dev
             )
             result_sequential.append(loss.item())
 
-        result_sequential = torch.tensor(result_sequential, device=result_vmap.device)  # ty: ignore[unresolved-attribute]
+        result_sequential = torch.tensor(result_sequential, device=result_vmap.device)
 
         # GPU vmap and sequential loops accumulate floating-point reductions in
         # different orders, yielding abs diffs up to ~1.3e-4 and rel diffs up to

tests/utils/test_framework_utils.py

@@ -527,7 +527,7 @@ def test_raises_on_incompatible_batch_sizes(self):
     def test_raises_on_scalar_input(self):
         scalar = np.float64(1.0)
         with pytest.raises(ValueError, match="ndim >= 1"):
-            match_batch(scalar, target_batch_size=2)  # type: ignore[no-matching-overload]
+            match_batch(scalar, target_batch_size=2)  # ty: ignore[no-matching-overload]
 
     def test_1d_array(self):
         array = np.array([42.0])
@@ -549,4 +549,4 @@ def test_preserves_dtype(self):
 
     def test_raises_on_unsupported_type(self):
         with pytest.raises(TypeError, match="unsupported array type"):
-            match_batch([1, 2, 3], target_batch_size=2)  # type: ignore[no-matching-overload]
+            match_batch([1, 2, 3], target_batch_size=2)  # ty: ignore[no-matching-overload]