From e23bcd7df32e15240c536bb81ae23662f2798674 Mon Sep 17 00:00:00 2001
From: Guilhem Lavaux <lavaux@iap.fr>
Date: Fri, 5 Jan 2018 19:14:43 +0100
Subject: [PATCH] Add smooth particle density wrapper for simple3DFilter

---
 python/cosmotool/__init__.py  |  1 +
 python/cosmotool/config.py.in |  2 +-
 python/cosmotool/smooth.py    | 54 +++++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 python/cosmotool/smooth.py

diff --git a/python/cosmotool/__init__.py b/python/cosmotool/__init__.py
index 63a339e..c5b1c79 100644
--- a/python/cosmotool/__init__.py
+++ b/python/cosmotool/__init__.py
@@ -15,6 +15,7 @@ except:
 from .simu import loadRamsesAll, simpleWriteGadget, SimulationBare
 from .timing import time_block, timeit, timeit_quiet
 from .bispectrum import bispectrum, powerspectrum
+from .smooth import smooth_particle_density
 
 try:
   from .fftw import CubeFT
diff --git a/python/cosmotool/config.py.in b/python/cosmotool/config.py.in
index a59d25e..7603b85 100644
--- a/python/cosmotool/config.py.in
+++ b/python/cosmotool/config.py.in
@@ -1 +1 @@
-install_prefix=@CMAKE_INSTALL_PREFIX@
+install_prefix="@CMAKE_INSTALL_PREFIX@"
diff --git a/python/cosmotool/smooth.py b/python/cosmotool/smooth.py
new file mode 100644
index 0000000..b28d329
--- /dev/null
+++ b/python/cosmotool/smooth.py
@@ -0,0 +1,54 @@
+from .config import install_prefix
+import subprocess
+import os
+from tempfile import TemporaryDirectory
+import h5py as h5
+import numpy as np
+import weakref
+
+def smooth_particle_density(position, velocities=None, radius=1e6, boxsize=None, resolution=128, center=None ):
+  if len(position.shape) != 2:
+     raise ValueError("Invalid position array shape")
+
+  if velocities is None:
+    if position.shape[1] != 6:
+      raise ValueError("Position must be phase space if no velocities are given")
+
+
+  if boxsize is None:
+    raise ValueError("Need a boxsize")
+
+  cx,cy,cz=center
+  with TemporaryDirectory() as tmpdir:
+    h5_file = os.path.join(tmpdir, 'particles.h5')
+    with h5.File(h5_file, mode="w") as f:
+      data = f.create_dataset('particles', shape=(position.shape[0],7), dtype=np.float32)
+      data[:,:3] = position[:,:3]
+      if velocities is not None: 
+        data[:,3:6] = velocities[:,:3]
+      else:
+        data[:,3:6] = position[:,3:]
+      data[:,6] = 1
+
+
+    ret =  \
+       subprocess.run([
+           os.path.join(install_prefix,'bin','simple3DFilter'),
+           h5_file,
+           str(radius),
+           str(boxsize),
+           str(resolution),
+           str(cx), str(cy), str(cz)
+       ], cwd=tmpdir)
+
+    
+    f0 = h5.File(os.path.join(tmpdir,'fields.h5'), mode="r")
+    def cleanup_f0():
+      f0.close()
+
+    class Dict(dict):
+      pass
+
+    t = Dict(rho=f0['density'], p=[f0['p0'], f0['p1'], f0['p2']])
+    weakref.finalize(t, cleanup_f0)
+    return t