Coverage for pesummary/tests/fetch_test.py: 100.0%

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

2 

3from pesummary.io import read 

4from pesummary.core.fetch import download_dir, download_and_read_file 

5from pesummary.gw.fetch import fetch_open_samples 

6import numpy as np 

7import os 

8 

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

10 

11 

12def test_download_and_read_file(): 

13 """Test that the `pesummary.core.fetch.download_and_read_file` function 

14 works as expected 

15 """ 

16 import requests 

17 import tempfile 

18 from pathlib import Path 

19 import shutil 

20 

21 tmp1 = Path(tempfile.TemporaryDirectory(prefix=".", dir=".").name).name 

22 tmp2 = Path(tempfile.TemporaryDirectory(prefix=".", dir=".").name).name 

23 os.mkdir(tmp1) 

24 os.mkdir(tmp2) 

25 data = download_and_read_file( 

26 "https://dcc.ligo.org/public/0157/P1800370/005/GW170608_GWTC-1.hdf5", 

27 outdir=tmp1 

28 ) 

29 _data = requests.get( 

30 "https://dcc.ligo.org/public/0157/P1800370/005/GW170608_GWTC-1.hdf5" 

31 ) 

32 with open(f"{tmp2}/GW170608_posterior_samples.h5", "wb") as f: 

33 f.write(_data.content) 

34 data2 = read(f"{tmp2}/GW170608_posterior_samples.h5") 

35 np.testing.assert_almost_equal( 

36 np.array(data.samples), np.array(data2.samples) 

37 ) 

38 shutil.rmtree(tmp1) 

39 shutil.rmtree(tmp2) 

40 

41 

42def test_download_and_keep_file(): 

43 """Test that when the `read=False` kwarg is passed to the 

44 download_and_read_file function the filename is returned 

45 """ 

46 file_name = download_and_read_file( 

47 "https://dcc.ligo.org/public/0157/P1800370/005/GW170817_GWTC-1.hdf5", 

48 outdir=".", read_file=False 

49 ) 

50 assert os.path.isfile(file_name) 

51 

52 

53def test_fetch_tarball_and_keep(): 

54 """Test that the `pesummary.gw.fetch.fetch_open_samples` function is able to 

55 fetch, unpack and keep a tarball 

56 """ 

57 directory_name = fetch_open_samples( 

58 "GW190424_180648", read_file=False, outdir=".", unpack=True, 

59 catalog="GWTC-2", download_kwargs={"timeout": 60} 

60 ) 

61 assert os.path.isdir("./GW190424_180648") 

62 assert os.path.isdir(directory_name) 

63 

64 

65def test_fetch_tarball_and_keep_single_file(): 

66 """Test that the `pesummary.gw.fetch.fetch_open_samples` function is able to 

67 fetch, unpack and keep a single file stored in a tarball 

68 """ 

69 file_name = fetch_open_samples( 

70 "GW190424_180648", read_file=False, outdir=".", unpack=True, 

71 path="GW190424_180648.h5", catalog="GWTC-2", 

72 download_kwargs={"timeout": 60} 

73 ) 

74 assert os.path.isfile("./GW190424_180648.h5") 

75 assert os.path.isfile(file_name) 

76 

77 

78def test_fetch_and_open_tarball(): 

79 """Test that a `pesummary.gw.fetch.fetch_open_samples` function is able to 

80 fetch, unpack and read a single file stored in a tarball 

81 """ 

82 import pesummary.gw.file.formats.pesummary 

83 

84 f = fetch_open_samples( 

85 "GW190424_180648", read_file=True, outdir=".", unpack=True, 

86 path="GW190424_180648.h5", catalog="GWTC-2", 

87 download_kwargs={"timeout": 60} 

88 ) 

89 assert isinstance(f, pesummary.gw.file.formats.pesummary.PESummary) 

90 

91 

92def test_fetch_open_samples(): 

93 """Test that the `pesummary.gw.fetch.fetch_open_samples` function works as 

94 expected 

95 """ 

96 data = fetch_open_samples( 

97 "GW190412", catalog="GWTC-2.1-confident", 

98 download_kwargs={"timeout": 60} 

99 ) 

100 data2 = download_and_read_file( 

101 "https://zenodo.org/api/records/6513631/files/" + 

102 "IGWN-GWTC2p1-v2-GW190412_053044_PEDataRelease_mixed_cosmo.h5/content", 

103 outdir=download_dir, 

104 read_file=True, 

105 download_kwargs=dict( 

106 cache=True, 

107 pkgname="pesummary", 

108 ) 

109 ) 

110 for num in range(len(data.labels)): 

111 np.testing.assert_almost_equal( 

112 np.array(data.samples[num]), np.array(data2.samples[num]) 

113 ) 

114 

115 

116def test_fetch_open_strain(): 

117 """Test that the `pesummary.gw.fetch.fetch_open_strain` function works as 

118 expected 

119 """ 

120 from pesummary.gw.fetch import fetch_open_strain 

121 from gwpy.timeseries import TimeSeries 

122 data = fetch_open_strain( 

123 "GW190412", IFO="H1", channel="H1:GWOSC-4KHZ_R1_STRAIN", 

124 sampling_rate=4096 

125 ) 

126 path = download_and_read_file( 

127 "https://www.gw-openscience.org/eventapi/html/GWTC-2/GW190412/v3/" + 

128 "H-H1_GWOSC_4KHZ_R1-1239082247-32.gwf", 

129 outdir=download_dir, 

130 read_file=False, 

131 download_kwargs=dict( 

132 cache=True, 

133 pkgname="pesummary", 

134 ) 

135 ) 

136 data2 = TimeSeries.read( 

137 path, channel="H1:GWOSC-4KHZ_R1_STRAIN" 

138 ) 

139 np.testing.assert_almost_equal(data.value, data2.value) 

140 np.testing.assert_almost_equal(data.times.value, data2.times.value)