Create RestCapableHybridTopologyFactory

perses/annihilation/lambda_protocol.py

@@ -211,6 +211,9 @@ def plot_functions(self,n=50):
         plt.show()
 
 class RESTProtocol(object):
+    """
+    Lambda protocol to be used with perses.annihilation.rest.RESTTopologyFactory (which enables rest at the endstates only, not over the alchemical protocol)
+    """
     default_functions = {'solute_scale': lambda beta0, beta : beta / beta0,
                          'inter_scale' : lambda beta0, beta : np.sqrt(beta / beta0),
                          'steric_scale' : lambda beta0, beta : beta / beta0 - 1,
@@ -219,6 +222,74 @@ class RESTProtocol(object):
     def __init__(self):
         self.functions = RESTProtocol.default_functions
 
+class RESTCapableLambdaProtocol(object):
+    """
+    Lambda protocols to be used with perses.annihilation.relative.RESTCapableHybridTopologyFactory (which enables rest during the alchemical protocol). 
+
+    `default_functions` - default protocol to be used for running with alchemical and rest scaling at the same time. Scales the old energies linearly from 1 to 0, the new energies linearly from 0 to 1, and the rest region linearly such that sqrt(beta / beta0) is reached half way through the protocol.
+
+    `no_alchemy_functions` - default protocol to be used for running with rest scaling at one of the endstates (no alchemy). Scales the rest region linearly such that sqrt(beta / beta0) is reached half way through the protocol. lambda_alchemical_* should be set to either 0 or 1 (see RESTCapableRelativeAlchemicalState.set_alchemical_parameters()).   
+
+    """
+    default_functions = {'lambda_rest_bonds': lambda x, beta0, beta : -2 * (1 - np.sqrt(beta / beta0)) * x + 1 if x < 0.5 else 2 * (1 - np.sqrt(beta / beta0)) * x - 1 + 2 * np.sqrt(beta / beta0),
+                         'lambda_rest_angles': lambda x, beta0, beta : -2 * (1 - np.sqrt(beta / beta0)) * x + 1 if x < 0.5 else 2 * (1 - np.sqrt(beta / beta0)) * x - 1 + 2 * np.sqrt(beta / beta0),
+                         'lambda_rest_torsions':lambda x, beta0, beta : -2 * (1 - np.sqrt(beta / beta0)) * x + 1 if x < 0.5 else 2 * (1 - np.sqrt(beta / beta0)) * x - 1 + 2 * np.sqrt(beta / beta0),
+                         'lambda_rest_electrostatics': lambda x, beta0, beta : -2 * (1 - np.sqrt(beta / beta0)) * x + 1 if x < 0.5 else 2 * (1 - np.sqrt(beta / beta0)) * x - 1 + 2 * np.sqrt(beta / beta0),
+                         'lambda_rest_electrostatics_exceptions': lambda x, beta0, beta : -2 * (1 - np.sqrt(beta / beta0)) * x + 1 if x < 0.5 else 2 * (1 - np.sqrt(beta / beta0)) * x - 1 + 2 * np.sqrt(beta / beta0),
+                         'lambda_rest_sterics':lambda x, beta0, beta : -2 * (1 - np.sqrt(beta / beta0)) * x + 1 if x < 0.5 else 2 * (1 - np.sqrt(beta / beta0)) * x - 1 + 2 * np.sqrt(beta / beta0),
+                         'lambda_rest_sterics_exceptions': lambda x, beta0, beta : -2 * (1 - np.sqrt(beta / beta0)) * x + 1 if x < 0.5 else 2 * (1 - np.sqrt(beta / beta0)) * x - 1 + 2 * np.sqrt(beta / beta0),
+                         'lambda_alchemical_bonds_old': lambda x: 1 - x,
+                         'lambda_alchemical_bonds_new': lambda x: x,
+                         'lambda_alchemical_angles_old': lambda x: 1 - x,
+                         'lambda_alchemical_angles_new': lambda x: x,
+                         'lambda_alchemical_torsions_old': lambda x: 1 - x,
+                         'lambda_alchemical_torsions_new': lambda x: x,
+                         'lambda_alchemical_electrostatics_old': lambda x: 1 - x,
+                         'lambda_alchemical_electrostatics_new': lambda x: x,
+                         'lambda_alchemical_electrostatics_exceptions_old': lambda x: 1 - x,
+                         'lambda_alchemical_electrostatics_exceptions_new': lambda x: x,
+                         'lambda_alchemical_electrostatics_reciprocal': lambda x: x,
+                         'lambda_alchemical_sterics_old': lambda x: 1 - x,
+                         'lambda_alchemical_sterics_new': lambda x: x,
+                         'lambda_alchemical_sterics_exceptions_old': lambda x: 1 - x,
+                         'lambda_alchemical_sterics_exceptions_new': lambda x: x
+                         }
+
+    no_alchemy_functions = {'lambda_rest_bonds': lambda x, beta0, beta : (np.sqrt(beta / beta0) - 1) * x + 1,
+                                    'lambda_rest_angles': lambda x, beta0, beta : (np.sqrt(beta / beta0) - 1) * x + 1,
+                                    'lambda_rest_torsions': lambda x, beta0, beta : (np.sqrt(beta / beta0) - 1) * x + 1,
+                                    'lambda_rest_electrostatics': lambda x, beta0, beta : (np.sqrt(beta / beta0) - 1) * x + 1,
+                                    'lambda_rest_electrostatics_exceptions': lambda x, beta0, beta : (np.sqrt(beta / beta0) - 1) * x + 1,
+                                    'lambda_rest_sterics': lambda x, beta0, beta : (np.sqrt(beta / beta0) - 1) * x + 1,
+                                    'lambda_rest_sterics_exceptions': lambda x, beta0, beta : (np.sqrt(beta / beta0) - 1) * x + 1,
+                                    'lambda_alchemical_bonds_old': lambda x: 1 - x,
+                                    'lambda_alchemical_bonds_new': lambda x: x,
+                                    'lambda_alchemical_angles_old': lambda x: 1 - x,
+                                    'lambda_alchemical_angles_new': lambda x: x,
+                                    'lambda_alchemical_torsions_old': lambda x: 1 - x,
+                                    'lambda_alchemical_torsions_new': lambda x: x,
+                                    'lambda_alchemical_electrostatics_old': lambda x: 1 - x,
+                                    'lambda_alchemical_electrostatics_new': lambda x: x,
+                                    'lambda_alchemical_electrostatics_exceptions_old': lambda x: 1 - x,
+                                    'lambda_alchemical_electrostatics_exceptions_new': lambda x: x,
+                                    'lambda_alchemical_electrostatics_reciprocal': lambda x: x,
+                                    'lambda_alchemical_sterics_old': lambda x: 1 - x,
+                                    'lambda_alchemical_sterics_new': lambda x: x,
+                                    'lambda_alchemical_sterics_exceptions_old': lambda x: 1 - x,
+                                    'lambda_alchemical_sterics_exceptions_new': lambda x: x
+        }
+
+    def __init__(self, functions='default'):
+        if functions == 'default':
+            self.functions = RESTCapableLambdaProtocol.default_functions
+        elif functions == 'no-alchemy':
+            self.functions = RESTCapableLambdaProtocol.no_alchemy_functions
+        else:
+            raise Exception("User defined lambda protocols are not yet supported")
+
+        # TODO: Do I want to subclass LambdaProtocol to get its methods?
+
+
 
 class RelativeAlchemicalState(AlchemicalState):
     """
@@ -266,12 +337,15 @@ def set_alchemical_parameters(self, global_lambda,
 
 class RESTState(AlchemicalState):
     """
-    REST State to handle all lambda parameters required for REST2 implementation.
+    AlchemicalState to handle all lambda parameters required for running REST at the endstates with
+    perses.annihilation.rest.RESTTopologyFactory.
 
     Attributes
     ----------
     solute_scale : solute scaling parameter
     inter_scale : inter-region scaling parameter
+    electrostatic_scale : electrostatics scaling parameter
+    steric_scale : steric scaling parameter
     """
 
     class _LambdaParameter(AlchemicalState._LambdaParameter):
@@ -302,3 +376,108 @@ def set_alchemical_parameters(self,
        for parameter_name in lambda_protocol.functions:
            lambda_value = lambda_protocol.functions[parameter_name](beta0, beta)
            setattr(self, parameter_name, lambda_value)
+
+class RESTCapableRelativeAlchemicalState(AlchemicalState):
+    """
+    AlchemicalState to handle all lambda parameters required for running REST during the alchemical transformation with
+    perses.annihilation.relative.RESTCapableHybridTopologyFactory.
+
+    Attributes
+    ----------
+    lambda_rest_bonds
+        controls scaling of the rest region's bond energy
+    lambda_rest_angles
+        controls scaling of the rest region's angle energy
+    lambda_rest_torsions
+        controls scaling of the rest region's torsion energy
+    lambda_rest_electrostatics
+        controls scaling of the rest region's electrostatics energy
+    lambda_rest_electrostatics_exceptions
+        controls scaling of the rest region's electrostatics exceptions energy
+    lambda_rest_sterics
+        controls scaling of the rest region's sterics energy
+    lambda_rest_sterics_exceptions
+        controls scaling of the rest region's sterics exceptions energy
+    lambda_alchemical_bonds_old
+        controls alchemical scaling of the old bond energy
+    lambda_alchemical_bonds_new
+        controls alchemical scaling of the the new bond energy
+    lambda_alchemical_angles_old
+        controls alchemical scaling of the old angle energy
+    lambda_alchemical_angles_new
+        controls alchemical scaling of the new angle energy
+    lambda_alchemical_torsions_old
+        controls alchemical scaling of the old torsion energy
+    lambda_alchemical_torsions_new
+        controls alchemical scaling of the new torsion energy
+    lambda_alchemical_electrostatics_old
+        controls alchemical scaling of the old electrostatics energy
+    lambda_alchemical_electrostatics_new
+        controls alchemical scaling of the new electrostatics energy
+    lambda_alchemical_electrostatics_exceptions_old
+        controls alchemical scaling of the old electrostatics exceptions energy
+    lambda_alchemical_electrostatics_exceptions_new
+        controls alchemical scaling of the new electrostatics exceptions energy
+    lambda_alchemical_electrostatics_reciprocal
+        controls alchemical scaling of the reciprocal space energy
+    lambda_alchemical_sterics_old
+        controls alchemical scaling of the old sterics energy
+    lambda_alchemical_sterics_new
+        controls alchemical scaling of the new sterics energy
+    lambda_alchemical_sterics_exceptions_old
+        controls alchemical scaling of the old sterics exceptions energy
+    lambda_alchemical_sterics_exceptions_new
+        controls alchemical scaling of the new sterics exceptions energy
+    """
+
+    class _LambdaParameter(AlchemicalState._LambdaParameter):
+        @staticmethod
+        def lambda_validator(self, instance, parameter_value):
+            if parameter_value is None:
+                return parameter_value
+            return float(parameter_value)
+
+    lambda_rest_bonds = _LambdaParameter('lambda_rest_bonds')
+    lambda_rest_angles = _LambdaParameter('lambda_rest_angles')
+    lambda_rest_torsions = _LambdaParameter('lambda_rest_torsions')
+    lambda_rest_electrostatics = _LambdaParameter('lambda_rest_electrostatics')
+    lambda_rest_electrostatics_exceptions = _LambdaParameter('lambda_rest_electrostatics_exceptions')
+    lambda_rest_sterics = _LambdaParameter('lambda_rest_sterics')
+    lambda_rest_sterics_exceptions = _LambdaParameter('lambda_rest_sterics_exceptions')
+    lambda_alchemical_bonds_old = _LambdaParameter('lambda_alchemical_bonds_old')
+    lambda_alchemical_bonds_new = _LambdaParameter('lambda_alchemical_bonds_new')
+    lambda_alchemical_angles_old = _LambdaParameter('lambda_alchemical_angles_old')
+    lambda_alchemical_angles_new = _LambdaParameter('lambda_alchemical_angles_new')
+    lambda_alchemical_torsions_old = _LambdaParameter('lambda_alchemical_torsions_old')
+    lambda_alchemical_torsions_new = _LambdaParameter('lambda_alchemical_torsions_new')
+    lambda_alchemical_electrostatics_old = _LambdaParameter('lambda_alchemical_electrostatics_old')
+    lambda_alchemical_electrostatics_new = _LambdaParameter('lambda_alchemical_electrostatics_new')
+    lambda_alchemical_electrostatics_exceptions_old = _LambdaParameter('lambda_alchemical_electrostatics_exceptions_old')
+    lambda_alchemical_electrostatics_exceptions_new = _LambdaParameter('lambda_alchemical_electrostatics_exceptions_new')
+    lambda_alchemical_electrostatics_reciprocal = _LambdaParameter('lambda_alchemical_electrostatics_reciprocal')
+    lambda_alchemical_sterics_old = _LambdaParameter('lambda_alchemical_sterics_old')
+    lambda_alchemical_sterics_new = _LambdaParameter('lambda_alchemical_sterics_new')
+    lambda_alchemical_sterics_exceptions_old = _LambdaParameter('lambda_alchemical_sterics_exceptions_old')
+    lambda_alchemical_sterics_exceptions_new = _LambdaParameter('lambda_alchemical_sterics_exceptions_new')
+
+    def set_alchemical_parameters(self, global_lambda, beta0, beta, lambda_protocol=RESTCapableLambdaProtocol(), endstate=None):
+        """Set each lambda value according to the lambda_functions protocol.
+        The undefined parameters (i.e. those being set to None) remain
+        undefined.
+
+        Parameters
+        ----------
+        lambda_value : float
+            The new value for all defined parameters.
+        """
+        self.global_lambda = global_lambda
+        for parameter_name in lambda_protocol.functions:
+            if 'rest' in parameter_name:
+                lambda_value = lambda_protocol.functions[parameter_name](global_lambda, beta0, beta)
+            else:
+                if endstate is None:
+                    lambda_value = lambda_protocol.functions[parameter_name](global_lambda)
+                else:
+                    assert endstate in [0, 1], f"`endstate` should be 0 or 1, but was {endstate}"
+                    lambda_value = lambda_protocol.functions[parameter_name](endstate)
+            setattr(self, parameter_name, lambda_value)