Add numpy variants of extract_variable tests

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()