Files
lammps-gran-kokkos/unittest/python/python-numpy.py
2020-08-28 13:53:36 -04:00

43 lines
1.7 KiB
Python

import sys,os,unittest
from lammps import lammps, LMP_STYLE_GLOBAL, LMP_STYLE_ATOM, LMP_TYPE_VECTOR, LMP_TYPE_SCALAR
class PythonNumpy(unittest.TestCase):
def setUp(self):
machine = None
if 'LAMMPS_MACHINE_NAME' in os.environ:
machine=os.environ['LAMMPS_MACHINE_NAME']
self.lmp = lammps(name=machine, cmdargs=['-nocite', '-log','none', '-echo','screen'])
def tearDown(self):
del self.lmp
def testExtractCompute(self):
self.lmp.command("region box block 0 2 0 2 0 2")
self.lmp.command("create_box 1 box")
self.lmp.command("create_atoms 1 single 1.0 1.0 1.0")
self.lmp.command("create_atoms 1 single 1.0 1.0 1.5")
self.lmp.command("compute coordsum all reduce sum x y z")
natoms = int(self.lmp.get_natoms())
self.assertEqual(natoms,2)
values = self.lmp.numpy.extract_compute("coordsum", LMP_STYLE_GLOBAL, LMP_TYPE_VECTOR)
self.assertEqual(len(values), 3)
self.assertEqual(values[0], 2.0)
self.assertEqual(values[1], 2.0)
self.assertEqual(values[2], 2.5)
def testExtractComputePerAtom(self):
self.lmp.command("region box block 0 2 0 2 0 2")
self.lmp.command("create_box 1 box")
self.lmp.command("create_atoms 1 single 1.0 1.0 1.0")
self.lmp.command("create_atoms 1 single 1.0 1.0 1.5")
self.lmp.command("compute ke all ke/atom")
natoms = int(self.lmp.get_natoms())
self.assertEqual(natoms,2)
values = self.lmp.numpy.extract_compute("ke", LMP_STYLE_ATOM, LMP_TYPE_VECTOR)
self.assertEqual(len(values), 2)
self.assertEqual(values[0], 0.0)
self.assertEqual(values[1], 0.0)
if __name__ == "__main__":
unittest.main()