-
Notifications
You must be signed in to change notification settings - Fork 8
/
utils.py
214 lines (176 loc) · 6.11 KB
/
utils.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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
import collections
import matplotlib.pyplot as plt
import matplotlib as mpl
margin_l_absolute = 8. * 0.1
margin_r_absolute = 8. * 0.02
margin_sep_absolute = 8. * 0.02
margin_sep_small_absolute = 8. * 0.005
margin_t_absolute = 8. * 0.02
margin_t_absolute_extra = 8. * 0.12
margin_b_absolute = 8. * 0.08
colors = [
"#CC002E",
"orange",
"mediumblue",
"darkgreen",
"c",
"#be96ff",
"0.65",
"black",
]
our_cmap = mpl.colors.LinearSegmentedColormap.from_list(
name="rascascal", colors=["#CC002E", "0.85", "mediumblue"], N=201
)
our_cmap_r = mpl.colors.LinearSegmentedColormap.from_list(
name="rascascal_r", colors=["mediumblue", "0.85", "#CC002E"], N=201
)
def calculate_height(
n_panels=2, width=8., panel_aspect_ratio=1., extra_top_space=False
):
if isinstance(n_panels, collections.Sequence):
n_panels_h, n_panels_v = n_panels
else:
n_panels_h = n_panels
n_panels_v = 1
# Determine top margin
_margin_t_absolute = (
margin_t_absolute_extra if extra_top_space else margin_t_absolute
)
# Calculate horizontal margins. Units: relative to width.
margin_l = margin_l_absolute / width
margin_r = margin_r_absolute / width
margin_l_subsequent = margin_l
if n_panels_h > 2:
margin_l_subsequent = margin_r
margin_sep = margin_sep_absolute / width
if n_panels_h > 2:
margin_sep = 0
margin_sep_total = margin_r + margin_sep + margin_l_subsequent
panel_width = (
1. - margin_l - margin_r - (n_panels_h - 1) * margin_sep_total
) / n_panels_h
# Calculate absolute height
panel_height_absolute = panel_width * width / panel_aspect_ratio
height = (
n_panels_v * (panel_height_absolute + _margin_t_absolute + margin_b_absolute)
+ (n_panels_v - 1) * margin_sep_absolute
)
# Calculate horizontal margins. Units: relative to width.
panel_height = panel_height_absolute / height
margin_t = _margin_t_absolute / height
margin_b = margin_b_absolute / height
margin_sep_total = margin_t + margin_b + margin_sep_absolute / height
# Return height
return height
def adjust_margins(n_panels=2, width=8., panel_aspect_ratio=1., extra_top_space=False):
if isinstance(n_panels, collections.Sequence):
n_panels_h, n_panels_v = n_panels
else:
n_panels_h = n_panels
n_panels_v = 1
# Determine top margin
_margin_t_absolute = (
margin_t_absolute_extra if extra_top_space else margin_t_absolute
)
# Calculate horizontal margins. Units: relative to width.
margin_l = margin_l_absolute / width
margin_r = margin_r_absolute / width
margin_l_subsequent = margin_l
if n_panels_h > 2:
margin_l_subsequent = margin_r
margin_sep = margin_sep_absolute / width
if n_panels_h > 2:
margin_sep = 0
margin_sep_total = margin_r + margin_sep + margin_l_subsequent
panel_width = (
1. - margin_l - margin_r - (n_panels_h - 1) * margin_sep_total
) / n_panels_h
# Calculate wspace argument of subplots_adjust
wspace = margin_sep_total / panel_width
# Calculate absolute height
panel_height_absolute = panel_width * width / panel_aspect_ratio
height = (
n_panels_v * (panel_height_absolute + _margin_t_absolute + margin_b_absolute)
+ (n_panels_v - 1) * margin_sep_absolute
)
# Calculate horizontal margins. Units: relative to width.
panel_height = panel_height_absolute / height
margin_t = _margin_t_absolute / height
margin_b = margin_b_absolute / height
margin_sep_total = margin_t + margin_b + margin_sep_absolute / height
# Calculate wspace argument of subplots_adjust
hspace = margin_sep_total / panel_height
# Set margins
plt.subplots_adjust(
left=margin_l,
right=1. - margin_r,
bottom=margin_b,
top=1. - margin_t,
wspace=wspace,
hspace=hspace,
)
def cmap_color(x, xmin=0., xmax=1.):
# cmap = mpl.cm.RdYlBu
return our_cmap((x - xmin) / (xmax - xmin))
def viridis_color(x, xmin=0., xmax=1.):
cmap = mpl.cm.viridis
return cmap((x - xmin) / (xmax - xmin))
def custom_figure(
n_panels=2,
width=8.,
panel_aspect_ratio=1.,
extra_top_space=False,
reduce_vertical_sep=False,
):
if isinstance(n_panels, collections.Sequence):
n_panels_h, n_panels_v = n_panels
else:
n_panels_h = n_panels
n_panels_v = 1
# Determine top margin
_margin_t_absolute = (
margin_t_absolute_extra if extra_top_space else margin_t_absolute
)
# Calculate horizontal margins. Units: relative to width.
margin_l = margin_l_absolute / width
margin_r = margin_r_absolute / width
margin_l_subsequent = margin_l
if n_panels_h > 2:
margin_l_subsequent = margin_r
margin_sep = margin_sep_absolute / width
if n_panels_h > 2:
margin_sep = 0
margin_sep_total = margin_r + margin_sep + margin_l_subsequent
panel_width = (
1. - margin_l - margin_r - (n_panels_h - 1) * margin_sep_total
) / n_panels_h
# Calculate wspace argument of subplots_adjust
wspace = margin_sep_total / panel_width
# Calculate absolute height
panel_height_absolute = panel_width * width / panel_aspect_ratio
height = (
n_panels_v * (panel_height_absolute + _margin_t_absolute + margin_b_absolute)
+ (n_panels_v - 1) * margin_sep_absolute
)
# Calculate vertical margins. Units: relative to width.
panel_height = panel_height_absolute / height
margin_t = _margin_t_absolute / height
margin_b = margin_b_absolute / height
if reduce_vertical_sep:
margin_sep_total = margin_sep_small_absolute / height
else:
margin_sep_total = margin_t + margin_b + margin_sep_absolute / height
# Calculate wspace argument of subplots_adjust
hspace = margin_sep_total / panel_height
# New figure
fig = plt.figure(figsize=(width, height))
# Adjust margins
plt.subplots_adjust(
left=margin_l,
right=1. - margin_r,
bottom=margin_b,
top=1. - margin_t,
wspace=wspace,
hspace=hspace,
)
return fig