diff --git a/disk_postprocess.py b/disk_postprocess.py new file mode 100644 index 0000000..3c961a5 --- /dev/null +++ b/disk_postprocess.py @@ -0,0 +1,664 @@ +import sys +import numpy as np +import os +import pymses +import pylab as P +import glob as glob +import pandas as pd +import pickle as pickle + + +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 ramses output 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 + if ( + not force + and len(glob.glob(directory_out + "/prop_disk_" + str(num).zfill(5) + ".save")) + != 0 + ): + return + + # Compute the bins array + lrad = np.log10(rad_ext) + rad = np.logspace(lrad - 2.0, lrad, num=nbin) + + # 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 or 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.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.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 = rad_disk * norm_rad + + cs_rad = np.zeros(nbin - 1) + temp_rad = np.zeros(nbin - 1) + press_rad = np.zeros(nbin - 1) + rho_rad = np.zeros(nbin - 1) + coldens_rad = np.zeros(nbin - 1) + v_az_rad = np.zeros(nbin - 1) + v_rad_rad = np.zeros(nbin - 1) + + for i in range(nbin - 1): + mask_bin = (rdisk_AU > rad[i]) or (rdisk_AU < rad[i + 1]) + + press_rad[i] = np.sum(press_disk[mask] * dvol_disk[mask]) / np.sum( + dvol_disk[mask] + ) + rho_rad[i] = np.sum(rho_disk[mask] * dvol_disk[mask]) / np.sum(dvol_disk[mask]) + temp_rad[i] = press_rad[i] / rho_rad[i] + + # TODO vérifier unités + # 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] * dvol_disk[mask]) + * (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] * rho_disk[mask] * dvol_disk[mask] + ) / np.sum(rho_disk[mask] * dvol_disk[mask]) + + v_rad_rad[i] = np.sum( + v_rad_disk[mask] * rho_disk[mask] * dvol_disk[mask] + ) / np.sum(rho_disk[mask] * dvol_disk[mask]) + + alpha_rey_rad[i] = ( + ( + np.sum( + v_az_disk[mask] + * v_rad_disk[mask] + * rho_disk[mask] + * dvol_disk[mask] + ) + / np.sum(dvol_disk[mask] * press_disk[mask]) + - v_az_rad[i] * v_rad_rad[i] * rho_rad[i] / press_rad[i] + ) + * v_az_rad[i] + / abs(v_az_rad[i]) + ) + + prop_disk = { + "time": time, + "rad_AU": rad[0 : nbin - 1], + "center": pos_mass, + "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, + } + + # 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 + + # store the results + name_save = directory_out + "/prop_disk_" + str(num).zfill(5) + ".save" + 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: + print("no pickle file for disk properties. Run single_z_disc_prop") + return + f = open(name_save, "r") + prop_disk = pickle.load(f) + f.close() + + P.clf() + P.plot( + np.log10(prop_disc["rad_AU"]), + np.log10(prop_disc["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_disc["rad_AU"]), + np.log10(prop_disc["temp"]), + color="k", + linewidth=2, + ) + P.ylabel(r"$\log(T) \, (K)$") + P.xlabel("disc 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_disc["rad_AU"]), ((prop_disc["v_rad"])), color="k", linewidth=2) + P.plot((prop_disc["rad_AU"]), (abs(prop_disc["v_az"])), color="r", linewidth=2) + P.plot((prop_disc["rad_AU"]), ((prop_disc["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("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_disc["rad_AU"]), + np.log10(prop_disc["coldens"]), + color="k", + linewidth=2, + ) + P.ylabel(r"$\log(N) \, (cm^{-2})$") + P.xlabel("disc 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_disc["rad_AU"], prop_disc["alpha_rey"], color="b", linewidth=2) + + P.legend((r"$\alpha _{Rey}$"), loc="upper right") + + P.ylabel(r"$\alpha}$") + P.xlabel("disc 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")