diff --git a/python/lammps.py b/python/lammps.py index 9f7e02e8e456bdadc574576cfdc247b670a05cc5..285596caa6260385b2271b888865213669c7b112 100644 --- a/python/lammps.py +++ b/python/lammps.py @@ -348,8 +348,11 @@ class AtomList(object): def __init__(self, lammps_wrapper_instance): self.lmp = lammps_wrapper_instance self.natoms = self.lmp.system.natoms + self.dimensions = self.lmp.system.dimensions def __getitem__(self, index): + if self.dimensions == 2: + return Atom2D(self.lmp, index + 1) return Atom(self.lmp, index + 1) @@ -397,6 +400,26 @@ class Atom(object): return self.lmp.eval("q[%d]" % self.index) +class Atom2D(Atom): + def __init__(self, lammps_wrapper_instance, index): + super(Atom2D, self).__init__(lammps_wrapper_instance, index) + + @property + def position(self): + return (self.lmp.eval("x[%d]" % self.index), + self.lmp.eval("y[%d]" % self.index)) + + @property + def velocity(self): + return (self.lmp.eval("vx[%d]" % self.index), + self.lmp.eval("vy[%d]" % self.index)) + + @property + def force(self): + return (self.lmp.eval("fx[%d]" % self.index), + self.lmp.eval("fy[%d]" % self.index)) + + class PyLammps(object): """ More Python-like wrapper for LAMMPS (e.g., for iPython)