Coverage for gwcelery/tasks/first2years.py: 100%

1"""Create mock events from the "First Two Years" paper.""" 

2import io 

3import random 

4from importlib import resources 

5 

6import lal 

7import numpy as np 

8from celery import group 

9from celery.utils.log import get_task_logger 

10from ligo.lw import lsctables, utils 

11from ligo.lw.table import Table 

12from ligo.skymap.io.events.ligolw import ContentHandler 

13 

14from .. import app 

15from ..data import first2years as data_first2years 

16from . import first2years_external, gracedb 

17from .core import get_last 

18 

19log = get_task_logger(__name__) 

20 

21 

22def pick_coinc(): 

23 """Pick a coincidence from the "First Two Years" paper.""" 

24 with resources.files( 

25 data_first2years).joinpath('gstlal.xml.gz').open('rb') as f: 

26 xmldoc = utils.load_fileobj(f, contenthandler=ContentHandler) 

27 root, = xmldoc.childNodes 

28 

29 # Remove unneeded tables 

30 for name in ( 

31 'filter', # lsctables.FilterTable removed from ligo.lw 

32 lsctables.SegmentTable.tableName, 

33 lsctables.SegmentDefTable.tableName, 

34 lsctables.SimInspiralTable.tableName, 

35 lsctables.SummValueTable.tableName, 

36 lsctables.SearchSummVarsTable.tableName): 

37 root.removeChild(Table.get_table(xmldoc, name)) 

38 

39 coinc_inspiral_table = table = lsctables.CoincInspiralTable.get_table( 

40 xmldoc) 

41 

42 # Determine event with most recent sideral time 

43 gps_time_now = lal.GPSTimeNow() 

44 gmsts = np.asarray([lal.GreenwichMeanSiderealTime(_.end) for _ in table]) 

45 gmst_now = lal.GreenwichMeanSiderealTime(gps_time_now) 

46 div, rem = divmod(gmst_now - gmsts, 2 * np.pi) 

47 i = np.argmin(rem) 

48 new_gmst = div[i] * 2 * np.pi + gmsts[i] 

49 old_time = table[i].end 

50 new_time = lal.LIGOTimeGPS() 

51 result = lal.GreenwichMeanSiderealTimeToGPS(new_gmst, new_time) 

52 result.disown() 

53 del result 

54 delta_t = new_time - old_time 

55 target_coinc_event_id = int(table[i].coinc_event_id) 

56 

57 # Remove unneeded rows 

58 table[:] = [row for row in table 

59 if int(row.coinc_event_id) == target_coinc_event_id] 

60 target_end_time = table[0].end 

61 

62 coinc_table = table = lsctables.CoincTable.get_table(xmldoc) 

63 table[:] = [row for row in table 

64 if int(row.coinc_event_id) == target_coinc_event_id] 

65 

66 table = lsctables.CoincMapTable.get_table(xmldoc) 

67 table[:] = [row for row in table 

68 if int(row.coinc_event_id) == target_coinc_event_id] 

69 target_sngl_inspirals = frozenset(row.event_id for row in table) 

70 

71 sngl_inspiral_table = table = lsctables.SnglInspiralTable.get_table(xmldoc) 

72 table[:] = [row for row in table if row.event_id in target_sngl_inspirals] 

73 

74 table = lsctables.ProcessTable.get_table(xmldoc) 

75 table[:] = [row for row in table if row.program == 'gstlal_inspiral'] 

76 target_process_ids = frozenset(row.process_id for row in table) 

77 

78 table = lsctables.SearchSummaryTable.get_table(xmldoc) 

79 table[:] = [row for row in table if target_end_time in row.out_segment 

80 and row.process_id in target_process_ids] 

81 target_process_ids = frozenset(row.process_id for row in table) 

82 

83 table = lsctables.ProcessTable.get_table(xmldoc) 

84 table[:] = [row for row in table if row.process_id in target_process_ids] 

85 

86 table = lsctables.ProcessParamsTable.get_table(xmldoc) 

87 table[:] = [row for row in table if row.process_id in target_process_ids] 

88 

89 # Shift event times 

90 for row in coinc_inspiral_table: 

91 row.end += delta_t 

92 for row in sngl_inspiral_table: 

93 row.end += delta_t 

94 row.end_time_gmst = lal.GreenwichMeanSiderealTime(row.end) 

95 

96 # The old version of gstlal used to produce the "First Two Years" data set 

97 # stored likelihood in the coinc_event.likelihood column, but newer 

98 # versions store the *natural log* of the likelihood here. The p_astro 

99 # calculation requires this to be log likelihood. 

100 for row in coinc_table: 

101 row.likelihood = np.log(row.likelihood) 

102 

103 # Gstlal stores the template's SVD bank index in the Gamma1 column. 

104 # Fill this in so that we can calculate p_astro 

105 # (see :mod:`gwcelery.tasks.p_astro_gstlal`). 

106 for row in sngl_inspiral_table: 

107 row.Gamma1 = 16 

108 

109 coinc_xml = io.BytesIO() 

110 utils.write_fileobj(xmldoc, coinc_xml) 

111 return coinc_xml.getvalue() 

112 

113 

114def _jitter_snr(coinc_bytes): 

115 coinc_xml = io.BytesIO(coinc_bytes) 

116 xmldoc = utils.load_fileobj(coinc_xml, contenthandler=ContentHandler) 

117 

118 coinc_inspiral_table = lsctables.CoincInspiralTable.get_table(xmldoc) 

119 

120 # Add a tiny amount of jitter in SNR so that uploads have random 

121 # preferred event precedence. 

122 for row in coinc_inspiral_table: 

123 row.snr += random.gauss(0, 1e-9) 

124 

125 coinc_xml = io.BytesIO() 

126 utils.write_fileobj(xmldoc, coinc_xml) 

127 return coinc_xml.getvalue() 

128 

129 

130@app.task(shared=False) 

131def _vet_event(superevents): 

132 if superevents: 

133 gracedb.create_signoff.s( 

134 random.choice(['NO', 'OK']), 

135 'If this had been a real gravitational-wave event candidate, ' 

136 'then an on-duty scientist would have left a comment here on ' 

137 'data quality and the status of the detectors.', 

138 'ADV', superevents[0]['superevent_id'] 

139 ).apply_async() 

140 

141 

142@app.task(shared=False) 

143def _log_event_and_return_query(event): 

144 graceid = event['graceid'] 

145 log.info('uploaded as %s', graceid) 

146 return 'MDC event: {}'.format(graceid) 

147 

148 

149@app.on_after_finalize.connect 

150def setup_periodic_tasks(sender, **kwargs): 

151 """Register periodic tasks. 

152 

153 See 

154 https://docs.celeryproject.org/en/stable/userguide/periodic-tasks.html. 

155 """ 

156 sender.add_periodic_task(3600.0, upload_event) 

157 sender.add_periodic_task(7200.0, first2years_external.upload_snews_event) 

158 

159 

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

161def upload_event(): 

162 """Upload a random event from the "First Two Years" paper. 

163 

164 After 2 minutes, randomly either retract or confirm the event to send a 

165 retraction or initial notice respectively. 

166 """ 

167 coinc = pick_coinc() 

168 num = 16 if app.conf['mock_events_simulate_multiple_uploads'] else 1 

169 

170 ( 

171 group( 

172 gracedb.create_event.si( 

173 _jitter_snr(coinc), 'MDC', 'gstlal', 'CBC' 

174 ) for _ in range(num) 

175 ) 

176 | 

177 get_last.s() 

178 | 

179 _log_event_and_return_query.s() 

180 | 

181 gracedb.get_superevents.s().set(countdown=600) 

182 | 

183 _vet_event.s() 

184 ).apply_async() 

185 

186 # Return gpstime to create an external MDC event with similar time 

187 coinc_xml = io.BytesIO(coinc) 

188 xmldoc = utils.load_fileobj(coinc_xml, contenthandler=ContentHandler) 

189 coinc_inspiral_table = lsctables.CoincInspiralTable.get_table(xmldoc) 

190 for row in coinc_inspiral_table: 

191 # The first value always returns the correct gpstime 

192 return float(row.end)