1# (c) Stefan Countryman, 2019 

2 

3""" 

4Utilities for working with LVC skymaps 

5""" 

6 

7import re 

8import datetime 

9from dateutil.parser import parse 

10from llama.com.gracedb import GraceDb 

11 

12SKYMAP_BASENAMES = { 

13 'cWB', 

14 'bayestar', 

15 'lalinference', 

16 'lalinference.offline', 

17 'LALInference', 

18 'LALInference.offline', 

19} 

20# keep this in the order you want to check in; first match is returned. being 

21# lazy, sorry. 

22SKYMAP_FILE_EXTENSIONS = ( 

23 '.multiorder.fits', 

24 '.fits.gz', 

25 '.fits', 

26) 

27SKYMAP_FILENAMES = {name+ext for name in SKYMAP_BASENAMES 

28 for ext in SKYMAP_FILE_EXTENSIONS} 

29 

30 

31class SkymapFilename(str): 

32 """ 

33 A skymap filename on GraceDB with built-in parsing of skymap file along 

34 with convenience methods used for selecting and canonicalizing skymap 

35 filenames. 

36 """ 

37 

38 def __init__(self, string): 

39 """Same as ``str.__init__`` but with checks on input filename 

40 validity.""" 

41 assert len(self.split(',')) in {1, 2}, "Filename can have only 1 comma" 

42 fname = self.filename 

43 if not any(fname.endswith(ext) for ext in SKYMAP_FILE_EXTENSIONS): 

44 raise TypeError(("file extension must be one of {}; instead, " 

45 "got: {}").format(SKYMAP_FILE_EXTENSIONS, string)) 

46 ver = self.version 

47 assert ver is None or ver >= 0 

48 super().__init__() 

49 

50 @property 

51 def filename(self): 

52 """The filename of this skymap, including the file extension, with the 

53 version number removed. 

54 

55 >>> skymap = SkymapFilename("bayestar.multiorder.fits") 

56 >>> skymap.filename 

57 'bayestar.multiorder.fits' 

58 

59 >>> skymap0 = SkymapFilename("bayestar.multiorder.fits,0") 

60 >>> skymap.filename 

61 'bayestar.multiorder.fits' 

62 """ 

63 if ',' in self: 

64 return type(self)(self.split(',')[0]) 

65 return self 

66 

67 @property 

68 def version(self): 

69 """The version of this skymap filename as an integer (since many 

70 versions of the same filename can be uploaded). Returns ``None`` if not 

71 specified. 

72 

73 >>> skymap = SkymapFilename("bayestar.multiorder.fits") 

74 >>> skymap.version == None 

75 True 

76 

77 >>> skymap0 = SkymapFilename("bayestar.multiorder.fits,0") 

78 >>> skymap0.version 

79 0 

80 """ 

81 if ',' in self: 

82 return int(self.split(',')[1]) 

83 return None 

84 

85 @property 

86 def extension(self): 

87 """The file extension, drawn from ``SKYMAP_FILE_EXTENSIONS``. 

88 

89 >>> skymap = SkymapFilename("bayestar.multiorder.fits") 

90 >>> skymap.extension 

91 '.multiorder.fits' 

92 

93 >>> skymap0 = SkymapFilename("bayestar.multiorder.fits,0") 

94 >>> skymap0.extension 

95 '.multiorder.fits' 

96 """ 

97 for ext in SKYMAP_FILE_EXTENSIONS: 

98 if self.filename.endswith(ext): 

99 return ext 

100 raise AssertionError("Should not have been able to initialize a " 

101 "``SkymapFilename`` called ", str(self)) 

102 

103 @property 

104 def basename(self): 

105 """Remove the filename extension and just get the base name. Since this 

106 base name corresponds to the information contained in the skymap (vs. 

107 the format of that information, given by ``SkymapFilename.extension``), 

108 this should be a good indicator of the contents of the skymap. 

109 

110 >>> skymap = SkymapFilename("bayestar.multiorder.fits") 

111 >>> skymap.basename 

112 'bayestar' 

113 

114 >>> skymap0 = SkymapFilename("bayestar.multiorder.fits,0") 

115 >>> skymap0.basename 

116 'bayestar' 

117 """ 

118 return str(self.filename[:-len(self.extension)]) 

119 

120 def canonicalize(self, graceid, canonical_date=None): 

121 """Return the canonicalized filename (including version information) 

122 for this skymap filename for the given ``graceid`` at the given 

123 ``date`` (instance of ``datetime.datetime`` or ``None`` to use the 

124 current time). If this filename already contains a version, just return 

125 ``self``. Queries GraceDB to resolve this information. Raises a 

126 ``ligo.gracedb.rest.HTTPError`` if a GraceDB query fails to connect. 

127 

128 For this event, a second version (v1) of ``bayestar.fits`` was created 

129 at 2019-05-10 04:03:40 UTC; canonicalizing at an earlier date should 

130 return v0 (the original) while a later date should return v1 (the 

131 update): 

132 

133 >>> from dateutil.parser import parse 

134 >>> skymap = SkymapFilename("bayestar.fits") 

135 >>> skymap.canonicalize("S190510g", parse("2019-05-10 04:03:30 UTC")) 

136 'bayestar.fits,0' 

137 >>> skymap.canonicalize("S190510g", parse("2019-05-10 04:03:50 UTC")) 

138 'bayestar.fits,1' 

139 

140 If the skymap name was already canonicalized, then the original value 

141 will be returned: 

142 

143 >>> canonical = SkymapFilename("bayestar.fits,0") 

144 >>> canonical == canonical.canonicalize("foo") 

145 True 

146 

147 If the file does not exist at the requested time, then the version 

148 should be assumed to be the zeroth version (since this will be the 

149 first available version): 

150 

151 >>> nonexistent = SkymapFilename("bayestar.fits") 

152 >>> unix0 = datetime.datetime.fromtimestamp(0) 

153 >>> nonexistent.canonicalize("S190510g", unix0) 

154 'bayestar.fits,0' 

155 """ 

156 if ',' in self: 

157 return self 

158 if canonical_date is None: 

159 canonical_date = datetime.datetime.utcnow() 

160 logs = GraceDb().logs(graceid).json()['log'] 

161 versions = sorted( 

162 [ 

163 l for l in logs 

164 if ( 

165 l['filename'] == self.filename and 

166 ( 

167 parse(l['created']).timestamp() < 

168 canonical_date.timestamp() 

169 ) 

170 ) 

171 ], 

172 key=lambda log: log['N'] 

173 ) 

174 if versions: 

175 ver = versions[-1]['file_version'] 

176 else: 

177 ver = 0 

178 return type(self)(self+','+str(ver)) 

179 

180 

181def skymap_filenames(filenames): 

182 """Return only filenames from ``filenames`` that could be skymap 

183 filenames.""" 

184 return [f for f in filenames 

185 if any(f == n or re.findall(n+r',\d+$', f, flags=re.MULTILINE) 

186 for n in SKYMAP_FILENAMES)] 

187 

188 

189def available_skymaps(graceid): 

190 """Get a list of available skymap logs in ascending order of creation time 

191 from GraceDB for a given ``graceid``.""" 

192 logs = GraceDb().logs(graceid).json() 

193 skymaps = [l for l in logs['log'] if skymap_filenames([l['filename']])] 

194 return sorted(skymaps, key=lambda s: s['N'])