# coding: utf-8 import numpy as np import tables import pickle import astropy.units as u import pandas as pd from struct import pack from skimage.morphology import medial_axis import os from distutils.file_util import copy_file from functools import partial from scipy.stats import linregress from astrophysix.simdm.results import Snapshot import pymses import pymses.utils.regions as reg from pymses.analysis import ( Camera, FractionOperator, MaxLevelOperator, ScalarOperator, raytracing, slicing, splatting, cube3d, ) from pymses.filters import CellsToPoints, RegionFilter from pymses.sources.hop.hop import HOP try: from fil_finder import FilFinder2D except ModuleNotFoundError: print("WARNING: No filFinder support") from scipy import fft from scripts import pspec as pspec from utils.units import U from baseprocessor import ( HDF5Container, Rule, norm_getter, oct_norm_getter, simple_getter, vect_getter, oct_vect_getter, ) from utils.runselector import RunSelector # Getters def mass_func(dset): dx = dset["dx"] return dset["rho"] * dx**3 # Mass function def vol_func(dset): return dset["dx"] ** 3 # Volume function def getter_T(dset): return dset["P"] / dset["rho"] # Temperature def getter_P(dset): return dset["P"] def getter_abs_cos_vB(dset): B_norm = np.sqrt(np.sum(dset["Br"] ** 2, axis=1)) v_norm = np.sqrt(np.sum(dset["vel"] ** 2, axis=1)) # Compute the dot product in each cell dot_prod = np.einsum("ij,ij->i", dset["vel"], dset["Br"]) return np.abs(dot_prod) / (v_norm * B_norm) def getter_B_int(dset): B_norm = np.sqrt(np.sum(dset["Br"] ** 2, axis=1)) return B_norm def getter_rho(dset): return dset["rho"] def getter_v_norm(dset): v_norm = np.sqrt(np.sum(dset["vel"] ** 2, axis=1)) return v_norm # Helpers def mean_by_bins( x, y, bins=100, logbins=False, weights=None, range=None, ): """ Compute the mean of y in bins of x Parameters ---------- x, y : np.array of same dimensions bins : int, number of bins logbins : bool, if true, the bins will be logaritmically distributed weights : np.array, same size as y, default None. range : tuple of float, default None. """ mask = np.isfinite(x) & np.isfinite(y) x = x[mask].flatten() y = y[mask].flatten() if logbins: if range is None: range = (np.min(x[x > 0]), np.max(x)) x_bins = np.logspace(np.log10(range[0]), np.log10(range[1]), bins, base=10) else: if range is None: range = (np.min(x), np.max(x)) x_bins = np.linspace(range[0], range[1], bins) # For each cell, bin_number contains the number of the bins it belongs to bin_number = np.zeros(len(y)) # Go through the min value of x of each bin for x_min in x_bins[1:-1]: bin_number = bin_number + (x > x_min).astype(int) # Compute the mean in each bin y_mean = np.zeros(len(x_bins) - 1) for i in np.arange(len(y_mean)): mask_bin = bin_number == i if weights is None: y_mean[i] = np.mean(y[mask_bin]) else: y_mean[i] = np.sum(y[mask_bin] * weights[mask_bin]) / np.sum( weights[mask_bin] ) # Get the center of each bin if logbins: centers = 10 ** (0.5 * (np.log10(x_bins[1:]) + np.log10(x_bins[:-1]))) else: centers = 0.5 * (x_bins[1:] + x_bins[:-1]) return centers, y_mean # Filament helpers def find_center(distance, skeleton, i_center, j_center, i, j): """ Given a distance array, find the cells at a center of a filament at a given postion """ if skeleton[i, j]: i_center[i, j], j_center[i, j] = i, j return i, j elif i_center[i, j] or j_center[i, j]: return i_center[i, j], j_center[i, j] else: i_neigh = np.array([i - 1, i, i + 1]) i_neigh = i_neigh[(i_neigh > 0) & (i_neigh < distance.shape[0])] j_neigh = np.array([j - 1, j, j + 1]) j_neigh = j_neigh[(j_neigh > 0) & (j_neigh < distance.shape[1])] ii_neigh, jj_neigh = np.meshgrid(i_neigh, j_neigh) d_neigh = distance[ii_neigh, jj_neigh] ind_max = np.unravel_index(np.argmax(d_neigh), d_neigh.shape) i_max, j_max = ii_neigh[ind_max], jj_neigh[ind_max] if i_max == i and j_max == j: i_center[i, j], j_center[i, j] = i, j else: i_center[i, j], j_center[i, j] = find_center( distance, skeleton, i_center, j_center, i_max, j_max ) return i_center[i, j], j_center[i, j] # Power spectrum helpers def pspec2d(map2D): """ Computes the 2D powerspectrum of a 2D map Parameters ---------- map2D : square map to compute the fft from """ # Resolution of the map n = map2D.shape[0] assert map2D.shape[1] == n # Create bin array for the wavenumber k # (Take into account the symmetry) kbins = np.linspace(0, n // 2, n // 4) k_alias = np.arange(n, dtype=np.float64) k = np.where(k_alias >= n // 2, k_alias - n, k_alias) kx, ky = np.meshgrid(k, k, indexing="ij") # Compute map of k kmap = np.sqrt(kx**2 + ky**2) # Compute fft fmap = fft.fft2(map2D) # Compute the power map from the fft pmap = pspec.pcube(fmap) # Use the power map and the fft to compute the powerspectrum # This is typically an histogram of k weighted by the fourier transform value power, kbins, power2, fbins = pspec.pspectrum(pmap, kmap, kbins, 1, 0) # Return bin center and power spectrum return 0.5 * (kbins[1:] + kbins[:-1]), power def degrade_map(dmap, nnew, integrate=False): """Degrade a data dmap to a coarser resolution Parameters: ----------- cube (numpy.ndarray, ndim=2) input data cube, lengths in each direction must integrate (bool, default False) if True, values are added instead of averaged Return value: ------------- degraded_cube (numpy.ndarray, ndim=3) output data cube, 2**lvl values in each direction """ assert dmap.ndim == 2 nold = dmap.shape[0] assert nnew <= nold nsum, rem = divmod(nold, nnew) assert rem == 0 # For each direction, we split the corresponding axis (length nold) into # 2 axes, the first one being the axis for the output dmap (length nnew), # and the second one the axis to average or integrate over (fine cells). # reshape creates a view, no copy is involved v = dmap.reshape((nnew, nsum, nnew, nsum), order="C") # Even indexes are kept, odd ones are to be summed dmap_new = np.einsum("iajb->ij", v) if not integrate: dmap_new /= nsum return dmap_new class SnapshotProcessor(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 # Pymses unit key of amr fiels unit_key = { "rho": "unit_density", "vel": "unit_velocity", "Br": "unit_mag", "Bl": "unit_mag", "P": "unit_pressure", "g": {"unit_gravpot": 1, "unit_length": -1}, "phi": "unit_gravpot", "mass": "unit_mass", "epoch": "unit_time", "size": "unit_length", "dx": "unit_length", "pos": "unit_length", "id": U.none, "level": U.none, } # Units sinks unit_sinks = { "id": U.none, "msink": U.Msun, "dmfsink": U.Msun, "x": "unit_density", "y": "unit_density", "z": "unit_density", "vx": "unit_velocity", "vy": "unit_velocity", "vz": "unit_velocity", "rot_period": "unit_time", "lx": {"unit_mass": 1, "unit_length": 2, "unit_time": -1}, "ly": {"unit_mass": 1, "unit_length": 2, "unit_time": -1}, "lz": {"unit_mass": 1, "unit_length": 2, "unit_time": -1}, "acc_rate": {"unit_mass": 1, "unit_time": -1}, "acc_lum": U.none, "age": U.year, "int_lum": U.none, "Teff": U.K, "level": U.none, "mbh": "unit_mass", "tform": "unit_time", "del_mass": U.Msun, "rho_gas": "unit_density", "cs**2": {"unit_velocity": 2}, "etherm": {"unit_mass": 1, "unit_velocity": 2}, "vx_gas": "unit_velocity", "vy_gas": "unit_velocity", "vz_gas": "unit_velocity", } G = 1.0 # Gravitational constant cells_loaded = False parts_loaded = False def __init__( self, path=".", num=None, path_out=".", params=None, tag=None, selector=None, unit_time=U.year, ): """ Creates the basic structures needed for the outputs Parameters ---------- path : str, path of where the outputs are located path_out : str, where to store postprocessings results params : str or params instance tag : str, distinguish this postprocessing set selector : RunSelector instance unit_time : Unit instance, used for astrophysix """ self.path = path self.run = os.path.basename(path) self.num = num self.log_id = "snap({}, {})".format(self.run, self.num) super(SnapshotProcessor, self).__init__(path, path_out, params, tag) # Open outfile if not self.params.out.tag == "": tag_name = self.params.out.tag + "_" else: tag_name = "" if self.params.out.ext_subfolder: subfolder = "/h5" else: subfolder = "" self.filename = f"{self.path_out}{subfolder}/postproc_{tag_name}{num:05}.h5" self.cells_filename = ( f"{self.path_out}{subfolder}/cells_{tag_name}{num:05}.h5" # legacy only ) self.parts_filename = ( f"{self.path_out}{subfolder}/parts_{tag_name}{num:05}.h5" # legacy only ) self.snap_filename = ( f"{self.path_out}{subfolder}/snap_{num:05}.h5" # new hdf5 snap format ) self.pspec_filename = f"{self.path_out}{subfolder}/pspec_{tag_name}{num:05}.h5" self.filaments_filename = ( f"{self.path_out}/{subfolder}filaments_{tag_name}{num:05}.h5" ) if not os.path.exists(f"{self.path_out}{subfolder}"): os.makedirs(f"{self.path_out}{subfolder}") # Create selector object if selector is None: dirpath = os.path.dirname(path) if dirpath == "": dirpath = "." selector = RunSelector( dirpath, self.run, self.num, self.params.input.nml_filename, ) self.info = selector.info[self.run][self.num] self.namelist = selector.namelist[self.run] # Save important info files if self.params.out.copy_info: info_src = f"{self.path}/output_{self.num:05}/info_{self.num:05}.txt" info_dest = f"{self.path_out}/info/info_{self.num:05}.txt" if os.path.exists(info_src): os.makedirs(os.path.dirname(info_dest), exist_ok=True) copy_file(info_src, info_dest, update=1) # Get box length self.lbox = self.info["boxlen"] # Get time self.time = self.info["time"] # Set post processing attributes self.open() 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 = self.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 = self.time self.close() if os.path.exists(self.filaments_filename): with open(self.filaments_filename, "rb") as f: self.fil = pickle.load(f) else: self.fil = None # Convert time unit if isinstance(unit_time, str): factor = self.get_nml(unit_time) unit_time = U.Unit( name=os.path.basename(unit_time), base_unit=self.info["unit_time"], coeff=factor, ) time_in_right_unit = self.time * self.info["unit_time"].express(unit_time) if self.params.astrophysix.generate: self.snapshot = Snapshot( name=str(self.num), description="", time=(time_in_right_unit, unit_time), directory_path=self.path, data_reference="OUTPUT_{}".format(self.num), ) try: self.init_pymses() except IOError: self.logger.error("Pymses not initialized", exc_info=1) self.def_rules() def check_variables(self, path): # Check if variables are in output name_conv = { "rho": "density", "vel": "velocity", "P": "pressure", "Br": "B", "Bl": "B", "mass": "mass", "id": "identity", "level": "levelp", "epoch": "epoch", } hydro_filename = f"{path}/output_{self.num:05}/hydro_file_descriptor.txt" if os.path.exists(hydro_filename): hydro_file = open(hydro_filename) hlines = hydro_file.readlines() hydro_var = np.unique( list(map(lambda s: s.split(",")[1][1:].split("_")[0], hlines[2:])) ) else: hydro_var = [] part_filename = f"{path}/output_{self.num:05}/part_file_descriptor.txt" if os.path.exists(part_filename): hydro_file = open(part_filename) part_file = open(f"{path}/output_{self.num:05}/part_file_descriptor.txt") plines = part_file.readlines() part_var = np.unique( list(map(lambda s: s.split(",")[1][1:].split("_")[0], plines[2:])) ) else: part_var = [] has_grav = os.path.exists( f"{path}/output_{self.num:05}/grav_{self.num:05}.out00001" ) def is_available(available_vars, pymsesrc, var): if var in ["g", "phi"]: if not has_grav: self.logger.debug(f"Variable {var} not in output") if var not in pymsesrc: self.logger.debug(f"Variable {var} not in pymsesrc") return has_grav and var in pymsesrc else: if var in name_conv: if name_conv[var] not in available_vars: self.logger.debug(f"Variable {var} not in output") if var not in pymsesrc: self.logger.debug(f"Variable {var} not in pymsesrc") return name_conv[var] in available_vars and var in pymsesrc else: self.logger.warning(f"Variable {var} is unknown") return False self.params.pymses.variables = list( filter( partial( is_available, hydro_var, pymses.rcConfig.Ramses.amr_fields.field_name_list, ), self.params.pymses.variables, ) ) self.params.pymses.part_variables = list( filter( partial( is_available, part_var, ["vel", "mass", "id", "level", "epoch"] ), self.params.pymses.part_variables, ) ) if ( "Br" in pymses.rcConfig.Ramses.amr_fields.field_name_list and "B" not in hydro_var ): pymses.rcConfig.Ramses.amr_fields.remove_field("P") pymses.rcConfig.Ramses.amr_fields.remove_field("Br") pymses.rcConfig.Ramses.amr_fields.remove_field("Bl") pymses.rcConfig.Ramses.amr_fields.add_scalar( name="P", ivar=4, file_prefix="hydro" ) if not has_grav and "g" in pymses.rcConfig.Ramses.amr_fields.field_name_list: pymses.rcConfig.Ramses.amr_fields.remove_field("g") pymses.rcConfig.Ramses.amr_fields.remove_field("phi") def init_pymses(self): # If ratarmount was used if os.path.exists(f"{self.path}/output_{self.num:05}/output_{self.num:05}"): path = f"{self.path}/output_{self.num:05}" else: path = self.path self._ro = pymses.RamsesOutput( path, self.num, order=self.params.pymses.order, verbose=self.params.pymses.verbose, check_endianness=False, ) if self.params.pymses.check_variables: self.check_variables(path) self._amr = self._ro.amr_source(self.params.pymses.variables) self._part = self._ro.particle_source(self.params.pymses.part_variables) if self.params.pymses.filter: self.min_coords = np.array(self.params.pymses.min_coords) self.max_coords = np.array(self.params.pymses.max_coords) box = reg.Box((self.min_coords, self.max_coords)) amr_filt = RegionFilter(box, self._amr) self._amr = amr_filt part_filt = RegionFilter(box, self._part) self._part = part_filt # Density operator self._rho_op = ScalarOperator( lambda dset: dset["rho"], self._ro.info["unit_density"] ) # Density ray tracer if self.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) if not self.params.pymses.multiprocessing: self._rt.disable_multiprocessing() # Set the extent of the image self._radius = 0.5 / self.params.pymses.zoom if self.params.pymses.filter: center = (self.max_coords + self.min_coords) / 2.0 im_extent = np.array( [ self.min_coords[0], self.max_coords[0], self.min_coords[1], self.max_coords[1], ] ) distance = (self.max_coords[2] - self.min_coords[2]) / 2.0 else: center = self.params.pymses.center im_extent = np.array( [ (-self._radius + center[0]), (self._radius + center[0]), (-self._radius + center[1]), (self._radius + center[1]), ] ) distance = self._radius # Initialize cameras self._cam = {} for ax_los in self._ax_nb: # los = line of sight ax_v = self._axes_v[ax_los] real_ax = ax_los if real_ax == "y": real_ax = [0, -1, 0] # we look from -y to keep z upwards self._cam[ax_los] = Camera( center=self.params.pymses.center, line_of_sight_axis=real_ax, region_size=[im_extent[1] - im_extent[0], im_extent[3] - im_extent[2]], distance=distance, far_cut_depth=distance, up_vector=ax_v, map_max_size=self.params.pymses.map_size, ) # Initialize HDF5 group try: self.open() if "/maps" not 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 except Exception() as e: self.logger.error("Error in HDF5", exc_info=1) raise e finally: self.close() def load_data(self, points_src, filename, save, keys=None, group="data"): """ Load data from the source file in the memory. (Long and memory heavy) """ loaded = False if os.path.exists(filename): self.logger.debug(f"Found hdf5, loading {filename}.") hdf5 = tables.open_file(filename, mode="r") if group in hdf5.root: node = hdf5.get_node(f"/{group}") loaded = True elif "data" in hdf5.root: self.logger.warning( f"{filename} has no {group} group, but I found the group data." ) node = hdf5.get_node("/data") loaded = True else: self.logger.warning(f"{filename} has no {group} group") if loaded: data = {} if keys is None: keys = node._v_children for key in keys: if key in node._v_children: data[key] = node[key].read() else: self.logger.warning( f"Key {key} is missing, I will try a full reload" ) loaded = False break hdf5.close() if not loaded: self.logger.debug("No hdf5, loading from ramses data.") data_pymses = points_src.flatten() data = {} for key in data_pymses.fields: data[key] = data_pymses[key] try: data["dx"] = data_pymses.get_sizes() except AttributeError: pass data["pos"] = data_pymses.points self.logger.info("Ramses data loaded.") if save: self.save_data(data, filename, group) return data def save_data(self, data, filename, group, overwrite=False, units=None): self.logger.debug(f"Writing {filename}") hdf5 = tables.open_file(filename, mode="a") units = units or self.unit_key try: nb_written = 0 for key in data: if len(data[key] > 0): if f"/{group}/{key}" in hdf5 and overwrite: hdf5.remove_node(f"/{group}/{key}") elif f"/{group}/{key}" not in hdf5: hdf5.create_array( f"/{group}", key, data[key], "", createparents=True ) unit = self._get_units(units[key]) hdf5.get_node(f"/{group}/{key}").unit = unit nb_written += 1 else: self.logger.warning("Empty key") if "namelist" not in hdf5.root._v_attrs: hdf5.root._v_attrs.namelist = self.namelist.data.todict() if "info " not in hdf5.root._v_attrs: hdf5.root._v_attrs.info = self.info finally: hdf5.close() self.logger.info( f"{filename} successfully written with {nb_written}/{len(data.keys())} updated fields" ) def load_parts(self, keys=None, force_new_filename=False, save=True): if not self.parts_loaded: self.logger.debug("Loading particles") if os.path.exists(self.parts_filename) and not force_new_filename: filename = self.parts_filename else: filename = self.snap_filename self.parts = self.load_data( self._part, filename, self.params.process.save_parts and save, keys=keys, group="parts", ) self.parts_loaded = True self.logger.info("Particles loaded") def unload_parts(self): """ Free space in the memory by telling the garbage collectors that self.parts is not needed """ if self.parts_loaded: del self.parts self.parts_loaded = False self.logger.info("Particles unloaded") def load_cells(self, keys=None, force_new_filename=False, save=True): """ 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: self.logger.debug("Loading cells") if os.path.exists(self.cells_filename) and not force_new_filename: filename = self.cells_filename else: filename = self.snap_filename cells_src = CellsToPoints(self._amr) self.cells = self.load_data( cells_src, filename, save, keys=keys, group="cells", ) self.cells_loaded = True self.logger.info("Cells loaded") def unload_cells(self): """ Free space in the memory by telling the garbage collector that self.cells is not needed """ if self.cells_loaded: del self.cells self.cells_loaded = False self.logger.info("Cells unloaded") def load_destructured(self, save=True): self.load_cells(save=save) self.load_parts(save=save) def unload_destructured(self): self.unload_cells() self.unload_parts() def convert_hdf5(self, filename=None): self.load_destructured(save=False) self.load_sinks() sinks = {key: self.sinks[key].values for key in self.sinks} if filename is None: filename = self.snap_filename self.save_data(self.cells, filename, "cells") self.save_data(self.parts, filename, "parts") self.save_data(sinks, filename, "sinks", units=self.unit_sinks) self.unload_destructured() def get_nml(self, nml_key): if self.namelist is not None: value = self.namelist[nml_key] else: raise AttributeError("No namelist associated with this snapshot") return value def getter_pos_disk(self, dset): """ Returns the position in normalized and centered units """ try: pos = dset.points except AttributeError: pos = dset["pos"] pos = pos - np.array(self.params.disk.center) return pos def getter_vect_r(self, dset, name_vect): """Radial component of a vector""" r = self.getter_pos_disk(dset)[:, :2] ur = np.transpose((np.transpose(r) / np.sqrt(np.sum(r**2, axis=1)))) return np.einsum("ij, ij -> i", dset[name_vect][:, :2], ur) def getter_vect_phi(self, dset, name_vect): """Azimuthal component of a vector""" r = self.getter_pos_disk(dset)[:, :2] r_norm = np.sqrt(np.sum(r**2, axis=1)) rot = np.array([[0, -1], [1, 0]]) uphi = np.transpose(np.einsum("ij, kj -> ik", rot, r) / r_norm) vect = dset[name_vect][:, :2] return np.einsum("ij,ij -> i", vect, uphi) def oct_getter_pos_disk(self, dset): """ Returns the position in normalized units centered on the position of the star """ pos = dset.get_cell_centers() pos = pos - np.array(self.params.disk.center) return pos def oct_getter_vect_r(self, dset, name_vect): """Radial component of a vector""" r = self.oct_getter_pos_disk(dset)[:, :, :2] ur = np.transpose( (np.transpose(r, (2, 0, 1)) / np.sqrt(np.sum(r**2, axis=2))), (1, 2, 0) ) return np.einsum("ikj, ikj -> ik", dset[name_vect][:, :, :2], ur) def oct_getter_vect_phi(self, dset, name_vect): """Azimuthal component of a vector""" r = self.oct_getter_pos_disk(dset)[:, :, :2] r_norm = np.sqrt(np.sum(r**2, axis=2)) rot = np.array([[0, -1], [1, 0]]) uphi = np.transpose(np.einsum("ij, klj -> ikl", rot, r) / r_norm, (1, 2, 0)) vect = dset[name_vect][:, :, :2] return np.einsum("ikj,ikj -> ik", vect, uphi) def oct_getter_vr(self, dset): return self.oct_getter_vect_r(dset, "vel") def oct_getter_vphi(self, dset): """Azimuthal velocity""" return self.oct_getter_vect_phi(dset, "vel") def datacube(self, getter, level=None, unit=U.none): """ Return a datacube of the source box. Parameters ---------- getter : callable A callable that extract the wanted data from a pymses dataset unit : U.Unit Unit of the resulting dataset Returns ------- A numpy array containing the slice """ unit = self._get_units(unit) operator = ScalarOperator(getter, unit) extractor = cube3d.CubeExtractor(self._amr, operator) if level is None: level = self.get_nml("amr_params/levelmin") size = 1.0 cam = Camera( center=self.params.pymses.center, line_of_sight_axis="z", region_size=[size, size], distance=size / 2.0, far_cut_depth=size / 2.0, up_vector="y", map_max_size=2**level, ) cube = extractor.process(cam, cube_size=1.0, resolution=2**level).data return cube def slice(self, getter, ax_los="z", z=0.0, unit=U.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 : U.Unit Unit of the resulting dataset Returns ------- A numpy array containing the slice """ unit = self._get_units(unit) op = ScalarOperator(getter, unit) datamap = slicing.SliceMap(self._amr, self._cam[ax_los], op, z=z) return datamap.map.T def ax_avg( self, oct_getter, ax_los, unit=U.none, mass_weighted=True, surf_qty=False ): """ Map of the average of a quantity (given by getter) along an axis (ax_los) Returns 2D array if getter returns a scalar quantity If surf_qty is set (projection mode), mass_weighted is ignored """ unit = self._get_units(unit) if surf_qty: op = ScalarOperator(oct_getter, unit) else: if mass_weighted: def num(dset): value = oct_getter(dset) rho = getter_rho(dset) return rho * value def denum(dset): return getter_rho(dset) else: # Volume weighted def num(dset): value = oct_getter(dset) return value def denum(dset): return 1.0 op = FractionOperator(num, denum, unit) if self.params.pymses.fft: rt = splatting.SplatterProcessor(self._amr, self._ro.info, op) else: rt = raytracing.RayTracer(self._amr, self._ro.info, op) if not self.params.pymses.multiprocessing: rt.disable_multiprocessing() datamap = rt.process(self._cam[ax_los], surf_qty=surf_qty) return datamap.map.T def _get_axis(self, axis): if isinstance(axis, str): axis = self._ax_nb[axis] self.load_cells() return np.sort(np.unique(self.cells["pos"][:, axis])) def _plane_avg_uniform(self, getter, axis, unit=U.none, mass_weighted=False): """ Profile of the average of a quantity (given by getter) perpendicular to an axis WARNING : This version only works on an uniform grid, need of a box version for AMR Returns 1D array if getter returns a scalar quantity """ unit = self._get_units(unit) self.load_cells() if isinstance(axis, str): axis = self._ax_nb[axis] axis_data = self.cells["pos"][:, axis] value = getter(self.cells) df = pd.DataFrame({"axis": axis_data}) if mass_weighted: mass = mass_func(self.cells) tot_mass = np.sum(mass) df["value"] = value * mass / tot_mass else: df["value"] = value if self.params.process.unload_cells: self.unload_cells() df.sort_values("axis", inplace=True) return df.groupby("axis").mean().values[:, 0] def sum(self, getter, mass_weighted=True): """ Global sum of the quantity returned by getter (variable must be extensive) Returns a scalar (or a vector if the quantity returned by getter is a getter, eg. speed) """ self.load_cells() value = getter(self.cells) if mass_weighted: mass = mass_func(self.cells) # Transpose (.T) is for vectorial values data = np.sum((mass * value.T).T, axis=0) / np.sum(mass) else: data = np.sum(value, axis=0) if self.params.process.unload_cells: self.unload_cells() return data def vol_avg(self, getter, mass_weighted=True): """ Global volumic (or mass_weighted) average of the quantity returned by getter Returns a scalar (or a vector if the quantity returned by getter is a getter, eg. speed) """ self.load_cells() value = getter(self.cells) if mass_weighted: weight = mass_func(self.cells) else: weight = vol_func(self.cells) # Transpose (.T) is for vectorial values data = np.sum((weight * value.T).T, axis=0) / np.sum(weight) if self.params.process.unload_cells: self.unload_cells() return data def vol_pdf( self, getter, bins=100, old_unit=None, unit=None, logbins=False, weight_func=vol_func, **kwargs, ): self.load_cells() data = getter(self.cells) if old_unit is not None and unit is not None: data *= old_unit.express(unit) if logbins: data = np.log10(data) weights = weight_func(self.cells) if self.params.process.unload_cells: self.unload_cells() values, edges = np.histogram( data, bins, weights=weights, density=True, **kwargs ) centers = 0.5 * (edges[1:] + edges[:-1]) return (np.stack([values, centers]), {"logbins": logbins}) def _Brho(self, bins=100, logbins=True): """ Mean of B in rho bins """ self.load_cells() B = getter_B_int(self.cells) rho = getter_rho(self.cells) centers, B_mean = mean_by_bins(rho, B, bins, logbins) if self.params.process.unload_cells: self.unload_cells() return ({"rho": centers, "B": B_mean}, {"logbins": logbins}) def mean_by_bins( self, name_x, name_y, ax_los, group="/maps/", bins=100, logbins=True ): """ Compute the mean of y in bins of x, where x, y are to array of the hdf5 file Parameters ---------- x_name, y_name : str, path of x and y in the hdf5 file bins : int, number of bins logbins : bool, if true, the bins will be logaritmically distributed """ x = self.get_value(group + name_x + "_" + ax_los) y = self.get_value(group + name_y + "_" + ax_los) centers, values = mean_by_bins(x, y, bins, logbins) # return ({os.path.basename(name_x) : centers, os.path.basename(name_y) : y_mean}, # {"logbins" : logbins}) return (np.stack([values, centers]), {"logbins": logbins}) def _Ek_Eb_rho(self, bins=100, logbins=True): """ Mean of Ek/Eb in rho bins """ self.load_cells() mean_speed = self.get_value("/globals/mwa_speed", unit=U.km_s) vel_fluct = (self.cells)["vel"] * self.info["unit_velocity"].express( U.km_s ) - mean_speed B_norm = getter_B_int(self.cells) B_norm = B_norm * self.info["unit_mag"].express(U.T) v_norm = np.sqrt( np.sum((vel_fluct * 10 ** (3)) ** 2, axis=1) ) # v_norm [m/s] et vel_fluct [km/s] rho = getter_rho(self.cells) rho_kg_m3 = rho * self.info["unit_density"].express(U.kg_m3) eb = 0.5 * (B_norm) ** 2 / (4 * np.pi * 10 ** (-7)) # mettre le bon mu ek = 0.5 * v_norm**2 * rho_kg_m3 rapport = ek / eb if logbins: rho_bins = np.logspace( np.log10(np.min(rho)), np.log10(np.max(rho)), bins, base=10 ) else: rho_bins = np.linspace(np.min(rho), np.max(rho), bins) # For each cell, bin_number contains the number of the bins it belongs to bin_number = np.zeros(len(B_norm)) # Go through the min value of rho of each bin for rho_min in rho_bins[:-1]: bin_number = bin_number + (rho > rho_min).astype(int) # Compute the mean in each bin ek_eb = np.zeros(len(rho_bins) - 1) for i in range(len(ek_eb)): ek_eb[i] = np.mean(rapport[bin_number == i]) # Get the center of each bin if logbins: centers = 10 ** (0.5 * (np.log10(rho_bins[1:]) + np.log10(rho_bins[:-1]))) else: centers = 0.5 * (rho_bins[1:] + rho_bins[:-1]) if self.params.process.unload_cells: self.unload_cells() return ({"rho": centers, "Ek_Eb_rho": ek_eb}, {"logbins": logbins}) def cos_vfluct_B(self): "return the cos of the angle between the magnetic field and the velocity fluctuation field" mean_speed = self.get_value("/globals/mwa_speed") def getter_cos_vfluct_B(dset): vel_fluct = dset["vel"] - mean_speed B_norm = np.sqrt(np.sum(dset["Br"] ** 2, axis=1)) v_norm = np.sqrt(np.sum(vel_fluct**2, axis=1)) # Compute the dot product in each cell dot_prod = np.einsum("ij,ij->i", vel_fluct, dset["Br"]) return np.abs(dot_prod) / (v_norm * B_norm) return self.vol_pdf(getter_cos_vfluct_B) def _mwa_sigma(self, axes=["x", "y", "z"]): mw_speed = self.get_value("/globals/mwa_speed") 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 _B_int(self, ax_los, z=0.0): """ Slice ont the intensity of the magnetic field """ B_op = ScalarOperator( lambda dset: np.sqrt(np.sum(dset["Br"] ** 2, axis=1)), self._ro.info["unit_mag"], ) dmap_B = (slicing.SliceMap(self._amr, self._cam[ax_los], B_op, z=z)).map.T return dmap_B 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.get_value("/maps/slice_rho_{}".format(ax_los)) return dmap_P / dmap_rho def _levels(self, ax_los): self._amr.set_read_levelmax(self.params.pymses.levelmax) level_op = MaxLevelOperator() rt_level = raytracing.RayTracer(self._amr, self._ro.info, level_op) if not self.params.pymses.multiprocessing: rt_level.disable_multiprocessing() datamap = rt_level.process(self._cam[ax_los], surf_qty=True) return datamap.map.T def _jeans(self, ax_los): dmap_T = self.get_value("/maps/T_" + ax_los) dmap_rho = self.get_value("/maps/slice_rho_" + ax_los) dmap_jeans = np.sqrt(np.pi * dmap_T / dmap_rho) return dmap_jeans def _jeans_ratio(self, ax_los): dmap_jeans = self.get_value("/maps/jeans_" + ax_los) dmap_levels = self.get_value("/maps/levels_" + ax_los) dmap_jeans_ratio = dmap_jeans * 2 ** (dmap_levels) return dmap_jeans_ratio def _omega_average(self, ax_los): # Operator to compute the angular speed times rho lbox = self.lbox def omega_rho_func(dset): pos = self.oct_getter_pos_disk(dset) xx = pos[:, :, 0] * lbox yy = pos[:, :, 1] * lbox rc2 = xx**2 + yy**2 # square of cylindrical radius vx = dset["vel"][:, :, 0] vy = dset["vel"][:, :, 1] omega_rho = dset["rho"] vyx = vy * xx vxy = vx * yy omega_rho = omega_rho * (vyx - vxy) / rc2 return omega_rho # Operator to compute the angular speed omega_unit = self._ro.info["unit_velocity"] / ( self._ro.info["unit_length"] / lbox ) omega_op = FractionOperator( omega_rho_func, lambda dset: dset["rho"], omega_unit, ) # Ray tracer for the angular speed rt_omega = raytracing.RayTracer(self._amr, self._ro.info, omega_op) if not self.params.pymses.multiprocessing: rt_omega.disable_multiprocessing() dmap_omega = rt_omega.process(self._cam[ax_los]).map.T return dmap_omega def _toomreQ_disk( self, ax_los, omega_approx=False, G1_units=False, coarsen_factor=1 ): """ Compute the Toomre Q parameter """ # Get maps if G1_units: G = 1.0 cs2 = self.get_value("/maps/T_mwavg_z") coldens = self.get_value("/maps/coldens_z") omega = self.get_value("/maps/omega_mwavg_z") space_coeff = self.lbox else: G = U.G.express(U.kg**-1 * U.m**3 * U.s**-2) cs2 = self.get_value("/maps/T_mwavg_z", unit=U.m**2 * U.s**-2) coldens = self.get_value("/maps/coldens_z", unit=U.kg * U.m**-2) omega = self.get_value("/maps/omega_mwavg_z", unit=U.s**-1) space_coeff = self.info["unit_length"].express(U.m) map_size = self.params.pymses.map_size if coarsen_factor > 1: map_size = map_size // coarsen_factor omega = degrade_map(omega, map_size) coldens = degrade_map(coldens, map_size) cs2 = degrade_map(cs2, map_size) # Compute Q if omega_approx: # Use angular frequency as epiciclic frequency (true if the disk is Keplerian) kappa = omega else: # Get coordinates im_extent = np.array(self.save.root.maps._v_attrs.im_extent) * space_coeff center = np.array(self.params.disk.center) * space_coeff # Physical size of cells dx = (im_extent[1] - im_extent[0]) / map_size dy = (im_extent[3] - im_extent[2]) / map_size # Physical coordinates of the center of the cells x = np.linspace(im_extent[0], im_extent[1], map_size) + 0.5 * dx - center[0] y = np.linspace(im_extent[2], im_extent[3], map_size) + 0.5 * dy - center[1] xx, yy = np.meshgrid(x, y) rr = np.sqrt(xx**2 + yy**2) # Compute kappa**2 = (2omega/R)*d(R**2*omega)/dR Gy, Gx = np.gradient(rr**2 * omega, dy, dx, edge_order=2) Gr = (xx * Gx + yy * Gy) / rr kappa_square = (2 * omega / rr) * Gr kappa = np.sqrt(kappa_square) map_Q = (np.sqrt(cs2) * kappa) / (np.pi * G * coldens) return map_Q def _radial_bins(self, ax_los="z"): """ Computes radial bins (for disk) """ center = np.array(self.params.disk.center) * self.lbox im_extent = np.array(self.save.root.maps._v_attrs.im_extent) * self.lbox # radius of the corner of the box plus a margin rad_of_box = ( np.sqrt((im_extent[1] - center[0]) ** 2 + (im_extent[3] - center[1]) ** 2) + 0.1 ) bin_in = self.params.disk.bin_in bin_out = self.params.disk.bin_out nb_bin = self.params.disk.nb_bin # radial bins if self.params.disk.binning in ["log", "logarithmic"]: 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.params.disk.binning in ["lin", "linear"]: rad_bins = np.linspace(bin_in, bin_out, num=nb_bin) else: raise RuntimeWarning( f"Invalid parameter {self.params.disk.binning} for disk binning method, using linear binning" ) 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 _radial_centers(self, ax_los="z"): radial_bins = self.get_value("/radial/radial_bins_" + ax_los) bin_centers = 0.5 * (radial_bins[1:] + radial_bins[:-1]) return bin_centers def _rr(self, ax_los="z"): """ Computes the radius from the center """ im_extent = np.array(self.save.root.maps._v_attrs.im_extent) * self.lbox map_size = self.params.pymses.map_size center = np.array(self.params.disk.center) * self.lbox # Physical size of cells dx = (im_extent[1] - im_extent[0]) / map_size dy = (im_extent[3] - im_extent[2]) / map_size # Physical coordinates of the center of the cells x = np.linspace(im_extent[0], im_extent[1], map_size) + 0.5 * dx y = np.linspace(im_extent[2], im_extent[3], map_size) + 0.5 * dy xx, yy = np.meshgrid(x, y) # Physical radius rr = np.sqrt((xx - center[0]) ** 2 + (yy - center[1]) ** 2) return rr def _bins_on_map(self, ax_los="z"): rad_bins = self.get_value("/radial/radial_bins_" + ax_los) rr = self.get_value("/maps/rr_" + ax_los) # 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="z", mass_weighted=False): radial_bins = self.get_value("/radial/radial_bins_" + ax_los) bins_on_map = self.get_value("/maps/bins_on_map_" + ax_los) dmap = self.get_value("/maps/" + name + "_" + ax_los) if mass_weighted: coldens = self.get_value("/maps/coldens_" + ax_los) # mean of all the cells in the bin mean_bin = np.zeros(len(radial_bins) - 1) for j in range(len(radial_bins) - 1): mask_bin = bins_on_map == j if mass_weighted: weight = coldens[mask_bin] mean_bin[j] = np.nanmean(dmap[mask_bin] * weight) mean_bin[j] = mean_bin[j] / np.mean(weight) else: mean_bin[j] = np.nanmean(dmap[mask_bin]) return mean_bin def _rad_avg_map(self, name, ax_los="z", mass_weighted=False): radial_bins = self.get_value("/radial/radial_bins_" + ax_los) bins_on_map = self.get_value("/maps/bins_on_map_" + ax_los) rr = self.get_value("/maps/rr_" + ax_los) if mass_weighted: mean_bin = self.get_value("/radial/rad_mwavg_" + name + "_" + ax_los) else: mean_bin = self.get_value("/radial/rad_avg_" + name + "_" + ax_los) # Add value for border mean_bin = np.concatenate((mean_bin, [mean_bin[-1]])) 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="z", mass_weighted=False): dmap = self.get_value("/maps/" + name + "_" + ax_los) if mass_weighted: avg_map = self.get_value("/maps/rad_mwavg_map_" + name + "_" + ax_los) else: avg_map = self.get_value("/maps/rad_avg_map_" + name + "_" + ax_los) return dmap / avg_map def _dispersion_rad(self, name, ax_los="z"): """ Computes the dispersion in radial bins of the quantity `name` Parameters ---------- name : name of 2D map of a scalar quantity ax_los : axis of the line of sight """ # 1. Get full storage names for the map and its azimuthal avererage map_name = f"/maps/{name}_{ax_los}" avg_map_name = f"/maps/rad_avg_map_{name}_{ax_los}" # 2. Get maps from the storage dmap = self.get_value(map_name) avgmap = self.get_value(avg_map_name) # 3. Get radial bins and centers bins_on_map = self.get_value(f"/maps/bins_on_map_{ax_los}") centers = self.get_value(f"/radial/radial_centers_{ax_los}") # 4. Initialize dispersion array sigma = np.zeros(len(centers)) # 5. Compute RMS in each bin for j in range(len(centers)): mask_bin = bins_on_map == j sigma[j] = np.sqrt( np.sum((dmap[mask_bin] - avgmap[mask_bin]) ** 2) / np.sum(mask_bin) ) # 6. Returns computed RMS return sigma def _rad_fluct_pdf(self, name, ax_los="z", mass_weighted=False): if mass_weighted: fluct_map = self.get_value("/maps/mwfluct_" + name + "_" + ax_los) else: fluct_map = self.get_value("/maps/fluct_" + name + "_" + ax_los) rr = self.get_value("/maps/rr_" + ax_los) mask_pdf = ( (rr > self.params.disk.rmin_pdf) & (rr < self.params.disk.rmax_pdf) & (fluct_map > 0) ) values, edges = np.histogram( np.log10(fluct_map[mask_pdf].flatten()), self.params.pdf.nb_bin, range=self.params.pdf.range, density=True, ) centers = 0.5 * (edges[1:] + edges[:-1]) return np.stack([values, centers]) def _fit_pdf(self, name, ax_los="z"): pdf = self.get_value("/hist/fluct_pdf_" + name + "_" + ax_los) values, centers = pdf mask_fit = ( (centers > self.params.pdf.xmin_fit) & (centers < self.params.pdf.xmax_fit) & (values > self.params.pdf.fit_cut * np.max(values)) ) (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 _alpha_disk(self, ax_los="z"): "Map of the Rayleigh contribution to the Shakura&Sunaev alpha parameter for disks" assert ax_los == "z" # Mean part T_avg = self.get_value("/maps/rad_avg_map_T_mwavg_z") radial_bins = self.get_value("/radial/radial_bins_" + ax_los) mean_bin_vr = self.get_value("/radial/rad_avg_" + "slice_velr" + "_" + ax_los) mean_bin_vphi = self.get_value( "/radial/rad_avg_" + "slice_velphi" + "_" + ax_los ) mean_bin_vr = np.concatenate((mean_bin_vr, [mean_bin_vr[-1]])) mean_bin_vphi = np.concatenate((mean_bin_vphi, [mean_bin_vr[-1]])) # Fluct part def getter_alpha_num(dset): r = np.sqrt( np.sum((self.lbox * self.oct_getter_pos_disk(dset)) ** 2, axis=2) ) bins = np.zeros(r.shape, dtype=int) for r0 in radial_bins[1:-1]: bins = bins + (r >= r0).astype(int) vr_mean = mean_bin_vr[bins] vphi_mean = mean_bin_vphi[bins] # use linear interpolation # v = ((r[i+1] - r)v[i] + (r - r[i])v[i + 1]) / (r[i+1] - r[i]) vr_mean = (radial_bins[bins + 1] - r) * vr_mean vr_mean = vr_mean + (r - radial_bins[bins]) * mean_bin_vr[bins + 1] vr_mean = vr_mean / (radial_bins[bins + 1] - radial_bins[bins]) vphi_mean = (radial_bins[bins + 1] - r) * vphi_mean vphi_mean = vphi_mean + (r - radial_bins[bins]) * mean_bin_vphi[bins + 1] vphi_mean = vphi_mean / (radial_bins[bins + 1] - radial_bins[bins]) vr = self.oct_getter_vr(dset) vphi = self.oct_getter_vphi(dset) alpha = (vphi - vphi_mean) * (vr - vr_mean) return alpha alpha_f = ( self.ax_avg(getter_alpha_num, "z", unit=U.none, mass_weighted=True) / T_avg ) # alpha alpha = (2.0 / 3) * alpha_f return alpha def _alpha_grav(self, ax_los="z"): "Map of the gravitational contribution to the Shakura&Sunaev alpha parameter for disks" assert ax_los == "z" T_avg = self.get_value("/maps/rad_avg_map_T_mwavg_z") coldens = self.get_value("/maps/rad_avg_map_coldens_z") def getter_alpha_grav(dset): r2 = np.sum((self.lbox * self.oct_getter_pos_disk(dset)) ** 2, axis=2) e2 = (1.0 / 256.0) ** 2 gstar = -self.G * self.params.disk.mass_star / (e2 + r2) gr = self.oct_getter_vect_r(dset, "g") - gstar gphi = self.oct_getter_vect_phi(dset, "g") return gr * gphi / (4 * np.pi * self.G) alpha_g = self.ax_avg(getter_alpha_grav, "z", unit=U.none, surf_qty=True) / ( coldens * T_avg ) # alpha alpha_g = (2.0 / 3) * alpha_g return alpha_g def load_sinks(self): csv_name = f"{self.path}/output_{self.num:05}/sink_{self.num:05}.csv" if not os.path.exists(csv_name): # If ratarmount was used csv_name = f"{self.path}/output_{self.num:05}/output_{self.num:05}/sink_{self.num:05}.csv" if not os.path.exists(csv_name): self.sinks = {} return {} f = open(csv_name) first_line = f.readline() second_line = f.readline() f.close() if len(first_line) > 0 and first_line[1] == "#": header = first_line[3:-2].split(",") units = second_line[3:-2].split(",") df = pd.read_csv(csv_name, header=None, names=header, skiprows=2) for key in header: if units[header.index(key)] == "m": df[key] = df[key] * self.info["unit_mass"].express(U.Msun) if "M" in df: df["msink"] = df["M"] else: header = [ "id", "msink", "dmfsink", "x", "y", "z", "vx", "vy", "vz", "rot_period", "lx", "ly", "lz", "acc_rate", "acc_lum", "age", "int_lum", "Teff", ] df = pd.read_csv(csv_name, header=None, names=header) self.sinks = df return {key: df[key].values for key in df} def pspec(self, overwrite_file=False, **kwargs): outfile = self.pspec_filename if overwrite_file or not os.path.exists(self.pspec_filename): pspec.pspec(repo=self.path, iouts=[self.num], outfile=outfile, **kwargs) def _write_particles(self): """Ensure particles are written in the hdf5 file""" if not os.path.exists(self.parts_filename) and not self.parts_loaded: self.load_parts() self.unload_parts() return np.array([self.parts_filename]) def _write_cells(self): """Ensure cells are written in the hdf5 file""" if not os.path.exists(self.cells_filename) and not self.cells_loaded: self.load_cells() self.unload_cells() return np.array([self.cells_filename]) def _filaments(self): datamap_name = self.params.filaments.datamap verbose = self.params.filaments.verbose rmin_frac = self.params.filaments.rmin rmax_frac = self.params.filaments.rmax size_thresh = self.params.filaments.size_thresh skel_thresh = self.params.filaments.skel_thresh branch_thresh = self.params.filaments.branch_thresh glob_thresh = self.params.filaments.glob_thresh datamap = self.get_value("/maps/" + datamap_name + "_z") shape = datamap.shape x = np.arange(shape[0]) - shape[0] / 2 y = np.arange(shape[1]) - shape[1] / 2 xx, yy = np.meshgrid(x, y) rr = np.sqrt(xx**2 + yy**2) rmin = int(rmin_frac * shape[0]) rmax = int(rmax_frac * shape[0]) mask = (rr >= rmin) & (rr <= rmax) datamap[np.logical_not(mask)] = np.nan self.fil = FilFinder2D(datamap, distance=1 * u.cm, beamwidth=1 * u.pix) self.fil.preprocess_image(flatten_percent=95) self.fil.create_mask( verbose=verbose, smooth_size=1 * u.pix, adapt_thresh=4 * u.pix, size_thresh=size_thresh * u.pix**2, glob_thresh=glob_thresh, fill_hole_size=0.1 * u.pix**2, ) self.fil.medskel(verbose=verbose) self.fil.analyze_skeletons( skel_thresh=skel_thresh * u.pix, branch_thresh=branch_thresh * u.pix, relintens_thresh=0.1, ) self.fil.exec_rht() self.fil.find_widths() with open(self.filaments_filename, "wb") as f: pickle.dump(self.fil, f, pickle.HIGHEST_PROTOCOL) return True def _filaments_center(self): """ Fill an array with center postion for each cell in a filament """ fil = self.fil mask = fil.mask.copy() _, distance = medial_axis(mask, return_distance=True) skel = fil.skeleton i_center = np.zeros(distance.shape, dtype=int) j_center = np.zeros(distance.shape, dtype=int) x_mask, y_mask = np.where(mask) for k in range(len(x_mask)): find_center(distance, skel, i_center, j_center, x_mask[k], y_mask[k]) return np.stack([i_center, j_center]) def _filaments_forces(self): """ Compute forces within a filament (for disks), within the slice at z=0 """ GM = self.G * self.params.disk.mass_star # Mass parameter # Find center of filaments i_center, j_center = self._filaments_center() # Get slices and projections at z = 0 vphi = self.get_value("/maps/slice_velphi_z") gr = self.get_value("/maps/slice_gr_z") Pz = self.get_value("/maps/slice_P_z") rho = self.get_value("/maps/slice_rho_z") vr = self.get_value("/maps/slice_velr_z") # Get coordinates im_extent = np.array(self.save.root.maps._v_attrs.im_extent) * self.lbox map_size = self.params.pymses.map_size center = np.array(self.params.disk.center) * self.lbox # Physical size of cells dx = (im_extent[1] - im_extent[0]) / map_size dy = (im_extent[3] - im_extent[2]) / map_size # Physical coordinates of the center of the cells x = np.linspace(im_extent[0], im_extent[1], map_size) + 0.5 * dx - center[0] y = np.linspace(im_extent[2], im_extent[3], map_size) + 0.5 * dy - center[1] xx, yy = np.meshgrid(x, y) rr = np.sqrt(xx**2 + yy**2) # Rotational support R = vphi**2 / rr # Equilibrium Gvry, Gvrx = np.gradient(vr, dy, dx, edge_order=2) gradvr = (xx * Gvrx + yy * Gvry) / rr dvr = vr * gradvr # Complete derivative # Thermal support GPy, GPx = np.gradient(Pz, dy, dx, edge_order=2) gradPr = (xx * GPx + yy * GPy) / rr fP = -gradPr / rho # Gravitational field e2 = (1.0 / 512) ** 2 gstar = -GM / (rr**2 + e2) # Substract gravitational field from the star Rdisk = R + gstar gdisk = gr - gstar # Forces at the center of filaments Rdisk_center = Rdisk[i_center, j_center] gr_center = gdisk[i_center, j_center] dvr_center = dvr[i_center, j_center] # Forces for the filaments equilibrium Rfil = Rdisk - Rdisk_center gfil = gdisk - gr_center fPfil = fP dvr_fil = dvr - dvr_center return {"gfil": gfil, "Rfil": Rfil, "fPfil": fPfil, "dvr": dvr_fil} def make_clump_hop(self, threshold_density=1e2): """ Apply HOP algorithm Args: threshold_density (float): select only cells over threshold """ # Selection of cells mask = ( self.cells["rho"] * self.info["unit_density"].express(U.H_cc) > threshold_density ) ncells = np.sum(mask) # fill the matrice with ID, x,y,z and masses of particles cells_group = np.zeros((ncells, 6)) cells_group[:, 0] = np.arange(ncells) # index position = self.cells["pos"][mask] * self.info["unit_length"].express(U.pc) cells_group[:, 1:4] = position # position density = self.cells["rho"][mask] * self.info["unit_density"].express( U.Msun / U.pc**3 ) size = self.cells["dx"][mask] * self.info["unit_length"].express(U.pc) cells_group[:, 4] = density * size**3 # mass # save file.txt head = str(ncells) np.savetxt( self.filename[:-3] + "_hop.txt", cells_group[:, :-1], fmt="%10d %.10e %.10e %.10e %.10e", header=head, delimiter=" ", comments=" ", ) # save file.den f = open(self.filename[:-3] + "_hop.den", "wb") f.write(pack("I", ncells)) self.cells["rho"][mask].astype("f").tofile(f) f.close() # exec HOP algo h = HOP(self.filename[:-3] + "_hop.txt", os.path.dirname(self.filename)) h.process_hop() # get the igroup array group_ids = h.get_group_ids() # sort it and apply the sorting to the coordinates # this means that the particules of group 1 are written first then of group 2 etc... ind_sort = np.argsort(group_ids) cells_group = cells_group[ind_sort] cells_group[6] = group_ids[ind_sort] # Make it a pandas' DataFrame cells_group = pd.DataFrame( cells_group, header=["id", "x", "y", "z", "mass", "group"] ) self.clumps = cells_group return cells_group def clump_properties(self, nograv=False, log_min=None, log_max=None): cells_group = self.make_clump_hop() cells_group.groupby("group") def def_rules(self): self.rules = { # Base rules "coldens": Rule( self._coldens, "Column density", "/maps", unit=self.info["unit_density"] * self.info["unit_length"], ), "T_mwavg": Rule( partial( self.ax_avg, getter_T, mass_weighted=True, unit=self.info["unit_pressure"] / self.info["unit_density"], ), "Ax mass-weighted averaged azimuthal temperature", "/maps", unit=self.info["unit_pressure"] / self.info["unit_density"], ), "omega_mwavg": Rule( partial(self._omega_average), "Ax mass-weighted averaged azimuthal of angular frequency", "/maps", unit=self.info["unit_velocity"] / (self.info["unit_length"] / self.lbox), ), "alpha_disk": Rule( self._alpha_disk, "Map of the Shakura&Sunaev alpha parameter for disks", "/maps", unit=U.none, dependencies=[ "avg_map_slice_rho_avg", "avg_map_T_mwavg", "avg_map_slice_velr", "avg_map_slice_velphi", ], ), "alpha_grav": Rule( self._alpha_grav, "Map of the graviational contrib to\ Shakura&Sunaev alpha parameter for disks", "/maps", unit=U.none, dependencies=["avg_map_coldens", "avg_map_T_mwavg"], ), "T": Rule( self._temperature, "Temperature slice", "/maps", dependencies=["slice_rho"], unit=self.info["unit_pressure"] / self.info["unit_density"], ), "levels": Rule(self._levels, "Max level within line of sight", "/maps"), "jeans": Rule( self._jeans, "Jeans length slice", "/maps", dependencies=["slice_rho", "T"], ), "jeans_ratio": Rule( self._jeans_ratio, "Jeans' length divided by the max resolution", "/maps", dependencies=["jeans", "levels"], ), "Q": Rule( self._toomreQ_disk, "Toomre Q parameter for a Keplerian disk", "/maps", dependencies=["coldens", "T_mwavg", "omega_mwavg"], ), "load_sinks_rule": Rule( self.load_sinks, group="/datasets", unit=self.unit_sinks, ), "pspec": Rule(self.pspec, "Power spectrum", "/hdf5"), "write_particles": Rule(self._write_particles, "Particles file", "/hdf5"), "write_cells": Rule(self._write_cells, "Cells file", "/hdf5"), "filaments": Rule( self._filaments, "Filaments", "/datasets", dependencies={self.params.filaments.datamap: "z"}, ), "filaments_forces": Rule( self._filaments_forces, "Filaments", "/datasets", dependencies={ "filaments": None, "slice_velphi": "z", "slice_gr": "z", "slice_P": "z", "slice_rho": "z", "slice_velr": "z", }, ), # Helpers "radial_bins": Rule( self._radial_bins, "Radial bins", "/radial", unit=self.info["unit_length"] / self.lbox, ), "radial_centers": Rule( self._radial_centers, "Centers of radial bins", "/radial", dependencies=["radial_bins"], unit=self.info["unit_length"] / self.lbox, ), "rr": Rule( self._rr, "Coordinate map", "/maps", unit=self.info["unit_length"] / self.lbox, ), "bins_on_map": Rule( self._bins_on_map, "Convert map coordinates to bins", "/maps", dependencies=["radial_bins", "rr"], ), "B_int": Rule( self._B_int, "Magnetic intensity slice", "/maps", dependencies=["slice_rho"], unit=self.info["unit_mag"], ), # PDF "rho_pdf": Rule( partial(self.vol_pdf, partial(simple_getter, "rho"), logbins=True), "Global rho-PDF", "/hist", unit=self.info["unit_density"], ), "rho_pdf_mw": Rule( partial( self.vol_pdf, partial(simple_getter, "rho"), weight_func=mass_func, logbins=True, ), "Global rho-PDF", "/hist", unit=self.info["unit_density"], ), "T_pdf": Rule( partial(self.vol_pdf, getter_T, logbins=True), "Global T-PDF", "/hist", unit=self.info["unit_pressure"] / self.info["unit_density"], ), "cos_pdf": Rule( partial(self.cos_vfluct_B), "Global cos fluctuation-PDF", "/hist", dependencies=["mwa_speed"], unit=U.none, ), "Brho": Rule( self._Brho, "Average of B as a function of rho", "/datasets", unit={"rho": self.info["unit_density"], "B": self.info["unit_mag"]}, ), "Ek_Eb_rho": Rule( self._Ek_Eb_rho, "Average of Ek/Eb as a function of rho", "/datasets", dependencies=["mwa_speed"], unit={"rho": self.info["unit_density"], "Ek_Eb_rho": U.none}, ), # Profiles "axis": Rule( partial(self._get_axis), "Axis", "/profile", unit=self.info["unit_length"], ), "rho_prof": Rule( partial(self._plane_avg_uniform, partial(simple_getter, "rho")), "Rho profile", "/profile", unit=self.info["unit_density"], dependencies=["axis"], ), # globals "time_num": Rule( lambda: self.info["time"], "Time", "/globals", unit=self.info["unit_time"], ), "mass": Rule( partial(self.sum, mass_func, mass_weighted=False), "Total mass", "/globals", unit=self.info["unit_density"] * self.info["unit_length"] ** 3, ), "mwa_speed": Rule( partial(self.vol_avg, partial(simple_getter, "vel")), "Mass weighted speed average", "/globals", unit=self.info["unit_velocity"], ), "mwa_sigma": Rule( self._mwa_sigma, "Mass weighted speed average", "/globals", dependencies={"mwa_speed": None}, unit=self.info["unit_velocity"], ), "mwa_B_int": Rule( partial(self.vol_avg, getter_B_int), "Mass weighted Magnetic intensity average", "/globals", unit=self.info["unit_mag"], ), } averageable_maps = [ "coldens", "Q", "T", "T_mwavg", "alpha_disk", "alpha_grav", ] # Generic rules directly from Ramses fields for field in self.params.pymses.variables: def generic_rule(name, getter, unit, oct_getter=None): if oct_getter is None: oct_getter = getter self.rules["slice_" + name] = Rule( partial(self.slice, getter, z=0.0, unit=unit), "{} slice".format(name), "/maps", unit=unit, ) self.rules[name + "_mwavg"] = Rule( partial(self.ax_avg, oct_getter, mass_weighted=True, unit=unit), "Ax mass-weighted averaged {}".format(name), "/maps", unit=unit, ) self.rules[name + "_avg"] = Rule( partial(self.ax_avg, oct_getter, mass_weighted=False, unit=unit), "Ax averaged {}".format(name), "/maps", unit=unit, ) averageable_maps.append("slice_" + name) averageable_maps.append(name + "_mwavg") averageable_maps.append(name + "_avg") # special for vectors if field in ["g", "vel", "Bl", "Br"]: # Components for i, dir in enumerate(["x", "y", "z"]): generic_rule( field + dir, partial(vect_getter, field, i), self.unit_key[field], oct_getter=partial(oct_vect_getter, field, i), ) # Radial generic_rule( field + "r", partial(self.getter_vect_r, name_vect=field), self.unit_key[field], oct_getter=partial(self.oct_getter_vect_r, name_vect=field), ) # Othoradial generic_rule( field + "phi", partial(self.getter_vect_phi, name_vect=field), self.unit_key[field], oct_getter=partial(self.oct_getter_vect_phi, name_vect=field), ) # Norm generic_rule( field + "_norm", partial(norm_getter, field), self.unit_key[field], oct_getter=partial(oct_norm_getter, field), ) else: generic_rule(field, partial(simple_getter, field), self.unit_key[field]) # radial average and other for name in averageable_maps: unit = self.rules[name].unit self.rules["rad_avg_" + name] = Rule( partial(self._rad_avg, name), "Azimuthal average of {}".format(name), "/radial", dependencies=["radial_centers", "bins_on_map", name], unit=unit, ) self.rules["rad_mwavg_" + name] = Rule( partial(self._rad_avg, name, mass_weighted=True), "Mass weighted azimuthal average of {}".format(name), "/radial", dependencies=["coldens", "radial_centers", "bins_on_map", name], unit=unit, ) self.rules["rad_avg_map_" + name] = Rule( partial(self._rad_avg_map, name), "Interpolated map of azimuthal average of {}".format(name), "/maps", dependencies=["radial_centers", "bins_on_map", "rad_avg_" + name], unit=unit, ) self.rules["rad_mwavg_map_" + name] = Rule( partial(self._rad_avg_map, name, mass_weighted=True), "Interpolated map of azimuthal average of {}".format(name), "/maps", dependencies=["radial_centers", "bins_on_map", "rad_mwavg_" + name], unit=unit, ) self.rules["fluct_" + name] = Rule( partial(self._fluct_map, name), "Fluctuation wrt to average of {}".format(name), "/maps", dependencies=[name, "rad_avg_map_" + name], ) self.rules["dispersion_rad_" + name] = Rule( partial(self._dispersion_rad, name), "radial RMS of {}".format(name), "/radial", dependencies=[name, "rad_avg_map_" + name], unit=unit, ) self.rules["mwfluct_" + name] = Rule( partial(self._fluct_map, name, mass_weigthed=True), "Fluctuation wrt to mass weighted average of {}".format(name), "/maps", dependencies=[name, "rad_mwavg_map_" + name], ) self.rules["fluct_pdf_" + name] = Rule( partial(self._rad_fluct_pdf, name), "Probability density function of {} fluctuations".format(name), "/hist", dependencies=["rr", "fluct_" + name], ) self.rules["fluct_mwpdf_" + name] = Rule( partial(self._rad_fluct_pdf, name, mass_weigthed=True), "Probability density function of {} mass weighted fluctuations".format( name ), "/hist", dependencies=["rr", "mwfluct_" + name], ) self.rules["fit_fluct_pdf_" + name] = Rule( partial(self._fit_pdf, name), "Fit the PDF of {} fluctuations".format(name), "/hist", dependencies=["fluct_pdf_" + name], ) for name_bin in averageable_maps: unit_bin = self.rules[name_bin].unit if name_bin is not name: self.rules["mbb_" + name + "_" + name_bin] = Rule( partial(self.mean_by_bins, name_bin, name), "Mean of {} by bins of {}".format(name, name_bin), "/hist", dependencies=[name, name_bin], unit=[unit, unit_bin], ) self._gen_rule_transform("fluct_coldens", np.nanmax, "max", group="/globals") super(SnapshotProcessor, 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