diff --git a/python/cosmotool/ctpv.py b/python/cosmotool/ctpv.py new file mode 100644 index 0000000..e97efe4 --- /dev/null +++ b/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]) + + + diff --git a/src/loadGadget.cpp b/src/loadGadget.cpp index 4a3952e..4ccbe00 100644 --- a/src/loadGadget.cpp +++ b/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]; } }