# 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 Pfrom scipy.stats import linregress import matplotlib.patches as mpatches from matplotlib.collections import PatchCollection from pp_params 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 Plotter: """ 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_out='.', filename=None, pp_params=Params()): self.pp_params = pp_params # Determining output directory if path_out is None: if filename is None: self.path_out='.' else: self.path_out = os.path.dirname(filename) else: self.path_out = path_out # Find HDF5 file if filename is None: if not pp_params.out.tag == '': tag_name = '_' + pp_params.out.tag else : tag_name = '' self.filename = (self.path_out + '/postproc_' + tag_name + format(num,'05') + '.h5') else: self.filename = filename def plot_list(self, to_plot_list, axes, overwrite=False): self._file_out = tables.open_file(self.filename, mode="r") maps = self._file_out.root.maps for ax_los in axes: for name in to_plot_list: name_full = name + '_' + ax_los plot_filename = self._find_filename(name_full) if overwrite or not os.path.exists(plot_filename): self._plot_map(name, ax_los) P.savefig(plot_filename) else: print("Data for {} is already computed, skipping...".format(name_full)) self._file_out.close() def _plot_map(self, name, ax_los): P.figure() ax_h = self._axes_h[ax_los] ax_v = self._axes_v[ax_los] im_extent = self._file_out.root.maps._v_attrs.im_extent radius = self._file_out.root.maps._v_attrs.radius center = self._file_out.root.maps._v_attrs.center if (name == 'Q' and not ax_los == 'z') or name == 'levels' or name=='speed': return dmap = self._file_out.get_node('/maps/{}_{}'.format(name, ax_los)).read() if name == 'Q' : im = P.imshow(dmap, extent=im_extent, origin='lower', cmap='RdBu', norm=mpl.colors.LogNorm(), vmin=0.01, vmax=100.) elif name == 'jeans_ratio' : im = P.imshow(dmap, extent=im_extent, origin='lower', cmap='RdBu', norm=mpl.colors.LogNorm(), vmin=0.1, vmax=1000.) else: im = P.imshow(dmap, extent=im_extent, origin='lower', norm=mpl.colors.LogNorm()) P.locator_params(axis=ax_h, nbins=pp.params.plot.ntick) P.locator_params(axis=ax_v, nbins=pp.params.plot.ntick) if(self.pp_params.put_title): pass P.xlabel(self._ax_title[ax_h]) P.ylabel(self._ax_title[ax_v]) cbar = P.colorbar(im) if name == 'coldens': cbar.set_label(r'$\Sigma$ (code)') if pp.params.set_lim: im.set_clim(0.01, 100) # if 'levels' in names: # map_level = self._file_out.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=im_extent, # origin='lower', # colors='white', # 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 name == 'rho': cbar.set_label(r'$\rho$ (code)') if 'speed' in names: dmap_vh = self._file_out.get_node('/maps/{}{}_{}'.format('v', ax_h, ax_los)).read() dmap_vv = self._file_out.get_node('/maps/{}{}_{}'.format('v', ax_v, ax_los)).read() vel_red = self.pp_params.plot.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)) 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 name == 'T': cbar.set_label(r'$T (code)$') elif name == 'Q': cbar.set_label(r'$Q$') elif name == 'jeans': cbar.set_label(r'Jeans\'s lenght') else: cbar.set_label(name) def _find_filename(self, name_full): num = self._file_out.root._v_attrs.num return (self.path_out + '/' + name_full + '_' + self.pp_params.out.tag + '_' + format(num,'05') + self.pp_params.plot.out_ext) 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()