openmm/wrappers/python/tests/TestPdbFile.py

import tempfile
import unittest
from openmm.app import *
from openmm import *
from openmm.unit import *
import openmm.app.element as elem
from io import StringIO


class TestPdbFile(unittest.TestCase):
    """Test the PDB file parser"""

    def test_Triclinic(self):
        """Test parsing a file that describes a triclinic box."""
        pdb = PDBFile('systems/triclinic.pdb')
        self.assertEqual(len(pdb.positions), 8)
        expectedPositions = [
            Vec3(1.744, 2.788, 3.162),
            Vec3(1.048, 0.762, 2.340),
            Vec3(2.489, 1.570, 2.817),
            Vec3(1.027, 1.893, 3.271),
            Vec3(0.937, 0.825, 0.009),
            Vec3(2.290, 1.887, 3.352),
            Vec3(1.266, 1.111, 2.894),
            Vec3(0.933, 1.862, 3.490)]*nanometers
        for (p1, p2) in zip(expectedPositions, pdb.positions):
            self.assertVecAlmostEqual(p1, p2)
        expectedVectors = [
            Vec3(2.5, 0, 0),
            Vec3(0.5, 3.0, 0),
            Vec3(0.7, 0.9, 3.5)]*nanometers
        for (v1, v2) in zip(expectedVectors, pdb.topology.getPeriodicBoxVectors()):
            self.assertVecAlmostEqual(v1, v2, 1e-4)
        self.assertVecAlmostEqual(Vec3(2.5, 3.0, 3.5)*nanometers, pdb.topology.getUnitCellDimensions(), 1e-4)
        for atom in pdb.topology.atoms():
            if atom.index < 4:
                self.assertEqual(elem.chlorine, atom.element)
                self.assertEqual('Cl', atom.name)
                self.assertEqual('Cl', atom.residue.name)
            else:
                self.assertEqual(elem.sodium, atom.element)
                self.assertEqual('Na', atom.name)
                self.assertEqual('Na', atom.residue.name)

    def test_WriteFile(self):
        """Write a file, read it back, and make sure it matches the original."""
        def compareFiles(pdb1, pdb2):
            self.assertEqual(len(pdb2.positions), 8)
            for (p1, p2) in zip(pdb1.positions, pdb2.positions):
                self.assertVecAlmostEqual(p1, p2)
            for (v1, v2) in zip(pdb1.topology.getPeriodicBoxVectors(), pdb2.topology.getPeriodicBoxVectors()):
                self.assertVecAlmostEqual(v1, v2, 1e-4)
            self.assertVecAlmostEqual(pdb1.topology.getUnitCellDimensions(), pdb2.topology.getUnitCellDimensions(), 1e-4)
            for atom1, atom2 in zip(pdb1.topology.atoms(), pdb2.topology.atoms()):
                self.assertEqual(atom1.element, atom2.element)
                self.assertEqual(atom1.name, atom2.name)
                self.assertEqual(atom1.residue.name, atom2.residue.name)

        pdb1 = PDBFile('systems/triclinic.pdb')

        # First try writing to an open file object.

        output = StringIO()
        PDBFile.writeFile(pdb1.topology, pdb1.positions, output)
        input = StringIO(output.getvalue())
        pdb2 = PDBFile(input)
        output.close()
        input.close()
        compareFiles(pdb1, pdb2)

        # Now try a filename.

        with tempfile.TemporaryDirectory() as tempdir:
            filename = os.path.join(tempdir, 'temp.pdb')
            PDBFile.writeFile(pdb1.topology, pdb1.positions, filename)
            pdb2 = PDBFile(filename)
            compareFiles(pdb1, pdb2)

    def test_BinaryStream(self):
        """Test reading a stream that was opened in binary mode."""
        with open('systems/triclinic.pdb', 'rb') as infile:
            pdb = PDBFile(infile)
        self.assertEqual(len(pdb.positions), 8)

    def test_ExtraParticles(self):
        """Test reading, and writing and re-reading of a file containing extra particle atoms."""
        pdb = PDBFile('systems/tip5p.pdb')
        for atom in pdb.topology.atoms():
            if atom.index > 2:
                self.assertEqual(None, atom.element)
        output = StringIO()
        PDBFile.writeFile(pdb.topology, pdb.positions, output)
        input = StringIO(output.getvalue())
        pdb = PDBFile(input)
        for atom in pdb.topology.atoms():
            if atom.index > 2:
                self.assertEqual(None, atom.element)

    def test_AltLocs(self):
        """Test reading a file that includes AltLocs"""
        for filename in ['altlocs.pdb', 'altlocs2.pdb']:
            pdb = PDBFile(f'systems/{filename}')
            self.assertEqual(1, pdb.topology.getNumResidues())
            self.assertEqual(19, pdb.topology.getNumAtoms())
            self.assertEqual(19, len(pdb.positions))
            self.assertEqual('ILE', list(pdb.topology.residues())[0].name)

    def test_FormalCharges(self):
        """Test reading, and writing and re-reading of a file containing formal charges."""
        pdb = PDBFile('systems/formal-charges.pdb')
        for atom in pdb.topology.atoms():
            if atom.index == 8:
                self.assertEqual(+1, atom.formalCharge)
            elif atom.index == 29:
                self.assertEqual(-1, atom.formalCharge)
            else:
                self.assertEqual(None, atom.formalCharge)
        output = StringIO()
        PDBFile.writeFile(pdb.topology, pdb.positions, output)
        input = StringIO(output.getvalue())
        pdb = PDBFile(input)
        for atom in pdb.topology.atoms():
            if atom.index == 8:
                self.assertEqual(+1, atom.formalCharge)
            elif atom.index == 29:
                self.assertEqual(-1, atom.formalCharge)
            else:
                self.assertEqual(None, atom.formalCharge)

    def test_LargeFile(self):
        """Write and read a file with more than 100,000 atoms"""
        topology = Topology()
        chain = topology.addChain('A')
        for i in range(20000):
            res = topology.addResidue('MOL', chain)
            atoms = []
            for j in range(6):
                atoms.append(topology.addAtom(f'AT{j}', elem.carbon, res))
            for j in range(5):
                topology.addBond(atoms[j], atoms[j+1])
        positions = [Vec3(0, 0, 0)]*topology.getNumAtoms()

        # The model has 20,000 residues and 120,000 atoms.

        output = StringIO()
        PDBFile.writeFile(topology, positions, output)
        input = StringIO(output.getvalue())
        pdb = PDBFile(input)
        output.close()
        input.close()
        self.assertEqual(len(positions), len(pdb.positions))
        self.assertEqual(topology.getNumAtoms(), pdb.topology.getNumAtoms())
        self.assertEqual(topology.getNumResidues(), pdb.topology.getNumResidues())
        self.assertEqual(topology.getNumChains(), pdb.topology.getNumChains())
        self.assertEqual(topology.getNumBonds(), pdb.topology.getNumBonds())
        for atom in pdb.topology.atoms():
            self.assertEqual(str(atom.index+1), atom.id)
        for res in pdb.topology.residues():
            self.assertEqual(str(res.index+1), res.id)

        # Make sure the CONECT records were interpreted correctly.

        bonds = set()
        for atom1, atom2 in topology.bonds():
            bonds.add(tuple(sorted((atom1.index, atom2.index))))
        for atom1, atom2 in pdb.topology.bonds():
            assert tuple(sorted((atom1.index, atom2.index))) in bonds

    def assertVecAlmostEqual(self, p1, p2, tol=1e-7):
        unit = p1.unit
        p1 = p1.value_in_unit(unit)
        p2 = p2.value_in_unit(unit)
        scale = max(1.0, norm(p1),)
        for i in range(3):
            diff = abs(p1[i]-p2[i])/scale
            self.assertTrue(diff < tol)


if __name__ == '__main__':
    unittest.main()