Adding CIF file support

README.md

@@ -35,29 +35,35 @@ WaterFlow/
 
 ## Data Preparation
 
-### Input PDB Files
+### Input Structure Files
 
-WaterFlow expects PDB files in a specific directory structure:
+WaterFlow reads PDB or mmCIF files, and expects them in a specific directory structure:
 
 ```
 <base_pdb_dir>/
 ├── 1abc/
-│   └── 1abc_final.pdb
+│   └── 1abc_final.cif      # .cif or .pdb
 ├── 2xyz/
 │   └── 2xyz_final.pdb
 └── ...
 ```
 
-Each PDB should have `_final` suffix and contain:
+Each structure should have the `_final` suffix and contain:
 - Protein atoms (used as conditioning context)
 - Water molecules (HOH residues, used as ground truth)
 
+**Format resolution:** entries in a split file are bare IDs (`6eey_final`) with no extension.
+For each entry WaterFlow looks in `<base_pdb_dir>/<pdb_id>/` and **prefers
+`<pdb_id>_final.cif` when it exists**, otherwise falls back to `<pdb_id>_final.pdb`. Both
+formats parse to identical atom counts, so the choice does not change the resulting graph.
+If neither file exists, reading the structure raises an error naming the missing path.
+
 ### Data Processing Pipeline
 
-WaterFlow processes PDB files through several stages to create training-ready graph representations:
+WaterFlow processes structure files through several stages to create training-ready graph representations:
 
-**PDB Parsing**
-- Uses Biotite to extract protein atoms, water molecules (HOH residues), and ligands
+**Structure Parsing**
+- Uses Biotite to extract protein atoms, water molecules (HOH residues), and ligands, dispatching on file extension (`.cif` via `CIFFile`, otherwise `PDBFile`)
 - "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)
@@ -316,7 +322,7 @@ EDIA measures how well an atom's position is supported by the experimental elect
 
 **Configuration:**
 - EDIA filtering is enabled by default 
-- The EDIA data lives in the `json` file of the format `<pdb_id>_final.json` in the same directory as the `pdb` file, and is obtained from PDB-REDO.
+- The EDIA data lives in the `json` file of the format `<pdb_id>_final.json` in the same directory as the structure file, and is obtained from PDB-REDO.
 - Use `--no_filter_by_edia` to explicitly disable EDIA filtering
 
 </details>

scripts/generate_esm_embeddings.py

@@ -43,11 +43,15 @@
 
 from src.constants import ONE_TO_THREE, THREE_TO_ONE
 from src.dataset import parse_asu_with_biotite
-from src.utils import normalize_ins_code, parse_split_file, setup_logging_for_tqdm
+from src.utils import (
+    normalize_ins_code,
+    parse_split_file,
+    setup_logging_for_tqdm,
+)
 
 
 def compute_esm_embeddings(
-    pdb_path: Path,
+    struc_path: Path,
     model: ESM3,
 ) -> dict | None:
     """
@@ -61,17 +65,17 @@ def compute_esm_embeddings(
     How ESM parses: (https://github.com/evolutionaryscale/esm/blob/main/esm/utils/structure/protein_chain.py)
 
     Args:
-        pdb_path: Path to PDB file
+        struc_path: Path to structure file (PDB/CIF)
         model: Loaded ESM3 model
 
     Returns:
         Dict with 'residue_embeddings', 'sequence', 'num_residues', or None on error
     """
     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(struc_path))
         if len(protein_atoms) == 0:
-            raise ValueError(f"No protein atoms found in {pdb_path}")
+            raise ValueError(f"No protein atoms found in {struc_path}")
 
         # Extract ground truth sequence before mutating the array
         key_to_resname = {}
@@ -115,7 +119,7 @@ def compute_esm_embeddings(
         protein = ESMProtein.from_protein_complex(complex_obj)
 
         if not protein.sequence or protein.sequence.replace("|", "") == "":
-            raise ValueError(f"ESM returned empty sequence for {pdb_path}")
+            raise ValueError(f"ESM returned empty sequence for {struc_path}")
 
         with torch.no_grad():
             protein_tensor = model.encode(protein)
@@ -140,7 +144,7 @@ def compute_esm_embeddings(
         # Validate: length mismatch means embeddings won't align with residues
         if len(esm_seq) != num_residues:
             raise ValueError(
-                f"Length mismatch after sanitization for {pdb_path}! "
+                f"Length mismatch after sanitization for {struc_path}! "
                 f"Biotite: {num_residues}, ESM: {len(esm_seq)}"
             )
 
@@ -151,7 +155,7 @@ def compute_esm_embeddings(
         }
 
     except Exception as e:
-        logger.error(f"Error computing embeddings for {pdb_path}: {e}")
+        logger.error(f"Error computing embeddings for {struc_path}: {e}")
         return None
 
 
@@ -239,16 +243,10 @@ def main() -> None:
     failures = []
 
     for entry in tqdm(entries, desc="Computing ESM embeddings"):
-        pdb_path = entry["pdb_path"]
         cache_key = entry["cache_key"]
         cache_path = esm_cache_dir / f"{cache_key}.pt"
 
-        if not pdb_path.exists():
-            logger.error(f"PDB file not found: {pdb_path}")
-            failures.append((cache_key, "PDB file not found"))
-            continue
-
-        result = compute_esm_embeddings(pdb_path, model)
+        result = compute_esm_embeddings(entry["struc_path"], model)
 
         if result is not None:
             result["pdb_id"] = entry["pdb_id"]

scripts/generate_slae_embeddings.py

@@ -394,17 +394,11 @@ def main() -> None:
         batch_info = []
 
         for entry in entry_batch:
-            pdb_path = entry["pdb_path"]
             cache_key = entry["cache_key"]
 
-            if not pdb_path.exists():
-                logger.error(f"PDB file not found: {pdb_path}")
-                failures.append((cache_key, "PDB file not found"))
-                continue
-
             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(entry["struc_path"]))
                 # coords: (num_residues, 37, 3) - atom37 coordinates
                 # residue_type: (num_residues,) - residue type indices
                 # chains: (num_residues,) - chain IDs