Noe Brucy
481 lines
18 KiB
Python
481 lines
18 KiB
Python
# 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 pp_params import *
|
|
from postprocessor import Rule, BaseProcessor
|
|
|
|
|
|
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, arg):
|
|
return self.process_fn(arg)
|
|
|
|
def is_valid(self, arg):
|
|
save = self.postproc.save
|
|
valid = True
|
|
for dep in self.dependencies:
|
|
if not arg is None:
|
|
valid = valid and dep + '_' + str(arg) in save
|
|
else:
|
|
valid = valid and dep in save
|
|
return arg in self.args_ok and valid 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, filename=None, path_out='.', num=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
|
|
|
|
self.log_id = "[plot {}] ".format(self.pp_params.out.tag)
|
|
self.def_rules()
|
|
|
|
def plot(self, to_plot_list, args=[None], overwrite=False):
|
|
"""
|
|
Plot the data in to_plot_list and save them
|
|
"""
|
|
self.process(to_plot_list, args, overwrite, False)
|
|
|
|
def _process_single(self, name, rule, arg, overwrite=False, overwrite_dep=False, just_done=[]):
|
|
done = self._plot_rule(name, rule, arg, overwrite)
|
|
return []
|
|
|
|
def _plot_rule(self, name, rule, arg, overwrite):
|
|
if not arg is None:
|
|
name_full = rule.group + '/' + name + '_' + str(arg)
|
|
else:
|
|
name_full = rule.group + '/' + name
|
|
|
|
if rule.is_valid(arg):
|
|
plot_filename = self._find_filename(name_full)
|
|
if overwrite or not os.path.exists(plot_filename):
|
|
rule.plot(arg)
|
|
P.tight_layout(pad=1)
|
|
P.savefig(plot_filename)
|
|
P.close()
|
|
self._log("{} plotted".format(name_full), "SUCCESS")
|
|
else:
|
|
self._log("Plot {} is already done, skipping...".format(name_full))
|
|
else:
|
|
self._log("{} is not valid in this context".format(name_full), "ERROR")
|
|
|
|
|
|
def _find_filename(self, name_full):
|
|
if not self.pp_params.out.tag == '':
|
|
tag_name = '_' + self.pp_params.out.tag
|
|
else :
|
|
tag_name = ''
|
|
|
|
if 'num' in self.save.root._v_attrs:
|
|
num = self.save.root._v_attrs.num
|
|
return (self.path_out + '/' + name_full +
|
|
tag_name + '_' + format(num,'05') +
|
|
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=[]):
|
|
P.figure()
|
|
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, ax_los='z', label=None, xlog=False, ylog=False):
|
|
P.figure()
|
|
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):
|
|
P.figure()
|
|
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, linearfit=False):
|
|
P.figure()
|
|
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:
|
|
xlabel = name_x._vattrs.label
|
|
else:
|
|
xlabel = name_x
|
|
|
|
if ylabel is None:
|
|
if 'label' in node_y:
|
|
ylabel = name_y._vattrs.label
|
|
else:
|
|
ylabel = name_y
|
|
|
|
x = node_x.read()
|
|
if node_y._v_attrs.CLASS == 'ARRAY':
|
|
y = node_y.read()
|
|
P.plot(x, y, fmt='*')
|
|
else:
|
|
y = node_y.mean.read()
|
|
yerr = node_y.std.read()
|
|
P.errorbar(x, y, yerr=y, fmt='*')
|
|
|
|
P.xlabel(xlabel)
|
|
P.ylabel(ylabel)
|
|
|
|
if linearfit:
|
|
if node_y._v_attrs.CLASS == 'ARRAY':
|
|
(a, b, rho, _, stderr) = linregress(node_x.read(), node_y.read())
|
|
else:
|
|
c = polyfit(node_x.read(), node_y.mean.read(), 1,
|
|
w = [1.0 / ty for ty in node_y.std.read()], full=False)
|
|
b, a = c[0], c[1]
|
|
|
|
P.plot(x, a*y + b, '--', linewidth=1.5)
|
|
|
|
P.grid()
|
|
|
|
|
|
|
|
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", ['/maps/coldens']),
|
|
'rho' : PlotRule(self, lambda ax_los:
|
|
self._plot_map('rho', ax_los, label=r'$\rho$ (code)'),
|
|
"Density slice", ['/maps/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", ['/maps/coldens', '/maps/levels']),
|
|
'rho_v' : PlotRule(self, lambda ax_los:
|
|
self._plot_map('rho', ax_los, label=r'$\rho$ (code)',
|
|
overlays=[self._overlay_speed]),
|
|
"Density slice", ['/maps/rho', '/maps/speed_h', '/maps/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=['/maps/jeans_ratio', '/maps/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=['/maps/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/radial_bins', '/radial/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=['/maps/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=['/hist/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], dependencies=['/comp/beta', '/comp/avg_pdf_slope_coldens']),
|
|
'sink_mass' : PlotRule(self, partial(self._plot, '/series/time', '/series/sinks_mass',
|
|
linearfit=True),
|
|
args_ok=[None], dependencies=['/series/time', '/series/sinks_mass'])
|
|
})
|
|
|
|
|
|
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()
|