Add smooth particle density wrapper for simple3DFilter

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

python/cosmotool/config.py.in

@@ -1 +1 @@
-install_prefix=@CMAKE_INSTALL_PREFIX@
+install_prefix="@CMAKE_INSTALL_PREFIX@"

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