Merge pull request #80 from JohannesKarwou/potential_shapes

Potential shapes

transformato/analysis.py

@@ -872,7 +872,7 @@ def show_summary(self):
             self.plot_complex_free_energy()
             self.plot_waterbox_free_energy()
 
-        energy_estimate, uncertanty = self.end_state_free_energy_difference
+        energy_estimate, uncertainty = self.end_state_free_energy_difference
         print(
-            f"Free energy to common core: {energy_estimate} [kT] with uncertanty: {uncertanty} [kT]."
+            f"Free energy to common core: {energy_estimate} [kT] with uncertainty: {uncertainty} [kT]."
         )

transformato/restraints.py

@@ -44,7 +44,8 @@ def __init__(
         k: float = 3,
         shape: str = "harmonic",
         wellsize: float = 0.05,
-        **kwargs, # NOTE: WHY is this necessary?
+
+        **kwargs
     ):
         """Class representing a restraint to apply to the system.
 
@@ -57,15 +58,15 @@ def __init__(
             selligand,selprotein (MDAnalysis selection string): MDAnalysis selection strings
             k: the force (=spring) constant applied to the potential energy formula. See the 'System Setup' section for details.
             pdbpath: the path to the pdbfile underlying the topology analysis
-            shape: 'harmonic' or 'flatbottom'. Defines the shape of the harmonic energy potential.
-            wellsize: Defines the well-size in a flat-bottom potential. Defaults to 0.1 nanometers.
-            kwargs: Catcher for unneeded restraint_args entry
+            shape: one of 'harmonic', 'flatbottom', 'flatbottom-oneside-sharp' or 'flatbottom-twoside'. Defines the shape of the harmonic energy potential.
+            wellsize: Defines the well-size in a two-sided flat-bottom potential. Defaults to 0.05 nanometers.
+            kwargs: Catcher for additional restraint_args
 
         """
 
         self.shape = shape
 
-        if self.shape not in ["harmonic", "flatbottom"]:
+        if self.shape not in ["harmonic","flatbottom", "flatbottom-oneside-sharp","flatbottom-oneside","flatbottom-twoside"]:
             raise AttributeError(
                 f"Invalid potential shape specified for restraint: {self.shape}"
             )
@@ -76,7 +77,8 @@ def __init__(
 
         self.force_constant = k
         self.shape = shape
-        self.wellsize = wellsize
+        self.wellsize = wellsize * angstrom
+        self.kwargs=kwargs
 
     def createForce(self, common_core_names):
         """Actually creates the force, after dismissing all idxs not in the common core from the molecule.
@@ -109,6 +111,11 @@ def createForce(self, common_core_names):
 
         self.initial_distance = np.linalg.norm(self.g1pos - self.g2pos)
 
+        if not "r0" in self.kwargs.keys(): # default value - equilibrium distance is the initial distance. Otherwise, overriden with the r0 specified
+            self.r0 = self.initial_distance * angstrom
+        else:
+            self.r0 = self.kwargs["r0"] * angstrom
+
         if self.shape == "harmonic":
             # create force with harmonic potential
             self.force = CustomCentroidBondForce(2, "0.5*k*(distance(g1,g2)-r0)^2")
@@ -118,29 +125,40 @@ def createForce(self, common_core_names):
             self.force.addGroup(self.g2_openmm)
 
             self.force.addBond(
-                [0, 1], [self.force_constant, self.initial_distance / 10]
+                [0, 1], [self.force_constant, self.r0]
             )
 
             logger.info(
                 f"""Restraint force (centroid/bonded, shape is {self.shape}, initial distance: {self.initial_distance}, k={self.force_constant}"""
             )
-        elif self.shape == "flatbottom":
-            # create force with flat-bottom potential using openmmtools
-            from openmmtools.forces import FlatBottomRestraintForce
-
-            self.force = FlatBottomRestraintForce(
-                spring_constant=self.force_constant,
-                well_radius=self.initial_distance * angstrom,
-                restrained_atom_indices1=self.g1_openmm,
-                restrained_atom_indices2=self.g2_openmm,
+        elif self.shape in ["flatbottom","flatbottom-oneside"]:
+            # create force with flat-bottom potential identical to the one used in openmmtools
+            logger.debug("creating flatbottom-1side-soft")
+            self.force = CustomCentroidBondForce(
+                2, "step(distance(g1,g2)-r0) * (k/2)*(distance(g1,g2)-r0)^2"
             )
 
-            self.force.setUsesPeriodicBoundaryConditions(periodic=True)
+            self._add_flatbottom_parameters()
 
-            logger.info(
-                f"Restraint force (centroid/bonded, shape is {self.shape}, initial distance: {self.initial_distance}, k={self.force_constant} the finall parameters are: {self.force.getBondParameters(0)}"
+        elif self.shape == "flatbottom-oneside-sharp":
+            # create force with flat-bottom potential
+            logger.debug("creating flatbottom-1side-sharp")
+            self.force = CustomCentroidBondForce(
+                2, "step(distance(g1,g2)-r0) * (k/2)*(distance(g1,g2))^2"
+            )
+
+            self._add_flatbottom_parameters()
+
+        elif self.shape == "flatbottom-twoside":
+            # create force with flat-bottom potential
+            logger.debug("creating flatbottom-2side-sharp")
+            self.force = CustomCentroidBondForce(
+                2, "step(abs(distance(g1,g2)-r0)-w)*k*(distance(g1,g2)-r0)^2"
             )
 
+            self._add_flatbottom_parameters()
+
+
         else:
             raise NotImplementedError(f"Cannot create potential of shape {self.shape}")
 
@@ -156,6 +174,31 @@ def createForce(self, common_core_names):
     def get_force(self):
         return self.force
 
+    def _add_flatbottom_parameters(self):
+        """Takes care of the initial setup of an openMM CustomCentroidBondForce with a flatbottom shape
+        
+        Args:
+            force: An openMM CustomCentroidBondForce object
+        """
+
+        self.force.addPerBondParameter("k")
+        self.force.addPerBondParameter("r0")
+
+        self.force.addGroup(self.g1_openmm)
+        self.force.addGroup(self.g2_openmm)
+
+        if "twoside" in self.shape:     # two - sided flatbottoms gain an additional parameter for the wellsize
+            self.force.addPerBondParameter("w")
+            self.force.addBond(
+                [0, 1], [self.force_constant, self.r0, self.wellsize]
+            )
+        else:
+            self.force.addBond(
+                [0, 1], [self.force_constant, self.r0]
+            )
+
+        logger.info(f"Restraint force (centroid/bonded), shape is {self.shape}, Parameters: {self.force.getBondParameters(0)} ")
+
     def applyForce(self, system):
         """Applies the force to the openMM system
 

transformato/tests/test_restraints.py

@@ -173,9 +173,11 @@ def test_integration():
     pdbpath = PATH_2OJ9
     with open(
         f"{get_testsystems_dir()}/config/test-2oj9-restraints.yaml", "r"
+
     ) as stream:
         try:
             configuration = yaml.safe_load(stream)
+            logger.debug(configuration)
         except yaml.YAMLError as exc:
             print(exc)
 
@@ -232,6 +234,9 @@ def test_integration():
             state = simulation.context.getState(getEnergy=True, groups={i})
             logger.info("Force contributions before steps:")
             logger.info(f"{f}::{state.getPotentialEnergy()}")
+            if f.getEnergyFunction() == "step(abs(distance(g1,g2)-r0)-w)*k*(distance(g1,g2)-r0)^2":  # flatbottom-twosided, has set r0
+                assert 0.9 in f.getBondParameters(0)[1] # assert the manually set r0 value is found in the bond parameters
+
     logger.info("Simulation stepping")
     simulation.step(9)
     temp_results = []
@@ -264,11 +269,27 @@ def test_integration():
         if isinstance(f, CustomCentroidBondForce):
             logger.debug(f.getBondParameters(0))
             if (
-                f.getPerBondParameterName(0) == "k"
-            ):  # harmonic shape. openmmTools calls theirs "K"
+                f.getEnergyFunction() == "0.5*k*(distance(g1,g2)-r0)^2"
+            ):  # harmonic shape.
                 f.setBondParameters(0, [0, 1], [25, 5])
-            else:  # flatbottom shape:
+
+            elif (
+            f.getEnergyFunction() == "step(distance(g1,g2)-r0) * (k/2)*(distance(g1,g2)-r0)^2"
+            ):  # flatbottom-oneside-soft shape:
+                f.setBondParameters(0, [0, 1], [250, 0.01])
+
+            elif (
+            f.getEnergyFunction() == "step(distance(g1,g2)-r0) * (k/2)*(distance(g1,g2))^2"
+            ):  # flatbottom-oneside-sharp:
                 f.setBondParameters(0, [0, 1], [250, 0.01])
+
+            elif (
+                f.getEnergyFunction() == "step(abs(distance(g1,g2)-r0)-w)*k*(distance(g1,g2)-r0)^2"
+                ):  # flatbottom-twosided, has an additional wellsize parameter:
+                f.setBondParameters(0, [0, 1], [250, 12,1])
+
+
+
             f.updateParametersInContext(simulation.context)
 
     for i, f in enumerate(system.getForces()):