mirror of
https://github.com/vale981/master-thesis
synced 2025-03-06 10:31:37 -05:00
601 lines
16 KiB
Python
601 lines
16 KiB
Python
import matplotlib
|
||
import matplotlib.pyplot as plt
|
||
import numpy as np
|
||
from contextlib import contextmanager
|
||
import utilities as ut
|
||
from hopsflow.util import EnsembleValue
|
||
from hops.util.utilities import (
|
||
relative_entropy,
|
||
relative_entropy_single,
|
||
entropy,
|
||
trace_distance,
|
||
)
|
||
|
||
try:
|
||
import hiro_models.model_auxiliary as aux
|
||
except:
|
||
aux = None
|
||
|
||
|
||
def cm2inch(*tupl):
|
||
inch = 2.54
|
||
if isinstance(tupl[0], tuple):
|
||
return tuple(i / inch for i in tupl[0])
|
||
else:
|
||
return tuple(i / inch for i in tupl)
|
||
|
||
|
||
###############################################################################
|
||
# Plot Porn #
|
||
###############################################################################
|
||
|
||
|
||
def wrap_plot(f):
|
||
def wrapped(*args, ax=None, setup_function=plt.subplots, **kwargs):
|
||
fig = None
|
||
if not ax:
|
||
fig, ax = setup_function()
|
||
|
||
ret_val = f(*args, ax=ax, **kwargs)
|
||
return (fig, ax, ret_val) if ret_val else (fig, ax)
|
||
|
||
return wrapped
|
||
|
||
|
||
def get_figsize(width="thesis", fraction=1, subplots=(1, 1)):
|
||
"""Set figure dimensions to avoid scaling in LaTeX.
|
||
|
||
Parameters
|
||
----------
|
||
width: float or string
|
||
Document width in points, or string of predined document type
|
||
fraction: float, optional
|
||
Fraction of the width which you wish the figure to occupy
|
||
subplots: array-like, optional
|
||
The number of rows and columns of subplots.
|
||
Returns
|
||
-------
|
||
fig_dim: tuple
|
||
Dimensions of figure in inches
|
||
"""
|
||
if width == "thesis":
|
||
width_pt = 330.62111
|
||
elif width == "poster":
|
||
width_pt = 957.13617
|
||
elif width == "beamer":
|
||
width_pt = 307.28987
|
||
else:
|
||
width_pt = width
|
||
|
||
# Width of figure (in pts)
|
||
fig_width_pt = width_pt * fraction
|
||
# Convert from pt to inches
|
||
inches_per_pt = 1 / 72.27
|
||
|
||
# Golden ratio to set aesthetic figure height
|
||
# https://disq.us/p/2940ij3
|
||
golden_ratio = (5 ** 0.5 - 1) / 2
|
||
|
||
# Figure width in inches
|
||
fig_width_in = fig_width_pt * inches_per_pt
|
||
# Figure height in inches
|
||
fig_height_in = fig_width_in * golden_ratio * (subplots[0] / subplots[1])
|
||
|
||
return (fig_width_in, fig_height_in)
|
||
|
||
|
||
# def get_figsize(
|
||
# columnwidth: float,
|
||
# wf: float = 0.5,
|
||
# hf: float = (5.0 ** 0.5 - 1.0) / 2.0,
|
||
# ) -> tuple[float, float]:
|
||
# """
|
||
# :param wf: Width fraction in columnwidth units.
|
||
# :param hf: Weight fraction in columnwidth units.
|
||
|
||
# Set by default to golden ratio.
|
||
|
||
# :param columnwidth: width of the column in latex.
|
||
|
||
# Get this from LaTeX using \showthe\columnwidth
|
||
|
||
# :returns: The ``[fig_width,fig_height]`` that should be given to
|
||
# matplotlib.
|
||
# """
|
||
# fig_width_pt = columnwidth * wf
|
||
# inches_per_pt = 1.0 / 72.27 # Convert pt to inch
|
||
# fig_width = fig_width_pt * inches_per_pt # width in inches
|
||
# fig_height = fig_width * hf # height in inches
|
||
# return [fig_width, fig_height]
|
||
|
||
|
||
@wrap_plot
|
||
def plot_complex(x, y, *args, ax=None, label="", absolute=False, **kwargs):
|
||
label = label + ", " if (len(label) > 0) else ""
|
||
ax.plot(
|
||
x,
|
||
abs(y.real) if absolute else y.real,
|
||
*args,
|
||
label=f"{label}absolute real part",
|
||
**kwargs,
|
||
)
|
||
ax.plot(
|
||
x,
|
||
abs(y.imag) if absolute else y.imag,
|
||
*args,
|
||
label=f"{label}absolute imag part",
|
||
**kwargs,
|
||
)
|
||
ax.legend()
|
||
|
||
|
||
@wrap_plot
|
||
def plot_convergence(
|
||
x,
|
||
y,
|
||
reference=None,
|
||
ax=None,
|
||
label="",
|
||
slice=None,
|
||
linestyle="-",
|
||
bath=None,
|
||
):
|
||
if bath is not None:
|
||
y = y.for_bath(bath)
|
||
|
||
label = label + ", " if (len(label) > 0) else ""
|
||
slice = (0, -1) if not slice else slice
|
||
y_final = y[-1]
|
||
|
||
for i in range(len(y) - 1 if reference is None else len(y)):
|
||
current_value = y[i]
|
||
consistency = current_value.consistency(
|
||
y_final if reference is None else reference
|
||
)
|
||
|
||
line = ax.plot(
|
||
x,
|
||
current_value.value,
|
||
label=f"{label}$N={current_value.N}$ $({consistency:.1f}\%)$",
|
||
alpha=current_value.N / y.N,
|
||
linestyle=linestyle if i == (len(y) - 1) else ":",
|
||
)
|
||
|
||
if reference is None:
|
||
plot_with_σ(
|
||
x,
|
||
y_final,
|
||
label=f"{label}$N={y.N}$",
|
||
linestyle=linestyle,
|
||
ax=ax,
|
||
)
|
||
else:
|
||
ax.fill_between(
|
||
x,
|
||
y_final.value - y_final.σ,
|
||
y_final.value + y_final.σ,
|
||
color=lighten_color(line[0].get_color(), 0.5),
|
||
)
|
||
|
||
return None
|
||
|
||
|
||
def lighten_color(color, amount=0.5):
|
||
"""
|
||
Lightens the given color by multiplying (1-luminosity) by the given amount.
|
||
Input can be matplotlib color string, hex string, or RGB tuple.
|
||
|
||
Examples:
|
||
>> lighten_color('g', 0.3)
|
||
>> lighten_color('#F034A3', 0.6)
|
||
>> lighten_color((.3,.55,.1), 0.5)
|
||
"""
|
||
import matplotlib.colors as mc
|
||
import colorsys
|
||
|
||
try:
|
||
c = mc.cnames[color]
|
||
except:
|
||
c = color
|
||
c = colorsys.rgb_to_hls(*mc.to_rgb(c))
|
||
return colorsys.hls_to_rgb(c[0], 1 - amount * (1 - c[1]), c[2])
|
||
|
||
|
||
def fancy_error(x, y, err, ax, **kwargs):
|
||
line = ax.plot(
|
||
x,
|
||
y,
|
||
**kwargs,
|
||
)
|
||
|
||
err = ax.fill_between(
|
||
x,
|
||
y + err,
|
||
y - err,
|
||
color=lighten_color(line[0].get_color(), 0.5),
|
||
alpha=0.5,
|
||
)
|
||
|
||
return line, err
|
||
|
||
|
||
@wrap_plot
|
||
def plot_with_σ(
|
||
x,
|
||
y,
|
||
ax=None,
|
||
transform=lambda y: y,
|
||
bath=None,
|
||
strobe_frequency=None,
|
||
strobe_tolerance=1e-3,
|
||
**kwargs,
|
||
):
|
||
err = (y.σ[bath] if bath is not None else y.σ).real
|
||
y_final = transform(y.value[bath] if bath is not None else y.value)
|
||
|
||
strobe_mode = strobe_frequency is not None
|
||
strobe_indices = None
|
||
strobe_times = None
|
||
strobe_style = dict(linestyle="none", marker="o", markersize=2) | kwargs
|
||
|
||
if strobe_mode:
|
||
strobe_times, strobe_indices = ut.strobe_times(
|
||
x, strobe_frequency, strobe_tolerance
|
||
)
|
||
|
||
return ax.errorbar(
|
||
strobe_times,
|
||
y_final[strobe_indices],
|
||
err[strobe_indices],
|
||
**strobe_style,
|
||
)
|
||
|
||
return fancy_error(x, y_final, err, ax=ax, **kwargs)
|
||
|
||
|
||
@wrap_plot
|
||
def plot_diff_vs_sigma(
|
||
x,
|
||
y,
|
||
reference,
|
||
ax=None,
|
||
label="",
|
||
transform=lambda y: y,
|
||
ecolor="yellow",
|
||
ealpha=0.5,
|
||
ylabel=None,
|
||
bath=None,
|
||
):
|
||
label = label + ", " if (len(label) > 0) else ""
|
||
if bath is not None:
|
||
y = y.for_bath(bath)
|
||
reference = reference.for_bath(bath)
|
||
|
||
ax.fill_between(
|
||
x,
|
||
0,
|
||
reference.σ,
|
||
color=ecolor,
|
||
alpha=ealpha,
|
||
label=rf"{label}$\sigma\, (N={reference.N})$",
|
||
)
|
||
|
||
for i in range(len(y)):
|
||
current = y[i]
|
||
|
||
not_last = current.N < y[-1].N
|
||
consistency = current.consistency(reference)
|
||
diff = abs(current - reference)
|
||
|
||
ax.plot(
|
||
x,
|
||
diff.value,
|
||
label=rf"{label}$N={current.N}$ $({consistency:.1f}\%)$",
|
||
alpha=consistency / 100 if not_last else 1,
|
||
linestyle=":" if not_last else "-",
|
||
color=None if not_last else "red",
|
||
)
|
||
|
||
if ylabel is not None:
|
||
if ylabel[0] == "$":
|
||
ylabel = ylabel[1:-1]
|
||
else:
|
||
ylabel = rf"\text{{ {ylabel} }}"
|
||
|
||
ax.set_ylabel(rf"$|{{{ylabel}}}_{{\mathrm{{ref}}}}-{{{ylabel}}}_{{N_i}}|$")
|
||
|
||
|
||
def plot_interaction_consistency(
|
||
models, reference=None, label_fn=lambda model: f"$ω_c={model.ω_c:.2f}$", **kwargs
|
||
):
|
||
fig, ax = plt.subplots()
|
||
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)
|
||
|
||
self_consistency = (diff.value < diff.σ).sum() / len(diff.value[0]) * 100
|
||
|
||
if reference:
|
||
diff = abs(interaction_ref - reference_energy)
|
||
final_consistency = (
|
||
(diff.value < diff.σ).sum() / len(diff.value[0]) * 100
|
||
)
|
||
|
||
_, _, (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_σ(
|
||
data.time[:],
|
||
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()
|
||
|
||
return fig, ax
|
||
|
||
|
||
def plot_interaction_consistency_development(
|
||
models, reference=None, label_fn=lambda model: f"$ω_c={model.ω_c:.2f}$", **kwargs
|
||
):
|
||
fig, (ax, ax2, ax3) = plt.subplots(nrows=1, ncols=3)
|
||
|
||
ax2.set_xscale("log")
|
||
ax2.set_yscale("log")
|
||
|
||
ax3.set_xscale("log")
|
||
ax3.set_yscale("log")
|
||
|
||
if reference is not None:
|
||
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:
|
||
interaction_ref = model.interaction_energy_from_conservation(data, **kwargs)
|
||
if reference is not None:
|
||
energy = reference_energy
|
||
else:
|
||
energy = model.interaction_energy(data, **kwargs)
|
||
|
||
diff = abs(interaction_ref - energy)
|
||
|
||
ns, values, σs = [], [], []
|
||
for N, val, σ in diff.aggregate_iterator:
|
||
ns.append(N)
|
||
values.append((val < σ).sum() / len(val[0]) * 100)
|
||
σs.append(σ.max())
|
||
|
||
σ_ref, σ_int = [], []
|
||
|
||
for _, _, σ in interaction_ref.aggregate_iterator:
|
||
σ_ref.append(σ.max())
|
||
|
||
for _, _, σ in energy.aggregate_iterator:
|
||
σ_int.append(σ.max())
|
||
|
||
ax.plot(
|
||
ns,
|
||
values,
|
||
linestyle="--",
|
||
marker=".",
|
||
markersize=2,
|
||
label=label_fn(model),
|
||
)
|
||
ax2.plot(ns, σs, label=label_fn(model))
|
||
ax3.plot(
|
||
ns,
|
||
σ_ref,
|
||
label=label_fn(model) + " (from conservation)",
|
||
linestyle="--",
|
||
)
|
||
ax3.plot(ns, σ_int, label=label_fn(model) + " (direct)")
|
||
|
||
ax.axhline(68, linestyle="-.", color="grey", alpha=0.5)
|
||
ax.set_xlabel("$N$")
|
||
ax2.set_xlabel("$N$")
|
||
ax3.set_xlabel("$N$")
|
||
ax2.set_ylabel("Maximum Allowed Deviation")
|
||
ax3.set_ylabel("Maximum $σ$")
|
||
ax.set_ylabel(("" if reference else "Self-") + r"Consistency [$\%$]")
|
||
ax.legend()
|
||
ax2.legend()
|
||
ax3.legend()
|
||
|
||
return fig, [ax, ax2, ax3]
|
||
|
||
|
||
def plot_consistency_development(value, reference):
|
||
fig, (ax, ax2) = plt.subplots(nrows=2, sharex=True)
|
||
|
||
ax.set_xscale("log")
|
||
ax2.set_xscale("log")
|
||
ax2.set_yscale("log")
|
||
|
||
diff = abs(value - reference)
|
||
|
||
ns, values, σs, max_diff = [], [], [], []
|
||
|
||
for N, val, σ in diff.aggregate_iterator:
|
||
ns.append(N)
|
||
values.append(((val < σ).sum() / len(val[0])) * 100)
|
||
σs.append(σ.mean())
|
||
max_diff.append(val.mean())
|
||
|
||
values = np.array(values)
|
||
ns = np.array(ns)
|
||
where_consistent = values >= 68
|
||
|
||
values_consistent = values.copy()
|
||
values_consistent[~where_consistent] = np.nan
|
||
|
||
values_inconsistent = values.copy()
|
||
values_inconsistent[where_consistent] = np.nan
|
||
|
||
ax.plot(
|
||
ns,
|
||
values_consistent,
|
||
linestyle="none",
|
||
marker=".",
|
||
markersize=5,
|
||
)
|
||
|
||
ax.plot(
|
||
ns,
|
||
values_inconsistent,
|
||
linestyle="none",
|
||
marker=".",
|
||
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$")
|
||
|
||
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}}$")
|
||
|
||
ax2.legend()
|
||
|
||
return fig, [ax, ax2]
|
||
|
||
|
||
def plot_flow_bcf(models, label_fn=lambda model: f"$ω_c={model.ω_c:.2f}$", **kwargs):
|
||
fig, ax = plt.subplots()
|
||
for model in models:
|
||
with aux.get_data(model) as data:
|
||
flow = model.bath_energy_flow(data, **kwargs)
|
||
_, _, (line, _) = plot_with_σ(
|
||
data.time[:],
|
||
flow,
|
||
ax=ax,
|
||
label=label_fn(model),
|
||
bath=0,
|
||
transform=lambda y: -y,
|
||
)
|
||
|
||
ax.plot(
|
||
data.time[:],
|
||
-model.L_expect * model.bcf_scale * model.bcf(data.time[:]).imag,
|
||
linestyle="--",
|
||
color=line[0].get_color(),
|
||
)
|
||
|
||
return fig, ax
|
||
|
||
|
||
def plot_energy_overview(
|
||
model,
|
||
ensemble_args=None,
|
||
# system=True,
|
||
# bath=True,
|
||
# total=True,
|
||
# flow=True,
|
||
# interaction=True,
|
||
**kwargs,
|
||
):
|
||
if not ensemble_args:
|
||
ensemble_args = {}
|
||
|
||
fig, ax = plt.subplots()
|
||
ax.set_ylabel("Energy")
|
||
ax.set_xlabel(r"$\tau$")
|
||
|
||
with aux.get_data(model) as data:
|
||
system_energy = model.system_energy(data, **ensemble_args)
|
||
bath_energy = model.bath_energy(data, **ensemble_args)
|
||
interaction_energy = model.interaction_energy(data, **ensemble_args)
|
||
flow = model.bath_energy_flow(data, **ensemble_args)
|
||
|
||
plot_with_σ(model.t, system_energy, ax=ax, label="System", **kwargs)
|
||
|
||
num_baths = flow.num_baths
|
||
for bath in range(num_baths):
|
||
plot_with_σ(
|
||
model.t, flow, bath=bath, ax=ax, label=f"Flow {bath+1}", **kwargs
|
||
)
|
||
plot_with_σ(
|
||
model.t, bath_energy, bath=bath, ax=ax, label=f"Bath {bath+1}", **kwargs
|
||
)
|
||
|
||
plot_with_σ(
|
||
model.t,
|
||
interaction_energy,
|
||
ax=ax,
|
||
bath=bath,
|
||
label=f"Interaction {bath+1}",
|
||
**kwargs,
|
||
)
|
||
|
||
total = model.total_energy(data, **ensemble_args)
|
||
plot_with_σ(
|
||
model.t,
|
||
total,
|
||
ax=ax,
|
||
label="Total",
|
||
linestyle="--",
|
||
**kwargs,
|
||
)
|
||
|
||
return (fig, ax)
|
||
|
||
|
||
@wrap_plot
|
||
def plot_coherences(model, ax=None):
|
||
with aux.get_data(model) as data:
|
||
plot_with_σ(
|
||
model.t,
|
||
EnsembleValue(
|
||
(
|
||
0,
|
||
np.abs(np.array(data.rho_t_accum.mean)[:, 0, 1]),
|
||
np.array(data.rho_t_accum.ensemble_std)[:, 0, 1],
|
||
)
|
||
),
|
||
ax=ax,
|
||
)
|
||
|
||
|
||
@wrap_plot
|
||
def plot_distance_measures(model, strobe_indices, ax=None):
|
||
with aux.get_data(model) as data:
|
||
plot_with_σ(
|
||
model.t, EnsembleValue(relative_entropy(data, strobe_indices[-1])), ax=ax
|
||
)
|
||
plot_with_σ(
|
||
model.t, EnsembleValue(trace_distance(data, strobe_indices[-1])), ax=ax
|
||
)
|
||
|
||
|
||
@wrap_plot
|
||
def plot_σ_development(ensemble_value, ax=None, **kwargs):
|
||
ax.plot(
|
||
ensemble_value.Ns,
|
||
[σ.mean() for σ in ensemble_value.σs],
|
||
marker=".",
|
||
linestyle="dotted",
|
||
**kwargs,
|
||
)
|
||
|
||
ax.set_ylabel("mean $σ$")
|
||
ax.set_xlabel("$N$")
|