# coding: utf-8 import sys import numpy as np import os import pymses import matplotlib as mpl if os.environ.get("DISPLAY", "") == "": print("No display found. Using non-interactive Agg backend") mpl.use("Agg") import pylab as P import glob as glob import pickle as pickle from pymses.sources.ramses import output from pymses.analysis import Camera, raytracing, slicing from pymses.filters import CellsToPoints from pymses.analysis import ScalarOperator, FractionOperator, MaxLevelOperator # extension for out files out_ext = ".jpeg" P.rcParams["image.cmap"] = "plasma" P.rcParams["savefig.dpi"] = 400 def make_image_disk( path, num, path_out=None, order="<", force=False, tag="", vel_red=20, map_size=512, put_title=True, cpuamr=False, pos_star=np.array([1.0, 1.0, 1.0]), interactive=False, ): """ Make several useful image of an output of a simulation Parameters ---------- path path of the Ramses output num Ramses output number path_out path of the pipeline outputb order '<' or '>' TODO force if set, erase any existing pipeline output files tag string to add to the output name vel_red number of point where velocity should be plot in the slices map_size size of the map cpuamr plot also levels and cpus at each step """ ro = pymses.RamsesOutput(path, num, order=order) amr = ro.amr_source(["rho", "vel", "P"]) rad = 0.5 center = [0.5, 0.5, 0.5] make_image_aux( amr, ro, center, rad, num, path, force=force, path_out=path_out, map_size=map_size, vel_red=vel_red, tag=tag, cpuamr=cpuamr, put_title=put_title, pos_star=pos_star, interactive=interactive, ) def make_image_aux( amr, ro, center, radius, num, path, force=False, path_out=None, map_size=512, vel_red=20, tag="", cpuamr=False, pos_star=np.array([1.0, 1.0, 1.0]), put_title=True, interactive=False, ): """ Make several useful image of an output of a simulation, auxillary function Parameters ---------- amr ro pymses.RamsesOutput object center 3D array for coordinates center num output number path_out path of the pipeline output force if set, erase any existing pipeline output files tag string to add to the output name vel_red number of point where velocity should be plot in the slices map_size size of the map cpuamr plot also levels and cpus at each step """ lbox = ro.info["boxlen"] # boxlen in codeunits lbox_units = lbox G = 1.0 # Gravitational constant ntick = 6 title_ax = {"x": "x (code)", "y": "y (code)", "z": "z (code)"} time = ro.info["time"] # time in codeunits title = "t=" + str(time)[0:5] + " (code)" if path_out is not None: directory = path_out else: directory = path name = directory + "/coldens_z" + "_" + tag + "_" + format(num, "05") + out_ext if len(glob.glob(name)) == 1 and not force: return rho_op = ScalarOperator(lambda dset: dset["rho"], ro.info["unit_density"]) rt = raytracing.RayTracer(amr, ro.info, rho_op) axes_los = ["x", "y", "z"] # Line of sight axes axes_h = ["y", "x", "x"] # Horizontal axes axes_v = ["z", "z", "y"] # Vertical axes ax_nb = {"x": 0, "y": 1, "z": 2} image_names = ["coldens", "rho", "T"] for i, ax_los in enumerate(axes_los): ax_h = axes_h[i] ax_v = axes_v[i] cam = Camera( center=center, line_of_sight_axis=ax_los, region_size=[2.0 * radius, 2.0 * radius], distance=radius, far_cut_depth=radius, up_vector=ax_v, map_max_size=map_size, ) datamap = rt.process(cam, surf_qty=True) # Column density dmap_col = datamap.map.T * lbox map_col = np.log10(dmap_col) if interactive: P.figure() else: P.close() 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=ax_h, nbins=ntick) P.locator_params(axis=ax_v, nbins=ntick) if put_title: P.title(title) P.xlabel(title_ax[ax_h]) P.ylabel(title_ax[ax_v]) cbar = P.colorbar(im) cbar.set_label(r"$log(N)$ (code)") name = directory + "/coldens_" + ax_los + "_" + tag + "_" + format(num, "05") name_im = name + out_ext if interactive: P.figure() else: P.savefig(name_im) P.close() # Rho slice dmap_rho = slicing.SliceMap(amr, cam, rho_op, z=0.0) map_rho = np.log10(dmap_rho.map) map_rho = map_rho.T vh_op = ScalarOperator( lambda dset: dset["vel"][:, ax_nb[ax_h]], ro.info["unit_velocity"] ) dmap_vh = slicing.SliceMap(amr, cam, vh_op, z=0.0) map_vh_red = dmap_vh.map[ ::vel_red, ::vel_red ] # take only a subset of velocities map_vh_red = map_vh_red.T vv_op = ScalarOperator( lambda dset: dset["vel"][:, ax_nb[ax_v]], ro.info["unit_velocity"] ) dmap_vv = slicing.SliceMap(amr, cam, vv_op, z=0.0) map_vv_red = dmap_vv.map[::vel_red, ::vel_red] map_vv_red = map_vv_red.T 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=ax_h, nbins=ntick) P.locator_params(axis=ax_v, nbins=ntick) nh = map_vh_red.shape[0] nv = map_vv_red.shape[1] vec_h = ( np.arange(nh) * 2.0 / nh * radius - radius + center[0] + radius / nh ) * lbox_units vec_v = ( np.arange(nv) * 2.0 / nv * radius - radius + center[1] + radius / nv ) * lbox_units hh, vv = np.meshgrid(vec_h, vec_v) max_v = np.max(np.sqrt(map_vh_red ** 2 + map_vv_red ** 2)) Q = P.quiver(hh, vv, map_vh_red, map_vv_red, units="width") P.quiverkey( Q, 0.7, 0.95, max_v, r"$" + str(max_v)[0:4] + "$ (code)", labelpos="E", coordinates="figure", ) if put_title: P.title(title) P.xlabel(title_ax[ax_h]) P.ylabel(title_ax[ax_v]) cbar = P.colorbar(im) cbar.set_label(r"$log(n)$ (code)") name = directory + "/rho_" + ax_los + "_" + tag + "_" + format(num, "05") name_im = name + out_ext if interactive: P.figure() else: P.savefig(name_im) P.close() P_op = ScalarOperator(lambda dset: dset["P"], ro.info["unit_pressure"]) dmap_P = slicing.SliceMap(amr, cam, P_op, z=0.0) dmap_T = dmap_P.map.T / dmap_rho.map.T map_T = np.log10(dmap_T) 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(title) P.xlabel(title_ax[ax_h]) P.ylabel(title_ax[ax_v]) cbar = P.colorbar(im) cbar.set_label(r"$log(T) \, (K)$") name = directory + "/T_" + ax_los + "_" + tag + "_" + format(num, "05") name_im = name + out_ext if interactive: P.figure() else: P.savefig(name_im) P.close() # Toomre parameter if ax_los == "z": def omega_rho_func(dset): pos = dset.get_cell_centers() pos = pos - (pos_star / lbox) xx = pos[:, :, 0] yy = pos[:, :, 1] rc = np.sqrt(xx ** 2 + yy ** 2) # cylindrical radius vx = dset["vel"][:, :, 0] vy = dset["vel"][:, :, 1] omega_rho = 1.0 / rc ** 2 omega_rho = omega_rho * dset["rho"] vyx = vy * xx vxy = vx * yy omega_rho = omega_rho * (vyx - vxy) return omega_rho omega_op = FractionOperator( omega_rho_func, lambda dset: dset["rho"], 1.0 / ro.info["unit_time"] ) cs_op = FractionOperator( lambda dset: dset["P"], lambda dset: dset["rho"], ro.info["unit_velocity"], ) rt_omega = raytracing.RayTracer(amr, ro.info, omega_op) rt_cs = raytracing.RayTracer(amr, ro.info, cs_op) dmap_omega = rt_omega.process(cam) dmap_cs = rt_cs.process(cam) dmap_Q = (lbox * dmap_cs.map.T) * dmap_omega.map.T / (np.pi * G * dmap_col) map_Q = np.log10(dmap_Q) im = P.imshow( map_Q, 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(title) P.xlabel(title_ax[ax_h]) P.ylabel(title_ax[ax_v]) cbar = P.colorbar(im) cbar.set_label(r"$log(Q)$") name = directory + "/Q_" + ax_los + "_" + tag + "_" + format(num, "05") name_im = name + out_ext if interactive: P.figure() else: P.savefig(name_im) P.close() if cpuamr: level_op = MaxLevelOperator() amr.set_read_levelmax(20) rt_level = raytracing.RayTracer(amr, ro.info, level_op) datamap = rt_level.process(cam, surf_qty=True) map_level = datamap.map.T 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(title) P.xlabel(title_ax[ax_h]) P.ylabel(title_ax[ax_v]) cbar = P.colorbar(im) cbar.set_label(r"level") name = directory + "/level_" + ax_los + "_" + tag + "_" + format(num, "05") name_im = name + out_ext if interactive: P.figure() else: P.savefig(name_im) P.close() cpu_op = ScalarOperator( lambda dset: dset.icpu * (np.ones(dset["P"].shape)), ro.info["unit_pressure"], ) rt_cpu = raytracing.RayTracer(amr, ro.info, cpu_op) datamap = rt_cpu.process(cam, surf_qty=True) map_cpu = datamap.map.T 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(title) P.xlabel(title_ax[ax_h]) P.ylabel(title_ax[ax_v]) cbar = P.colorbar(im) cbar.set_label(r"cpu") name = directory + "/cpu_" + ax_los + "_" + tag + "_" + format(num, "05") name_im = name + out_ext if interactive: P.figure() else: P.savefig(name_im) P.close() def disk_prop( path_in, num, path_out=None, force=False, nb_bin=20, rad_ext=1.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) time = ro.info["time"] # time in codeunits # 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 * lbox # 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 G = 1.0 cs = np.sqrt(cells["P"] / cells["rho"]) # 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] v_kepl = np.sqrt(mass_star * G / rc_disk) total_mass_disk = np.sum(rho_disk * dvol_disk) # Initialize binned quantities 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_kepl_rad = np.zeros(nb_bin - 1) v_rad_rad = np.zeros(nb_bin - 1) alpha_rey_rad = np.zeros(nb_bin - 1) Q_kepl_rad = np.zeros(nb_bin - 1) for i in range(nb_bin - 1): mask_bin = (rc_disk > rad[i]) & (rc_disk < rad[i + 1]) print( "Bin #{} : {} cells between {} and {}".format( i, 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] = np.sum(press_disk[mask_bin] * dvol_disk[mask_bin]) / np.sum( rho_disk[mask_bin] * dvol_disk[mask_bin] ) # 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) ** 3 / ((rad[i + 1] - rad[i]) * (rad[i + 1] + rad[i]) * np.pi) ) 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]) ) v_kepl_rad[i] = np.sum( v_kepl[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin] ) / np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin]) # 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 v_kepl_rad = v_kepl_rad Q_kepl_rad = cs_rad * v_az_rad / (np.pi * G * coldens_rad * rad[0 : nb_bin - 1]) prop_disk = { "time": time, "tot_mass": total_mass_disk, "rad": rad[0 : nb_bin - 1], "center": pos_star, "alpha_rey": alpha_rey_rad, "v_rad": v_rad_rad, "v_az": v_az_rad, "v_kepl": v_kepl_rad, "coldens": coldens_rad, "rho": rho_rad, "press": press_rad, "temp": temp_rad, "cs": cs_rad, "Q_kepl": Q_kepl_rad, } # store the results f = open(name_save, "w") pickle.dump(prop_disk, f) f.close() def plot_disk_prop(path, num, force=False, tag="", interactive=False): """ 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 """ # 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: raise ("no pickle file for disk properties. Run single_z_disk_prop") f = open(name_save, "r") prop_disk = pickle.load(f) f.close() # Check if the output file exists, and exit if it is the case name_save = path + "/rho_disk_r_" + str(num).zfill(5) + out_ext if not force and len(glob.glob(name_save)) != 0: return time = prop_disk["time"] mass = prop_disk["tot_mass"] title = "t=" + str(time)[0:5] + " (code)" if interactive: P.figure() else: P.close() P.xscale("log") P.plot(prop_disk["rad"], np.log10(prop_disk["rho"]), color="k", linewidth=2) P.ylabel(r"$\log(n) \, (code)$") P.xlabel("disk radius") P.title(title) if interactive: P.figure() else: P.savefig(path + "/rho_disk_r_" + str(num).zfill(5) + out_ext) P.close() P.xscale("log") P.plot(prop_disk["rad"], np.log10(prop_disk["temp"]), color="k", linewidth=2) P.ylabel(r"$\log(T) \, (K)$") P.xlabel("disk radius") P.title(title) if interactive: P.figure() else: P.savefig(path + "/T_disk_r_" + str(num).zfill(5) + out_ext) P.close() P.xscale("log") P.yscale("symlog", linthreshy=0.01) P.plot((prop_disk["rad"]), ((prop_disk["v_rad"])), color="k", linewidth=2) P.plot((prop_disk["rad"]), ((prop_disk["v_kepl"])), color="b", linewidth=2) P.plot((prop_disk["rad"]), (abs(prop_disk["v_az"])), color="r", linewidth=2) P.plot((prop_disk["rad"]), ((prop_disk["cs"])), color="c", linewidth=2) P.legend((r"$v_r$", r"$v_{kepl}$", r"$v_\phi$", r"$c_s$"), loc="upper right") P.ylabel(r"$V \, (km s^{-1})$") P.xlabel("disk radius") if interactive: P.figure() else: P.savefig(path + "/V_disk_r_" + str(num).zfill(5) + out_ext) P.close() P.xscale("log") P.plot(prop_disk["rad"], np.log10(prop_disk["coldens"]), color="k", linewidth=2) P.ylabel(r"$\log(N) \, (cm^{-2})$") P.xlabel("disk radius ") P.title(title) if interactive: P.figure() else: P.savefig(path + "/coldens_disk_r_" + str(num).zfill(5) + out_ext) P.close() # Alpha P.xscale("log") P.yscale("log") P.ylim([1e-5, 1.0]) P.plot(prop_disk["rad"], abs(prop_disk["alpha_rey"]), color="b", linewidth=2) P.plot(prop_disk["rad"], abs(prop_disk["alpha_rey"]), color="b", linewidth=2) P.ylabel(r"$\alpha$") P.xlabel("disk radius ") P.title(title) if interactive: P.figure() else: P.savefig(path + "/alpha_disk_r_" + str(num).zfill(5) + out_ext) P.close() # Q P.ylim([0, 10.0]) P.plot(prop_disk["rad"], abs(prop_disk["Q_kepl"]), color="b", linewidth=2) P.ylabel(r"$Q$") P.xlabel("disk radius ") P.title(title + ", mass of disk = {} (code)".format(mass)) if interactive: pass else: P.savefig(path + "/Q_r_" + str(num).zfill(5) + out_ext) P.close() if pdf: P.savefig(path + "/Q_r_" + str(num).zfill(5) + ".pdf") P.savefig(path + "/Q_r_" + str(num).zfill(5) + ".jpeg")