From 7aa45ea816a2872933b744241a3b4e684840798c Mon Sep 17 00:00:00 2001 From: Richard Berger Date: Mon, 23 Nov 2020 14:13:58 -0500 Subject: [PATCH] Add numpy variants of extract_variable tests --- unittest/python/python-numpy.py | 35 ++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/unittest/python/python-numpy.py b/unittest/python/python-numpy.py index 46794590f4..9f1d5f12c8 100644 --- a/unittest/python/python-numpy.py +++ b/unittest/python/python-numpy.py @@ -1,5 +1,7 @@ import sys,os,unittest -from lammps import lammps, LAMMPS_INT, LMP_STYLE_GLOBAL, LMP_STYLE_LOCAL, LMP_STYLE_ATOM, LMP_TYPE_VECTOR, LMP_TYPE_SCALAR, LMP_TYPE_ARRAY +from lammps import lammps, LAMMPS_INT, LMP_STYLE_GLOBAL, LMP_STYLE_LOCAL, \ + LMP_STYLE_ATOM, LMP_TYPE_VECTOR, LMP_TYPE_SCALAR, LMP_TYPE_ARRAY, \ + LMP_VAR_ATOM from ctypes import c_void_p try: @@ -178,5 +180,36 @@ class PythonNumpy(unittest.TestCase): self.assertIn(1, neighbors_i) self.assertNotIn(0, neighbors_j) + def test_extract_variable_equalstyle(self): + self.lmp.command("variable a equal 100") + a = self.lmp.numpy.extract_variable("a") + self.assertEqual(a, 100) + + self.lmp.command("variable a equal 3.14") + a = self.lmp.numpy.extract_variable("a") + self.assertEqual(a, 3.14) + + def test_extract_variable_atomstyle(self): + self.lmp.command("units lj") + self.lmp.command("atom_style atomic") + self.lmp.command("atom_modify map array") + self.lmp.command("boundary f f f") + self.lmp.command("region box block 0 2 0 2 0 2") + self.lmp.command("create_box 1 box") + + x = [ + 1.0, 1.0, 1.0, + 1.0, 1.0, 1.5 + ] + + types = [1, 1] + + self.assertEqual(self.lmp.create_atoms(2, id=None, type=types, x=x), 2) + self.lmp.command("variable a atom x*x+y*y+z*z") + a = self.lmp.numpy.extract_variable("a", "all", LMP_VAR_ATOM) + self.assertIs(type(a), numpy.ndarray) + self.assertEqual(a[0], x[0]*x[0]+x[1]*x[1]+x[2]*x[2]) + self.assertEqual(a[1], x[3]*x[3]+x[4]*x[4]+x[5]*x[5]) + if __name__ == "__main__": unittest.main()