diff --git a/params.py b/params.py new file mode 100644 index 0000000..e7ecb17 --- /dev/null +++ b/params.py @@ -0,0 +1,49 @@ +# coding: utf-8 + +class PlotParams: + """ + Plot parameters + """ + out_ext = '.jpeg' # extension for plots + put_title = False # Add a title to plot + ntick = 6 # Number of ticks for maps + set_lim = True # Set default limits + vel_red = 40 # Take point each vel_red for velocities + + + +class DiskParams: + """ + Disk speficic parameters + """ + on = False # Enable specific disk analysis + pos_star = np.array([1., 1., 1.]) # Position of the central star + + +class PymsesParams: + """ + Parameters for Pymses reader + """ + order = '<' # In which order the output are read + fft = False # Quick and dirty rendering using FFT + +class OutputParams: + """ + Parameters for post processing + """ + center = [0.5, 0.5, 0.5] # Center of the image + zoom = 1. # Zoom of the image + map_size = 512 # Size of the computed maps in pixel + + tag = "" # Tag for the image + + +class Params: + """ + Strutured parameters for the post processing + """ + disk = DiskParams() + pymses = PymsesParams() + out = OutputParams() + plot = PlotParams() + diff --git a/plotter.py b/plotter.py new file mode 100644 index 0000000..46dd738 --- /dev/null +++ b/plotter.py @@ -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() diff --git a/postprocessor.py b/postprocessor.py new file mode 100644 index 0000000..c75e676 --- /dev/null +++ b/postprocessor.py @@ -0,0 +1,276 @@ +# coding: utf-8 + +import sys +import os + +import tables +import pymses +import numpy as np +from numpy.polynomial.polynomial import polyfit +from pymses.sources.ramses import output +from pymses.sources.hop.file_formats import * +from pymses.analysis import Camera, raytracing, slicing, splatting +from pymses.filters import CellsToPoints +from pymses.analysis import ScalarOperator, FractionOperator, MaxLevelOperator + +from pp_params import * + +class Rule: + + def __init__(self, process, description, group='', dependencies=[], axes=['x', 'y', 'z'], + is_valid=lambda save, ax:True): + self.process_fn = process + self.dependencies = dependencies + self.is_valid_add = is_valid + self.group = group + self.axes = axes + self.description = description + + def process(self, ax_los): + return self.process_fn(ax_los) + + def is_valid(self, save, ax): + valid = True + for dep in self.dependencies: + valid = valid and self.group + '/' + dep + '_' + ax in save + return ax in self.axes and valid and self.is_valid_add(save, ax) + +class PostProcessor: + """ + This class enable to compute and save derived quantities from the raw output + """ + + # 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 + + G = 1. # Gravitational constant + + + def __init__(self, path, num, path_out=None, pp_params=Params()): + """ + Creates the basic structures needed for the outputs + """ + + # TODO : Make possible to load the HDF5 file even without the original file + self.pp_params = pp_params + + # Determining output directory + if (path_out is None): + path_out = path + + # Open outfile + if not pp_params.out.tag == '': + tag_name = '_' + pp_params.out.tag + else : + tag_name = '' + + self.filename = (path_out + '/postproc_' + + tag_name + format(num,'05') + '.h5') + self.save = tables.open_file(self.filename, mode="a", + title=os.path.basename(path) + format(num,'05')) + + # Ramses Output + self._ro = pymses.RamsesOutput(path, num, order=pp_params.pymses.order) + self._amr = self._ro.amr_source(["rho","vel","P"]) + + # Density operator + self._rho_op = ScalarOperator(lambda dset: dset["rho"], self._ro.info["unit_density"]) + + # Density ray tracer + if(pp_params.pymses.fft): + self._rt = splatting.SplatterProcessor(self._amr, self._ro.info, self._rho_op) + else: + self._rt = raytracing.RayTracer(self._amr, self._ro.info, self._rho_op) + + # Set the extend of the image + self._radius = 0.5 / pp_params.out.zoom + self._lbox = self._ro.info['boxlen'] + center = pp_params.out.center + im_extent = [(- self._radius + center[0]) * self._lbox, + ( self._radius + center[0]) * self._lbox, + (- self._radius + center[1]) * self._lbox, + ( self._radius + center[1]) * self._lbox] + + # Get time + time = self._ro.info['time'] # time in codeunits + + # Set post processing attributes + self.save.root._v_attrs.num = num + self.save.root._v_attrs.lbox = self._lbox + self.save.root._v_attrs.time = time + + self.save.root.maps._v_attrs.center = center + self.save.root.maps._v_attrs.radius = self._radius + self.save.root.maps._v_attrs.im_extent = im_extent + + + # Initialize cameras + self._cam = dict() + for ax_los in self._ax_nb : # los = line of sight + ax_h = self._axes_h[ax_los] + ax_v = self._axes_v[ax_los] + + self._cam[ax_los] = Camera(center=pp_params.out.center, + line_of_sight_axis=ax_los, + region_size=[2.*self._radius, 2.*self._radius], + distance=self._radius, + far_cut_depth=self._radius, + up_vector=ax_v, + map_max_size=pp_params.out.map_size) + self.save.close() + self.def_rules() + + def process(self, to_process_list, axes, overwrite=False): + """ + Render the data in to_process_list and save them + """ + self.save = tables.open_file(self.filename, mode="a") + for name in to_process_list: + if name in self.rules: + rule = self.rules[name] + for ax_los in axes: + # Solve dependencies + for dep in rule.dependencies: + if dep in self.rules: + rule_dep = self.rules[dep] + self._process_rule(dep, rule_dep, ax_los, overwrite) + else: + print("ERROR: Dependency {} for {} is unknown".format(dep, name)) + # Process rule + self._process_rule(name, rule, ax_los, overwrite) + else: + print("ERROR: {} is unknown".format(name)) + self.save.close() + + def _process_rule(self, name, rule, ax_los, overwrite): + name_full = rule.group + '/' + name + '_' + ax_los + if rule.is_valid(self.save, ax_los): + if overwrite or not name_full in self.save: + data = rule.process(ax_los) + self._save_data(name_full, data, rule.description) + else: + print("Data for {} is already computed, skipping...".format(name_full)) + else: + print("ERROR: {} is not valid in this context".format(name_full)) + + + def _save_data(self, name_full, data, description): + """ + Save data in the HDF5 structure, overwrite if necessary + """ + if name_full in self.save: + node = self.save.get_node(name_full) + del node + self.save.create_array(os.path.dirname(name_full), os.path.basename(name_full), + data, description, createparents=True) + + def _coldens(self, ax_los): + datamap = self._rt.process(self._cam[ax_los], surf_qty=True) + return datamap.map.T * self._lbox + + def _rho(self, ax_los): + datamap_rho = slicing.SliceMap(self._amr, self._cam[ax_los], self._rho_op, z=0.) + return (datamap_rho.map).T + + def _speed_h(self, ax_los): + vh_op = ScalarOperator(lambda dset: dset["vel"][:, self._ax_nb[self._axes_h[ax_los]]], + self._ro.info["unit_velocity"]) + dmap_vh = slicing.SliceMap(self._amr, self._cam[ax_los], vh_op, z=0.).map.T + return dmap_vh + + def _speed_v(self, ax_los): + vv_op = ScalarOperator(lambda dset: dset["vel"][:, self._ax_nb[self._axes_v[ax_los]]], + self._ro.info["unit_velocity"]) + dmap_vv = slicing.SliceMap(self._amr, self._cam[ax_los], vv_op, z=0.).map.T + return dmap_vv + + def _temperature(self, ax_los): + P_op = ScalarOperator(lambda dset: dset["P"], self._ro.info["unit_pressure"]) + dmap_P = (slicing.SliceMap(self._amr, self._cam[ax_los], P_op, z=0.)).map.T + dmap_rho = self.save.get_node("/maps/rho_{}".format(ax_los)).read() + return dmap_P/dmap_rho + + def _levels(self, ax_los): + self._amr.set_read_levelmax(20) + level_op = MaxLevelOperator() + rt_level = raytracing.RayTracer(self._amr, self._ro.info, level_op) + datamap = rt_level.process(self._cam[ax_los], surf_qty=True) + return datamap.map.T + + def _jeans(self, ax_los): + dmap_T = self.save.get_node('/maps/T_' + ax_los).read() + dmap_rho = self.save.get_node('/maps/rho_' + ax_los).read() + dmap_jeans = np.sqrt(np.pi * dmap_T / dmap_rho) + return dmap_jeans + + def _jeans_ratio(self, ax_los): + dmap_jeans = self.save.get_node('/maps/jeans_' + ax_los).read() + dmap_levels = self.save.get_node('/maps/levels_' + ax_los).read() + dmap_jeans_ratio = dmap_jeans * 2**(dmap_levels) + return dmap_jeans_ratio + + def _toomreQ_disk(self, ax_los): + """ + Compute the Toomre Q parameter in a Keplerian disk + """ + + # Operator to compute the angular speed times rho + def omega_rho_func(dset): + pos = dset.get_cell_centers() + pos = pos - (self.pp_params.disk.pos_star / self._lbox) + xx = pos[:, :, 0] + yy = pos[:, :, 1] + rc = np.sqrt(xx**2 + yy**2) # cylindrical radius + vx = dset["vel"][:, :, 0] + vy = dset["vel"][:, :, 1] + omega_rho = (1. / rc**2) + omega_rho = omega_rho * dset["rho"] + vyx = vy * xx + vxy = vx * yy + omega_rho = omega_rho * (vyx - vxy) + return omega_rho + + # Operator to compute the angular speed + omega_op = FractionOperator(omega_rho_func, lambda dset: dset["rho"], + 1. / self._ro.info["unit_time"]) + + # Operator to compute the sound speed + cs_op = FractionOperator(lambda dset: dset["P"], + lambda dset: dset["rho"], self._ro.info["unit_velocity"]) + + # Ray tracer for the angular speed + rt_omega = raytracing.RayTracer(self._amr, self._ro.info, omega_op) + + # Ray tracer for the sound speed + if self.pp_params.pymses.fft: + rt_cs = splatting.SplatterProcessor(self._amr, ro.info, cs_op, surf_qty=False) + else : + rt_cs = raytracing.RayTracer(self._amr, self._ro.info, cs_op) + + dmap_omega = rt_omega.process(self._cam[ax_los]) + dmap_cs = rt_cs.process(self._cam[ax_los]) + dmap_col = self.save.root.maps.coldens_z.read() + map_Q = (self._lbox * dmap_cs.map.T) * dmap_omega.map.T / (np.pi * self.G * dmap_col) + + return map_Q + + def def_rules(self): + self.rules = { + 'coldens' : Rule(self._coldens, "Column density", '/maps'), + 'rho' : Rule(self._rho, "Density slice", '/maps'), + 'speed_h' : Rule(self._speed_h, "Horizontal speed slice wrt the line of sight", '/maps'), + 'speed_v' : Rule(self._speed_v, "Vertical speed slice wrt the line of sight", '/maps'), + 'T' : Rule(self._temperature, "Temperature slice", '/maps', dependencies=['rho']), + 'levels' : Rule(self._levels, "Max level within line of sight", '/maps'), + 'jeans' : Rule(self._jeans, "Jeans lenght slice", '/maps', dependencies=['rho', 'T']), + 'jeans_ratio' : Rule(self._jeans_ratio, "Jeans' lenght divided by the max resolution", + '/maps', dependencies=['jeans', 'levels']), + 'Q' : Rule(self._toomreQ_disk, "Toomre Q parameter for a Keplerian disk", '/maps', + dependencies=['coldens'], axes=['z'], + is_valid=lambda save, axe: self.pp_params.disk.on) + + } + +