Merge pull request #15 from rbberger/pylammps-improvements

PyLammps improvements

python/lammps.py

@@ -309,6 +309,9 @@ class OutputCapture(object):
 
   def __exit__(self, type, value, tracebac):
     os.dup2(self.stdout, self.stdout_fd)
+    os.close(self.stdout)
+    os.close(self.stdout_pipe_read)
+    os.close(self.stdout_pipe_write)
 
   # check if we have more to read from the pipe
   def more_data(self, pipe):
@@ -350,8 +353,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)
 
 
@@ -382,6 +388,12 @@ class Atom(object):
             self.lmp.eval("y[%d]" % self.index),
             self.lmp.eval("z[%d]" % self.index))
 
+  @position.setter
+  def position(self, value):
+     self.lmp.set("atom", self.index, "x", value[0])
+     self.lmp.set("atom", self.index, "y", value[1])
+     self.lmp.set("atom", self.index, "z", value[2])
+
   @property
   def velocity(self):
     return (self.lmp.eval("vx[%d]" % self.index),
@@ -399,6 +411,31 @@ 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))
+
+  @position.setter
+  def position(self, value):
+     self.lmp.set("atom", self.index, "x", value[0])
+     self.lmp.set("atom", self.index, "y", value[1])
+
+  @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)
@@ -608,6 +645,10 @@ class PyLammps(object):
       with OutputCapture() as capture:
         self.lmp.command(' '.join(cmd_args))
         output = capture.output
+
+      if 'verbose' in kwargs and kwargs['verbose']:
+        print(output)
+
       lines = output.splitlines()
 
       if len(lines) > 1:
@@ -662,3 +703,7 @@ class IPyLammps(PyLammps):
     self.write_dump(*cmd_args)
     from IPython.core.display import Image
     return Image('snapshot.png')
+
+  def video(self, filename):
+    from IPython.display import HTML
+    return HTML("<video controls><source src=\"" + filename + "\"></video>")