-
Notifications
You must be signed in to change notification settings - Fork 28
/
grid_plotters.py
253 lines (216 loc) · 7.94 KB
/
grid_plotters.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
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
import autofit as af
import matplotlib
backend = af.conf.instance.visualize.get("figures", "backend", str)
matplotlib.use(backend)
from matplotlib import pyplot as plt
import numpy as np
import itertools
from autolens import exc
from autolens.plotters import plotter_util
def plot_grid(
grid,
colors=None,
axis_limits=None,
points=None,
lines=None,
as_subplot=False,
units="arcsec",
kpc_per_arcsec=None,
figsize=(12, 8),
pointsize=5,
pointcolor="k",
xyticksize=16,
cmap="jet",
cb_ticksize=10,
cb_fraction=0.047,
cb_pad=0.01,
cb_tick_values=None,
cb_tick_labels=None,
title="Grid",
titlesize=16,
xlabelsize=16,
ylabelsize=16,
symmetric_around_centre=True,
output_path=None,
output_format="show",
output_filename="grid",
):
"""Plot a grid of (y,x) Cartesian coordinates as a scatter plot of points.
Parameters
-----------
grid : data_type.array.aa.Grid
The (y,x) coordinates of the grid, in an array of shape (total_coordinates, 2).
axis_limits : []
The axis limits of the figure on which the grid is plotted, following [xmin, xmax, ymin, ymax].
points : []
A set of points that are plotted in a different colour for emphasis (e.g. to show the mappings between \
different planes).
as_subplot : bool
Whether the grid is plotted as part of a subplot, in which case the grid figure is not opened / closed.
units : str
The units of the y / x axis of the plots, in arc-seconds ('arcsec') or kiloparsecs ('kpc').
kpc_per_arcsec : float
The conversion factor between arc-seconds and kiloparsecs, required to plot the units in kpc.
figsize : (int, int)
The size of the figure in (rows, columns).
pointsize : int
The size of the points plotted on the grid.
xyticksize : int
The font size of the x and y ticks on the figure axes.
title : str
The text of the title.
titlesize : int
The size of of the title of the figure.
xlabelsize : int
The fontsize of the x axes label.
ylabelsize : int
The fontsize of the y axes label.
output_path : str
The path on the hard-disk where the figure is output.
output_filename : str
The filename of the figure that is output.
output_format : str
The format the figue is output:
'show' - display on computer screen.
'png' - output to hard-disk as a png.
"""
plotter_util.setup_figure(figsize=figsize, as_subplot=as_subplot)
grid = convert_grid_units(
grid_arcsec=grid, units=units, kpc_per_arcsec=kpc_per_arcsec
)
if colors is not None:
plt.cm.get_cmap(cmap)
plt.scatter(
y=np.asarray(grid[:, 0]),
x=np.asarray(grid[:, 1]),
c=colors,
s=pointsize,
marker=".",
cmap=cmap,
)
if colors is not None:
plotter_util.set_colorbar(
cb_ticksize=cb_ticksize,
cb_fraction=cb_fraction,
cb_pad=cb_pad,
cb_tick_values=cb_tick_values,
cb_tick_labels=cb_tick_labels,
)
plotter_util.set_title(title=title, titlesize=titlesize)
set_xy_labels(
units=units,
kpc_per_arcsec=kpc_per_arcsec,
xlabelsize=xlabelsize,
ylabelsize=ylabelsize,
xyticksize=xyticksize,
)
set_axis_limits(
axis_limits=axis_limits,
grid=grid,
symmetric_around_centre=symmetric_around_centre,
)
plot_points(grid=grid, points=points, pointcolor=pointcolor)
plotter_util.plot_lines(line_lists=lines)
plt.tick_params(labelsize=xyticksize)
plotter_util.output_figure(
array=None,
as_subplot=as_subplot,
output_path=output_path,
output_filename=output_filename,
output_format=output_format,
)
plotter_util.close_figure(as_subplot=as_subplot)
def convert_grid_units(grid_arcsec, units, kpc_per_arcsec):
"""Convert the grid from its input units (arc-seconds) to the input unit (e.g. retain arc-seconds) or convert to \
another set of units (kiloparsecs).
Parameters
-----------
grid_arcsec : ndarray or data_type.array.aa.Grid
The (y,x) coordinates of the grid in arc-seconds, in an array of shape (total_coordinates, 2).
units : str
The units of the y / x axis of the plots, in arc-seconds ('arcsec') or kiloparsecs ('kpc').
kpc_per_arcsec : float
The conversion factor between arc-seconds and kiloparsecs, required to plot the units in kpc.
"""
if units in "arcsec" or kpc_per_arcsec is None:
return grid_arcsec
elif units in "kpc":
return grid_arcsec * kpc_per_arcsec
def set_xy_labels(units, kpc_per_arcsec, xlabelsize, ylabelsize, xyticksize):
"""Set the x and y labels of the figure, and set the fontsize of those labels.
The x and y labels are always the distance scales, thus the labels are either arc-seconds or kpc and depend on the \
units the figure is plotted in.
Parameters
-----------
units : str
The units of the y / x axis of the plots, in arc-seconds ('arcsec') or kiloparsecs ('kpc').
kpc_per_arcsec : float
The conversion factor between arc-seconds and kiloparsecs, required to plot the units in kpc.
xlabelsize : int
The fontsize of the x axes label.
ylabelsize : int
The fontsize of the y axes label.
xyticksize : int
The font size of the x and y ticks on the figure axes.
"""
if units in "arcsec" or kpc_per_arcsec is None:
plt.xlabel("x (arcsec)", fontsize=xlabelsize)
plt.ylabel("y (arcsec)", fontsize=ylabelsize)
elif units in "kpc":
plt.xlabel("x (kpc)", fontsize=xlabelsize)
plt.ylabel("y (kpc)", fontsize=ylabelsize)
else:
raise exc.PlottingException(
"The units supplied to the plotted are not a valid string (must be pixels | "
"arcsec | kpc)"
)
plt.tick_params(labelsize=xyticksize)
def set_axis_limits(axis_limits, grid, symmetric_around_centre):
"""Set the axis limits of the figure the grid is plotted on.
Parameters
-----------
axis_limits : []
The axis limits of the figure on which the grid is plotted, following [xmin, xmax, ymin, ymax].
"""
if axis_limits is not None:
plt.axis(axis_limits)
elif symmetric_around_centre:
ymin = np.min(grid[:, 0])
ymax = np.max(grid[:, 0])
xmin = np.min(grid[:, 1])
xmax = np.max(grid[:, 1])
x = np.max([np.abs(xmin), np.abs(xmax)])
y = np.max([np.abs(ymin), np.abs(ymax)])
axis_limits = [-x, x, -y, y]
plt.axis(axis_limits)
def plot_points(grid, points, pointcolor):
"""Plot a subset of points in a different color, to highlight a specifc region of the grid (e.g. how certain \
pixels map between different planes).
Parameters
-----------
grid : ndarray or data_type.array.aa.Grid
The (y,x) coordinates of the grid, in an array of shape (total_coordinates, 2).
points : []
A set of points that are plotted in a different colour for emphasis (e.g. to show the mappings between \
different planes).
pointcolor : str or None
The color the points should be plotted. If None, the points are iterated through a cycle of colors.
"""
if points is not None:
if pointcolor is None:
point_colors = itertools.cycle(["y", "r", "k", "g", "m"])
for point_set in points:
plt.scatter(
y=np.asarray(grid[point_set, 0]),
x=np.asarray(grid[point_set, 1]),
s=8,
color=next(point_colors),
)
else:
for point_set in points:
plt.scatter(
y=np.asarray(grid[point_set, 0]),
x=np.asarray(grid[point_set, 1]),
s=8,
color=pointcolor,
)