Added paraview python toolkit

python/cosmotool/ctpv.py

@@ -0,0 +1,106 @@
+import numpy as np
+from contextlib import contextmanager
+
+class ProgrammableDensityLoad(object):
+
+  @staticmethod
+  def main_script(source, density, extents=None, aname="default"):
+    import vtk
+    from vtk.util import numpy_support
+  
+    Nx, Ny, Nz = density.shape
+
+    ido = source.GetOutput()
+    ido.SetDimensions(Nx, Ny, Nz)
+    if not extents is None:
+      origin = extents[:6:2]
+      spacing = (extents[1]-extents[0])/Nx, (extents[3]-extents[2])/Ny, (extents[5]-extents[4])/Nz
+    else:
+      origin = (-1, -1, -1)
+      spacing = 2.0 / Nx, 2.0/Ny, 2.0/Nz
+      
+    ido.SetOrigin(*origin)
+    ido.SetSpacing(*spacing)
+    ido.SetExtent([0,Nx-1,0,Ny-1,0,Nz-1])
+    arr = numpy_support.numpy_to_vtk(density.transpose().astype(np.float64).ravel(), deep=1)
+    arr.SetName(aname)
+    ido.GetPointData().AddArray(arr)
+    
+  @staticmethod
+  def request_information(source, density=None, dims=None):
+    import vtk
+    
+    Nx = Ny = Nz = None
+    if not density is None:
+      Nx, Ny, Nz = density.shape
+    elif not dims is None:
+      Nx, Ny, Nz = dims
+    else:
+      raise ValueError("Need at least a density or dims")
+
+    source.GetExecutive().GetOutputInformation(0).Set(vtk.vtkStreamingDemandDrivenPipeline.WHOLE_EXTENT(), 0, Nx-1, 0, Ny-1, 0, Nz-1)
+
+  @staticmethod
+  def prepare_timesteps_info(algorithm, timesteps):
+
+    def SetOutputTimesteps(algorithm, timesteps):
+      executive = algorithm.GetExecutive()
+      outInfo = executive.GetOutputInformation(0)
+      outInfo.Remove(executive.TIME_STEPS())
+      for timestep in timesteps:
+        outInfo.Append(executive.TIME_STEPS(), timestep)
+      outInfo.Remove(executive.TIME_RANGE())
+      outInfo.Append(executive.TIME_RANGE(), timesteps[0])
+      outInfo.Append(executive.TIME_RANGE(), timesteps[-1])
+
+    SetOutputTimesteps(algorithm, timesteps)
+
+  @staticmethod
+  @contextmanager
+  def get_timestep(algorithm):
+
+    def GetUpdateTimestep(algorithm):
+      """Returns the requested time value, or None if not present"""
+      executive = algorithm.GetExecutive()
+      outInfo = executive.GetOutputInformation(0)
+      if not outInfo.Has(executive.UPDATE_TIME_STEP()):
+          return None
+      return outInfo.Get(executive.UPDATE_TIME_STEP())
+
+    # This is the requested time-step. This may not be exactly equal to the 
+    # timesteps published in RequestInformation(). Your code must handle that
+    # correctly
+    req_time = GetUpdateTimestep(algorithm)
+
+    output = algorithm.GetOutput()
+
+    yield req_time
+    
+    # Now mark the timestep produced.
+    output.GetInformation().Set(output.DATA_TIME_STEP(), req_time)
+
+    
+def load_borg(pdo, restart_name, mcmc_name, info=False, aname="BORG"):
+  import h5py as h5
+  
+  with h5.File(restart_name) as f:
+    N0 = f["/info/scalars/N0"][:]
+    N1 = f["/info/scalars/N1"][:]
+    N2 = f["/info/scalars/N2"][:]
+    L0 = f["/info/scalars/L0"][:]
+    L1 = f["/info/scalars/L1"][:]
+    L2 = f["/info/scalars/L2"][:]
+    c0 = f["/info/scalars/corner0"][:]
+    c1 = f["/info/scalars/corner1"][:]
+    c2 = f["/info/scalars/corner2"][:]
+
+  if not info:
+    with h5.File(mcmc_name) as f:
+      d = f["/scalars/BORG_final_density"][:]+1
+                
+    ProgrammableDensityLoad.main_script(pdo, d, extents=[c0,c0+L0,c1,c1+L1,c2,c2+L2], aname=aname)
+  else:
+    ProgrammableDensityLoad.request_information(pdo, dims=[N0,N1,N2])
+    
+    
+

src/loadGadget.cpp

@@ -395,7 +395,7 @@ SimuData *CosmoTool::loadGadgetMulti(const char *fname, int id,
             } else {
               for(int n = 0; n < h.npart[k]; n++)
 	        {
-            if ((n%1000000)==0) cout << n << endl;
+//            if ((n%1000000)==0) cout << n << endl;
                   data->Mass[l++] = h.mass[k];
                 }
             }