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