Upload refine.py

scripts/refine.py

@@ -1,78 +1,48 @@
 import os
-from Bio.PDB import PDBParser, Superimposer, PDBIO # type: ignore
-from pdbfixer import PDBFixer # type: ignore
-from openmm.app import PDBFile, ForceField, Simulation, Modeller # type: ignore
-from openmm import LangevinIntegrator # type: ignore
-from openmm.unit import kelvin, picosecond # type: ignore
+from Bio.PDB import PDBParser, Superimposer
+from openmm.app import PDBFile, ForceField, Simulation, Modeller
+from openmm import LangevinIntegrator
+from openmm.unit import kelvin, picosecond
 
-def polish_design(target_pdb_id, uploaded_file_obj):
-    """
-    Refines the uploaded ESM-folded PDB against the target template.
-    uploaded_file_obj: This is the temp file path provided by Gradio.
-    """
-    # 1. Get the local template path
-    target_path = f"data/{target_pdb_id.lower()}.pdb"
-    # Gradio passes a 'temp' file object; we get the string path
-    design_path = uploaded_file_obj.name 
-
-    parser = PDBParser(QUIET=True)
+def polish_design(target_pdb_id, uploaded_file_path):
+    # 1. Setup paths
+    target_path = os.path.join("data", f"{target_pdb_id.lower()}.pdb")
     
-    # 2. ALIGNMENT (The Matchmaker)
+    # 2. ALIGNMENT (Biopython)
+    parser = PDBParser(QUIET=True)
     target_struct = parser.get_structure("target", target_path)
-    design_struct = parser.get_structure("design", design_path)
+    design_struct = parser.get_structure("design", uploaded_file_path)
     
     sup = Superimposer()
-    # Grabbing Alpha Carbons for the alignment
-    target_atoms = [a for a in target_struct.get_atoms() if a.get_name() == 'CA']
-    design_atoms = [a for a in design_struct.get_atoms() if a.get_name() == 'CA']
+    t_atoms = [a for a in target_struct.get_atoms() if a.get_name() == 'CA']
+    d_atoms = [a for a in design_struct.get_atoms() if a.get_name() == 'CA']
     
-    # Snap the design onto the template and calculate RMSD
-    sup.set_atoms(target_atoms[:len(design_atoms)], design_atoms)
+    sup.set_atoms(t_atoms[:len(d_atoms)], d_atoms)
     sup.apply(design_struct.get_atoms())
-    current_rmsd = sup.rms
+    rmsd = sup.rms
+
+    # 3. PHYSICS (OpenMM Modeller replaces PDBFixer)
+    pdb = PDBFile(uploaded_file_path)
+    modeller = Modeller(pdb.topology, pdb.positions)
+    
+    # Standardize Forcefield
+    forcefield = ForceField('amber14-all.xml', 'amber14/tip3p.xml')
+    
+    # Add Hydrogens (Standardizes the Nitrogen typenames)
+    modeller.addHydrogens(forcefield, pH=7.4)
     
-    # 3. PHYSICS CLEANUP (Resolving Clashes)
-    # PDBFixer solves the 'Nitrogen' error by standardizing atom types
-    fixer = PDBFixer(design_path)
-    fixer.findMissingResidues()
-    fixer.findMissingAtoms()
-    fixer.addMissingAtoms()
-    fixer.addMissingHydrogens(7.4)
+    # Run the wiggle to kill the 8 clashes
+    system = forcefield.createSystem(modeller.topology)
+    integrator = LangevinIntegrator(300*kelvin, 1/picosecond, 0.002*picosecond)
+    simulation = Simulation(modeller.topology, system, integrator)
+    simulation.context.setPositions(modeller.positions)
     
-    # Try to do energy minimization, but fall back to just using fixed structure if it fails
-    output_dir = "generated/shuttle"
-    os.makedirs(output_dir, exist_ok=True)
-    output_filename = os.path.join(output_dir, "refined.pdb")
+    simulation.minimizeEnergy(maxIterations=100)
     
-    try:
-        # Use Modeller to properly handle terminal residues
-        modeller = Modeller(fixer.topology, fixer.positions)
-        
-        # Setup the Forcefield (AMBER14)
-        forcefield = ForceField('amber14-all.xml', 'amber14/tip3p.xml')
-        # Add hydrogens with forcefield to get proper terminal capping
-        modeller.addHydrogens(forcefield=forcefield)
-        
-        # Create system - now with properly capped terminals
-        system = forcefield.createSystem(modeller.topology)
-        integrator = LangevinIntegrator(300*kelvin, 1/picosecond, 0.002*picosecond)
-        simulation = Simulation(modeller.topology, system, integrator)
-        simulation.context.setPositions(modeller.positions)
-        
-        # Energy Minimization: The 'Wiggle' that removes clashes
-        simulation.minimizeEnergy(maxIterations=100)
-        
-        # Export minimized structure
-        with open(output_filename, "w") as f:
-            PDBFile.writeFile(
-                simulation.topology, 
-                simulation.context.getState(getPositions=True).getPositions(), 
-                f # type: ignore
-            )
-    except Exception as e:
-        # Fallback: if OpenMM fails (e.g., due to incomplete terminal residues),
-        # just save the PDBFixer-processed structure without energy minimization
-        print(f"Warning: Energy minimization failed ({str(e)}), using PDBFixer output only")
-        PDBFile.writeFile(fixer.topology, fixer.positions, output_filename)
+    # 4. EXPORT
+    output_name = "Broteinshake_Lead_Refined.pdb"
+    with open(output_name, "w") as f:
+        PDBFile.writeFile(simulation.topology, 
+                         simulation.context.getState(getPositions=True).getPositions(), f)
     
-    return output_filename, current_rmsd
+    return output_name, rmsd