class TestDeepPotAPBC(unittest.TestCase):
@classmethod
def setUpClass(self):
self.dp_original = DeepPot(FROZEN_MODEL)
self.dp_compressed = DeepPot(COMPRESSED_MODEL)
self.coords = np.array([
12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, 3.36,
3.00, 1.81, 3.51, 2.51, 2.60, 4.27, 3.22, 1.56
])
self.atype = [0, 1, 1, 0, 1, 1]
self.box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.])
def test_attrs(self):
self.assertEqual(self.dp_original.get_ntypes(), 2)
self.assertAlmostEqual(self.dp_original.get_rcut(),
6.0,
places=default_places)
self.assertEqual(self.dp_original.get_type_map(), ['O', 'H'])
self.assertEqual(self.dp_original.get_dim_fparam(), 0)
self.assertEqual(self.dp_original.get_dim_aparam(), 0)
self.assertEqual(self.dp_compressed.get_ntypes(), 2)
self.assertAlmostEqual(self.dp_compressed.get_rcut(),
6.0,
places=default_places)
self.assertEqual(self.dp_compressed.get_type_map(), ['O', 'H'])
self.assertEqual(self.dp_compressed.get_dim_fparam(), 0)
self.assertEqual(self.dp_compressed.get_dim_aparam(), 0)
def test_1frame(self):
ee0, ff0, vv0 = self.dp_original.eval(self.coords,
self.box,
self.atype,
atomic=False)
ee1, ff1, vv1 = self.dp_compressed.eval(self.coords,
self.box,
self.atype,
atomic=False)
# check shape of the returns
nframes = 1
natoms = len(self.atype)
self.assertEqual(ee0.shape, (nframes, 1))
self.assertEqual(ff0.shape, (nframes, natoms, 3))
self.assertEqual(vv0.shape, (nframes, 9))
self.assertEqual(ee1.shape, (nframes, 1))
self.assertEqual(ff1.shape, (nframes, natoms, 3))
self.assertEqual(vv1.shape, (nframes, 9))
# check values
np.testing.assert_almost_equal(ff0, ff1, default_places)
np.testing.assert_almost_equal(ee0, ee1, default_places)
np.testing.assert_almost_equal(vv0, vv1, default_places)
def test_1frame_atm(self):
ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords,
self.box,
self.atype,
atomic=True)
ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(self.coords,
self.box,
self.atype,
atomic=True)
# check shape of the returns
nframes = 1
natoms = len(self.atype)
self.assertEqual(ee0.shape, (nframes, 1))
self.assertEqual(ff0.shape, (nframes, natoms, 3))
self.assertEqual(vv0.shape, (nframes, 9))
self.assertEqual(ae0.shape, (nframes, natoms, 1))
self.assertEqual(av0.shape, (nframes, natoms, 9))
self.assertEqual(ee1.shape, (nframes, 1))
self.assertEqual(ff1.shape, (nframes, natoms, 3))
self.assertEqual(vv1.shape, (nframes, 9))
self.assertEqual(ae1.shape, (nframes, natoms, 1))
self.assertEqual(av1.shape, (nframes, natoms, 9))
# check values
np.testing.assert_almost_equal(ff0, ff1, default_places)
np.testing.assert_almost_equal(ae0, ae1, default_places)
np.testing.assert_almost_equal(av0, av1, default_places)
np.testing.assert_almost_equal(ee0, ee1, default_places)
np.testing.assert_almost_equal(vv0, vv1, default_places)
def test_2frame_atm(self):
coords2 = np.concatenate((self.coords, self.coords))
box2 = np.concatenate((self.box, self.box))
ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(coords2,
box2,
self.atype,
atomic=True)
ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(coords2,
box2,
self.atype,
atomic=True)
# check shape of the returns
nframes = 2
natoms = len(self.atype)
self.assertEqual(ee0.shape, (nframes, 1))
self.assertEqual(ff0.shape, (nframes, natoms, 3))
self.assertEqual(vv0.shape, (nframes, 9))
self.assertEqual(ae0.shape, (nframes, natoms, 1))
self.assertEqual(av0.shape, (nframes, natoms, 9))
self.assertEqual(ee1.shape, (nframes, 1))
self.assertEqual(ff1.shape, (nframes, natoms, 3))
self.assertEqual(vv1.shape, (nframes, 9))
self.assertEqual(ae1.shape, (nframes, natoms, 1))
self.assertEqual(av1.shape, (nframes, natoms, 9))
# check values
np.testing.assert_almost_equal(ff0, ff1, default_places)
np.testing.assert_almost_equal(ae0, ae1, default_places)
np.testing.assert_almost_equal(av0, av1, default_places)
np.testing.assert_almost_equal(ee0, ee1, default_places)
np.testing.assert_almost_equal(vv0, vv1, default_places)
class TestDeepPotAPBC(unittest.TestCase):
def setUp(self):
convert_pbtxt_to_pb(
str(tests_path / os.path.join("infer", "deeppot.pbtxt")),
"deeppot.pb")
self.dp = DeepPot("deeppot.pb")
self.coords = np.array([
12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, 3.36,
3.00, 1.81, 3.51, 2.51, 2.60, 4.27, 3.22, 1.56
])
self.atype = [0, 1, 1, 0, 1, 1]
self.box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.])
self.expected_e = np.array([
-9.275780747115504710e+01, -1.863501786584258468e+02,
-1.863392472863538103e+02, -9.279281325486221021e+01,
-1.863671545232153903e+02, -1.863619822847602165e+02
])
self.expected_f = np.array([
-3.034045420701179663e-01, 8.405844663871177014e-01,
7.696947487118485642e-02, 7.662001266663505117e-01,
-1.880601391333554251e-01, -6.183333871091722944e-01,
-5.036172391059643427e-01, -6.529525836149027151e-01,
5.432962643022043459e-01, 6.382357912332115024e-01,
-1.748518296794561167e-01, 3.457363524891907125e-01,
1.286482986991941552e-03, 3.757251165286925043e-01,
-5.972588700887541124e-01, -5.987006197104716154e-01,
-2.004450304880958100e-01, 2.495901655353461868e-01
])
self.expected_v = np.array([
-2.912234126853306959e-01, -3.800610846612756388e-02,
2.776624987489437202e-01, -5.053761003913598976e-02,
-3.152373041953385746e-01, 1.060894290092162379e-01,
2.826389131596073745e-01, 1.039129970665329250e-01,
-2.584378792325942586e-01, -3.121722367954994914e-01,
8.483275876786681990e-02, 2.524662342344257682e-01,
4.142176771106586414e-02, -3.820285230785245428e-02,
-2.727311173065460545e-02, 2.668859789777112135e-01,
-6.448243569420382404e-02, -2.121731470426218846e-01,
-8.624335220278558922e-02, -1.809695356746038597e-01,
1.529875294531883312e-01, -1.283658185172031341e-01,
-1.992682279795223999e-01, 1.409924999632362341e-01,
1.398322735274434292e-01, 1.804318474574856390e-01,
-1.470309318999652726e-01, -2.593983661598450730e-01,
-4.236536279233147489e-02, 3.386387920184946720e-02,
-4.174017537818433543e-02, -1.003500282164128260e-01,
1.525690815194478966e-01, 3.398976109910181037e-02,
1.522253908435125536e-01, -2.349125581341701963e-01,
9.515545977581392825e-04, -1.643218849228543846e-02,
1.993234765412972564e-02, 6.027265332209678569e-04,
-9.563256398907417355e-02, 1.510815124001868293e-01,
-7.738094816888557714e-03, 1.502832772532304295e-01,
-2.380965783745832010e-01, -2.309456719810296654e-01,
-6.666961081213038098e-02, 7.955566551234216632e-02,
-8.099093777937517447e-02, -3.386641099800401927e-02,
4.447884755740908608e-02, 1.008593228579038742e-01,
4.556718179228393811e-02, -6.078081273849572641e-02
])
def tearDown(self):
os.remove("deeppot.pb")
def test_attrs(self):
self.assertEqual(self.dp.get_ntypes(), 2)
self.assertAlmostEqual(self.dp.get_rcut(), 6.0, places=default_places)
self.assertEqual(self.dp.get_type_map(), ['O', 'H'])
self.assertEqual(self.dp.get_dim_fparam(), 0)
self.assertEqual(self.dp.get_dim_aparam(), 0)
def test_1frame(self):
ee, ff, vv = self.dp.eval(self.coords,
self.box,
self.atype,
atomic=False)
# check shape of the returns
nframes = 1
natoms = len(self.atype)
self.assertEqual(ee.shape, (nframes, 1))
self.assertEqual(ff.shape, (nframes, natoms, 3))
self.assertEqual(vv.shape, (nframes, 9))
# check values
np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(),
default_places)
expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1)
np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(),
default_places)
expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1)
np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(),
default_places)
def test_1frame_atm(self):
ee, ff, vv, ae, av = self.dp.eval(self.coords,
self.box,
self.atype,
atomic=True)
# check shape of the returns
nframes = 1
natoms = len(self.atype)
self.assertEqual(ee.shape, (nframes, 1))
self.assertEqual(ff.shape, (nframes, natoms, 3))
self.assertEqual(vv.shape, (nframes, 9))
self.assertEqual(ae.shape, (nframes, natoms, 1))
self.assertEqual(av.shape, (nframes, natoms, 9))
# check values
np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(),
default_places)
np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(),
default_places)
np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(),
default_places)
expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1)
np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(),
default_places)
expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1)
np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(),
default_places)
def test_2frame_atm(self):
coords2 = np.concatenate((self.coords, self.coords))
box2 = np.concatenate((self.box, self.box))
ee, ff, vv, ae, av = self.dp.eval(coords2,
box2,
self.atype,
atomic=True)
# check shape of the returns
nframes = 2
natoms = len(self.atype)
self.assertEqual(ee.shape, (nframes, 1))
self.assertEqual(ff.shape, (nframes, natoms, 3))
self.assertEqual(vv.shape, (nframes, 9))
self.assertEqual(ae.shape, (nframes, natoms, 1))
self.assertEqual(av.shape, (nframes, natoms, 9))
# check values
expected_f = np.concatenate((self.expected_f, self.expected_f), axis=0)
expected_e = np.concatenate((self.expected_e, self.expected_e), axis=0)
expected_v = np.concatenate((self.expected_v, self.expected_v), axis=0)
np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(),
default_places)
np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(),
default_places)
np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(),
default_places)
expected_se = np.sum(expected_e.reshape([nframes, -1]), axis=1)
np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(),
default_places)
expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis=1)
np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(),
default_places)
class TestDeepPotAPBCExcludeTypes(unittest.TestCase):
@classmethod
def setUpClass(self):
self.dp_original = DeepPot(FROZEN_MODEL)
self.dp_compressed = DeepPot(COMPRESSED_MODEL)
self.coords = np.array([
12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, 3.36,
3.00, 1.81, 3.51, 2.51, 2.60, 4.27, 3.22, 1.56
])
self.atype = [0, 1, 1, 0, 1, 1]
self.box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.])
@classmethod
def tearDownClass(self):
_file_delete(INPUT)
_file_delete(FROZEN_MODEL)
_file_delete(COMPRESSED_MODEL)
_file_delete("out.json")
_file_delete("compress.json")
_file_delete("checkpoint")
_file_delete("model.ckpt.meta")
_file_delete("model.ckpt.index")
_file_delete("model.ckpt.data-00000-of-00001")
_file_delete("model.ckpt-100.meta")
_file_delete("model.ckpt-100.index")
_file_delete("model.ckpt-100.data-00000-of-00001")
_file_delete("model-compression/checkpoint")
_file_delete("model-compression/model.ckpt.meta")
_file_delete("model-compression/model.ckpt.index")
_file_delete("model-compression/model.ckpt.data-00000-of-00001")
_file_delete("model-compression")
_file_delete("input_v2_compat.json")
_file_delete("lcurve.out")
def test_attrs(self):
self.assertEqual(self.dp_original.get_ntypes(), 2)
self.assertAlmostEqual(self.dp_original.get_rcut(),
6.0,
places=default_places)
self.assertEqual(self.dp_original.get_type_map(), ['O', 'H'])
self.assertEqual(self.dp_original.get_dim_fparam(), 0)
self.assertEqual(self.dp_original.get_dim_aparam(), 0)
self.assertEqual(self.dp_compressed.get_ntypes(), 2)
self.assertAlmostEqual(self.dp_compressed.get_rcut(),
6.0,
places=default_places)
self.assertEqual(self.dp_compressed.get_type_map(), ['O', 'H'])
self.assertEqual(self.dp_compressed.get_dim_fparam(), 0)
self.assertEqual(self.dp_compressed.get_dim_aparam(), 0)
def test_1frame(self):
ee0, ff0, vv0 = self.dp_original.eval(self.coords,
self.box,
self.atype,
atomic=False)
ee1, ff1, vv1 = self.dp_compressed.eval(self.coords,
self.box,
self.atype,
atomic=False)
# check shape of the returns
nframes = 1
natoms = len(self.atype)
self.assertEqual(ee0.shape, (nframes, 1))
self.assertEqual(ff0.shape, (nframes, natoms, 3))
self.assertEqual(vv0.shape, (nframes, 9))
self.assertEqual(ee1.shape, (nframes, 1))
self.assertEqual(ff1.shape, (nframes, natoms, 3))
self.assertEqual(vv1.shape, (nframes, 9))
# check values
np.testing.assert_almost_equal(ff0, ff1, default_places)
np.testing.assert_almost_equal(ee0, ee1, default_places)
np.testing.assert_almost_equal(vv0, vv1, default_places)
def test_1frame_atm(self):
ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords,
self.box,
self.atype,
atomic=True)
ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(self.coords,
self.box,
self.atype,
atomic=True)
# check shape of the returns
nframes = 1
natoms = len(self.atype)
self.assertEqual(ee0.shape, (nframes, 1))
self.assertEqual(ff0.shape, (nframes, natoms, 3))
self.assertEqual(vv0.shape, (nframes, 9))
self.assertEqual(ae0.shape, (nframes, natoms, 1))
self.assertEqual(av0.shape, (nframes, natoms, 9))
self.assertEqual(ee1.shape, (nframes, 1))
self.assertEqual(ff1.shape, (nframes, natoms, 3))
self.assertEqual(vv1.shape, (nframes, 9))
self.assertEqual(ae1.shape, (nframes, natoms, 1))
self.assertEqual(av1.shape, (nframes, natoms, 9))
# check values
np.testing.assert_almost_equal(ff0, ff1, default_places)
np.testing.assert_almost_equal(ae0, ae1, default_places)
np.testing.assert_almost_equal(av0, av1, default_places)
np.testing.assert_almost_equal(ee0, ee1, default_places)
np.testing.assert_almost_equal(vv0, vv1, default_places)
def test_2frame_atm(self):
coords2 = np.concatenate((self.coords, self.coords))
box2 = np.concatenate((self.box, self.box))
ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(coords2,
box2,
self.atype,
atomic=True)
ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(coords2,
box2,
self.atype,
atomic=True)
# check shape of the returns
nframes = 2
natoms = len(self.atype)
self.assertEqual(ee0.shape, (nframes, 1))
self.assertEqual(ff0.shape, (nframes, natoms, 3))
self.assertEqual(vv0.shape, (nframes, 9))
self.assertEqual(ae0.shape, (nframes, natoms, 1))
self.assertEqual(av0.shape, (nframes, natoms, 9))
self.assertEqual(ee1.shape, (nframes, 1))
self.assertEqual(ff1.shape, (nframes, natoms, 3))
self.assertEqual(vv1.shape, (nframes, 9))
self.assertEqual(ae1.shape, (nframes, natoms, 1))
self.assertEqual(av1.shape, (nframes, natoms, 9))
# check values
np.testing.assert_almost_equal(ff0, ff1, default_places)
np.testing.assert_almost_equal(ae0, ae1, default_places)
np.testing.assert_almost_equal(av0, av1, default_places)
np.testing.assert_almost_equal(ee0, ee1, default_places)
np.testing.assert_almost_equal(vv0, vv1, default_places)
class TestDeepPotRPBC(unittest.TestCase):
def setUp(self):
convert_pbtxt_to_pb(
str(tests_path / os.path.join("infer", "deeppot-r.pbtxt")),
"deeppot.pb")
self.dp = DeepPot("deeppot.pb")
self.coords = np.array([
12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, 3.36,
3.00, 1.81, 3.51, 2.51, 2.60, 4.27, 3.22, 1.56
])
self.atype = [0, 1, 1, 0, 1, 1]
self.box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.])
self.expected_e = np.array([
-9.320909762801588272e+01, -1.868020345400987878e+02,
-1.868011172371355997e+02, -9.320868430396934912e+01,
-1.868010398844378415e+02, -1.868016706555875999e+02
])
self.expected_f = np.array([
6.385312846474267391e-04, -6.460452911141417731e-03,
-5.652405655332678417e-04, -7.516468794343579736e-03,
1.128804614240160216e-03, 5.531937784564192051e-03,
1.914138124904981664e-03, 5.601819906021693503e-03,
-5.131359585752605541e-03, -4.847104424804288617e-03,
1.992071550328819614e-03, -4.028159855157302516e-03,
1.236340684486603517e-03, -5.373955841338794344e-03,
8.312829460571366513e-03, 8.574563125108854156e-03,
3.111712681889538742e-03, -4.120007238692381148e-03
])
self.expected_v = np.array([
5.844056241889131371e-03, 4.663973497239899614e-04,
-2.268382127762904633e-03, 4.663973497239897988e-04,
2.349338784202595950e-03, -6.908546513234039253e-04,
-2.268382127762904633e-03, -6.908546513234039253e-04,
2.040499248150800561e-03, 4.238130266437327605e-03,
-1.539867187443782223e-04, -2.393101333240631613e-03,
-1.539867187443782223e-04, 4.410341945447907377e-04,
9.544239698119633068e-06, -2.393101333240631613e-03,
9.544239698119578858e-06, 1.877785959095269654e-03,
5.798992562057291543e-03, 6.943392552230453693e-04,
-1.180376879311998773e-03, 6.943392552230453693e-04,
1.686725132156275536e-03, -1.461632060145726542e-03,
-1.180376879311998556e-03, -1.461632060145726325e-03,
1.749543733794208444e-03, 7.173915604192910439e-03,
3.903218041111061569e-04, -5.747400467123527524e-04,
3.903218041111061569e-04, 1.208289706621179949e-03,
-1.826828914132010932e-03, -5.747400467123527524e-04,
-1.826828914132011148e-03, 2.856960586657185906e-03,
4.067553030177322240e-03, -3.267469855253819430e-05,
-6.980667859103454904e-05, -3.267469855253830272e-05,
1.387653029234650918e-03, -2.096820720698671855e-03,
-6.980667859103444062e-05, -2.096820720698671855e-03,
3.218305506720191278e-03, 4.753992590355240674e-03,
1.224911338353675992e-03, -1.683421934571502484e-03,
1.224911338353676209e-03, 7.332113564901583539e-04,
-1.025577052190138451e-03, -1.683421934571502484e-03,
-1.025577052190138234e-03, 1.456681925652047018e-03
])
def tearDown(self):
os.remove("deeppot.pb")
def test_attrs(self):
self.assertEqual(self.dp.get_ntypes(), 2)
self.assertAlmostEqual(self.dp.get_rcut(), 6.0, places=default_places)
self.assertEqual(self.dp.get_type_map(), ['O', 'H'])
self.assertEqual(self.dp.get_dim_fparam(), 0)
self.assertEqual(self.dp.get_dim_aparam(), 0)
# def test_1frame(self):
# ee, ff, vv, ae, av = self.dp.eval(self.coords, self.box, self.atype, atomic = True)
# np.savetxt('ee.out', ae.reshape([1, -1]), delimiter=',')
# np.savetxt('ff.out', ff.reshape([1, -1]), delimiter=',')
# np.savetxt('vv.out', av.reshape([1, -1]), delimiter=',')
def test_1frame(self):
ee, ff, vv = self.dp.eval(self.coords,
self.box,
self.atype,
atomic=False)
# check shape of the returns
nframes = 1
natoms = len(self.atype)
self.assertEqual(ee.shape, (nframes, 1))
self.assertEqual(ff.shape, (nframes, natoms, 3))
self.assertEqual(vv.shape, (nframes, 9))
# check values
np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(),
default_places)
expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1)
np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(),
default_places)
expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1)
np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(),
default_places)
def test_1frame_atm(self):
ee, ff, vv, ae, av = self.dp.eval(self.coords,
self.box,
self.atype,
atomic=True)
# check shape of the returns
nframes = 1
natoms = len(self.atype)
self.assertEqual(ee.shape, (nframes, 1))
self.assertEqual(ff.shape, (nframes, natoms, 3))
self.assertEqual(vv.shape, (nframes, 9))
self.assertEqual(ae.shape, (nframes, natoms, 1))
self.assertEqual(av.shape, (nframes, natoms, 9))
# check values
np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(),
default_places)
np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(),
default_places)
np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(),
default_places)
expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1)
np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(),
default_places)
expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1)
np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(),
default_places)
def test_2frame_atm(self):
coords2 = np.concatenate((self.coords, self.coords))
box2 = np.concatenate((self.box, self.box))
ee, ff, vv, ae, av = self.dp.eval(coords2,
box2,
self.atype,
atomic=True)
# check shape of the returns
nframes = 2
natoms = len(self.atype)
self.assertEqual(ee.shape, (nframes, 1))
self.assertEqual(ff.shape, (nframes, natoms, 3))
self.assertEqual(vv.shape, (nframes, 9))
self.assertEqual(ae.shape, (nframes, natoms, 1))
self.assertEqual(av.shape, (nframes, natoms, 9))
# check values
expected_f = np.concatenate((self.expected_f, self.expected_f), axis=0)
expected_e = np.concatenate((self.expected_e, self.expected_e), axis=0)
expected_v = np.concatenate((self.expected_v, self.expected_v), axis=0)
np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(),
default_places)
np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(),
default_places)
np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(),
default_places)
expected_se = np.sum(expected_e.reshape([nframes, -1]), axis=1)
np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(),
default_places)
expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis=1)
np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(),
default_places)