Files
pipeline/disk_postprocess.py
T
2020-12-14 16:17:00 +01:00

659 lines
19 KiB
Python

# coding: utf-8
import sys
import numpy as np
import os
import pymses
import pylab as P
import glob as glob
import pickle as pickle
import module_extract as me
from pymses.filters import CellsToPoints
from pymses.sources.ramses import output
from pymses.analysis import Camera, raytracing, slicing
def make_image_temp(
amr,
ro,
center,
radius,
num,
path,
i_im=0,
force=False,
path_out=None,
col_dens_only=False,
map_size=512,
vel_red=20,
im_pos=0.0,
tag="",
cpuamr=False,
put_title=True,
mass_norm=1.0,
AU_units=False,
):
"""
Make several useful image of an output of a simulation
Parameters
----------
amr
ro pymses.RamsesOutput object
center 3D array for coordinates center
num output number
"""
lbox = ro.info["boxlen"] # boxlen in codeunits (=>pc)
(
AU,
pc,
Ms,
Myr,
scale_n,
scale_d,
scale_t,
scale_l,
scale_v,
scale_T2,
scale_ener,
scale_mag,
microG,
km_s,
Cwnm,
scale_mass,
unit_col,
lbox_pc,
) = normalisation(ro)
ntick = 6
if AU_units:
units = pc / AU
lbox_units = lbox * units
titre_x = "x (AU)"
titre_y = "y (AU)"
titre_z = "z (AU)"
else:
units = 1.0
lbox_units = lbox
titre_x = "x (pc)"
titre_y = "y (pc)"
titre_z = "z (pc)"
lbox_cm = lbox * pc # lbox in cm
time = ro.info["time"] # time in codeunits
time = time * scale_t / Myr
titre = "t=" + str(time)[0:5] + " (Myr)"
if path_out is not None:
directory = path_out
else:
directory = path
rt = raytracing.RayTracer(amr, ro.info, rho_op)
rho_op = ScalarOperator(lambda dset: dset["rho"], ro.info["unit_density"])
ax_names = ["x", "y", "z"]
up_axes = ["z", "z", "x"]
other_axes = ["z", "x", "y"]
image_names = ["coldens", "rho", "T"]
for i, ax_name in enumerate(ax_names):
cam = Camera(
center=center,
line_of_sight_axis=ax_name,
region_size=[2.0 * radius, 2.0 * radius],
distance=radius,
far_cut_depth=radius,
up_vector=up_axes[i],
map_max_size=map_size,
)
datamap = rt.process(cam, surf_qty=True)
map_col = np.log10(datamap.map.T * lbox_cm)
P.clf()
im = P.imshow(
map_col,
extent=[
(-radius + center[0]) * lbox_units,
(radius + center[0]) * lbox_units,
(-radius + center[1]) * lbox_units,
(radius + center[1]) * lbox_units,
],
origin="lower",
)
P.locator_params(axis=other_axes[i], nbins=ntick)
P.locator_params(axis=up_axes[i], nbins=ntick)
if put_title:
P.title(titre)
P.xlabel(titre_x)
P.ylabel(titre_y)
cbar = P.colorbar(im)
cbar.set_label(r"$log(N) \, cm^{-2}$")
name = directory + "/coldens_" + ax_name + "_" + tag + +"_" + format(num, "05")
name_im = name + ".jpeg"
P.savefig(name_im)
dmap_rho = slicing.SliceMap(amr, cam, rho_op, z=0.0)
map_rho = np.log10(dmap_rho.map)
map_rho = map_rho.T
vx_op = ScalarOperator(lambda dset: dset["vel"][:, 0], ro.info["unit_velocity"])
dmap_vx = slicing.SliceMap(amr, cam_z, vx_op, z=0.0)
map_vx_red = dmap_vx.map[::vel_red, ::vel_red]
map_vx_red = map_vx_red.T
vy_op = ScalarOperator(lambda dset: dset["vel"][:, 1], ro.info["unit_velocity"])
dmap_vy = slicing.SliceMap(amr, cam, vy_op, z=0.0)
map_vy_red = dmap_vy.map[::vel_red, ::vel_red]
map_vy_red = map_vy_red.T
P.clf()
im = P.imshow(
map_rho,
extent=[
(-radius + center[0]) * lbox_units,
(radius + center[0]) * lbox_units,
(-radius + center[1]) * lbox_units,
(radius + center[1]) * lbox_units,
],
origin="lower",
)
P.locator_params(axis=other_axes[i], nbins=ntick)
P.locator_params(axis=up_axes[i], nbins=ntick)
nx = map_vx_red.shape[0]
ny = map_vx_red.shape[1]
vec_x = (
np.arange(nx) * 2.0 / nx * radius - radius + center[0] + radius / nx
) * lbox_units
vec_y = (
np.arange(ny) * 2.0 / ny * radius - radius + center[1] + radius / nx
) * lbox_units
xx, yy = np.meshgrid(vec_x, vec_y)
map_vx_red = map_vx_red * scale_v / km_s
map_vy_red = map_vy_red * scale_v / km_s
max_v = np.max(np.sqrt(map_vx_red ** 2 + map_vy_red ** 2))
Q = P.quiver(xx, yy, map_vx_red, map_vy_red, units="width")
P.quiverkey(
Q,
0.7,
0.95,
max_v,
r"$" + str(max_v)[0:4] + "\, km \, s^{-1}$",
labelpos="E",
coordinates="figure",
)
if put_title:
P.title(titre)
P.xlabel(titre_x)
P.ylabel(titre_y)
cbar = P.colorbar(im)
cbar.set_label(r"$log(n) \, (cm^{-3})$")
name = directory + "/rho_z" + "_" + tag + str(i_im) + "_" + format(num, "05")
name_im = name + ".jpeg"
P.savefig(name_im)
P_op = ScalarOperator(lambda dset: dset["P"], ro.info["unit_pressure"])
dmap_P = slicing.SliceMap(amr, cam_z, P_op, z=0.0)
P.clf()
map_T = np.log10(dmap_P.map.T / dmap_rho.map.T * scale_T2)
im = P.imshow(
map_T,
extent=[
(-radius + center[0]) * lbox_units,
(radius + center[0]) * lbox_units,
(-radius + center[1]) * lbox_units,
(radius + center[1]) * lbox_units,
],
origin="lower",
)
P.locator_params(axis="x", nbins=ntick)
P.locator_params(axis="y", nbins=ntick)
if put_title:
P.title(titre)
P.xlabel(titre_x)
P.ylabel(titre_y)
cbar = P.colorbar(im)
cbar.set_label(r"$log(T) \, (K)$")
name = directory + "/T_z" + "_" + tag + str(i_im) + "_" + format(num, "05")
name_im = name + ".jpeg"
P.savefig(name_im)
if cpuamr:
level_op = MaxLevelOperator()
amr.set_read_levelmax(20)
rt = raytracing.RayTracer(amr, ro.info, level_op)
datamap = rt.process(cam_z, surf_qty=True)
map_level = datamap.map.T
P.clf()
im = P.imshow(
map_level,
extent=[
(-radius + center[0]) * lbox_units,
(radius + center[0]) * lbox_units,
(-radius + center[1]) * lbox_units,
(radius + center[1]) * lbox_units,
],
origin="lower",
)
P.locator_params(axis="x", nbins=ntick)
P.locator_params(axis="y", nbins=ntick)
if put_title:
P.title(titre)
P.xlabel(titre_x)
P.ylabel(titre_y)
cbar = P.colorbar(im)
cbar.set_label(r"level")
name = (
directory + "/level_z" + "_" + tag + str(i_im) + "_" + format(num, "05")
)
name_im = name + ".jpeg"
P.savefig(name_im)
if ps:
P.savefig(name + ".ps")
if descrip is not None:
dd = {}
dd.update(
{
"name": "mean AMR level in the z-direction"
+ " ("
+ tag
+ str(i_im)
+ ")"
}
)
dd.update({"type": "image"})
dd.update({"description": "Mean AMR level in the z-direction."})
dd.update({"display": 0})
dd.update({"item": i_im})
dd.update({"tag": tag})
data = {"jpeg": name_im}
dd.update({"data": data})
list_im.append(dd)
cpu_op = ScalarOperator(
lambda dset: dset.icpu * (np.ones(dset["P"].shape)),
ro.info["unit_pressure"],
)
rt = raytracing.RayTracer(amr, ro.info, cpu_op)
datamap = rt.process(cam_z, surf_qty=True)
map_cpu = datamap.map.T
P.clf()
im = P.imshow(
map_cpu,
extent=[
(-radius + center[0]) * lbox_units,
(radius + center[0]) * lbox_units,
(-radius + center[1]) * lbox_units,
(radius + center[1]) * lbox_units,
],
origin="lower",
)
P.locator_params(axis="x", nbins=ntick)
P.locator_params(axis="y", nbins=ntick)
if put_title:
P.title(titre)
P.xlabel(titre_x)
P.ylabel(titre_y)
cbar = P.colorbar(im)
cbar.set_label(r"level")
name = (
directory + "/cpu_z" + "_" + tag + str(i_im) + "_" + format(num, "05")
)
name_im = name + ".jpeg"
P.savefig(name_im)
def disk_prop(
path_in,
num,
path_out=None,
force=False,
nb_bin=20,
rad_ext=100.0,
mass_star=1.0,
pos_star=np.array([1.0, 1.0, 1.0]),
):
"""
Compute properties of a disk in the plane (0,x,y)
with a protostar at the center of the box
The region of the disk is defined by its scale height
Parameters
----------
path_in path of the input data files (output of ramses)
num id of the output
path_out optional path to the output files
force if set, redo ouptut even if already done
nb_bin Number of radial bins
rad_ext Outer radius of the disk
pos_star position of the central protostar
mass_star mass of the central protostar
"""
# Set th output directory
if path_out is not None:
directory_out = path_out
else:
directory_out = path_in
# Check if the output file exists, and exit if it is the case
name_save = directory_out + "/prop_disk_" + str(num).zfill(5) + ".save"
if not force and len(glob.glob(name_save)) != 0:
return
# Compute the bins array
lrad = np.log10(rad_ext)
rad = np.logspace(lrad - 2.0, lrad, num=nb_bin)
# Get Ramses data
ro = pymses.RamsesOutput(path_in, num)
lbox = ro.info["boxlen"] # boxlen in codeunits (=>pc)
(
AU,
pc,
Ms,
Myr,
scale_n,
scale_d,
scale_t,
scale_l,
scale_v,
scale_T2,
scale_ener,
scale_mag,
microG,
km_s,
Cwnm,
scale_mass,
unit_col,
lbox_pc,
) = me.normalisation(ro)
time = ro.info["time"] * scale_t / Myr
# Get array of cell positions
amr = ro.amr_source(["rho", "vel", "Br", "Bl", "P"])
cell_source = CellsToPoints(amr)
cells = cell_source.flatten()
dx = cells.get_sizes()
pos = cells.points
# Get positions in the frame of the protostar
pos = pos - pos_star
# Get cylindrical radius
rc = np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2)
# Get velocities
vel = cells["vel"]
# Get radial component of velocity
norm_pos = rc
norm_pos[rc == 0] = 1.0e-10 # Avoid division per 0
v_rad = (pos[:, 0] * vel[:, 0] + pos[:, 1] * vel[:, 1]) / norm_pos
# Get azimuthal component of velocity
v_az = (pos[:, 0] * vel[:, 1] - pos[:, 1] * vel[:, 0]) / norm_pos
# Select cells that are actually in the disk, ie within the scale height
# TODO Check units
G = 6.8e-8
cs = np.sqrt(cells["P"] / cells["rho"]) * scale_v # sound velocity
height = cs * np.sqrt(rc ** 3 / (G * mass_star))
mask_pos = np.abs(pos[:, 2]) < height # condition on position
mask_dens = cells["rho"] > 1.0e6 # condition on density
mask = mask_pos | mask_dens
print("Number of selected cells ", np.sum(mask))
pos_disk = pos[mask]
rc_disk = rc[mask]
vel_disk = vel[mask]
rho_disk = cells["rho"][mask] # density
press_disk = cells["P"][mask] # pressure
dx_disk = dx[mask]
dvol_disk = dx_disk ** 3
v_rad_disk = v_rad[mask]
v_az_disk = v_az[mask]
# TODO Check what do that does
nzoom = 9
eps = 0.5 ** nzoom
# map_coldens , map_w13, xedges, yedges = me.make_hierarch_map(pos_disk_x,pos_disk_y,pos_disk_z,dx_disk,rho_disk,rho_disk,eps,center=[0.,0.,0.],make_image=do_plot,path_out=directory_out,tag='xy_'+ str(num).zfill(5))
# map_coldens , map_w13, xedges, yedges = me.make_hierarch_map(pos_disk_z,pos_disk_x,pos_disk_y,dx_disk,rho_disk,rho_disk,eps,center=[0.,0.,0.],make_image=do_plot,path_out=directory_out,tag='xz_'+ str(num).zfill(5))
# Initialize binned quantities
norm_rad = lbox * scale_l / AU # radius in AU
rdisk_AU = rc_disk * norm_rad
cs_rad = np.zeros(nb_bin - 1)
temp_rad = np.zeros(nb_bin - 1)
press_rad = np.zeros(nb_bin - 1)
rho_rad = np.zeros(nb_bin - 1)
coldens_rad = np.zeros(nb_bin - 1)
v_az_rad = np.zeros(nb_bin - 1)
v_rad_rad = np.zeros(nb_bin - 1)
alpha_rey_rad = np.zeros(nb_bin - 1)
for i in range(nb_bin - 1):
mask_bin = (rdisk_AU > rad[i]) & (rdisk_AU < rad[i + 1])
print(
"Bin {} cells between {} and {} AU".format(
np.sum(mask_bin), rad[i], rad[i + 1]
)
)
press_rad[i] = np.sum(press_disk[mask_bin] * dvol_disk[mask_bin]) / np.sum(
dvol_disk[mask_bin]
)
rho_rad[i] = np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin]) / np.sum(
dvol_disk[mask_bin]
)
temp_rad[i] = press_rad[i] / rho_rad[i]
# TODO verifier unites
# Surface of a bin : S = dr * 2 * pi * r with
# dr = rad[i + 1] - rad[i] and r = (rad[i + 1] + rad[i]) / 2.
coldens_rad[i] = (
np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin])
* (lbox * pc) ** 3
/ ((rad[i + 1] - rad[i]) * (rad[i + 1] + rad[i]) * np.pi * AU ** 2)
)
v_az_rad[i] = np.sum(
v_az_disk[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin]
) / np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin])
v_rad_rad[i] = np.sum(
v_rad_disk[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin]
) / np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin])
alpha_rey_rad[i] = (
(
np.sum(
v_az_disk[mask_bin]
* v_rad_disk[mask_bin]
* rho_disk[mask_bin]
* dvol_disk[mask_bin]
)
/ np.sum(dvol_disk[mask_bin] * press_disk[mask_bin])
- v_az_rad[i] * v_rad_rad[i] * rho_rad[i] / press_rad[i]
)
* v_az_rad[i]
/ abs(v_az_rad[i])
)
# Convert to good units (TODO check)
cs_rad = np.sqrt(temp_rad) * scale_v / km_s
temp_rad = temp_rad * scale_T2
press_rad = press_rad * scale_v ** 2 * scale_d
v_az_rad = v_az_rad * scale_v / km_s
v_rad_rad = v_rad_rad * scale_v / km_s
prop_disk = {
"time": time,
"rad_AU": rad[0 : nb_bin - 1],
"center": pos_star,
"alpha_rey": alpha_rey_rad,
"v_rad": v_rad_rad,
"v_az": v_az_rad,
"coldens": coldens_rad,
"rho": rho_rad,
"press": press_rad,
"temp": temp_rad,
"cs": cs_rad,
}
# store the results
f = open(name_save, "w")
pickle.dump(prop_disk, f)
f.close()
def plot_disk_prop(path, num, force=False, pdf=False, tag=""):
"""
Plot properties of a disk
num id of the ramses output
path path to the properties file
force if set, redo plots even if already done
pdf if set, do output in pdf as well
"""
# Load property file
name_save = path + "/prop_disk_" + str(num).zfill(5) + ".save"
# Check if the properties file exists
if len(glob.glob(name_save)) == 0:
throw("no pickle file for disk properties. Run single_z_disk_prop")
f = open(name_save, "r")
prop_disk = pickle.load(f)
f.close()
P.clf()
P.plot(
np.log10(prop_disk["rad_AU"]),
np.log10(prop_disk["rho"]),
color="k",
linewidth=2,
)
P.ylabel(r"$\log(n) \, (cm^{-3})$")
P.xlabel("disk radius (AU)")
if pdf:
P.savefig(path + "/rho_disk_r_" + str(num).zfill(5) + ".pdf")
P.savefig(path + "/rho_disk_r_" + str(num).zfill(5) + ".jpeg")
P.clf()
P.plot(
np.log10(prop_disk["rad_AU"]),
np.log10(prop_disk["temp"]),
color="k",
linewidth=2,
)
P.ylabel(r"$\log(T) \, (K)$")
P.xlabel("disk radius (AU)")
if pdf:
P.savefig(path + "/T_disk_r_" + str(num).zfill(5) + ".pdf")
P.savefig(path + "/T_disk_r_" + str(num).zfill(5) + ".jpeg")
P.clf()
P.xscale("log")
P.yscale("symlog", linthreshy=0.01)
P.plot((prop_disk["rad_AU"]), ((prop_disk["v_rad"])), color="k", linewidth=2)
P.plot((prop_disk["rad_AU"]), (abs(prop_disk["v_az"])), color="r", linewidth=2)
P.plot((prop_disk["rad_AU"]), ((prop_disk["cs"])), color="c", linewidth=2)
P.legend((r"$v_r$", r"$v_\phi$", r"$c_s$"), loc="upper right")
P.ylabel(r"$V \, (km s^{-1})$")
P.xlabel("disk radius (AU)")
if pdf:
P.savefig(path + "/V_disk_r_" + str(num).zfill(5) + ".pdf")
P.savefig(path + "/V_disk_r_" + str(num).zfill(5) + ".jpeg")
P.legend((r"$v_r$", r"$v_\phi$", r"$c_s$"), loc="upper right")
P.ylabel(r"$V \, (km s^{-1})$")
P.xlabel("disc radius (AU)")
if pdf:
P.savefig(path + "V_disk_r_" + str(num).zfill(5) + ".pdf")
P.savefig(path + "V_disk_r_" + str(num).zfill(5) + ".jpeg")
P.clf()
P.plot(
np.log10(prop_disk["rad_AU"]),
np.log10(prop_disk["coldens"]),
color="k",
linewidth=2,
)
P.ylabel(r"$\log(N) \, (cm^{-2})$")
P.xlabel("disk radius (AU)")
if pdf:
P.savefig(path + "/coldens_disk_r_" + str(num).zfill(5) + ".pdf")
P.savefig(path + "/coldens_disk_r_" + str(num).zfill(5) + ".jpeg")
P.clf()
P.xscale("log")
P.yscale("symlog", linthreshy=0.001)
P.plot(prop_disk["rad_AU"], prop_disk["alpha_rey"], color="b", linewidth=2)
P.plot(prop_disc["rad_AU"], prop_disc["alpha_rey"], color="b", linewidth=2)
P.ylabel(r"$\alpha}$")
P.xlabel("disk radius (AU)")
if pdf:
P.savefig(path + "/alpha_disk_r_" + str(num).zfill(5) + ".pdf")
P.savefig(path + "/alpha_disk_r_" + str(num).zfill(5) + ".jpeg")
if pdf:
P.savefig(path + "alpha_disk_r_" + str(num).zfill(5) + ".pdf")
P.savefig(path + "alpha_disk_r_" + str(num).zfill(5) + ".jpeg")