diff --git a/python/energy_flow_proper/figsaver.py b/python/energy_flow_proper/figsaver.py index a014bcd..25d9a1e 100644 --- a/python/energy_flow_proper/figsaver.py +++ b/python/energy_flow_proper/figsaver.py @@ -2,34 +2,48 @@ import matplotlib import matplotlib.pyplot as plt import os from pathlib import Path -import tikzplotlib + from contextlib import contextmanager import numpy as np +try: + import tikzplotlib + +except: + tikzplotlib = None # dirty + fig_path = Path(os.getcwd()) / "figures" val_path = Path(os.getcwd()) / "values" -def export_fig(name, fig=None, aspect="golden", **kwargs): +def export_fig(name, fig=None, x_scaling=1, y_scaling=0.4, tikz=True, **kwargs): fig_path.mkdir(parents=True, exist_ok=True) if fig is None: fig = plt.gcf() + w, _ = fig.get_size_inches() + fig.set_size_inches((w * x_scaling, w * y_scaling)) + + if not fig.get_constrained_layout(): + fig.tight_layout() + fig.canvas.draw() # fig.savefig(fig_path / f"{name}.pdf") # fig.savefig(fig_path / f"{name}.svg") # fig.savefig(fig_path / f"{name}.pgf") + fig.savefig(fig_path / f"{name}.pdf") - tikzplotlib.clean_figure() - fig.tight_layout() - tikzplotlib.save( - fig_path / f"{name}.tex", - figure=fig, - axis_width="\\figW", - axis_height="\\figH", - **kwargs, - ) + if tikz and tikzplotlib: + tikzplotlib.clean_figure() + fig.tight_layout() + tikzplotlib.save( + fig_path / f"{name}.tex", + figure=fig, + axis_width="\\figW", + axis_height="\\figH", + **kwargs, + ) def scientific_round(val, *err, retprec=False): @@ -120,11 +134,22 @@ def tex_value(val, err=None, unit=None, prefix="", suffix="", prec=0, save=None) ############################################################################### MPL_RC = { - "lines.linewidth": 1.5, + "lines.linewidth": 0.5, "text.color": "black", "axes.labelcolor": "black", "xtick.color": "black", "ytick.color": "black", + "figure.figsize": (5.2, 3.2), + "text.usetex": True, + "font.family": "serif", + "font.serif": ["Palatino"], + "font.size": 7, + "axes.labelsize": 7, + "axes.titlesize": 7, + "legend.fontsize": 7, + "xtick.labelsize": 5, + "ytick.labelsize": 5, + "figure.constrained_layout.use": True, } MPL_RC_POSTER = { diff --git a/python/energy_flow_proper/plot_utils.py b/python/energy_flow_proper/plot_utils.py index 2a082bc..6a8c84d 100644 --- a/python/energy_flow_proper/plot_utils.py +++ b/python/energy_flow_proper/plot_utils.py @@ -10,6 +10,7 @@ from hops.util.utilities import ( entropy, trace_distance, ) +from matplotlib.gridspec import GridSpec try: import hiro_models.model_auxiliary as aux @@ -306,35 +307,41 @@ def plot_diff_vs_sigma( def plot_interaction_consistency( - models, reference=None, label_fn=lambda model: f"$ω_c={model.ω_c:.2f}$", **kwargs + models, + reference=None, + label_fn=lambda model: f"$ω_c={model.ω_c:.2f}$", + inset=None, + bath=0, + **kwargs, ): - fig, ax = plt.subplots() + fig, (ax, ax2) = plt.subplots(ncols=2) + + ax3 = None + if inset: + slc, bounds = inset + ax3 = ax.inset_axes(bounds) + if reference: with aux.get_data(reference) as data: reference_energy = reference.interaction_energy(data, **kwargs) for model in models: with aux.get_data(model) as data: - energy = model.interaction_energy(data, **kwargs) - interaction_ref = model.interaction_energy_from_conservation(data, **kwargs) - diff = abs(interaction_ref - energy) + energy = model.interaction_energy(data, **kwargs).for_bath(bath) + interaction_ref = model.interaction_energy_from_conservation( + data, **kwargs + ).for_bath(bath) - self_consistency = (diff.value < diff.σ).sum() / len(diff.value[0]) * 100 + self_consistency = energy.consistency(interaction_ref) + normalizer = 1 / abs(energy.value).max() if reference: - diff = abs(interaction_ref - reference_energy) - final_consistency = ( - (diff.value < diff.σ).sum() / len(diff.value[0]) * 100 - ) + final_consistency = reference_energy.consistency(interaction_ref) _, _, (line, _) = plot_with_σ( data.time[:], energy, ax=ax, - label=label_fn(model) - + fr", (${self_consistency:.0f}\%$" - + (fr", ${final_consistency:.0f}\%$)" if reference else ")"), - bath=0, ) plot_with_σ( @@ -342,21 +349,58 @@ def plot_interaction_consistency( interaction_ref, ax=ax, linestyle="--", - bath=0, color=lighten_color(line[0].get_color(), 0.8), ) - ax.set_xlabel("$τ$") - ax.set_ylabel(r"$\langle H_\mathrm{I}\rangle$") - ax.legend() + plot_with_σ( + data.time[:], + (energy - interaction_ref) * normalizer, + label=label_fn(model) + + fr", (${self_consistency:.0f}\%$" + + (fr", ${final_consistency:.0f}\%$)" if reference else ")"), + ax=ax2, + color=line[0].get_color(), + ) - return fig, ax + if inset: + plot_with_σ( + data.time[slc], + energy.slice(slc), + ax=ax3, + color=line[0].get_color(), + ) + + plot_with_σ( + data.time[slc], + interaction_ref.slice(slc), + ax=ax3, + linestyle="--", + color=lighten_color(line[0].get_color(), 0.8), + ) + + ax.set_xlabel(r"$\tau$") + ax2.set_xlabel(r"$\tau$") + ax.set_ylabel(r"$\langle H_\mathrm{I}\rangle$") + ax2.set_ylabel( + r"$\langle (H_\mathrm{I}\rangle - \langle H_\mathrm{I}\rangle_\mathrm{ref}) / \langle H_\mathrm{I}\rangle_\mathrm{max}$" + ) + ax2.legend() + + return fig, (ax, ax2, ax3) def plot_interaction_consistency_development( - models, reference=None, label_fn=lambda model: f"$ω_c={model.ω_c:.2f}$", **kwargs + models, + reference=None, + label_fn=lambda model: rf"$\omega_c={model.ω_c:.2f}$", + **kwargs, ): - fig, (ax, ax2, ax3) = plt.subplots(nrows=1, ncols=3) + fig = plt.figure(constrained_layout=True) + gs = GridSpec(2, 2, figure=fig) + + ax = fig.add_subplot(gs[0, :]) + ax2 = fig.add_subplot(gs[1, 0]) + ax3 = fig.add_subplot(gs[1, 1]) ax2.set_xscale("log") ax2.set_yscale("log") @@ -376,23 +420,25 @@ def plot_interaction_consistency_development( else: energy = model.interaction_energy(data, **kwargs) - diff = abs(interaction_ref - energy) + normalizer = 1 / abs(energy.value).max() + diff = abs(interaction_ref - energy) * normalizer - ns, values, σs = [], [], [] + ns, values, σs, diffs = [], [], [], [] for N, val, σ in diff.aggregate_iterator: ns.append(N) values.append((val < σ).sum() / len(val[0]) * 100) - σs.append(σ.max()) + σs.append(σ.mean() * normalizer) + diffs.append(val.mean() * normalizer) σ_ref, σ_int = [], [] for _, _, σ in interaction_ref.aggregate_iterator: - σ_ref.append(σ.max()) + σ_ref.append(σ.mean() * normalizer) for _, _, σ in energy.aggregate_iterator: - σ_int.append(σ.max()) + σ_int.append(σ.mean() * normalizer) - ax.plot( + lines = ax.plot( ns, values, linestyle="--", @@ -400,31 +446,64 @@ def plot_interaction_consistency_development( markersize=2, label=label_fn(model), ) - ax2.plot(ns, σs, label=label_fn(model)) + ax2.plot( + ns, + σs, + label=label_fn(model), + color=lines[0].get_color(), + ) + + ax2.plot( + ns, + diffs, + label=label_fn(model), + linestyle="--", + color=lines[0].get_color(), + ) + ax3.plot( ns, σ_ref, label=label_fn(model) + " (from conservation)", linestyle="--", + color=lines[0].get_color(), ) ax3.plot(ns, σ_int, label=label_fn(model) + " (direct)") ax.axhline(68, linestyle="-.", color="grey", alpha=0.5) ax.set_xlabel("$N$") + ax.set_title("Consistency") ax2.set_xlabel("$N$") ax3.set_xlabel("$N$") - ax2.set_ylabel("Maximum Allowed Deviation") - ax3.set_ylabel("Maximum $σ$") + ax2.set_ylabel(r"$[|\langle H_I\rangle|_\mathrm{max}]$") + ax2.set_title(r"Mean Actual vs. Allowed deviation.") + ax3.set_ylabel(r"$[|\langle H_I\rangle|_\mathrm{max}]$") + ax3.set_title(r"Statistical Error of $\langle H_I\rangle$.") ax.set_ylabel(("" if reference else "Self-") + r"Consistency [$\%$]") ax.legend() - ax2.legend() - ax3.legend() + + straight_line = matplotlib.lines.Line2D( + [], [], color="black", label=r"$\langle \sigma_\Delta\rangle$" + ) + dashed_line = matplotlib.lines.Line2D( + [], [], color="black", label=r"$\langle\Delta\rangle$", linestyle="--" + ) + ax2.legend(handles=[straight_line, dashed_line]) + + straight_line = matplotlib.lines.Line2D([], [], color="black", label=r"Direct") + dashed_line = matplotlib.lines.Line2D( + [], [], color="black", label=r"Energy Conservation", linestyle="--" + ) + ax3.legend(handles=[dashed_line, straight_line]) return fig, [ax, ax2, ax3] -def plot_consistency_development(value, reference): - fig, (ax, ax2) = plt.subplots(nrows=2, sharex=True) +def plot_consistency_development(value, reference, fig=None): + if fig: + (ax, ax2) = fig.subplots(nrows=2, sharex=True) + else: + fig, (ax, ax2) = (fig or plt).subplots(nrows=2, sharex=True) ax.set_xscale("log") ax2.set_xscale("log") @@ -466,14 +545,17 @@ def plot_consistency_development(value, reference): markersize=5, ) - ax2.plot(ns, np.array(σs) / abs(reference).max(), label=r"$\langle σ\rangle$") - ax2.plot(ns, np.array(max_diff) / abs(reference).max(), label=r"$\langle Δ\rangle$") + ax2.plot(ns, np.array(σs) / abs(reference).max(), label=r"$\langle \sigma\rangle$") + ax2.plot( + ns, np.array(max_diff) / abs(reference).max(), label=r"$\langle \Delta\rangle$" + ) ax.axhline(68, linestyle="-.", color="grey", alpha=0.5) - ax.set_xlabel("$N$") ax2.set_xlabel("$N$") - ax.set_ylabel(r"Consistency [$\%$]") - ax2.set_ylabel(r"$\mathrm{Deviation}/{J_\mathrm{max}}$") + ax.set_ylabel(r"[$\%$]") + ax.set_title(r"Consistency") + ax2.set_ylabel(r"$[{J_\mathrm{max}}]$") + ax2.set_title(r"Deviation") ax2.legend() @@ -589,13 +671,80 @@ def plot_distance_measures(model, strobe_indices, ax=None): @wrap_plot def plot_σ_development(ensemble_value, ax=None, **kwargs): + # norm = abs(ensemble_value.value).max() ax.plot( ensemble_value.Ns, [σ.mean() for σ in ensemble_value.σs], marker=".", + markersize=6, linestyle="dotted", **kwargs, ) - ax.set_ylabel("mean $σ$") + ax.set_ylabel("mean normalized $σ$") ax.set_xlabel("$N$") + + +def plot_multi_energy_overview( + models, + label_fn=lambda m: fr"$\omega_c={m.ω_c:.1f},\,\alpha(0)={m.bcf(0).real:.2f}$", +): + fig, ((ax_sys, ax_flow), (ax_inter, ax_bath)) = plt.subplots( + nrows=2, ncols=2, sharex=True + ) + + flow_ins = ax_flow.inset_axes([0.4, 0.20, 0.5, 0.55]) + ax_sys.set_ylabel(r"$\langle H_\mathrm{S}\rangle$") + ax_flow.set_ylabel("$J$") + ax_inter.set_ylabel(r"$\langle H_\mathrm{I}\rangle$") + ax_bath.set_ylabel(r"$\langle H_\mathrm{B}\rangle$") + + ax_sys.set_xlabel(r"$\tau$") + ax_flow.set_xlabel(r"$\tau$") + ax_inter.set_xlabel(r"$\tau$") + ax_bath.set_xlabel(r"$\tau$") + + handles = [] + for model in models: + with aux.get_data(model) as data: + sys_energy = model.system_energy(data, **ensemble_arg) + flow = model.bath_energy_flow(data, **ensemble_arg).for_bath(0) + inter = model.interaction_energy(data, **ensemble_arg).for_bath(0) + bath = model.bath_energy(data, **ensemble_arg).for_bath(0) + + _, _, (lines, _) = plot_with_σ( + data.time[:], + sys_energy, + ax=ax_sys, + label=label_fn(model), + ) + handles.append(lines[0]) + + plot_with_σ( + data.time[:], + flow, + ax=ax_flow, + ) + + mask = np.logical_and(7 <= data.time[:], data.time[:] <= 9) + plot_with_σ( + data.time[mask], + flow.slice(mask), + ax=flow_ins, + ) + + plot_with_σ( + data.time[:], + inter, + ax=ax_inter, + ) + + plot_with_σ( + data.time[:], + bath, + ax=ax_bath, + ) + + ax_sys.legend( + # ncol=2, + )