Upload generator.py

scripts/generator.py

@@ -5,131 +5,118 @@ import warnings
 import sys
 import shutil
 
-def run_broteinshake_generator(pdb_path, fixed_chains, variable_chains, num_seqs=20, temp=0.1):
+def sync_protein_metadata(jsonl_path, dict_path):
     """
-    Generalized first-half pipeline for protein redesign.
-    
-    Args:
-        pdb_path: Path to the target complex (e.g., 'data/3KAS.pdb').
-        fixed_chains: Chains to remain unchanged (e.g., 'A'). Empty for single-chain proteins.
-        variable_chains: Chains to be redesigned/repainted (e.g., 'B'). For single-chain, this is the only chain.
+    Automated metadata sanitization to prevent KeyError (e.g., 'seq_chain_C').
+    Prunes non-proteogenic chain IDs from the dictionary before the run.
     """
-    # 1. Setup project identifiers and directories
+    if not os.path.exists(jsonl_path) or not os.path.exists(dict_path):
+        return
+
+    # 1. Identify chains that actually have proteogenic sequence data
+    valid_chains_map = {}
+    with open(jsonl_path, 'r') as f:
+        for line in f:
+            entry = json.loads(line)
+            name = entry['name']
+            # Only keep chains that have a 'seq_chain_X' entry in the JSONL
+            valid = {k.split('_')[-1] for k in entry.keys() if k.startswith('seq_chain_')}
+            valid_chains_map[name] = valid
+
+    # 2. Clean the chain ID dictionary
+    with open(dict_path, 'r') as f:
+        chain_id_dict = json.load(f)
+
+    for pdb_name, configs in chain_id_dict.items():
+        if pdb_name in valid_chains_map:
+            valid = valid_chains_map[pdb_name]
+            # configs[0] = redesign list, configs[1] = fixed list
+            original_chains = set(configs[0] + configs[1])
+            
+            chain_id_dict[pdb_name] = [
+                [c for c in configs[0] if c in valid],
+                [c for c in configs[1] if c in valid]
+            ]
+            
+            # Diagnostic feedback
+            removed = original_chains - valid
+            if removed:
+                print(f"🧹 Sanitizer: Pruned non-protein chains from metadata: {removed}")
+
+    # 3. Overwrite with cleaned metadata for ProteinMPNN
+    with open(dict_path, 'w') as f:
+        json.dump(chain_id_dict, f)
+
+def run_broteinshake_generator(pdb_path, fixed_chains, variable_chains, num_seqs=20, temp=0.1):
+    # 1. Setup identifiers and directories
     pdb_name = os.path.basename(pdb_path).split('.')[0]
     output_dir = f"./generated/{pdb_name}" 
     os.makedirs(output_dir, exist_ok=True)
     
-    # Get the project root directory (where ProteinMPNN should be)
     script_dir = os.path.dirname(os.path.abspath(__file__))
     project_root = os.path.dirname(script_dir)
     proteinmpnn_dir = os.path.join(project_root, "ProteinMPNN")
     
-    # Clone ProteinMPNN if it doesn't exist (for HuggingFace Spaces deployment)
     if not os.path.exists(proteinmpnn_dir):
         print("ProteinMPNN not found, cloning repository...")
-        subprocess.run(
-            ["git", "clone", "https://github.com/dauparas/ProteinMPNN.git"],
-            cwd=project_root,
-            check=True,
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL
-        )
+        subprocess.run(["git", "clone", "https://github.com/dauparas/ProteinMPNN.git"], cwd=project_root, check=True)
     
     mpnn_script = os.path.join(proteinmpnn_dir, "protein_mpnn_run.py")
     
-    # 2. Check if single-chain protein (no fixed chains means single-chain)
+    # 2. Handle Single vs Multi-Chain Logic 
     if not fixed_chains or len(fixed_chains) == 0:
-        # Single-chain protein: use direct PDB path command
-        # For single-chain, variable_chains should be the only chain (e.g., "A")
         chain_to_design = variable_chains[0] if variable_chains else "A"
-        
         mpnn_cmd = (
-            f"python -W ignore {mpnn_script} "
-            f"--pdb_path {pdb_path} "
-            f"--pdb_path_chains {chain_to_design} "
-            f"--out_folder {output_dir} "
-            f"--num_seq_per_target {num_seqs} "
-            f"--sampling_temp {temp} "
-            f"--seed 42 "
-            f"--batch_size 1"
+            f"python -W ignore {mpnn_script} --pdb_path {pdb_path} --pdb_path_chains {chain_to_design} "
+            f"--out_folder {output_dir} --num_seq_per_target {num_seqs} --sampling_temp {temp} --seed 42 --batch_size 1"
         )
-        
-        print(f"🚀 Designing sequences for {pdb_name} (single-chain mode)...")
-        print(f"✏️ Redesigning chain: {chain_to_design}")
+        print(f"🚀 Designing {pdb_name} (Single-chain: {chain_to_design})...")
     else:
-        # Multi-chain protein: use JSONL-based command
-        # 2. Parse the PDB into JSONL format for the model
-        pdb_dir = os.path.dirname(os.path.abspath(pdb_path))
-        if not pdb_dir:
-            pdb_dir = "."
+        # Multi-chain setup
+        pdb_dir = os.path.dirname(os.path.abspath(pdb_path)) or "."
         jsonl_path = os.path.join(output_dir, "parsed_pdbs.jsonl")
-        
         parse_script = os.path.join(proteinmpnn_dir, "helper_scripts", "parse_multiple_chains.py")
         
-        parse_cmd = f"python -W ignore {parse_script} --input_path={pdb_dir}/ --output_path={jsonl_path}"
-        subprocess.run(parse_cmd, shell=True, check=True, stderr=subprocess.DEVNULL)
+        # Step A: Parse PDB to JSONL
+        subprocess.run(f"python -W ignore {parse_script} --input_path={pdb_dir}/ --output_path={jsonl_path}", shell=True, check=True)
 
-        # Update the name in parsed JSONL to include "_clones"
+        # Step B: Create initial Chain Dictionary
         pdb_name_clones = f"{pdb_name}_clones"
+        # Fix: ensure the name in JSONL matches the dict key
         with open(jsonl_path, 'r') as f:
             jsonl_data = json.loads(f.readline())
         jsonl_data['name'] = pdb_name_clones
         with open(jsonl_path, 'w') as f:
             f.write(json.dumps(jsonl_data) + '\n')
 
-        # 3. Generate the Chain Configuration JSONs (The 'Engine' Logic)
-        # Format: {"name": [masked_chains_list, visible_chains_list]}
-        # masked_chains = chains to redesign, visible_chains = chains to keep fixed
-        masked_chains_list = [c for c in variable_chains]
-        visible_chains_list = [c for c in fixed_chains]
-        chain_id_dict = {pdb_name_clones: [masked_chains_list, visible_chains_list]}
-        
         chain_id_json = os.path.join(output_dir, "chain_id_dict.json")
+        chain_id_dict = {pdb_name_clones: [[c for c in variable_chains], [c for c in fixed_chains]]}
         with open(chain_id_json, 'w') as f:
             json.dump(chain_id_dict, f)
 
-        # 4. Execute optimized ProteinMPNN command for multi-chain
+        # Step C: AUTOMATED CLEANING - Prunes ghost chains like 'C'
+        sync_protein_metadata(jsonl_path, chain_id_json)
+
+        # Step D: Final Execution Command
         mpnn_cmd = (
-            f"python -W ignore {mpnn_script} "
-            f"--jsonl_path {jsonl_path} "
-            f"--chain_id_jsonl {chain_id_json} "
-            f"--out_folder {output_dir} "
-            f"--num_seq_per_target {num_seqs} "
-            f"--sampling_temp {temp} "
-            f"--seed 42"
+            f"python -W ignore {mpnn_script} --jsonl_path {jsonl_path} --chain_id_jsonl {chain_id_json} "
+            f"--out_folder {output_dir} --num_seq_per_target {num_seqs} --sampling_temp {temp} --seed 42"
         )
-        
-        print(f"🚀 Designing sequences for {pdb_name}...")
-        print(f"🔒 Fixed: {fixed_chains} | ✏️ Redesigning: {variable_chains}")
+        print(f"🚀 Designing {pdb_name}... (Fixed: {fixed_chains} | Redesign: {variable_chains})")
     
-    # Suppress warnings by redirecting stderr
+    # 3. Execute with suppressed warnings
     env = os.environ.copy()
     env['PYTHONWARNINGS'] = 'ignore'
     subprocess.run(mpnn_cmd, shell=True, check=True, env=env, stderr=subprocess.DEVNULL)
     
-    # For single-chain proteins, ProteinMPNN saves as {pdb_name}.fa
-    # Rename it to {pdb_name}_clones.fa for consistency
-    if not fixed_chains or len(fixed_chains) == 0:
-        seqs_dir = os.path.join(output_dir, "seqs")
-        old_file = os.path.join(seqs_dir, f"{pdb_name}.fa")
-        new_file = os.path.join(seqs_dir, f"{pdb_name}_clones.fa")
-        if os.path.exists(old_file) and not os.path.exists(new_file):
-            os.rename(old_file, new_file)
-            print(f"📝 Renamed {pdb_name}.fa → {pdb_name}_clones.fa")
-    
-    print(f"✅ Success! Fold the top sequences at https://esmatlas.com/resources?action=fold")
+    print(f"✅ Success! Design complete for {pdb_name}.")
 
 if __name__ == "__main__":
-    import sys
     if len(sys.argv) < 4:
         print("Usage: python scripts/generator.py <pdb_path> <fixed_chains> <variable_chains> [num_seqs] [temp]")
-        print("Example: python scripts/generator.py data/3kas.pdb A B")
         sys.exit(1)
-    
-    pdb_path = sys.argv[1]
-    fixed_chains = sys.argv[2]
-    variable_chains = sys.argv[3]
-    num_seqs = int(sys.argv[4]) if len(sys.argv) > 4 else 20
-    temp = float(sys.argv[5]) if len(sys.argv) > 5 else 0.1
-    
-    run_broteinshake_generator(pdb_path, fixed_chains, variable_chains, num_seqs, temp)
+    run_broteinshake_generator(
+        sys.argv[1], sys.argv[2], sys.argv[3], 
+        int(sys.argv[4]) if len(sys.argv) > 4 else 20, 
+        float(sys.argv[5]) if len(sys.argv) > 5 else 0.1
+    )