Files
pipeline/plotter.py
T
2020-12-14 16:59:00 +01:00

1685 lines
55 KiB
Python

# 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
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
P.rcParams["image.cmap"] = "plasma"
P.rcParams["savefig.dpi"] = 400
tex_params = {"text.latex.preamble": r"\usepackage{amsmath}"}
P.rcParams.update(tex_params)
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 exetute 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)
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,
**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, 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()
self.save = None
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:
done = self.comp.process(
dep, dep_arg, overwrite, overwrite_dep=self.overwrite_dep
)
self.just_done.extend(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,
movie=False,
from_cells=False,
**kwargs,
):
"""
Open storage and figure if needed before processing a rule
"""
if not arg is None:
name_full = (
name
+ "_"
+ str(arg)
.replace(" ", "")
.replace("[", "")
.replace("]", "")
.replace(",", "_")
.replace("'", "")
.replace("/", "")
)
else:
name_full = name
if rule.is_valid(arg):
if rule.kind == "classic" or rule.kind == "cells":
if "select" in kwargs:
select = kwargs.pop("select")
runs, nums = self.selector.select(**select)
elif "runs" in kwargs:
runs = kwargs.pop("runs")
if isinstance(runs, RunSelector):
nums = runs.nums
runs = runs.runs
else:
nums = self.nums
else:
runs = self.runs
nums = self.nums
i = 0
for run in runs:
files = []
for num in nums[run]:
plot_filename = self._find_filename(name_full, run, num)
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")
try:
self._plot_rule(
rule,
save,
arg,
plot_filename,
overwrite,
ax=ax[i],
run=run,
**kwargs,
)
except TypeError as e:
if str(e) in [
"'LocatableAxes' object does not support indexing",
"'AxesSubplot' object does not support indexing",
"'AxesSubplot' object is not subscriptable",
"'Axes' object is not subscriptable",
"'LocatableAxes' object is not subscriptable",
]:
self._plot_rule(
rule,
save,
arg,
plot_filename,
overwrite,
ax=ax,
run=run,
**kwargs,
)
elif ax is None:
fig = P.figure()
self._plot_rule(
rule,
save,
arg,
plot_filename,
overwrite,
ax=P.gca(),
run=run,
**kwargs,
)
else:
raise
finally:
save.close()
i = i + 1
files.append(plot_filename)
else:
if "select" in kwargs and not "runs" in kwargs:
select = kwargs.pop("select")
runs, nums = self.selector.select(**select)
if not rule.kind == "runs":
kwargs["runs"] = runs
elif rule.kind == "runs" and "runs" in kwargs:
runs = kwargs.pop("runs")
else:
runs = self.runs
if ax is None:
ax = P.gca()
if rule.kind == "series" and len(runs) == 1:
run = self.runs[0]
plot_filename = self._find_filename(name_full, run)
else:
plot_filename = self._find_filename(name_full)
save = tables.open_file(self.comp.filename, "r")
try:
if rule.kind == "runs":
for i, run in enumerate(runs):
try:
self._plot_rule(
rule,
save,
arg,
plot_filename,
overwrite,
ax=ax[i],
run=run,
**kwargs,
)
except TypeError as e:
if str(e) in [
"'LocatableAxes' object does not support indexing",
"'AxesSubplot' object does not support indexing",
"'AxesSubplot' object is not subscriptable",
"'Axes' object is not subscriptable",
"'LocatableAxes' object is not subscriptable",
]:
self._plot_rule(
rule,
save,
arg,
plot_filename,
overwrite,
ax=ax,
run=run,
**kwargs,
)
else:
self._plot_rule(
rule, save, arg, plot_filename, overwrite, ax, **kwargs
)
finally:
save.close()
else:
self._log("{} is not valid in this context".format(name_full), "ERROR")
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):
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()
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):
"""
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):
"""
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 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 _plot_map(
self,
name,
ax_los,
run,
label=None,
unit=None,
unit_coeff=1.0,
overlays=[],
overlays_kwargs=[],
title=None,
put_title=True,
nml_key=None,
put_time=True,
time_unit=cst.Myr,
unit_space=cst.pc,
cmap="plasma",
norm="log",
put_cbar=True,
autoscale=True,
**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)
dmap = dmap * unit_old.express(unit)
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)
P.xlabel(self._ax_title[ax_h] + unit_str(unit_space))
P.ylabel(self._ax_title[ax_v] + unit_str(unit_space))
try:
cbar = P.colorbar(im, cax=P.gca().cax)
except AttributeError:
cbar = P.colorbar()
if not label is None:
cbar.set_label(label)
if put_title:
title = self._label_run(run, node, title, nml_key)
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(time_unit), time_unit.latex
)
if len(title) > 0:
title = title + " | " + time_str
else:
title = time_str
P.title(title)
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)
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,
time_unit=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)
if put_title:
title = self._label_run(run, node, title, nml_key)
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(time_unit), time_unit.latex
)
if len(title) > 0:
title = title + " | " + time_str
else:
title = time_str
P.title(title)
if label is None:
label = title
P.plot(bin_centers, mean_bin, label=label, **kwargs)
P.plot(bin_centers, mean_bin, label=title, **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,
title=None,
nml_key=None,
put_time=True,
time_unit=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 ...)
"""
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
xlabel, unit_old, unit = self._ax_label_unit(node, label, unit, unit_coeff)
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)
title = self._label_run(run, node, title, nml_key)
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(time_unit), time_unit.latex
)
if len(title) > 0:
title = title + " | " + time_str
else:
title = time_str
P.title(title)
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)
if label == None:
label = title
if kind == "bar":
width = centers[1] - centers[0]
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'")
if not label is None:
P.xlabel(xlabel)
if not ylabel is None:
P.ylabel(ylabel)
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",
)
P.ylim([None, 1.0])
if not fit is None:
self._overlay_fit(
centers, values, kind=fit, ls="--", lw=1.5, label=fitlabel
)
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,
ylog=False,
fit=None,
fitlabel=None,
smooth=0,
nml_key=None,
run=None,
runs=None,
yerr=None,
yerr_kind="std",
sigma_err=2.0,
grid=False,
put_time=False,
time_unit=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
"""
if not node_arg is None:
name_x, name_y = name_x + "_" + node_arg, name_y + "_" + node_arg
node_x = self.save.get_node(name_x)
node_y = self.save.get_node(name_y)
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)
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
)
P.xlabel(xlabel)
P.ylabel(ylabel)
if grid:
P.grid()
if ylog:
P.yscale("log")
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(time_unit), time_unit.latex
)
if label is not None and len(label) > 0:
label = label + " | " + time_str
else:
label = time_str
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]
if smooth > 0:
y = gaussian_filter1d(y, sigma=smooth)
if not run is None:
label = self._label_run(run, node_y, label, nml_key)
if colors is None:
(base_line,) = P.plot(x, y, 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(time_unit)
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:
if isinstance(yerr, str):
yerr = self.save.get_node(yerr).read()
base_line, _, _ = P.errorbar(x, y, yerr=yerr, label=label, **kwargs)
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],
)
if not run is None:
label = self._label_run(run, node_y, label, nml_key)
if yerr_kind is None:
yerr = None
(base_line,) = P.plot(x, y, label=label, **kwargs)
else:
base_line, _, _ = P.errorbar(
x, y, yerr=[y - yerr_min, yerr_max - y], label=label, **kwargs
)
else:
if runs is None:
runs = self.runs
for i, run in enumerate(runs):
x_run, y_run = node_x[run], node_y[run]
x = x_run.read() * xunit_old.express(xunit)
y = y_run.read() * yunit_old.express(yunit)
mask = np.isfinite(x) & np.isfinite(y)
x, y = x[mask], y[mask]
if smooth > 0:
y = gaussian_filter1d(y, sigma=smooth)
label_run = self._label_run(run, y_run, label, nml_key)
if colors is None:
(base_line,) = P.plot(x, y, label=label_run, **kwargs)
else:
if nml_color is None:
color = colors[i % len(colors)]
else:
nml = self.comp.get_nml(nml_color, run)
try:
color = colors[nml]
except:
color = colors(nml)
(base_line,) = P.plot(x, y, label=label_run, color=color, **kwargs)
if legend is None:
legend = True
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["fluct_" + name] = PlotRule(
self,
partial(
self._plot_map, "fluct_" + name, vmin=0.01, vmax=100, 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")
# Othoradial
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()