TestBytes.py 1.26 KB
Newer Older
Robert McGibbon's avatar
Robert McGibbon committed
1
2
3
4
5
6
7
8
import unittest
import simtk.openmm as mm


class TestBytes(unittest.TestCase):
    def test_createCheckpoint(self):
        system = mm.System()
        system.addParticle(1.0)
Robert McGibbon's avatar
Robert McGibbon committed
9
10
        refPositions = [(0,0,0)]

11
12
        platform = mm.Platform.getPlatformByName('Reference')
        context = mm.Context(system, mm.VerletIntegrator(0), platform)
Robert McGibbon's avatar
Robert McGibbon committed
13
14
15
        context.setPositions(refPositions)
        chk = context.createCheckpoint()
        # check that the return value of createCheckpoint is of type bytes (non-unicode)
Robert McGibbon's avatar
Robert McGibbon committed
16
17
        assert isinstance(chk, bytes)

Robert McGibbon's avatar
Robert McGibbon committed
18
19
20
21
22
23
24
25
26
        # set the positions to something random then reload the checkpoint, and
        # make sure that the positions get restored correctly
        context.setPositions([(12345, 12345, 123451)])
        context.loadCheckpoint(chk)
        newPositions = context.getState(getPositions=True).getPositions()._value

        assert newPositions == refPositions

        # try encoding the checkpoint in utf-8. OpenMM should be able to handle this too
27
28
29
        context.setPositions([(12345, 12345, 123451)])
        context.loadCheckpoint(chk.decode('utf-8'))
        newPositions = context.getState(getPositions=True).getPositions()._value
Robert McGibbon's avatar
Robert McGibbon committed
30
31
32
33

        assert newPositions == refPositions


Robert McGibbon's avatar
Robert McGibbon committed
34
35
if __name__ == '__main__':
    unittest.main()
Robert McGibbon's avatar
Robert McGibbon committed
36
37