Make openmm try refinement a few times before giving up

ABlooper/ABlooper.py

@@ -34,6 +34,8 @@ def __init__(self, pdb_file, chains=("H", "L"), model=None, refine=False, refine
         :param pdb_file: File of IMGT numbered antibody structure in .pdb format. Must include heavy and light chain
         :param chains: Name of heavy and light chain to be remodelled in the file
         :param model: Trained neural network. If left as None a pretrained network is used.
+        :bool refine: If the refinement and side-chain prediction steps should done.
+        :param refine_method: Package used to do relaxation.
         """
         self.pdb_file = pdb_file
         self.chains = chains

ABlooper/openmm_refine.py

@@ -1,16 +1,23 @@
 import pdbfixer
 import tempfile
 from io import StringIO
-from simtk.openmm import app, LangevinIntegrator, CustomExternalForce
+from simtk.openmm import app, LangevinIntegrator, CustomExternalForce, OpenMMException
 from simtk import unit
 
-
 ENERGY = unit.kilocalories_per_mole
 LENGTH = unit.angstroms
 spring_unit = ENERGY / (LENGTH ** 2)
 
 
-def openmm_refine(pdb_txt, CDR_definitions, spring_constant=10):
+def relax_CDR_loops(pdb_txt, CDR_definitions, spring_constant=10):
+    """ Function to perform restrained energy minimisation on the CDR loops.
+
+    :param pdb_txt: ABlooper prediction stored as text in PDB format.
+    :param CDR_definitions: How we are defining the CDR loops.
+    :param spring_constant: Strength of the spring keeping backbone atoms in place.
+    :param attempts: Number of attempts to run minimization before giving up.
+    """
+
     with tempfile.NamedTemporaryFile("wt") as tmp:
         tmp.writelines(pdb_txt)
         tmp.flush()
@@ -38,7 +45,7 @@ def openmm_refine(pdb_txt, CDR_definitions, spring_constant=10):
 
     # Keep atoms close to initial prediction
     force = CustomExternalForce("0.5 * k * ((x-x0)^2 + (y-y0)^2 + (z-z0)^2)")
-    force.addGlobalParameter("k", spring_constant*spring_unit)
+    force.addGlobalParameter("k", spring_constant * spring_unit)
     for p in ["x0", "y0", "z0"]:
         force.addPerParticleParameter(p)
 
@@ -68,4 +75,23 @@ def openmm_refine(pdb_txt, CDR_definitions, spring_constant=10):
     out_txt = out_handle.getvalue()
     out_handle.close()
 
-    return [x+"\n" for x in out_txt.split("\n") if x[:3] in ["ATO", "TER"]]
+    return [x + "\n" for x in out_txt.split("\n") if x[:3] in ["ATO", "TER"]]
+
+
+def openmm_refine(pdb_txt, CDR_definitions, spring_constant=10, attempts=5):
+    """ Function to do energy minimization with openmm.
+
+    Sometimes (around 5 in 1000) OpenMM will return a 'Particle position is NaN' error.
+    Rerunning appears to fix this, so I have set it so it tries a few times before giving up
+    """
+
+    for i in range(1,attempts):
+        try:
+            output = relax_CDR_loops(pdb_txt, CDR_definitions, spring_constant)
+        except OpenMMException:
+            print("Relaxation Failed: Trying again. Number of attempts: {}".format(i))
+        else:
+            return output
+
+    return relax_CDR_loops(pdb_txt, CDR_definitions, spring_constant)
+

setup.py

@@ -4,7 +4,7 @@
     long_description = fh.read()
 setup(
     name='ABlooper',
-    version='1.0.8',
+    version='1.0.9',
     description='Set of functions to predict CDR structure',
     license='BSD 3-clause license',
     maintainer='Brennan Abanades',