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

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 has_ns, has_remnant, has_massgap = em_bright.source_classification_pe( 

111 filename, num_eos_draws=10000, eos_seed=0 

112 ) 

113 data = json.dumps({ 

114 'HasNS': has_ns, 

115 'HasRemnant': has_remnant, 

116 'HasMassGap': has_massgap 

117 }) 

118 return data 

119 

120 

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

122def source_properties(mass1, mass2, spin1z, spin2z, snr): 

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

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

125 lower mass gap for the detected event. 

126 

127 Parameters 

128 ---------- 

129 mass1 : float 

130 Primary mass in solar masses 

131 mass2 : float 

132 Secondary mass in solar masses 

133 spin1z : float 

134 Dimensionless primary aligned spin component 

135 spin2z : float 

136 Dimensionless secondary aligned spin component 

137 snr : float 

138 Signal to noise ratio 

139 

140 Returns 

141 ------- 

142 str 

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

144 and `HasMassGap`` probabilities 

145 

146 Examples 

147 -------- 

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

149 '{"HasNS": 1.0, "HasRemnant": 1.0, "HasMassGap"}' 

150 """ 

151 from ligo.em_bright import em_bright 

152 p_ns, p_em, p_mg = em_bright.source_classification( 

153 mass1, mass2, spin1z, spin2z, snr 

154 ) 

155 

156 data = json.dumps({ 

157 'HasNS': p_ns, 

158 'HasRemnant': p_em, 

159 'HasMassGap': p_mg 

160 }) 

161 return data