-
Notifications
You must be signed in to change notification settings - Fork 15
Expand file tree
/
Copy pathbase_mpl_plotter.py
More file actions
53 lines (45 loc) · 1.89 KB
/
Copy pathbase_mpl_plotter.py
File metadata and controls
53 lines (45 loc) · 1.89 KB
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
from matplotlib import pyplot as plt
from matplotlib import gridspec
class MPLPlotter:
def new_ax(self, xlabel=None, ylabel=None):
"""Return a new matplotlib axis optionally with the given x and y labels"""
fig, ax = plt.subplots()
if xlabel:
ax.set_xlabel(xlabel)
if ylabel:
ax.set_ylabel(ylabel)
return ax
def new_two_panel_axes(self, n_bottom=1, n_top=1, emphasis="top"):
"""Return the axes handles for a bottom and top panel.
Args:
n_top (int): 1 for a single y-axis, 2 for left and right y-axes on top panel
n_bottom (int): 1 for a single y-axis, 2 for left and right y-axes on bottom
emphasis (str or None): "top" for bigger top panel, "bottom" for bigger
bottom panel, None for equal-sized panels
Returns list of axes: top left, bottom left(, bottom right)(, top right)
"""
self.new_ax() # necessary to avoid deleting an open figure, I don't know why
if emphasis == "top":
gs = gridspec.GridSpec(5, 1)
# gs.update(hspace=0.025)
axes = [plt.subplot(gs[0:3, 0])]
axes += [plt.subplot(gs[3:5, 0])]
elif emphasis == "bottom":
gs = gridspec.GridSpec(5, 1)
# gs.update(hspace=0.025)
axes = [plt.subplot(gs[0:2, 0])]
axes += [plt.subplot(gs[2:5, 0])]
else:
gs = gridspec.GridSpec(6, 1)
# gs.update(hspace=0.025)
axes = [plt.subplot(gs[0:3, 0])]
axes += [plt.subplot(gs[3:6, 0])]
axes[0].xaxis.set_label_position("top")
axes[0].tick_params(
axis="x", top=True, bottom=False, labeltop=True, labelbottom=False
)
if n_bottom == 2:
axes += [axes[1].twinx()]
if n_top == 2:
axes += [axes[0].twinx()]
return axes