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): 

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 ligo_skymap_plot.main([fitsfile.name, '-o', pngfile.name, 

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

130 return pngfile.read() 

131 

132 

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

134@closing_figures() 

135def plot_volume(filecontents): 

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

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

138 """ 

139 # Explicitly use a non-interactive Matplotlib backend. 

140 plt.switch_backend('agg') 

141 

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

143 NamedTemporaryFile(content=filecontents) as fitsfile, \ 

144 handling_system_exit(): 

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

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

147 return pngfile.read() 

148 

149 

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

151def flatten(filecontents, filename): 

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

153 more common IMPLICIT indexing using the command-line tool 

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

155 """ 

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

157 tempfile.TemporaryDirectory() as tmpdir, \ 

158 handling_system_exit(): 

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

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

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

162 

163 

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

165def unflatten(filecontents, filename): 

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

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

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

169 """ 

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

171 tempfile.TemporaryDirectory() as tmpdir, \ 

172 handling_system_exit(): 

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

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

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

176 

177 

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

179def skymap_from_samples(samplefilecontents, superevent_id, instruments): 

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

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

182 tempfile.TemporaryDirectory() as tmpdir, \ 

183 handling_system_exit(): 

184 ligo_skymap_from_samples.main([ 

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

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

187 samplefile.name]) 

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

189 return f.read() 

190 

191 

192@app.task(shared=False) 

193@closing_figures() 

194def plot_coherence(filecontents): 

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

196 

197 Parameters 

198 ---------- 

199 contents : str, bytes 

200 The contents of the FITS file. 

201 

202 Returns 

203 ------- 

204 png : bytes 

205 The contents of a PNG file. 

206 

207 Notes 

208 ----- 

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

210 

211 """ 

212 # Explicitly use a non-interactive Matplotlib backend. 

213 plt.switch_backend('agg') 

214 

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

216 NamedTemporaryFile(content=filecontents) as fitsfile: 

217 header = fits.getheader(fitsfile, 1) 

218 try: 

219 header['LOGBCI'] 

220 except KeyError: 

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

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

223 return pngfile.read() 

224 

225 

226@igwn_alert.handler('superevent', 

227 'mdc_superevent', 

228 shared=False) 

229def handle_plot_coherence(alert): 

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

231 Bayes factor. 

232 

233 Notes 

234 ----- 

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

236 

237 """ 

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

239 return # not for us 

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

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

242 external_skymaps.COMBINED_SKYMAP_FILENAME_MULTIORDER): 

243 return # not for us 

244 

245 graceid = alert['uid'] 

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

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

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

249 

250 ( 

251 gracedb.download.s(fv, graceid) 

252 | 

253 plot_coherence.s() 

254 | 

255 gracedb.upload.s( 

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

257 message=( 

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

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

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

261 tags=['sky_loc'] 

262 ) 

263 ).delay()