noe/pipeline
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Ax-aware plotting, More PDFs, improve format

Browse Source
This commit is contained in:
Noe Brucy
2020-02-19 17:51:23 +01:00
parent a87abeb52d
commit 8d7c5296cc
5 changed files with 281 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plotter.py
+133 -73
View File
@@ -49,7 +49,7 @@ class Plotter(Aggregator, BaseProcessor):
_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$ [code]", "y": r"$y$ [code]", "z": r"$z$ [code]"}
_ax_title = {"x": r"$x$", "y": r"$y$", "z": r"$z$"}
G = 1.0 # Gravitational constant
@@ -58,6 +58,7 @@ class Plotter(Aggregator, BaseProcessor):
"beta": "$\\beta$",
"beta_cool": "$\\beta_{c}$",
"dens0": "$n_0$",
"coldens0": "$\Sigma_0$",
"sfr_avg_window": "window",
"comp_frac": "$\\zeta$",
}
@@ -118,7 +119,7 @@ class Plotter(Aggregator, BaseProcessor):
or not os.path.exists(plot_filename)
)
def _process_rule(self, name, rule, arg, overwrite=False, **kwargs):
def _process_rule(self, name, rule, arg, overwrite=False, ax=None, **kwargs):
if not arg is None:
name_full = name + "_" + str(arg)
else:
@@ -126,8 +127,16 @@ class Plotter(Aggregator, BaseProcessor):
if rule.is_valid(arg):
if rule.kind == "classic":
for run in self.runs:
for i, num in enumerate(self.nums[run]):
try:
runs = kwargs.pop("runs")
except KeyError:
runs = self.runs
if ax is None:
ax = [P.subplots(1, 1)[1] for run in runs for num in self.nums[run]]
i = 0
for run in runs:
for num in self.nums[run]:
plot_filename = self._find_filename(name_full, run, num)
save = tables.open_file(self.pp[run][num].filename)
try:
@@ -137,12 +146,27 @@ class Plotter(Aggregator, BaseProcessor):
arg,
plot_filename,
overwrite,
ax=ax[i],
run=run,
**kwargs
)
except TypeError:
self._plot_rule(
rule,
save,
arg,
plot_filename,
overwrite,
ax=ax,
run=run,
**kwargs
)
finally:
save.close()
i = i + 1
else:
if ax is None:
ax = P.gca()
if rule.kind == "series" and len(self.runs) == 1:
run = self.runs[0]
plot_filename = self._find_filename(name_full, run)
@@ -151,25 +175,16 @@ class Plotter(Aggregator, BaseProcessor):
save = tables.open_file(self.comp.filename, "r")
try:
self._plot_rule(
rule,
save,
arg,
plot_filename,
overwrite,
open_figure=not self.pp_params.out.interactive,
**kwargs
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, open_figure=True, **kwargs
):
def _plot_rule(self, rule, save, arg, plot_filename, overwrite, ax, **kwargs):
P.sca(ax)
if self._needs_computation(overwrite, plot_filename):
if open_figure:
P.figure()
rule.plot(save, arg, **kwargs)
P.tight_layout(pad=1)
if not self.pp_params.out.interactive:
@@ -283,19 +298,26 @@ class Plotter(Aggregator, BaseProcessor):
unit=None,
unit_coeff=1.0,
overlays=[],
overlays_kwargs=[],
title=None,
nml_key=None,
put_time=True,
time_unit=cst.Myr,
unit_space=cst.pc,
cmap="plasma",
norm="log",
put_cbar=True,
autoscale=True,
**kwargs
):
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()
@@ -319,10 +341,13 @@ class Plotter(Aggregator, BaseProcessor):
P.locator_params(axis=ax_h, nbins=self.pp_params.plot.ntick)
P.locator_params(axis=ax_v, nbins=self.pp_params.plot.ntick)
P.xlabel(self._ax_title[ax_h])
P.ylabel(self._ax_title[ax_v])
P.xlabel(self._ax_title[ax_h] + unit_str(unit_space))
P.ylabel(self._ax_title[ax_v] + unit_str(unit_space))
cbar = P.colorbar(im)
try:
cbar = P.colorbar(im, cax=P.gca().cax)
except AttributeError:
cbar = P.colorbar()
if not label is None:
cbar.set_label(label)
@@ -331,7 +356,7 @@ class Plotter(Aggregator, BaseProcessor):
if put_time:
time = self.save.root._v_attrs.time * self.comp.info["unit_time"]
time_str = "time = {:.6g} {}".format(
time_str = self.pp_params.plot.time_fmt.format(
time.express(time_unit), time_unit.latex
)
if len(title) > 0:
@@ -341,8 +366,11 @@ class Plotter(Aggregator, BaseProcessor):
P.title(title)
for plot_overlay in overlays:
plot_overlay(ax_los)
for i, plot_overlay in enumerate(overlays):
try:
plot_overlay(ax_los, **overlays_kwargs[i])
except:
plot_overlay(ax_los)
def _overlay_levels(self, ax_los):
map_level = self.save.get_node("/maps/{}_{}".format("levels", ax_los)).read()
@@ -367,7 +395,7 @@ class Plotter(Aggregator, BaseProcessor):
P.clabel(cont, cont.levels[cont.levels < 11], inline=1, fontsize=8.0, fmt="%1d")
def _overlay_speed(self, ax_los, unit=cst.km_s, unit_coeff=1.0):
def _overlay_speed(self, ax_los, unit=cst.km_s, unit_coeff=1.0, key_v=None):
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))
@@ -397,7 +425,8 @@ class Plotter(Aggregator, BaseProcessor):
Q = P.quiver(hh, vv, map_vh_red, map_vv_red, units="width", color="grey")
label, unit_old, unit = self._ax_label_unit(dmap_vh_node, "", unit, unit_coeff)
key_v = (max_v + min_v) / 2.0
if key_v is None:
key_v = (max_v + min_v) / 2.0
P.quiverkey(
Q,
0.6,
@@ -438,24 +467,35 @@ class Plotter(Aggregator, BaseProcessor):
nml_key=None,
put_time=True,
time_unit=cst.Myr,
xlog=None,
ylog=False,
kind="bar",
colors=None,
nml_color=None,
**kwargs
):
node = self.save.get_node("/hist/" + name)
label, unit_old, unit = self._ax_label_unit(node, label, unit, unit_coeff)
values, centers = node.read() * unit_old.express(unit)
width = centers[1] - centers[0]
if xlog is None:
try:
xlog = node._v_attrs_.logbins
except:
xlog = False
P.bar(centers, values, width, log=ylog, **kwargs)
P.grid()
label, unit_old, unit = self._ax_label_unit(node, label, unit, unit_coeff)
values, centers = node.read()
if xlog:
centers = centers + np.log10(unit_old.express(unit))
else:
centers = centers * unit_old.express(unit)
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 = "time = {:.6g} {}".format(
time_str = self.pp_params.out.time_fmt.format(
time.express(time_unit), time_unit.latex
)
if len(title) > 0:
@@ -465,10 +505,32 @@ class Plotter(Aggregator, BaseProcessor):
P.title(title)
color = None
if 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 kind == "bar":
width = centers[1] - centers[0]
P.bar(centers, values, width, log=ylog, color=color, label=title, **kwargs)
elif kind == "step":
if ylog:
P.yscale("log")
P.step(centers, values, where="mid", color=color, label=title, **kwargs)
else:
raise ValueError("kind must be 'bar' or 'step'")
P.grid()
if not label is None:
P.xlabel(label)
if "/hist/fit_" + name + "_" + ax_los in self.save:
if not ax_los is None and "/hist/fit_" + name + "_" + ax_los in self.save:
slope = node.attrs.slope
origin = node.attrs.origin
P.plot(
@@ -495,10 +557,10 @@ class Plotter(Aggregator, BaseProcessor):
fit=None,
fitlabel=None,
smooth=0,
sigma_err=2.0,
nml_key=None,
runs=None,
yerr_kind="std",
sigma_err=2.0,
colors=None,
nml_color=None,
**kwargs
@@ -530,43 +592,33 @@ class Plotter(Aggregator, BaseProcessor):
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":
yerr = node_y.std.read() * yunit_old.express(yunit) * sigma_err
mask = np.isfinite(x) & np.isfinite(y) & np.isfinite(yerr)
x, y, yerr = x[mask], y[mask], yerr[mask]
if smooth > 0:
y = gaussian_filter1d(y, sigma=smooth)
base_line, _, _ = P.errorbar(x, y, yerr=yerr, label=label, **kwargs)
elif yerr_kind in ["min_max", "95per"]:
if 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)
yerr = yerr_max - yerr_min
mask = (
np.isfinite(x)
& np.isfinite(y)
& np.isfinite(yerr_min)
& np.isfinite(yerr_max)
)
x, y, yerr, yerr_min, yerr_max = (
x[mask],
y[mask],
yerr[mask],
yerr_min[mask],
yerr_max[mask],
)
base_line, _, _ = P.errorbar(
x, y, yerr=[y - yerr_min, yerr_max - y], label=label, **kwargs
)
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)
else:
mask = np.isfinite(y)
x, y = x[mask], y[mask]
if smooth > 0:
y = gaussian_filter1d(y, sigma=smooth)
(base_line,) = P.plot(x, y, "*", **kwargs)
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],
)
base_line, _, _ = P.errorbar(
x, y, yerr=[y - yerr_min, yerr_max - y], label=label, **kwargs
)
else:
if runs is None:
runs = self.runs
@@ -584,10 +636,12 @@ class Plotter(Aggregator, BaseProcessor):
else:
if nml_color is None:
color = colors[i % len(colors)]
(base_line,) = P.plot(x, y, label=label_run, **kwargs)
else:
nml = self.comp.get_nml(nml_color, run)
color = colors[nml]
try:
color = colors[nml]
except:
color = colors(nml)
(base_line,) = P.plot(x, y, label=label_run, color=color, **kwargs)
P.legend()
@@ -632,8 +686,8 @@ class Plotter(Aggregator, BaseProcessor):
if yerr is None:
(a, b, rho, _map_rule, stderr) = linregress(np.log10(x), np.log10(y))
self._log(
"Power law fit y = x^({}) * 10^({}) with R^2 = {} and error is {}".format(
a, b, rho, stderr
"Power law fit y = x^({}) * {} with R^2 = {} and error is {}".format(
a, 10 ** b, rho, stderr
)
)
else:
@@ -651,8 +705,8 @@ class Plotter(Aggregator, BaseProcessor):
b, a = c[0], c[1]
residual = errfunc(c, np.log10(x), np.log10(y), yerr / y)
self._log(
"Power law fit y = x^({}) * 10^({}) with residual {}".format(
a, b, residual
"Power law fit y = x^({}) * {} with residual {}".format(
a, 10 ** b, residual
)
)
if label is None:
@@ -746,6 +800,12 @@ class Plotter(Aggregator, BaseProcessor):
"$\rho$-PDF",
dependencies=["rho_pdf"],
),
"T_pdf": PlotRule(
self, partial(self._plot_hist, "T_pdf"), "T-PDF", dependencies=["T_pdf"]
),
"P_pdf": PlotRule(
self, partial(self._plot_hist, "P_pdf"), "P-PDF", dependencies=["P_pdf"]
),
}
averageables = ["coldens", "rho", "T", "Q"]