diff --git a/python/cosmotool/ctpv.py b/python/cosmotool/ctpv.py index e97efe4..820141c 100644 --- a/python/cosmotool/ctpv.py +++ b/python/cosmotool/ctpv.py @@ -1,6 +1,28 @@ import numpy as np from contextlib import contextmanager +class ProgrammableParticleLoad(object): + + @staticmethod + def main_script(source, particles, aname="default"): + import vtk + from vtk.util import numpy_support as ns + + out = source.GetOutput() + vv = vtk.vtkPoints() + + assert len(particles.shape) == 2 + assert particles.shape[1] == 3 + + vv.SetData(ns.numpy_to_vtk(np.ascontiguousarray(particles.astype(np.float64)), deep=1)) + vv.SetDataTypeToDouble() + + out.SetPoints(vv) + + @staticmethod + def request_information(source): + pass + class ProgrammableDensityLoad(object): @staticmethod @@ -103,4 +125,18 @@ def load_borg(pdo, restart_name, mcmc_name, info=False, aname="BORG"): ProgrammableDensityLoad.request_information(pdo, dims=[N0,N1,N2]) +def load_borg_galaxies(pdo, restart_name, cid=0, info=False, aname="Galaxies"): + import h5py as h5 + + with h5.File(restart_name) as f: + gals = f['/info/galaxy_catalog_%d/galaxies' % cid] + ra = gals['phi'][:] + dec = gals['theta'][:] + r = gals['r'][:] + if not info: + x = r * np.cos(ra)*np.cos(dec) + y = r * np.sin(ra)*np.cos(dec) + z = r * np.sin(dec) + parts = np.array([x,y,z]).transpose() + ProgrammableParticleLoad.main_script(pdo, parts)