Noe Brucy
51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
from snapshotprocessor import SnapshotProcessor, U
|
|
|
|
|
|
def get_velocity_cubes(pp, unit=None):
|
|
velcubes = [None, None, None]
|
|
for i in range(3):
|
|
velcubes[i] = pp.datacube(getter=lambda dset: dset["vel"][..., i])
|
|
if unit is not None:
|
|
velcubes[i] *= pp.info["unit_velocity"].express(unit)
|
|
return velcubes
|
|
|
|
|
|
def get_density_cube(pp, unit=None):
|
|
dens_cube = pp.datacube(getter=lambda dset: dset["rho"])
|
|
if unit is not None:
|
|
dens_cube *= pp.info["unit_density"].express(unit)
|
|
return dens_cube
|
|
|
|
|
|
def write_data(filename, vel, dens):
|
|
# write fields to ramses frig readable ascii file
|
|
f = open(filename, "w")
|
|
dummy = 1
|
|
size = vel[0].shape[0]
|
|
f.write(
|
|
"{:8}{:13.5f}{:13.5f}{:13.5f}{:13.5f}\n".format(
|
|
size, dummy, dummy, dummy, dummy
|
|
)
|
|
)
|
|
vx, vy, vz = vel
|
|
# This strange order matches the one in the galbox condinit
|
|
for z in range(size):
|
|
for y in range(size):
|
|
for x in range(size):
|
|
f.write(
|
|
"{:13.5f}{:13.5f}{:13.5f}{:13.5f}\n".format(
|
|
vx[x, y, z], vy[x, y, z], vz[x, y, z], dens[x, y, z]
|
|
)
|
|
)
|
|
|
|
|
|
def extract_from_pp(pp):
|
|
vel = get_velocity_cubes(pp, unit=U.km_s)
|
|
dens = get_density_cube(pp, unit=U.H_cc)
|
|
write_data(f"{pp.path_out}/{pp.run}_{pp.num}_velrho.data", vel, dens)
|
|
|
|
|
|
def extract(path, snap_number, params=None):
|
|
pp = SnapshotProcessor(path, snap_number, params=params)
|
|
extract_from_pp(pp)
|