Coverage for gwcelery/tasks/em_bright.py: 67%

1"""Qualitative source properties for CBC events.""" 

2import io 

3import json 

4 

5from celery.utils.log import get_task_logger 

6from matplotlib import pyplot as plt 

7 

8from .. import app 

9from ..util import NamedTemporaryFile, closing_figures 

10from . import gracedb, igwn_alert 

11from .p_astro import _format_prob 

12 

13log = get_task_logger(__name__) 

14 

15 

16@igwn_alert.handler('superevent', 

17 'mdc_superevent', 

18 shared=False) 

19def handle(alert): 

20 """IGWN alert handler to plot and upload a visualization 

21 of every ``em_bright.json``. 

22 """ 

23 filename = 'em_bright.json' 

24 graceid = alert['uid'] 

25 if alert['alert_type'] == 'log' and alert['data']['filename'] == filename: 

26 ( 

27 gracedb.download.si(filename, graceid) 

28 | 

29 plot.s() 

30 | 

31 gracedb.upload.s( 

32 filename.replace('.json', '.png'), 

33 graceid, 

34 message=( 

35 'Source properties visualization from ' 

36 '<a href="/api/superevents/{graceid}/files/{filename}">' 

37 '{filename}</a>').format( 

38 graceid=graceid, filename=filename), 

39 tags=['em_follow', 'em_bright', 'public'] 

40 ) 

41 ).delay() 

42 

43 

44@app.task(shared=False) 

45@closing_figures() 

46def plot(contents): 

47 """Make a visualization of the source properties. 

48 

49 Examples 

50 -------- 

51 .. plot:: 

52 :include-source: 

53 

54 >>> from gwcelery.tasks import em_bright 

55 >>> contents = '{"HasNS": 0.9137, "HasRemnant": 0.0, "HasMassGap": 0.0}' # noqa E501 

56 >>> em_bright.plot(contents) 

57 """ 

58 # Explicitly use a non-interactive Matplotlib backend. 

59 plt.switch_backend('agg') 

60 

61 properties = json.loads(contents) 

62 outfile = io.BytesIO() 

63 

64 properties = dict(sorted(properties.items(), reverse=True)) 

65 probs, names = list(properties.values()), list(properties.keys()) 

66 

67 with plt.style.context('seaborn-white'): 

68 fig, ax = plt.subplots(figsize=(3, 1)) 

69 ax.barh(names, probs) 

70 ax.barh(names, [1.0 - p for p in probs], 

71 color='lightgray', left=probs) 

72 for i, prob in enumerate(probs): 

73 ax.annotate(_format_prob(prob), (0, i), (4, 0), 

74 textcoords='offset points', ha='left', va='center') 

75 ax.set_xlim(0, 1) 

76 ax.set_xticks([]) 

77 ax.tick_params(left=False) 

78 for side in ['top', 'bottom', 'right']: 

79 ax.spines[side].set_visible(False) 

80 fig.tight_layout() 

81 fig.savefig(outfile, format='png') 

82 return outfile.getvalue() 

83 

84 

85@app.task(shared=False, queue='em-bright') 

86def em_bright_posterior_samples(posterior_file_content): 

87 """Returns the probability of having a NS component and remnant 

88 using Bilby posterior samples. 

89 

90 Parameters 

91 ---------- 

92 posterior_file_content : hdf5 posterior file content 

93 

94 Returns 

95 ------- 

96 str 

97 JSON formatted string storing ``HasNS``, ``HasRemnant``, 

98 and ``HasMassGap`` probabilities 

99 

100 Examples 

101 -------- 

102 >>> em_bright_posterior_samples(GraceDb().files('S190930s', 

103 ... 'Bilby.posterior_samples.hdf5').read()) 

104 {"HasNS": 0.014904901243599122, "HasRemnant": 0.0, "HasMassGap": 0.0} 

105 

106 """ 

107 from ligo.em_bright import em_bright 

108 with NamedTemporaryFile(content=posterior_file_content) as samplefile: 

109 filename = samplefile.name 

110 r = em_bright.source_classification_pe( 

111 filename, num_eos_draws=10000, eos_seed=0 

112 ) 

113 has_ssm, has_ns, has_remnant, has_massgap = r 

114 data = json.dumps({ 

115 'HasNS': has_ns, 

116 'HasRemnant': has_remnant, 

117 'HasMassGap': has_massgap, 

118 'HasSSM': has_ssm 

119 }) 

120 return data 

121 

122 

123@app.task(shared=False, queue='em-bright') 

124def source_properties(mass1, mass2, spin1z, spin2z, snr, 

125 pipeline='gstlal', search='allsky'): 

126 """Returns the probability of having a NS component, the probability of 

127 having non-zero disk mass, and the probability of any component being the 

128 lower mass gap for the detected event. 

129 

130 Parameters 

131 ---------- 

132 mass1 : float 

133 Primary mass in solar masses 

134 mass2 : float 

135 Secondary mass in solar masses 

136 spin1z : float 

137 Dimensionless primary aligned spin component 

138 spin2z : float 

139 Dimensionless secondary aligned spin component 

140 snr : float 

141 Signal to noise ratio 

142 pipeline_search : tuple 

143 The pipeline and the search as a tuple. This is used 

144 to select the appropriate classifiers in ``ligo.em-bright`` 

145 for ``SSM`` search only. This is unused for ``AllSky`` 

146 searches. 

147 

148 Returns 

149 ------- 

150 str 

151 JSON formatted string storing ``HasNS``, ``HasRemnant``, 

152 ``HasMassGap`` probabilities for ``AllSky`` searches, and 

153 ``HasSSM``, ``HasNS``, ``HasMassGap`` probabilities for 

154 ``SSM`` searches. 

155 

156 Examples 

157 -------- 

158 >>> em_bright.source_properties(2.0, 1.0, 0.0, 0.0, 10.) 

159 '{"HasNS": 1.0, "HasRemnant": 1.0, "HasMassGap": 0.0}' 

160 >>> em_bright.source_properties(2.0, 1.0, 0.0, 0.0, 10., 

161 ... pipeline='gstlal', search='ssm') 

162 '{"HasSSM": 0.52, "HasNS": 0.9199999999999999, "HasMassGap": 0.08}' 

163 """ 

164 from ligo.em_bright import em_bright 

165 if search == 'ssm': 

166 chirp_mass = (mass1 * mass2) ** (3. / 5.) 

167 chirp_mass /= (mass1 + mass2) ** (1. / 5.) 

168 p_ssm, p_ns, p_mg = em_bright.source_classification_ssm( 

169 mass1, mass2, spin1z, spin2z, chirp_mass, 

170 snr, pipeline 

171 ) 

172 data = json.dumps({ 

173 'HasSSM': p_ssm, 

174 'HasNS': p_ns, 

175 'HasMassGap': p_mg, 

176 }) 

177 else: 

178 p_ns, p_em, p_mg = em_bright.source_classification( 

179 mass1, mass2, spin1z, spin2z, snr 

180 ) 

181 data = json.dumps({ 

182 'HasNS': p_ns, 

183 'HasRemnant': p_em, 

184 'HasMassGap': p_mg, 

185 }) 

186 return data