Added simulation test, and cleaned up tests a bit.

We still only test one PDB file, however.

tests/test_build.py

-from Bio.PDB import PDBList
-from simtk import unit
-from pdbfixer.pdbfixer import *
+
+import pdbfixer
+import simtk.openmm 
+import Bio.PDB
+
 import os
 import sys
+import numpy
 
 def test_build():
     # These are tough PDB codes from http://www.umass.edu/microbio/chime/pe_beta/pe/protexpl/badpdbs.htm
@@ -10,6 +13,7 @@ def test_build():
     pdbcodes = ['1VII'] # DEBUG: just use one
 
     # Set up PDB retrieval.
+    from Bio.PDB import PDBList
     pdblist = PDBList(server=PDBList.alternative_download_url)
 
     success = True
@@ -29,6 +33,7 @@ def test_build():
 
         # PDB setup parameters.
         # TODO: Try several combinations?
+        from simtk import unit
         pH = 7.0
         ionic = 50.0 * unit.millimolar
         box = 10.0 * unit.angstrom
@@ -39,6 +44,8 @@ def test_build():
         outfile = open(output_pdb_filename, 'w')
 
         try:
+            from pdbfixer.pdbfixer import PDBFixer, PdbStructure
+            from simtk.openmm import app
             fixer = PDBFixer(PdbStructure(infile))
             fixer.findMissingResidues()
             fixer.findNonstandardResidues()
@@ -63,3 +70,5 @@ def test_build():
     if not success:
         raise Exception("build test failed on one or more PDB files.")
 
+if __name__ == '__main__':
+    test_build()

tests/test_simulate.py

-from Bio.PDB import PDBList
-from simtk import unit
-from pdbfixer.pdbfixer import *
+
+import pdbfixer
+import simtk.openmm 
+import Bio.PDB
+
 import os
 import sys
+import numpy
+
+def simulate(pdbcode, pdb_filename):
+    from simtk.openmm import app
+    import simtk.openmm as mm
+    from simtk import unit
+    from sys import stdout
+
+    # Load the PDB file.
+    pdb = app.PDBFile(pdb_filename)
+    
+    # Set up implicit solvent forcefield.
+    forcefield = app.ForceField('amber99sbildn.xml', 'amber99_obc.xml')
+    
+    # Create the system.
+    system = forcefield.createSystem(pdb.topology, nonbondedMethod=app.NoCutoff, constraints=app.HBonds)
+
+    # Create an integrator.
+    integrator = mm.LangevinIntegrator(300*unit.kelvin, 91.0/unit.picoseconds, 1.0*unit.femtoseconds)
+
+    # Create a context.
+    context = mm.Context(system, integrator)
+    context.setPositions(pdb.positions)
+    
+    # Check to make sure energy is finite.
+    state = context.getState(getEnergy=True)
+    potential = state.getPotentialEnergy() / unit.kilocalories_per_mole
+    if numpy.isnan(potential):
+        raise Exception("Initial energy for %s is NaN." % pdbcode)
+
+    # Minimize.
+    tolerance = 1.0 * unit.kilocalories_per_mole / unit.angstroms
+    maxIterations = 50
+    mm.LocalEnergyMinimizer.minimize(context, tolerance, maxIterations)
 
-def test_build():
+    # Check to make sure energy is finite.
+    state = context.getState(getEnergy=True)
+    potential = state.getPotentialEnergy() / unit.kilocalories_per_mole
+    if numpy.isnan(potential):
+        raise Exception("Energy for %s is NaN after minimization." % pdbcode)
+
+    # Simulate.
+    nsteps = 500
+    integrator.step(nsteps)
+
+    # Check to make sure energy is finite.
+    state = context.getState(getEnergy=True)
+    potential = state.getPotentialEnergy() / unit.kilocalories_per_mole
+    if numpy.isnan(potential):
+        raise Exception("Energy for %s is NaN after simulation." % pdbcode)
+
+    del context, integrator
+
+    print "Simulation completed: potential = %.3f kcal/mol" % potential
+
+    return
+
+def test_simulate():
     # These are tough PDB codes from http://www.umass.edu/microbio/chime/pe_beta/pe/protexpl/badpdbs.htm
     pdbcodes = ['1AS5', '1CBN', '1DPO', '1IGY', '1HAG', '1IAO', '4CPA', '1QCQ']
     pdbcodes = ['1VII'] # DEBUG: just use one
 
     # Set up PDB retrieval.
+    from Bio.PDB import PDBList
     pdblist = PDBList(server=PDBList.alternative_download_url)
 
     success = True
@@ -29,6 +88,7 @@ def test_build():
 
         # PDB setup parameters.
         # TODO: Try several combinations?
+        from simtk import unit
         pH = 7.0
         ionic = 50.0 * unit.millimolar
         box = 10.0 * unit.angstrom
@@ -39,6 +99,8 @@ def test_build():
         outfile = open(output_pdb_filename, 'w')
 
         try:
+            from pdbfixer.pdbfixer import PDBFixer, PdbStructure
+            from simtk.openmm import app
             fixer = PDBFixer(PdbStructure(infile))
             fixer.findMissingResidues()
             fixer.findNonstandardResidues()
@@ -52,6 +114,9 @@ def test_build():
             infile.close()
             outfile.close()
             
+            # Test simulating this with OpenMM.
+            simulate(pdbcode, output_pdb_filename)
+
             # Delete input file.
             os.remove(input_pdb_filename)
             os.remove(output_pdb_filename)
@@ -63,3 +128,5 @@ def test_build():
     if not success:
         raise Exception("build test failed on one or more PDB files.")
 
+if __name__ == '__main__':
+    test_simulate()