import openpathsampling.engines.openmm as omm_engine
import openpathsampling as paths
from nose.tools import assert_equal
from test_helpers import data_filename

import openmmtools as omt

import logging

logging.getLogger('openpathsampling.initialization').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.ensemble').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.storage').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.netcdfplus').setLevel(logging.CRITICAL)

class testOpenMMSnapshot(object):
    def setup(self):
        test_system = omt.testsystems.AlanineDipeptideVacuum()
        self.template = omm_engine.snapshot_from_testsystem(test_system)
        self.engine = omm_engine.Engine(
            topology=self.template.topology,
            system=test_system.system,
            integrator=omt.integrators.VVVRIntegrator()
        )
        self.n_atoms = self.engine.topology.n_atoms
        self.engine.current_snapshot = self.template

    def test_masses_from_file(self):
        masses = self.template.masses
        assert_equal(len(masses), self.n_atoms)

    def test_masses_from_simulation(self):
        sim_snap = self.engine.current_snapshot
        masses = sim_snap.masses
        assert_equal(len(masses), self.n_atoms)

    def test_n_degrees_of_freedom(self):
        assert_equal(self.engine.current_snapshot.n_degrees_of_freedom, 51)

@@ -114,16 +114,37 @@
 rt_equal
+, assert_almost_equal
 %0Afrom te
@@ -174,16 +174,35 @@
 filename
+, assert_close_unit
 %0A%0Aimport
@@ -220,16 +220,58 @@
 s as omt
+%0Aimport simtk.unit as u%0Aimport numpy as np
 %0A%0Aimport
@@ -1486,8 +1486,799 @@
 om, 51)%0A
+%0A    def test_instantaneous_temperature(self):%0A        vel_unit = u.nanometers / u.picoseconds%0A        new_velocities = %5B%5B1.0, 0.0, 0.0%5D%5D * self.n_atoms * vel_unit%0A        self.engine.current_snapshot = omm_engine.Snapshot.construct(%0A            coordinates=self.engine.current_snapshot.coordinates,%0A            box_vectors=self.engine.current_snapshot.box_vectors,%0A            velocities=new_velocities,%0A            engine=self.engine%0A        )%0A        test_snap = self.engine.current_snapshot%0A        expected_ke = sum(%0A            %5Bm * vel_unit**2 for m in test_snap.masses%5D, %0A            0.0*u.joule%0A        )%0A        n_dofs = 51.0  # see test above%0A        assert_close_unit(test_snap.instantaneous_temperature,%0A                          expected_ke / u.BOLTZMANN_CONSTANT_kB / n_dofs)%0A

test for OpenMM instantaneous temperature