add extract_velcube

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)
+    
+