Coverage for frank/plot.py: 91%

1# Frankenstein: 1D disc brightness profile reconstruction from Fourier data 

2# using non-parametric Gaussian Processes 

3# 

4# Copyright (C) 2019-2020 R. Booth, J. Jennings, M. Tazzari 

5# 

6# This program is free software: you can redistribute it and/or modify 

7# it under the terms of the GNU General Public License as published by 

8# the Free Software Foundation, either version 3 of the License, or 

9# (at your option) any later version. 

10# 

11# This program is distributed in the hope that it will be useful, 

12# but WITHOUT ANY WARRANTY; without even the implied warranty of 

13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

14# GNU General Public License for more details. 

15# 

16# You should have received a copy of the GNU General Public License 

17# along with this program. If not, see <https://www.gnu.org/licenses/> 

18# 

19"""This module contains plotting routines for visualizing and analyzing 

20Frankenstein fits. 

21""" 

22 

23import numpy as np 

24import matplotlib.pyplot as plt 

25 

26# Suppress some benign warnings 

27import warnings 

28warnings.filterwarnings('ignore', '.*compatible with tight_layout.*') 

29warnings.filterwarnings('ignore', '.*handles with labels found.*') 

30 

31from frank.utilities import sweep_profile 

32 

33 

34def plot_deprojection_effect(u, v, up, vp, vis, visp, axes): 

35 """ 

36 Overplot projected and deprojected (u, v) coordinates; 

37 projected and deprojected visibility amplitudes 

38 (here 'projection' refers to correcting collectively for the source 

39 inclination, position angle and phase offset) 

40 

41 Parameters 

42 ---------- 

43 u, v : array 

44 Projected (u, v) coordinates 

45 up, vp : array 

46 Deprojected (u, v) coordinates 

47 vis : array 

48 Projected visibilities (either the real or imaginary component) 

49 visp : array 

50 Deprojected visibilities (either the real or imaginary component) 

51 axes : list[Axes, Axes] 

52 Three axes on which to plot 

53 """ 

54 

55 ax0, ax1 = axes 

56 

57 ax0.plot(u, v, '+', c='#23E1DB', label='Projected') 

58 ax0.plot(up, vp, 'x', c='#D14768', label='Deprojected') 

59 

60 # Projected baselines 

61 bs = np.hypot(u, v) 

62 # Deprojected baselines 

63 bsp = np.hypot(up, vp) 

64 

65 ax1.plot(bs, vis, '+', c='#23E1DB', label='Projected') 

66 ax1.plot(bsp, visp, 'x', c='#D14768', label='Deprojected') 

67 

68 ax0.legend(loc='best') 

69 ax1.legend(loc='best') 

70 

71 

72def plot_brightness_profile(fit_r, fit_i, ax, dist=None, low_uncer=None, 

73 high_uncer=None, **kwargs): 

74 """ 

75 Plot a brightness profile (and optionally a confidence inverval) as a 

76 function of disc radius, I(r) 

77 

78 Parameters 

79 ---------- 

80 fit_r : array 

81 Radial data coordinates 

82 fit_i : array 

83 Brightness values at fit_r 

84 ax : Matplotlib `~.axes.Axes` class 

85 Axis on which to plot 

86 dist : float, optional, default = None, unit = [AU] 

87 Distance to source. If not None, a new `ax` will be created to 

88 show an upper x-axis in [AU] for the plot on the current `ax` 

89 low_uncer : Negative (i.e., below MAP) uncertainty on fit_i 

90 high_uncer : Positive (i.e., above MAP) uncertainty on fit_i 

91 

92 Returns 

93 ------- 

94 ax_new : Matplotlib `~.axes.Axes` class 

95 Only if dist is not None, the second x-axis, ax_new will be returned 

96 

97 """ 

98 

99 if dist: 

100 ax_new = ax.twiny() 

101 ax_new.spines['top'].set_color('#1A9E46') 

102 ax_new.tick_params(axis='x', which='both', colors='#1A9E46') 

103 ax_new.plot(fit_r * dist, fit_i, **kwargs) 

104 ax_new.set_xlabel('r [AU]', color='#1A9E46') 

105 

106 return ax_new 

107 

108 else: 

109 if low_uncer is not None: 

110 if high_uncer is None: 

111 high_uncer = low_uncer * 1. 

112 ax.fill_between(fit_r, fit_i - low_uncer, fit_i + high_uncer, **kwargs) 

113 

114 ax.plot(fit_r, fit_i, **kwargs) 

115 

116 ax.axhline(0, c='c', ls='--', zorder=10) 

117 

118 if 'label' in kwargs: 

119 ax.legend(loc='best') 

120 

121 

122def plot_vis_quantity(baselines, vis_quantity, ax, vis_quantity_err=None, 

123 **kwargs): 

124 r""" 

125 Plot a visibility domain quantity (e.g., observed visibilities, a frank fit, 

126 residual visibilities, a power spectrum) as a function of baseline 

127 

128 Parameters 

129 ---------- 

130 baselines : array 

131 Baseline data coordinates `b` 

132 vis_quantity : array 

133 A generic quantity `Q` to plot as a function of baselines `b`, Q(b) 

134 ax : Matplotlib `~.axes.Axes` class 

135 Axis on which to plot 

136 vis_quantity_err : array, optional, default = None 

137 Uncertainty on vis_quantity values 

138 """ 

139 

140 # If input arrays are masked with invalid values ('--'), replace those 

141 # masked values with NaN 

142 if np.ma.is_masked(baselines): 

143 baselines = np.ma.array(baselines).filled(np.nan) 

144 vis_quantity = np.ma.array(vis_quantity).filled(np.nan) 

145 if vis_quantity_err is not None: 

146 vis_quantity_err = np.ma.array(vis_quantity_err).filled(np.nan) 

147 

148 if vis_quantity_err is not None: 

149 ax.errorbar(baselines, vis_quantity, yerr=vis_quantity_err, **kwargs) 

150 else: 

151 ax.plot(baselines, vis_quantity, **kwargs) 

152 

153 ax.axhline(0, c='c', ls='--', zorder=10) 

154 

155 if 'label' in kwargs: 

156 ax.legend(loc='best') 

157 

158 

159def plot_vis_hist(binned_vis, ax, rescale=None, **kwargs): 

160 r""" 

161 Plot a histogram of visibilities using a precomputed binning 

162 

163 Parameters 

164 ---------- 

165 binned_vis : UVDataBinner object 

166 Pre-binned visibilities (see utilities.UVDataBinner) 

167 ax : Matplotlib `~.axes.Axes` class 

168 Axis on which to plot 

169 rescale : float, default=None 

170 Constant by which to rescale x-values 

171 """ 

172 

173 edges = np.concatenate([binned_vis.bin_edges[0].data, 

174 binned_vis.bin_edges[1].data[-1:]]) 

175 

176 # alter x-axis units 

177 if rescale is not None: 

178 edges /= rescale 

179 

180 counts = binned_vis.bin_counts.data 

181 

182 ax.hist(0.5 * (edges[1:] + edges[:-1]), edges, weights=counts, alpha=.5, **kwargs) 

183 

184 if 'label' in kwargs: 

185 ax.legend(loc='best') 

186 

187 

188def plot_convergence_criterion(profile_iter, N_iter, ax, **kwargs): 

189 r""" 

190 Plot the following convergence criterion for a Frankenstein fit, 

191 :math:`{\rm max}(|I_i - I_{i-1}|) / {\rm max}(I_i)`, 

192 where :math:`I_i` is the brightness profile at iteration :math:`i` 

193 

194 Parameters 

195 ---------- 

196 profile_iter : list, shape = (N_iter, N_coll) 

197 Brightness profile reconstruction over N_iter iterations. 

198 N_coll is the number of collocation points, i.e., the number of grid 

199 points at which the profile is defined 

200 N_iter : int 

201 Total number of iterations in the fit 

202 ax : Matplotlib `~.axes.Axes` class 

203 Axis on which to plot 

204 

205 """ 

206 

207 convergence_criterion = [] 

208 for i in range(N_iter): 

209 this_conv_cri = np.max(np.abs(profile_iter[i] - profile_iter[i-1])) / \ 

210 np.max(profile_iter[i]) 

211 convergence_criterion.append(this_conv_cri) 

212 

213 ax.plot(range(0, N_iter), convergence_criterion, **kwargs) 

214 

215 

216def make_colorbar(ax, vmin, vmax, cmap, label, loc=3, bbox_x=.05, bbox_y=.175): 

217 """ 

218 Custom format to place a colorbar in an inset 

219 

220 Parameters 

221 ---------- 

222 ax : Matplotlib `~.axes.Axes` class 

223 Axis in which to inset the colorbar 

224 vmin, vmax : int 

225 Lower and upper bounds of colorbar scale 

226 cmap : plt.cm colormap 

227 Colormap to apply to the colorbar 

228 label : string 

229 Label for colorbar 

230 loc : int, one of [1, 2, 3, 4], default = 3 

231 Quadrant of colorbar in ax 

232 bbox_x, bbox_y : float, default = 0.05 and 0.175 

233 x- and y-value where the colorbar is placed 

234 """ 

235 

236 from mpl_toolkits.axes_grid1.inset_locator import inset_axes 

237 axins1 = inset_axes(ax, width="50%", height="5%", loc=loc, 

238 bbox_to_anchor=(bbox_x, bbox_y, 1, 1), 

239 bbox_transform=ax.transAxes) 

240 sm = plt.cm.ScalarMappable(cmap=cmap, 

241 norm=plt.Normalize(vmin=vmin, vmax=vmax)) 

242 cbar = plt.colorbar(sm, cax=axins1, orientation="horizontal") 

243 cbar.set_label(label) 

244 axins1.xaxis.set_ticks_position("bottom") 

245 

246 

247def plot_iterations(x, iters, n_iter, ax, 

248 cmap=plt.cm.cool, # pylint: disable=no-member 

249 bbox_x=-.02, bbox_y=-.1, **kwargs): 

250 r""" 

251 Plot a fit quantity (e.g., the brightness profile or power spectrum) over a 

252 range of the fit's iterations 

253 

254 Parameters 

255 ---------- 

256 x : array 

257 x-values at which to plot (e.g., radii for a brightness profile or 

258 baselines for a power spectrum) 

259 iters : list, shape = (n_iter, N_coll) 

260 Iterations to plot (e.g., brightness profiles or power spectra at each 

261 of n_iter iterations) 

262 n_iter : list, of the form [start_iteration, stop_iteration] 

263 Range of iterations in the fit over which to plot iters 

264 ax : Matplotlib `~.axes.Axes` class 

265 Axis on which to plot 

266 cmap : plt.cm colormap, default=plt.cm.cool 

267 Colormap to apply to the overplotted profiles 

268 bbox_x, bbox_y : float, default = -0.02 and -0.1 

269 x- and y-value where the colorbar is placed 

270 """ 

271 if n_iter[0] >= n_iter[1] or n_iter[1] > len(iters): 

272 raise ValueError("Require: n_iter[0] < n_iter[1] and" 

273 " n_iter[1] <= len(iters)") 

274 

275 iter_range = range(n_iter[0], n_iter[1]) 

276 for i in iter_range: 

277 ax.plot(x, iters[i], c=cmap(i / len(iter_range)), **kwargs) 

278 ax.plot(x, iters[-1], ':', c='k', label='Last iteration', **kwargs) 

279 

280 make_colorbar(ax, vmin=n_iter[0], vmax=n_iter[1], cmap=cmap, 

281 label='Iteration', loc=1, bbox_x=bbox_x, bbox_y=bbox_y) 

282 

283 ax.legend(loc='best') 

284 

285 

286def plot_2dsweep(r, I, ax, cmap='inferno', norm=None, xmax=None, ymax=None, 

287 dr=None, plot_colorbar=True, project=False, phase_shift=False, 

288 geom=None, cbar_label=r'I [Jy sr$^{-1}$]', **kwargs): 

289 r""" 

290 Plot a radial profile swept over :math:`2 \pi` to produce an image 

291 

292 Parameters 

293 ---------- 

294 r : array 

295 Radial coordinates at which the 1D brightness profile is defined. 

296 The assumed unit (for the x- and y-label) is arcsec 

297 I : array 

298 Brightness values at r. The assumed unit (for the colorbar) is Jy / sr 

299 ax : Matplotlib `~.axes.Axes` class 

300 Axis on which to plot 

301 cmap : Matplotlib colormap, default = 'inferno' 

302 Colormap to apply to the 2D image 

303 norm : Matplotlib `colors.Normalize` class 

304 Colormap normalization for the image and colorbar. 

305 xmax, ymax : float or None (default) 

306 Value setting the x- and y-bounds of the image (same units as r). The 

307 positive and negative bounds are both set to this value (modulo sign). 

308 If not provided, these will be set to r.max() 

309 dr : float, optional, default = None 

310 Pixel size (same units as r). If not provided, it will be set at the 

311 same spatial scale as r 

312 plot_colorbar: bool, default = True 

313 Whether to plot a colorbar beside the image 

314 project : bool, default = False 

315 Whether to project the swept profile by the supplied geom 

316 phase_shift : bool, default = False 

317 Whether to phase shift the projected profile by the supplied geom. 

318 If False, the source will be centered in the image 

319 geom : SourceGeometry object, default=None 

320 Fitted geometry (see frank.geometry.SourceGeometry). Here we use 

321 geom.inc [deg], geom.PA [deg], geom.dRA [arcsec], geom.dDec [arcsec] if 

322 project=True 

323 cbar_label : string, default = r'I [Jy sr$^{-1}$]' 

324 Colorbar axis label 

325 """ 

326 

327 I2D, xmax_computed, ymax_computed = sweep_profile(r, I, 

328 xmax=xmax, ymax=ymax, 

329 dr=dr, 

330 project=project, 

331 phase_shift=phase_shift, 

332 geom=geom) 

333 

334 if xmax is None: 

335 xmax = xmax_computed 

336 if ymax is None: 

337 ymax = ymax_computed 

338 

339 if norm is None: 

340 import matplotlib.colors as mpl_cs 

341 norm = mpl_cs.Normalize(vmin=I2D.min(), vmax=I2D.max()) 

342 

343 ax.imshow(I2D, origin='lower', extent=(xmax, -1.0 * xmax, -1.0 * ymax, ymax), 

344 cmap=cmap, norm=norm, **kwargs 

345 ) 

346 

347 # Set a normalization and colormap for the colorbar 

348 if plot_colorbar: 

349 from matplotlib import cm 

350 m = cm.ScalarMappable(norm=norm, cmap=cmap) 

351 m.set_array([]) 

352 

353 cbar = plt.colorbar(m, ax=ax, orientation='vertical', shrink=.7) 

354 cbar.set_label(cbar_label)