-
Notifications
You must be signed in to change notification settings - Fork 2
/
plawt.py
133 lines (112 loc) · 5.98 KB
/
plawt.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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import matplotlib as mpl
import matplotlib.pyplot as plt
from numpy import ndarray
def plot(*plotStructs):
if len(plotStructs) == 0:
raise TypeError("Missing 1 required argument")
globalParams = plotStructs[0] # set all global properties using the first plotStruct
plt.style.use(globalParams['style']) if 'style' in globalParams else None
# setup subplots
sharex = globalParams['sharex'] if 'sharex' in globalParams else False
sharey = globalParams['sharey'] if 'sharey' in globalParams else False
gridspec_kw = globalParams['gridspec_kw'] if 'gridspec_kw' in globalParams else None
# setup figure and adjustments
# https://matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.subplots_adjust
hspace = globalParams['hspace'] if 'hspace' in globalParams else 0.2
wspace = globalParams['wspace'] if 'wspace' in globalParams else 0.2
top = globalParams['top'] if 'top' in globalParams else 0.9
bottom = globalParams['bottom'] if 'bottom' in globalParams else 0.1
left = globalParams['left'] if 'left' in globalParams else 0.125
right = globalParams['right'] if 'right' in globalParams else 0.9
title = globalParams['title'] if 'title' in globalParams else ''
titlesize = globalParams['titlesize'] if 'titlesize' in globalParams else mpl.rcParams['figure.titlesize']
nrows = globalParams['nrows'] if 'nrows' in globalParams else 1
ncols = globalParams['ncols'] if 'ncols' in globalParams else 1
isGridSpecified = 'nrows' in globalParams or 'ncols' in globalParams
if len(plotStructs) != 1 and not isGridSpecified:
nrows = len(plotStructs)
fig, axesarray = plt.subplots(nrows=nrows, ncols=ncols, sharex=sharex, sharey=sharey, gridspec_kw=gridspec_kw)
axes = axesarray.flatten() if type(axesarray) is ndarray else axesarray
fig.set_size_inches(globalParams['figsize']) if 'figsize' in globalParams else None
fig.subplots_adjust(hspace=hspace, wspace=wspace, top=top, bottom=bottom, right=right, left=left)
fig.suptitle(title, fontsize=titlesize)
# hack for common x and y labels on subplots
if sharex or sharey:
fig.add_subplot(111, frameon=False)
plt.tick_params(labelcolor='none', top='off', bottom='off', left='off', right='off')
xlabel = globalParams['xlabel'] if 'xlabel' in globalParams else ''
ylabel = globalParams['ylabel'] if 'ylabel' in globalParams else ''
plt.xlabel(xlabel) if sharex else None
plt.ylabel(ylabel) if sharey else None
plt.gca().xaxis.label.set_size(globalParams['xlabelsize']) if 'xlabelsize' in globalParams else None
plt.gca().yaxis.label.set_size(globalParams['ylabelsize']) if 'ylabelsize' in globalParams else None
if not isinstance(axes, ndarray):
axes = (axes,) # this makes axes iterable even if there's just one axes
for i, ax in enumerate(axes):
try: # show blank axis if user mispecifies nrows or ncols, like matplotlib
plotStruct = plotStructs[i]
except IndexError:
plotStruct = {}
for key in plotStruct:
val = plotStruct[key]
if type(key) == int: # its a line
x = val['x']
y = val['y'] if 'y' in val else None
linestyle = val['line'] if 'line' in val else ''
drawstyle = val['draw'] if 'draw' in val else 'default'
label = val['label'] if 'label' in val else ''
markersize = val['markersize'] if 'markersize' in val else 6.0
markersize = val['ms'] if 'ms' in val else markersize
linewidth = val['linewidth'] if 'linewidth' in val else 1.0
linewidth = val['lw'] if 'lw' in val else linewidth
alpha = val['alpha'] if 'alpha' in val else 1.0
if y is not None:
line, = ax.plot(x, y, linestyle)
else:
line, = ax.plot(y, linestyle)
mfc = val['markerfacecolor'] if 'markerfacecolor' in val else line.get_markerfacecolor()
mfc = val['mfc'] if 'mfc' in val else line.get_markerfacecolor()
line.set_drawstyle(drawstyle)
line.set_label(label)
line.set_markersize(markersize)
line.set_linewidth(linewidth)
line.set_alpha(alpha)
line.set_markerfacecolor(mfc)
elif type(key) == str: # its a property
ax.set_xlabel(val) if key == 'xlabel' and not sharex else None
ax.set_ylabel(val) if key == 'ylabel' and not sharey else None
ax.set_yscale(val) if key == 'set_yscale' else None
ax.set_xscale(val) if key == 'set_xscale' else None
ax.grid() if key == 'grid' and val else None
ax.set_xlim(val) if key == 'xlim' else None
ax.set_ylim(val) if key == 'ylim' else None
ax.set_aspect(val) if key == 'aspect' else None
ax.legend(**val) if key == 'legend' else None
ax.minorticks_on() if key == 'minorticks' and val else None
ax.xaxis.label.set_size(globalParams['xlabelsize']) if 'xlabelsize' in plotStruct else None
ax.yaxis.label.set_size(globalParams['ylabelsize']) if 'xlabelsize' in plotStruct else None
if key == 'tick_params':
val = plotStruct['tick_params']
if type(val) is list:
for operation in val:
ax.tick_params(**operation)
else:
ax.tick_params(**val)
if key == 'subtitle':
fontsize = plotStruct['fontsize'] if 'fontsize' in plotStruct else mpl.rcParams['axes.titlesize']
loc = plotStruct['subloc'] if 'subloc' in plotStruct else 'center'
fontdict = plotStruct['subtitledict'] if 'subtitledict' in plotStruct else None
ax.set_title(val, fontdict=fontdict, fontsize=fontsize, loc=loc)
if key == 'text':
if not isinstance(val, list):
val = [val,]
for v in val:
textx, texty, s = v.pop('x'), v.pop('y'), v.pop('s')
ax.text(textx, texty, s, **v)
plt.tight_layout(**globalParams['tight_layout']) if 'tight_layout' in globalParams else None
fig.subplots_adjust(hspace=hspace, wspace=wspace, top=top, bottom=bottom, right=right, left=left)
# Storing and cleanup of plot
fig.savefig(globalParams['filename']) if 'filename' in globalParams else None
plt.show() if 'show' in globalParams and globalParams['show'] else None
plt.close(fig) if 'keepOpen' not in globalParams or not globalParams['keepOpen'] else None
return fig # in case people wanna do extra things