55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
|
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
|