diff --git a/python/lammps/core.py b/python/lammps/core.py index 269b716c13..d981243503 100644 --- a/python/lammps/core.py +++ b/python/lammps/core.py @@ -1842,6 +1842,8 @@ class lammps(object): """ nlocal = self.extract_setting('nlocal') + if len(eatom) < nlocal: + raise Exception('per-atom energy list length must be at least nlocal') ceatom = (nlocal*c_double)(*eatom) with ExceptionCheck(self): return self.lib.lammps_fix_external_set_energy_peratom(self.lmp, fix_id.encode(), ceatom) @@ -1862,6 +1864,10 @@ class lammps(object): # copy virial data to C compatible buffer nlocal = self.extract_setting('nlocal') + if len(vatom) < nlocal: + raise Exception('per-atom virial first dimension must be at least nlocal') + if len(vatom[0]) != 6: + raise Exception('per-atom virial second dimension must be 6') vbuf = (c_double * 6) vptr = POINTER(c_double) c_virial = (vptr * nlocal)() diff --git a/python/lammps/numpy_wrapper.py b/python/lammps/numpy_wrapper.py index 7752d7902e..29a6f9d74e 100644 --- a/python/lammps/numpy_wrapper.py +++ b/python/lammps/numpy_wrapper.py @@ -283,6 +283,9 @@ class numpy_wrapper: """ import numpy as np nlocal = self.lmp.extract_setting('nlocal') + if len(eatom) < nlocal: + raise Exception('per-atom energy dimension must be at least nlocal') + c_double_p = POINTER(c_double) value = eatom.astype(np.double) return self.lmp.lib.lammps_fix_external_set_energy_peratom(self.lmp.lmp, fix_id.encode(), @@ -307,7 +310,13 @@ class numpy_wrapper: :param eatom: per-atom potential energy :type: numpy.array """ - raise Exception('fix_external_set_virial_peratom() not yet implemented for NumPy arrays') + import numpy as np + nlocal = self.lmp.extract_setting('nlocal') + if len(vatom) < nlocal: + raise Exception('per-atom virial first dimension must be at least nlocal') + if len(vatom[0]) != 6: + raise Exception('per-atom virial second dimension must be 6') + # -------------------------------------------------------------------------