master-thesis/python/energy_flow_proper/05_gaussian_two_baths/figsaver.py
2022-07-04 16:36:29 +02:00

350 lines
8.7 KiB
Python

import os
from pathlib import Path
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from contextlib import contextmanager
fig_path = Path(os.getcwd()) / "figures"
val_path = Path(os.getcwd()) / "values"
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)
def export_fig(name, fig=None, style=None):
fig_path.mkdir(parents=True, exist_ok=True)
if fig is None:
fig = plt.gcf()
with hiro_style(style):
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")
return fig
def scientific_round(val, *err, retprec=False):
"""Scientifically rounds the values to the given errors."""
val, err = np.asarray(val), np.asarray(err)
if len(err.shape) == 1:
err = np.array([err])
err = err.T
err = err.T
if err.size == 1 and val.size > 1:
err = np.ones_like(val) * err
if len(err.shape) == 0:
err = np.array([err])
if val.size == 1 and err.shape[0] > 1:
val = np.ones_like(err) * val
i = np.floor(np.log10(err))
first_digit = (err // 10 ** i).astype(int)
prec = (-i + np.ones_like(err) * (first_digit <= 3)).astype(int)
prec = np.max(prec, axis=1)
def smart_round(value, precision):
value = np.round(value, precision)
if precision <= 0:
value = value.astype(int)
return value
if val.size > 1:
rounded = np.empty_like(val)
rounded_err = np.empty_like(err)
for n, (value, error, precision) in enumerate(zip(val, err, prec)):
rounded[n] = smart_round(value, precision)
rounded_err[n] = smart_round(error, precision)
if retprec:
return rounded, rounded_err, prec
else:
return rounded, rounded_err
else:
prec = prec[0]
if retprec:
return (smart_round(val, prec), *smart_round(err, prec)[0], prec)
else:
return (smart_round(val, prec), *smart_round(err, prec)[0])
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 tex_value(val, err=None, unit=None, prefix="", suffix="", prec=0, save=None):
"""Generates LaTeX output of a value with units and error."""
if err:
val, err, prec = scientific_round(val, err, retprec=True)
else:
val = np.round(val, prec)
if prec == 0:
val = int(val)
if err:
err = int(err)
val_string = fr"{val:.{prec}f}" if prec > 0 else str(val)
if err:
val_string += fr"\pm {err:.{prec}f}" if prec > 0 else str(err)
ret_string = r"\(" + prefix
if unit is None:
ret_string += val_string
else:
ret_string += fr"\SI{{{val_string}}}{{{unit}}}"
ret_string += suffix + r"\)"
if save is not None:
val_path.mkdir(parents=True, exist_ok=True)
with open(val_path / f"{save}.tex", "w") as f:
f.write(ret_string)
return ret_string
###############################################################################
# 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)
MPL_RC = {
"font.family": "serif",
"text.usetex": False,
"pgf.rcfonts": False,
"lines.linewidth": 1,
"font.size": 10,
"axes.labelsize": 10,
"axes.titlesize": 8,
"legend.fontsize": 8,
"xtick.labelsize": 8,
"ytick.labelsize": 8,
}
MPL_RC_POSTER = {
"font.family": "sans",
"text.usetex": False,
"pgf.rcfonts": False,
"lines.linewidth": 2,
"font.size": 17.28,
"axes.labelsize": 17.28,
"axes.titlesize": 17.28,
"legend.fontsize": 17.28,
"xtick.labelsize": 14.4,
"ytick.labelsize": 14.4,
}
@contextmanager
def hiro_style(style=None):
with plt.style.context("ggplot"):
with matplotlib.rc_context(MPL_RC_POSTER if style == "poster" else MPL_RC):
yield True
plt.style.use("ggplot")
matplotlib.rcParams.update(MPL_RC)
@wrap_plot
def plot_complex(x, y, *args, ax=None, label="", **kwargs):
label = label + ", " if (len(label) > 0) else ""
ax.plot(x, y.real, *args, label=f"{label}real part", **kwargs)
ax.plot(x, y.imag, *args, label=f"{label}imag part", **kwargs)
ax.legend()
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_convergence(
x,
y,
ax=None,
label="",
transform=lambda y: y,
slice=None,
linestyle="-",
bath=None,
color=None,
**kwargs,
):
slice = (0, -1) if not slice else slice
ys = y[slice[0] : slice[1]]
label = label + (", " if (len(label) > 0 and len(ys) > 0) else "")
for n, val, _ in ys:
plt.plot(
x,
transform(val[bath] if bath is not None else val),
label=f"{label}$N={n}$",
alpha=n / y[-1][0],
linestyle=":",
)
fancy_error(
x,
transform(y[-1][1][bath] if bath is not None else y[-1][1]),
(y[-1][2][bath] if bath is not None else y[-1][2]).real,
label=f"{label}" + (fr"$N={y[-1][0]}$" if len(ys) > 1 else ""),
linestyle=linestyle,
color=color,
ax=ax,
**kwargs,
)
return None
@wrap_plot
def plot_diff_vs_sigma(
x,
y,
reference,
ax=None,
label="",
transform=lambda y: y,
ecolor="yellow",
ylabel=None,
bath=None,
):
label = label + ", " if (len(label) > 0) else ""
y = y if bath is None else [(n, val[bath], err[bath]) for n, val, err in y]
ref_traj = transform(reference)
ref_err = np.abs(y[-1][2].real)
ax.fill_between(
x,
0,
ref_err,
color=ecolor,
label=fr"{label}$\sigma\, (N={y[-1][0]})$",
)
for n, val, err in y:
diff = np.abs(transform(val) - ref_traj)
within = (diff < ref_err).sum() / y[-1][2].size
not_last = n < y[-1][0]
ax.plot(
x,
diff,
label=fr"{label}$N={n}$ $\Delta<\sigma = {within * 100:.1f}\%$",
alpha=within 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 = fr"\text{{ {ylabel} }}"
ax.set_ylabel(fr"$|{{{ylabel}}}_{{\mathrm{{ref}}}}-{{{ylabel}}}_{{N_i}}|$")