1# (c) Stefan Countryman 2016-2018 

2 

3""" 

4``FileHandler`` classes related to working with skymaps provided by the 

5LIGO-Virgo Collaboration (LVC). 

6""" 

7 

8import logging 

9import tempfile 

10from llama.filehandler import ( 

11 GenerateOnceMixin, 

12 GenerationError, 

13 JSONFile, 

14) 

15from llama.filehandler.mixins import OnlineVetoMixin 

16from llama.files import healpix 

17from llama.utils import write_gzip 

18from llama.files.skymap_info import SkymapInfo 

19from llama.files.slack import SlackReceiptLlama 

20from llama.com.gracedb import GraceDb, HTTPError 

21from llama.files.lvc_skymap.utils import ( 

22 SKYMAP_FILENAMES, 

23 skymap_filenames, 

24 available_skymaps, 

25) 

26from llama import detectors 

27 

28LOGGER = logging.getLogger(__name__) 

29 

30 

31class SkymapNotFoundError(IOError): 

32 """ 

33 Raised when a valid skymap cannot be found or downloaded from 

34 GraceDb. 

35 """ 

36 

37 

38@healpix.LvcHEALPixSkyMapFileHandler.set_class_attributes 

39class LvcSkymapFits( 

40 GenerateOnceMixin, 

41 healpix.LvcHEALPixSkyMapFileHandler, 

42 OnlineVetoMixin, 

43): 

44 """The skymap suggested by the LVC Initial GCN Notice. 

45 """ 

46 

47 FILENAME = "lvc_skymap.fits.gz" 

48 DEPENDENCIES = (SkymapInfo,) 

49 DETECTORS = (detectors.LVC,) 

50 

51 def _generate(self): # pylint: disable=W0221 

52 """Get the skymap filename on GraceDB from lvc_initial.xml.""" 

53 client = GraceDb() 

54 LOGGER.debug('Loaded GraceDb client successfully.') 

55 gracedbfilename = SkymapInfo(self).skymap_filename 

56 fname = gracedbfilename.filename 

57 file_version = gracedbfilename.version 

58 graceid = SkymapInfo(self).graceid 

59 # sometimes filenames come with version suffixes in the form of a comma 

60 # followed by the version number, e.g. bayestar.fits.gz,0 

61 LOGGER.debug('Acquired fname: %s', fname) 

62 LOGGER.debug('File version: %s', file_version) 

63 try: 

64 # make sure the file exists before fetching it 

65 available = available_skymaps(graceid) 

66 name_match = [l for l in available if l['filename'] == fname] 

67 if not name_match: 

68 msg = (f"No versions of '{fname}' available for {graceid}. " 

69 f"Requested filename: '{gracedbfilename}'") 

70 LOGGER.debug(msg) 

71 raise GenerationError(msg) 

72 if file_version not in [l['file_version'] for l in name_match]: 

73 msg = (f"Version {file_version} of '{fname}' not available " 

74 f"for {graceid}. Requested filename: " 

75 f"'{gracedbfilename}' Available versions: {name_match}") 

76 LOGGER.debug(msg) 

77 raise GenerationError(msg) 

78 res = client.files(graceid, gracedbfilename) 

79 except HTTPError as exception: 

80 LOGGER.debug(("Could not fetch '{}' from gracedb for " 

81 "{}.").format(gracedbfilename, self)) 

82 raise GenerationError(exception.message) 

83 with tempfile.TemporaryFile(mode='rb+', suffix=self.FILENAME) as tmp: 

84 tmp.write(res.read()) 

85 tmp.seek(0) 

86 if not write_gzip(tmp, self.fullpath): 

87 LOGGER.warning('%s is not a gzip file! zipping now...', fname) 

88 if not (fname.endswith('.fits') or fname.endswith('.fits.gz')): 

89 LOGGER.error('Unexpected filename extension encountered while') 

90 LOGGER.error('downloading skymap: %s', fname) 

91 raise ValueError('bad Skymap name; must end with .fits or .gz ' + 

92 'but got: {}'.format(fname)) 

93 

94 

95@healpix.HEALPixSkyMapFileHandler.set_class_attributes 

96class LvcSkymapHdf5(healpix.HEALPixSkyMapFileHandler): 

97 """An HDF5-formatted copy of the initial LVC skymap. Loads more quickly.""" 

98 

99 DEPENDENCIES = (LvcSkymapFits,) 

100 DETECTORS = (detectors.LVC,) 

101 FILENAME = 'lvc_skymap.hdf5' 

102 

103 def _generate(self): # pylint: disable=W0221 

104 self.write_healpix(self.DEPENDENCIES[0](self).get_healpix()) 

105 

106 

107# pylint: disable=too-many-ancestors 

108@SlackReceiptLlama.upload_this() 

109@JSONFile.set_class_attributes 

110class LvcDistancesJson(JSONFile): 

111 """ 

112 A simple JSON file with the reconstructed distances to an event. 

113 Contains the mean reconstructed distance as ``distmean`` and the standard 

114 deviation as ``diststd``. 

115 """ 

116 

117 DEPENDENCIES = (LvcSkymapFits,) 

118 FILENAME = 'lvc_skymap_distances.json' 

119 

120 def _generate(self): 

121 import numpy as np 

122 from ligo.skymap.io.fits import read_sky_map 

123 skymap = read_sky_map(LvcSkymapFits(self).fullpath, moc=True) 

124 self._write_json([{"distmean": skymap.meta.get('distmean', np.nan), 

125 "diststd": skymap.meta.get('diststd', np.nan)}]) 

126 

127 @property 

128 def distmean(self): 

129 """The mean distance to the event; a probability-weighted average of 

130 reconstructed distances. Extracted from the original LVC skymap.""" 

131 try: 

132 return self.read_json()[0]['distmean'] 

133 except KeyError: 

134 return None 

135 

136 @property 

137 def diststd(self): 

138 """The probability-weighted standard deviation of the reconstructed 

139 distance to the event. Extracted from the original LVC skymap.""" 

140 try: 

141 return self.read_json()[0]['diststd'] 

142 except KeyError: 

143 return None 

144 

145 

146__all__ = [ 

147 'SkymapNotFoundError', 

148 'LvcSkymapFits', 

149 'LvcDistancesJson', 

150 'LvcSkymapHdf5', 

151 'SKYMAP_FILENAMES', 

152 'skymap_filenames', 

153 'available_skymaps', 

154]