Coverage for gwcelery/tasks/skymaps.py: 89%

1"""Annotations for sky maps.""" 

2import os 

3import tempfile 

4 

5from astropy import table 

6from astropy.io import fits 

7from celery import group 

8from celery.exceptions import Ignore 

9from ligo.skymap.tool import (ligo_skymap_flatten, ligo_skymap_from_samples, 

10 ligo_skymap_plot, ligo_skymap_plot_coherence, 

11 ligo_skymap_plot_volume, ligo_skymap_unflatten) 

12from matplotlib import pyplot as plt 

13 

14from .. import app 

15from ..jinja import env 

16from ..util.cmdline import handling_system_exit 

17from ..util.matplotlib import closing_figures 

18from ..util.tempfile import NamedTemporaryFile 

19from . import external_skymaps, gracedb, igwn_alert 

20 

21 

22@app.task(ignore_result=True, shared=False) 

23def annotate_fits_tuple(filecontents_versioned_filename, graceid, tags): 

24 filecontents, versioned_filename = filecontents_versioned_filename 

25 annotate_fits(filecontents, versioned_filename, graceid, tags) 

26 

27 

28@app.task(ignore_result=True, shared=False) 

29def annotate_fits(filecontents, versioned_filename, graceid, tags): 

30 """Perform annotations on a sky map. 

31 

32 This function downloads a FITS file and then generates and uploads all 

33 derived images as well as an HTML dump of the FITS header. 

34 """ 

35 multiorder_extension = '.multiorder.fits' 

36 flat_extension = '.fits' 

37 

38 if multiorder_extension in versioned_filename: 

39 extension = multiorder_extension 

40 multiorder = True 

41 else: 

42 extension = flat_extension 

43 multiorder = False 

44 

45 filebase, _, _ = versioned_filename.partition(extension) 

46 

47 header_msg = ( 

48 'FITS headers for <a href="/api/superevents/{graceid}/files/' 

49 '{versioned_filename}">{versioned_filename}</a>').format( 

50 graceid=graceid, versioned_filename=versioned_filename) 

51 allsky_msg = ( 

52 'Mollweide projection of <a href="/api/superevents/{graceid}/files/' 

53 '{versioned_filename}">{versioned_filename}</a>').format( 

54 graceid=graceid, versioned_filename=versioned_filename) 

55 volume_msg = ( 

56 'Volume rendering of <a href="/api/superevents/{graceid}/files/' 

57 '{versioned_filename}">{versioned_filename}</a>').format( 

58 graceid=graceid, versioned_filename=versioned_filename) 

59 flatten_msg = ( 

60 'Flat-resolution FITS file created from ' 

61 '<a href="/api/superevents/{graceid}/files/' 

62 '{versioned_filename}">{versioned_filename}</a>').format( 

63 graceid=graceid, versioned_filename=versioned_filename) 

64 

65 group( 

66 fits_header.s(versioned_filename) 

67 | 

68 gracedb.upload.s( 

69 filebase + '.html', graceid, header_msg, tags), 

70 

71 plot_allsky.s() 

72 | 

73 gracedb.upload.s( 

74 filebase + '.png', graceid, allsky_msg, tags), 

75 

76 annotate_fits_volume.s( 

77 filebase + '.volume.png', graceid, volume_msg, tags), 

78 

79 *( 

80 [ 

81 flatten.s(f'{filebase}.fits.gz') 

82 | 

83 gracedb.upload.s( 

84 f'{filebase}.fits.gz', graceid, flatten_msg, tags) 

85 ] if multiorder else [] 

86 ) 

87 ).delay(filecontents) 

88 

89 

90def is_3d_fits_file(filecontents): 

91 """Determine if a FITS file has distance information.""" 

92 with NamedTemporaryFile(content=filecontents) as fitsfile: 

93 return 'DISTNORM' in table.Table.read(fitsfile.name).colnames 

94 

95 

96@app.task(ignore_result=True, shared=False) 

97def annotate_fits_volume(filecontents, *args): 

98 """Perform annotations that are specific to 3D sky maps.""" 

99 if is_3d_fits_file(filecontents): 

100 ( 

101 plot_volume.s(filecontents) 

102 | 

103 gracedb.upload.s(*args) 

104 ).apply_async() 

105 

106 

107@app.task(shared=False) 

108def fits_header(filecontents, filename): 

109 """Dump FITS header to HTML.""" 

110 template = env.get_template('fits_header.jinja2') 

111 with NamedTemporaryFile(content=filecontents) as fitsfile, \ 

112 fits.open(fitsfile.name) as hdus: 

113 return template.render(filename=filename, hdus=hdus) 

114 

115 

116@app.task(shared=False) 

117@closing_figures() 

118def plot_allsky(filecontents, ra=None, dec=None): 

119 """Plot a Mollweide projection of a sky map using the command-line tool 

120 :doc:`ligo-skymap-plot <ligo.skymap:tool/ligo_skymap_plot>`. 

121 """ 

122 # Explicitly use a non-interactive Matplotlib backend. 

123 plt.switch_backend('agg') 

124 

125 with NamedTemporaryFile(mode='rb', suffix='.png') as pngfile, \ 

126 NamedTemporaryFile(content=filecontents) as fitsfile, \ 

127 handling_system_exit(): 

128 if ra is not None and dec is not None: 

129 ligo_skymap_plot.main([fitsfile.name, '-o', pngfile.name, 

130 '--annotate', '--radec', str(ra), str(dec)]) 

131 else: 

132 ligo_skymap_plot.main([fitsfile.name, '-o', pngfile.name, 

133 '--annotate', '--contour', '50', '90']) 

134 return pngfile.read() 

135 

136 

137@app.task(priority=1, queue='openmp', shared=False) 

138@closing_figures() 

139def plot_volume(filecontents): 

140 """Plot a 3D volume rendering of a sky map using the command-line tool 

141 :doc:`ligo-skymap-plot-volume <ligo.skymap:tool/ligo_skymap_plot_volume>`. 

142 """ 

143 # Explicitly use a non-interactive Matplotlib backend. 

144 plt.switch_backend('agg') 

145 

146 with NamedTemporaryFile(mode='rb', suffix='.png') as pngfile, \ 

147 NamedTemporaryFile(content=filecontents) as fitsfile, \ 

148 handling_system_exit(): 

149 ligo_skymap_plot_volume.main([fitsfile.name, '-o', 

150 pngfile.name, '--annotate']) 

151 return pngfile.read() 

152 

153 

154@app.task(shared=False, queue='highmem') 

155def flatten(filecontents, filename): 

156 """Convert a HEALPix FITS file from multi-resolution UNIQ indexing to the 

157 more common IMPLICIT indexing using the command-line tool 

158 :doc:`ligo-skymap-flatten <ligo.skymap:tool/ligo_skymap_flatten>`. 

159 """ 

160 with NamedTemporaryFile(content=filecontents) as infile, \ 

161 tempfile.TemporaryDirectory() as tmpdir, \ 

162 handling_system_exit(): 

163 outfilename = os.path.join(tmpdir, filename) 

164 ligo_skymap_flatten.main([infile.name, outfilename]) 

165 return open(outfilename, 'rb').read() 

166 

167 

168@app.task(shared=False, queue='highmem') 

169def unflatten(filecontents, filename): 

170 """Convert a HEALPix FITS file to multi-resolution UNIQ indexing from 

171 the more common IMPLICIT indexing using the command-line tool 

172 :doc:`ligo-skymap-unflatten <ligo.skymap:tool/ligo_skymap_unflatten>`. 

173 """ 

174 with NamedTemporaryFile(content=filecontents) as infile, \ 

175 tempfile.TemporaryDirectory() as tmpdir, \ 

176 handling_system_exit(): 

177 outfilename = os.path.join(tmpdir, filename) 

178 ligo_skymap_unflatten.main([infile.name, outfilename]) 

179 return open(outfilename, 'rb').read() 

180 

181 

182@app.task(shared=False, queue='multiprocessing') 

183def skymap_from_samples(samplefilecontents, superevent_id, instruments): 

184 """Generate multi-resolution FITS file from samples.""" 

185 with NamedTemporaryFile(content=samplefilecontents) as samplefile, \ 

186 tempfile.TemporaryDirectory() as tmpdir, \ 

187 handling_system_exit(): 

188 ligo_skymap_from_samples.main([ 

189 '-j', '--seed', '150914', '--maxpts', '5000', '--objid', 

190 superevent_id, '--instruments', *instruments, '-o', tmpdir, 

191 samplefile.name]) 

192 with open(os.path.join(tmpdir, 'skymap.fits'), 'rb') as f: 

193 return f.read() 

194 

195 

196@app.task(shared=False) 

197@closing_figures() 

198def plot_coherence(filecontents): 

199 """IGWN alert handler to plot the coherence Bayes factor. 

200 

201 Parameters 

202 ---------- 

203 contents : str, bytes 

204 The contents of the FITS file. 

205 

206 Returns 

207 ------- 

208 png : bytes 

209 The contents of a PNG file. 

210 

211 Notes 

212 ----- 

213 Under the hood, this just calls :meth:`plot_bayes_factor`. 

214 

215 """ 

216 # Explicitly use a non-interactive Matplotlib backend. 

217 plt.switch_backend('agg') 

218 

219 with NamedTemporaryFile(mode='rb', suffix='.png') as pngfile, \ 

220 NamedTemporaryFile(content=filecontents) as fitsfile: 

221 header = fits.getheader(fitsfile, 1) 

222 try: 

223 header['LOGBCI'] 

224 except KeyError: 

225 raise Ignore('FITS file does not have a LOGBCI field') 

226 ligo_skymap_plot_coherence.main([fitsfile.name, '-o', pngfile.name]) 

227 return pngfile.read() 

228 

229 

230@igwn_alert.handler('superevent', 

231 'mdc_superevent', 

232 shared=False) 

233def handle_plot_coherence(alert): 

234 """IGWN alert handler to plot and upload a visualization of the coherence 

235 Bayes factor. 

236 

237 Notes 

238 ----- 

239 Under the hood, this just calls :meth:`plot_coherence`. 

240 

241 """ 

242 if alert['alert_type'] != 'log': 

243 return # not for us 

244 if not alert['data']['filename'].endswith('.fits') or \ 

245 (alert['data']['filename'] == 

246 external_skymaps.COMBINED_SKYMAP_FILENAME_MULTIORDER): 

247 return # not for us 

248 

249 graceid = alert['uid'] 

250 f = alert['data']['filename'] 

251 v = alert['data']['file_version'] 

252 fv = '{},{}'.format(f, v) 

253 

254 ( 

255 gracedb.download.s(fv, graceid) 

256 | 

257 plot_coherence.s() 

258 | 

259 gracedb.upload.s( 

260 f.replace('.fits', '.coherence.png'), graceid, 

261 message=( 

262 f'Bayes factor for coherence vs. incoherence of ' 

263 f'<a href="/api/superevents/{graceid}/files/{fv}">' 

264 f'{fv}</a>'), 

265 tags=['sky_loc'] 

266 ) 

267 ).delay()