# 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) }