# coding: utf-8 import sys 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') from matplotlib.patches import Polygon import pylab as P from scipy.stats import linregress import matplotlib.patches as mpatches from matplotlib.collections import PatchCollection 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 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) class Plotter(BaseProcessor): """ This class loads derived quantities and plot them """ # 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)'} G = 1. # Gravitational constant def __init__(self, path, runs, nums, path_out=None, pp_params=default_params(), tag=None): if 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 # 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 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) # Get comparator object self.comp = Comparator(path, runs, nums, path_out, pp_params) # save infos self.runs = runs self.nums = nums self.log_id = "[plot {}] ".format(self.pp_params.out.tag) 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) else: name_full = name if rule.is_valid(arg): 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() else: 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) 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 _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 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) elif not run is None: 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 def _plot_map(self, name, ax_los, label=None, cmap='plasma', vmin=None, vmax=None, overlays=[]): 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() 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) P.xlabel(self._ax_title[ax_h]) P.ylabel(self._ax_title[ax_v]) cbar = P.colorbar(im) if not label is None: cbar.set_label(label) 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() # 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=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'); def _overlay_speed(self, ax_los): 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() 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] 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', color='grey') P.quiverkey(Q, 0.6, 0.98, max_v, r'$'+str(max_v)[0:4]+'$ (code)', 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() P.grid() P.xlabel(r'$r$') if xlog: P.xscale('log') if ylog: 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): 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) P.grid() if not label is None: P.xlabel(label) 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.ylim([None, 1.]) 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): 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) P.xlabel(xlabel) P.ylabel(ylabel) P.grid() 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: 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) else: yerr = None if runs is None: runs = self.runs for run in runs: x_run, y_run = node_x[run], node_y[run] x = x_run.read() * xunit_old.express(xunit) 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) P.plot(x, y, label=label_run, **kwargs) P.legend() if linearfit: _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] P.plot(x, a*y + b, '--', linewidth=1.5) 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']) } 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['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.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']) }) class InteractiveGUI: """ This is a matplotlib interactive session to restrain analysis to a specific area """ 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()