1# (c) Stefan Countryman 2016-2018 

2 

3""" 

4More complex plotting commands for joint events in more human-readable formats, 

5e.g. the GW skymap plots from O2 with their overlaid crosses marking neutrino 

6positions. 

7""" 

8 

9# TODO actually import this into the healpix skymaps or otherwise make it 

10# available for joint plots 

11from llama.files.healpix.skymap import MollviewMixin 

12from llama.filehandler import FileHandler 

13from llama.filehandler.mixins import FileExtensionMixin 

14 

15DPI = 200 

16WIDTH = 6 # [inches] 

17HEIGHT = 4 # [inches] 

18OUTLINE_STROKE = 1.2 # thickness of outline surrounding text 

19TITLE_OUTLINE_STROKE = 2 # thickness of outline surrounding plot titles 

20OUTLINE_COLOR = (1, 1, 1, 1) # outlines of text and neutrino markers 

21LEFT_SHIFT_COEFF = 1/20. # quadratic curve dec labels away from meridian 

22LEFT_SHIFT_BASE = -20 # baseline shift from leftmost meridian [deg] 

23XMARGIN = 0.4 # additional margin in x (make sure text fits) [inches] 

24TOPMARGIN = -0.0 # additional margin at top [inches] 

25BOTTOMMARGIN = -0.0 # additional margin at bottom [inches] 

26NEUTRINO_COLOR = 'green' 

27 

28 

29def default_cmap(): 

30 """Get the matplotlib colormap.""" 

31 from matplotlib import cm 

32 # cmap = 'magma_r' # colormap to use (matplotlib.org/users/colormaps.html) 

33 # colormap to use (matplotlib.org/users/colormaps.html). 

34 # NOTE that due to a little bug in healpy 

35 # (stackoverflow.com/questions/34023932), you'll get a hideous background color 

36 # (corresponding to the min. value color for your cmap) behind the plot. this 

37 # gets fixed with the line after the next one. 

38 

39 # O2 classic color scheme 

40 cmap = cm.gist_heat_r 

41 

42 # cmap = cm.BuPu 

43 # cmap = cm.bone_r 

44 # cmap = cm.YlGnBu 

45 # cmap = cm.gnuplot2_r 

46 # cmap = cm.CMRmap_r 

47 cmap.set_under((0, 0, 0, 0)) # transparent background 

48 return cmap 

49 

50 

51DEC_X_OFFSET = -0.37 # [inches] 

52DEC_Y_OFFSET = -0.02 # [inches] 

53RA_X_OFFSET = 0 # [inches] 

54RA_Y_OFFSET = 0.06 # [inches] 

55N_X_OFFSET = 0.08 # [inches] 

56N_Y_OFFSET = 0.08 # [inches] 

57CENTRAL_LON = 180 # longitude to place at center of skymap [deg] 

58CENTRAL_LAT = 0 # latitude to place at center of skymap [deg] 

59PSI = 0 # additional rotation of skymap about center 

60DELTA_PARALLEL = 15 # space between parallels in the graticule [deg] 

61DELTA_MERIDIAN = 30 # space between meridians in the graticule [deg] 

62GRATICULE_COLOR = "#B0B0B0" 

63GRATICULE_LABEL_COLOR = (0.2, 0.2, 0.2) 

64TITLES = ( 

65 '1-detector Skymap (LHO)', 

66 '2-detector Skymap (LHO, LLO)', 

67 '3-detector Skymap (LHO, LLO, Virgo)', 

68) 

69DIRNAMES = ('1det', '2det', '3det') 

70DEFAULT_SCATTER_MARKER = "$\\bigodot$" # bullseye hides less GW density 

71MERIDIAN_FONT_SIZE = 11 

72FONT_SIZE = 14 # matplotlib font size 

73UNCERTAINTY_ALPHA = 0.4 # opacity in [0,1] for scatterplot uncertainty discs 

74DEFAULT_PLOT_EXTENSION = 'pdf' # file type to save plots as 

75 

76 

77def outline_effect(): 

78 """Get a ``matplotlib.patheffects.withStroke`` effect that outlines text 

79 nicely to improve plot readability.""" 

80 from matplotlib.patheffects import withStroke 

81 return withStroke( 

82 linewidth=OUTLINE_STROKE, 

83 foreground=OUTLINE_COLOR, 

84 ) 

85 

86 

87# pylint: disable=too-many-arguments 

88def plot_healpix( 

89 skymap, 

90 title=None, 

91 unit="Probability Density", 

92 central_longitude=CENTRAL_LON, 

93 central_latitude=CENTRAL_LAT, 

94 rotation_about_center=PSI, 

95 dpi=DPI, 

96 cbar=None, 

97 cmap=None, 

98 alpha=None, 

99 **kwargs 

100): 

101 """Plot a healpix skymap in a LLAMA-specific style. Works on skymap objects 

102 with a HEALpy-esque ``mollview`` method. Returns a figure object for this 

103 skymap plot. No colorbar by default. Extra kwargs are passed to 

104 ``skymap.mollview``. Define the opacity of the skymap with ``alpha`` 

105 (must be in [0,1]). In particular, use ``hold=True`` to use the current 

106 axis for this plot.""" 

107 import numpy as np 

108 import healpy as hp 

109 from matplotlib.patheffects import withStroke 

110 from matplotlib.transforms import ScaledTranslation 

111 from matplotlib import pyplot as plt 

112 if cmap is None: 

113 cmap = default_cmap() 

114 # make the fonts bigger than the default 10pt 

115 plt.rcParams['font.size'] = FONT_SIZE 

116 if alpha: 

117 if alpha < 0 or alpha > 1: 

118 raise ValueError("``alpha`` must be in [0,1].") 

119 # TODO create a custom cmap based off of ``cmap`` and ``alpha`` 

120 # arguments with ``alpha`` opacity and redefine ``cmap`` as that. 

121 skymap.mollview( 

122 unit=unit, 

123 rot=(central_longitude, central_latitude, rotation_about_center), 

124 cbar=cbar, 

125 cmap=cmap, 

126 **kwargs 

127 ) 

128 # size the figure 

129 fig = plt.gcf() 

130 fig.set_size_inches(WIDTH, HEIGHT) 

131 fig.set_dpi(dpi) 

132 if title is not None: 

133 plt.title( 

134 title, 

135 path_effects=[ 

136 withStroke( 

137 linewidth=TITLE_OUTLINE_STROKE, 

138 foreground='w', 

139 ), 

140 ], 

141 ) 

142 # add extra margin so that the declination labels fit 

143 x_1, x_2, y_1, y_2 = plt.axis() 

144 plt.axis((x_1-XMARGIN, x_2+XMARGIN, y_1-TOPMARGIN, y_2+BOTTOMMARGIN)) 

145 # Make a matplotlib translation to offset text labels by a bit so that 

146 # they don't cover up the scatter plot markers they are labeling. See: 

147 # https://matplotlib.org/users/transforms_tutorial.html 

148 transdata = plt.gca().transData 

149 dectrans = transdata + ScaledTranslation(DEC_X_OFFSET, DEC_Y_OFFSET, 

150 fig.dpi_scale_trans) 

151 ratrans = transdata + ScaledTranslation(RA_X_OFFSET, RA_Y_OFFSET, 

152 fig.dpi_scale_trans) 

153 # add a graticule 

154 hp.visufunc.graticule( 

155 DELTA_PARALLEL, 

156 DELTA_MERIDIAN, 

157 color=GRATICULE_COLOR, 

158 ) 

159 # label declinations 

160 for dec in np.arange(DELTA_PARALLEL-90, 90, DELTA_PARALLEL): 

161 hp.projtext( 

162 # curve dec labels away from leftmost meridian 

163 LEFT_SHIFT_BASE+abs(dec*LEFT_SHIFT_COEFF)**2, 

164 dec, 

165 # label dec, right-pad to avoid hitting plot 

166 str(dec) + '$\\degree$', 

167 lonlat=True, 

168 color=GRATICULE_LABEL_COLOR, 

169 transform=dectrans, 

170 horizontalalignment='right', 

171 verticalalignment='center', 

172 path_effects=[outline_effect()], 

173 ) 

174 # label right ascensions 

175 for right_ascension in np.arange(DELTA_MERIDIAN, 360, DELTA_MERIDIAN): 

176 hp.projtext( 

177 right_ascension, 

178 0, 

179 str(right_ascension) + '$\\degree$', 

180 lonlat=True, 

181 color=GRATICULE_LABEL_COLOR, 

182 transform=ratrans, 

183 horizontalalignment='center', 

184 verticalalignment='bottom', 

185 fontsize=MERIDIAN_FONT_SIZE, 

186 path_effects=[outline_effect()], 

187 ) 

188 return fig 

189 

190 

191# pylint: disable=too-many-arguments 

192def add_scatterplot( 

193 right_ascensions, 

194 declinations, 

195 marker=DEFAULT_SCATTER_MARKER, 

196 color=NEUTRINO_COLOR, 

197 zorder=1000, 

198 **kwargs 

199): 

200 """Add a scatterplot to the current HEALpy axis. Useful for plotting 

201 neutrinos etc.; remaining ``kwargs`` are passed to 

202 ``healpy.projscatter``. Returns the current figure.""" 

203 import numpy as np 

204 import healpy as hp 

205 from matplotlib.colors import ColorConverter 

206 from matplotlib.transforms import ScaledTranslation 

207 from matplotlib import pyplot as plt 

208 # make the fonts bigger than the default 10pt 

209 plt.rcParams['font.size'] = FONT_SIZE 

210 events = np.array([right_ascensions, declinations]) 

211 # canonicalize input colors as RGBA 

212 color = ColorConverter.to_rgba(color) 

213 # uncertainty_color = color[0:3] + (UNCERTAINTY_ALPHA,) 

214 

215 # add an 'x' at each event location 

216 hp.projscatter( 

217 events, 

218 lonlat=True, 

219 edgecolor=OUTLINE_COLOR, 

220 linewidth=0.5, 

221 facecolor=color, 

222 rasterized=False, 

223 marker=marker, 

224 s=(2*plt.rcParams['lines.markersize']) ** 2, # double default size 

225 zorder=zorder, 

226 **kwargs 

227 ) 

228 # Make a matplotlib translation to offset text labels by a bit so that 

229 # they don't cover up the scatter plot markers they are labeling. See: 

230 # https://matplotlib.org/users/transforms_tutorial.html 

231 fig = plt.gcf() 

232 transdata = plt.gca().transData 

233 ntrans = transdata + ScaledTranslation(N_X_OFFSET, N_Y_OFFSET, 

234 fig.dpi_scale_trans) 

235 # label events 

236 for i, event in enumerate(events.transpose()): 

237 hp.projtext( 

238 event[0], 

239 event[1], 

240 str(i+1), 

241 lonlat=True, 

242 color=color, 

243 transform=ntrans, 

244 zorder=zorder, 

245 path_effects=[outline_effect()], 

246 ) 

247 return fig 

248 

249 

250class JointSkymapScatter(FileHandler, MollviewMixin, FileExtensionMixin): 

251 """ 

252 Plot a base skymap but with a scatterplot added on top. This is a bit 

253 different from joint skymaps that contain information on pixelwise 

254 likelihood products since it is only meant to be a human-readable 

255 representation. 

256 

257 Required Class Attributes 

258 ------------------------- 

259 The following class attributes must be defined in subclasses, either 

260 manually or programmatically (see: ``FileHandler.set_class_attributes`` and 

261 its implementations in subclasses). 

262 

263 ``POINT_SOURCES`` 

264 ~~~~~~~~~~~~~~~~~ 

265 A dictionary mapping ``TriggerList`` and ``HEALPixPSF`` subclass instances 

266 (i.e. pointlike-source lists) to a dict of style properties that can be 

267 fed to ``matplotlib.pyplot.scatterplot`` as kwargs. In particular, you MUST 

268 define ``marker`` and ``color`` in order to distinguish between different 

269 types of sources. 

270 

271 ``BACKGROUND_SKYMAP`` 

272 ~~~~~~~~~~~~~~~~~~~~~ 

273 The ``HEALPixRepresentableFileHandler`` containing skymaps that should 

274 be used as the continuous background (against which point-sources will 

275 be plotted). 

276 """ 

277 

278 POINT_SOURCES = None 

279 BACKGROUND_SKYMAP = None 

280 

281 _REQUIRED = ("POINT_SOURCES", "BACKGROUND_SKYMAP") 

282 

283 @classmethod 

284 def set_class_attributes(cls, subclass): 

285 """See ``FileHandler.set_class_attributes``; this method additionally 

286 sets the ``DEPENDENCIES`` and ``FILENAME`` attributes based on 

287 ``subclass.POINT_SOURCES`` and ``subclass.BACKGROUND_SKYMAP``.""" 

288 subclass.DEPENDENCIES = tuple( 

289 [subclass.BACKGROUND_SKYMAP] + 

290 sorted( 

291 subclass.POINT_SOURCES.keys(), 

292 key=lambda s: s.DETECTORS[0].name 

293 ) 

294 ) 

295 assert all(len(d.DETECTORS) == 1 for d in subclass.DEPENDENCIES) 

296 subclass.FILENAME = 'scatterplot_{}.{}'.format( 

297 '-'.join( 

298 d.DETECTORS[0].abbrev.lower() 

299 for d in subclass.DEPENDENCIES 

300 ), 

301 subclass.FILEEXT, 

302 ) 

303 super().set_class_attributes(subclass) 

304 return subclass 

305 

306 def mollview(self, **kwargs): 

307 """Return a Mollweide projection of this skymap as a ``matplotlib`` 

308 figure. Remaining ``**kwargs`` are passed to ``healpy.mollview``.""" 

309 # TODO make sure we can handle MULTIPLE background skymaps. 

310 import numpy as np 

311 skymap = self.BACKGROUND_SKYMAP(self).get_healpix()[0] 

312 if 'title' not in kwargs: 

313 if skymap.title is not None: 

314 kwargs['title'] = skymap.title 

315 else: 

316 names = ([self.BACKGROUND_SKYMAP.DETECTORS[0].name] + 

317 [p.DETECTORS[0].name for p in self.POINT_SOURCES]) 

318 kwargs['title'] = f"{'/'.join(names)} {self.eventid}" 

319 fig = plot_healpix( 

320 skymap, 

321 unit=skymap.unit, 

322 **kwargs 

323 ) 

324 for filehandler, plot_properties in self.POINT_SOURCES.items(): 

325 point_sources = np.array(filehandler(self).source_locations()) 

326 if len(point_sources): 

327 add_scatterplot( 

328 right_ascensions=point_sources[:, 0], 

329 declinations=point_sources[:, 1], 

330 **plot_properties 

331 ) 

332 return fig 

333 

334 def _generate(self): 

335 self.mollview().savefig(self.fullpath)