New object oriented version using HDF5

This commit is contained in:
Noe Brucy
2019-10-23 13:48:39 +02:00
parent 7311eb3329
commit 4ff3e9f605
3 changed files with 654 additions and 0 deletions
+329
View File
@@ -0,0 +1,329 @@
# 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()