# 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 scipy.stats import linregress from functools import partial from abc import ABCMeta, abstractmethod from pp_params import * class Rule: def __init__(self, postproc, process, description='', group='', dependencies=[], args_ok=['x', 'y', 'z'], is_valid=lambda arg:True): self.postproc = postproc self.process_fn = process self.dependencies = dependencies self.is_valid_add = is_valid self.group = group self.args_ok = args_ok self.description = description def process(self, arg): if not arg is None: return self.process_fn(arg) else: return self.process_fn() def is_valid(self, arg): save = self.postproc.save valid = True for dep in self.dependencies: rule_dep = self.postproc.rules[dep] if not arg is None: valid = valid and rule_dep.group + '/' + dep + '_' + str(arg) in save else: valid = valid and rule_dep.group + '/' + dep in save return arg in self.args_ok and valid and self.is_valid_add(arg) class BaseProcessor: """ Base class for processors, should not be instanciated """ __metaclass__ = ABCMeta log_id = "" rules = dict() filename = "" @abstractmethod def __init__(self): self.def_rules() def _log(self, string, status=""): if len(status) > 0: print(status + ": " + self.log_id + string) else: print(self.log_id + string) def process(self, to_process_list, args=[None], overwrite=False, overwrite_dep=None): """ Render the data in to_process_list and save them """ if overwrite_dep is None: overwrite_dep = overwrite self.overwrite_dep = overwrite_dep just_done = [] # Computations done within this call 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 arg in args: just_done = self._process_single(name, rule, arg, overwrite, just_done) else: self._log("{} is unknown, allowed rules are {}".format(name, self.rules.keys()), "ERROR") self.save.close() def _process_single(self, name, rule, arg, overwrite=False, just_done=[]): # Solve dependencies for dep in rule.dependencies: if dep in self.rules: rule_dep = self.rules[dep] just_done = self._process_single(dep, rule_dep, arg, self.overwrite_dep, just_done) else: self._log("Dependency {} for {} is unknown".format(dep, name), "ERROR") # Process rule done = self._process_rule(name, rule, arg, overwrite, just_done) return just_done + [done] def _process_rule(self, name, rule, arg, overwrite=False, just_done=[]): if not arg is None: name_full = rule.group + '/' + name + '_' + str(arg) else: name_full = rule.group + '/' + name if rule.is_valid(arg): if not name_full in just_done: if overwrite or not name_full in self.save: self._log("Processing {}".format(name_full)) data = rule.process(arg) self._save_data(name_full, data, rule.description) self._log("Data for {} computed".format(name_full), "SUCCESS") return name_full else: self._log("Data for {} is already computed, skipping...".format(name_full)) else: self._log("{} is not valid in this context".format(name_full), "ERROR") def _save_data(self, name_full, data, description): """ Save data in the HDF5 structure, overwrite if necessary """ if name_full in self.save: self.save.remove_node(name_full, recursive=True) if not type(data) == dict: self.save.create_array(os.path.dirname(name_full), os.path.basename(name_full), data, description, createparents=True) else: for key in data: if type(description) == dict: self.save.create_array(name_full, key, data[key], description[key], createparents=True) else: self.save.create_array(name_full, key, data[key], description, createparents=True) @abstractmethod def def_rules(self): pass class PostProcessor(BaseProcessor): """ 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=None, num=None, path_out=None, pp_params=Params()): """ Creates the basic structures needed for the outputs """ if not path is None and not num is None: # 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.path = path self.run = os.path.basename(path) self.num = num 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.dir = os.path.dirname(path) self.save.root._v_attrs.run = os.path.basename(path) self.save.root._v_attrs.num = num self.save.root._v_attrs.lbox = self._lbox self.save.root._v_attrs.time = time if not '/maps' in self.save: self.save.create_group('/', 'maps', '2D maps') 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._add_metadata() self.save.close() self.log_id = "[{}, {}] ".format(self.run, self.num) self.def_rules() def _add_metadata(self): """ Add additional metadata to the file """ # Beta for the beta cooling if not (self.pp_params.disk.beta is None or self.pp_params.disk.beta == False): if type(self.pp_params.disk.beta) == int: self.save.root._v_attrs.beta = self.pp_params.disk.beta else: self.save.root._v_attrs.beta = int(self.save.root._v_attrs.run.split('_')[1][4:]) 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 _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, self._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 _radial_bins(self, _): pos_star = self.pp_params.disk.pos_star im_extent = self.save.root.maps._v_attrs.im_extent # radius of the corner of the box plus a margin rad_of_box = np.sqrt((im_extent[1] - pos_star[0])**2 + (im_extent[3] - pos_star[1])**2) + 0.1 bin_in = self.pp_params.disk.bin_in bin_out = self.pp_params.disk.bin_out nb_bin = self.pp_params.disk.nb_bin # radial bins if self.pp_params.disk.binning == 'log': lrad_in = np.log10(bin_in) lrad_ext = np.log10(bin_out) rad_bins = np.logspace(lrad_in, lrad_ext, num=nb_bin) elif self.pp_params.disk.binning == 'lin': rad_bins = np.linspace(bin_in, bin_out, num=nb_bin) # Add boundaries rad_bins = np.concatenate(([0.], rad_bins, [rad_of_box])) return rad_bins def _rr(self, _): im_extent = self.save.root.maps._v_attrs.im_extent map_size = self.pp_params.out.map_size pos_star = self.pp_params.disk.pos_star x = np.linspace(im_extent[0], im_extent[1], map_size) y = np.linspace(im_extent[2], im_extent[3], map_size) xx, yy = np.meshgrid(x, y) rr = np.sqrt((xx - pos_star[0])**2 + (yy - pos_star[1])**2) return rr def _bins_on_map(self, ax_los): rad_bins = self.save.get_node('/radial/radial_bins_' + ax_los).read() rr = self.save.get_node('/maps/rr_' + ax_los).read() # Find appropriate bin for each coordinate set bins = np.zeros(rr.shape, dtype=int) for r in rad_bins[1:]: bins = bins + (rr >= r).astype(int) return bins def _rad_avg(self, name, ax_los): radial_bins = self.save.get_node('/radial/radial_bins_' + ax_los).read() bins_on_map = self.save.get_node('/maps/bins_on_map_' + ax_los).read() dmap = self.save.get_node('/maps/' + name + '_' + ax_los).read() # mean of all the cells in the bin mean_bin = np.zeros(len(radial_bins) - 1) for j in range(len(radial_bins) - 1): mean_bin[j] = np.mean(dmap[bins_on_map == j]) return mean_bin def _rad_avg_map(self, name, ax_los): radial_bins = self.save.get_node('/radial/radial_bins_' + ax_los).read() bins_on_map = self.save.get_node('/maps/bins_on_map_' + ax_los).read() rr = self.save.get_node('/maps/rr_' + ax_los).read() mean_bin = self.save.get_node('/radial/rad_avg_' + name + '_' + ax_los).read() # Add value for border mean_bin = np.concatenate(([mean_bin[0]], mean_bin)) rr_flat = rr.flatten() bins_on_map_flat = bins_on_map.flatten() # Compute the map azimuthally averaged # use linear interpolation to improve accuracy avg_flat = (radial_bins[bins_on_map_flat + 1] - rr_flat) * mean_bin[bins_on_map_flat] avg_flat = avg_flat + (rr_flat - radial_bins[bins_on_map_flat]) * mean_bin[bins_on_map_flat + 1] avg_flat = avg_flat / (radial_bins[bins_on_map_flat + 1] - radial_bins[bins_on_map_flat]) avg_map = np.reshape(avg_flat, rr.shape) return avg_map def _fluct_map(self, name, ax_los): dmap = self.save.get_node('/maps/' + name + '_' + ax_los).read() avg_map = self.save.get_node('/maps/avg_map_' + name + '_' + ax_los).read() return dmap / avg_map def _pdf(self, name, ax_los): fluct_map = self.save.get_node('/maps/fluct_' + name + '_' + ax_los).read() rr = self.save.get_node('/maps/rr_' + ax_los).read() mask_pdf = (rr > self.pp_params.disk.rmin_pdf) & (rr < self.pp_params.disk.rmax_pdf) nb_cells = np.sum(mask_pdf.flatten()) values, edges = np.histogram(np.log10(fluct_map[mask_pdf].flatten()), self.pp_params.pdf.nb_bin, weights = np.ones(nb_cells) / nb_cells) centers = 0.5 * (edges[1:] + edges[:-1]) return np.stack([values, centers]) def _fit_pdf(self, name, ax_los): pdf = self.save.get_node('/hist/pdf_' + name + '_' + ax_los) values, centers = pdf.read() mask_fit = ((centers > self.pp_params.pdf.xmin_fit) & (centers < self.pp_params.pdf.xmax_fit) & (values > 0)) (slope, origin, correlation, _, stderr) = linregress(centers[mask_fit], np.log10(values[mask_fit])) pdf.attrs.slope = slope pdf.attrs.origin = origin pdf.attrs.correlation = correlation pdf.attrs.stderr = stderr pdf.attrs.var = np.var return True def _sinks(self): csv_name = self.path + '/output_' + str(self.num).zfill(5) + '/sink_' + str(self.num).zfill(5) + '.csv' sinks = np.loadtxt(csv_name, delimiter=',') header = ['Id', 'M', 'dmf', 'x', 'y', 'z', 'vx', 'vy', 'vz', 'rot_period', 'lx', 'ly', 'lz', 'acc_rate', 'acc_lum', 'age', 'int_lum', 'Teff'] if len(sinks) == 0: sinks = np.zeros(len(header)) sinks_dict = dict() for key, a in zip(header, sinks): sinks_dict[key] = a return sinks_dict def def_rules(self): self.rules = { # Base rules 'coldens' : Rule(self, self._coldens, "Column density", '/maps'), 'rho' : Rule(self, self._rho, "Density slice", '/maps'), 'speed_h' : Rule(self, self._speed_h, "Horizontal speed slice wrt the line of sight", '/maps'), 'speed_v' : Rule(self, self._speed_v, "Vertical speed slice wrt the line of sight", '/maps'), 'T' : Rule(self, self._temperature, "Temperature slice", '/maps', dependencies=['rho']), 'levels' : Rule(self, self._levels, "Max level within line of sight", '/maps'), 'jeans' : Rule(self, self._jeans, "Jeans lenght slice", '/maps', dependencies=['rho', 'T']), 'jeans_ratio' : Rule(self, self._jeans_ratio, "Jeans' lenght divided by the max resolution", '/maps', dependencies=['jeans', 'levels']), 'Q' : Rule(self, self._toomreQ_disk, "Toomre Q parameter for a Keplerian disk", '/maps', dependencies=['coldens'], args_ok=['z'], is_valid=lambda _: self.pp_params.disk.on), 'sinks' : Rule(self, self._sinks, group="/datasets", args_ok=[None], description={'Id': '', 'M':'[Msol]', 'dmf':'[Msol]', 'x': '', 'y': '', 'z': '', 'vx': '', 'vy': '', 'vz': '', 'rot_period':'[y]', 'lx':'|l|', 'ly':'|l|', 'lz':'|l|', 'acc_rate':'[Msol/y]', 'acc_lum':'[Lsol]', 'age':'[y]', 'int_lum':'[Lsol]', 'Teff':'[K]'}), # Helpers 'radial_bins' : Rule(self, self._radial_bins, "Radial bins", '/radial', args_ok=['z']), 'rr' : Rule(self, self._rr, "Coordinate map", '/maps', args_ok=['z']), 'bins_on_map' : Rule(self, self._bins_on_map, "Convert map coordinates to bins", '/maps', dependencies=['radial_bins', 'rr'], args_ok=['z']) } # Average and other averageables = ['coldens', 'rho', 'T', 'Q'] for name in averageables: self.rules['rad_avg_' + name] = Rule(self, partial(self._rad_avg, name), "Azimuthal average of {}".format(name), '/radial', dependencies=['radial_bins', 'bins_on_map', name], args_ok=['z']) self.rules['avg_map_' + name] = Rule(self, partial(self._rad_avg_map, name), "Interpolated map of azimuthal average of {}".format(name), '/maps', dependencies=['radial_bins', 'bins_on_map', 'rr', 'rad_avg_' + name], args_ok=['z']) self.rules['fluct_' + name] = Rule(self, partial(self._fluct_map, name), "Fluctuation wrt to average of {}".format(name), '/maps', dependencies=[name, 'avg_map_' + name], args_ok=['z']) self.rules['pdf_' + name] = Rule(self, partial(self._pdf, name), "Probability density function of {} fluctuations".format(name), '/hist', dependencies=['rr', 'fluct_' + name], args_ok=['z']) self.rules['fit_pdf_' + name] = Rule(self, partial(self._fit_pdf, name), "Fit the PDF of {} fluctuations".format(name), '/hist', dependencies=['pdf_' + name], args_ok=['z']) class Comparator(BaseProcessor): """ Do comparaison between outputs and runs """ def __init__(self, path, runs, nums, path_out=None, pp_params=Params()): """ Creates the basic structures needed for the outputs """ 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 + '/comp' + tag_name + '.h5') self.save = tables.open_file(self.filename, mode="a", title="Comparaison file") # Get postprocesor objets for each run self.pp_runs = dict() if not type(nums) == dict: nums_tmp = nums nums = dict() for run in runs: nums[run] = nums_tmp for run in runs: path_run = path + '/' + run path_out_run = path_out + '/' + run self.pp_runs[run] = dict() for num in nums[run]: self.pp_runs[run][num] = PostProcessor(path_run, num, path_out=path_out_run, pp_params=pp_params) # save metadata self.save.root._v_attrs.runs = runs self.save.root._v_attrs.nums = nums # log info self.log_id = "[comp {}] ".format(self.pp_params.out.tag) self.save.close() self.def_rules() def _time_series(self, name, getter): nums = self.save.root._v_attrs.nums series = dict() for run in self.save.root._v_attrs.runs: series[run] = np.zeros(len(nums[run])) for i, num in enumerate(nums[run]): series[run][i] = getter(self.pp_runs[run][num]) return series def _comp(self, name, getter): runs = self.save.root._v_attrs.runs nums = self.save.root._v_attrs.nums prop = np.zeros(len(runs)) for i, run in enumerate(runs): num = nums[run][0] prop[i] = getter(self.pp_runs[run][num]) return prop def _time_avg(self, name): runs = self.save.root._v_attrs.runs mean = np.zeros(len(runs)) std = np.zeros(len(runs)) for i, run in enumerate(runs): serie = self.save.get_node('/series/' + name + '/' + run).read() mean[i] = np.mean(serie) std[i] = np.std(serie) return {"mean": mean, "std": std} def _get_attr(self, attr_name, pp): h5file = tables.open_file(pp.filename, "r") attr = h5file.root._v_attrs[attr_name] h5file.close() return attr def _get_pdf_slope(self, name, pp): pp.process(['fit_pdf_' + name], ['z'], overwrite=self.overwrite_dep) h5file = tables.open_file(pp.filename, "r") pdf = h5file.get_node('/hist/pdf_' + name +'_z') slope = pdf.attrs.slope h5file.close() return slope def _get_sinks_mass(self, pp): pp.process(['sinks'], overwrite=self.overwrite_dep) h5file = tables.open_file(pp.filename, "r") sinks_mass = h5file.get_node('/datasets/sinks/M').read() h5file.close() return np.sum(sinks_mass) def def_rules(self): averageables = ['coldens', 'rho', 'T', 'Q'] self.rules = { 'beta' : Rule(self, lambda arg: self._comp("beta", partial(self._get_attr, 'beta')), group='/comp', args_ok = [None]), 'time_pdf_slope' : Rule(self, lambda name: self._time_series("pdf_slope_" + name, partial(self._get_pdf_slope, name)), group='/series', args_ok = averageables), 'time_sinks_mass' : Rule(self, partial(self._time_series, "sinks", self._get_sinks_mass), group='/series', args_ok=[None]), 'time' : Rule(self, partial(self._time_series, "time", partial(self._get_attr, 'time')), group='/series', args_ok=[None]), 'avg_pdf_slope' : Rule(self, lambda name: self._time_avg("time_pdf_slope_" + name), group='/comp', dependencies=['time_pdf_slope'], args_ok=averageables, description={"mean": "Temporal average", "std": "Standard deviation"}) } def get_time(path, num): """ Return the time of the output (code units) Parameters ---------- num output number path_out path of the pipeline output Returns ------- time the time of the output (code units) """ try: f = open(path + '/output_' + str(num).zfill(5) + '/info_' + str(num).zfill(5) + '.txt') for line in f: ls = line.split() if len(ls) > 1 and ls[0] == 'time': time = float(ls[2]) break # ro = pymses.RamsesOutput(path, num, order='>') # time = ro.info['time'] # time in codeunits f.close() return time except IOError as e: print(e) return np.nan