From 04ad4ff2664a50becc9cef76734f4e8d83552b1c Mon Sep 17 00:00:00 2001
From: Guilhem Lavaux <lavaux@iap.fr>
Date: Wed, 8 Jun 2016 21:26:59 +0200
Subject: [PATCH] Load particles

---
 python/cosmotool/ctpv.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

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)