New Features: Open Source OpenMM relaxing and antigen trimming for memory efficiency.

IgGM/deploy/ab_design.py

@@ -101,7 +101,7 @@ def infer(self, chains, task='design', *args, **kwargs):
         outputs = self.forward(inputs, *args, **kwargs)
         return inputs, outputs
 
-    def infer_pdb(self, chains, filename, relax=False, task='design', *args, **kwargs):
+    def infer_pdb(self, chains, filename, relax=False, relax_open=False, task='design', *args, **kwargs):
         inputs, outputs = self.infer(chains, task, *args, **kwargs)
         complex_id = inputs["base"]["complex_id"]
         raw_seqs = {}
@@ -110,8 +110,8 @@ def infer_pdb(self, chains, filename, relax=False, task='design', *args, **kwarg
             raw_seqs[chain_id] = raw_seq
 
         self._output_to_fasta(inputs, outputs, filename[:-4] + ".fasta")
-        if task == 'design' or task == 'fr_design':
-            self._output_to_pdb(inputs, outputs, filename, relax=relax)
+        if task == 'design' or task == 'fr_design' or task == 'affinity_maturation':
+            self._output_to_pdb(inputs, outputs, filename, relax=relax, relax_open=relax_open)
 
     def __sample_cm_ss2ss(self, prot_data_curr, idx_step, inputs_addi, chunk_size=None, temperature=1.0):
         """Sample amino-acid sequences & backbone structures w/ CM."""

IgGM/deploy/base_designer.py

@@ -13,7 +13,7 @@
 from IgGM.protein.parser import PdbParser
 from IgGM.protein.prot_constants import RESD_NAMES_1C, N_ATOMS_PER_RESD
 from IgGM.protein.utils import init_qta_params
-from IgGM.utils import calc_trsl_vec, get_tmp_dpath, Rosetta_relax
+from IgGM.utils import calc_trsl_vec, get_tmp_dpath, Rosetta_relax, OpenMM_relax
 
 
 class BaseDesigner(BaseModel):
@@ -214,7 +214,7 @@ def _output_to_fasta(inputs, outputs, filename):
         export_fasta(sequences, ids=ids, output=filename)
 
     @staticmethod
-    def _output_to_pdb(inputs, outputs, filename, relax=False):
+    def _output_to_pdb(inputs, outputs, filename, relax=False, relax_open=False):
         """Build a dict of protein structure data."""
         pred_info = 'REMARK 250 Structure predicted by IgGM\n'
         ligand_id = inputs["base"]["ligand_id"]
@@ -264,12 +264,19 @@ def _output_to_pdb(inputs, outputs, filename, relax=False):
         PdbParser.save_multimer(prot_data, filename, pred_info=pred_info)
         # remove the tmp directory tmp_path
         shutil.rmtree(tmp_path)
-        if relax:
-            try:
-                Rosetta_relax(filename)
-            except Exception as e:
-                logging.error(f'Error during Rosetta relaxation: {e}')
-                logging.info('Relaxation failed, skipping...')
+        if relax or relax_open:
+            if relax:
+                try:
+                    Rosetta_relax(filename)
+                except Exception as e:
+                    logging.error(f'Error during Rosetta relaxation: {e}')
+                    logging.info('Relaxation failed, skipping...')
+            elif relax_open:
+                try:
+                    OpenMM_relax(filename)
+                except Exception as e:
+                    logging.error(f'Error during OpenMM relaxation: {e}')
+                    logging.info('Relaxation failed, skipping...')
 
         logging.info(f'PDB file generated: {filename}')
 

IgGM/utils/__init__.py

@@ -6,18 +6,23 @@
 from .file import jload, jdump, get_tmp_dpath, download_file
 from .env import seed_all_rng, setup_logger, setup
 from .diff_util import ss2ptr, ptr2ss, so3_scale, intp_prob_mat_dsct, intp_trsl_mat, intp_rota_tns, IsotropicGaussianSO3, IGSO3Buffer, rota2quat, replace_with_mask, calc_trsl_vec
+from .openmm_relax import OpenMM_relax
 
 # Relax the input pdb file
 def Rosetta_relax(pdb_file):
     import os
-    import pyrosetta
-    pyrosetta.init()
-    from pyrosetta.rosetta.core.select import residue_selector as selections
-    from pyrosetta import pose_from_pdb, create_score_function
-    from pyrosetta.rosetta.core.pack.task import TaskFactory, operation
-    from pyrosetta.rosetta.core.select.movemap import MoveMapFactory, move_map_action
-    from pyrosetta.rosetta.protocols.minimization_packing import PackRotamersMover
-    from pyrosetta.rosetta.protocols.relax import FastRelax
+    try:
+        import pyrosetta
+        pyrosetta.init()
+        from pyrosetta.rosetta.core.select import residue_selector as selections
+        from pyrosetta import pose_from_pdb, create_score_function
+        from pyrosetta.rosetta.core.pack.task import TaskFactory, operation
+        from pyrosetta.rosetta.core.select.movemap import MoveMapFactory, move_map_action
+        from pyrosetta.rosetta.protocols.minimization_packing import PackRotamersMover
+        from pyrosetta.rosetta.protocols.relax import FastRelax
+    except ImportError:
+        print("[ERROR] PyRosetta not found. Please install via `pip install pyrosetta-installer` and setup, or use --relax_open for OpenMM.")
+        return
 
     print(f'Rosetta processing {pdb_file} for Relax')
 

IgGM/utils/openmm_relax.py

@@ -0,0 +1,129 @@
+import os
+import sys
+import tempfile
+
+def OpenMM_relax(pdb_file):
+    """
+    Relaxes a PDB structure using OpenMM and PDBFixer.
+    Replaces the proprietary PyRosetta relaxation.
+
+    Strategy:
+    1. Fix residues/atoms with PDBFixer.
+    2. Remove Heterogens (Ligands/Water) to ensure Amber14 compatibility.
+    3. Save to temp PDB and STRIP CONECTs to clean topology.
+    4. Reload, Add Hydrogens (Modeller), Minimize.
+    """
+    print(f"[OpenMM] Processing {pdb_file} for Relax")
+
+    try:
+        from pdbfixer import PDBFixer
+        from openmm.app import PDBFile, Modeller, ForceField, Simulation, PME, HBonds, NoCutoff
+        from openmm import LangevinIntegrator, unit, Platform
+    except ImportError:
+        print("[ERROR] OpenMM or PDBFixer not found. Please install them via conda: `conda install -c conda-forge openmm pdbfixer`")
+        return
+
+    # 1. Load and Fix PDB
+    try:
+        print("   Initializing PDBFixer...")
+        fixer = PDBFixer(filename=pdb_file)
+
+        print("   Removing heterogens (ligands/water)...")
+        fixer.removeHeterogens(keepWater=False)
+
+        print("   Fixing topology (missing residues/atoms)...")
+        fixer.findMissingResidues()
+        fixer.findMissingAtoms()
+        fixer.addMissingAtoms()
+        # skip fixer.addMissingHydrogens - we use Modeller later
+    except Exception as e:
+        print(f"   [ERROR] Error validating/fixing PDB: {e}")
+        return
+
+    # 2. Save Intermediate and Strip CONECTs
+    print("   Stripping bad CONECT records...")
+    temp_fd, temp_path = tempfile.mkstemp(suffix=".pdb")
+    os.close(temp_fd)
+
+    try:
+        # Write fixed PDB to temp
+        with open(temp_path, 'w') as f:
+            PDBFile.writeFile(fixer.topology, fixer.positions, f)
+
+        # Read back and strip CONECT lines
+        with open(temp_path, 'r') as f:
+            lines = f.readlines()
+
+        clean_lines = [line for line in lines if not line.startswith("CONECT")]
+
+        # Write clean PDB (reusing temp path)
+        with open(temp_path, 'w') as f:
+            f.writelines(clean_lines)
+
+        # 3. Reload Clean PDB
+        print("   Reloading clean topology...")
+        pdb = PDBFile(temp_path)
+
+    except Exception as e:
+        print(f"   [ERROR] Error in Strip-CONECT step: {e}")
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+        return
+    finally:
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+
+    # 4. Setup Modeller and ForceField
+    try:
+        # Standard choice: amber14-all.xml and implicit solvent like implicit/gbn2.xml
+        forcefield = ForceField('amber14-all.xml', 'implicit/gbn2.xml')
+
+        modeller = Modeller(pdb.topology, pdb.positions)
+
+        # 5. Add Hydrogens using ForceField
+        print("   Adding hydrogens...")
+        modeller.addHydrogens(forcefield)
+
+        # 6. Create System
+        system = forcefield.createSystem(modeller.topology, nonbondedMethod=NoCutoff, 
+                                         constraints=HBonds)
+
+        # Integrator
+        integrator = LangevinIntegrator(300*unit.kelvin, 1.0/unit.picosecond, 0.002*unit.picoseconds)
+
+        # Platform
+        try:
+            platform = Platform.getPlatformByName('CUDA')
+            prop = {'CudaPrecision': 'mixed'}
+        except Exception:
+            try:
+                platform = Platform.getPlatformByName('OpenCL')
+                prop = {}
+            except:
+                platform = Platform.getPlatformByName('CPU')
+                prop = {}
+
+        print(f"   Using platform: {platform.getName()}")
+
+        simulation = Simulation(modeller.topology, system, integrator, platform, prop)
+        simulation.context.setPositions(modeller.positions)
+
+        # 7. Minimize Energy
+        state0 = simulation.context.getState(getEnergy=True)
+        e0 = state0.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole)
+        print(f"   Initial Energy: {e0:.2f} kJ/mol")
+
+        print("   Minimizing energy...")
+        simulation.minimizeEnergy()
+
+        # 8. Save Relaxed PDB
+        state = simulation.context.getState(getPositions=True, getEnergy=True)
+        e1 = state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole)
+        print(f"   Final Energy:   {e1:.2f} kJ/mol")
+        with open(pdb_file, 'w') as f:
+            PDBFile.writeFile(simulation.topology, state.getPositions(), f)
+
+        print(f"[SUCCESS] OpenMM relaxed PDB saved to: {pdb_file}")
+
+    except Exception as e:
+        print(f"   [ERROR] OpenMM Relax failed: {e}")

IgGM/utils/registry.py

@@ -27,7 +27,7 @@ def get_registry(name):
 
 class Registry(dict):
     """
-    registry helper，
+    registry helper,
     Eg. creeting a registry:
         some_registry = Registry({'default': default_module})