Merge branch 'master' of bitbucket.org:glavaux/cosmotool

python/cosmotool/ctpv.py

@@ -26,6 +26,55 @@ class ProgrammableParticleLoad(object):
   def request_information(source):
    pass
 
+
+
+class ProgrammableParticleHistoryLoad(object):
+
+  @staticmethod
+  def main_script(source, particles, velocities=None, aname="default",addtime=False):
+    import vtk
+    from vtk.util import numpy_support as ns
+  
+    out = source.GetOutput()
+    vv = vtk.vtkPoints()
+
+    assert len(particles.shape) == 3
+    assert particles.shape[2] == 3
+    
+    if not velocities is None:
+      for i,j in zip(velocities.shape,particles.shape):
+        assert i==j
+
+    Ntime,Npart,_ = particles.shape
+
+    vv.SetData(ns.numpy_to_vtk(np.ascontiguousarray(particles.reshape((Ntime*Npart,3)).astype(np.float64)), deep=1))
+    vv.SetDataTypeToDouble()
+
+    out.Allocate(1,1)
+    out.SetPoints(vv)
+
+    if not velocities is None:
+      print("Adding velocities")
+      vel_vtk = ns.numpy_to_vtk(np.ascontiguousarray(velocities.reshape((Ntime*Npart,3)).astype(np.float64)), deep=1)
+      vel_vtk.SetName("velocities")
+      out.GetPointData().AddArray(vel_vtk)
+
+    if addtime:
+      timearray = np.arange(Ntime)[:,None].repeat(Npart, axis=1).reshape(Ntime*Npart)
+      timearray = ns.numpy_to_vtk(np.ascontiguousarray(timearray.astype(np.float64)), deep=1)
+      timearray.SetName("timearray")
+      out.GetPointData().AddArray(timearray)
+
+    out.InsertNextCell(vtk.VTK_VERTEX, particles.shape[0], range(particles.shape[0]))
+    for p in range(Npart):
+      out.InsertNextCell(vtk.VTK_LINE, Ntime, range(p, p + Npart*Ntime, Npart) )
+
+
+  @staticmethod
+  def request_information(source):
+   pass
+
+
 class ProgrammableDensityLoad(object):
 
   @staticmethod