# coding: utf-8 import sys import os import pymses import numpy as np import matplotlib as mpl if os.environ.get('DISPLAY','') == '': print('No display found. Using non-interactive Agg backend') mpl.use('Agg') from matplotlib.patches import Polygon import pylab as P import glob as glob try: import cPickle as pickle except: print("cPickle not found, using pickle") import pickle import tables from scipy.stats import linregress from numpy.polynomial.polynomial import polyfit from pymses.sources.ramses import output from pymses.sources.hop.file_formats import * from pymses.analysis import Camera, raytracing, slicing, splatting from pymses.filters import CellsToPoints from pymses.analysis import ScalarOperator, FractionOperator, MaxLevelOperator from scipy.stats import gaussian_kde import matplotlib.patches as mpatches from matplotlib.collections import PatchCollection import module_extract as me # extension for out files out_ext = '.jpeg' P.rcParams['image.cmap']='plasma' P.rcParams['savefig.dpi']=400 params= {'text.latex.preamble' : [r'\usepackage{amsmath}']} P.rcParams.update(params) def get_time(path, num): """ Return the time of the output (code units) Parameters ---------- num output number path_out path of the pipeline output Returns ------- time the time of the output (code units) """ try: f = open(path + '/output_' + str(num).zfill(5) + '/info_' + str(num).zfill(5) + '.txt') for line in f: ls = line.split() if len(ls) > 1 and ls[0] == 'time': time = float(ls[2]) break # ro = pymses.RamsesOutput(path, num, order='>') # time = ro.info['time'] # time in codeunits f.close() return time except IOError as e: print(e) return np.nan def compute_image_data(path, num, radius=0.5, path_out=None, order='<', save_data=True, map_size=512, tag='', axes_los=['x', 'y', 'z'], images=['coldens', 'rho', 'speed', 'Q', 'T', 'level', 'cpu'], pos_star=np.array([1., 1., 1.]), put_title=True, force=False, fft=False): """ Make several useful image of an output of a simulation, auxillary function Parameters ---------- 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 map_size size of the map """ ro = pymses.RamsesOutput(path, num, order=order) amr = ro.amr_source(["rho", "vel", "P"]) center = [0.5, 0.5, 0.5] lbox = ro.info["boxlen"] # boxlen in codeunits G = 1.0 # Gravitational constant im_extent = [ (-radius + center[0]) * lbox, (radius + center[0]) * lbox, (-radius + center[1]) * lbox, (radius + center[1]) * lbox, ] time = ro.info["time"] # time in codeunits title = "t=" + str(time)[0:5] + " (code)" # Determining output directory if path_out is not None: directory = path_out else: directory = path # Checking for existing files name_save = directory + '/maps_disk' + '_' + tag + '_' + format(num,'05') + '.save' if (len(glob.glob(name_save))==1 and not force): return # Prepare saving data maps_disk = {'time' : time, 'im_extent' : im_extent, 'center' : center, 'radius' : radius, 'lbox' : lbox, 'images': images, 'axes_los': axes_los} # Prepare raytracing rho_op = ScalarOperator(lambda dset: dset["rho"], ro.info["unit_density"]) rt = None if fft: rt = splatting.SplatterProcessor(amr, ro.info, rho_op) else: rt = raytracing.RayTracer(amr, ro.info, rho_op) # Prepare axes ax_nb = {"x": 0, "y": 1, "z": 2} # Associated number of each axes axes_h = {"x": "y", "y": "x", "z": "x"} # Associated horizontal axe axes_v = {"x": "z", "y": "z", "z": "y"} # Associated vertical axe for ax_los in axes_los: ax_h = axes_h[ax_los] ax_v = axes_v[ax_los] 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, ) # Column density if "coldens" in images: datamap = rt.process(cam, surf_qty=True) dmap_col = datamap.map.T * lbox maps_disk["coldens_" + ax_los] = dmap_col # Rho slice if "rho" in images: datamap_rho = slicing.SliceMap(amr, cam, rho_op, z=0.0) dmap_rho = (datamap_rho.map).T maps_disk["rho_" + ax_los] = dmap_rho if "speed" in images: 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) 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) maps_disk["v" + ax_h + "_" + ax_los] = (dmap_vh.map).T maps_disk["v" + ax_v + "_" + ax_los] = (dmap_vv.map).T if "T" in images: P_op = ScalarOperator(lambda dset: dset["P"], ro.info["unit_pressure"]) dmap_P = (slicing.SliceMap(amr, cam, P_op, z=0.0)).map.T dmap_T = dmap_P / dmap_rho maps_disk["T_" + ax_los] = dmap_T # Toomre parameter if "Q" in images and 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) if fft: rt_cs = splatting.SplatterProcessor(amr, ro.info, cs_op, surf_qty=False) else: rt_cs = raytracing.RayTracer(amr, ro.info, cs_op) dmap_omega = rt_omega.process(cam) dmap_cs = rt_cs.process(cam) map_Q = (lbox * dmap_cs.map.T) * dmap_omega.map.T / (np.pi * G * dmap_col) maps_disk["Q_" + ax_los] = map_Q # Levels if "levels" in images: amr.set_read_levelmax(20) level_op = MaxLevelOperator() rt_level = raytracing.RayTracer(amr, ro.info, level_op) datamap = rt_level.process(cam, surf_qty=True) map_level = datamap.map.T maps_disk["levels_" + ax_los] = map_level # Cpus if "cpu" in images: 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 maps_disk["cpu_" + ax_los] = map_cpu if save_data: f = open(name_save, "w") pickle.dump(maps_disk, f) f.close() return maps_disk def plot_maps(path, num, force=False, vel_red=40, tag='', images=None, axes_los=None, maps_disk=None, interactive=False, put_title=True, set_lim=True) : """ 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 path of the pipeline output force if set, erase any existing pipeline output files tag string to add to the output name vel_red Take point each vel_red for velocities map_size size of the map """ path_out = path # Load maps file print("load maps file") name_maps = path + "/maps_disk" + "_" + tag + "_" + format(num, "05") + ".save" if maps_disk is None: if len(glob.glob(name_maps)) == 0: raise IOError("no pickle file for disk maps. Run make_image_disk() first") f = open(name_maps, "r") maps_disk = pickle.load(f) f.close() print("maps file loaded") time = maps_disk["time"] im_extent = maps_disk["im_extent"] center = maps_disk["center"] radius = maps_disk["radius"] lbox = maps_disk["lbox"] # Plot parameters title = "t=" + str(time)[0:5] + " (code)" ntick = 6 title_ax = {"x": r"$x$ (code)", "y": r"$y$ (code)", "z": r"$z$ (code)"} # Determine output directory if path_out is not None: directory = path_out else: directory = path # Determine what to plot if images == None: images = maps_disk["images"] if axes_los == None: axes_los = maps_disk["axes_los"] # Prepare axes axes_h = {"x": "y", "y": "x", "z": "x"} # Associated horizontal axe axes_v = {"x": "z", "y": "z", "z": "y"} # Associated vertical axe P.close() for ax_los in axes_los: ax_h = axes_h[ax_los] ax_v = axes_v[ax_los] for image in images: if ( (image == "Q" and not ax_los == "z") or image == "levels" or image == "speed" ): continue if (image == 'jeans'): maps_disk['jeans_' + ax_los] = np.sqrt(np.pi * maps_disk['T_' + ax_los] / maps_disk['rho_' + ax_los]) if (image == 'jeans_ratio'): maps_disk['jeans_ratio_' + ax_los] = maps_disk['jeans_' + ax_los] * 2**(maps_disk['levels_' + ax_los]) map_disk = maps_disk[image + '_' + ax_los] if image == 'Q' : im = P.imshow(map_disk, extent=im_extent, origin='lower', cmap='RdBu', norm=mpl.colors.LogNorm(), vmin=0.01, vmax=100.) elif image == 'jeans_ratio' : im = P.imshow(map_disk, extent=im_extent, origin='lower', cmap='RdBu', norm=mpl.colors.LogNorm(), vmin=0.1, vmax=1000.) else: im = P.imshow( map_disk, extent=im_extent, origin="lower", norm=mpl.colors.LogNorm(), ) 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) if image == 'coldens': cbar.set_label(r'$\Sigma$ (code)') if set_lim: im.set_clim(0.01, 100) if 'levels' in images: map_level = maps_disk['levels_' + ax_los] # Computing linewidths levels_ar = np.arange(np.min(map_level), np.max(map_level) + 1) lw = np.ones(levels_ar.size) * 2 lvl_th = 8 # Level threeshold for reducing linewidths lw[levels_ar >= lvl_th] = lw[levels_ar >= lvl_th]**(lvl_th - levels_ar[levels_ar >= lvl_th]) lw[levels_ar < lvl_th] = 1. cont = P.contour(map_level, extent=im_extent, origin='lower', colors='k', linewidths=lw, levels=levels_ar) cont.levels = cont.levels + 1 P.clabel(cont, cont.levels[cont.levels < 11], inline=1, fontsize=8., fmt='%1d') elif image == 'rho': cbar.set_label(r'$\rho$ (code)') if 'speed' in images: dmap_vh = maps_disk['v' + ax_h + '_' + ax_los] dmap_vv = maps_disk['v' + ax_v + '_' + ax_los] map_vh_red = dmap_vh[::vel_red,::vel_red] # take only a subset of velocities map_vv_red = dmap_vv[::vel_red,::vel_red] nh = map_vh_red.shape[0] nv = map_vv_red.shape[1] vec_h = (np.arange(nh)*2./nh*radius - radius + center[0] + radius/nh) * lbox vec_v = (np.arange(nv)*2./nv*radius - radius + center[1] + radius/nv) * lbox 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') elif image == 'T': cbar.set_label(r'$T (code)$') elif image == 'Q': cbar.set_label(r'$Q$') elif image == 'jeans': cbar.set_label(r'Jeans\'s lenght') else: cbar.set_label(image) name = ( directory + "/" + image + "_" + ax_los + "_" + tag + "_" + format(num, "05") + out_ext ) if interactive: P.figure() else: P.tight_layout(pad=1) P.savefig(name) P.close() return maps_disk 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]), binning="log", ): """ 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 the 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 if binning == "log": lrad = np.log10(rad_ext) rad = np.logspace(lrad - 2.0, lrad, num=nb_bin) elif binning == "lin": rad = np.linspace(0.0, rad_ext, num=nb_bin) else: raise ValueError( "Incorrect binning specification, binnning should be 'lin' or 'log'" ) # 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", "g", "phi"]) cell_source = CellsToPoints(amr) cells = cell_source.flatten() dx = cells.get_sizes() * lbox 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 # Gravitational field g = cells["g"] g_rad = (pos[:, 0] * g[:, 0] + pos[:, 1] * g[:, 1]) / norm_pos g_az = (pos[:, 0] * g[:, 1] - pos[:, 1] * g[:, 0]) / norm_pos # Select cells that are actually in the disk, ie within the scale height G = 1. 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"] > 0.01 * np.mean(cells["rho"]) # condition on density mask_vel = abs(v_rad / v_az) < 1.0 # condition on speed mask = mask_pos # & mask_dens & mask_vel 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) # height_disk = height[mask] g_rad_disk = g_rad[mask] g_az_disk = g_az[mask] total_mass_disk = np.sum(rho_disk * dvol_disk) total_mass = np.sum(cells["rho"] * dx ** 3) # 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) alpha_rey_rad_bis = np.zeros(nb_bin - 1) alpha_grav_rad = np.zeros(nb_bin - 1) Q_kepl_rad = np.zeros(nb_bin - 1) height_rad = np.zeros(nb_bin - 1) vol_rad = np.zeros(nb_bin - 1) # Volume of a bin surf_rad = np.zeros(nb_bin - 1) # Surface of a bin mass_rad = np.zeros(nb_bin - 1) # Mass of a bin 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])) vol_rad[i] = np.sum(dvol_disk[mask_bin]) mass_rad[i] = np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin]) press_rad[i] = np.sum(press_disk[mask_bin] * dvol_disk[mask_bin]) / vol_rad[i] rho_rad[i] = np.sum(rho_disk[mask_bin] * dvol_disk[mask_bin]) / vol_rad[i] temp_rad[i] = np.sum(press_disk[mask_bin] * dvol_disk[mask_bin]) / mass_rad[i] # Surface of a bin : S = dr * 2 * pi * r with # dr = rad[i + 1] - rad[i] and r = (rad[i + 1] + rad[i]) / 2. surf_rad[i] = (rad[i + 1] - rad[i]) * (rad[i + 1] + rad[i]) * np.pi coldens_rad[i] = mass_rad[i] / surf_rad[i] v_az_rad[i] = ( np.sum(v_az_disk[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin]) / mass_rad[i] ) v_rad_rad[i] = ( np.sum(v_rad_disk[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin]) / mass_rad[i] ) try: height_rad[i] = ( np.max(pos_disk[:, 2][mask_bin]) - np.min(pos_disk[:, 2][mask_bin]) ) / 2.0 except ValueError: height_rad[i] = 0 alpha_rey_rad[i] = (2.0 / 3) * ( ( 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]) ) alpha_grav_rad[i] = (2.0 / 3) * ( np.sum( g_az_disk[mask_bin] * g_rad_disk[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin] ) / (4 * np.pi * G) / np.sum(dvol_disk[mask_bin] * press_disk[mask_bin]) * coldens_rad[i] ) v_kepl_rad[i] = ( np.sum(v_kepl[mask_bin] * rho_disk[mask_bin] * dvol_disk[mask_bin]) / mass_rad[i] ) # Derived quantities cs_rad = np.sqrt(temp_rad) Q_kepl_rad = cs_rad * v_az_rad / (np.pi * G * coldens_rad * rad[0 : nb_bin - 1]) # Means mask_mean = (0.1 < rad[0 : nb_bin - 1]) & (rad[0 : nb_bin - 1] < 0.2) mass_mean = np.sum(mass_rad[mask_mean]) alpha_rey_mean = np.sum(alpha_rey_rad[mask_mean] * mass_rad[mask_mean]) / mass_mean alpha_grav_mean = ( np.sum(alpha_grav_rad[mask_mean] * mass_rad[mask_mean]) / mass_mean ) Q_mean = np.sum(Q_kepl_rad[mask_mean] * mass_rad[mask_mean]) / mass_mean Q_min = np.nanmin(Q_kepl_rad) # store the results prop_disk = { "time": time, "mass_disk": total_mass_disk, "mass_box": total_mass, "rad": rad[:-1], "center": pos_star, "alpha_rey": alpha_rey_rad, "alpha_rey_mean": alpha_rey_mean, "alpha_grav": alpha_grav_rad, "alpha_grav_mean": alpha_grav_mean, "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, "Q_mean": Q_mean, "Q_min": Q_min, "height": height_rad, } f = open(name_save, "w") pickle.dump(prop_disk, f) f.close() def plot_disk_prop( path, num, plots=["rho", "T", "V", "coldens", "Q", "alpha", "H"], force=False, tag="", interactive=False, put_title=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 disk_prop() first") 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["mass_disk"] rad = prop_disk["rad"] for plot in plots: title = "t=" + str(time)[0:5] + " (code)" P.grid() P.xlabel(r'$r$') if plot == 'rho': P.xscale('log') P.yscale('log') P.plot(rad, prop_disk['rho'], color='k', linewidth=2) P.ylabel(r'$\rho \, (code)$') elif plot == 'T': P.xscale('log') P.yscale('log') P.plot(rad,prop_disk['temp'],color='k',linewidth=2) P.ylabel(r'$T \, (K)$') elif plot == 'V': P.xscale('log') P.yscale('symlog',linthreshy=0.01) P.plot(rad, prop_disk['v_rad'], color='k', linewidth=2) P.plot(rad, prop_disk['v_kepl'], color='b', linewidth=2) P.plot(rad, abs(prop_disk['v_az']), color='r', linewidth=2) P.plot(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 \ (code)$') elif plot == 'coldens': P.xscale('log') P.yscale('log') P.plot(rad,prop_disk['coldens'],color='k',linewidth=2) P.ylabel(r'$\Sigma\ (code)$') elif plot == 'alpha': alpha_rey_mean, alpha_grav_mean = prop_disk['alpha_rey_mean'], prop_disk['alpha_grav_mean'] P.xlim([1e-2, 0.25]) P.yscale("log") P.ylim([1e-7, 1.0]) P.plot( rad, abs(prop_disk["alpha_rey"]), "b", linewidth=2, label=r"$\alpha_{Reynolds}$", ) P.plot(rad, abs(alpha_rey_mean * np.ones(len(rad))), "b:", linewidth=1) P.plot( rad, abs(prop_disk["alpha_grav"]), "r", linewidth=2, label=r"$\alpha_{grav}$", ) P.plot(rad, abs(alpha_grav_mean * np.ones(len(rad))), "r:", linewidth=1) P.plot( rad, abs(prop_disk["alpha_rey"]) + abs(prop_disk["alpha_grav"]), "g--", linewidth=2, label=r"$\alpha_{tot}$", ) P.legend() P.ylabel(r'$\alpha$') title = title + r', $\bar{\alpha}_{Reynolds} = %.1e, \bar{\alpha}_{grav} = %.1e$' % (alpha_rey_mean, alpha_grav_mean) elif plot == 'Q': P.ylim([0, 5.]) P.xlim([0, 0.25]) P.yticks(np.arange(0., 5, 1.0)) P.plot(rad,abs(prop_disk['Q_kepl']),color='b',linewidth=2) # P.plot(rad, abs(prop_disk['Q_mean']) * np.ones(len(rad)), 'b:', linewidth=1) P.ylabel(r'$Q$') title = title elif plot == 'H': P.plot(rad,abs(prop_disk['height']/ rad),color='b',linewidth=2) P.ylabel(r'H ratio') title = title + ', mass of box = {} (code)'.format( prop_disk['mass_box']) if put_title: P.title(title) if interactive: P.figure() else: P.tight_layout(pad=1) P.savefig(path + '/' + plot + '_disk_r_'+ tag + '_' + str(num).zfill(5) + out_ext) P.close() def sigma_pdf(fluct_maps, mask_flat, put_title=False, title='', nb_bin_hist=50, tag=''): # Sigma-PDF dcoldens = np.log10(fluct_maps['coldens']).flatten() nb_cells = np.sum(mask_flat) P.grid() P.yscale('log') P.ylim([0.5 / nb_cells, 1.]) P.xlabel(r'$\log(\Sigma / \bar{\Sigma})$') P.ylabel(r'$\mathcal{P}_\Sigma$') if put_title: P.title(title) values, edges, _ = P.hist(dcoldens[mask_flat], nb_bin_hist, range=(-1, 3), weights = np.ones(nb_cells) / nb_cells) centers = 0.5 * (edges[1:] + edges[:-1]) # Variance var = np.var(dcoldens[mask_flat]) # Compute the slope of the right part of the histogramm mask_fit = (centers > 0.) & (centers < 1.25) & (values > 0) if (np.sum(mask_fit > 0)): (a, b, rho, _, stderr) = linregress(centers[mask_fit], np.log10(values[mask_fit])) P.plot(centers, 10**(a*centers + b), '--', linewidth=2) print("pdf a=%e, b=%e, rho=%e, var=%e"% (a, b, rho, var)) try : beta = int(tag.split('_')[1][4:]) except ValueError : beta = 0 fit = {'beta': beta,'slope': a, 'origin': b, 'correlation': rho, 'stderr':stderr, 'var':var} return fit def plot_dcsdrho(fluct_maps, mask_flat, put_title=False, title='', nb_bin_hist=50, tag=''): drho = fluct_maps['rho'].flatten() dcs = fluct_maps['cs'].flatten() nb_points = np.sum(mask_flat) size_hist = int(np.sqrt(nb_points)) P.hist2d(np.log10(drho[mask_flat]), np.log10(dcs[mask_flat]), (size_hist, size_hist), norm=mpl.colors.LogNorm()) P.xlabel(r'$\log(\rho / \bar{\rho})$') P.ylabel(r'$\log(c_s / \bar{c_s})$') P.ylim([-0.6, 0.6]) P.xlim([-1, 1.]) (a, b, rho, _, stderr) = linregress(np.log10(drho[mask_flat]), np.log10(dcs[mask_flat])) P.plot(np.log10(drho[mask_flat]), a*np.log10(drho[mask_flat]) + b, '--', linewidth=2) print("cs/rho a=%e, b=%e, rho=%e"% (a, b, rho)) try : beta = int(tag.split('_')[1][4:]) except ValueError : beta = 0 fit = {'beta': beta,'slope': a, 'origin': b, 'correlation': rho, 'stderr':stderr} def disk_pdf(path, num, maps_disk, pos_star=[1., 1.], force=False, interactive=False, nb_bin_hist=50, tag='', rad_min=0.075, rad_max=0.3, put_title=True, do_speed=True): # Load property file name_prop = path + "/prop_disk_" + str(num).zfill(5) + ".save" # Check if the properties file exists if len(glob.glob(name_prop)) == 0: raise IOError("no pickle file for disk properties. Run disk_prop() first") f = open(name_prop, "r") prop_disk = pickle.load(f) f.close() # Check if the job has already been done if not force: try: slope = prop_disk["fit"]["slope"] print("PDF already computed, slope = {}, exiting ...".format(slope)) return except KeyError: pass # Load maps file print("load maps file") name_maps = path + "/maps_disk" + "_" + tag + "_" + format(num, "05") + ".save" # Check if the maps file exists if maps_disk is None: if len(glob.glob(name_maps)) == 0: raise IOError("no pickle file for disk maps. Run make_image_disk() first") f = open(name_maps, "r") maps_disk = pickle.load(f) f.close() print("maps file loaded") # Properties time = prop_disk['time'] im_extent = maps_disk['im_extent'] title = None if (put_title): title = tag.split('_')[1] + ' t='+ str(time)[0:5] +' (code)' # Load coldens coldens_map = maps_disk["coldens_z"] # radius of the corner of the box plus a margin rad_box = ( np.sqrt((im_extent[1] - pos_star[0]) ** 2 + (im_extent[3] - pos_star[1]) ** 2) + 0.1 ) # radial bins rad_bins = prop_disk["rad"] rad_bins = 0.5 * (rad_bins[0:-1] + rad_bins[1:]) rad_bins = np.concatenate(([0.0], rad_bins, [rad_box])) x = np.linspace(im_extent[0], im_extent[1], coldens_map.shape[0]) y = np.linspace(im_extent[2], im_extent[3], coldens_map.shape[1]) xx, yy = np.meshgrid(x, y) rr = np.sqrt((xx - pos_star[0])**2 + (yy - pos_star[1])**2) # Find appropriate bin for each coordinate set bins = np.zeros(rr.shape, dtype=int) for r in rad_bins[1:]: bins = bins + (rr >= r).astype(int) bins_flat = bins.flatten() rr_flat = rr.flatten() # Mask selecting the zone of interest mask_map = (rr > rad_min) & (rr < rad_max) # Additionnal maps rho_map = maps_disk["rho_z"] cs_map = np.sqrt(maps_disk["T_z"]) vx_map = maps_disk['vx_z'] vy_map = maps_disk['vy_z'] v_map = np.sqrt(vx_map**2 + vy_map**2) xx_star = xx - pos_star[0] yy_star = yy - pos_star[1] vrad_map = (vx_map * xx_star + vy_map * yy_star) / rr vaz_map = (vy_map * xx_star - vx_map * yy_star) / rr maps = { "coldens": coldens_map, "rho": rho_map, "cs": cs_map, "v": v_map, "vaz": vaz_map, } avg_maps = {} fluct_maps = {} for cur_map in maps: map_arr = maps[cur_map] # mean of all the cells in the bin mean_bin = np.zeros(len(rad_bins) - 1) for j in range(len(rad_bins) - 1): mean_bin[j] = np.mean(map_arr[bins == j]) # save for later # prop_disk[cur_map + '_rmean'] = mean_bin # Add value for borders mean_bin = np.concatenate(([mean_bin[0]], mean_bin)) # Compute the map azimuthally averaged # use linear interpolation to improve accuracy avg_flat = (rad_bins[bins_flat + 1] - rr_flat) * mean_bin[bins_flat] avg_flat = avg_flat + (rr_flat - rad_bins[bins_flat]) * mean_bin[bins_flat + 1] avg_flat = avg_flat / (rad_bins[bins_flat + 1] - rad_bins[bins_flat]) avg_maps[cur_map] = np.reshape(avg_flat, rr.shape) # Select zone of interest # avg_maps[cur_map][np.logical_not(mask_map)] = np.nan # Compute fluctuation fluct_maps[cur_map] = map_arr / avg_maps[cur_map] if not interactive: plot_fluctuations_map(path, num, maps, fluct_maps, avg_maps, im_extent, mask_map, put_title, title, nb_bin_hist, tag, prop_disk) # Save on disk f = open(name_prop,'w') pickle.dump(prop_disk, f) f.close() return (xx, yy, fluct_maps) def plot_fluctuations_map(path, num, maps, fluct_maps, avg_maps, im_extent, mask_map, put_title, title, nb_bin_hist, tag, prop_disk, interactive=False): mask_flat = mask_map.flatten() # Sigma-PDF fit = sigma_pdf(fluct_maps, mask_flat, put_title, title, nb_bin_hist, tag) if interactive: P.figure() else: P.tight_layout(pad=1) P.savefig(path + "/dcol_hist_" + tag + "_" + str(num).zfill(5) + out_ext) P.close() # rho-PDF drho = fluct_maps['rho'].flatten() nb_cells = np.sum(mask_flat) P.grid() P.yscale('log') P.ylim([0.5 / nb_cells, 1.]) P.xlabel(r'$\log(\rho / \bar{\rho})$') P.ylabel(r'$\mathcal{P}_\rho$') if put_title: P.title(title) values, edges, _ = P.hist(drho[mask_flat], nb_bin_hist, # range=(-1, 6), weights = np.ones(nb_cells) / nb_cells) centers = 0.5 * (edges[1:] + edges[:-1]) if interactive: P.figure() else: P.tight_layout(pad=1) P.savefig(path + '/drho_hist_' + tag + '_' + str(num).zfill(5) + out_ext) P.close() # Derived quantities dcoldens = np.log10(fluct_maps['coldens']).flatten() dcs = fluct_maps['cs'].flatten() dv = fluct_maps['v'].flatten() dmach = abs(maps['v'] - avg_maps['v']) / maps['cs'] dmach[dmach > 10.] = 10. fluct_maps['mach'] = dmach dmach_mean = np.mean(dmach[mask_map & (dmach < 5.)].flatten()) dcoldens_max = np.max(dcoldens[mask_flat]) # Save general properties prop_disk['dcoldens_max'] = dcoldens_max prop_disk['dmach_mean'] = dmach_mean print("dmach_mean = {}".format(dmach_mean)) # Fluctuations plots # dcs = f(drho) prop_disk['fit_cs'] = plot_dcsdrho(fluct_maps, mask_flat, put_title, title, nb_bin_hist, tag) P.colorbar() if interactive: P.figure() else: P.tight_layout(pad=1) P.savefig(path + "/dcs_" + tag + "_" + str(num).zfill(5) + out_ext) P.close() # dv = f(drho) P.hist2d( np.log10(drho[mask_flat]), dv[mask_flat], (1000, 1000), norm=mpl.colors.LogNorm(), ) P.xlabel(r"$\log(\rho / \bar{\rho})$") P.ylabel(r"$\log(v_\varphi / \bar{v_\varphi})$") P.colorbar() if interactive: P.figure() else: P.tight_layout(pad=1) P.savefig(path + "/dv_" + tag + "_" + str(num).zfill(5) + out_ext) P.close() # Fluctuations maps for cur_map in fluct_maps: fluct_map = fluct_maps[cur_map] if cur_map == 'coldens': im = P.imshow(fluct_map, norm=mpl.colors.LogNorm(), extent=im_extent, origin='lower', cmap='viridis') label = r'$log(\Sigma/\bar{\Sigma})$' elif cur_map == 'cs': im = P.imshow(np.log10(fluct_map), extent=im_extent, origin='lower', cmap='RdBu_r', vmin=-0.6, vmax=0.6) label = r'$log(c_s/\bar{c_s})$' elif cur_map == 'rho': im = P.imshow(np.log10(fluct_map), extent=im_extent, origin='lower', cmap='RdBu_r', vmin=-2., vmax=2.) label = r'$log(\rho/\bar{\rho})$' elif cur_map == 'vaz': im = P.imshow(fluct_map, extent=im_extent, origin='lower', cmap='RdBu_r') label = r'$v_\varphi / \bar{v_\varphi}$' elif cur_map == 'mach': im = P.imshow(fluct_map, extent=im_extent, origin='lower', cmap='RdBu_r') label = r'$|v - \bar{v}| / c_s$' if put_title: P.title(title) P.xlabel(r"$x$") P.ylabel(r"$y$") cbar = P.colorbar(im) cbar.set_label(label) name = path + "/d" + cur_map + "map_" + tag + "_" + format(num, "05") name_im = name + out_ext if interactive: P.figure() else: P.tight_layout(pad=1) P.savefig(name_im) P.close() def compare(path, runs, nums, path_out=None, force=False, interactive=False, Q_in_name=True, pdf=False, gamma_ad=5./3., tag=''): """ Compare time averaged properties of a disk in several simulations nums id or array of ids of the ramses output runs list of runs to consider path path to the properties file force if set, redo plots even if already done interactive interactive mode, to use in a %pylab ipython shell """ if type(nums) == int: nums = [nums] if(type(nums)==dict): nums_name = '_' else: nums_name = "_" + "_".join(str(num).zfill(5) for num in [nums[0], nums[-1]]) # Initialize arrays var_names = ['alpha_rey_mean', 'alpha_grav_mean', 'Q', 'time'] std_names = ['alpha_rey_mean', 'alpha_grav_mean', 'Q'] plot_names = ['alphaQ', 'alphaQall'] if Q_in_name: var_names.append('Q0') plot_names.append('alphaQ0') if pdf: var_names = var_names + ['beta', 'kappa', 'var', 'dmach', 'gamma'] std_names = std_names + ['kappa', 'var', 'dmach', 'gamma'] plot_names = plot_names + ['alphabeta', 'kappabeta', 'gammabeta', 'varbeta', 'dmachvarall'] var_data = dict() std_data = dict() all_data = dict() for name in var_names: var_data[name] = np.zeros(len(runs)) all_data[name] = [] for name in std_names: std_data[name] = np.zeros(len(runs)) for i, run in enumerate(runs): path_run = path + "/" + run nb_outputs = 0 if(type(nums)==dict): nums_run = nums[run] else: nums_run = nums run_data = dict() for name in var_names: run_data[name] = [] for num in nums_run: try: # Load property file name_save = path_run + "/prop_disk_" + str(num).zfill(5) + ".save" # Check if the properties file exists if len(glob.glob(name_save)) == 0: raise IOError( "no pickle file for disk properties {}. Run disk_prop() first".format( name_save ) ) f = open(name_save, "r") prop_disk = pickle.load(f) f.close() for name in var_names: if pdf: fit = prop_disk['fit'] fit_cs = prop_disk['fit_cs'] if name == 'alpha_rey_mean': data = abs(prop_disk['alpha_rey_mean']) elif name == 'alpha_grav_mean': data = abs(prop_disk['alpha_grav_mean']) elif name == 'Q': data = prop_disk['Q_min'] elif name == 'beta': data = fit['beta'] elif name == 'kappa': data = - fit['slope'] elif name == 'var': data = fit['var'] elif name == 'dmach': data = prop_disk['dmach_mean'] elif name == 'gamma': data = 2 * fit_cs['slope'] + 1 elif name == 'Q0': data = float(run.split('_')[2][1:]) print(data) else: data = prop_disk[name] run_data[name].append(data) all_data[name].append(data) nb_outputs = nb_outputs + 1 print(run, num, nb_outputs) except (IOError, KeyError) as e: print(run, num, e, nb_outputs) if nb_outputs > 0: for name in var_names: var_data[name][i] = np.mean(run_data[name]) for name in std_names: std_data[name][i] = np.std(run_data[name]) else: for name in var_names: var_data[name][i] = np.nan for name in std_names: std_data[name][i] = np.nan for name in var_names: all_data[name] = np.array(all_data[name]) for pname in plot_names: P.grid() print("Ploting "+ pname) if pname == 'alphaQ' or pname == 'alphaQ0': # alpha = f(Q) if pname == 'alphaQ': P.xlabel(r'$Q_{min}$') Q = var_data['Q'] xerr = std_data['Q'] else: P.xlabel(r'$Q_{0}$') Q = var_data['Q0'] xerr = None P.yscale('log') P.ylim([1e-6,1.]) P.xlim([np.nanmin(Q)- 0.2, np.nanmax(Q) + 0.2]) P.errorbar(Q, abs(var_data['alpha_rey_mean']), xerr=xerr, yerr=std_data['alpha_rey_mean'], fmt='*--', label=r"$\bar{\alpha}_{Reynolds}$") # P.errorbar(Q, abs(var_data['alpha_grav_mean']), xerr=xerr, yerr=std_data['alpha_grav_mean'], # fmt='o--', label=r"$\bar{\alpha}_{grav}$") P.legend() P.ylabel(r'$\bar{\alpha}$') elif pname == 'kappabeta': P.errorbar(var_data['beta'], var_data['kappa'], yerr=std_data['kappa'], fmt='*') P.ylabel(r'$\kappa = - dlog\mathcal{P}_\Sigma / d\log\left(\Sigma/\overline{\Sigma}\right)$') P.xlabel(r'$\beta$') P.xlim([None, np.max(var_data['beta']) + 0.3]) (a, b, rho, _, stderr) = linregress(var_data['beta'], var_data['kappa']) #P.plot(var_data['beta'], a*var_data['beta'] + b, '--', linewidth=1.5) c = polyfit(var_data['beta'], var_data['kappa'], 1, w = [1.0 / ty for ty in std_data['kappa']], full=False) P.plot(var_data['beta'], c[1]*var_data['beta'] + c[0], '--', linewidth=1.5) print("slope #1 : a=%e, b=%e, rho^2=%e"% (a,b,rho**2)) rho_p = np.sqrt(np.sum((var_data['kappa'] - c[1]*var_data['beta'] - c[0])**2)) print("slope #2 : a=%e, b=%e, err^2=%e"% (c[1],c[0],rho_p**2)) elif pname =='gammabeta': # gamma = f(beta) P.errorbar(var_data['beta'], var_data['gamma'], yerr=std_data['gamma'], fmt='*') P.xlim([None, np.max(var_data['beta']) + 0.3]) P.ylabel(r'$\Gamma_{eff}$') P.xlabel(r'$\beta$') (a, b, rho, _, stderr) = linregress(var_data['beta'], var_data['gamma']) # P.plot(var_data['beta'], a*var_data['beta'] + b, '--', linewidth=1.5) print("gamma : a=%e, b=%e, rho^2=%e"% (a,b,rho**2)) elif pname == 'varbeta': P.errorbar(var_data['beta'], var_data['var'], yerr=std_data['var'], fmt='*') P.ylabel(r'$Var(\log(\Sigma / \bar{\Sigma}))$') P.xlabel(r'$\beta$') (a, b, rho, _, stderr) = linregress(var_data['beta'], var_data['var']) P.plot(var_data['beta'], a*var_data['beta'] + b, '--', linewidth=1.5) print("a=%e, b=%e, rho^2=%e"% (a,b,rho**2)) elif pname == 'dmachvarall': mask = all_data['dmach'] < 1 all_var = all_data['var'][mask] all_beta = all_data['beta'][mask] all_dmach = all_data['dmach'][mask] beta_max = np.max(var_data['beta']) beta_min = np.min(var_data['beta']) P.grid() cmap = mpl.cm.get_cmap('RdYlBu', beta_max - beta_min + 1) P.scatter(all_dmach, np.exp(all_var), c=all_beta, vmin=beta_min-0.5, vmax=beta_max+0.5, cmap=cmap) cbar = P.colorbar() cbar.set_ticks(var_data['beta']) cbar.set_label(r'$\beta$') P.xlabel(r'$<(v - \bar{v}) / c_s>$') P.ylabel(r'$\exp(Var(\log(\Sigma / \bar{\Sigma}))$') if pname == 'alphaQall' : Q = all_data['Q'] alpha_rey = abs(all_data['alpha_rey_mean']) time = all_data['time'] P.yscale('log') P.ylim([1e-7,1.]) P.scatter(Q, alpha_rey, c=time) P.xlabel(r'$Q_{min}$') P.ylabel(r'$\bar{\alpha}_{Reynolds}$') cbar = P.colorbar() cbar.set_label(r'time (code)') if pname == 'alphabeta' : alpha_th = (4./9.) / (gamma_ad * (gamma_ad - 1.) * var_data['beta']) mask = var_data['beta'] >= 10 P.errorbar(var_data['beta'][mask], abs(var_data['alpha_rey_mean'][mask]), yerr=std_data['alpha_rey_mean'][mask], fmt='*--', label=r"$\bar{\alpha}_{Reynolds}$") # P.plot(var_data['beta'], abs(var_data['alpha_grav_mean']), '*--', label=r"$\bar{\alpha}_{grav}$") P.plot(var_data['beta'][mask], abs(alpha_th[mask]), ':', label=r"$\bar{\alpha}_{th}$") P.legend() P.ylabel(r'$\bar{\alpha}$') P.xlabel(r'$\beta$') (a, b, rho, _, stderr) = linregress(np.log(var_data['beta']), np.log(var_data['alpha_rey_mean'])) print("log(alpha)= a*log(beta) + b, a=%e, b=%e, rho^2=%e"% (a,b,rho**2)) if interactive: P.figure() else: P.tight_layout(pad=1) P.savefig(path_out + '/'+ pname + '_' + tag + nums_name + out_ext) P.close() # alpha = f(beta) # # P.yscale('log') # P.ylim([None,0.3]) # P.grid() # # var of the pdf = f(var_data['beta']) # P.grid() # # theoritical alpha (Gammie 2001) # alpha_th = (4./9.) / (gamma_ad * (gamma_ad - 1.) * var_data['beta']) # P.plot(var_data['beta'], abs(alpha_rey_mean), 'o-.', label=r"$\bar{\alpha}_{Reynolds}$") # P.plot(var_data['beta'], abs(alpha_grav), '*--', label=r"$\bar{\alpha}_{grav}$") # P.plot(var_data['beta'], abs(alpha_th), ':', label=r"$\bar{\alpha}_{th}$") # P.legend() # P.ylabel(r'$\bar{\alpha}$') # P.xlabel(r'$\beta_{cool}$') # (a, b, rho, _, stderr) = linregress(np.log(var_data['beta']), np.log(alpha_rey_mean)) # print("log(alpha)= a*log(beta) + b, a=%e, b=%e, rho^2=%e"% (a,b,rho**2)) def evolution(path, nums, force=False, interactive=False, pdf=False, tag=''): """ Plot properties over time path path to the properties file nums list of id of the ramses output force if set, redo plots even if already done interactive interactive mode, to use in a %pylab ipython shell """ # Initialize arrays var_names = ['time', 'alpha_rey_mean', 'alpha_grav_mean', 'Q_min', 'Q_mean', 'mass_disk', 'mass_box'] plot_names = ['alpha', 'Q', 'mass'] plot_labels = [r'$\alpha$', r'$Q$', 'mass (code)'] if pdf: var_names = var_names + ['kappa', 'var', 'dmach_mean', 'dcoldens_max', 'gamma'] plot_names = plot_names + ['kappa', 'var', 'dmach_mean', 'dcoldens_max', 'gamma'] plot_labels = plot_labels + [r'$\kappa$', r'$Var(\log(\Sigma/\overline{\Sigma})$', 'dmach_mean', r'$\max\left(\log(\Sigma/\overline{\Sigma})\right)$', r'$\Gamma_{eff}$'] plot_labels = dict(zip(plot_names, plot_labels)) var_data = dict() for name in var_names: var_data[name] = np.zeros(len(nums)) for i, num in enumerate(nums): # Load property file name_prop = path + "/prop_disk_" + str(num).zfill(5) + ".save" # Check if the properties file exists if len(glob.glob(name_prop)) == 0: raise ("no pickle file for disk properties. Run disk_prop() first") f = open(name_prop, "r") prop_disk = pickle.load(f) f.close() for name in var_names: try: if name == 'kappa': var_data[name][i] = - prop_disk['fit']['slope'] elif name == 'var': var_data[name][i] = prop_disk['fit']['var'] elif name == 'gamma': var_data[name][i] = 2 * prop_disk['fit_cs']['slope'] + 1 else: var_data[name][i] = prop_disk[name] except KeyError as e: print('[{}, {}] Key error {}'.format(path, num, e)) var_data[name][i] = np.nan time = var_data['time'] for pname in plot_names: if pname == 'alpha': P.yscale('log') P.plot(time, abs(var_data['alpha_rey_mean']), 'o-.', label=r"$\bar{\alpha}_{Reynolds}$") P.plot(time, abs(var_data['alpha_grav_mean']), '*--', label=r"$\bar{\alpha}_{grav}$") P.legend() elif pname == 'Q': P.plot(time, var_data['Q_min'], 'o-.', label=r"$Q_{min}$") P.plot(time, var_data['Q_mean'], '*--', label=r"$\overline{Q}$") P.legend() elif pname == 'mass': P.plot(time, var_data['mass_disk'], 'o-.', label=r"$M_{disk}$") P.plot(time, var_data['mass_box'], '*--', label=r"$M_{box}$") P.legend() else: P.plot(time, var_data[pname], '*--') if pname =='dcoldens_max': P.plot(time, 1.5 * np.ones(len(time)), 'r-', lw=4) P.ylabel(plot_labels[pname]) P.xlabel(r'time (code)') P.grid() if interactive: P.figure() else: P.tight_layout(pad=1) P.savefig(path + '/' + pname + 'time_' + tag + out_ext) P.close() def clump_study(path, num, path_out, tag, force=False, rho_thres=1e3, lvl_thres=10): name = tag + '_' + str(num).zfill(5) hop_save = name + '_hop' + '_prop_struct.save' if not (len(glob.glob(path_out + '/' + hop_save))==1) or force: me.make_clump_hop(path, num, name + '_hop', rho_thres, lvl_thres, [0.5, 0.5, 0.5], 1, path_out=path_out + '/', path_hop='./', force=True, gcomp=False) hop_save = me.clump_properties(name + '_hop', path, num, path_out=path_out + '/', gcomp=False) f = open(path_out + '/' + hop_save) cp = pickle.load(f) f.close() ener = cp['grav_vir'] + cp['therm_vir'] * 3./2 rad = np.sqrt((cp['pos_n_max'][:, 0] - 1.)**2 + (cp['pos_n_max'][:, 1] - 1.)**2) mask = (ener < 0) & (rad > 0.07) pos = cp['pos_n_max'][mask] color = np.log10(cp['n_max'][mask]) print("Load maps files") f = open(path_out + '/maps_disk_' + tag + '_' + str(num).zfill(5) + '.save') m = pickle.load(f) f.close() print("Mpas files loaded") P.scatter(pos[:, 0],pos[:, 1], c = color) im = P.imshow(m['coldens_z'], extent=m['im_extent'], origin='lower', norm=mpl.colors.LogNorm(), vmin=0.01, vmax=100) P.xlabel(r'$x$') P.ylabel(r'$y$') cbar = P.colorbar(im) cbar.set_label(r"$\Sigma$") P.tight_layout(pad=1) P.title("time = {} (code)".format(m['time'])) P.savefig(path_out + '/clumps_' + tag + '_' + str(num).zfill(5) + out_ext) P.close() # fig, ax = P.subplots(1, 1) # offsets = list(zip(pos[:, 0], 2 - pos[:, 1])) # size = cp['size_iner2'][mask] # ax.add_collection(EllipseCollection(widths=size, heights=size, angles=0, units='xy', # facecolors=P.cm.hsv(color), # offsets=offsets, transOffset=ax.transData)) # P.imshow(np.log10(m['coldens_z']), extent=m['im_extent'], cmap='plasma') # P.savefig(path_out + '/clumps_size_' + tag + '_' + str(num).zfill(5) + out_ext) class interactive_plots: def onbuttonrelease(self, event): """Deal with click events""" button = ['left','middle','right'] toolbar = P.get_current_fig_manager().toolbar if toolbar.mode=='zoom rect' and event.inaxes == self.ax_col: print("zooming ") xlim = self.ax_col.get_xlim() ylim = self.ax_col.get_ylim() self.reset_mask() elif self.add_mask and event.inaxes == self.ax_col: self.plot_side() P.draw() def onbuttonpress(self, event): """Deal with click events""" button = ['left','middle','right'] toolbar = P.get_current_fig_manager().toolbar if toolbar.mode!='': print("You clicked on something, but toolbar is in mode {:s}.".format(toolbar.mode)) print(self.add_mask) if self.add_mask and toolbar.mode=='' and event.inaxes == self.ax_col: ix, iy = event.xdata, event.ydata print("Add patch {}, {}".format(ix, iy)) xlim = self.ax_col.get_xlim() ylim = self.ax_col.get_ylim() radius = 0.05 * min(abs(xlim[1] - xlim[0]), abs(ylim[1] - ylim[0])) circle = mpatches.Circle([ix, iy], radius, color='black', alpha=0.1, ec="none") self.circles.append(circle) self.ax_col.add_artist(circle) self.ax_col.draw_artist(circle) self.patch_mask = self.patch_mask | ((self.xx - ix)**2 + (self.yy -iy)**2 < radius**2) #self.plot_side() def onkeypress(self, event): """whenever a key is pressed""" if not event.inaxes: return if event.key == 't': self.add_mask = not self.add_mask print("Add mode is {}".format(self.add_mask)) elif event.key == 'r': self.reset_mask() def plot_side(self): if (self.add_mask): mask = (self.patch_mask & self.mask).flatten() else: mask = self.mask.flatten() self.ax_gamma.clear() P.sca(self.ax_gamma) plot_dcsdrho(self.fluct_maps, mask, tag=self.tag) self.ax_pdf.clear() P.sca(self.ax_pdf) sigma_pdf(self.fluct_maps, mask, tag=self.tag, nb_bin_hist=self.args.pdf_nb_bin) def reset_mask(self): xlim = self.ax_col.get_xlim() ylim = self.ax_col.get_ylim() self.mask = (self.xx >= xlim[0]) & (self.xx <= xlim[1]) & (self.yy >= ylim[0]) & (self.yy <= ylim[1]) self.patch_mask = np.full(self.mask.shape, False) for circle in self.circles: circle.remove() self.circles = [] self.plot_side() P.draw() def __init__(self, args, path, num, maps_disk=None, tag='', set_lim=True) : """ Interactive plotting Parameters ---------- num output number path path of the pipeline output """ self.args = args self.add_mask = False self.circles = [] self.tag = tag path_out = path # Load maps file print("load maps file") name_maps = path + '/maps_disk' + '_' + tag + '_' + format(num,'05') + '.save' if maps_disk is None: if (len(glob.glob(name_maps)) == 0): raise IOError('no pickle file for disk maps {}. Run make_image_disk() first'.format(name_maps)) f = open(name_maps,'r') maps_disk = pickle.load(f) f.close() print("maps file loaded") im_extent = maps_disk['im_extent'] fig = P.figure(); self.ax_col = P.subplot(1, 2, 1) coldens = maps_disk['coldens_z'] im = self.ax_col.imshow(coldens, extent=im_extent, origin='lower', norm=mpl.colors.LogNorm()) if set_lim: im.set_clim(0.01, 100) self.ax_col.set_xlabel(r'$x$') self.ax_col.set_ylabel(r'$y$') self.xx, self.yy, self.fluct_maps = disk_pdf(path, num, maps_disk, tag=self.tag, force=True, put_title=False, interactive=True) coord_flat = zip(self.xx.flatten(), self.yy.flatten()) self.ax_gamma = P.subplot(2, 2, 2) self.ax_pdf = P.subplot(2,2,4) xlim = self.ax_col.get_xlim() ylim = self.ax_col.get_ylim() self.mask = (self.xx >= xlim[0]) & (self.xx <= xlim[1]) & (self.yy >= ylim[0]) & (self.yy <= ylim[1]) self.patch_mask = np.full(self.mask.shape, False) self.plot_side() fig.canvas.mpl_connect('button_release_event', self.onbuttonrelease) fig.canvas.mpl_connect('button_press_event', self.onbuttonpress) fig.canvas.mpl_connect('key_press_event', self.onkeypress) P.tight_layout() P.show()