-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathfigure_grid.py
34 lines (27 loc) · 1.03 KB
/
figure_grid.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
import matplotlib.pyplot as plt
class figure_grid():
def next_subplot(self, **kwargs):
if self.subplots:
self.after_each()
self.subplots += 1
return self.fig.add_subplot(self.rows, self.cols, self.subplots, **kwargs)
def each_subplot(self):
for _ in range(self.rows * self.cols):
yield self.next_subplot()
def title(self, title, fontsize=14, y=1.05, **kwargs):
self.fig.suptitle(title, y=y, fontsize=fontsize, **kwargs)
def __init__(self, rows, cols, rowheight=3, rowwidth=12, after_each=lambda: None):
self.rows = rows
self.cols = cols
self.fig = plt.figure(figsize=(rowwidth, rowheight*self.rows))
self.subplots = 0
if after_each == 'legend':
after_each = lambda: plt.legend(loc='best')
self.after_each = after_each
def __enter__(self):
return self
def __exit__(self, _type, _value, _traceback):
self.after_each()
plt.tight_layout()
plt.show()
next = next_subplot