Noe Brucy
46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
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.path_out}/{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)
|
|
|
|
|