# coding: utf-8 from baseprocessor import * mass_func = lambda dset: dset["rho"] * dset["dx"] ** 3 # Mass function vol_func = lambda dset: dset["dx"] ** 3 # Volume function getter_T = lambda dset: dset["P"] / dset["rho"] # Temperature class PostProcessor(HDF5Container): """ 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.0 # Gravitational constant cells_loaded = False def __init__(self, path=None, num=None, path_out=None, pp_params=None, tag=None): """ Creates the basic structures needed for the outputs """ super(PostProcessor, self).__init__(path, path_out, pp_params, tag) # Open outfile if not self.pp_params.out.tag == "": tag_name = self.pp_params.out.tag + "_" else: tag_name = "" self.filename = path_out + "/postproc_" + tag_name + format(num, "05") + ".h5" self.cells_filename = ( path_out + "/cells_" + tag_name + format(num, "05") + ".h5" ) if not os.path.exists(path_out): os.makedirs(path_out) self.open() # 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, verbose=pp_params.pymses.verbose ) self._amr = self._ro.amr_source(self.pp_params.pymses.variables) self.info = self._ro.info.copy() # 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.pymses.zoom self._lbox = self.info["boxlen"] center = pp_params.pymses.center im_extent = [ (-self._radius + center[0]), (self._radius + center[0]), (-self._radius + center[1]), (self._radius + center[1]), ] # 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.unit_length = self.info["unit_length"] 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 = {} 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.pymses.center, line_of_sight_axis=ax_los, region_size=[2.0 * self._radius, 2.0 * self._radius], distance=self._radius, far_cut_depth=self._radius, up_vector=ax_v, map_max_size=pp_params.pymses.map_size, ) self.close() self.log_id = "[{}, {}] ".format(self.run, self.num) self.def_rules() def load_cells(self): """ Load all cells from the source file in the memory. Cells will be accessible trough self.cells (/!\ Long and memory heavy) """ if not self.cells_loaded: if os.path.exists(self.cells_filename): cells_hdf5 = tables.open_file(self.cells_filename, mode="r") try: node = cells_hdf5.get_node("/cells") self.cells = {} for key in node._v_children: self.cells[key] = cells_hdf5.get_node("/cells/" + key).read() finally: cells_hdf5.close() else: cell_source = CellsToPoints(self._amr) cells_pymses = cell_source.flatten() self.cells = {} for key in cells_pymses.fields: self.cells[key] = cells_pymses[key] self.cells["dx"] = cells_pymses.get_sizes() if self.pp_params.process.save_cells: cells_hdf5 = tables.open_file(self.cells_filename, mode="w") try: for key in self.cells: cells_hdf5.create_array( "/cells", key, self.cells[key], "", createparents=True ) finally: cells_hdf5.close() self.cells_loaded = True def unload_cells(self): """ Free space in the memory by telling the garbage collectors that self.cells is not needed """ if self.cells_loaded: del self.cells self.cells_loaded = False def _slice(self, getter, ax_los="z", z=0, unit=cst.none): """ Slice process function. Return a slice of the source box. Parameters ---------- getter : callable A callable that extract the wanted data from a pymses dataset ax_los : string The axis perpendicular to the slice plane z : float Coordinate of the slice on the ax_los axis unit : cst.Unit Unit of the resulting dataset Returns ------- A numpy array containing the slice """ op = ScalarOperator(getter, unit) datamap = slicing.SliceMap(self._amr, self._cam[ax_los], op, z=z) return datamap.map.T def _ax_avg( self, getter, ax_los, unit=cst.none, mass_weighted=True, surf_qty=False ): """ Map of the average of a quantity (given by getter) along an axis (ax_los) """ if mass_weighted: def num(cells): value = getter(cells) mass = mass_func(cells) # Transpose (.T) is for vectorial values return (mass * value.T).T op = FractionOperator(num, mass_function, unit) else: op = ScalarOperator(getter, unit) if pp_params.pymses.fft: rt = splatting.SplatterProcessor(self._amr, self._ro.info, op) else: rt = raytracing.RayTracer(self._amr, self._ro.info, op) datamap = rt.process(self._cam[ax_los], surf_qty=surf_qty) return datamap.map.T def _vol_avg(self, getter, mass_weighted=True): self.load_cells() value = getter(self.cells) if mass_weighted: mass = mass_func(self.cells) # Transpose (.T) is for vectorial values return np.sum((mass * value.T).T, axis=0) / np.sum(mass) else: return np.sum(value, axis=0) def _vol_pdf(self, getter, bins=100, logbins=False, weight_func=vol_func): self.load_cells() data = getter(self.cells) if logbins: data = np.log10(data) weights = weight_func(self.cells) values, edges = np.histogram(data, bins, weights=weights) centers = 0.5 * (edges[1:] + edges[:-1]) return (np.stack([values, centers]), {"logbins": logbins}) def _mwa_sigma(self, axes=["x", "y", "z"]): mw_speed = self.save.get_node("/globals/mwa_speed").read() if axes == ["x", "y", "z"]: def getter(dset): return np.sum((dset["vel"] - mw_speed) ** 2, axis=1) else: def getter(dset): sigma_squared = 0.0 for ax in axes: ax_nb = self._ax_nb[ax] sigma_sq_ax = (dset["vel"][:, ax_nb] - mw_speed[ax_nb]) ** 2 sigma_squared = sigma_squared + sigma_sq_ax return sigma_squared return np.sqrt(self._vol_avg(getter, mass_weighted=True)) def _coldens(self, ax_los): datamap = self._rt.process(self._cam[ax_los], surf_qty=True) return datamap.map.T def _rho(self, ax_los, z=0.0): datamap_rho = slicing.SliceMap(self._amr, self._cam[ax_los], self._rho_op, z=z) return (datamap_rho.map).T def _speed_h(self, ax_los, z=0.0): 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=z).map.T return dmap_vh def _speed_v(self, ax_los, z=0.0): 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=z).map.T return dmap_vv def _temperature(self, ax_los, z=0.0): 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=z)).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(self.pp_params.pymses.levelmax) 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.0 / 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.0 / 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, _): """ Computes radial bins (for disk) """ 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.0], rad_bins, [rad_of_box])) return rad_bins def _rr(self, _): """ Computes the radius from the center """ im_extent = self.save.root.maps._v_attrs.im_extent map_size = self.pp_params.pymses.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 _rad_fluct_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 _clumps(self): name = self.path_out + "/" + self.tag + "_" + str(self.num).zfill(5) hop_save = name + "_hop" + "_prop_struct.save" me.make_clump_hop( self.path, self.num, name + "_hop", self.pp_params.hop.rho_thres, self.pp_params.hop.lvl_thres, [0.5, 0.5, 0.5], 1, path_out=path_out + "/", path_hop="./", force=True, gcomp=False, ) hop_save = me.clump_properties( name + "_hop", path, num, path_out=path_out + "/", gcomp=False ) f = open(path_out + "/" + hop_save) hop_data = pickle.load(f) f.close() return hop_data 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 = {} 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", unit=self.info["unit_density"] * self.info["unit_length"], ), "rho": Rule( self, self._rho, "Density slice", "/maps", unit=self.info["unit_density"], ), "speed_h": Rule( self, self._speed_h, "Horizontal speed slice wrt the line of sight", "/maps", unit=self.info["unit_velocity"], ), "speed_v": Rule( self, self._speed_v, "Vertical speed slice wrt the line of sight", "/maps", unit=self.info["unit_velocity"], ), "T": Rule( self, self._temperature, "Temperature slice", "/maps", dependencies=["rho"], unit=self.info["unit_temperature"], ), "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"], is_valid=lambda _: self.pp_params.disk.on, ), "sinks": Rule( self, self._sinks, group="/datasets", unit={ "Id": cst.none, "M": cst.Msun, "dmf": cst.Msun, "x": "", "y": "", "z": "", "vx": "", "vy": "", "vz": "", "rot_period": "[y]", "lx": "|l|", "ly": "|l|", "lz": "|l|", "acc_rate": "[Msol/y]", "acc_lum": "[Lsol]", "age": cst.year, "int_lum": "[Lsol]", "Teff": cst.K, }, ), "clumps": Rule(self, self._clumps, group="/datasets"), # Helpers "radial_bins": Rule(self, self._radial_bins, "Radial bins", "/radial"), "rr": Rule(self, self._rr, "Coordinate map", "/maps"), "bins_on_map": Rule( self, self._bins_on_map, "Convert map coordinates to bins", "/maps", dependencies=["radial_bins", "rr"], ), # PDF "rho_pdf": Rule( self, partial(self._vol_pdf, partial(simple_getter, "rho"), logbins=True), "Global rho-PDF", "/hist", unit=self.info["unit_density"], ), "T_pdf": Rule( self, partial(self._vol_pdf, getter_T, logbins=True), "Global T-PDF", "/hist", unit=self.info["unit_temperature"], ), "P_pdf": Rule( self, partial(self._vol_pdf, getter_T, logbins=True), "Global P-PDF", "/hist", unit=self.info["unit_pressure"], ), # globals "time_num": Rule( self, lambda: self.info["time"], "Time", "/globals", unit=self.info["unit_time"], ), "mwa_speed": Rule( self, partial(self._vol_avg, partial(simple_getter, "vel")), "Mass weighted speed average", "/globals", unit=self.info["unit_velocity"], ), "mwa_sigma": Rule( self, self._mwa_sigma, "Mass weighted speed average", "/globals", dependencies={"mwa_speed": None}, unit=self.info["unit_velocity"], ), } # 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], ) 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], ) self.rules["fluct_" + name] = Rule( self, partial(self._fluct_map, name), "Fluctuation wrt to average of {}".format(name), "/maps", dependencies=[name, "avg_map_" + name], ) self.rules["pdf_" + name] = Rule( self, partial(self._rad_fluct_pdf, name), "Probability density function of {} fluctuations".format(name), "/hist", dependencies=["rr", "fluct_" + name], ) self.rules["fit_pdf_" + name] = Rule( self, partial(self._fit_pdf, name), "Fit the PDF of {} fluctuations".format(name), "/hist", dependencies=["pdf_" + name], ) self._gen_rule_transform("fluct_coldens", np.max, "max", group="/globals") super(PostProcessor, self).def_rules() 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 f.close() return time except IOError as e: print(e) return np.nan