Merge pull request #109 from choderalab/samplers

Fix sampler bugs, cut down NCMC testing time, add alkanes test system

.travis.yml

@@ -36,11 +36,11 @@ script:
   - conda remove --yes parmed parmed-dev && conda install --yes parmed-dev
   - conda remove --yes openmm openmm-dev && conda install --yes openmm-dev
   # Run tests
-  - cd devtools && nosetests $PACKAGENAME --nocapture --verbosity=3 --with-doctest --with-timer && cd ..
+  - cd devtools && nosetests --processes=-1 --process-timeout=1800 $PACKAGENAME --nocapture --verbosity=3 --with-doctest --with-timer && cd ..
 env:
   matrix:
     - python=2.7  CONDA_PY=27 OPENEYE_CHANNEL="-i https://pypi.anaconda.org/openeye/channel/main/simple"
-    - python=3.4  CONDA_PY=34 OPENEYE_CHANNEL="-i https://pypi.anaconda.org/openeye/channel/main/simple"
+    #- python=3.4  CONDA_PY=34 OPENEYE_CHANNEL="-i https://pypi.anaconda.org/openeye/channel/main/simple"
     #- python=3.5  CONDA_PY=35 OPENEYE_CHANNEL="-i https://pypi.anaconda.org/openeye/channel/main/simple"
 
   global:

perses/annihilation/ncmc_switching.py

@@ -16,6 +16,10 @@
 default_nsteps = 1
 default_timestep = 1.0 * unit.femtoseconds
 
+class NaNException(Exception):
+    def __init__(self, *args, **kwargs):
+        super(NaNException,self).__init__(*args,**kwargs)
+
 class NCMCEngine(object):
     """
     NCMC switching engine
@@ -75,14 +79,23 @@ def __init__(self, temperature=default_temperature, functions=default_functions,
         self.constraint_tolerance = constraint_tolerance
         self.platform = platform
 
-    def _getAvailableParameters(self, system):
+    @property
+    def beta(self):
+        kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA
+        kT = kB * self.temperature
+        beta = 1.0 / kT
+        return beta
+
+    def _getAvailableParameters(self, system, prefix='lambda'):
         """
-        Return a list of available context parameters defined in the system
+        Return a list of available alchemical context parameters defined in the system
 
         Parameters
         ----------
         system : simtk.openmm.System
             The system for which available context parameters are to be determined
+        prefix : str, optional, default='lambda'
+            Prefix required for parameters to be returned.
 
         Returns
         -------
@@ -95,7 +108,9 @@ def _getAvailableParameters(self, system):
             force = system.getForce(force_index)
             if hasattr(force, 'getNumGlobalParameters'):
                 for parameter_index in range(force.getNumGlobalParameters()):
-                    parameters.append(force.getGlobalParameterName(parameter_index))
+                    parameter_name = force.getGlobalParameterName(parameter_index)
+                    if parameter_name[0:(len(prefix)+1)] == (prefix + '_'):
+                        parameters.append(parameter_name)
         return parameters
 
     def _computeAlchemicalCorrection(self, unmodified_system, alchemical_system, initial_positions, final_positions, direction='insert'):
@@ -167,15 +182,11 @@ def computePotentialEnergy(system, positions):
             return potential
 
         # Compute correction from transforming real system to/from alchemical system
-        kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA
-        kT = kB * self.temperature
-        beta = 1.0 / kT
-
         if direction == 'delete':
             alchemical_potential_correction = computePotentialEnergy(alchemical_system, initial_positions) - computePotentialEnergy(unmodified_system, initial_positions)
         elif direction == 'insert':
             alchemical_potential_correction = computePotentialEnergy(unmodified_system, final_positions) - computePotentialEnergy(alchemical_system, final_positions)
-        logP_alchemical_correction = -beta * alchemical_potential_correction
+        logP_alchemical_correction = -self.beta * alchemical_potential_correction
 
         return logP_alchemical_correction
 
@@ -260,6 +271,32 @@ def integrate(self, topology_proposal, initial_positions, direction='insert', pl
         # Create alchemical system.
         [unmodified_system, alchemical_system] = self.make_alchemical_system(topology_proposal, direction=direction)
 
+        # DEBUG: Compute initial potential of unmodified system.
+        integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
+        if self.platform is not None:
+            context = openmm.Context(unmodified_system, integrator, self.platform)
+        else:
+            context = openmm.Context(unmodified_system, integrator)
+        context.setPositions(initial_positions)
+        context.applyConstraints(integrator.getConstraintTolerance())
+        unmodified_potential = self.beta * context.getState(getEnergy=True).getPotentialEnergy()
+        del context, integrator
+        if np.isnan(unmodified_potential):
+            raise NaNException("Initial potential of unmodified system is NaN")
+
+        # DEBUG: Compute initial potential of alchemical system without changing alchemical lambda.
+        integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
+        if self.platform is not None:
+            context = openmm.Context(alchemical_system, integrator, self.platform)
+        else:
+            context = openmm.Context(alchemical_system, integrator)
+        context.setPositions(initial_positions)
+        context.applyConstraints(integrator.getConstraintTolerance())
+        alchemical_potential = self.beta * context.getState(getEnergy=True).getPotentialEnergy()
+        del context, integrator
+        if np.isnan(alchemical_potential):
+            raise NaNException("Initial potential of alchemical system is NaN")
+
         # Select subset of switching functions based on which alchemical parameters are present in the system.
         available_parameters = self._getAvailableParameters(alchemical_system)
         functions = { parameter_name : self.functions[parameter_name] for parameter_name in self.functions if (parameter_name in available_parameters) }
@@ -279,18 +316,43 @@ def integrate(self, topology_proposal, initial_positions, direction='insert', pl
         # Set velocities to temperature and apply velocity constraints.
         context.setVelocitiesToTemperature(self.temperature)
         context.applyVelocityConstraints(integrator.getConstraintTolerance())
-        # Store initial potential if 'insert'
+
+        # Set initial context parameters.
         if direction == 'insert':
             for parameter_name in available_parameters:
                 context.setParameter(parameter_name, 0)
-            potential = context.getState(getEnergy=True).getPotentialEnergy()
-        # Only take a single integrator step since all switching steps are unrolled in NCMCAlchemicalIntegrator.
-        integrator.step(1)
-        # Store final potential if 'insert'
-        if direction == 'delete':
+        elif direction == 'delete':
+            for parameter_name in available_parameters:
+                context.setParameter(parameter_name, 1)
+
+        # Compute initial potential of alchemical state.
+        initial_potential = self.beta * context.getState(getEnergy=True).getPotentialEnergy()
+        if np.isnan(initial_potential):
+            raise NaNException("Initial potential of 'insert' operation is NaN (unmodified potential was %.3f kT, alchemical potential was %.3f kT before changing lambda)" % (unmodified_potential, alchemical_potential))
+
+        # Take a single integrator step since all switching steps are unrolled in NCMCAlchemicalIntegrator.
+        try:
+            integrator.step(1)
+        except Exception as e:
+            # Trap NaNs as a special exception (allowing us to reject later, if desired)
+            if e.msg == "Particle coordinate is nan":
+                raise NaNException(str(e))
+            else:
+                raise e
+
+        # Set final context parameters.
+        if direction == 'insert':
+            for parameter_name in available_parameters:
+                context.setParameter(parameter_name, 1)
+        elif direction == 'delete':
             for parameter_name in available_parameters:
                 context.setParameter(parameter_name, 0)
-            potential = context.getState(getEnergy=True).getPotentialEnergy()
+
+        # Compute final potential of alchemical state.
+        final_potential = self.beta * context.getState(getEnergy=True).getPotentialEnergy()
+        if np.isnan(final_potential):
+            raise NaNException("Final potential of 'delete' operation is NaN")
+
         # Store final positions and log acceptance probability.
         final_positions = context.getState(getPositions=True).getPositions(asNumpy=True)
         logP_NCMC = integrator.getLogAcceptanceProbability()
@@ -308,6 +370,12 @@ def integrate(self, topology_proposal, initial_positions, direction='insert', pl
         # Clean up alchemical system.
         del alchemical_system
 
+        # Select whether to return initial or final potential.
+        if direction == 'insert':
+            potential = initial_potential
+        elif direction == 'delete':
+            potential = final_potential
+
         # Return
         return [final_positions, logP, potential]
 

perses/rjmc/geometry.py

@@ -543,6 +543,8 @@ def _choose_torsion(self, atoms_with_positions, atom_for_proposal):
         Pick an eligible torsion uniformly
         """
         eligible_torsions = self._get_torsions(atoms_with_positions, atom_for_proposal)
+        if len(eligible_torsions) == 0:
+            raise Exception("No eligible torsions found for placing atom %s." % str(atom_for_proposal))
         torsion_idx = np.random.randint(0, len(eligible_torsions))
         torsion_selected = eligible_torsions[torsion_idx]
         return torsion_selected, np.log(1.0/len(eligible_torsions))

perses/rjmc/topology_proposal.py

@@ -75,8 +75,10 @@ class TopologyProposal(object):
         Number of atoms in the new system
     natoms_old : int
         Number of atoms in the old system
-    chemical_state_key : str
-        The current chemical state
+    old_chemical_state_key : str
+        The previous chemical state key
+    new_chemical_state_key : str
+        The proposed chemical state key
     metadata : dict
         additional information of interest about the state
     """
@@ -255,14 +257,14 @@ def propose(self, current_system, current_topology, current_metadata=None):
         if metadata == None:
             metadata = dict()
         # old_chemical_state_key : str
-        old_chemical_state_key = self.compute_state_key(current_topology)
+        old_chemical_state_key = self.compute_state_key(old_topology)
 
         # chain_id : str
         chain_id = self._chain_id
         # save old indeces for mapping -- could just directly save positions instead
         # modeller : simtk.openmm.app.Modeller
-        current_positions = np.zeros((current_topology.getNumAtoms(), 3))
-        modeller = app.Modeller(current_topology, current_positions)
+        current_positions = np.zeros((old_topology.getNumAtoms(), 3))
+        modeller = app.Modeller(old_topology, current_positions)
         # atom : simtk.openmm.app.topology.Atom
         for atom in modeller.topology.atoms():
             # atom.old_index : int
@@ -300,7 +302,22 @@ def propose(self, current_system, current_topology, current_metadata=None):
         # new_system : simtk.openmm.System
         new_system = self._system_generator.build_system(new_topology)
 
-        return TopologyProposal(new_topology=new_topology, new_system=new_system, old_topology=old_topology, old_system=current_system, old_chemical_state_key=old_chemical_state_key, new_chemical_state_key=new_chemical_state_key, logp_proposal=0.0, new_to_old_atom_map=atom_map)
+        # Create TopologyProposal.
+        topology_proposal = TopologyProposal(new_topology=new_topology, new_system=new_system, old_topology=old_topology, old_system=current_system, old_chemical_state_key=old_chemical_state_key, new_chemical_state_key=new_chemical_state_key, logp_proposal=0.0, new_to_old_atom_map=atom_map)
+
+        # Check to make sure no out-of-bounds atoms are present in new_to_old_atom_map
+        natoms_old = topology_proposal.old_system.getNumParticles()
+        natoms_new = topology_proposal.new_system.getNumParticles()
+        if not set(topology_proposal.new_to_old_atom_map.values()).issubset(range(natoms_old)):
+            msg = "Some old atoms in TopologyProposal.new_to_old_atom_map are not in span of old atoms (1..%d):\n" % natoms_old
+            msg += str(topology_proposal.new_to_old_atom_map)
+            raise Exception(msg)
+        if not set(topology_proposal.new_to_old_atom_map.keys()).issubset(range(natoms_new)):
+            msg = "Some new atoms in TopologyProposal.new_to_old_atom_map are not in span of old atoms (1..%d):\n" % natoms_new
+            msg += str(topology_proposal.new_to_old_atom_map)
+            raise Exception(msg)
+
+        return topology_proposal
 
     def _choose_mutation_from_allowed(self, modeller, chain_id, allowed_mutations):
         """
@@ -658,13 +675,27 @@ class SystemGenerator(object):
         Allows specification of various aspects of system creation.
     metadata : dict, optional
         Metadata associated with the SystemGenerator.
+    use_antechamber : bool, optional, default=True
+        If True, will add the GAFF residue template generator.
     """
 
-    def __init__(self, forcefields_to_use, forcefield_kwargs=None, metadata=None):
+    def __init__(self, forcefields_to_use, forcefield_kwargs=None, metadata=None, use_antechamber=True):
         self._forcefield_xmls = forcefields_to_use
         self._forcefield_kwargs = forcefield_kwargs if forcefield_kwargs is not None else {}
         self._forcefield = app.ForceField(*self._forcefield_xmls)
-        self._forcefield.registerTemplateGenerator(forcefield_generators.gaffTemplateGenerator)
+        if use_antechamber:
+            self._forcefield.registerTemplateGenerator(forcefield_generators.gaffTemplateGenerator)
+
+    def getForceField(self):
+        """
+        Return the associated ForceField object.
+
+        Returns
+        -------
+        forcefield : simtk.openmm.app.ForceField
+            The current ForceField object.
+        """
+        return self._forcefield
 
     def build_system(self, new_topology):
         """
@@ -676,13 +707,27 @@ def build_system(self, new_topology):
         new_topology : simtk.openmm.app.Topology object
             The topology of the system
 
-
         Returns
         -------
         new_system : openmm.System
             A system object generated from the topology
         """
-        system = self._forcefield.createSystem(new_topology, **self._forcefield_kwargs)
+        try:
+            system = self._forcefield.createSystem(new_topology, **self._forcefield_kwargs)
+        except Exception as e:
+            from simtk import unit
+            nparticles = sum([1 for atom in new_topology.atoms()])
+            positions = unit.Quantity(np.zeros([nparticles,3], np.float32), unit.angstroms)
+            # Write PDB file of failed topology
+            from simtk.openmm.app import PDBFile
+            outfile = open('BuildSystem-failure.pdb', 'w')
+            pdbfile = PDBFile.writeFile(new_topology, positions, outfile)
+            outfile.close()
+            msg = str(e)
+            msg += "\n"
+            msg += "PDB file written as 'BuildSystem-failure.pdb'"
+            raise Exception(msg)
+
         return system
 
     @property
@@ -749,7 +794,6 @@ def propose(self, current_system, current_topology, current_metadata=None):
         smiles_new, _ = self._topology_to_smiles(new_topology)
         assert smiles_new == proposed_mol_smiles
 
-
         #map the atoms between the new and old molecule only:
         mol_atom_map, alignment_logp = self._get_mol_atom_map(current_mol, proposed_mol)