ESM Encoder Class Creation

.github/workflows/build.yml

@@ -1,10 +1,16 @@
 name: Build
 
 on:
-  # push:
-  #   branches: [main]
-  # pull_request:
-  #   branches: [main]
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+    paths-ignore:
+      - '**.md'
+      - 'docs/**'
+      - 'figures/**'
+      - '.gitignore'
+      - 'LICENSE'
   workflow_dispatch:
 
 jobs:

.github/workflows/lint.yml

@@ -1,8 +1,8 @@
 name: Lint
 
 on:
-  # pull_request:
-  #   branches: [main]
+  pull_request:
+    branches: [main]
   workflow_dispatch:
 
 jobs:
@@ -37,7 +37,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v6
-      
+
       - name: Setup UV and python version
         uses: astral-sh/setup-uv@v7
 

pyproject.toml

@@ -230,9 +230,21 @@ python_version = "3.12"
 warn_return_any = true
 warn_unused_configs = true
 
+[tool.ty.rules]
+# uvx runs ty in isolation without project dependencies - ignore missing imports
+unresolved-import = "ignore"
+# PyTorch Geometric has incomplete type stubs - dynamic attributes on Batch/Data
+unresolved-attribute = "ignore"
+# PyG's MessagePassing.message() override pattern is intentional
+invalid-method-override = "ignore"
+# PyG's propagate() uses **kwargs which ty doesn't understand
+missing-argument = "ignore"
+# Dict.get() returns union types that don't narrow well
+invalid-argument-type = "ignore"
+
 [tool.ruff.lint]
 fixable = ["I001", "F401", "UP"]
-ignore = ["E402", "E721", "E731", "E741", "F722", "F821", "UP015", "UP037"]
+ignore = ["E402", "E501", "E721", "E731", "E741", "F722", "F821", "UP015", "UP037"]
 select = ["E", "F", "I001", "UP"]
 
 [tool.ruff.lint.flake8-import-conventions.extend-aliases]

scripts/generate_esm_embeddings.py

@@ -54,8 +54,8 @@ def compute_esm_embeddings(
     Compute ESM3 residue-level embeddings using an in-memory sanitized structure.
 
     This bypasses ESM's default behavior of dropping HETATM non-canonicals.
-    It extracts the sequence, strips HETATM flags, renames modified residues 
-    (e.g., MSE -> MET) and unknowns (-> UNK), and feeds the buffer directly 
+    It extracts the sequence, strips HETATM flags, renames modified residues
+    (e.g., MSE -> MET) and unknowns (-> UNK), and feeds the buffer directly
     to ESM3 so that all residues receive a structural embedding.
 
     How ESM parses: (https://github.com/evolutionaryscale/esm/blob/main/esm/utils/structure/protein_chain.py)
@@ -89,7 +89,7 @@ def compute_esm_embeddings(
                 key_to_resname[key] = res_name_arr[i]
 
         unique_res_keys = list(key_to_resname.keys())
-        
+
         biotite_seq = [
             THREE_TO_ONE.get(key_to_resname[key], "X") for key in unique_res_keys
         ]
@@ -270,4 +270,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    main()

scripts/generate_slae_embeddings.py

@@ -27,6 +27,7 @@
         'pdb_id': str,
     }
 """
+
 from __future__ import annotations
 
 import argparse
@@ -193,8 +194,8 @@ def compute_slae_embeddings_batch(
     """
     Compute and align SLAE node embeddings from atom37 coords for batched structures.
 
-    Embeddings are aligned to the target geometry's atom order. Non-canonical atoms 
-    lacking SLAE representations are zero-padded. All input lists share the 
+    Embeddings are aligned to the target geometry's atom order. Non-canonical atoms
+    lacking SLAE representations are zero-padded. All input lists share the
     same length (batch size).
 
     Args:
@@ -217,7 +218,11 @@ def compute_slae_embeddings_batch(
     slae_atom_info_list = []
 
     for coords, residue_type, residue_id, chains, ins_code in zip(
-        coords_list, residue_type_list, residue_id_list, chains_list, ins_code_list,
+        coords_list,
+        residue_type_list,
+        residue_id_list,
+        chains_list,
+        ins_code_list,
         strict=True,
     ):
         # create PyG Data object with atom37 coords (featurizer will convert to flat)
@@ -430,13 +435,14 @@ def main() -> None:
         # compute embeddings for batch
         if batch_data:
             try:
-                (coords_list,
-                 residue_type_list,
-                 residue_id_list,
-                 chains_list,
-                 ins_code_list,
-                 geometry_atom_info_list
-                 ) = zip(*batch_data)
+                (
+                    coords_list,
+                    residue_type_list,
+                    residue_id_list,
+                    chains_list,
+                    ins_code_list,
+                    geometry_atom_info_list,
+                ) = zip(*batch_data)
 
                 embeddings_list = compute_slae_embeddings_batch(
                     list(coords_list),