diff --git a/extract_velcube.py b/extract_velcube.py new file mode 100644 index 0000000..0174174 --- /dev/null +++ b/extract_velcube.py @@ -0,0 +1,45 @@ +from snapshotprocessor import SnapshotProcessor, U +import pandas as pd +import os + +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.run}_{pp.num}_velrho.data", vel, dens) + + +def extract(path, snap_number): + pp = SnapshotProcessor(path, snap_number, params="../turbox_params.yml") + extract_from_pp(pp) + +