-
Notifications
You must be signed in to change notification settings - Fork 17
/
plotting.py
119 lines (102 loc) · 4.17 KB
/
plotting.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import os
import matplotlib.pyplot as plt
import warnings
DEFAULT_WIDTH = 6.0
DEFAULT_HEIGHT = 1.5
SIZE_SMALL = 9 # Caption size in the pml book
def latexify(
width_scale_factor=1,
height_scale_factor=1,
fig_width=None,
fig_height=None,
font_size=SIZE_SMALL,
):
f"""
width_scale_factor: float, DEFAULT_WIDTH will be divided by this number, DEFAULT_WIDTH is page width: {DEFAULT_WIDTH} inches.
height_scale_factor: float, DEFAULT_HEIGHT will be divided by this number, DEFAULT_HEIGHT is {DEFAULT_HEIGHT} inches.
fig_width: float, width of the figure in inches (if this is specified, width_scale_factor is ignored)
fig_height: float, height of the figure in inches (if this is specified, height_scale_factor is ignored)
font_size: float, font size
"""
if "LATEXIFY" not in os.environ:
warnings.warn("LATEXIFY environment variable not set, not latexifying")
return
if fig_width is None:
fig_width = DEFAULT_WIDTH / width_scale_factor
if fig_height is None:
fig_height = DEFAULT_HEIGHT / height_scale_factor
# use TrueType fonts so they are embedded
# https://stackoverflow.com/questions/9054884/how-to-embed-fonts-in-pdfs-produced-by-matplotlib
# https://jdhao.github.io/2018/01/18/mpl-plotting-notes-201801/
plt.rcParams["pdf.fonttype"] = 42
# Font sizes
# SIZE_MEDIUM = 14
# SIZE_LARGE = 24
# https://stackoverflow.com/a/39566040
plt.rc("font", size=font_size) # controls default text sizes
plt.rc("axes", titlesize=font_size) # fontsize of the axes title
plt.rc("axes", labelsize=font_size) # fontsize of the x and y labels
plt.rc("xtick", labelsize=font_size) # fontsize of the tick labels
plt.rc("ytick", labelsize=font_size) # fontsize of the tick labels
plt.rc("legend", fontsize=font_size) # legend fontsize
plt.rc("figure", titlesize=font_size) # fontsize of the figure title
# latexify: https://nipunbatra.github.io/blog/visualisation/2014/06/02/latexify.html
plt.rcParams["backend"] = "ps"
plt.rc("text", usetex=True)
plt.rc("font", family="serif")
plt.rc("figure", figsize=(fig_width, fig_height))
def is_latexify_enabled():
"""
returns true if LATEXIFY environment variable is set
"""
return "LATEXIFY" in os.environ
def _get_fig_name(fname_full):
fname_full = fname_full.replace("_latexified", "")
LATEXIFY = "LATEXIFY" in os.environ
# extention = "_latexified.pdf" if LATEXIFY else ".png"
extention = "_latexified.pdf" if LATEXIFY else ".pdf"
if fname_full[-4:] in [".png", ".pdf", ".jpg"]:
fname = fname_full[:-4]
warnings.warn(
f"renaming {fname_full} to {fname}{extention} because LATEXIFY is {LATEXIFY}",
)
else:
fname = fname_full
return fname + extention
def savefig(
f_name, tight_layout=True, tight_bbox=False, pad_inches=0.0, *args, **kwargs
):
if len(f_name) == 0:
return
if "FIG_DIR" not in os.environ:
warnings.warn("set FIG_DIR environment variable to save figures")
return
fig_dir = os.environ["FIG_DIR"]
# Auto create the directory if it doesn't exist
if not os.path.exists(fig_dir):
os.makedirs(fig_dir)
fname_full = os.path.join(fig_dir, f_name)
fname_full = _get_fig_name(fname_full)
print("saving image to {}".format(fname_full))
if tight_layout:
plt.tight_layout(pad=pad_inches)
print("Figure size:", plt.gcf().get_size_inches())
if tight_bbox:
# This changes the size of the figure
plt.savefig(
fname_full, pad_inches=pad_inches, bbox_inches="tight", *args, **kwargs
)
else:
plt.savefig(fname_full, pad_inches=pad_inches, *args, **kwargs)
if "DUAL_SAVE" in os.environ:
if fname_full.endswith(".pdf"):
fname_full = fname_full[:-4] + ".png"
else:
fname_full = fname_full[:-4] + ".pdf"
if tight_bbox:
# This changes the size of the figure
plt.savefig(
fname_full, pad_inches=pad_inches, bbox_inches="tight", *args, **kwargs
)
else:
plt.savefig(fname_full, pad_inches=pad_inches, *args, **kwargs)