Adding ligand processing to dataset and encoders.

.gitignore

@@ -17,3 +17,5 @@ build/
 venv/
 notebooks/
 figures/
+
+STYLE_GUIDE.md

README.md

@@ -57,7 +57,8 @@ Each PDB should have `_final` suffix and contain:
 WaterFlow processes PDB files through several stages to create training-ready graph representations:
 
 **PDB Parsing**
-- Uses Biotite to extract protein atoms and water molecules (HOH residues)
+- Uses Biotite to extract protein atoms, water molecules (HOH residues), and ligands
+- "Ligand" means every non-protein, non-water heavy atom: small molecules, ions, cofactors, and nucleic acids. Included by default; disable with `--no-include_ligands`
 - Modified residues are retained during structure parsing and geometry preprocessing
 - When generating ESM embeddings, modified residues are mapped to encoder-compatible amino acid identities (e.g., MSE→M/MET, SEC→U/SEC)
 - Hydrogen atoms are excluded
@@ -70,7 +71,9 @@ WaterFlow processes PDB files through several stages to create training-ready gr
 - Mate atoms are stored separately for proper handling during training
 
 **Graph Representation**
-- Node types: `protein` (ASU + symmetry mates), `water` (ground truth)
+- Node types: `protein` (ASU + symmetry mates + ligands), `water` (ground truth)
+- ASU ligand atoms are appended after ASU and mate atoms and carry the boolean `is_ligand` mask plus `residue_index = -1` (they have no residue embedding, so residue pooling masks them out)
+- `is_ligand` marks **ASU ligands only**. Symmetry-mate generation is currently unfiltered, so mate nodes can include HETATM and water atoms that `is_ligand` does not mark — see `TODO(mates)` in `ProteinWaterDataset._preprocess_one`. Don't treat `is_ligand` as an exhaustive ligand selector
 - Edge types (defined in `src/constants.py`):
   - `('protein', 'pp', 'protein')`: protein-protein edges
   - `('protein', 'pw', 'water')`: protein to water
@@ -100,17 +103,25 @@ Preprocessed data is cached under `--processed_dir` in a three-layer architectur
 
 ```
 <processed_dir>/
-├── geometry/              # Graph structures (or geometry_mates/ when include_mates=True)
+├── geometry/              # Graph structures; see cache directory naming below
 │   └── <pdb_id>_final.pt
 │       - protein_pos: centered protein coordinates (N, 3)
 │       - protein_x: element one-hot encoding (N, 16)
 │       - protein_res_idx: residue indices for grouping
+│       - is_ligand: bool mask marking the appended ASU ligand atoms (N,)
 │       - water_pos, water_x: water coordinates and features
 │       - num_asu_protein: ASU atom count (mate boundary metadata)
 │       # Note: When include_mates=True, mate atoms are concatenated into
-│       # protein_pos/protein_x. Recover boundaries via:
-│       #   ASU atoms = protein_pos[:num_asu_protein]
-│       #   Mate atoms = protein_pos[num_asu_protein:]
+│       # protein_pos/protein_x, and ASU ligand atoms are appended after those.
+│       # Node order is [ASU protein | mates | ASU ligands]. Recover blocks via:
+│       #   ASU protein atoms = protein_pos[:num_asu_protein]
+│       #   ASU ligand atoms  = protein_pos[is_ligand]          # always last
+│       #   Mate atoms        = protein_pos[num_asu_protein:][~is_ligand[num_asu_protein:]]
+│       #
+│       # is_ligand marks ASU ligands ONLY -- it is not an exhaustive ligand
+│       # selector. The mate block is unfiltered (see TODO(mates) in
+│       # _preprocess_one), so mate atoms may include HETATM/ligand/water atoms
+│       # that are NOT marked by is_ligand.
 ├── esm/                   # ESM embeddings (per-residue)
 │   └── <pdb_id>_final.pt
 │       - residue_embeddings: ESM3 embeddings (N_res, embed_dim)
@@ -122,10 +133,27 @@ Preprocessed data is cached under `--processed_dir` in a three-layer architectur
         - atom37_coords: standard atom37 coordinates (N_res, 37, 3)
 ```
 
+**Cache Directory Naming:**
+
+The geometry cache directory name encodes the flags that change which nodes get cached, so
+configs that produce different graphs never share a directory:
+
+| `--include_mates` | `--include_ligands` | Directory |
+|---|---|---|
+| true | true (default) | `geometry_mates/` |
+| true | false | `geometry_mates_noligands/` |
+| false | true | `geometry/` |
+| false | false | `geometry_noligands/` |
+
+The base name comes from `--geometry_cache_name` (default `geometry`).
+
 **Cache Generation Notes:**
 - Geometry cache is generated automatically when `preprocess=True` (default)
 - ESM/SLAE caches require running the respective `generate_*_embeddings.py` scripts first
 - Preprocessing failures are logged to `<geometry_dir>/preprocessing_failures.log`
+- Geometry caches built before ligand support lack the `is_ligand` field and will fail to
+  load with a `KeyError`. Delete the geometry cache directory and let it regenerate — the
+  cached graphs are stale, not merely missing a field
 
 ## Environment Setup
 
@@ -238,6 +266,8 @@ To resume training from a checkpoint, you can load the model weights and optimiz
 | `--scheduler` | `cosine` | LR scheduler: `cosine`, `step`, or `none` |
 | `--warmup_steps` | `0` | Linear warmup steps |
 | `--processed_dir` | `~/flow_cache/` | Cache directory for preprocessed data |
+| `--include_mates` | `false` | Include symmetry mate atoms as protein nodes |
+| `--include_ligands` | `true` | Include ligand/ion/cofactor/nucleic acid heavy atoms as protein nodes. Negate with `--no-include_ligands` |
 | `--save_dir` | `../flow_checkpoints` | Directory to save checkpoints |
 | `--save_every` | `10` | Save checkpoint every N epochs |
 | `--eval_every` | `5` | Run evaluation every N epochs |

pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
   "e3nn",
   "esm",
   "biotite",
-  "pymol-open-source",
+  "pymol-open-source-whl>=3.1.0.4",
   "scipy",
   "pandas",
   "numpy",

scripts/generate_esm_embeddings.py

@@ -69,7 +69,7 @@ def compute_esm_embeddings(
     """
     try:
         # Load ground truth atoms using geometry cache parser in src/dataset.py
-        protein_atoms, _ = parse_asu_with_biotite(str(pdb_path))
+        protein_atoms, _, _ = parse_asu_with_biotite(str(pdb_path))
         if len(protein_atoms) == 0:
             raise ValueError(f"No protein atoms found in {pdb_path}")
 

scripts/generate_slae_embeddings.py

@@ -1,6 +1,10 @@
 """
 Precompute SLAE embeddings for protein structures and save to separate cache files.
 
+NOTE: This SLAE encoder is legacy and is NOT currently used. We primarily use the
+ESM encoder (see scripts/generate_esm_embeddings.py). This script is retained for
+reference/reproducibility only.
+
 This script:
 1. Reads a split file containing PDB entries
 2. For each entry, loads the PDB and converts to atom37 representation
@@ -400,7 +404,7 @@ def main() -> None:
 
             try:
                 # protein_atoms: biotite AtomArray with num_atoms atoms
-                protein_atoms, _ = parse_asu_with_biotite(str(pdb_path))
+                protein_atoms, _, _ = parse_asu_with_biotite(str(pdb_path))
                 # coords: (num_residues, 37, 3) - atom37 coordinates
                 # residue_type: (num_residues,) - residue type indices
                 # chains: (num_residues,) - chain IDs

scripts/train.py

@@ -104,6 +104,16 @@ def parse_args():
         action="store_true",
         help="Include symmetry mate atoms as protein nodes",
     )
+    p.add_argument(
+        "--include_ligands",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help=(
+            "Include ligand, ion, cofactor and nucleic acid heavy atoms as protein "
+            "nodes. Disabling appends '_noligands' to the geometry cache directory, "
+            "so the two configs cache separately."
+        ),
+    )
     p.add_argument(
         "--duplicate_single_sample",
         type=int,
@@ -365,6 +375,7 @@ def _build_dataset_config(args: argparse.Namespace) -> tuple[dict, dict, dict]:
         "base_pdb_dir": args.base_pdb_dir,
         "geometry_cache_name": args.geometry_cache_name,
         "include_mates": args.include_mates,
+        "include_ligands": args.include_ligands,
         "sample_cache_size": args.sample_cache_size,
         "cache_load_mmap": args.cache_load_mmap,
         **quality_kwargs,

src/constants.py

@@ -6,6 +6,13 @@
 # Node feature dimensions
 NODE_FEATURE_DIM = 16  # Default node scalar feature dimension
 
+# Native widths of the cached embeddings produced by scripts/generate_*_embeddings.py.
+# These are fixed by the upstream models, not tunable: ESM3-open emits 1536-wide
+# per-residue vectors, SLAE emits 128-wide per-atom vectors. Cached encoders take the
+# width as a required config key (embedding_dim); these are the values to pass.
+ESM_EMBEDDING_DIM = 1536
+SLAE_EMBEDDING_DIM = 128
+
 # RBF (Radial Basis Function) parameters
 NUM_RBF = 16  # Number of RBF basis functions
 RBF_CUTOFF = 8.0  # Distance cutoff in Angstroms for RBF encoding

src/dataset.py

@@ -48,16 +48,21 @@ def element_onehot(symbols: list[str]) -> Tensor:
 
 def parse_asu_with_biotite(
     path: str,
-) -> tuple[bts.AtomArray, bts.AtomArray]:
+) -> tuple[bts.AtomArray, bts.AtomArray, bts.AtomArray]:
     """
-    Parse PDB file and extract protein and water atoms.
+    Parse PDB file and extract protein, water, and ligand atoms.
 
     Args:
         path: Path to PDB file
 
     Returns:
-        Tuple of (protein_atoms, water_atoms) as biotite AtomArrays.
-        Hydrogen atoms are excluded.
+        Tuple of (protein_atoms, water_atoms, ligand_atoms) as biotite AtomArrays.
+        Hydrogen atoms are excluded. ligand_atoms contains every non-protein,
+        non-water heavy atom: small-molecule ligands, ions, cofactors, AND
+        non-amino-acid polymers such as nucleic acids (DNA/RNA). It is deliberately
+        NOT restricted to HETATM records -- nucleic acids are written as ATOM
+        records but are kept here as context (their surfaces, especially the
+        phosphate backbone, order nearby water).
 
     Notes:
         - model=1: Uses first model in PDB (standard for X-ray structures)
@@ -73,10 +78,15 @@ def parse_asu_with_biotite(
     protein_mask = bts.filter_amino_acids(atoms)
     water_mask = (atoms.res_name == "HOH") | (atoms.res_name == "WAT")
 
+    # "ligand" here is broad: every non-protein, non-water heavy atom.
+    # includes small-molecule ligands, ions, cofactors and even nucleic acids
+    ligand_mask = ~protein_mask & ~water_mask
+
     protein_atoms = atoms[protein_mask]
     water_atoms = atoms[water_mask]
+    ligand_atoms = atoms[ligand_mask]
 
-    return protein_atoms, water_atoms
+    return protein_atoms, water_atoms, ligand_atoms
 
 
 def get_crystal_contacts_pymol(
@@ -678,11 +688,12 @@ def __init__(
         self,
         pdb_list_file: str,
         processed_dir: str,
+        base_pdb_dir: str,
         encoder_type: str = "gvp",
-        base_pdb_dir: str = "/sb/wankowicz_lab/data/srivasv/pdb_redo_data",
         cutoff: float = 8.0,
         max_neighbors: int = 256,
         include_mates: bool = True,
+        include_ligands: bool = True,
         geometry_cache_name: str = "geometry",
         preprocess: bool = True,
         duplicate_single_sample: int = 1,
@@ -704,18 +715,25 @@ def __init__(
         Args:
             pdb_list_file: Text file with lines like "<pdb_id>_final"
             processed_dir: Cache root directory. Geometry caches are stored in
-                           {processed_dir}/{geometry_cache_name}[_mates] and embedding
-                           caches in {processed_dir}/{encoder_name}.
+                           {processed_dir}/{geometry_cache_name}[_mates][_noligands]
+                           and embedding caches in {processed_dir}/{encoder_name}.
+            base_pdb_dir: Base directory containing PDB subdirectories
             encoder_type: Encoder used downstream ('gvp', 'slae', or 'esm').
                           Embeddings are loaded only for the selected type.
-            base_pdb_dir: Base directory containing PDB subdirectories
             cutoff: Distance cutoff for PP edges and crystal contacts (Angstroms)
             max_neighbors: Maximum neighbors per node for radius graph construction.
             include_mates: If True, include symmetry mate atoms as protein nodes
-            geometry_cache_name: Base name for geometry cache directory. When
-                                 include_mates=True, "_mates" is appended automatically.
-                                 Default is "geometry", resulting in "geometry/" or
-                                 "geometry_mates/" subdirectories.
+            include_ligands: If True (default), include every non-protein,
+                             non-water heavy atom (small-molecule ligands, ions,
+                             cofactors, and nucleic acids) as protein-type nodes.
+                             They are appended after protein (and mate) atoms with a
+                             boolean is_ligand mask and residue_index = -1.
+            geometry_cache_name: Base name for geometry cache directory. Flags that
+                                 change the cached node set are appended to it:
+                                 "_mates" when include_mates=True, "_noligands" when
+                                 include_ligands=False. Default is "geometry", yielding
+                                 "geometry_mates/" for the default config or e.g.
+                                 "geometry_mates_noligands/" with ligands excluded.
             preprocess: If True, run preprocessing on missing cached files
             duplicate_single_sample: If dataset has 1 sample, duplicate it this many times
             Quality checks (always active):
@@ -749,8 +767,12 @@ def __init__(
             raise ValueError("max_neighbors must be >= 1")
 
         self.cache_dir = Path(processed_dir)
-        # Directory-based separation: geometry/ vs geometry_mates/
+        # Directory-based separation: geometry/ vs geometry_mates/. Both flags change
+        # the cached node set, so both are encoded in the directory name -- otherwise
+        # toggling one would silently reuse geometry built under the other setting.
         cache_suffix = "_mates" if include_mates else ""
+        if not include_ligands:
+            cache_suffix += "_noligands"
         self.geometry_dir = self.cache_dir / f"{geometry_cache_name}{cache_suffix}"
         self.base_pdb_dir = Path(base_pdb_dir)
         self.cutoff = cutoff
@@ -761,6 +783,7 @@ def __init__(
         else:
             self.embedding_dir = None
         self.include_mates = include_mates
+        self.include_ligands = include_ligands
         self.duplicate_single_sample = duplicate_single_sample
 
         self.max_com_dist = max_com_dist
@@ -900,7 +923,7 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
         """
         pdb_path = str(entry["pdb_path"])
 
-        protein_atoms, water_atoms = parse_asu_with_biotite(pdb_path)
+        protein_atoms, water_atoms, ligand_atoms = parse_asu_with_biotite(pdb_path)
 
         # check inter-chain interactions for multi-chain proteins
         chain_valid, chain_reason, _ = check_chain_interactions(
@@ -1043,6 +1066,15 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
             water_x = torch.zeros((0, len(ELEMENT_VOCAB) + 1), dtype=torch.float32)
 
         # process symmetry mate atoms
+        #
+        # TODO(mates): the mate atom set is unfiltered and inconsistent with the ASU
+        # path. get_crystal_contacts_pymol runs symexp over the whole object and
+        # selects "sym* and interface" with no polymer filter, so mate_atoms carries
+        # het atoms and waters as well as protein, and every one of them becomes a
+        # protein-type node below. Consequences: mate ligands are already included
+        # but never marked in is_ligand (unlike ASU ligands) and are not gated by
+        # include_ligands; mate waters -- symmetry images of the prediction target --
+        # enter as protein context. Fix in dev_crystal_mates.
         mate_coords = crystal_data["mate_coords"]
         if mate_coords.shape[0] > 0:
             mate_pos = torch.tensor(mate_coords, dtype=torch.float32) - center
@@ -1079,6 +1111,30 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
             final_protein_x = protein_x
             final_protein_res_idx = protein_res_idx
 
+        # Append ASU ligand atoms after protein (and mate) atoms when enabled.
+        # is_ligand mask marks which protein-type nodes are ligand atoms.
+        # Ligands always go last so num_asu_protein and mate counts are unaffected,
+        # preserving ESM/SLAE embedding alignment via _pad_atom_embeddings_for_mates.
+        # Only ASU ligands are handled here -- mate het atoms come in unfiltered via
+        # the mate block above, see TODO(mates) there.
+        if self.include_ligands and len(ligand_atoms) > 0:
+            ligand_pos = torch.tensor(ligand_atoms.coord, dtype=torch.float32) - center
+            ligand_elements = [str(e).upper() for e in ligand_atoms.element]
+            ligand_x = element_onehot(ligand_elements)
+            final_protein_pos = torch.cat([final_protein_pos, ligand_pos], dim=0)
+            final_protein_x = torch.cat([final_protein_x, ligand_x], dim=0)
+            # Ligand atoms get residue_index = -1 (sentinel; no residue embedding).
+            # The is_ligand mask identifies them; residue-pooling masks out these
+            # negative indices before any scatter (see GVPEncoder._pool_by_residue).
+            ligand_res_idx = torch.full((len(ligand_atoms),), -1, dtype=torch.long)
+            final_protein_res_idx = torch.cat(
+                [final_protein_res_idx, ligand_res_idx], dim=0
+            )
+            is_ligand = torch.zeros(final_protein_pos.size(0), dtype=torch.bool)
+            is_ligand[-len(ligand_atoms) :] = True
+        else:
+            is_ligand = torch.zeros(final_protein_pos.size(0), dtype=torch.bool)
+
         # Compute PP edges and features
         if final_protein_pos.size(0) > 0:
             pp_edge_index = radius_graph(
@@ -1109,6 +1165,7 @@ def _preprocess_one(self, entry: dict, cache_path: Path):
                 "protein_pos": final_protein_pos,
                 "protein_x": final_protein_x,
                 "protein_res_idx": final_protein_res_idx,
+                "is_ligand": is_ligand,
                 "water_pos": water_pos,
                 "water_x": water_x,
                 # PP topology and features (precomputed)
@@ -1212,6 +1269,7 @@ def __getitem__(self, idx: int) -> HeteroData:
         protein_pos = cached["protein_pos"]
         protein_x = cached["protein_x"]
         protein_res_idx = cached["protein_res_idx"]
+        is_ligand = cached["is_ligand"]
         pp_edge_index = cached["pp_edge_index"]
         pp_edge_unit_vectors = cached["pp_edge_unit_vectors"]
         pp_edge_rbf = cached["pp_edge_rbf"]
@@ -1233,6 +1291,7 @@ def __getitem__(self, idx: int) -> HeteroData:
         data["protein"].x = protein_x
         data["protein"].pos = protein_pos
         data["protein"].residue_index = protein_res_idx
+        data["protein"].is_ligand = is_ligand
         data["protein"].num_nodes = protein_pos.size(0)
         data["protein"].num_residues = num_residues
         data["protein"].num_protein_residues = num_protein_residues
@@ -1271,6 +1330,7 @@ def __getitem__(self, idx: int) -> HeteroData:
 def get_dataloader(
     pdb_list_file: str,
     processed_dir: str,
+    base_pdb_dir: str,
     batch_size: int = 8,
     shuffle: bool = True,
     num_workers: int = 8,
@@ -1287,6 +1347,7 @@ def get_dataloader(
         processed_dir: Cache root directory. Uses:
                       - {processed_dir}/geometry for geometry caches
                       - {processed_dir}/{encoder_name} for embedding caches
+        base_pdb_dir: Base directory containing PDB subdirectories
         encoder_type: Encoder used downstream ('gvp', 'slae', or 'esm').
                       Embeddings are loaded only for this type.
         batch_size: Number of graphs per batch
@@ -1307,7 +1368,10 @@ def get_dataloader(
         - Then batch_size works normally
     """
     dataset = ProteinWaterDataset(
-        pdb_list_file=pdb_list_file, processed_dir=processed_dir, **dataset_kwargs
+        pdb_list_file=pdb_list_file,
+        processed_dir=processed_dir,
+        base_pdb_dir=base_pdb_dir,
+        **dataset_kwargs,
     )
 
     loader = DataLoader(