Coverage for pesummary/gw/fetch.py: 76.5%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

3from pathlib import Path 

4 

5from pesummary.core.fetch import ( 

6 download_and_read_file, _download_authenticated_file 

7) 

8from pesummary.utils.decorators import deprecation 

9from gwosc.api import fetch_event_json 

10 

11__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"] 

12 

13 

14def fetch(url, download_kwargs={}, **kwargs): 

15 """Download and read files from LIGO authenticated URLs 

16 

17 Parameters 

18 ---------- 

19 url: str 

20 url you wish to download 

21 download_kwargs: dict, optional 

22 optional kwargs passed to _download_autheticated_file 

23 **kwargs: dict, optional 

24 additional kwargs passed to pesummary.io.read function 

25 """ 

26 if "idp" not in download_kwargs.keys(): 

27 download_kwargs["idp"] = "LIGO" 

28 return download_and_read_file( 

29 url, download_kwargs=download_kwargs, 

30 _function=_download_authenticated_file, **kwargs 

31 ) 

32 

33 

34def _DCC_url( 

35 event, type="posterior", catalog=None, sampling_rate=16384, format="gwf", 

36 duration=32, IFO="L1", version=None, 

37): 

38 """Return the url for posterior samples stored on the DCC for a given event 

39 

40 Parameters 

41 ---------- 

42 event: str 

43 name of the event you wish to return posterior samples for 

44 type: str, optional 

45 type of data you wish to query. Default "posterior" 

46 catalog: str, optional 

47 Name of catalog that hosts the event. Default None 

48 sampling_rate: int, optional 

49 sampling rate of strain data you wish to download. Only used when 

50 type="strain". Default 16384 

51 format: str, optional 

52 format of strain data you wish to download. Only used when 

53 type="strain". Default "gwf" 

54 duration: int, optional 

55 duration of strain data you wish to download. Only used when 

56 type="strain". Default 32 

57 IFO: str, optional 

58 detector strain data you wish to download. Only used when type="strain". 

59 Default 'L1' 

60 version: str, optional 

61 version of the file to download. Default None 

62 """ 

63 if type not in ["posterior", "strain"]: 

64 raise ValueError( 

65 "Unknown data type: '{}'. Must be either 'posterior' or " 

66 "'strain'.".format(type) 

67 ) 

68 data, = fetch_event_json( 

69 event, catalog=catalog, version=version 

70 )["events"].values() 

71 url = None 

72 if type == "posterior": 

73 for key, item in data["parameters"].items(): 

74 if "_pe_" in key: 

75 url = item["data_url"] 

76 break 

77 elif type == "strain": 

78 strain = data["strain"] 

79 for _strain in strain: 

80 cond = ( 

81 _strain["sampling_rate"] == sampling_rate 

82 and _strain["format"] == format 

83 and _strain["duration"] == duration 

84 and _strain["detector"] == IFO 

85 ) 

86 if cond: 

87 url = _strain["url"] 

88 if url is None: 

89 raise RuntimeError("Failed to find data URL for {}".format(event)) 

90 return url 

91 

92 

93@deprecation( 

94 "The 'fetch_open_data' function has changed its name to " 

95 "'fetch_open_samples' and 'fetch_open_data' may not be supported in future " 

96 "releases. Please update" 

97) 

98def fetch_open_data(event, **kwargs): 

99 """Download and read publically available gravitational wave posterior 

100 samples 

101 

102 Parameters 

103 ---------- 

104 event: str 

105 name of the gravitational wave event you wish to download data for 

106 """ 

107 return fetch_open_samples(event, **kwargs) 

108 

109 

110def _fetch_open_data( 

111 event, type="posterior", catalog=None, version=None, sampling_rate=16384, 

112 format="gwf", duration=32, IFO="L1", **kwargs 

113): 

114 """Download and read publcally available gravitational wave data 

115 

116 Parameters 

117 ---------- 

118 event: str 

119 name of the gravitational wave event you wish to download data for 

120 type: str, optional 

121 type of data you wish to download. Default "posterior" 

122 catalog: str, optional 

123 Name of catalog that hosts the event. Default None 

124 version: str, optional 

125 Version of the file to download. Default None 

126 sampling_rate: int, optional 

127 sampling rate of strain data you wish to download. Only used when 

128 type="strain". Default 16384 

129 format: str, optional 

130 format of strain data you wish to download. Only used when 

131 type="strain". Default "gwf" 

132 duration: int, optional 

133 duration of strain data you wish to download. Only used when 

134 type="strain". Default 32 

135 IFO: str, optional 

136 detector strain data you wish to download. Only used when type="strain". 

137 Default 'L1' 

138 """ 

139 try: 

140 url = _DCC_url( 

141 event, type=type, catalog=catalog, sampling_rate=sampling_rate, 

142 format=format, duration=duration, IFO=IFO, version=version 

143 ) 

144 except RuntimeError: 

145 raise ValueError( 

146 "Unknown URL for {}. If the URL is known, please run " 

147 "download_and_read_file(URL)".format(event) 

148 ) 

149 if type == "strain": 

150 kwargs.update({"IFO": IFO}) 

151 return download_and_read_file(url, **kwargs) 

152 

153 

154def fetch_open_samples(event, **kwargs): 

155 """Download and read publically available gravitational wave posterior 

156 samples 

157 

158 Parameters 

159 ---------- 

160 event: str 

161 name of the gravitational wave event you wish to download data for 

162 **kwargs: dict, optional 

163 all additional kwargs passed to _fetch_open_data 

164 """ 

165 # fetch posterior data 

166 out = _fetch_open_data(event, type="posterior", **kwargs) 

167 

168 # if asked to read the data, or unpack a tarball, just return it now 

169 if ( 

170 kwargs.get("read_file", True) 

171 or kwargs.get("unpack", False) 

172 ): 

173 return out 

174 

175 # otherwise, if Zenodo returned a file without a suffix, we need to add one 

176 # see https://git.ligo.org/gwosc/client/-/issues/95 

177 out = Path(out) 

178 if not out.suffix: 

179 out = out.rename(out.with_suffix(".h5")) 

180 return str(out) 

181 

182 

183def fetch_open_strain(event, format="gwf", **kwargs): 

184 """Download and read publically available gravitational wave strain data 

185 

186 Parameters 

187 ---------- 

188 event: str 

189 name of the gravitational wave event you wish to download data for 

190 format: str, optional 

191 format of strain data you wish to download. Default "gwf" 

192 **kwargs: dict, optional 

193 all additional kwargs passed to _fetch_open_data 

194 """ 

195 _kwargs = kwargs.copy() 

196 _kwargs["format"] = "gwf" 

197 return _fetch_open_data(event, type="strain", **_kwargs)