From 7d5aa9d9112ca380855e6e79016d2fde27faabf4 Mon Sep 17 00:00:00 2001 From: Noe Brucy Date: Sun, 1 Dec 2019 23:46:17 +0100 Subject: [PATCH] Add run selector + unit handling + multiprocessing + examples --- .gitignore | 1 + auto_pp.py | 38 + auto_pp.sh | 8 + clump.py | 7 - disk_postprocess.py | 1680 ------------------------------------------- mypool.py | 63 ++ pipeline.py | 443 ++++++------ pipeline_disk.py | 397 ---------- plotter.py | 920 +++++++++++++++--------- postprocessor.py | 1396 +++++++++++++++++++++++------------ pp_params.py | 15 +- run_selector.py | 166 +++++ units.py | 71 ++ 13 files changed, 2088 insertions(+), 3117 deletions(-) create mode 100644 .gitignore create mode 100644 auto_pp.py create mode 100755 auto_pp.sh delete mode 100644 clump.py delete mode 100644 disk_postprocess.py create mode 100644 mypool.py delete mode 100644 pipeline_disk.py create mode 100644 run_selector.py create mode 100644 units.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d5e992d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/.ipynb_checkpoints/ diff --git a/auto_pp.py b/auto_pp.py new file mode 100644 index 0000000..0b7f9a5 --- /dev/null +++ b/auto_pp.py @@ -0,0 +1,38 @@ +from plotter import * + +# Turb +in_dir = "/home/nbrucy/simus/ismfeed/turb" +out_dir = "/home/nbrucy/visus/ismfeed/turb" + +nml_key = ["turb_params/turb_rms", "turb_params/comp_frac"] + +cond1_5 = [("cloud_params/dens0", "=", 1.5), ("turb_params/turb_rms", "!=", 72000)] +cond6 = ("cloud_params/dens0", "=", 6) + +nproc = 15 + + +pl = Plotter( + in_dir, + namelist_cond=cond1_5, + in_nums="all", + sort_run_by=nml_key, + path_out=out_dir, + tag="sfr_turb", +) +pl6 = Plotter( + in_dir, + namelist_cond=cond6, + in_nums="all", + sort_run_by=nml_key, + path_out=out_dir, + tag="sfr_turb_n6", +) + +pl.pp_params.process.num_process = nproc +pl6.pp_params.process.num_process = nproc + + +for plotter in [pl, pl6]: + plotter.process(["sigma"]) + plotter.process(["coldens_l", "rho_v"], ["x, y, z"]) diff --git a/auto_pp.sh b/auto_pp.sh new file mode 100755 index 0000000..aa90116 --- /dev/null +++ b/auto_pp.sh @@ -0,0 +1,8 @@ +#! /bin/bash + +while true; do + date >> ~/pp.log + python auto_pp.py >> ~/pp.log 2>&1 + printf "\n\n" >> ~/pp.log + sleep 2000 +done diff --git a/clump.py b/clump.py deleted file mode 100644 index 9fd5237..0000000 --- a/clump.py +++ /dev/null @@ -1,7 +0,0 @@ -import module_extract as me -import pickle -import matplotlib.pyplot as plt -import numpy as np -from matplotlib.collections import EllipseCollection - - diff --git a/disk_postprocess.py b/disk_postprocess.py deleted file mode 100644 index ae05c7d..0000000 --- a/disk_postprocess.py +++ /dev/null @@ -1,1680 +0,0 @@ -# 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() diff --git a/mypool.py b/mypool.py new file mode 100644 index 0000000..331784a --- /dev/null +++ b/mypool.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +import multiprocessing + +# We must import this explicitly, it is not imported by the top-level +# multiprocessing module. +import multiprocessing.pool +import time + +from random import randint + + +class NoDaemonProcess(multiprocessing.Process): + # make 'daemon' attribute always return False + def _get_daemon(self): + return False + + def _set_daemon(self, value): + pass + + daemon = property(_get_daemon, _set_daemon) + + +# We sub-class multiprocessing.pool.Pool instead of multiprocessing.Pool +# because the latter is only a wrapper function, not a proper class. +class MyPool(multiprocessing.pool.Pool): + Process = NoDaemonProcess + + +def sleepawhile(t): + print("Sleeping %i seconds..." % t) + time.sleep(t) + return t + + +def work(num_procs): + print("Creating %i (daemon) workers and jobs in child." % num_procs) + pool = multiprocessing.Pool(num_procs) + + result = pool.map(sleepawhile, [randint(1, 5) for x in range(num_procs)]) + + # The following is not really needed, since the (daemon) workers of the + # child's pool are killed when the child is terminated, but it's good + # practice to cleanup after ourselves anyway. + pool.close() + pool.join() + return result + + +def test(): + print("Creating 5 (non-daemon) workers and jobs in main process.") + pool = MyPool(5) + + result = pool.map(work, [randint(1, 5) for x in range(5)]) + + pool.close() + pool.join() + print(result) + + +if __name__ == "__main__": + test() diff --git a/pipeline.py b/pipeline.py index 67dc8c4..d641b08 100644 --- a/pipeline.py +++ b/pipeline.py @@ -16,163 +16,174 @@ fake_pp = PostProcessor() parser = argparse.ArgumentParser() -input_args = parser.add_argument_group('input', "Input selection") -input_args.add_argument("runs", help='name of runs', nargs='*', - default=[]) +input_args = parser.add_argument_group("input", "Input selection") +input_args.add_argument("runs", help="name of runs", nargs="*", default=[]) -input_args.add_argument("-ip", "--input_path", - help="specify input directory", - default="/home/nbrucy/simus/") -input_args.add_argument("-p", "--project", - help="specify project name (directory within the input directory)", - default="disk") +input_args.add_argument( + "-ip", "--input_path", help="specify input directory", default="/home/nbrucy/simus/" +) +input_args.add_argument( + "-p", + "--project", + help="specify project name (directory within the input directory)", + default="disk", +) -input_args.add_argument("-c", "--config", - help="Path of a default config file", - default=None) +input_args.add_argument( + "-c", "--config", help="Path of a default config file", default=None +) -input_args.add_argument("-wo", "--which_inputs", - choices=['all', 'id', 'time'], - help="Select inputs by time range, id range or all of them", - default='all') -input_args.add_argument("-b", "--begin", - help="id of first input", - type=int, - default=1) -input_args.add_argument("-e", "--end", - help="id of last input", - type=int, - default=100) -input_args.add_argument("-s", "--step", - help="step between two input", - type=int, - default=1) -input_args.add_argument("-tb", "--time_begin", - help="time of first input", - type=float, - default=0.) -input_args.add_argument("-te", "--time_end", - help="time of last input", - type=float, - default=6.) +input_args.add_argument( + "-wo", + "--which_inputs", + choices=["all", "id", "time"], + help="Select inputs by time range, id range or all of them", + default="all", +) +input_args.add_argument("-b", "--begin", help="id of first input", type=int, default=1) +input_args.add_argument("-e", "--end", help="id of last input", type=int, default=100) +input_args.add_argument( + "-s", "--step", help="step between two input", type=int, default=1 +) +input_args.add_argument( + "-tb", "--time_begin", help="time of first input", type=float, default=0.0 +) +input_args.add_argument( + "-te", "--time_end", help="time of last input", type=float, default=6.0 +) -input_args.add_argument("-w", "--watch", - help="wait and watch for missing inputs", - action='store_true') -input_args.add_argument("--skip", - help="skip failed loadings", - action='store_true') -input_args.add_argument("-wt", "--waiting_time", - help="time between to successive try when watching new inputs (in second)", - type=int, - default=120) -input_args.add_argument("-af", "--allowed_failures", - help="number of allowed failures when waiting", - default=30) +input_args.add_argument( + "-w", "--watch", help="wait and watch for missing inputs", action="store_true" +) +input_args.add_argument("--skip", help="skip failed loadings", action="store_true") +input_args.add_argument( + "-wt", + "--waiting_time", + help="time between to successive try when watching new inputs (in second)", + type=int, + default=120, +) +input_args.add_argument( + "-af", + "--allowed_failures", + help="number of allowed failures when waiting", + default=30, +) -output_args = parser.add_argument_group('output', "Output configuration") +output_args = parser.add_argument_group("output", "Output configuration") -output_args.add_argument("-op", "--output_path", - help="specify output directory", - default='/home/nbrucy/visus/') +output_args.add_argument( + "-op", + "--output_path", + help="specify output directory", + default="/home/nbrucy/visus/", +) -output_args.add_argument("--tag", - help="Add a special tag on output filemanes", - default='') +output_args.add_argument( + "--tag", help="Add a special tag on output filemanes", default="" +) -output_args.add_argument("--interactive", - help="Interactive mode : do not save data", - action='store_true') +output_args.add_argument( + "--interactive", help="Interactive mode : do not save data", action="store_true" +) -output_args.add_argument("-owr", "--overwrite", - help="Overwrite outputs", - action='store_true') +output_args.add_argument( + "-owr", "--overwrite", help="Overwrite outputs", action="store_true" +) -output_args.add_argument("-owrd", "--overwrite_dependencies", - help="Overwrite outputs for dependencies", - action='store_true', - default=False) +output_args.add_argument( + "-owrd", + "--overwrite_dependencies", + help="Overwrite outputs for dependencies", + action="store_true", + default=False, +) -pp_args = parser.add_argument_group('postproc', "Post Processing configuration") +pp_args = parser.add_argument_group("postproc", "Post Processing configuration") -pp_args.add_argument("--process", - help="Individual rules to apply", - choices=fake_pp.rules.keys(), - default=[], - nargs='*') +pp_args.add_argument( + "--process", + help="Individual rules to apply", + choices=fake_pp.rules.keys(), + default=[], + nargs="*", +) -pp_args.add_argument("-pargs", "--process_args", - help="Args to give to process rules", - default=['x', 'y', 'z'], - nargs='*') +pp_args.add_argument( + "-pargs", + "--process_args", + help="Args to give to process rules", + default=["x", "y", "z"], + nargs="*", +) -pp_args.add_argument("--compare", - help="Time and inter run comparaison", - default=[], - nargs='*') +pp_args.add_argument( + "--compare", help="Time and inter run comparaison", default=[], nargs="*" +) -pp_args.add_argument("-cargs", "--compare_args", - help="Args to give to process rules", - default=[None], - nargs='*') +pp_args.add_argument( + "-cargs", + "--compare_args", + help="Args to give to process rules", + default=[None], + nargs="*", +) -pp_args.add_argument("--plot", - help="Plot rules", - default=[], - nargs='*') +pp_args.add_argument("--plot", help="Plot rules", default=[], nargs="*") -pp_args.add_argument("-plargs", "--plot_args", - help="Args to give to plot rules", - default=[None], - nargs='*') +pp_args.add_argument( + "-plargs", + "--plot_args", + help="Args to give to plot rules", + default=[None], + nargs="*", +) -pp_args.add_argument("-d", "--disk", - help="Specify this for disk simulations", - action='store_true') -pp_args.add_argument("--fft", - help="use quick and dirty fft rendering", - action='store_true') -pp_args.add_argument("--zoom", - help="zoom", - type=float, - default=2.) -pp_args.add_argument("-ms", "--map_size", - help="size of the maps created in he map mode (in pixel)", - type=int, - default=1024) +pp_args.add_argument( + "-d", "--disk", help="Specify this for disk simulations", action="store_true" +) +pp_args.add_argument( + "--fft", help="use quick and dirty fft rendering", action="store_true" +) +pp_args.add_argument("--zoom", help="zoom", type=float, default=2.0) +pp_args.add_argument( + "-ms", + "--map_size", + help="size of the maps created in he map mode (in pixel)", + type=int, + default=1024, +) -pp_args.add_argument("--nb_bin", - help="Number of bins for azimuthal averages", - type=int, - default=50) -pp_args.add_argument("--pdf_nb_bin", - help="Number of bins for PDF", - type=int, - default=50) -pp_args.add_argument("--binning", - help="Kind of binning (logarithmic or linear)", - choices=['log', 'lin'], - default='log') +pp_args.add_argument( + "--nb_bin", help="Number of bins for azimuthal averages", type=int, default=50 +) +pp_args.add_argument( + "--pdf_nb_bin", help="Number of bins for PDF", type=int, default=50 +) +pp_args.add_argument( + "--binning", + help="Kind of binning (logarithmic or linear)", + choices=["log", "lin"], + default="log", +) -plot_args = parser.add_argument_group('plot', "Plot configuration") +plot_args = parser.add_argument_group("plot", "Plot configuration") -plot_args.add_argument("--colormap", - help="Colormap used", - choices=P.colormaps(), - default='plasma') -plot_args.add_argument("--format", - help="Format of the plot images", - choices=['png', 'jpeg', 'pdf', 'svg', 'ps'], - default='jpeg') -plot_args.add_argument("--dpi", - help="Resolution of the plot images", - type=int, - default=400) -plot_args.add_argument("--beamer", - help="Beamer mode", - action='store_true') +plot_args.add_argument( + "--colormap", help="Colormap used", choices=P.colormaps(), default="plasma" +) +plot_args.add_argument( + "--format", + help="Format of the plot images", + choices=["png", "jpeg", "pdf", "svg", "ps"], + default="jpeg", +) +plot_args.add_argument( + "--dpi", help="Resolution of the plot images", type=int, default=400 +) +plot_args.add_argument("--beamer", help="Beamer mode", action="store_true") args = parser.parse_args() @@ -182,16 +193,16 @@ storage_in = args.input_path storage_out = args.output_path if args.config is None: - pp_params = default_params() + pp_params = default_params() else: - pp_params = load_params(args.config) + pp_params = load_params(args.config) pp_params.out.zoom = args.zoom -pp_params.out.tag = args.tag -pp_params.out.map_size = args.map_size +pp_params.out.tag = args.tag +pp_params.out.map_size = args.map_size pp_params.out.interactive = args.interactive -pp_params.pymses.fft = args.fft +pp_params.pymses.fft = args.fft pp_params.disk.on = args.disk @@ -202,94 +213,114 @@ pp_params.pdf.nb_bin = args.pdf_nb_bin # extension for out files P.style.use("seaborn-deep") -if args.format == 'pdf': - P.style.use("~/.config/matplotlib/pdf.mplstyle") +if args.format == "pdf": + P.style.use("~/.config/matplotlib/pdf.mplstyle") if args.beamer: - P.rcParams['font.family'] = 'sans-serif' - P.rcParams['figure.figsize'] = (7, 4.5) + P.rcParams["font.family"] = "sans-serif" + P.rcParams["figure.figsize"] = (7, 4.5) # Plot properties -P.rcParams['image.cmap'] = args.colormap -P.rcParams['savefig.dpi'] = args.dpi -P.rcParams['lines.linewidth'] = 2 -P.rcParams['lines.markersize'] = 10 +P.rcParams["image.cmap"] = args.colormap +P.rcParams["savefig.dpi"] = args.dpi +P.rcParams["lines.linewidth"] = 2 +P.rcParams["lines.markersize"] = 10 P.rcParams["errorbar.capsize"] = 4 # List of id that were successfully computed -nums_success = dict() +nums_success = {} # Go through all runs for run in runs: - path_suffix = project + '/' + run - path_in = storage_in + path_suffix - path_out = storage_out + path_suffix + path_suffix = project + "/" + run + path_in = storage_in + path_suffix + path_out = storage_out + path_suffix - if args.tag == '': - tag_run = run - else: - tag_run = run + '_' + args.tag + if args.tag == "": + tag_run = run + else: + tag_run = run + "_" + args.tag - if not os.path.exists(path_out): - os.makedirs(path_out) - try: - copy(path_in + '/disk.nml', path_out) - copy(path_in + '/output_00001/compilation.txt', path_out) - except: - pass + if not os.path.exists(path_out): + os.makedirs(path_out) + try: + copy(path_in + "/disk.nml", path_out) + copy(path_in + "/output_00001/compilation.txt", path_out) + except: + pass - nums_success[run] = [] + nums_success[run] = [] - if args.which_inputs in ['all', 'time'] : - names = glob.glob(path_in + "/output_[0-9][0-9][0-9][0-9][0-9]") - nums_all = [int(n.split('/')[-1].split('_')[1]) for n in names] - nums_all = np.sort(nums_all) - if args.which_inputs == 'all': - nums = nums_all - else: - time = [get_time(path_in, n) for n in nums_all] - nums = [n for i,n in enumerate(nums_all) if time[i] >= args.time_begin - and time[i] < args.time_end] - else: - nums = range(args.begin, args.end + 1, args.step) + if args.which_inputs in ["all", "time"]: + names = glob.glob(path_in + "/output_[0-9][0-9][0-9][0-9][0-9]") + nums_all = [int(n.split("/")[-1].split("_")[1]) for n in names] + nums_all = np.sort(nums_all) + if args.which_inputs == "all": + nums = nums_all + else: + time = [get_time(path_in, n) for n in nums_all] + nums = [ + n + for i, n in enumerate(nums_all) + if time[i] >= args.time_begin and time[i] < args.time_end + ] + else: + nums = range(args.begin, args.end + 1, args.step) - for num in nums: - failures = 0 - success = False + for num in nums: + failures = 0 + success = False - while not success: - try: - if len(args.process) > 0 : - pp = PostProcessor(run, num, pp_params=pp_params) - pp.process(args.process, args.process_args, - overwrite=args.overwrite, overwrite_dep=args.overwrite_dependencies) + while not success: + try: + if len(args.process) > 0: + pp = PostProcessor(run, num, pp_params=pp_params) + pp.process( + args.process, + args.process_args, + overwrite=args.overwrite, + overwrite_dep=args.overwrite_dependencies, + ) - # If we are here, success ! - success = True - nums_success[run].append(num) - except (ValueError, IOError, KeyError) as e: - print(e) - if(args.watch and failures < args.allowed_failures): - failures = failures + 1 - print("ERROR: Unable to proceed for run {} output {}.\ - Trying again in {} s ({} tries remaining)".format(run, num, - args.waiting_time, args.allowed_failures - failures)) - time.sleep(args.waiting_time) - elif args.skip: - break - else: - raise + # If we are here, success ! + success = True + nums_success[run].append(num) + except (ValueError, IOError, KeyError) as e: + print(e) + if args.watch and failures < args.allowed_failures: + failures = failures + 1 + print( + "ERROR: Unable to proceed for run {} output {}.\ + Trying again in {} s ({} tries remaining)".format( + run, + num, + args.waiting_time, + args.allowed_failures - failures, + ) + ) + time.sleep(args.waiting_time) + elif args.skip: + break + else: + raise path_in = storage_in + project path_out = storage_out + project if len(args.plot) > 0: - pl = Plotter(path_in, runs, nums_success, path_out=path_out, pp_params=pp_params) - pl.process(args.plot, args.plot_args, - overwrite=args.overwrite, overwrite_dep=args.overwrite_dependencies) + pl = Plotter(path_in, runs, nums_success, path_out=path_out, pp_params=pp_params) + pl.process( + args.plot, + args.plot_args, + overwrite=args.overwrite, + overwrite_dep=args.overwrite_dependencies, + ) if len(args.compare) > 0: - cc = Comparator(path_in, runs, nums_success, path_out=path_out, pp_params=pp_params) - cc.process(args.compare, args.compare_args, - overwrite=args.overwrite, overwrite_dep=args.overwrite_dependencies) - + cc = Comparator(path_in, runs, nums_success, path_out=path_out, pp_params=pp_params) + cc.process( + args.compare, + args.compare_args, + overwrite=args.overwrite, + overwrite_dep=args.overwrite_dependencies, + ) diff --git a/pipeline_disk.py b/pipeline_disk.py deleted file mode 100644 index d60c99f..0000000 --- a/pipeline_disk.py +++ /dev/null @@ -1,397 +0,0 @@ -# coding: utf-8 - -import os -import glob -from shutil import copy -import argparse -import time -import numpy as np -from functools import reduce - -from pp_params import * -from plotter import * -from postprocessor import * - - - -storage_in = '/home/nbrucy/simus/' -storage_out = '/home/nbrucy/visus/' - -parser = argparse.ArgumentParser() - -parser.add_argument("runs", help='name of runs', nargs='*', - default=[]) - - -parser.add_argument("-wo", "--which_outputs", - choices=['all', 'id', 'time'], - help="Select outputs by time range, id range or all of them", - default='all') -parser.add_argument("-b", "--begin", - help="id of first output", - type=int, - default=1) -parser.add_argument("-e", "--end", - help="id of last output", - type=int, - default=100) -parser.add_argument("-s", "--step", - help="step between two output", - type=int, - default=1) -parser.add_argument("-tb", "--time_begin", - help="time of first output", - type=float, - default=0.) -parser.add_argument("-te", "--time_end", - help="time of last output", - type=float, - default=6.) - - -parser.add_argument("-p", "--project", - help="specify project name", - default="disk") -parser.add_argument("-fr", "--force_redo", - help="redo plots even if the files already exist", - action='store_true') -parser.add_argument("-w", "--watch", - help="wait and watch for missing outputs", - action='store_true') -parser.add_argument("--skip", - help="skip failed loadings", - action='store_true') -parser.add_argument("-wt", "--waiting_time", - help="time between to successive try when watching new outputs (in second)", - type=int, - default=120) -parser.add_argument("-af", "--allowed_failures", - help="number of allowed failures when waiting", - default=30) -parser.add_argument("-i", "--interactive", - help="Interactive mode", - action='store_true') - - -parser.add_argument("--tag", - help="Add a special tag on output filemanes", - default='') - - - -parser.add_argument("-d", "--disk", - help="compute specific disk radial analysis", - action='store_true') -parser.add_argument("-pd", "--plot_disk", - help="plot specific disk radial analysis", - action='store_true') -parser.add_argument("-m", "--maps", - help="compute generic maps", - action='store_true') -parser.add_argument("-pm", "--plot_maps", - help="plot generic maps", - action='store_true') -parser.add_argument("-c", "--compare", - help="compare different runs", - action='store_true') -parser.add_argument("-ev", "--evolution", - help="plot evolution of quantities over time", - action='store_true') -parser.add_argument("--pdf", - help="plot pdf of fluctuations of column density", - action='store_true') -parser.add_argument("--pdf2", - help="plot patrick's PDF", - action='store_true') -parser.add_argument("--cpdf", - help="compare pdf", - action='store_true') -parser.add_argument("--clump", - help="Do clump study", - action='store_true') -parser.add_argument("--print_outputs", - help="print names of outputs", - action='store_true') - -parser.add_argument("--fft", - help="use quick and dirty fft rendering", - action='store_true') -parser.add_argument("--images", nargs='*', - default=['coldens', 'rho', 'speed', 'Q', 'T'], - choices=['coldens', 'rho', 'speed', 'Q', 'T', 'levels', 'cpu', 'jeans', 'jeans_ratio']) -parser.add_argument("--axes", nargs='*', - default=['x', 'y', 'z'], - choices=['x', 'y', 'z']) -parser.add_argument("--zoom", - help="zoom", - type=float, - default=2.) -parser.add_argument("-ms", "--mapsize", - help="size of the maps created in he map mode (in pixel)", - type=int, - default=1024) - - -parser.add_argument("--nb_bin", - help="Number of bins for azimuthal averages", - type=int, - default=50) -parser.add_argument("--pdf_nb_bin", - help="Number of bins for PDF", - type=int, - default=50) -parser.add_argument("--binning", - help="Kind of binning (logarithmic or linear)", - choices=['log', 'lin'], - default='log') -parser.add_argument("--rad_ext", - help="Value of the highest bin", - type=float, - default=1.) -parser.add_argument("-x", - help="x position of the central point", - type=float, - default=1.) -parser.add_argument("-y", - help="y position of the central point", - type=float, - default=1.) -parser.add_argument("-z", - help="z position of the central point", - type=float, - default=1.) - -parser.add_argument( - "-ta", "--time_avg", help="Plot time averaged comparaison", action="store_true" -) - - -parser.add_argument("--colormap", - help="Colormap used", - choices=dp.P.colormaps(), - default='plasma') -parser.add_argument("--format", - help="Format of the output", - choices=['png', 'jpeg', 'pdf', 'svg', 'ps'], - default='jpeg') -parser.add_argument("--dpi", - help="Resolution of the output", - type=int, - default=400) -parser.add_argument("--beamer", - help="Beamer mode", - action='store_true') -parser.add_argument("--Q_in_name", - help="Whether to use the name to have Q", - action='store_true') - - - -args = parser.parse_args() - -project = args.project -runs = args.runs - -first = args.begin -last = args.end -step = args.step - -rad = 0.5 / args.zoom - -# extension for out files -dp.out_ext = '.' + args.format -dp.P.style.use("seaborn-deep") -if args.format == 'pdf': - dp.P.style.use("~/.config/matplotlib/pdf.mplstyle") - - - -if args.beamer: - dp.P.rcParams['font.family'] = 'sans-serif' - dp.P.rcParams['figure.figsize'] = (7, 4.5) - -# Plot properties -dp.P.rcParams['image.cmap'] = args.colormap -dp.P.rcParams['savefig.dpi'] = args.dpi -dp.P.rcParams['lines.linewidth'] = 2 -dp.P.rcParams['lines.markersize'] = 10 -dp.P.rcParams["errorbar.capsize"] = 4 - -me.P = dp.P - -# List of id that were successfully computed -run_succeded = {} - - -# Care for dependencies -images = args.images -if 'jeans_ratio' in images and not 'jeans' in images : - images = ['jeans'] + images -if 'jeans_ratio' in images and not 'levels' in images : - images.append('levels') -if 'jeans' in images and not 'T' in images : - images.append('T') -if ('T' in images or 'jeans' in images) and not 'rho' in images : - images.append('rho') -if 'Q' in images and not 'coldens' in images : - images.append('coldens') - -# Go through all runs -for run in runs: - path_suffix = project + '/' + run - path_in = storage_in + path_suffix - path_out = storage_out + path_suffix - - if args.tag == '': - tag_run = run - else: - tag_run = run + '_' + args.tag - - if not os.path.exists(path_out): - os.makedirs(path_out) - try: - copy(path_in + '/disk.nml', path_out) - copy(path_in + '/output_00001/compilation.txt', path_out) - except: - pass - - run_succeded[run] = [] - - if args.which_outputs in ['all', 'time'] : - names = glob.glob(path_in + "/output_[0-9][0-9][0-9][0-9][0-9]") - nums_all = [int(n.split('/')[-1].split('_')[1]) for n in names] - nums_all = np.sort(nums_all) - if args.which_outputs == 'all': - nums = nums_all - else: - time = [dp.get_time(path_in, n) for n in nums_all] - nums = [n for i,n in enumerate(nums_all) if time[i] >= args.time_begin and time[i] < args.time_end] - else: - nums = range(first, last + 1, step) - - for num in nums: - failures = 0 - success = False - - while not success: - try: - maps_disk = None - if args.print_outputs: - print("[{}, {}]".format(run, num)) - if args.maps: - print("[{}, {}] computing maps".format(run, num)) - maps_disk = dp.compute_image_data(path_in, num, - radius=rad, - path_out=path_out, - tag=tag_run, - map_size=args.mapsize, - force=args.force_redo, - axes_los=args.axes, - images=images, - pos_star=np.array([args.x, args.y, args.z]), - fft=args.fft) - print("[{}, {}] maps computed".format(run, num)) - if args.plot_maps: - print("[{}, {}] plotting maps".format(run, num)) - maps_disk = dp.plot_maps(path_out, num, - maps_disk=maps_disk, - axes_los=args.axes, - images=images, - tag=tag_run, - force=args.force_redo, - interactive=args.interactive, - put_title=(not args.beamer)) - print("[{}, {}] maps plotted".format(run, num)) - if args.disk: - print("[{}, {}] computing disk props".format(run, num)) - dp.disk_prop(path_in, num, path_out=path_out, - nb_bin=args.nb_bin, - binning=args.binning, - rad_ext=args.rad_ext, - force=args.force_redo, - pos_star=np.array([args.x, args.y, args.z])) - print("[{}, {}] disk_props computed".format(run, num)) - if args.plot_disk: - print("[{}, {}] plotting disk props".format(run, num)) - dp.plot_disk_prop(path_out, num, tag=tag_run, - force=args.force_redo, - interactive=args.interactive, - put_title=(not args.beamer)) - print("[{}, {}] disk props plotted".format(run, num)) - if args.pdf: - print("[{}, {}] computing pdf #1".format(run, num)) - dp.disk_pdf(path_out, num, maps_disk, - pos_star=np.array([args.x, args.y, args.z]), - force=args.force_redo, - tag=tag_run, - nb_bin_hist=args.pdf_nb_bin, - interactive=args.interactive, - put_title=(not args.beamer)) - print("[{}, {}] pdf #1 computed".format(run, num)) - if args.pdf2: - print("[{}, {}] computing pdf #2".format(run, num)) - me.get_pdf(path_in, num, path_out=path_out, - force=args.force_redo, - tag=tag_run) - print("[{}, {}] pdf #2 computed".format(run, num)) - if args.clump: - print("[{}, {}] computing clumps".format(run, num)) - dp.clump_study(path_in, num, path_out, tag_run) - print("[{}, {}] clumps computed".format(run, num)) - if args.interactive: - print("[{}, {}] Interactive session".format(run, num)) - maps_disk = dp.interactive_plots(args, path_out, num, tag=tag_run) - # If we are here, success ! - success = True - run_succeded[run].append(num) - except (ValueError, IOError, KeyError) as e: - print(e) - if(args.watch and failures < args.allowed_failures): - failures = failures + 1 - print("Unable to proceed for run {} output {}. Trying again in {} s ({} tries remaining)".format(run, num, args.waiting_time, args.allowed_failures - failures)) - time.sleep(args.waiting_time) - elif args.skip: - break - else: - raise - if args.evolution: - print("[{}] plotting evolution".format(run)) - dp.evolution(path_out, run_succeded[run], - force=args.force_redo, - tag=args.tag, - interactive=args.interactive, - pdf=args.pdf or args.cpdf) - print("[{}] evolution plotted".format(run)) - -if args.compare: - path_suffix = project - path = storage_out + path_suffix - - if (args.time_avg): - dp.compare(path, runs, run_succeded, - path_out=path + '/comp', - force=args.force_redo, - interactive=args.interactive, - Q_in_name=args.Q_in_name, - tag=args.tag, - pdf=args.pdf or args.cpdf) - else: - for i in range(first, last + 1, step): - # Select usable runs - runs_ok = [run for run in runs if i in run_succeded[run]] - - try: - if len(runs_ok) > 1: - dp.compare(path, runs_ok, i, - path_out=path + '/comp', - force=args.force_redo, - interactive=args.interactive, - Q_in_name=args.Q_in_name, - tag=args.tag, - pdf=args.pdf or args.cpdf) - except (KeyError, IOError) as e: - print(e) - if (args.skip): - pass - else: - raise diff --git a/plotter.py b/plotter.py index 58514ff..dd31abd 100644 --- a/plotter.py +++ b/plotter.py @@ -5,9 +5,10 @@ import os import tables 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') + +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 from scipy.stats import linregress @@ -17,184 +18,291 @@ from functools import partial from numpy.polynomial.polynomial import polyfit from scipy.ndimage.filters import gaussian_filter1d -from pp_params import * from postprocessor import * -P.rcParams['image.cmap']='plasma' -P.rcParams['savefig.dpi']=400 +P.rcParams["image.cmap"] = "plasma" +P.rcParams["savefig.dpi"] = 400 -tex_params= {'text.latex.preamble' : [r'\usepackage{amsmath}']} +tex_params = {"text.latex.preamble": [r"\usepackage{amsmath}"]} P.rcParams.update(tex_params) class PlotRule(Rule): - def plot(self, save, arg, **kwargs): self.postproc.save = save return self.process_fn(arg, **kwargs) def is_valid(self, arg): - return arg in self.args_ok and self.is_valid_add(arg) + return self.is_valid_add(arg) -class Plotter(BaseProcessor): + +class Plotter(Aggregator, BaseProcessor): """ This class loads derived quantities and plot them """ + solve_self_dep = False + # Axes information - _ax_nb = {'x' : 0, 'y' : 1, 'z' : 2} # 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 - _ax_title = {'x' : r'$x$ (code)', 'y' : r'$y$ (code)', 'z' : r'$z$ (code)'} + _ax_nb = {"x": 0, "y": 1, "z": 2} # 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 + _ax_title = {"x": r"$x$ [code]", "y": r"$y$ [code]", "z": r"$z$ [code]"} - G = 1. # Gravitational constant + G = 1.0 # Gravitational constant - def __init__(self, path, runs, nums, path_out=None, pp_params=default_params(), tag=None): + label_convert = { + "turb_rms": "$f_{rms}$", + "beta": "$\\beta$", + "dens0": "$n_0$", + "sfr_avg_window": "window", + "comp_frac": "$\\zeta$", + } - if type(pp_params) == str: - self.pp_params = load_params(pp_params) - else : - self.pp_params = pp_params + value_convert = { + "sfr_avg_window": lambda x: "${:g}$ Myr".format(80 * x), + "comp_frac": lambda x: "${:g}$".format(1 - x), + } - if tag is not None: - self.pp_params.out.tag = tag + def __init__( + self, + path, + in_runs=None, + in_nums=None, + path_out=None, + pp_params=None, + selector=None, + tag=None, + **kwargs + ): - # Determining output directory - if (path_out is None): - self.path_out = path - else: - self.path_out = path_out + super(Plotter, self).__init__(path, path_out, pp_params, tag) - # Get postprocesor objets for each run - self.pp_runs = dict() - if not type(nums) == dict: - nums_tmp = nums - nums = dict() - for run in runs: - nums[run] = nums_tmp + # Select runs + if selector is None: + selector = RunSelector(path, in_runs, in_nums, self.pp_params, **kwargs) - for run in runs: - path_run = path + '/' + run - path_out_run = path_out + '/' + run - self.pp_runs[run] = dict() - for num in nums[run]: - self.pp_runs[run][num] = PostProcessor(path_run, num, path_out_run, pp_params) + # Save infos + self.path = path + self.runs = selector.runs + self.nums = selector.nums # Get comparator object - self.comp = Comparator(path, runs, nums, path_out, pp_params) - - # save infos - self.runs = runs - self.nums = nums + self.comp = Comparator( + path, self.runs, self.nums, path_out, self.pp_params, selector=selector + ) + # Get postprocesor objets for each run + self.pp_runs = self.comp.pp_runs + # Define log prefix self.log_id = "[plot {}] ".format(self.pp_params.out.tag) + + # Define rules self.def_rules() - def _process_single(self, name, rule, arg, overwrite=False, just_done=[], **kwargs): - # Solve dependencies - for dep in rule.dependencies: - if dep in self.comp.rules: - if type(rule.dependencies) == dict: - self.comp.process([dep], [rule.dependencies[dep]], - self.overwrite_dep) - else: - self.comp.process([dep], [arg], self.overwrite_dep, self.overwrite_dep) - else: - for run in self.runs: - for i, num in enumerate(self.nums[run]): - if type(rule.dependencies) == dict: - dep_arg = rule.dependencies[dep] - self.pp_runs[run][num].process([dep], [dep_arg], - self.overwrite_dep) - else: - self.pp_runs[run][num].process([dep], [arg], - self.overwrite_dep, - self.overwrite_dep) - - # Process rule - done = self._process_rule(name, rule, arg, overwrite, just_done, **kwargs) - return just_done + [done] - - def _process_rule(self, name, rule, arg, overwrite=False, just_done=[], **kwargs): - if not arg is None: - name_full = name + '_' + str(arg) + def _not_self_dep(self, name, dep, dep_arg, overwrite, **kwargs): + if dep in self.comp.rules: + done = self.comp.process(dep, dep_arg, overwrite, overwrite) + self.just_done.extend(done) else: - name_full = name + super(Plotter, self)._not_self_dep(name, dep, dep_arg, overwrite, **kwargs) + + def _process_rule(self, name, rule, arg, overwrite=False, **kwargs): + if not arg is None: + name_full = name + "_" + str(arg) + else: + name_full = name if rule.is_valid(arg): - if rule.kind == 'classic': + if rule.kind == "classic": for run in self.runs: for i, num in enumerate(self.nums[run]): - plot_filename = self._find_filename(name_full, run, num) - save = tables.open_file(self.pp_runs[run][num].filename) - try: - self._plot_rule(rule, save, arg, plot_filename, overwrite, **kwargs) - finally: - save.close() + plot_filename = self._find_filename(name_full, run, num) + save = tables.open_file(self.pp_runs[run][num].filename) + try: + self._plot_rule( + rule, + save, + arg, + plot_filename, + overwrite, + run=run, + **kwargs + ) + finally: + save.close() else: - if rule.kind == 'series' and len(self.runs) == 1: + if rule.kind == "series" and len(self.runs) == 1: run = self.runs[0] plot_filename = self._find_filename(name_full, run) else: plot_filename = self._find_filename(name_full) - save = tables.open_file(self.comp.filename, 'r') - try : - self._plot_rule(rule, save, arg, plot_filename, overwrite, **kwargs) + save = tables.open_file(self.comp.filename, "r") + try: + self._plot_rule( + rule, + save, + arg, + plot_filename, + overwrite, + open_figure=not self.pp_params.out.interactive, + **kwargs + ) finally: save.close() else: self._log("{} is not valid in this context".format(name_full), "ERROR") - def _plot_rule(self, rule, save, arg, plot_filename, overwrite, **kwargs): - if overwrite or not os.path.exists(plot_filename): - if not self.pp_params.out.interactive: - P.figure() - rule.plot(save, arg, **kwargs) - P.tight_layout(pad=1) - if not self.pp_params.out.interactive: - P.savefig(plot_filename) - P.close() - self._log("{} plotted".format(plot_filename), "SUCCESS") - else: - self._log("{} plotted".format(os.path.basename(plot_filename)), "SUCCESS") - else: - self._log("Plot {} is already done, skipping...".format(plot_filename)) - + def _plot_rule( + self, rule, save, arg, plot_filename, overwrite, open_figure=True, **kwargs + ): + if overwrite or not os.path.exists(plot_filename): + if open_figure: + P.figure() + rule.plot(save, arg, **kwargs) + P.tight_layout(pad=1) + if not self.pp_params.out.interactive: + P.savefig(plot_filename) + P.close() + self._log("{} plotted".format(plot_filename), "SUCCESS") + else: + self._log( + "{} plotted".format(os.path.basename(plot_filename)), "SUCCESS" + ) + else: + self._log("Plot {} is already done, skipping...".format(plot_filename)) def _find_filename(self, name_full, run=None, num=None): - if not self.pp_params.out.tag == '': - tag_name = '_' + self.pp_params.out.tag - else : - tag_name = '' + if not self.pp_params.out.tag == "": + tag_name = "_" + self.pp_params.out.tag + else: + tag_name = "" if not run is None and not num is None: - return (self.path_out + '/' + run + '/' + name_full + - tag_name + '_' + format(num,'05') + - self.pp_params.plot.out_ext) + return ( + self.path_out + + "/" + + run + + "/" + + name_full + + tag_name + + "_" + + format(num, "05") + + self.pp_params.plot.out_ext + ) elif not run is None: - return (self.path_out + '/' + run + '/' - + name_full + tag_name + self.pp_params.plot.out_ext) + return ( + self.path_out + + "/" + + run + + "/" + + name_full + + tag_name + + self.pp_params.plot.out_ext + ) else: - return self.path_out + '/' + name_full + tag_name + self.pp_params.plot.out_ext + return ( + self.path_out + "/" + name_full + tag_name + self.pp_params.plot.out_ext + ) + def _label_run(self, run, node, label, nml_key): + def get_label_nml(nml_key): + prop_name = os.path.basename(nml_key) + if prop_name in self.label_convert: + prop_label = self.label_convert[prop_name] + else: + prop_label = prop_name + prop_value = self.comp.get_nml(nml_key, run) + if prop_name in self.value_convert: + prop_value_str = self.value_convert[prop_name](prop_value) + else: + prop_value_str = "${:.6g}$".format(prop_value) + return r"{} = {}".format(prop_label, prop_value_str) + if nml_key is None and label is None: + label_run = r"{}".format(self.save.root._v_attrs.attrs[node.name].label) + elif not nml_key is None: + if not type(nml_key) == list: + nml_key = [nml_key] + label_run = ", ".join(map(get_label_nml, nml_key)) - def _plot_map(self, name, ax_los, label=None, cmap='plasma', vmin=None, vmax=None, overlays=[]): + if not label is None: + label_run = label + " (" + label_run + ")" + else: + label_run = label + return label_run + + def _ax_label_unit(self, node, label, unit, unit_coeff): + if label is None: + if "label" in node._v_attrs: + label = node._v_attrs.label + elif node._v_name in self.label_convert: + label = self.label_convert[node._v_name] + elif not node._v_title == "": + label = node._v_title + else: + label = node._v_name + + if "unit" in node._v_attrs: + unit_old = node._v_attrs.unit + else: + unit_old = cst.none + + if unit is None: + unit = unit_old + + if not unit_coeff == 1: + base = unit + unit = unit_coeff * unit + label = label + unit_str(unit, base=base) + else: + label = label + unit_str(unit) + + return label, unit_old, unit + + def _plot_map( + self, + name, + ax_los, + run, + label=None, + unit=None, + unit_coeff=1.0, + overlays=[], + title=None, + nml_key=None, + put_time=True, + cmap="plasma", + norm="log", + autoscale=True, + **kwargs + ): ax_h = self._axes_h[ax_los] ax_v = self._axes_v[ax_los] im_extent = self.save.root.maps._v_attrs.im_extent - dmap = self.save.get_node('/maps/{}_{}'.format(name, ax_los)).read() + node = self.save.get_node("/maps/{}_{}".format(name, ax_los)) + dmap = node.read() + + label, unit_old, unit = self._ax_label_unit(node, label, unit, unit_coeff) + + dmap = dmap * unit_old.express(unit) + + if norm == "log": + norm = mpl.colors.LogNorm() + elif norm == "linear": + norm = mpl.colors.NoNorm() + + if autoscale and not norm is None: + norm.autoscale(dmap) + + im = P.imshow( + dmap, extent=im_extent, origin="lower", norm=norm, cmap=cmap, **kwargs + ) - im = P.imshow(dmap, - extent=im_extent, - origin='lower', - cmap=cmap, - norm=mpl.colors.LogNorm()) - im.set_clim(vmin, vmax) P.locator_params(axis=ax_h, nbins=self.pp_params.plot.ntick) P.locator_params(axis=ax_v, nbins=self.pp_params.plot.ntick) @@ -206,75 +314,106 @@ class Plotter(BaseProcessor): if not label is None: cbar.set_label(label) + title = self._label_run(run, node, title, nml_key) + + if put_time: + time = self.save.root._v_attrs.time * self.comp.info["unit_time"] + time_str = "time = {:.6g} Myr".format(time.express(cst.Myr)) + if len(title) > 0: + title = title + " | " + time_str + else: + title = time_str + + P.title(title) + for plot_overlay in overlays: plot_overlay(ax_los) def _overlay_levels(self, ax_los): - map_level = self.save.get_node('/maps/{}_{}'.format('levels', ax_los)).read() + map_level = self.save.get_node("/maps/{}_{}".format("levels", ax_los)).read() # 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. + 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.0 - cont = P.contour(map_level, - extent=self.save.root.maps._v_attrs.im_extent, - origin='lower', - colors='grey', - linewidths=lw, - levels=levels_ar) + cont = P.contour( + map_level, + extent=self.save.root.maps._v_attrs.im_extent, + origin="lower", + colors="grey", + linewidths=lw, + levels=levels_ar, + ) cont.levels = cont.levels + 1 - P.clabel(cont, - cont.levels[cont.levels < 11], - inline=1, fontsize=8., fmt='%1d'); + P.clabel(cont, cont.levels[cont.levels < 11], inline=1, fontsize=8.0, fmt="%1d") - def _overlay_speed(self, ax_los): + def _overlay_speed(self, ax_los, unit=cst.km_s, unit_coeff=1.0): ax_h = self._axes_h[ax_los] ax_v = self._axes_v[ax_los] - dmap_vh = self.save.get_node('/maps/speed_h_{}'.format(ax_los)).read() - dmap_vv = self.save.get_node('/maps/speed_v_{}'.format(ax_los)).read() + dmap_vh_node = self.save.get_node("/maps/speed_h_{}".format(ax_los)) + dmap_vh = dmap_vh_node.read() + dmap_vv = self.save.get_node("/maps/speed_v_{}".format(ax_los)).read() vel_red = self.pp_params.plot.vel_red radius = self.save.root.maps._v_attrs.radius center = self.save.root.maps._v_attrs.center lbox = self.save.root._v_attrs.lbox - map_vh_red = dmap_vh[::vel_red,::vel_red] # take only a subset of velocities - map_vv_red = dmap_vv[::vel_red,::vel_red] + 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)) + vec_h = ( + np.arange(nh) * 2.0 / nh * radius - radius + center[0] + radius / nh + ) * lbox + vec_v = ( + np.arange(nv) * 2.0 / nv * radius - radius + center[1] + radius / nv + ) * lbox + hh, vv = np.meshgrid(vec_h, vec_v) + norm_v = np.sqrt(map_vh_red ** 2 + map_vv_red ** 2) + max_v = np.max(norm_v) + min_v = np.min(norm_v) - Q = P.quiver(hh, vv, map_vh_red, map_vv_red, units='width', color='grey') - P.quiverkey(Q, 0.6, 0.98, max_v, - r'$'+str(max_v)[0:4]+'$ (code)', labelpos='E', coordinates='figure') + Q = P.quiver(hh, vv, map_vh_red, map_vv_red, units="width", color="grey") + label, unit_old, unit = self._ax_label_unit(dmap_vh_node, "", unit, unit_coeff) + + key_v = (max_v + min_v) / 2.0 + P.quiverkey( + Q, + 0.6, + 0.98, + key_v, + r"${:g}$".format(key_v) + label, + labelpos="E", + coordinates="figure", + ) def _plot_radial(self, name, label=None, xlog=False, ylog=False): - radial_bins = self.save.get_node('/radial/radial_bins_' + ax_los).read() - bin_centers = 0.5 * (radial_bins[1:] + radial_bins[:-1]) - mean_bin = self.save.get_node('/radial/{}_{}'.format(name, ax_los)).read() + radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read() + bin_centers = 0.5 * (radial_bins[1:] + radial_bins[:-1]) + mean_bin = self.save.get_node("/radial/{}_{}".format(name, ax_los)).read() P.grid() - P.xlabel(r'$r$') + P.xlabel(r"$r$") if xlog: - P.xscale('log') + P.xscale("log") if ylog: - P.yscale('log') + P.yscale("log") P.plot(bin_centers, mean_bin) if not label is None: P.ylabel(label) - def _plot_hist(self, name, ax_los='z', label=None, ylog=False): + def _plot_hist(self, name, ax_los="z", label=None, ylog=False): - pdf = self.save.get_node('/hist/' + name + '_' + ax_los) + pdf = self.save.get_node("/hist/" + name + "_" + ax_los) values, centers = pdf.read() width = centers[1] - centers[0] P.bar(centers, values, width, log=ylog) @@ -282,69 +421,65 @@ class Plotter(BaseProcessor): if not label is None: P.xlabel(label) - if '/hist/fit_' + name + '_' + ax_los in self.save: + if "/hist/fit_" + name + "_" + ax_los in self.save: slope = pdf.attrs.slope origin = pdf.attrs.origin - P.plot(centers, 10**(slope*centers + origin), '--', linewidth=2, color='orange') + P.plot( + centers, + 10 ** (slope * centers + origin), + "--", + linewidth=2, + color="orange", + ) - P.ylim([None, 1.]) + P.ylim([None, 1.0]) - def _plot(self, name_x, name_y, xlabel=None, ylabel=None, - xunit=None, yunit=None, - linearfit=False, smooth=0, - nml_key=None, runs=None, **kwargs): + def _plot( + self, + name_x, + name_y, + xlabel=None, + ylabel=None, + label=None, + xunit=None, + yunit=None, + xunit_coeff=1.0, + yunit_coeff=1.0, + linearfit=False, + smooth=0, + nml_key=None, + runs=None, + **kwargs + ): node_x = self.save.get_node(name_x) node_y = self.save.get_node(name_y) - if xlabel is None: - if 'label' in node_x._v_attrs: - xlabel = node_x._v_attrs.label - else: - xlabel = os.path.basename(name_x) - - if ylabel is None: - if 'label' in node_y._v_attrs: - ylabel = node_y._v_attrs.label - else: - ylabel = os.path.basename(name_y) - - if 'unit' in node_x._v_attrs: - xunit_old = node_x._v_attrs.unit - else: - xunit_old = cst.none - - if 'unit' in node_y._v_attrs: - yunit_old = node_y._v_attrs.unit - else: - yunit_old = cst.none - - if xunit is None: - xunit = xunit_old - if yunit is None: - yunit = yunit_old - - xlabel = xlabel + unit_str(xunit) - ylabel = ylabel + unit_str(yunit) + xlabel, xunit_old, xunit = self._ax_label_unit( + node_x, xlabel, xunit, xunit_coeff + ) + ylabel, yunit_old, yunit = self._ax_label_unit( + node_y, ylabel, yunit, yunit_coeff + ) P.xlabel(xlabel) P.ylabel(ylabel) P.grid() - if node_y._v_attrs.CLASS == 'ARRAY': + if node_y._v_attrs.CLASS == "ARRAY": x = node_x.read() * xunit_old.express(xunit) y = node_y.read() * yunit_old.express(yunit) if smooth > 0: y = gaussian_filter1d(y, sigma=smooth) yerr = None - P.plot(x, y, fmt='*', **kwargs) - elif 'mean' in node_y: + P.plot(x, y, "*", **kwargs) + elif "mean" in node_y: x = node_x.read() * xunit_old.express(xunit) y = node_y.mean.read() * yunit_old.express(yunit) if smooth > 0: - y = gaussian_filter1d(y, sigma=smooth) - yerr = node_y.std.read() * yunit_old.express(yunit) - P.errorbar(x, y, yerr=yerr, fmt='*', **kwargs) + y = gaussian_filter1d(y, sigma=smooth) + yerr = node_y.std.read() * yunit_old.express(yunit) + P.errorbar(x, y, yerr=yerr, fmt="*", **kwargs) else: yerr = None if runs is None: @@ -355,12 +490,7 @@ class Plotter(BaseProcessor): y = y_run.read() * yunit_old.express(yunit) if smooth > 0: y = gaussian_filter1d(y, sigma=smooth) - if nml_key is None: - label_run = r"{}".format(self.save.root._v_attrs.attrs[y_run.name].label) - else: - prop_name = os.path.basename(nml_key) - prop_value = self.comp.get_nml(nml_key, run) - label_run = r"{} = {}".format(prop_name, prop_value) + label_run = self._label_run(run, y_run, label, nml_key) P.plot(x, y, label=label_run, **kwargs) P.legend() @@ -368,107 +498,206 @@ class Plotter(BaseProcessor): _overlay_linearfit(x, y, yerr) def _overlay_linearfit(x, y, yerr=None): - if yerr is None: - (a, b, rho, _, stderr) = linregress(x, y) - else: - c = polyfit(x, y, 1, - w = [1.0 / ty for ty in yerr], full=False) - b, a = c[0], c[1] + if yerr is None: + (a, b, rho, _, stderr) = linregress(x, y) + else: + c = polyfit(x, y, 1, w=[1.0 / ty for ty in yerr], full=False) + b, a = c[0], c[1] - P.plot(x, a*y + b, '--', linewidth=1.5) + P.plot(x, a * y + b, "--", linewidth=1.5) + + def overlay_kennicutt(self, n0, step): + P.grid(False) + ylim = P.ylim() + (tmin, tmax) = P.xlim() + tmax = tmax + 20 + ymax = P.ylim()[1] + ssfr_sun = 2.5e-9 + ssfr_ken = ssfr_sun * n0 ** 1.4 + + coeff = ssfr_ken * 1e6 * (self.comp.info["unit_length"].express(cst.pc)) ** 2 + for i in np.arange(tmin, max(tmax, tmin + ymax / coeff), step): + t = np.linspace(0, tmax, 1000) + P.plot(t + i, t * coeff, ls="--", lw=0.9, color="grey") + P.plot(t + tmin, (t + i - tmin) * coeff, ls="--", lw=0.9, color="grey") + P.xlim(tmin, tmax) + P.ylim(ylim) def def_rules(self): self.rules = { - 'coldens' : PlotRule(self, lambda ax_los: - self._plot_map('coldens', ax_los, - label=r'$\Sigma$ (code)', vmin=0.01, vmax=100), - "Column density", dependencies=['coldens']), - 'rho' : PlotRule(self, lambda ax_los: - self._plot_map('rho', ax_los, label=r'$\rho$ (code)'), - "Density slice", dependencies=['rho']), - 'coldens_l' : PlotRule(self, lambda ax_los: - self._plot_map('coldens', ax_los, - label=r'$\Sigma$ (code)', - vmin=0.01, vmax=100, - overlays=[self._overlay_levels]), - "Column density", - dependencies=['coldens', 'levels']), - 'rho_v' : PlotRule(self, lambda ax_los: - self._plot_map('rho', ax_los, label=r'$\rho$ (code)', - overlays=[self._overlay_speed]), - "Density slice", - dependencies= ['rho', 'speed_h', 'speed_v']), - 'jeans_ratio' : PlotRule(self, lambda ax_los: - self._plot_map('jeans_ratio', ax_los, vmin=0.1, vmax=100, - cmap='RdBu_r', - overlays=[self._overlay_levels]), - "Jeans' lenght divided by the max resolution", - dependencies=['jeans_ratio', 'levels']), - 'Q' : PlotRule(self, lambda ax_los: - self._plot_map('rho', ax_los, label=r'$Q$', - vmin=0.01, vmax=100, cmap='RdBu_r'), - "Toomre Q parameter for a Keplerian disk", - dependencies=['Q'], args_ok=['z']) - + "coldens": PlotRule( + self, + partial(self._plot_map, "coldens", label=r"$\Sigma$", unit=cst.coldens), + "Column density", + dependencies=["coldens"], + ), + "rho": PlotRule( + self, + partial(self._plot_map, "rho", label=r"$\rho$", unit=cst.Msun_pc3), + "Density slice", + dependencies=["rho"], + ), + "coldens_l": PlotRule( + self, + partial( + self._plot_map, + "coldens", + label=r"$\Sigma$", + unit=cst.coldens, + overlays=[self._overlay_levels], + ), + "Column density", + dependencies=["coldens", "levels"], + ), + "rho_v": PlotRule( + self, + partial( + self._plot_map, + "rho", + label=r"$\rho$", + unit=cst.Msun_pc3, + overlays=[self._overlay_speed], + ), + "Density slice", + dependencies=["rho", "speed_h", "speed_v"], + ), + "jeans_ratio": PlotRule( + self, + partial( + self._plot_map, + "jeans_ratio", + vmin=0.1, + vmax=100, + cmap="RdBu_r", + overlays=[self._overlay_levels], + ), + "Jeans' lenght divided by the max resolution", + dependencies=["jeans_ratio", "levels"], + ), + "Q": PlotRule( + self, + partial( + self._plot_map, + "rho", + label=r"$Q$", + vmin=0.01, + vmax=100, + cmap="RdBu_r", + ), + "Toomre Q parameter for a Keplerian disk", + dependencies=["Q"], + ), } - averageables = ['coldens', 'rho', 'T', 'Q'] + averageables = ["coldens", "rho", "T", "Q"] for name in averageables: - self.rules['rad_' + name] = PlotRule(self, - partial(self._plot_radial, - 'rad_avg_' + name, - label=name, - xlog=True, ylog=True), - "Azimuthal average of {}".format(name), - dependencies=['radial_bins', 'rad_avg_' + name], - args_ok=['z']) + self.rules["rad_" + name] = PlotRule( + self, + partial( + self._plot_radial, + "rad_avg_" + name, + label=name, + xlog=True, + ylog=True, + ), + "Azimuthal average of {}".format(name), + dependencies=["radial_bins", "rad_avg_" + name], + ) - self.rules['fluct_' + name] = PlotRule(self, - partial(self._plot_map, 'fluct_' + name, - vmin=0.01, vmax=100, cmap='RdBu_r', - label='{}/avg({})'.format(name, name)), - "Fluctuation wrt to average of {}".format(name), - dependencies=['fluct_' + name], - args_ok=['z']) - self.rules['pdf_' + name] = PlotRule(self, partial(self._plot_hist, 'pdf_' + name, ylog=True, - label='{}/avg({})'.format(name, name)), - "Probability density function of {} fluctuations".format(name), - dependencies=['pdf_' + name], - args_ok=['z']) + self.rules["fluct_" + name] = PlotRule( + self, + partial( + self._plot_map, + "fluct_" + name, + vmin=0.01, + vmax=100, + cmap="RdBu_r", + label="{}/avg({})".format(name, name), + ), + "Fluctuation of {}".format(name), + dependencies=["fluct_" + name], + ) + self.rules["pdf_" + name] = PlotRule( + self, + partial( + self._plot_hist, + "pdf_" + name, + ylog=True, + label="{}/avg({})".format(name, name), + ), + "Probability density function of {} fluctuations".format(name), + dependencies=["pdf_" + name], + ) - - self.rules.update({ - 'kappa_beta' : PlotRule(self, partial(self._plot, - '/comp/beta', - '/comp/avg_pdf_slope_coldens', - linearfit=True), - args_ok=[None], kind='comp', - dependencies=['beta', 'avg_pdf_slope_coldens']), - 'sink_mass' : PlotRule(self, partial(self._plot, - '/series/sinks_from_log/time', - '/series/sinks_from_log/mass_sink', - ylabel="Mass of sinks", - xunit=cst.Myr, - yunit=cst.Msun), - args_ok=[None], kind='series', - dependencies=['sinks_from_log']), - 'assfr' : PlotRule(self, partial(self._plot, - '/series/sfr_from_log/time', - '/series/sfr_from_log/sfr', - ylabel="Averaged surfacic SFR", - xunit=cst.Myr, - yunit=cst.ssfr), - args_ok=[None], kind='series', - dependencies=['sfr_from_log']), - 'issfr' : PlotRule(self, partial(self._plot, - '/series/sinks_from_log/time', - '/series/sinks_from_log/issfr', - ylabel="Surfacic SFR", - xunit=cst.Myr, - yunit=cst.ssfr), - args_ok=[None], kind='series', - dependencies=['issfr']) - }) + self.rules.update( + { + "kappa_beta": PlotRule( + self, + partial( + self._plot, + "/comp/beta", + "/comp/avg_pdf_slope_coldens", + linearfit=True, + ), + kind="comp", + dependencies=["beta", "avg_pdf_slope_coldens"], + ), + "sink_mass": PlotRule( + self, + partial( + self._plot, + "/series/sinks_from_log/time", + "/series/sinks_from_log/mass_sink", + xunit=cst.Myr, + yunit=cst.Msun, + ), + kind="series", + dependencies=["sinks_from_log"], + ), + "assfr": PlotRule( + self, + partial( + self._plot, + "/series/sfr_from_log/time", + "/series/sfr_from_log/sfr", + ylabel="Averaged surfacic SFR", + xunit=cst.Myr, + yunit=cst.ssfr, + ), + kind="series", + dependencies=["sfr_from_log"], + ), + "issfr": PlotRule( + self, + partial( + self._plot, + "/series/sinks_from_log/time", + "/series/sinks_from_log/issfr", + ylabel="Surfacic SFR", + xunit=cst.Myr, + yunit=cst.ssfr, + ), + kind="series", + dependencies=["issfr"], + ), + "sigma": PlotRule( + self, + partial( + self._plot, + "/series/time", + "/series/time_sigma", + ylabel="$\\sigma$", + xunit=cst.Myr, + yunit=cst.km_s, + ), + kind="comp", + dependencies=["time", "time_sigma"], + ), + "plot": PlotRule( + self, lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="comp" + ), + } + ) class InteractiveGUI: @@ -478,9 +707,9 @@ class InteractiveGUI: def onbuttonrelease(self, event): """Deal with click events""" - button = ['left','middle','right'] + button = ["left", "middle", "right"] toolbar = P.get_current_fig_manager().toolbar - if toolbar.mode=='zoom rect' and event.inaxes == self.ax_col: + 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() @@ -491,36 +720,44 @@ class InteractiveGUI: def onbuttonpress(self, event): """Deal with click events""" - button = ['left','middle','right'] + 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)) + 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: + 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") + 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() + 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': + if event.key == "t": self.add_mask = not self.add_mask print("Add mode is {}".format(self.add_mask)) - elif event.key == 'r': + elif event.key == "r": self.reset_mask() def plot_side(self): - if (self.add_mask): + if self.add_mask: mask = (self.patch_mask & self.mask).flatten() else: mask = self.mask.flatten() @@ -535,7 +772,12 @@ class InteractiveGUI: 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.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() @@ -543,70 +785,60 @@ class InteractiveGUI: self.plot_side() P.draw() - - - def __init__(self, args, path, num, - maps_disk=None, - tag='', - set_lim=True) : + def __init__(self, path, run, num, path_out=None, pp_params=default_params()): """ 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 + # Get plotter object + self.plot = Plotter(path, [run], [num], path_out, pp_params) - # Load maps file - print("load maps file") - name_maps = path + '/maps_disk' + '_' + tag + '_' + format(num,'05') + '.save' + im_extent = maps_disk["im_extent"] - 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(); + 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()) + 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.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) + 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) + 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.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) + 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() diff --git a/postprocessor.py b/postprocessor.py index 3c9df57..2a3beb0 100644 --- a/postprocessor.py +++ b/postprocessor.py @@ -16,42 +16,36 @@ 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 -import pymses.utils.constants as cst - -import f90nml -import bunch +from mypool import MyPool as Pool from functools import partial from abc import ABCMeta, abstractmethod -from pp_params import * +import bunch -def unit_str(unit): - if unit == cst.none: - return '' - elif len(unit.name) > 0 : - return ' ({})'.format(unit.name) - else: - return ' ' + str(unit) - -cst.ssfr = cst.create_unit('Msun/yr/pc^2', - base_unit=cst.Msun/cst.year/cst.pc**2, - descr="Surfacic SFR", - latex='M_{\astrosun}.yr^{-1}.pc^{-2}') +from run_selector import * +from units import * class Rule: - def __init__(self, postproc, process, description='', - group='', dependencies=[], args_ok=['x', 'y', 'z'], - is_valid=lambda arg:True, kind='classic', unit=cst.none): + def __init__( + self, + postproc, + process, + description="", + group="", + dependencies=[], + is_valid=lambda arg: True, + kind="classic", + unit=cst.none, + ): self.postproc = postproc self.process_fn = process self.dependencies = dependencies self.is_valid_add = is_valid self.group = group - self.args_ok = args_ok self.description = description - self.unit=unit + self.unit = unit self.kind = kind def process(self, arg, **kwargs): @@ -64,29 +58,44 @@ class Rule: save = self.postproc.save valid = True for dep in self.dependencies: - rule_dep = self.postproc.rules[dep] - if not arg is None: - valid = valid and rule_dep.group + '/' + dep + '_' + str(arg) in save - else: - valid = valid and rule_dep.group + '/' + dep in save - return arg in self.args_ok and valid and self.is_valid_add(arg) + if dep in self.postproc.rules: + rule_dep = self.postproc.rules[dep] + if not arg is None: + valid = ( + valid and rule_dep.group + "/" + dep + "_" + str(arg) in save + ) + else: + valid = valid and rule_dep.group + "/" + dep in save + return valid and self.is_valid_add(arg) class BaseProcessor: """ Base class for processors, should not be instanciated """ + __metaclass__ = ABCMeta log_id = "" rules = dict() - data = dict() - filename = "" - save = None + solve_self_dep = True - @abstractmethod - def __init__(self): - self.def_rules() + def __init__(self, path, path_out=None, pp_params=None, tag=None): + if pp_params is None: + self.pp_params = default_params() + elif type(pp_params) == str: + self.pp_params = load_params(pp_params) + else: + self.pp_params = pp_params + + if tag is not None: + self.pp_params.out.tag = tag + + # Determining output directory + if path_out is None: + self.path_out = path + else: + self.path_out = path_out def _log(self, string, status=""): if len(status) > 0: @@ -94,65 +103,92 @@ class BaseProcessor: else: print(self.log_id + string) - def open(self): - self.save = tables.open_file(self.filename, mode="a") - - def close(self): - self.save.close() - - def process(self, to_process_list, args=[None], overwrite=False, overwrite_dep=False, **kwargs): + def process( + self, + to_process_list, + args=[None], + overwrite=False, + overwrite_dep=False, + **kwargs + ): """ Render the data in to_process_list and save them """ + + if type(to_process_list) == str: + to_process_list = [to_process_list] + + if type(args) == str or args is None: + args = [args] + self.overwrite_dep = overwrite_dep - just_done = [] # Computations done within this call + self.just_done = [] # Computations done within this call for name in to_process_list: if name in self.rules: rule = self.rules[name] for arg in args: - just_done = self._process_single(name, rule, arg, overwrite, just_done, **kwargs) + self._solve_and_process_rule(name, rule, arg, overwrite, **kwargs) else: - self._log("{} is unknown, allowed rules are {}".format(name, self.rules.keys()), "ERROR") + self._log( + "{} is unknown, allowed rules are {}".format( + name, self.rules.keys() + ), + "ERROR", + ) - def _process_single(self, name, rule, arg, overwrite=False, just_done=[], **kwargs): + return self.just_done + + def _solve_and_process_rule(self, name, rule, arg, overwrite=False, **kwargs): + updated = self._solve_dependencies(name, rule, arg, overwrite, **kwargs) + overwrite_rule = overwrite or updated + self._process_rule(name, rule, arg, overwrite, **kwargs) + + def _solve_dependencies(self, name, rule, arg, overwrite=False, **kwargs): + + self.done_before_dep = len(self.just_done) + if type(rule.dependencies) == dict: + dep_arg = rule.dependencies[dep] + else: + dep_arg = arg # Solve dependencies for dep in rule.dependencies: - if dep in self.rules: + if self.solve_self_dep and dep in self.rules: rule_dep = self.rules[dep] - just_done = self._process_single(dep, rule_dep, arg, self.overwrite_dep, - just_done) + self._solve_and_process_rule(dep, rule_dep, dep_arg, self.overwrite_dep) else: - self._log("Dependency {} for {} is unknown".format(dep, name), "ERROR") - # Process rule - self.open() - try: - done = self._process_rule(name, rule, arg, overwrite, just_done, **kwargs) - finally: - self.close() - return just_done + [done] + self._not_self_dep(name, dep, dep_arg, self.overwrite_dep, **kwargs) - def _process_rule(self, name, rule, arg, overwrite=False, just_done=[], **kwargs): + # Whether dependencies where updated + return len(self.just_done) > self.done_before_dep + + def _not_self_dep(self, name, dep, dep_arg, overwrite, **kwargs): + self._log("Dependency {} for {} is unknown".format(dep, name), "ERROR") + + def _check_existing(self, overwrite, name_full): + return overwrite or not (name_full in self.save) + + def _process_rule(self, name, rule, arg, overwrite=False, **kwargs): if not arg is None: - name_full = rule.group + '/' + name + '_' + str(arg) + name_full = rule.group + "/" + name + "_" + str(arg) else: - name_full = rule.group + '/' + name + name_full = rule.group + "/" + name if rule.is_valid(arg): - if not name_full in just_done: - if (overwrite or - (not name_full in self.save and not name_full in self.data)): + if not name_full in self.just_done: + if self._check_existing(overwrite, name_full): self._log("Processing {}".format(name_full)) data = rule.process(arg, **kwargs) self._save_data(name_full, data, rule.description, rule.unit) self._log("Data for {} computed".format(name_full), "SUCCESS") - return name_full + self.just_done.append(name_full) else: - self._log("Data for {} is already computed, skipping...".format(name_full)) + self._log( + "Data for {} is already computed, skipping...".format(name_full) + ) else: self._log("{} is not valid in this context".format(name_full), "ERROR") - def _save_data(self, name_full, data, description, unit): """ Save data in the HDF5 structure, overwrite if necessary @@ -162,13 +198,19 @@ class BaseProcessor: if type(data) == dict: if type(description) == str: - self.save.create_group(os.path.dirname(name_full), - os.path.basename(name_full), - description, createparents=True) + self.save.create_group( + os.path.dirname(name_full), + os.path.basename(name_full), + description, + createparents=True, + ) else: - self.save.create_group(os.path.dirname(name_full), - os.path.basename(name_full), - '', createparents=True) + self.save.create_group( + os.path.dirname(name_full), + os.path.basename(name_full), + "", + createparents=True, + ) if not type(unit) == dict: self.save.get_node(name_full)._v_attrs.unit = unit @@ -176,134 +218,233 @@ class BaseProcessor: for key in data: if type(description) == dict: if type(unit) == dict: - self._save_data(name_full + '/' + key, data[key], description[key], unit[key]) + self._save_data( + name_full + "/" + key, + data[key], + description[key], + unit[key], + ) else: - self._save_data(name_full + '/' + key, data[key], description[key], unit) + self._save_data( + name_full + "/" + key, data[key], description[key], unit + ) else: if type(unit) == dict: - self._save_data(name_full + '/' + key, data[key], '', unit[key]) + self._save_data(name_full + "/" + key, data[key], "", unit[key]) else: - self._save_data(name_full + '/' + key, data[key], '', unit) + self._save_data(name_full + "/" + key, data[key], "", unit) else: - self.save.create_array(os.path.dirname(name_full), - os.path.basename(name_full), - data, description, createparents=True) + self.save.create_array( + os.path.dirname(name_full), + os.path.basename(name_full), + data, + description, + createparents=True, + ) self.save.get_node(name_full).attrs.unit = unit - @abstractmethod def def_rules(self): pass -class PostProcessor(BaseProcessor): + +class HDF5Container(BaseProcessor): + + filename = "" + save = None + opened = False + + def open(self): + if not self.opened: + self.save = tables.open_file(self.filename, mode="a") + self.opened = True + + def close(self): + if self.opened: + self.save.close() + self.opened = False + + def _process_rule(self, name, rule, arg, overwrite, **kwargs): + self.open() + try: + super(HDF5Container, self)._process_rule( + name, rule, arg, overwrite, **kwargs + ) + finally: + self.close() + + def get_value(self, node_name): + self.open() + try: + node = self.save.get_node(node_name) + if node._v_attrs.CLASS == "GROUP": + value = dict() + for child_name in node._v_children: + value[child_name] = self.get_value(node_name + "/" + child_name) + else: + value = node.read() + finally: + self.close() + return value + + def get_attribute(self, node_name, attr_name): + self.open() + try: + node = self.save.get_node(node_name) + attr = node._v_attrs[attr_name] + finally: + self.close() + return attr + + +def _map_rule(rule, arg, overwrite, path, path_out, pp_params, run_num): + pp = PostProcessor( + path + "/" + run_num[0], run_num[1], path_out + "/" + run_num[0], pp_params + ) + return pp.process(rule, arg, overwrite) + + +class Aggregator: + def _not_self_dep(self, name, dep, dep_arg, overwrite, **kwargs): + if "runs" in kwargs: + dep_runs = [run for run in self.runs if run in kwargs["runs"]] + else: + dep_runs = self.runs + + pps = [[self.pp_runs[run][num] for num in self.nums[run]] for run in dep_runs] + run_num = [(run, num) for run in dep_runs for num in self.nums[run]] + map_fn = partial( + _map_rule, dep, dep_arg, overwrite, self.path, self.path_out, self.pp_params + ) + + if self.pp_params.process.num_process > 1: + pool = Pool(processes=self.pp_params.process.num_process) + done = pool.map(map_fn, run_num) + pool.close() + pool.join() + else: + done = map(map_fn, run_num) + self.just_done.extend([item for li in done for item in li]) + + +def simple_getter(name, dset): + return dset[name] + + +mass_func = lambda dset: dset["rho"] * dset.get_sizes() ** 3 # Mass function + + +class PostProcessor(HDF5Container): """ This class enable to compute and save derived quantities from the raw output """ # Axes information - _ax_nb = {'x' : 0, 'y' : 1, 'z' : 2} # 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 + _ax_nb = {"x": 0, "y": 1, "z": 2} # 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 - G = 1. # Gravitational constant + G = 1.0 # Gravitational constant - def __init__(self, path=None, num=None, path_out=None, pp_params=default_params()): + cells_loaded = False + + def __init__( + self, + path=None, + num=None, + path_out=None, + pp_params=default_params(), + tag=None, + variables=["rho", "vel", "Br", "Bl", "P", "g", "phi"], + ): """ Creates the basic structures needed for the outputs """ - if not path is None and not num is None: - # TODO : Make possible to load the HDF5 file even without the original file + super(PostProcessor, self).__init__(path, path_out, pp_params, tag) - if type(pp_params) == str: - self.pp_params = load_params(pp_params) - else : - self.pp_params = pp_params + # Open outfile + if not pp_params.out.tag == "": + tag_name = pp_params.out.tag + "_" + else: + tag_name = "" - # Determining output directory - if (path_out is None): - path_out = path + self.filename = path_out + "/postproc_" + tag_name + format(num, "05") + ".h5" - # Open outfile - if not pp_params.out.tag == '': - tag_name = pp_params.out.tag + '_' - else : - tag_name = '' + if not os.path.exists(path_out): + os.makedirs(path_out) - self.filename = (path_out + '/postproc_' + - tag_name + format(num,'05') + '.h5') + self.open() - if not os.path.exists(path_out): - os.makedirs(path_out) - self.save = tables.open_file(self.filename, mode="a", - title=os.path.basename(path)+ '_' + format(num,'05')) + # Ramses Output + self.path = path + self.run = os.path.basename(path) + self.num = num + self._ro = pymses.RamsesOutput(path, num, order=pp_params.pymses.order) + self.variables = variables + self._amr = self._ro.amr_source(self.variables) + self.info = self._ro.info.copy() - # Ramses Output - self.path = path - self.run = os.path.basename(path) - self.num = num - self._ro = pymses.RamsesOutput(path, num, order=pp_params.pymses.order) - self._amr = self._ro.amr_source(["rho","vel","P"]) - self.info = self._ro.info.copy() + # Density operator + self._rho_op = ScalarOperator( + lambda dset: dset["rho"], self._ro.info["unit_density"] + ) - # Delete specific info - self.info.pop('dom_decomp_Hilbert_keys') - self.info.pop('nstep_coarse') - self.info.pop('dom_decomp') - self.info.pop('time') + # Density ray tracer + if pp_params.pymses.fft: + self._rt = splatting.SplatterProcessor( + self._amr, self._ro.info, self._rho_op + ) + else: + self._rt = raytracing.RayTracer(self._amr, self._ro.info, self._rho_op) - # Density operator - self._rho_op = ScalarOperator(lambda dset: dset["rho"], self._ro.info["unit_density"]) + # Set the extend of the image + self._radius = 0.5 / pp_params.out.zoom + self._lbox = self.info["boxlen"] + center = pp_params.out.center + im_extent = [ + (-self._radius + center[0]) * self._lbox, + (self._radius + center[0]) * self._lbox, + (-self._radius + center[1]) * self._lbox, + (self._radius + center[1]) * self._lbox, + ] - # Density ray tracer - if(pp_params.pymses.fft): - self._rt = splatting.SplatterProcessor(self._amr, self._ro.info, self._rho_op) - else: - self._rt = raytracing.RayTracer(self._amr, self._ro.info, self._rho_op) + # Get time + time = self._ro.info["time"] # time in codeunits - # Set the extend of the image - self._radius = 0.5 / pp_params.out.zoom - self._lbox = self.info['boxlen'] - center = pp_params.out.center - im_extent = [(- self._radius + center[0]) * self._lbox, - ( self._radius + center[0]) * self._lbox, - (- self._radius + center[1]) * self._lbox, - ( self._radius + center[1]) * self._lbox] + # Set post processing attributes + self.save.root._v_attrs.dir = os.path.dirname(path) + self.save.root._v_attrs.run = os.path.basename(path) + self.save.root._v_attrs.num = num + self.save.root._v_attrs.lbox = self._lbox + self.save.root._v_attrs.time = time - # Get time - time = self._ro.info['time'] # time in codeunits + if not "/maps" in self.save: + self.save.create_group("/", "maps", "2D maps") + self.save.root.maps._v_attrs.center = center + self.save.root.maps._v_attrs.radius = self._radius + self.save.root.maps._v_attrs.im_extent = im_extent - # Set post processing attributes - self.save.root._v_attrs.dir = os.path.dirname(path) - self.save.root._v_attrs.run = os.path.basename(path) - self.save.root._v_attrs.num = num - self.save.root._v_attrs.lbox = self._lbox - self.save.root._v_attrs.time = time + # Initialize cameras + self._cam = dict() + for ax_los in self._ax_nb: # los = line of sight + ax_h = self._axes_h[ax_los] + ax_v = self._axes_v[ax_los] - if not '/maps' in self.save: - self.save.create_group('/', 'maps', '2D maps') - self.save.root.maps._v_attrs.center = center - self.save.root.maps._v_attrs.radius = self._radius - self.save.root.maps._v_attrs.im_extent = im_extent + self._cam[ax_los] = Camera( + center=pp_params.out.center, + line_of_sight_axis=ax_los, + region_size=[2.0 * self._radius, 2.0 * self._radius], + distance=self._radius, + far_cut_depth=self._radius, + up_vector=ax_v, + map_max_size=pp_params.out.map_size, + ) - # Initialize cameras - self._cam = dict() - for ax_los in self._ax_nb : # los = line of sight - ax_h = self._axes_h[ax_los] - ax_v = self._axes_v[ax_los] + self._add_metadata() + self.close() - self._cam[ax_los] = Camera(center=pp_params.out.center, - line_of_sight_axis=ax_los, - region_size=[2.*self._radius, 2.*self._radius], - distance=self._radius, - far_cut_depth=self._radius, - up_vector=ax_v, - map_max_size=pp_params.out.map_size) - - self._add_metadata() - self.save.close() - - self.log_id = "[{}, {}] ".format(self.run, self.num) + self.log_id = "[{}, {}] ".format(self.run, self.num) self.def_rules() @@ -313,64 +454,119 @@ class PostProcessor(BaseProcessor): """ # Label of the run in the label.txt file - label_filename = self.path + '/' + self.pp_params.input.label_filename + label_filename = self.path + "/" + self.pp_params.input.label_filename if os.path.exists(label_filename): - label_file = open(label_filename, 'r') + label_file = open(label_filename, "r") self.label = label_file.readline() label_file.close() else: self.label = self.run self.save.root._v_attrs.label = self.label - def _coldens(self, ax_los): - datamap = self._rt.process(self._cam[ax_los], surf_qty=True) - return datamap.map.T * self._lbox + def load_cells(self): + if not self.cells_loaded: + cell_source = CellsToPoints(self._amr) + self.cells = cell_source.flatten() + self.cells_loaded = True - def _rho(self, ax_los): - datamap_rho = slicing.SliceMap(self._amr, self._cam[ax_los], self._rho_op, z=0.) + def unload_cells(self): + if self.cells_loaded: + del self.cells + self.cells_loaded = False + + def _slice(self, getter, ax_los="z", z=0, unit=cst.none): + op = ScalarOperator(getter, unit) + datamap = slicing.SliceMap(self._amr, self._cam[ax_los], op, z=z) + return datamap.map.T + + def _ax_avg( + self, getter, ax_los, unit=cst.none, mass_weighted=True, surf_qty=False + ): + if mass_weighted: + + def num(cells): + value = getter(cells) + mass = mass_func(cells) + # Transpose (.T) is for vectorial values + return (mass * value.T).T + + op = FractionOperator(num, mass_function, unit) + else: + op = ScalarOperator(getter, unit) + + if pp_params.pymses.fft: + rt = splatting.SplatterProcessor(self._amr, self._ro.info, op) + else: + rt = raytracing.RayTracer(self._amr, self._ro.info, op) + + datamap = rt.process(self._cam[ax_los], surf_qty=surf_qty) + return datamap.map.T + + def _vol_avg(self, getter, mass_weighted=True): + self.load_cells() + value = getter(self.cells) + if mass_weighted: + mass = mass_func(self.cells) + # Transpose (.T) is for vectorial values + return np.sum((mass * value.T).T, axis=0) / np.sum(mass) + else: + return np.sum(value, axis=0) + + def _mwa_sigma(self): + mw_speed = self.save.get_node("/globals/mwa_speed").read() + + def getter(dset): + return np.sum((dset["vel"] - mw_speed) ** 2, axis=1) + + return np.sqrt(self._vol_avg(getter, mass_weighted=True)) + + def _coldens(self, ax_los): + datamap = self._rt.process(self._cam[ax_los], surf_qty=True) + return datamap.map.T + + def _rho(self, ax_los, z=0.0): + datamap_rho = slicing.SliceMap(self._amr, self._cam[ax_los], self._rho_op, z=z) return (datamap_rho.map).T - def _speed_h(self, ax_los): - vh_op = ScalarOperator(lambda dset: dset["vel"][:, self._ax_nb[self._axes_h[ax_los]]], - self._ro.info["unit_velocity"]) - dmap_vh = slicing.SliceMap(self._amr, self._cam[ax_los], vh_op, z=0.).map.T + def _speed_h(self, ax_los, z=0.0): + vh_op = ScalarOperator( + lambda dset: dset["vel"][:, self._ax_nb[self._axes_h[ax_los]]], + self._ro.info["unit_velocity"], + ) + dmap_vh = slicing.SliceMap(self._amr, self._cam[ax_los], vh_op, z=z).map.T return dmap_vh - def _speed_v(self, ax_los): - vv_op = ScalarOperator(lambda dset: dset["vel"][:, self._ax_nb[self._axes_v[ax_los]]], - self._ro.info["unit_velocity"]) - dmap_vv = slicing.SliceMap(self._amr, self._cam[ax_los], vv_op, z=0.).map.T + def _speed_v(self, ax_los, z=0.0): + vv_op = ScalarOperator( + lambda dset: dset["vel"][:, self._ax_nb[self._axes_v[ax_los]]], + self._ro.info["unit_velocity"], + ) + dmap_vv = slicing.SliceMap(self._amr, self._cam[ax_los], vv_op, z=z).map.T return dmap_vv - def _temperature(self, ax_los): + def _temperature(self, ax_los, z=0.0): P_op = ScalarOperator(lambda dset: dset["P"], self._ro.info["unit_pressure"]) - dmap_P = (slicing.SliceMap(self._amr, self._cam[ax_los], P_op, z=0.)).map.T + dmap_P = (slicing.SliceMap(self._amr, self._cam[ax_los], P_op, z=z)).map.T dmap_rho = self.save.get_node("/maps/rho_{}".format(ax_los)).read() - return dmap_P/dmap_rho + return dmap_P / dmap_rho def _levels(self, ax_los): self._amr.set_read_levelmax(20) - level_op = MaxLevelOperator() - rt_level = raytracing.RayTracer(self._amr, self._ro.info, level_op) - datamap = rt_level.process(self._cam[ax_los], surf_qty=True) + level_op = MaxLevelOperator() + rt_level = raytracing.RayTracer(self._amr, self._ro.info, level_op) + datamap = rt_level.process(self._cam[ax_los], surf_qty=True) return datamap.map.T def _jeans(self, ax_los): - dmap_T = self.save.get_node('/maps/T_' + ax_los).read() - dmap_rho = self.save.get_node('/maps/rho_' + ax_los).read() + dmap_T = self.save.get_node("/maps/T_" + ax_los).read() + dmap_rho = self.save.get_node("/maps/rho_" + ax_los).read() dmap_jeans = np.sqrt(np.pi * dmap_T / dmap_rho) return dmap_jeans def _jeans_ratio(self, ax_los): - dmap_jeans = self.save.get_node('/maps/jeans_' + ax_los).read() - dmap_levels = self.save.get_node('/maps/levels_' + ax_los).read() - dmap_jeans_ratio = dmap_jeans * 2**(dmap_levels) - return dmap_jeans_ratio - - def _jeans_ratio(self, ax_los): - dmap_jeans = self.save.get_node('/maps/jeans_' + ax_los).read() - dmap_levels = self.save.get_node('/maps/levels_' + ax_los).read() - dmap_jeans_ratio = dmap_jeans * 2**(dmap_levels) + dmap_jeans = self.save.get_node("/maps/jeans_" + ax_los).read() + dmap_levels = self.save.get_node("/maps/levels_" + ax_los).read() + dmap_jeans_ratio = dmap_jeans * 2 ** (dmap_levels) return dmap_jeans_ratio def _toomreQ_disk(self, ax_los): @@ -383,78 +579,92 @@ class PostProcessor(BaseProcessor): pos = dset.get_cell_centers() pos = pos - (self.pp_params.disk.pos_star / self._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. / rc**2) - omega_rho = omega_rho * dset["rho"] + 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 # Operator to compute the angular speed - omega_op = FractionOperator(omega_rho_func, lambda dset: dset["rho"], - 1. / self._ro.info["unit_time"]) + omega_op = FractionOperator( + omega_rho_func, lambda dset: dset["rho"], 1.0 / self._ro.info["unit_time"] + ) # Operator to compute the sound speed - cs_op = FractionOperator(lambda dset: dset["P"], - lambda dset: dset["rho"], self._ro.info["unit_velocity"]) + cs_op = FractionOperator( + lambda dset: dset["P"], + lambda dset: dset["rho"], + self._ro.info["unit_velocity"], + ) # Ray tracer for the angular speed rt_omega = raytracing.RayTracer(self._amr, self._ro.info, omega_op) # Ray tracer for the sound speed if self.pp_params.pymses.fft: - rt_cs = splatting.SplatterProcessor(self._amr, self._ro.info, cs_op, surf_qty=False) - else : + rt_cs = splatting.SplatterProcessor( + self._amr, self._ro.info, cs_op, surf_qty=False + ) + else: rt_cs = raytracing.RayTracer(self._amr, self._ro.info, cs_op) dmap_omega = rt_omega.process(self._cam[ax_los]) dmap_cs = rt_cs.process(self._cam[ax_los]) dmap_col = self.save.root.maps.coldens_z.read() - map_Q = (self._lbox * dmap_cs.map.T) * dmap_omega.map.T / (np.pi * self.G * dmap_col) - + map_Q = ( + (self._lbox * dmap_cs.map.T) + * dmap_omega.map.T + / (np.pi * self.G * dmap_col) + ) return map_Q def _radial_bins(self, _): - pos_star = self.pp_params.disk.pos_star + pos_star = self.pp_params.disk.pos_star im_extent = self.save.root.maps._v_attrs.im_extent # radius of the corner of the box plus a margin - rad_of_box = np.sqrt((im_extent[1] - pos_star[0])**2 + (im_extent[3] - pos_star[1])**2) + 0.1 + rad_of_box = ( + np.sqrt( + (im_extent[1] - pos_star[0]) ** 2 + (im_extent[3] - pos_star[1]) ** 2 + ) + + 0.1 + ) - bin_in = self.pp_params.disk.bin_in + bin_in = self.pp_params.disk.bin_in bin_out = self.pp_params.disk.bin_out - nb_bin = self.pp_params.disk.nb_bin + nb_bin = self.pp_params.disk.nb_bin # radial bins - if self.pp_params.disk.binning == 'log': + if self.pp_params.disk.binning == "log": lrad_in = np.log10(bin_in) lrad_ext = np.log10(bin_out) rad_bins = np.logspace(lrad_in, lrad_ext, num=nb_bin) - elif self.pp_params.disk.binning == 'lin': + elif self.pp_params.disk.binning == "lin": rad_bins = np.linspace(bin_in, bin_out, num=nb_bin) # Add boundaries - rad_bins = np.concatenate(([0.], rad_bins, [rad_of_box])) + rad_bins = np.concatenate(([0.0], rad_bins, [rad_of_box])) return rad_bins def _rr(self, _): im_extent = self.save.root.maps._v_attrs.im_extent - map_size = self.pp_params.out.map_size - pos_star = self.pp_params.disk.pos_star + map_size = self.pp_params.out.map_size + pos_star = self.pp_params.disk.pos_star x = np.linspace(im_extent[0], im_extent[1], map_size) y = np.linspace(im_extent[2], im_extent[3], map_size) xx, yy = np.meshgrid(x, y) - rr = np.sqrt((xx - pos_star[0])**2 + (yy - pos_star[1])**2) + rr = np.sqrt((xx - pos_star[0]) ** 2 + (yy - pos_star[1]) ** 2) return rr def _bins_on_map(self, ax_los): - rad_bins = self.save.get_node('/radial/radial_bins_' + ax_los).read() - rr = self.save.get_node('/maps/rr_' + ax_los).read() + rad_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read() + rr = self.save.get_node("/maps/rr_" + ax_los).read() # Find appropriate bin for each coordinate set bins = np.zeros(rr.shape, dtype=int) @@ -463,9 +673,9 @@ class PostProcessor(BaseProcessor): return bins def _rad_avg(self, name, ax_los): - radial_bins = self.save.get_node('/radial/radial_bins_' + ax_los).read() - bins_on_map = self.save.get_node('/maps/bins_on_map_' + ax_los).read() - dmap = self.save.get_node('/maps/' + name + '_' + ax_los).read() + radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read() + bins_on_map = self.save.get_node("/maps/bins_on_map_" + ax_los).read() + dmap = self.save.get_node("/maps/" + name + "_" + ax_los).read() # mean of all the cells in the bin mean_bin = np.zeros(len(radial_bins) - 1) @@ -475,314 +685,461 @@ class PostProcessor(BaseProcessor): def _rad_avg_map(self, name, ax_los): - radial_bins = self.save.get_node('/radial/radial_bins_' + ax_los).read() - bins_on_map = self.save.get_node('/maps/bins_on_map_' + ax_los).read() - rr = self.save.get_node('/maps/rr_' + ax_los).read() - mean_bin = self.save.get_node('/radial/rad_avg_' + name + '_' + ax_los).read() + radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read() + bins_on_map = self.save.get_node("/maps/bins_on_map_" + ax_los).read() + rr = self.save.get_node("/maps/rr_" + ax_los).read() + mean_bin = self.save.get_node("/radial/rad_avg_" + name + "_" + ax_los).read() # Add value for border mean_bin = np.concatenate(([mean_bin[0]], mean_bin)) - rr_flat = rr.flatten() + rr_flat = rr.flatten() bins_on_map_flat = bins_on_map.flatten() - # Compute the map azimuthally averaged # use linear interpolation to improve accuracy - avg_flat = (radial_bins[bins_on_map_flat + 1] - rr_flat) * mean_bin[bins_on_map_flat] - avg_flat = avg_flat + (rr_flat - radial_bins[bins_on_map_flat]) * mean_bin[bins_on_map_flat + 1] - avg_flat = avg_flat / (radial_bins[bins_on_map_flat + 1] - radial_bins[bins_on_map_flat]) - avg_map = np.reshape(avg_flat, rr.shape) + avg_flat = (radial_bins[bins_on_map_flat + 1] - rr_flat) * mean_bin[ + bins_on_map_flat + ] + avg_flat = ( + avg_flat + + (rr_flat - radial_bins[bins_on_map_flat]) * mean_bin[bins_on_map_flat + 1] + ) + avg_flat = avg_flat / ( + radial_bins[bins_on_map_flat + 1] - radial_bins[bins_on_map_flat] + ) + avg_map = np.reshape(avg_flat, rr.shape) return avg_map def _fluct_map(self, name, ax_los): - dmap = self.save.get_node('/maps/' + name + '_' + ax_los).read() - avg_map = self.save.get_node('/maps/avg_map_' + name + '_' + ax_los).read() + dmap = self.save.get_node("/maps/" + name + "_" + ax_los).read() + avg_map = self.save.get_node("/maps/avg_map_" + name + "_" + ax_los).read() return dmap / avg_map def _pdf(self, name, ax_los): - fluct_map = self.save.get_node('/maps/fluct_' + name + '_' + ax_los).read() - rr = self.save.get_node('/maps/rr_' + ax_los).read() + fluct_map = self.save.get_node("/maps/fluct_" + name + "_" + ax_los).read() + rr = self.save.get_node("/maps/rr_" + ax_los).read() - mask_pdf = (rr > self.pp_params.disk.rmin_pdf) & (rr < self.pp_params.disk.rmax_pdf) + mask_pdf = (rr > self.pp_params.disk.rmin_pdf) & ( + rr < self.pp_params.disk.rmax_pdf + ) nb_cells = np.sum(mask_pdf.flatten()) - values, edges = np.histogram(np.log10(fluct_map[mask_pdf].flatten()), - self.pp_params.pdf.nb_bin, - weights = np.ones(nb_cells) / nb_cells) + values, edges = np.histogram( + np.log10(fluct_map[mask_pdf].flatten()), + self.pp_params.pdf.nb_bin, + weights=np.ones(nb_cells) / nb_cells, + ) centers = 0.5 * (edges[1:] + edges[:-1]) return np.stack([values, centers]) def _fit_pdf(self, name, ax_los): - pdf = self.save.get_node('/hist/pdf_' + name + '_' + ax_los) + pdf = self.save.get_node("/hist/pdf_" + name + "_" + ax_los) values, centers = pdf.read() - mask_fit = ((centers > self.pp_params.pdf.xmin_fit) & - (centers < self.pp_params.pdf.xmax_fit) & - (values > 0)) - (slope, origin, correlation, _, stderr) = linregress(centers[mask_fit], np.log10(values[mask_fit])) + mask_fit = ( + (centers > self.pp_params.pdf.xmin_fit) + & (centers < self.pp_params.pdf.xmax_fit) + & (values > 0) + ) + (slope, origin, correlation, _, stderr) = linregress( + centers[mask_fit], np.log10(values[mask_fit]) + ) - pdf.attrs.slope = slope - pdf.attrs.origin = origin + pdf.attrs.slope = slope + pdf.attrs.origin = origin pdf.attrs.correlation = correlation - pdf.attrs.stderr = stderr - pdf.attrs.var = np.var + pdf.attrs.stderr = stderr + pdf.attrs.var = np.var return True def _sinks(self): - csv_name = self.path + '/output_' + str(self.num).zfill(5) + '/sink_' + str(self.num).zfill(5) + '.csv' - sinks = np.loadtxt(csv_name, delimiter=',') - header = ['Id', 'M', 'dmf', 'x', 'y', 'z', 'vx', 'vy', 'vz', - 'rot_period', 'lx', 'ly', 'lz', - 'acc_rate', 'acc_lum', 'age', 'int_lum', 'Teff'] + csv_name = ( + self.path + + "/output_" + + str(self.num).zfill(5) + + "/sink_" + + str(self.num).zfill(5) + + ".csv" + ) + sinks = np.loadtxt(csv_name, delimiter=",") + header = [ + "Id", + "M", + "dmf", + "x", + "y", + "z", + "vx", + "vy", + "vz", + "rot_period", + "lx", + "ly", + "lz", + "acc_rate", + "acc_lum", + "age", + "int_lum", + "Teff", + ] if len(sinks) == 0: sinks = np.zeros(len(header)) - sinks_dict = dict() + sinks_dict = {} for key, a in zip(header, sinks): sinks_dict[key] = a return sinks_dict def def_rules(self): + self.rules = { # Base rules - 'coldens' : Rule(self, self._coldens, "Column density", '/maps'), - 'rho' : Rule(self, self._rho, "Density slice", '/maps'), - 'speed_h' : Rule(self, self._speed_h, - "Horizontal speed slice wrt the line of sight", '/maps'), - 'speed_v' : Rule(self, self._speed_v, - "Vertical speed slice wrt the line of sight", '/maps'), - 'T' : Rule(self, self._temperature, - "Temperature slice", '/maps', dependencies=['rho']), - 'levels' : Rule(self, self._levels, - "Max level within line of sight", '/maps'), - 'jeans' : Rule(self, self._jeans, - "Jeans lenght slice", '/maps', dependencies=['rho', 'T']), - 'jeans_ratio' : Rule(self, self._jeans_ratio, - "Jeans' lenght divided by the max resolution", - '/maps', dependencies=['jeans', 'levels']), - 'Q' : Rule(self, self._toomreQ_disk, - "Toomre Q parameter for a Keplerian disk", '/maps', - dependencies=['coldens'], args_ok=['z'], - is_valid=lambda _: self.pp_params.disk.on), - 'sinks' : Rule(self, self._sinks, group="/datasets", args_ok=[None], - description={'Id': '', 'M':'[Msol]', 'dmf':'[Msol]', - 'x': '', 'y': '', 'z': '', - 'vx': '', 'vy': '', 'vz': '', - 'rot_period':'[y]', - 'lx':'|l|', 'ly':'|l|', 'lz':'|l|', - 'acc_rate':'[Msol/y]', 'acc_lum':'[Lsol]', - 'age':'[y]', - 'int_lum':'[Lsol]', 'Teff':'[K]'}), - + "coldens": Rule( + self, + self._coldens, + "Column density", + "/maps", + unit=self.info["unit_density"] * self.info["unit_length"], + ), + "rho": Rule( + self, + self._rho, + "Density slice", + "/maps", + unit=self.info["unit_density"], + ), + "speed_h": Rule( + self, + self._speed_h, + "Horizontal speed slice wrt the line of sight", + "/maps", + unit=self.info["unit_velocity"], + ), + "speed_v": Rule( + self, + self._speed_v, + "Vertical speed slice wrt the line of sight", + "/maps", + unit=self.info["unit_velocity"], + ), + "T": Rule( + self, + self._temperature, + "Temperature slice", + "/maps", + dependencies=["rho"], + unit=self.info["unit_temperature"], + ), + "levels": Rule( + self, self._levels, "Max level within line of sight", "/maps" + ), + "jeans": Rule( + self, + self._jeans, + "Jeans lenght slice", + "/maps", + dependencies=["rho", "T"], + ), + "jeans_ratio": Rule( + self, + self._jeans_ratio, + "Jeans' lenght divided by the max resolution", + "/maps", + dependencies=["jeans", "levels"], + ), + "Q": Rule( + self, + self._toomreQ_disk, + "Toomre Q parameter for a Keplerian disk", + "/maps", + dependencies=["coldens"], + is_valid=lambda _: self.pp_params.disk.on, + ), + "sinks": Rule( + self, + self._sinks, + group="/datasets", + unit={ + "Id": cst.none, + "M": cst.Msun, + "dmf": cst.Msun, + "x": "", + "y": "", + "z": "", + "vx": "", + "vy": "", + "vz": "", + "rot_period": "[y]", + "lx": "|l|", + "ly": "|l|", + "lz": "|l|", + "acc_rate": "[Msol/y]", + "acc_lum": "[Lsol]", + "age": cst.year, + "int_lum": "[Lsol]", + "Teff": cst.K, + }, + ), # Helpers - 'radial_bins' : Rule(self, self._radial_bins, - "Radial bins", '/radial', args_ok=['z']), - 'rr' : Rule(self, self._rr, "Coordinate map", '/maps', args_ok=['z']), - 'bins_on_map' : Rule(self, self._bins_on_map, - "Convert map coordinates to bins", '/maps', - dependencies=['radial_bins', 'rr'], args_ok=['z']) + "radial_bins": Rule(self, self._radial_bins, "Radial bins", "/radial"), + "rr": Rule(self, self._rr, "Coordinate map", "/maps"), + "bins_on_map": Rule( + self, + self._bins_on_map, + "Convert map coordinates to bins", + "/maps", + dependencies=["radial_bins", "rr"], + ), + # globals + "time_num": Rule( + self, + lambda: self.info["time"], + "Time", + "/globals", + unit=self.info["unit_time"], + ), + "mwa_speed": Rule( + self, + partial(self._vol_avg, partial(simple_getter, "vel")), + "Mass weighted speed average", + "/globals", + unit=self.info["unit_velocity"], + ), + "mwa_sigma": Rule( + self, + self._mwa_sigma, + "Mass weighted speed average", + "/globals", + dependencies=["mwa_speed"], + unit=self.info["unit_velocity"], + ), } # Average and other - averageables = ['coldens', 'rho', 'T', 'Q'] + averageables = ["coldens", "rho", "T", "Q"] for name in averageables: - self.rules['rad_avg_' + name] = Rule(self, partial(self._rad_avg, name), - "Azimuthal average of {}".format(name), '/radial', - dependencies=['radial_bins', 'bins_on_map', name], - args_ok=['z']) + self.rules["rad_avg_" + name] = Rule( + self, + partial(self._rad_avg, name), + "Azimuthal average of {}".format(name), + "/radial", + dependencies=["radial_bins", "bins_on_map", name], + ) - self.rules['avg_map_' + name] = Rule(self, partial(self._rad_avg_map, name), - "Interpolated map of azimuthal average of {}".format(name), - '/maps', - dependencies=['radial_bins', 'bins_on_map', - 'rr', 'rad_avg_' + name], - args_ok=['z']) - self.rules['fluct_' + name] = Rule(self, partial(self._fluct_map, name), - "Fluctuation wrt to average of {}".format(name), - '/maps', - dependencies=[name, 'avg_map_' + name], - args_ok=['z']) - self.rules['pdf_' + name] = Rule(self, partial(self._pdf, name), - "Probability density function of {} fluctuations".format(name), - '/hist', - dependencies=['rr', 'fluct_' + name], - args_ok=['z']) + self.rules["avg_map_" + name] = Rule( + self, + partial(self._rad_avg_map, name), + "Interpolated map of azimuthal average of {}".format(name), + "/maps", + dependencies=["radial_bins", "bins_on_map", "rr", "rad_avg_" + name], + ) + self.rules["fluct_" + name] = Rule( + self, + partial(self._fluct_map, name), + "Fluctuation wrt to average of {}".format(name), + "/maps", + dependencies=[name, "avg_map_" + name], + ) + self.rules["pdf_" + name] = Rule( + self, + partial(self._pdf, name), + "Probability density function of {} fluctuations".format(name), + "/hist", + dependencies=["rr", "fluct_" + name], + ) - self.rules['fit_pdf_' + name] = Rule(self, partial(self._fit_pdf, name), - "Fit the PDF of {} fluctuations".format(name), - '/hist', - dependencies=['pdf_' + name], - args_ok=['z']) + self.rules["fit_pdf_" + name] = Rule( + self, + partial(self._fit_pdf, name), + "Fit the PDF of {} fluctuations".format(name), + "/hist", + dependencies=["pdf_" + name], + ) -class Comparator(BaseProcessor): + +class Comparator(Aggregator, HDF5Container): """ Do comparaison between outputs and runs """ - def __init__(self, path, runs, nums, path_out=None, pp_params=default_params()): + def __init__( + self, + path, + in_runs, + in_nums, + path_out=None, + pp_params=default_params(), + selector=None, + tag=None, + **kwargs + ): """ Creates the basic structures needed for the outputs """ - if type(pp_params) == str: - self.pp_params = load_params(pp_params) - else : - self.pp_params = pp_params - - # Determining output directory - if (path_out is None): - path_out = path + super(Comparator, self).__init__(path, path_out, pp_params, tag) # Open outfile - if not pp_params.out.tag == '': - tag_name = '_' + pp_params.out.tag - else : - tag_name = '' + if not self.pp_params.out.tag == "": + tag_name = "_" + self.pp_params.out.tag + else: + tag_name = "" - self.filename = (path_out + '/comp' + tag_name + '.h5') - self.save = tables.open_file(self.filename, mode="a", title="Comparaison file") + self.filename = path_out + "/comp" + tag_name + ".h5" + + # Select runs + if selector is None: + selector = RunSelector(path, in_runs, in_nums, self.pp_params, **kwargs) + + # Save infos + self.path = path + self.runs = selector.runs + self.nums = selector.nums # Get postprocesor objets for each run - self.pp_runs = dict() - if not type(nums) == dict: - nums_tmp = nums - nums = dict() - for run in runs: - nums[run] = nums_tmp + self.pp_runs = {} + attrs = {} - attrs = dict() - self.namelist = dict() + for run in self.runs: + path_run = path + "/" + run + path_out_run = path_out + "/" + run + self.pp_runs[run] = {} + for num in self.nums[run]: + self.pp_runs[run][num] = PostProcessor( + path_run, num, path_out=path_out_run, pp_params=self.pp_params + ) - for run in runs: - path_run = path + '/' + run - path_out_run = path_out + '/' + run - self.pp_runs[run] = dict() - for num in nums[run]: - self.pp_runs[run][num] = PostProcessor(path_run, num, path_out=path_out_run, - pp_params=pp_params) - - h5file = tables.open_file(self.pp_runs[run][nums[run][0]].filename, 'r') - attrs[run] = h5file.root._v_attrs - h5file.close() - path_nml = path_run + '/' + self.pp_params.input.nml_filename - self.namelist[run] = bunch.bunchify(f90nml.read(path_nml)) - - - self.info = self.pp_runs[runs[0]][nums[runs[0]][0]].info - - # save metadata - self.save.root._v_attrs.runs = runs - self.save.root._v_attrs.nums = nums - self.save.root._v_attrs.attrs = attrs - self.save.root._v_attrs.namelist = bunch.unbunchify(self.namelist) + self.namelist = selector.namelist + self.info = self.pp_runs[self.runs[0]][self.nums[self.runs[0]][0]].info # log info self.log_id = "[comp {}] ".format(self.pp_params.out.tag) - self.save.close() + # Define rules self.def_rules() - def _time_series(self, name, getter): - nums = self.save.root._v_attrs.nums - series = dict() - for run in self.save.root._v_attrs.runs: - series[run] = np.zeros(len(nums[run])) - for i, num in enumerate(nums[run]): + def _check_existing(self, overwrite, name_full): + if overwrite or not (name_full in self.save): + return True + elif not "runs" in self.save.get_node(name_full)._v_attrs: + return True + else: + saved_runs = self.save.get_node(name_full)._v_attrs.runs + return len([run for run in self.runs if not run in saved_runs]) == 0 + + def _save_data(self, name_full, data, description, unit): + super(Comparator, self)._save_data(name_full, data, description, unit) + self.save.get_node(name_full)._v_attrs.runs = self.runs + + def _time_series(self, getter): + series = {} + for run in self.runs: + series[run] = np.zeros(len(self.nums[run])) + for i, num in enumerate(self.nums[run]): series[run][i] = getter(run, num) return series - def _comp(self, name, getter): - runs = self.save.root._v_attrs.runs - nums = self.save.root._v_attrs.nums - prop = np.zeros(len(runs)) - for i, run in enumerate(runs): - num = nums[run][0] + def _comp(self, getter): + prop = np.zeros(len(self.runs)) + for i, run in enumerate(self.runs): + num = self.nums[run][0] prop[i] = getter(run, num) return prop - def _time_avg(self, name): - runs = self.save.root._v_attrs.runs - mean = np.zeros(len(runs)) - std = np.zeros(len(runs)) - for i, run in enumerate(runs): - serie = self.save.get_node('/series/' + name + '/' + run).read() - mean[i] = np.mean(serie) - std[i] = np.std(serie) - return {"mean": mean, "std": std} + def _time_avg(self, name, start=None, end=None, span=None, group="/series"): + mean = np.zeros(len(self.runs)) + median = np.zeros(len(self.runs)) + std = np.zeros(len(self.runs)) - def get_attr(self, attr_name, run, num): + for i, run in enumerate(self.runs): + serie = self.save.get_node(group + "/" + name + "/" + run).read() + time = self.save.get_node(group + "/time/" + run).read() + mask = abs(serie) != np.inf + + if not ((start, end, span) == (None, None, None)): + start_r, end_r = start, end + # Be sure that start_r and end_r are defined + if start_r is None: + if end_r is None: + end_r = time[-1] + start_r = end_r - span + elif end_r is None: + end_r = start_r + span + + mask = mask & (time >= start_r) & (time <= end_r) + + mean[i] = np.nanmean(serie[mask]) + median[i] = np.nanmedian(serie[mask]) + std[i] = np.nanstd(serie[mask]) + return {"runs": self.runs, "mean": mean, "std": std, "median": median} + + def get_attr(self, attr_name, run, num, node_name="/"): pp = self.pp_runs[run][num] - h5file = tables.open_file(pp.filename, "r") - attr = h5file.root._v_attrs[attr_name] - h5file.close() - return attr + return pp.get_attribute(node_name, attr_name) - def get_nml(self, nml_key, run): + def get_global(self, node_name, run, num, args=[None], unload_cells=False): + pp = self.pp_runs[run][num] + if unload_cells: + pp.unload_cells() + value = pp.get_value(node_name) + return value + + def get_nml(self, nml_key, run, num=0): + """ + num parameter is ignored + """ res = self.namelist[run] - for key in nml_key.split('/'): + for key in nml_key.split("/"): res = res[key] return res def get_pdf_slope(self, name, run, num): pp = self.pp_runs[run][num] - pp.process(['fit_pdf_' + name], ['z'], overwrite=self.overwrite_dep) - h5file = tables.open_file(pp.filename, "r") - pdf = h5file.get_node('/hist/pdf_' + name +'_z') - slope = pdf.attrs.slope - h5file.close() + pp.process(["fit_pdf_" + name], ["z"], overwrite=self.overwrite_dep) + slope = pp.get_attribute("/hist/pdf_" + name + "_z", "slope") return slope - def get_sinks_mass(self, run, num): - pp = self.pp_runs[run][num] - pp.process(['sinks'], overwrite=self.overwrite_dep) - h5file = tables.open_file(pp.filename, "r") - sinks_mass = h5file.get_node('/datasets/sinks/M').read() - h5file.close() - return np.sum(sinks_mass) - def _extract_sinks_from_log(self, series, log_filename, run): cmd_grep = "grep 'Number of sink' {} -A 2".format(log_filename) content = os.popen(cmd_grep).readlines() for i in range(0, len(content), 4): - series['nb_sink'][run].append(np.int(content[i].split('=')[1])) - series['mass_sink'][run].append(np.float(content[i + 1].split('=')[1])) - series['time'][run].append(np.float(content[i + 2].split('=')[1])) + series["nb_sink"][run].append(np.int(content[i].split("=")[1])) + series["mass_sink"][run].append(np.float(content[i + 1].split("=")[1])) + series["time"][run].append(np.float(content[i + 2].split("=")[1])) return series def _extract_sfr_from_log(self, series, log_filename, run): cmd_grep = "grep '\[SFR' {} ".format(log_filename) content = os.popen(cmd_grep).readlines() for i in range(0, len(content)): - time = np.float(content[i].split(']')[0].split('=')[1].split()[0]) - sfr = np.float(content[i].split(']')[1].split('=')[1].split()[0]) - series['time'][run].append(time) - series['sfr'][run].append(sfr) + time = np.float(content[i].split("]")[0].split("=")[1].split()[0]) + sfr = np.float(content[i].split("]")[1].split("=")[1].split()[0]) + series["time"][run].append(time) + series["sfr"][run].append(sfr) return series def _from_log(self, keys, extractor): - nums = self.save.root._v_attrs.nums + nums = self.nums # Initialize series - series = dict() + series = {} for key in keys: - series[key] = dict() + series[key] = {} - for run in self.save.root._v_attrs.runs: + for run in self.runs: # Initialize list for key in keys: series[key][run] = [] # get one preprocessor - num = nums[run][0] - pp = self.pp_runs[run][num] + path_run = self.path + "/" + run # Get list of run files - log_files = (pp.path + '/' - + self.pp_params.input.log_prefix + '*') + log_files = path_run + "/" + self.pp_params.input.log_prefix + "*" # Parse files for log_filename in glob.glob(log_files): @@ -793,77 +1150,155 @@ class Comparator(BaseProcessor): series[key][run] = np.array(series[key][run]) # Sort the arrays - ind_sort = series['time'][run].argsort() + ind_sort = series["time"][run].argsort() for key in series: series[key][run] = series[key][run][ind_sort] return series - - def _ssfr_from_mass_sink(self): - time_unit = self.save.get_node('/series/sinks_from_log/time')._v_attrs.unit - mass_unit = self.save.get_node('/series/sinks_from_log/mass_sink')._v_attrs.unit + def _ssfr_from_mass_sink(self, avg_window=None): + """ + avg_window in year + """ + time_unit = self.save.get_node("/series/sinks_from_log/time")._v_attrs.unit + mass_unit = self.save.get_node("/series/sinks_from_log/mass_sink")._v_attrs.unit # Surface of the box in pc^2 - surface = (self.info['unit_length'].express(cst.pc))**2 - # WARNING : We do not multiply by boxlen since already done in 'unit_length' - ssfr = dict() - for run in self.save.root._v_attrs.runs: - time = self.save.get_node('/series/sinks_from_log/time/' + run).read() + surface = (self.info["unit_length"].express(cst.pc)) ** 2 + # WARNING : We do not multiply by boxlen since already done in 'unit_length' (pymses) + ssfr = {} + for run in self.runs: + time = self.save.get_node("/series/sinks_from_log/time/" + run).read() time = time * time_unit.express(cst.year) - mass_sink = self.save.get_node('/series/sinks_from_log/mass_sink/' + run).read() + mass_sink = self.save.get_node( + "/series/sinks_from_log/mass_sink/" + run + ).read() mass_sink = mass_sink * mass_unit.express(cst.Msun) - sfr = (mass_sink[1:] - mass_sink[:-1]) / (time[1:] - time[:-1]) + if avg_window is None: + shift = 1 + else: + # We assume that the timestep do not vary a lot ... + shift = np.searchsorted(time, avg_window, side="left") + sfr = (mass_sink[shift:] - mass_sink[:-shift]) / ( + time[shift:] - time[:-shift] + ) ssfr[run] = np.zeros(len(mass_sink)) - ssfr[run][1:] = sfr / surface + ssfr[run][shift:] = sfr / surface return ssfr + def _gen_rule_avg(self, rule_src_name, subarray_name=None): + + rule_src = self.rules[rule_src_name] + + if subarray_name is None: + src_name = rule_src_name + group_src = rule_src.group + unit = rule_src.unit + descr = rule_src.description + else: + src_name = subarray_name + group_src = rule_src.group + "/" + rule_src_name + unit = rule_src.unit[subarray_name] + descr = rule_src.description[subarray_name] + + description = { + "runs": "List of runs", + "mean": "Temporal average of {}".format(descr), + "std": "Standard deviation of {}".format(descr), + "median": "Median of {}".format(descr), + } + name = "avg_" + src_name + + def fn(arg=None, **kwargs): + if arg is None: + return self._time_avg(src_name, group=group_src, **kwargs) + else: + return self._time_avg( + src_name + "_" + str(arg), group=group_src, **kwargs + ) + + self.rules[name] = Rule( + self, + fn, + group="/comp", + unit=unit, + description=description, + dependencies=[rule_src_name], + ) def def_rules(self): - averageables = ['coldens', 'rho', 'T', 'Q'] + averageables = ["coldens", "rho", "T", "Q"] self.rules = { # Read from log - 'sinks_from_log' : Rule(self, - partial(self._from_log, - ['time', 'mass_sink', 'nb_sink'], - self._extract_sinks_from_log), - group='/series', args_ok=[None], - unit={'time' : self.info['unit_time'], - 'mass_sink' : cst.Msun, - 'nb_sink' : cst.none}), - 'issfr' : Rule(self, - self._ssfr_from_mass_sink, - group='/series/sinks_from_log', args_ok=[None], - unit=cst.ssfr, - description="Instantaneous surfacic star formation rate", - dependencies=['sinks_from_log']), - 'sfr_from_log' : Rule(self, - partial(self._from_log, - ['time', 'sfr'], - self._extract_sfr_from_log), - group='/series', args_ok=[None], - unit={'time' : cst.year, - 'sfr' : cst.ssfr}, - description={'time': 'Time', - 'sfr' : 'Averaged surfacic star formation rate'}), + "sinks_from_log": Rule( + self, + partial( + self._from_log, + ["time", "mass_sink", "nb_sink"], + self._extract_sinks_from_log, + ), + group="/series", + unit={ + "time": self.info["unit_time"], + "mass_sink": cst.Msun, + "nb_sink": cst.none, + }, + description={ + "time": "Time", + "mass_sink": "Total mass of stars", + "nb_sink": "Number of stars", + }, + ), + "issfr": Rule( + self, + self._ssfr_from_mass_sink, + group="/series/sinks_from_log", + unit=cst.ssfr, + description="Instantaneous surfacic star formation rate", + dependencies=["sinks_from_log"], + ), + "sfr_from_log": Rule( + self, + partial(self._from_log, ["time", "sfr"], self._extract_sfr_from_log), + group="/series", + unit={"time": cst.year, "sfr": cst.ssfr}, + description={ + "time": "Time", + "sfr": "Averaged surfacic star formation rate", + }, + ), # Read from outputs - 'time' : Rule(self, partial(self._time_series, "time", - partial(self.get_attr, 'time')), - group='/series', args_ok=[None]), - 'time_pdf_slope' : Rule(self, - lambda name: - self._time_series("pdf_slope_" + name, - partial(self.get_pdf_slope, - name)), - group='/series', args_ok = averageables), - # Averages and run properties - 'avg_pdf_slope' : Rule(self, - lambda name: - self._time_avg("time_pdf_slope_" + name), - group='/comp', dependencies=['time_pdf_slope'], - args_ok=averageables, - description={"mean": "Temporal average", - "std": "Standard deviation"}), + "time": Rule( + self, + partial( + self._time_series, partial(self.get_global, "/globals/time_num") + ), + group="/series", + unit=self.info["unit_time"], + dependencies=["time_num"], + ), + "time_sigma": Rule( + self, + partial( + self._time_series, + partial(self.get_global, "/globals/mwa_sigma", unload_cells=True), + ), + group="/series", + unit=self.info["unit_velocity"], + dependencies=["time", "mwa_sigma"], + ), + # namelist + "nml": Rule( + self, + lambda nml_key: self._comp(partial(self.get_nml, nml_key)), + group="/comp", + ), } + for name in ["issfr", "time_sigma"]: + self._gen_rule_avg(name) + + self._gen_rule_avg("sinks_from_log", "mass_sink") + self._gen_rule_avg("sinks_from_log", "nb_sink") + self._gen_rule_avg("sfr_from_log", "sfr") def get_time(path, num): @@ -880,14 +1315,19 @@ def get_time(path, num): time the time of the output (code units) """ try: - f = open(path + '/output_' + str(num).zfill(5) + '/info_' + str(num).zfill(5) + '.txt') + 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': + 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: diff --git a/pp_params.py b/pp_params.py index bb3500e..88f4bb1 100644 --- a/pp_params.py +++ b/pp_params.py @@ -11,15 +11,19 @@ _dir_path = os.path.dirname(os.path.realpath(__file__)) # Add support for '1e3' kind of float _loader = yaml.SafeLoader _loader.add_implicit_resolver( - u'tag:yaml.org,2002:float', - re.compile(u'''^(?: + u"tag:yaml.org,2002:float", + re.compile( + u"""^(?: [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)? |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+) |\\.[0-9_]+(?:[eE][-+][0-9]+)? |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]* |[-+]?\\.(?:inf|Inf|INF) - |\\.(?:nan|NaN|NAN))$''', re.X), - list(u'-+0123456789.')) + |\\.(?:nan|NaN|NAN))$""", + re.X, + ), + list(u"-+0123456789."), +) def load_params(filename): @@ -27,5 +31,6 @@ def load_params(filename): para_disk = yaml.load(f, Loader=_loader) return bunch.bunchify(para_disk) + def default_params(): - return load_params(_dir_path + '/pp_params.yml') + return load_params(_dir_path + "/pp_params.yml") diff --git a/run_selector.py b/run_selector.py new file mode 100644 index 0000000..1a287c6 --- /dev/null +++ b/run_selector.py @@ -0,0 +1,166 @@ +# coding: utf-8 + +import os +import glob +from functools import partial +from pp_params import * +import f90nml + + +class RunSelector: + def __init__( + self, + path_in, + in_runs=None, + in_nums="all", + pp_params=default_params(), + number_run="[0-9]*_", + name_run="*", + namelist_cond={}, + sort_run_by=None, + time_min=None, + time_max=None, + ): + self.path_in = path_in + self.pp_params = pp_params + + self.namelist = {} + self.runs = self.get_runs( + in_runs, number_run, name_run, namelist_cond, sort_run_by + ) + + self.info = {} + for run in self.runs: + self.info[run] = {} + self.nums = {} + + if not type(in_nums) == dict: + nums_temp = in_nums + in_nums = {} + for run in self.runs: + in_nums[run] = nums_temp + + for run in self.runs: + self.nums[run] = self.get_nums(run, in_nums[run], time_min, time_max) + + def load_namelist(self, run): + path_run = self.path_in + "/" + run + path_nml = path_run + "/" + self.pp_params.input.nml_filename + return f90nml.read(path_nml) + + def get_nml_value(self, nml_key, run): + res = self.namelist[run] + for key in nml_key.split("/"): + res = res[key] + return res + + def get_runs( + self, + in_runs=None, + number_run="[0-9]*_", + name_run="*", + namelist_cond={}, + sort_run_by=None, + ): + def try_load_nml(run): + try: + self.namelist[run] = self.load_namelist(run) + success = True + except IOError: + success = False + return success + + runs = map( + os.path.basename, glob.glob(self.path_in + "/" + number_run + name_run) + ) + if not in_runs is None: + runs = filter(lambda n: n in runs, in_runs) + runs = filter(try_load_nml, runs) + + if type(namelist_cond) == tuple: + namelist_cond = [namelist_cond] + + for (nml_key, operator, operand) in namelist_cond: + value = {} + for run in runs: + value[run] = self.get_nml_value(nml_key, run) + if operator == "=": + runs = filter(lambda r: value[r] == operand, runs) + if operator == "!=": + runs = filter(lambda r: not value[r] == operand, runs) + elif operator == ">": + runs = filter(lambda r: value[r] > operand, runs) + elif operator == "<": + runs = filter(lambda r: value[r] < operand, runs) + elif operator == "in": + runs = filter(lambda r: value[r] in operand, runs) + + # Sort by the value in the namelist of sort_run_by + if not sort_run_by is None: + if type(sort_run_by) == str: + sort_run_by = [sort_run_by] + for nml_key in reversed(sort_run_by): + runs.sort(key=partial(self.get_nml_value, nml_key)) + + return runs + + def load_info(self, run, num): + info_file = open( + self.path_in + + "/" + + run + + "/" + + "output_" + + str(num).zfill(5) + + "/" + + "info_" + + str(num).zfill(5) + + ".txt", + "r", + ) + info = {} + for line in info_file.readlines(): + parsed = yaml.safe_load(line.replace("=", ":")) + if type(parsed) == dict: + info.update(parsed) + info_file.close() + return info + + def get_nums(self, run, in_nums=None, time_min=None, time_max=None): + def try_load_info(num): + try: + self.info[run][num] = self.load_info(run, num) + success = True + except IOError: + success = False + return success + + names = glob.glob( + self.path_in + "/" + run + "/output_[0-9][0-9][0-9][0-9][0-9]" + ) + nums = map(lambda n: int(n.split("/")[-1].split("_")[1]), names) + + if type(in_nums) == list: + nums = filter(lambda n: n in nums, in_nums) + + nums = np.sort(nums) + + if in_nums == "first": + i = 0 + while i < len(nums) and not try_load_info(nums[i]): + i = i + 1 + nums = [nums[i]] + elif in_nums == "last": + i = len(nums) - 1 + while i >= 0 and not try_load_info(nums[i]): + i = i - 1 + nums = [nums[i]] + else: + nums = filter(try_load_info, nums) + + if not time_min is None: + nums = filter(lambda n: self.info[run][n]["time"] >= time_min, nums) + if not time_max is None: + nums = filter(lambda n: self.info[run][n]["time"] <= time_max, nums) + + return nums diff --git a/units.py b/units.py new file mode 100644 index 0000000..d185b55 --- /dev/null +++ b/units.py @@ -0,0 +1,71 @@ +# coding: utf-8 + +import pymses.utils.constants as cst + + +def parse_exp_unit(u): + splitted = u.split("^") + name_u = cst.Unit.from_name(splitted[0]).latex + exp = "" + if len(splitted) > 1: + exp = "$^{" + str(splitted[1]) + "}$" + return name_u + exp + + +def convert_exp(number): + splitted = "{:.4g}".format(number).split("e") + if len(splitted) == 1: + return "${}$".format(splitted[0]) + else: + coeff = float(splitted[0]) + exp = int(splitted[1]) + exp_str = "10^{" + str(exp) + "}" + if coeff == 1.0: + return "$" + exp_str + "$" + else: + return "$" + str(coeff) + "\\times" + exp_str + "$" + + +def unit_str(unit, base=None, prefix=""): + if unit == cst.none: + return "" + elif not base is None: + coeff = unit.express(base) + return unit_str(base, prefix=convert_exp(coeff) + " ") + elif len(unit.latex) > 0: + if ("." in unit.latex or "^" in unit.latex) and not "$" in unit.latex: + base_str = ".".join(map(parse_exp_unit, unit.name.split("."))) + return r" [{}{}]".format(prefix, base_str) + else: + return r" [{}{}]".format(prefix, unit.latex) + elif len(unit.name) > 0: + try: + base_str = ".".join(map(parse_exp_unit, unit.name.split("."))) + u_str = r" [{}{}]".format(prefix, base_str) + except: + u_str = r" [{}{}]".format(prefix, unit.name) + return u_str + else: + base_str = ".".join( + map(parse_exp_unit, unit._decompose_base_units().split(".")) + ) + return r" [{}{} {}]".format(prefix, unit.coeff, base_str) + + +cst.coldens = cst.create_unit( + "Msun.pc^-2", base_unit=cst.Msun / cst.pc ** 2, descr="Column density" +) +cst.km_s = cst.create_unit("km.s^-1", base_unit=cst.km / cst.s, descr="Speed") + +cst.Msun_pc3 = cst.create_unit( + "Msun.pc^-3", base_unit=cst.Msun / cst.pc ** 3, descr="Density" +) + +cst.kg_m3 = cst.create_unit("kg.m^-3", base_unit=cst.kg / cst.m ** 3, descr="Density") + +cst.ssfr = cst.create_unit( + "Msun.yr^-1.pc^-2", + base_unit=cst.Msun / cst.year / cst.pc ** 2, + descr="Surfacic SFR", + latex="M$_{\odot}$.yr$^{-1}$.p$c^{-2}$", +)