# coding: utf-8 """ This software is a helper to use pysmes tools to read and analyse RAMSES Outputs. It's a rule based interface. This is the plotter module. @author Noé Brucy 2019-2020 """ import sys import os from functools import partial import tables import numpy as np from scipy.stats import linregress from numpy.polynomial.polynomial import polyfit from scipy.ndimage.filters import gaussian_filter1d from scipy import optimize from astrophysix.simdm.datafiles import Datafile, PlotType, PlotInfo from astrophysix.utils.file import FileType import matplotlib as mpl if os.environ.get("DISPLAY", "") == "": print("No display found. Using non-interactive Agg backend") mpl.use("Agg") import pylab as P from comparator import * import pspec_read import datetime filetype_from_ext = {ext: ft for ft in FileType for ext in ft.extension_list} def not_array_error(err): epy2 = "object does not support indexing" epy3 = "object is not subscriptable" return str(err)[-len(epy2) :] == epy2 or str(err)[-len(epy3) :] == epy3 class PlotRule(Rule): """ The rule class, speficic to plot. Add an extra method, plot, that take the reference to an open hdf5 file (from pytables) """ def plot(self, save, arg, **kwargs): """ Set the plotter's storage to 'save' and execute the rule Parameters ---------- save : opended pytables hdf5 file, where to find the data arg : main argument of the plotting function kargs : optional keyword arguments to the plotting function """ self.postproc.save = save return self.process_fn(arg, **kwargs) def datafile(self, arg): return Datafile( name=self.description + "_" + arg, description=self.description + " ({})".format(arg), ) class Plotter(Aggregator, BaseProcessor): """ This class loads derived quantities and plot them """ solve_self_dep = False # Axes information _ax_nb = {"x": 0, "y": 1, "z": 2} # Number of each axes _axes_h = {"x": "y", "y": "x", "z": "x"} # Associated horizontal axe _axes_v = {"x": "z", "y": "z", "z": "y"} # Associated vertical axe _ax_title = {"x": r"$x$", "y": r"$y$", "z": r"$z$"} G = 1.0 # Gravitational constant # Conversion table from namelist keys (from amses config file) into LaTex strings label_convert = { "turb_rms": "$f_{rms}$", "beta": "$\\beta$", "beta_cool": "$\\beta$", "dens0": "$n_0$", "coldens0": "$\Sigma_0$", "sfr_avg_window": "window", "bx_bound": "$B_0$", "levelmax": "$l_{\max}$", "levelmin": "$l_{\min}$", "comp_frac": "$1 - \\zeta$", } # Conversion table from namelist values (from amses config file) into LaTex strings value_convert = { "sfr_avg_window": lambda x: "${:g}$ Myr".format(80 * x), #'comp_frac' : lambda x: "${:g}$".format(1 - x), "bx_bound": lambda x: "${:g}$ $\mu G$".format(5.267501272979475 * x), } def __init__( self, path, in_runs=None, in_nums=None, path_out=None, pp_params=None, selector=None, tag=None, unit_time=cst.year, **kwargs, ): """ Create a new Plotter instance. Will select run and outputs via a RunSelector object. Parameters ---------- path : path to the main folder of the simulations (ex '~/simus/myproject') in_runs : list of the runs to consider (ex ['run1', 'run2']) in_nums : list or dict of the outputs numbers to consider (ex [3, 5] or {'run1' : [3, 5], 'run2' : [4, 6]) path_out : Path where the plot will be saved. By default set to `path` pp_params : Parameters for postprocessing. See pp_params module. selector : Existing instance of RunSelector, that selects runs and outputs. If set, in_runs and in_nums will be ignored tag : string to add in the output and data files. kwargs : Keyword arguments for RunSelector. """ super(Plotter, self).__init__(path, path_out, pp_params, tag) # Select runs if selector is None: self.selector = RunSelector( path, in_runs, in_nums, self.pp_params, **kwargs ) else: self.selector = selector # Save infos self.path = path self.runs = self.selector.runs self.nums = self.selector.nums # Get comparator object self.comp = Comparator( path, self.runs, self.nums, path_out, self.pp_params, unit_time=unit_time, selector=self.selector, ) # Get postprocesor objets for each run self.pp = self.comp.pp # Define log prefix self.log_id = "[plot {}] ".format(self.pp_params.out.tag) # Define rules self.def_rules() # generate astrophysix's simulations object self.gen_simus() self.save = None def gen_simus(self): self.simulations = {} simu_fmt = self.pp_params.astrophysix.simu_fmt descr_fmt = self.pp_params.astrophysix.descr_fmt tag = self.pp_params.out.tag for run in self.runs: pp = self.pp[run][self.nums[run][0]] nml = self.comp.namelist[run] name = simu_fmt.format(run=run, tag=tag, nml=nml) exec_time = str(datetime.datetime.fromtimestamp(os.stat(pp.path).st_ctime)) description = descr_fmt.format(run=run, tag=tag, nml=nml) simu = Simulation( simu_code=ramses, name=name, alias=name.upper(), description=description, directory_path=pp.path, execution_time=exec_time, ) for param in ramses.input_parameters: try: param_setting = ParameterSetting( input_param=param, value=self.comp.get_nml(param.key, run), visibility=ParameterVisibility.BASIC_DISPLAY, ) simu.parameter_settings.add(param_setting) except KeyError as e: self._log("key {} not found".format(e), "WARNING") except AttributeError as e: self._log("{}".format(e), "WARNING") self.simulations[run] = simu def _not_self_dep(self, name, dep, dep_arg, overwrite, **kwargs): """ Check if the dependency belongs to the plotter object or to another one (comp, pp, ..) """ if dep in self.comp.rules: result = self.comp.process( dep, dep_arg, overwrite, overwrite_dep=self.overwrite_dep ) if result is not None: self.just_done.append(done) else: super(Plotter, self)._not_self_dep(name, dep, dep_arg, overwrite, **kwargs) def _needs_computation(self, overwrite, plot_filename): """ Returns true if the plot needs to be redone """ return ( self.pp_params.out.interactive or overwrite or not os.path.exists(plot_filename) ) def _process_rule( self, name, rule, arg, overwrite=False, ax=None, from_cells=False, **kwargs ): """ Open storage and figure if needed before processing a rule """ # Set full name according to argument if not arg is None: name_full = ( name + "_" + str(arg) .replace(" ", "") .replace("[", "") .replace("]", "") .replace(",", "_") .replace("'", "") .replace("/", "") ) else: name_full = name # Exit if not valid if not rule.is_valid(arg): self._log("{} is not valid in this context".format(name_full), "ERROR") return # get filetype of the output filetype = filetype_from_ext[self.pp_params.out.ext] # Select runs and nums if "select" in kwargs: select = kwargs.pop("select") runs, nums = self.selector.select(**select) else: runs = self.runs nums = self.nums datafiles = [] # Several plots if rule.kind == "classic" or rule.kind == "cells": run_num = [(run, num) for run in runs for num in nums[run]] else: run_num = [(run, None) for run in runs] for i, (run, num) in enumerate(run_num): # Find filename plot_filename = self._find_filename(name_full, run, num) # Find ax try: real_ax = ax[i] except TypeError as e: if ax is None: fig, real_ax = P.subplots(1, 1) elif not_array_error(e): real_ax = ax else: raise # Find plot save if from_cells or rule.kind == "cells": if not os.exists(self.pp[run][num].cells_filename): self.pp[run][num].load_cells() self.pp[run][num].unload_cells() save = tables.open_file(self.pp[run][num].cells_filename) elif rule.kind == "classic": save = tables.open_file(self.pp[run][num].filename) else: save = tables.open_file(self.comp.filename, "r") # Call plot routine try: plot_info = self._plot_rule( rule, save, arg, plot_filename, overwrite, ax=real_ax, run=run, **kwargs, ) finally: save.close() # Save in astrophysix format df = rule.datafile(arg) df[filetype] = plot_filename if plot_info is not None: df.plot_info = plot_info if num is not None: snap = self.pp[run][num].snapshot if overwrite and df.name in snap.datafiles: del snap.datafiles[df.name] elif df.name not in snap.datafiles: snap.datafiles.add(df) if snap not in self.simulations[run].snapshots: self.simulations[run].snapshots.add(snap) datafiles.append(df) return datafiles def _plot_rule(self, rule, save, arg, plot_filename, overwrite, ax, **kwargs): """ Once all dependencies are met, actually process the rule """ P.sca(ax) if self._needs_computation(overwrite, plot_filename): plot_info = rule.plot(save, arg, **kwargs) if not self.pp_params.out.interactive: P.tight_layout(pad=1) if self.pp_params.out.save: P.savefig(plot_filename) self._log("{} plotted".format(plot_filename), "SUCCESS") else: self._log( "{} plotted".format(os.path.basename(plot_filename)), "SUCCESS" ) if not self.pp_params.out.interactive: P.close() return plot_info else: self._log("Plot {} is already done, skipping...".format(plot_filename)) def _find_filename(self, name_full, run=None, num=None, fmt=None): """ Determine a filename based on rule name, run, output and parameters """ tag_name = self.pp_params.out.tag if fmt is None and self.pp_params.out.fmt == "": if not self.pp_params.out.tag == "": tag_name = "_" + tag_name if not run is None and not num is None: fmt = "{out}/{run}/{name}{tag}_{run}_{num:05}{ext}" elif not run is None: fmt = "{out}/{run}/{name}{tag}_{run}{ext}" else: fmt = "{out}/{name}{tag}{ext}" elif fmt is None: fmt = self.pp_params.out.fmt nml = None if not run is None: nml = self.comp.namelist[run] return fmt.format( run=run, name=name_full, tag=tag_name, num=num, nml=nml, out=self.path_out, ext=self.pp_params.out.ext, ) def _label_run(self, run, node, label, nml_key, time=None): """ Set up a label for the run from the namelist and parameters """ def get_label_nml(nml_key): prop_name = os.path.basename(nml_key) if prop_name in self.label_convert: prop_label = self.label_convert[prop_name] else: prop_label = prop_name prop_value = self.comp.get_nml(nml_key, run) if prop_name in self.value_convert: prop_value_str = self.value_convert[prop_name](prop_value) elif type(prop_value) in [int, float]: prop_value_str = convert_exp(prop_value, digits=5) else: prop_value_str = str(prop_value) return r"{} = {}".format(prop_label, prop_value_str) if nml_key is None and label is None: if ( "attrs" in self.save.root._v_attrs and run in self.save.root._v_attrs.attrs ): label_run = r"{}".format(self.save.root._v_attrs.attrs[run].label) else: label_run = run elif not nml_key is None: if not type(nml_key) == list: nml_key = [nml_key] label_run = ", ".join(map(get_label_nml, nml_key)) if not label is None: label_run = label + " (" + label_run + ")" else: label_run = label return label_run def _ax_label_unit(self, node, label, unit, unit_coeff, put_units=True): """ Find appropriate labels for axis """ if label is None: if "label" in node._v_attrs: label = node._v_attrs.label elif node._v_name in self.label_convert: label = self.label_convert[node._v_name] elif not node._v_title == "": label = node._v_title else: label = node._v_name if "unit" in node._v_attrs: unit_old = node._v_attrs.unit else: unit_old = cst.none if unit is None: unit = unit_old if put_units: if not unit_coeff == 1: base = unit unit = unit_coeff * unit label = label + unit_str(unit, base=base) else: label = label + unit_str(unit) return label, unit_old, unit def _snapshot_title(self, run, node, title, nml_key, put_time, unit_time=cst.Myr): title = self._label_run(run, node, title, nml_key) if put_time: time = self.save.root._v_attrs.time * self.comp.info["unit_time"] u_str = unit_str(unit_time, format="{unit}") time_str = self.pp_params.plot.time_fmt.format( time.express(unit_time), u_str ) if len(title) > 0: title = title + " | " + time_str else: title = time_str return title def _plot_map( self, name, ax_los, run, xlabel=None, ylabel=None, label=None, unit=None, unit_coeff=1.0, overlays=[], overlays_kwargs=[], title=None, put_title=True, nml_key=None, put_time=True, unit_time=cst.Myr, put_units=True, unit_space=cst.pc, cmap="plasma", norm="log", put_cbar=True, autoscale=True, transform=None, **kwargs, ): """ Plot data on a map """ ax_h = self._axes_h[ax_los] ax_v = self._axes_v[ax_los] im_extent = self.save.root.maps._v_attrs.im_extent unit_length = self.save.root._v_attrs["unit_length"] im_extent = np.array(im_extent) * unit_length.express(unit_space) node = self.save.get_node("/maps/{}_{}".format(name, ax_los)) dmap = node.read() label, unit_old, unit = self._ax_label_unit( node, label, unit, unit_coeff, put_units ) dmap = dmap * unit_old.express(unit) if transform is not None: dmap = transform(dmap) if norm == "log": norm = mpl.colors.LogNorm() elif norm == "linear": norm = mpl.colors.NoNorm() if autoscale and not norm is None: norm.autoscale(dmap) im = P.imshow( dmap, extent=im_extent, origin="lower", norm=norm, cmap=cmap, **kwargs ) P.locator_params(axis="both", nbins=self.pp_params.plot.ntick) if xlabel is None: xlabel = self._ax_title[ax_h] if ylabel is None: ylabel = self._ax_title[ax_v] if put_units: xlabel = xlabel + unit_str(unit_space) ylabel = ylabel + unit_str(unit_space) P.xlabel(xlabel) P.ylabel(ylabel) try: cbar = P.colorbar(im, cax=P.gca().cax) except AttributeError: cbar = P.colorbar() if put_title: title = self._snapshot_title(run, node, title, nml_key, put_time, unit_time) P.title(title) if not label is None: cbar.set_label(label) for i, plot_overlay in enumerate(overlays): if plot_overlay in self.overlays: plot_overlay = self.overlays[plot_overlay] try: plot_overlay(ax_los, im_extent, **overlays_kwargs[i]) except: plot_overlay(ax_los, im_extent) return PlotInfo( plot_type=PlotType.IMAGE, xaxis_values=np.linspace(im_extent[0], im_extent[1], dmap.shape[0] + 1), yaxis_values=np.linspace(im_extent[2], im_extent[3], dmap.shape[1] + 1), values=dmap, xaxis_log_scale=False, yaxis_log_scale=False, values_log_scale=(norm == mpl.colors.LogNorm()), xaxis_label=xlabel, yaxis_label=ylabel, values_label=label, xaxis_unit=unit_space, yaxis_unit=unit_space, values_unit=unit, plot_title=title, ) def _overlay_contour( self, ax_los, im_extent, map_name, log=False, lvl_array=None, lw=None, lvl_th=None, lvl_max_lbl=np.inf, lvl_offset=0, lbl_fmt="%g", **kwargs, ): """ Add an overlay : contour of other map """ map_contour = self.save.get_node("/maps/{}_{}".format(map_name, ax_los)).read() if log: map_contour = np.log10(map_contour) # Computing linewidths mask_fin = np.isfinite(map_contour) if lvl_array is None: lvl_array = np.arange( np.min(map_contour[mask_fin]), np.max(map_contour[mask_fin]) + 1 ) if lw is None: lw = np.ones(lvl_array.size) * 2 if lvl_th: lw[lvl_array >= lvl_th] = lw[lvl_array >= lvl_th] ** ( lvl_th - lvl_array[lvl_array >= lvl_th] ) lw[lvl_array < lvl_th] = 1.0 cont = P.contour( map_contour, extent=im_extent, origin="lower", linewidths=lw, levels=lvl_array, **kwargs, ) # used levels lvls = np.array(cont.levels) + lvl_offset cont.levels = lvls P.clabel( cont, lvls[np.array(lvls) < lvl_max_lbl], inline=1, fontsize=8.0, fmt=lbl_fmt, ) def _overlay_levels(self, ax_los, im_extent, **kwargs): """ Add an overlay : AMR levels """ return self._overlay_contour( ax_los, im_extent, "levels", lbl_fmt="%1d", lvl_offset=1, lvl_th=8, lvl_max_lbl=11, **kwargs, ) def _overlay_speed( self, ax_los, im_extent, unit=cst.km_s, unit_coeff=1.0, key_v=None, **kwargs ): """ Add an overlay : velocity vector field """ ax_h = self._axes_h[ax_los] ax_v = self._axes_v[ax_los] dmap_vh_node = self.save.get_node("/maps/speed_h_{}".format(ax_los)) dmap_vh = dmap_vh_node.read() dmap_vv = self.save.get_node("/maps/speed_v_{}".format(ax_los)).read() label, unit_old, unit = self._ax_label_unit(dmap_vh_node, "", unit, unit_coeff) vel_red = self.pp_params.plot.vel_red radius = self.save.root.maps._v_attrs.radius center = self.save.root.maps._v_attrs.center lbox = self.save.root._v_attrs.lbox map_vh_red = dmap_vh[::vel_red, ::vel_red] * unit_old.express( unit ) # take only a subset of velocities map_vv_red = dmap_vv[::vel_red, ::vel_red] * unit_old.express(unit) # TODO : redo this with im_extent nh = map_vh_red.shape[0] nv = map_vv_red.shape[1] vec_h = ( np.arange(nh) * 2.0 / nh * radius - radius + center[0] + radius / nh ) * lbox vec_v = ( np.arange(nv) * 2.0 / nv * radius - radius + center[1] + radius / nv ) * lbox hh, vv = np.meshgrid(vec_h, vec_v) norm_v = np.sqrt(map_vh_red ** 2 + map_vv_red ** 2) max_v = np.max(norm_v) min_v = np.min(norm_v) Q = P.quiver(hh, vv, map_vh_red, map_vv_red, units="width", **kwargs) if key_v is None: key_v = (max_v + min_v) / 2.0 P.quiverkey( Q, 0.6, 0.98, key_v, r"${:g}$".format(key_v) + label, labelpos="E", coordinates="figure", ) def _overlay_B(self, ax_los, im_extent, **kwargs): """ Add an overlay : magnetic streamlines """ ax_h = self._axes_h[ax_los] ax_v = self._axes_v[ax_los] dmap_Bh_node = self.save.get_node("/maps/B_h_{}".format(ax_los)) dmap_Bh = dmap_Bh_node.read() dmap_Bv = self.save.get_node("/maps/B_v_{}".format(ax_los)).read() # TODO : redo this with im_extent vel_red = self.pp_params.plot.vel_red radius = self.save.root.maps._v_attrs.radius center = self.save.root.maps._v_attrs.center lbox = self.save.root._v_attrs.lbox map_Bh_red = dmap_Bh[::vel_red, ::vel_red] # take only a subset of velocities map_Bv_red = dmap_Bv[::vel_red, ::vel_red] nh = map_Bh_red.shape[0] nv = map_Bv_red.shape[1] vec_h = ( np.arange(nh) * 2.0 / nh * radius - radius + center[0] + radius / nh ) * lbox vec_v = ( np.arange(nv) * 2.0 / nv * radius - radius + center[1] + radius / nv ) * lbox hh, vv = np.meshgrid(vec_h, vec_v) P.streamplot(hh, vv, map_Bh_red, map_Bv_red, **kwargs) def _plot_radial( self, name, ax_los, run, ylabel=None, xlog=False, ylog=False, ytransform=None, label=None, title=None, nml_key=None, put_title=True, put_time=True, unit_time=cst.Myr, **kwargs, ): """ Plot a radial profile (for disks, OUTDATED) """ radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read() bin_centers = 0.5 * (radial_bins[1:] + radial_bins[:-1]) node = self.save.get_node("/radial/{}_{}".format(name, ax_los)) mean_bin = node.read() if ytransform is not None: mean_bin = ytransform(mean_bin) P.xlabel(r"$r$") if xlog: P.xscale("log") if ylog: P.yscale("log") if not ylabel is None: P.ylabel(ylabel) title = self._snapshot_title(run, node, title, nml_key, put_time, unit_time) if put_title: P.title(title) if label == None: label = title P.plot(bin_centers, mean_bin, label=label, **kwargs) def _plot_hist( self, name, ax_los=None, run=None, group="/hist/", xlabel=None, unit=None, unit_coeff=1.0, ytransform=None, label=None, put_title=True, title=None, nml_key=None, put_time=True, unit_time=cst.Myr, xlog=None, ylog=False, kind="bar", ylabel="$\mathcal{P}$", color=None, colors=None, nml_color=None, fit=None, fitlabel=None, **kwargs, ): """ Plot an histogram (PDF, etc ...) """ # Get node if not ax_los is None: name = name + "_" + ax_los node = self.save.get_node(group + name) if xlog is None: try: xlog = node._v_attrs_.logbins except: xlog = False # get label and units xlabel, unit_old, unit = self._ax_label_unit(node, label, unit, unit_coeff) # Read data if "mean" in node: index = node["runs"].read().index(run.encode()) values, centers = node["mean"].read()[index] else: values, centers = node.read() if xlog: centers = centers + np.log10(unit_old.express(unit)) else: centers = centers * unit_old.express(unit) if ytransform is not None: values = ytransform(values) width = centers[1] - centers[0] # Set title title = self._snapshot_title(run, node, title, nml_key, put_time, unit_time) if put_title: P.title(title) if label == None: label = title # Set colors if color is None and not colors is None: if nml_color is None: color = colors[run] else: nml = self.comp.get_nml(nml_color, run) try: color = colors[nml] except: color = colors(nml) # Actual plot if kind == "bar": P.bar(centers, values, width, log=ylog, color=color, label=label, **kwargs) elif kind == "step": if ylog: P.yscale("log") P.step(centers, values, where="mid", color=color, label=label, **kwargs) else: raise ValueError("kind must be 'bar' or 'step'") # put labels if not label is None: P.xlabel(xlabel) if not ylabel is None: P.ylabel(ylabel) # Also diplay fit, previously saved if ax_los is not None and "/hist/fit_" + name + "_" + ax_los in self.save: slope = node.attrs.slope origin = node.attrs.origin P.plot( centers, 10 ** (slope * centers + origin), "--", linewidth=2, color="orange", ) # or a new one if not fit is None: self._overlay_fit( centers, values, kind=fit, ls="--", lw=1.5, label=fitlabel ) # returns PlotInfo (for Galactica) edges = np.append(centers - width / 2.0, centers[-1] + width / 2.0) return PlotInfo( plot_type=PlotType.HISTOGRAM, xaxis_values=edges, yaxis_values=values, xaxis_log_scale=False, yaxis_log_scale=ylog, xaxis_label=xlabel, yaxis_label=ylabel, xaxis_unit=unit, yaxis_unit=cst.none, plot_title=title, ) def _plot( self, name_x, name_y, node_arg=None, xlabel=None, ylabel=None, label=None, xunit=None, yunit=None, xunit_coeff=1.0, yunit_coeff=1.0, fit=None, fitlabel=None, smooth=0, nml_key=None, run=None, yerr=None, yerr_kind="std", sigma_err=2.0, grid=False, put_time=False, unit_time=cst.Myr, colors=None, nml_color=None, legend=None, subname_x=None, subname_y=None, **kwargs, ): """ Generic plot routine, with name_x and name_y two path in the hdf5 file """ # Get proper hdf5 names if not node_arg is None: name_x, name_y = name_x + "_" + node_arg, name_y + "_" + node_arg # Get hdf5 nodes node_x = self.save.get_node(name_x) node_y = self.save.get_node(name_y) # If the actual data is in a,other file, fetch it if subname_x: hdf5_x = tables.open_file(node_x.read()) node_x = hdf5_x.get_node(subname_x) if subname_y: hdf5_y = tables.open_file(node_y.read()) node_y = hdf5_y.get_node(subname_y) # Find proper labels xlabel, xunit_old, xunit = self._ax_label_unit( node_x, xlabel, xunit, xunit_coeff ) ylabel, yunit_old, yunit = self._ax_label_unit( node_y, ylabel, yunit, yunit_coeff ) # If relevent, get time if put_time: time = self.save.root._v_attrs.time * self.comp.info["unit_time"] time_str = self.pp_params.plot.time_fmt.format( time.express(unit_time), unit_time.latex ) if label is not None and len(label) > 0: label = label + " | " + time_str else: label = time_str # Manage the different forms in which the data may be stored : # Possibilities are : plain array, dict of arrays (mean, std, ..) or dict of array (runs) if node_y._v_attrs.CLASS == "ARRAY": x = node_x.read() * xunit_old.express(xunit) y = node_y.read() * yunit_old.express(yunit) mask = np.isfinite(x) & np.isfinite(y) x, y = x[mask], y[mask] elif "mean" in node_y: x = node_x.read() * xunit_old.express(xunit) y = node_y.mean.read() * yunit_old.express(yunit) if yerr_kind == "std": std = node_y.std.read() * yunit_old.express(yunit) yerr_min = y - sigma_err * std yerr_max = y + sigma_err * std elif yerr_kind == "min_max": yerr_min = node_y.min.read() * yunit_old.express(yunit) yerr_max = node_y.max.read() * yunit_old.express(yunit) elif yerr_kind == "95per": yerr_min = node_y.q025.read() * yunit_old.express(yunit) yerr_max = node_y.q975.read() * yunit_old.express(yunit) elif yerr_kind == "68per": yerr_min = node_y.q16.read() * yunit_old.express(yunit) yerr_max = node_y.q84.read() * yunit_old.express(yunit) else: yerr_min = y yerr_max = y yerr = yerr_max - yerr_min mask = np.isfinite(x) & np.isfinite(y) & np.isfinite(yerr) x, y, yerr, yerr_min, yerr_max = ( x[mask], y[mask], yerr[mask], yerr_min[mask], yerr_max[mask], ) else: x, y = node_x[run], node_y[run] mask = np.isfinite(x) & np.isfinite(y) x, y = x[mask], y[mask] if isinstance(yerr, str): yerr = self.save.get_node(yerr).read() if smooth > 0: y = gaussian_filter1d(y, sigma=smooth) if not run is None: label = self._label_run(run, node_y, label, nml_key) # Look if special colors method is used if colors is None: if yerr is None: (base_line,) = P.plot(x, y, label=label, **kwargs) else: base_line, _, _ = P.errorbar(x, y, yerr=yerr, label=label, **kwargs) else: if nml_color is None: color = colors[run] elif nml_color == "time": time = ( self.save.root._v_attrs.time * self.comp.info["unit_time"] ).express(unit_time) color = colors(time) else: nml = self.comp.get_nml(nml_color, run) try: color = colors[nml] except: color = colors(nml) if yerr is None: (base_line,) = P.plot(x, y, label=label, color=color, **kwargs) else: base_line, _, _ = P.errorbar( x, y, yerr=yerr, color=color, label=label, **kwargs ) # Ax decorations P.xlabel(xlabel) P.ylabel(ylabel) if grid: P.grid() if legend: P.legend() if not fit is None: self._overlay_fit( x, y, yerr, kind=fit, ls="--", lw=1.5, color=base_line.get_color(), label=fitlabel, ) if subname_x: hdf5_x.close() if subname_y: hdf5_y.close() def _pspec(self, name, **kwargs): """ Plot power spectrum """ del kwargs["run"] file_pspec = self.save.get_node("/hdf5/pspec").read() num = self.save.root._v_attrs.num getattr(pspec_read, "pspec_" + name)(file_pspec, ".", num, **kwargs) def _overlay_fit(self, x, y, yerr=None, kind="linear", label=None, **kwargs): """ Add an overlay : fit a curve, linear or powerlaw """ if kind == "linear": if yerr is None: (a, b, rho, _map_rule, stderr) = linregress(x, y) self._log( "Linear fit y = {} x + {} with R^2 = {} and error is {}".format( a, b, rho, stderr ) ) if label is None: label = r"Linear fit with slope ${:.3g}$ and $R^2 = {:.3f}$".format( a, rho ) else: fit = polyfit(x, y, 1, w=[1.0 / ty for ty in yerr], full=True) c = fit[0] residual = fit[1][0][0] b, a = c[0], c[1] self._log( "Linear fit y = {} x + {} with residual {}".format(a, b, residual) ) if label is None: label = r"Linear fit with slope ${:.3g}$".format(a) P.plot(x, a * x + b, label=label, **kwargs) elif kind == "power_law": if yerr is None: (a, b, rho, _map_rule, stderr) = linregress(np.log10(x), np.log10(y)) self._log( "Power law fit y = x^({}) * {} with R^2 = {} and error is {}".format( a, 10 ** b, rho, stderr ) ) else: fitfunc = lambda p, x: p[0] + p[1] * x errfunc = lambda p, x, y, err: (y - fitfunc(p, x)) / err pinit = [1.0, -1.0] out = optimize.leastsq( errfunc, pinit, args=(np.log10(x), np.log10(y), yerr / y), full_output=1, ) c = out[0] b, a = c[0], c[1] residual = errfunc(c, np.log10(x), np.log10(y), yerr / y) self._log( "Power law fit y = x^({}) * {} with residual {}".format( a, 10 ** b, residual ) ) if label is None: label = r"Power-law fit with index {:.1f}".format(a) P.plot(x, (10 ** b) * x ** a, label=label, **kwargs) def overlay_kennicutt(self, n0, step): """ Add an overlay : Kennicutt mass accretion """ P.grid(False) ylim = P.ylim() (tmin, tmax) = P.xlim() tmax = tmax + 20 ymax = P.ylim()[1] ssfr_sun = 2.5e-9 ssfr_ken = ssfr_sun * n0 ** 1.4 coeff = ssfr_ken * 1e6 * (self.comp.info["unit_length"].express(cst.pc)) ** 2 for i in np.arange(tmin, max(tmax, tmin + ymax / coeff), step): t = np.linspace(0, tmax, 1000) P.plot(t + i, t * coeff, ls="--", lw=0.9, color="grey") P.plot(t + tmin, (t + i - tmin) * coeff, ls="--", lw=0.9, color="grey") P.xlim(tmin, tmax) P.ylim(ylim) def _gen_from_log(self, logrule, name, description="Generated"): self.rules[name] = PlotRule( self, partial( self._plot, "/series/" + logrule + "/time", "/series/" + logrule + "/" + name, xunit=cst.Myr, ), description=description, kind="series", dependencies=[logrule], ) def def_rules(self): """ This is where rules are defined """ self.rules = { # Generic rules "plot": PlotRule( self, lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="comp" ), "coldens": PlotRule( self, partial( self._plot_map, "coldens", label=r"$\Sigma$", # unit=cst.coldens ), "Column density map", dependencies=["coldens"], ), "T": PlotRule( self, partial( self._plot_map, "T", label=r"$T$", ), "Temperature map", dependencies=["T"], ), "alpha_disk": PlotRule( self, partial(self._plot_map, "alpha_disk", label=r"$\alpha$"), "Map of the Shakura&Sunaev alpha parameter for disks", dependencies=["alpha_disk"], ), "alpha_grav": PlotRule( self, partial(self._plot_map, "alpha_grav", label=r"$\alpha_g$"), "Map of the grav Shakura&Sunaev alpha parameter for disks", dependencies=["alpha_grav"], ), "vphi": PlotRule( self, partial( self._plot_map, "vphi", label=r"$v_\phi$", # unit=cst.km_s ), "Azimuthal speed", dependencies=["vphi"], ), "vr": PlotRule( self, partial( self._plot_map, "vr", label=r"$v_r$", # unit=cst.km_s ), "Radial speed", dependencies=["vr"], ), "rho": PlotRule( self, partial( self._plot_map, "rho", label=r"$\rho$", # unit=cst.Msun_pc3 ), "Density slice at s = 0, with s = x, y or z.", dependencies=["rho"], ), "coldens_l": PlotRule( self, partial( self._plot_map, "coldens", label=r"$\Sigma$", unit=cst.coldens, overlays=[self._overlay_levels], ), "Column density with level overlay", dependencies=["coldens", "levels"], ), "rho_v": PlotRule( self, partial( self._plot_map, "rho", label=r"$\rho$", unit=cst.Msun_pc3, overlays=[self._overlay_speed], ), "Density slice with speed overlay", dependencies=["rho", "speed_h", "speed_v"], ), "rho_B": PlotRule( self, partial( self._plot_map, "rho", label=r"$\rho$", unit=cst.Msun_pc3, overlays=[self._overlay_B], ), "Density slice with magnetic field overlay", dependencies=["rho", "B_h", "B_v"], ), "rho_B_vel": PlotRule( self, partial( self._plot_map, "rho", label=r"$\rho$", unit=cst.Msun_pc3, overlays=[self._overlay_B, self._overlay_speed], ), "Density slice with magnetic field and velocity overlay", dependencies=["rho", "B_h", "B_v", "speed_h", "speed_v"], ), "jeans_ratio": PlotRule( self, partial( self._plot_map, "jeans_ratio", vmin=0.1, vmax=100, cmap="RdBu_r", overlays=[self._overlay_levels], ), "Jeans' lenght divided by the max resolution", dependencies=["jeans_ratio", "levels"], ), "Q": PlotRule( self, partial( self._plot_map, "rho", label=r"$Q$", vmin=0.01, vmax=100, cmap="RdBu_r", ), "Toomre Q parameter for a Keplerian disk", dependencies=["Q"], ), "rho_pdf": PlotRule( self, partial(self._plot_hist, "rho_pdf"), "$\rho$-PDF", dependencies=["rho_pdf"], ), "rho_pdf_mw": PlotRule( self, partial(self._plot_hist, "rho_pdf_mw"), "Mass weighted $\rho$-PDF", dependencies=["rho_pdf_mw"], ), "cos_pdf": PlotRule( self, partial(self._plot_hist, "cos_pdf"), "cos-PDF", dependencies=["cos_pdf", "mwa_speed"], ), "avg_coldens_pdf": PlotRule( self, partial( self._plot_hist, "avg_time_coldens_pdf_z", group="/comp/", xlog=True, put_time=False, ), "Column density PDF, averaged in time", kind="runs", dependencies={"avg_time_coldens_pdf": "z"}, ), "T_pdf": PlotRule( self, partial(self._plot_hist, "T_pdf"), "T-PDF on a 2D slice", dependencies=["T_pdf"], ), "P_pdf": PlotRule( self, partial(self._plot_hist, "P_pdf"), "P-PDF on a 2D slice ", dependencies=["P_pdf"], ), "B_int": PlotRule( self, partial( self._plot_map, "B_int", label=r"$\mid \mathrm{B} \mid$", unit=cst.T ), "Magnetic intensity map", dependencies=["B_int"], ), "Brho": PlotRule( self, partial( self._plot, "/datasets/Brho/rho", "/datasets/Brho/B", label=r"$\mathrm{B} $", put_time=True, ), "Brho on a 2D slice ", dependencies=["Brho"], ), "Ek_Eb_rho": PlotRule( self, partial( self._plot, "/datasets/Ek_Eb_rho/rho", "/datasets/Ek_Eb_rho/Ek_Eb_rho", label=r"Ek/Eb", put_time=True, ), "Ek/Eb on a 2D slice ", dependencies=["Ek_Eb_rho", "mwa_speed"], ), "rho_prof": PlotRule( self, partial(self._plot, "/profile/axis", "/profile/rho_prof"), "Density profile", dependencies=["axis", "rho_prof"], ), "pspec": PlotRule(self, self._pspec, dependencies={"pspec": None}), "sbeta": PlotRule( self, partial( self._plot, "/comp/nml_cloud_params/beta_cool", "/comp/avg_time_pdf_slope_coldens", ), "Slope of the Sigma-PDF against cooling beta factor", kind="comp", dependencies={ "nml": "cloud_params/beta_cool", "avg_time_pdf_slope_coldens": None, }, ), "sbeta_onavg": PlotRule( self, partial( self._plot, "/comp/sbeta_onavg/beta", "/comp/sbeta_onavg/slope", yerr="/comp/sbeta_onavg/stderr", ), "Slope of the time averaged Sigma-PDF against cooling beta factor", kind="comp", dependencies=["sbeta_onavg"], ), "sink_mass": PlotRule( self, partial( self._plot, "/series/sinks_from_log/time", "/series/sinks_from_log/mass_sink", xunit=cst.Myr, yunit=cst.Msun, ), "Mass of the sinks as a function of time", kind="series", dependencies=["sinks_from_log"], ), "ssm": PlotRule( self, partial( self._plot, "/series/sinks_from_log/time", "/series/sinks_from_log/ssm", xunit=cst.Myr, yunit=cst.Msun / cst.pc ** 2, ), "Mass of the sinks as a function of time divided by surface", kind="series", dependencies=["ssm"], ), "assfr": PlotRule( self, partial( self._plot, "/series/sfr_from_log/time", "/series/sfr_from_log/sfr", ylabel="Averaged surfacic SFR", xunit=cst.Myr, yunit=cst.ssfr, ), kind="series", dependencies=["sfr_from_log"], ), "issfr": PlotRule( self, partial( self._plot, "/series/sinks_from_log/time", "/series/sinks_from_log/issfr", ylabel="Surfacic SFR", xunit=cst.Myr, yunit=cst.ssfr, ), kind="series", dependencies=["issfr"], ), "turb_rms": PlotRule( self, partial( self._plot, "/series/rms_from_log/time", "/series/rms_from_log/turb_rms", xunit=cst.Myr, ), "Turbulent RMS", kind="series", dependencies=["rms_from_log"], ), "turb_energy": PlotRule( self, partial( self._plot, "/series/rms_from_log/time", "/series/rms_from_log/turb_energy", xunit=cst.Myr, ), "Turbulent energy", kind="series", dependencies=["rms_from_log"], ), "turb_power": PlotRule( self, partial( self._plot, "/series/rms_from_log/time", "/series/rms_from_log/turb_power", xunit=cst.Myr, ), "Turbulent power", kind="series", dependencies=["turb_power"], ), "sigma": PlotRule( self, partial( self._plot, "/series/time", "/series/time_sigma", ylabel="$\\sigma$", xunit=cst.Myr, yunit=cst.km_s, ), "Velocity dispersion", kind="series", dependencies=["time_sigma"], ), "mwa_B_int": PlotRule( self, partial( self._plot, "/series/time", "/series/time_mwa_B_int", xunit=cst.Myr, yunit=cst.uG, ), "Magnetic intensity average", kind="series", dependencies=["time_mwa_B_int"], ), "mass": PlotRule( self, partial( self._plot, "/series/time", "/series/time_mass", xunit=cst.Myr, yunit=cst.Msun, ), "Total mass in the box", kind="series", dependencies=["time_mass"], ), "max_fluct_coldens": PlotRule( self, partial( self._plot, "/series/time", "/series/time_max_fluct_coldens_z", ylabel="$\\max(\Sigma/\overline{\Sigma})$", xunit=cst.Myr, ), "Maximal fluctuation of the column density against time", kind="series", dependencies={"time_max_fluct_coldens": "z"}, ), } averageables = [ "coldens", "rho", "T", "Q", "vr", "vphi", "rho_avg", "P_avg", "T_avg", "alpha_disk", "alpha_grav", ] for name in averageables: self.rules["rad_" + name] = PlotRule( self, partial(self._plot_radial, "rad_avg_" + name, xlog=True, ylog=True), "Azimuthal average of {}".format(name), dependencies=["radial_bins", "rad_avg_" + name], ) self.rules["avg_map_" + name] = PlotRule( self, partial(self._plot_map, "avg_map_" + name), "Map of the radial average of {}".format(name), dependencies=["avg_map_" + name], ) self.rules["fluct_" + name] = PlotRule( self, partial(self._plot_map, "fluct_" + name, cmap="RdBu_r"), "Fluctuation of {}".format(name), dependencies=["fluct_" + name], ) self.rules["pdf_" + name] = PlotRule( self, partial(self._plot_hist, "pdf_" + name, ylog=True), "Probability density function of {} fluctuations".format(name), dependencies=["fit_pdf_" + name], ) for name_bin in averageables: if name_bin is not name: group = "mbb_{}_{}".format(name, name_bin) self.rules["mbb_" + name + "_" + name_bin] = PlotRule( self, partial(self._plot_hist, group, ylabel=r"$\alpha$"), "Mean of {} by bins of {}".format(name, name_bin), dependencies=[group], ) for name in ["step", "mcons", "econs", "epot", "ekin", "eint", "emag"]: self._gen_from_log("cons_from_log", name) # Generic rules directly from Ramses fields for field in self.pp_params.pymses.variables: def generic_rule(name): self.rules["slice_" + name] = PlotRule( self, partial(self._plot_map, "slice_" + name), "{} slice".format(name), dependencies=["slice_" + name], ) self.rules[name + "_mwavg"] = PlotRule( self, partial(self._plot_map, name + "_mwavg"), "Ax mass-weighted averaged {}".format(name), dependencies=[name + "_mwavg"], ) self.rules[name + "_avg"] = PlotRule( self, partial(self._plot_map, name + "_avg"), "Ax averaged {}".format(name), dependencies=[name + "_avg"], ) # special for vectors if field in ["g", "vel"]: # Components for i, dir in enumerate(["x", "y", "z"]): generic_rule(field + "x") # Radial generic_rule(field + "r") # Orthoradial generic_rule(field + "phi") # Norm generic_rule(field + "_norm") else: generic_rule(field) # Dict of overlays self.overlays = { "B": self._overlay_B, "speed": self._overlay_speed, "levels": self._overlay_levels, "contour": self._overlay_contour, } super(Plotter, self).def_rules()