Coverage for event_normalizer/parsers.py: 45%

51 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-28 13:07 +0000

1"""Parsers for different data sources.""" 

2 

3from datetime import UTC, datetime 

4from typing import Any 

5 

6from .models import ( 

7 EventBurst, 

8 EventCoinc, 

9 EventEMBright, 

10 EventLabels, 

11 EventLinks, 

12 EventMLyBurst, 

13 EventPastro, 

14 EventSingleInspiral, 

15 NormalizedEvent, 

16) 

17from .transformers import Transformers 

18 

19 

20# ============================================================================ 

21# GraceDB, Kafka and File Parsers 

22# ============================================================================ 

23def parse_gracedb_event(data: dict[str, Any]) -> NormalizedEvent: 

24 """ 

25 Parse GraceDB REST API event response. 

26 

27 Extracts common fields and pipeline-specific data from extra_attributes. 

28 """ 

29 

30 uid = data.get("graceid") 

31 if not uid: 

32 raise ValueError("Missing 'graceid' in event data") 

33 

34 # Initialize a Transformer to deal with nested extra_attributes 

35 transformer = Transformers( 

36 group=data.get("group"), 

37 pipeline=data.get("pipeline"), 

38 search=data.get("search"), 

39 event=data, 

40 ) 

41 

42 # Get channel names 

43 channels_dict = transformer.transform_channels() 

44 

45 # Get entries from CoincInspiral table 

46 coinc_dict = transformer.transform_coinc_table( 

47 [ 

48 "mass", 

49 "mchirp", 

50 "minimum_duration", 

51 "snr", 

52 "ifos", 

53 "end_time", 

54 "end_time_ns", 

55 "false_alarm_rate", 

56 "combined_far", 

57 ] 

58 ) 

59 

60 # Get entries from MultiBurst table 

61 burst_dict = transformer.transform_burst_table( 

62 [ 

63 "mchirp", 

64 "snr", 

65 "duration", 

66 "ifos", 

67 "start_time", 

68 "start_time_ns", 

69 "strain", 

70 "peak_time", 

71 "peak_time_ns", 

72 "central_freq", 

73 "bandwidth", 

74 "amplitude", 

75 "confidence", 

76 "false_alarm_rate", 

77 "ligo_axis_ra", 

78 "ligo_axis_dec", 

79 "ligo_angle", 

80 "ligo_angle_sig", 

81 "single_ifo_times", 

82 "code", 

83 ] 

84 ) 

85 

86 # Get entries from MLyBurst table 

87 mly_dict = transformer.transform_mly_table( 

88 [ 

89 "central_freq", 

90 "bandwidth", 

91 "central_time", 

92 "detection_statistic", 

93 "duration", 

94 "bbh", 

95 "sglf", 

96 "sghf", 

97 "background", 

98 "glitch", 

99 "freq_correlation", 

100 "mass1", 

101 "mass2", 

102 "mtotal", 

103 "mchirp", 

104 "spin1z", 

105 "spin2z", 

106 "end_time", 

107 "end_time_ns", 

108 "template_duration", 

109 "SNR", 

110 "scores.coherency", 

111 "scores.coincidence", 

112 "scores.combined", 

113 ] 

114 ) 

115 

116 # Get SingleInspiral data per IFO 

117 single_inspiral_dict = transformer.transform_single_inspiral( 

118 [ 

119 "search", 

120 "end_time", 

121 "end_time_ns", 

122 "end_time_gmst", 

123 "impulse_time", 

124 "impulse_time_ns", 

125 "template_duration", 

126 "event_duration", 

127 "amplitude", 

128 "eff_distance", 

129 "coa_phase", 

130 "mass1", 

131 "mass2", 

132 "mchirp", 

133 "mtotal", 

134 "eta", 

135 "kappa", 

136 "chi", 

137 "tau0", 

138 "tau2", 

139 "tau3", 

140 "tau4", 

141 "tau5", 

142 "ttotal", 

143 "psi0", 

144 "psi3", 

145 "alpha", 

146 "alpha1", 

147 "alpha2", 

148 "alpha3", 

149 "alpha4", 

150 "alpha5", 

151 "alpha6", 

152 "beta", 

153 "f_final", 

154 "snr", 

155 "chisq", 

156 "chisq_dof", 

157 "bank_chisq", 

158 "bank_chisq_dof", 

159 "cont_chisq", 

160 "cont_chisq_dof", 

161 "sigmasq", 

162 "rsqveto_duration", 

163 "Gamma0", 

164 "Gamma1", 

165 "Gamma2", 

166 "Gamma3", 

167 "Gamma4", 

168 "Gamma5", 

169 "Gamma6", 

170 "Gamma7", 

171 "Gamma8", 

172 "Gamma9", 

173 "spin1x", 

174 "spin1y", 

175 "spin1z", 

176 "spin2x", 

177 "spin2y", 

178 "spin2z", 

179 ] 

180 ) 

181 

182 # Get labels 

183 labels = transformer.transform_labels() 

184 

185 # Get instruments (convert None to empty list for model validation) 

186 instruments = transformer.transform_instruments(data.get("instruments")) or [] 

187 

188 # Get links 

189 links = transformer.transform_links() 

190 

191 # Define NormalizedEvent 

192 event = NormalizedEvent( 

193 # Mandatory 

194 uid=uid, 

195 group=data.get("group"), 

196 pipeline=data.get("pipeline"), 

197 search=data.get("search"), 

198 far=data.get("far"), 

199 far_is_upper_limit=data.get("far_is_upper_limit"), 

200 instruments=instruments, 

201 H1_channel=channels_dict.get("H1_channel", "None"), 

202 L1_channel=channels_dict.get("L1_channel", "None"), 

203 V1_channel=channels_dict.get("V1_channel", "None"), 

204 K1_channel=channels_dict.get("K1_channel", "None"), 

205 reporting_latency=data.get("reporting_latency"), 

206 gpstime=data.get("gpstime"), 

207 last_updated=datetime.now(UTC), 

208 # Optional root fields 

209 alert_type=data.get("alert_type"), 

210 submitter=data.get("submitter"), 

211 offline=data.get("offline"), 

212 nevents=data.get("nevents"), 

213 likelihood=data.get("likelihood"), 

214 superevent=data.get("superevent"), 

215 created=data.get("created"), 

216 processing_status=data.get("processing_status"), 

217 content_id=data.get("content_id"), 

218 # CBC specific 

219 coinc=EventCoinc(**coinc_dict) if coinc_dict else None, 

220 # Burst specific 

221 burst=EventBurst(**burst_dict) if burst_dict else None, 

222 # MLy specific 

223 mly=EventMLyBurst(**mly_dict) if mly_dict else None, 

224 # Single Inspiral per IFO 

225 single_H1=EventSingleInspiral(**single_inspiral_dict["H1"]) 

226 if "H1" in single_inspiral_dict 

227 else None, 

228 single_L1=EventSingleInspiral(**single_inspiral_dict["L1"]) 

229 if "L1" in single_inspiral_dict 

230 else None, 

231 single_V1=EventSingleInspiral(**single_inspiral_dict["V1"]) 

232 if "V1" in single_inspiral_dict 

233 else None, 

234 single_K1=EventSingleInspiral(**single_inspiral_dict["K1"]) 

235 if "K1" in single_inspiral_dict 

236 else None, 

237 # Labels 

238 labels=EventLabels(**labels) if labels else None, 

239 # Links 

240 links=EventLinks(**links) if links else None, 

241 ) 

242 

243 return event 

244 

245 

246def parse_pastro(data: dict[str, Any]) -> EventPastro: 

247 """ 

248 Parse a p_astro.json payload. 

249 

250 Expected format: 

251 

252 { 

253 "Terrestrial": 0.999996409, 

254 "BNS": 0.000003591, 

255 "BBH": 0, 

256 "NSBH": 0 

257 } 

258 """ 

259 if not isinstance(data, dict): 

260 raise ValueError("p_astro payload must be a dictionary") 

261 

262 required_fields = ( 

263 "Terrestrial", 

264 "BNS", 

265 "BBH", 

266 "NSBH", 

267 ) 

268 

269 missing_fields = [field for field in required_fields if field not in data] 

270 

271 if missing_fields: 

272 raise ValueError( 

273 "p_astro payload is missing required fields: " + ", ".join(missing_fields) 

274 ) 

275 

276 return EventPastro( 

277 p_astro_terrestrial=data["Terrestrial"], 

278 p_astro_bns=data["BNS"], 

279 p_astro_bbh=data["BBH"], 

280 p_astro_nsbh=data["NSBH"], 

281 ) 

282 

283 

284def attach_pastro( 

285 event: NormalizedEvent, 

286 pastro: EventPastro | dict[str, Any], 

287) -> NormalizedEvent: 

288 """ 

289 Return a copy of a normalized event with p_astro information attached. 

290 

291 Parameters 

292 ---------- 

293 event: 

294 Existing normalized event. 

295 

296 pastro: 

297 Either an already validated EventPastro object or a raw p_astro.json 

298 dictionary. 

299 

300 Returns 

301 ------- 

302 NormalizedEvent 

303 A newly validated event containing the supplied p_astro values. 

304 

305 Notes 

306 ----- 

307 The source p_astro.json file does not contain an event identifier. 

308 The caller is responsible for associating the correct file with the 

309 correct NormalizedEvent. 

310 """ 

311 if isinstance(pastro, dict): 

312 pastro = parse_pastro(pastro) 

313 

314 if not isinstance(pastro, EventPastro): 

315 raise TypeError("pastro must be an EventPastro instance or a dictionary") 

316 

317 event_data = event.model_dump(mode="python") 

318 event_data["p_astro"] = pastro 

319 

320 return NormalizedEvent.model_validate(event_data) 

321 

322 

323def parse_embright(data: dict[str, Any]) -> EventEMBright: 

324 """ 

325 Parse a em_bright.json payload. 

326 

327 Expected format: 

328 

329 { 

330 "HasNS": 1, 

331 "HasRemnant": 1, 

332 "HasMassGap": 0.0678271766372453, 

333 "HasSSM": 0 

334 } 

335 """ 

336 if not isinstance(data, dict): 

337 raise ValueError("em_bright payload must be a dictionary") 

338 

339 required_fields = ("HasNS",) 

340 

341 missing_fields = [field for field in required_fields if field not in data] 

342 

343 if missing_fields: 

344 raise ValueError( 

345 "em_bright payload is missing required fields: " + ", ".join(missing_fields) 

346 ) 

347 

348 return EventEMBright( 

349 em_bright_has_ns=data["HasNS"], 

350 em_bright_has_remnant=data.get("HasRemnant"), 

351 em_bright_has_mass_gap=data.get("HasMassGap"), 

352 em_bright_has_ssm=data.get("HasSSM"), 

353 ) 

354 

355 

356def attach_embright( 

357 event: NormalizedEvent, 

358 embright: EventEMBright | dict[str, Any], 

359) -> NormalizedEvent: 

360 """ 

361 Return a copy of a normalized event with em_bright information attached. 

362 

363 Parameters 

364 ---------- 

365 event: 

366 Existing normalized event. 

367 

368 embright: 

369 Either an already validated EventEMBright object or a raw em_bright.json 

370 dictionary. 

371 

372 Returns 

373 ------- 

374 NormalizedEvent 

375 A newly validated event containing the supplied em_bright values. 

376 

377 Notes 

378 ----- 

379 The source em_bright.json file does not contain an event identifier. 

380 The caller is responsible for associating the correct file with the 

381 correct NormalizedEvent. 

382 """ 

383 if isinstance(embright, dict): 

384 embright = parse_embright(embright) 

385 

386 if not isinstance(embright, EventEMBright): 

387 raise TypeError("embright must be an EventEMBright instance or a dictionary") 

388 

389 event_data = event.model_dump(mode="python") 

390 event_data["em_bright"] = embright 

391 

392 return NormalizedEvent.model_validate(event_data) 

393 

394 

395# TODO: parse_kafka_alert