Began the work on a new pipeline using the new HDF5 version

This commit is contained in:
Noe Brucy
2019-11-04 16:05:09 +01:00
parent 4ff3e9f605
commit ed55a0cba1
4 changed files with 1083 additions and 220 deletions
+252 -101
View File
@@ -9,11 +9,15 @@ 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 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'
@@ -23,7 +27,22 @@ tex_params= {'text.latex.preamble' : [r'\usepackage{amsmath}']}
P.rcParams.update(tex_params)
class Plotter:
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
"""
@@ -36,7 +55,7 @@ class Plotter:
G = 1. # Gravitational constant
def __init__(self, path_out='.', filename=None, pp_params=Params()):
def __init__(self, filename=None, path_out='.', num=None, pp_params=Params()):
self.pp_params = pp_params
@@ -60,132 +79,264 @@ class Plotter:
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
self.log_id = "[plot {}] ".format(self.pp_params.out.tag)
self.def_rules()
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))
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)
self._file_out.close()
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):
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._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
im_extent = self.save.root.maps._v_attrs.im_extent
if (name == 'Q' and not ax_los == 'z') or name == 'levels' or name=='speed':
return
dmap = self.save.get_node('/maps/{}_{}'.format(name, ax_los)).read()
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
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 name == 'coldens':
cbar.set_label(r'$\Sigma$ (code)')
if not label is None:
cbar.set_label(label)
if pp.params.set_lim:
im.set_clim(0.01, 100)
for plot_overlay in overlays:
plot_overlay(ax_los)
# 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.
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=im_extent,
# origin='lower',
# colors='white',
# linewidths=lw,
# levels=levels_ar)
# cont.levels = cont.levels + 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')
elif name == 'rho':
cbar.set_label(r'$\rho$ (code)')
P.clabel(cont,
cont.levels[cont.levels < 11],
inline=1, fontsize=8., fmt='%1d');
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()
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
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))
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')
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')
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')
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:
cbar.set_label(name)
y = node_y.mean.read()
yerr = node_y.std.read()
P.errorbar(x, y, yerr=y, fmt='*')
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)
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: