Coverage for /home/runner/work/nr-catalog-tools/nr-catalog-tools/nrcatalogtools/utils.py: 52%

82 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-01 05:18 +0000

1import functools 

2import pathlib 

3import shutil 

4 

5import lal 

6import requests 

7 

8import sxs 

9 

10nrcatalog_cache_dir = pathlib.Path("~/.cache/").expanduser().resolve() 

11nr_group_tags = {} 

12nr_group_tags["SXS"] = "SXS" 

13nr_group_tags["RIT"] = "RIT" 

14nr_group_tags["MAYA"] = "MAYA" 

15nr_group_tags["UNKNOWN"] = "UNKNOWN" 

16 

17rit_catalog_info = {} 

18rit_catalog_info["cache_dir"] = nrcatalog_cache_dir / "RIT" 

19rit_catalog_info["metadata_dir"] = rit_catalog_info["cache_dir"] / "metadata" 

20rit_catalog_info["data_dir"] = rit_catalog_info["cache_dir"] / "data" 

21rit_catalog_info["url"] = "https://ccrgpages.rit.edu/~RITCatalog/" 

22rit_catalog_info["metadata_url"] = rit_catalog_info["url"] + "/Metadata/" 

23rit_catalog_info["data_url"] = rit_catalog_info["url"] + "/Data/" 

24rit_catalog_info["possible_resolutions"] = [100, 120, 88, 118, 130, 140, 144, 160, 200] 

25rit_catalog_info["metadata_file_fmts"] = [ 

26 "RIT:BBH:{:04d}-n{:3d}-id{:d}_Metadata.txt", 

27 "RIT:eBBH:{:04d}-n{:3d}-ecc_Metadata.txt", 

28] 

29rit_catalog_info["waveform_file_fmts"] = [ 

30 "ExtrapStrain_RIT-BBH-{:04d}-n{:3d}.h5", 

31 "ExtrapStrain_RIT-eBBH-{:04d}-n{:3d}.h5", 

32] 

33rit_catalog_info["max_id_val"] = 6 

34 

35maya_catalog_info = { 

36 "cache_dir": nrcatalog_cache_dir / "MAYA", 

37 "url": "https://raw.githubusercontent.com/cevans216/gt-waveform-catalog/master/h5files", 

38 "metadata_url": "https://raw.githubusercontent.com/cevans216/gt-waveform-catalog/master/catalog-table.txt", 

39} 

40maya_catalog_info["data_dir"] = maya_catalog_info["cache_dir"] / "data" 

41maya_catalog_info["metadata_dir"] = maya_catalog_info["cache_dir"] / "metadata" 

42maya_catalog_info["data_url"] = maya_catalog_info["url"] 

43 

44 

45def url_exists(link, num_retries=100): 

46 """Check if a given URL exists on the web. 

47 

48 Args: 

49 link : complete web URL 

50 

51 Returns: 

52 bool: True/False whether the URL could be found on WWW. 

53 """ 

54 requests.packages.urllib3.disable_warnings() 

55 for n in range(num_retries): 

56 try: 

57 response = requests.head(link, verify=False) 

58 if response.status_code == requests.codes.ok: 

59 return True 

60 else: 

61 return False 

62 except Exception: 

63 continue 

64 return False 

65 

66 

67def download_file(url, path, progress=False, if_newer=True): 

68 if url_exists(url): 

69 try: 

70 return sxs.utilities.downloads.download_file( 

71 url, path, progress=progress, if_newer=if_newer 

72 ) 

73 except Exception: 

74 requests.packages.urllib3.disable_warnings() 

75 for n in range(100): 

76 try: 

77 r = requests.get( 

78 url, verify=False, stream=True, allow_redirects=True 

79 ) 

80 break 

81 except Exception: 

82 continue 

83 if r.status_code != 200: 

84 print(f"An error occurred when trying to access <{url}>.") 

85 try: 

86 print(r.json()) 

87 except Exception: 

88 pass 

89 r.raise_for_status() 

90 raise RuntimeError() # Will only happen if the response was not strictly an error 

91 r.raw.read = functools.partial(r.raw.read, decode_content=True) 

92 path = pathlib.Path(path).expanduser().resolve() 

93 with path.open("wb") as f: 

94 shutil.copyfileobj(r.raw, f) 

95 return path 

96 

97 

98def call_with_timeout(myfunc, args=(), kwargs={}, timeout=5): 

99 """ 

100 This function calls user-provided `myfunc` with user-provided 

101 `args` and `kwargs` in a separate multiprocessing.Process. 

102 if the function evaluation takes more than `timeout` seconds, the 

103 `Process` is terminated and error raised. If it evalutes within 

104 `timeout` seconds, the results are fetched from the `Queue` and 

105 returned. 

106 """ 

107 

108 from multiprocessing import Process, Queue 

109 

110 def funcwrapper(p, *args, **kwargs): 

111 """ 

112 This thin wrapper calls the user-provided function, and puts 

113 its result into the multiprocessing `Queue` so that it can be 

114 obtained via `Queue().get()`. 

115 """ 

116 res = myfunc(*args, **kwargs) 

117 p.put(res) 

118 

119 queue = Queue() 

120 task = Process(target=funcwrapper, args=(queue, *args)) 

121 task.start() 

122 task.join(timeout=timeout) 

123 task.terminate() 

124 try: 

125 result = queue.get(timeout=0) 

126 return result 

127 except Exception: 

128 raise Exception("Timeout") 

129 

130 

131def time_to_physical(M): 

132 """Factor to convert time from dimensionless units to SI units 

133 

134 parameters 

135 ---------- 

136 M: mass of system in the units of solar mass 

137 

138 Returns 

139 ------- 

140 converting factor 

141 """ 

142 

143 return M * lal.MTSUN_SI 

144 

145 

146def amp_to_physical(M, D): 

147 """Factor to rescale strain to mass M and distance D convert from 

148 dimensionless units to SI units 

149 

150 parameters 

151 ---------- 

152 M: mass of the system in units of solar mass 

153 D: Luminosity distance in units of megaparsecs 

154 

155 Returns 

156 ------- 

157 Scaling factor 

158 """ 

159 

160 return lal.G_SI * M * lal.MSUN_SI / (lal.C_SI**2 * D * 1e6 * lal.PC_SI)