Coverage for pesummary/io/read.py: 64.1%

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

2 

3import os 

4import importlib 

5from pathlib import Path 

6from pesummary.core.file.formats.ini import read_ini 

7from pesummary.core.file.formats.pickle import read_pickle 

8from pesummary.gw.file.strain import StrainData 

9from pesummary.gw.file.skymap import SkyMap 

10 

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

12 

13OTHER = { 

14 "fits": SkyMap.from_fits, 

15 "ini": read_ini, 

16 "gwf": StrainData.read, 

17 "lcf": StrainData.read, 

18 "pickle": read_pickle 

19} 

20 

21 

22def _fetch(ff, function, **kwargs): 

23 """Copy a file from a remote location to a temporary folder ready for 

24 reading 

25 

26 Parameters 

27 ---------- 

28 ff: str 

29 path to file you wish to read 

30 function: func 

31 function you wish to use for fetching the file from a remote location 

32 """ 

33 filename = function(ff, read_file=False, **kwargs) 

34 path = str(filename) 

35 return path 

36 

37 

38def _fetch_from_url(url): 

39 """Copy file from url to a temporary folder ready for reading 

40 

41 Parameters 

42 ---------- 

43 url: str 

44 url of a result file you wish to load 

45 """ 

46 from pesummary.core.fetch import download_and_read_file 

47 return _fetch(url, download_and_read_file, download_kwargs={"timeout": 60}) 

48 

49 

50def _fetch_from_remote_server(ff): 

51 """Copy file from remote server to a temporary folder ready for reading 

52 

53 Parameters 

54 ---------- 

55 ff: str 

56 path to a results file on a remote server. Must be in the form 

57 {username}@{servername}:{path} 

58 """ 

59 from pesummary.core.fetch import scp_and_read_file 

60 return _fetch(ff, scp_and_read_file) 

61 

62 

63def read( 

64 path, package="gw", file_format=None, skymap=False, strain=False, cls=None, 

65 checkpoint=False, **kwargs 

66): 

67 """Read in a results file. 

68 

69 Parameters 

70 ---------- 

71 path: str 

72 path to results file. If path is on a remote server, add username and 

73 servername in the form {username}@{servername}:{path} 

74 package: str 

75 the package you wish to use 

76 file_format: str 

77 the file format you wish to use. Default None. If None, the read 

78 function loops through all possible options 

79 skymap: Bool, optional 

80 if True, path is the path to a fits file generated with `ligo.skymap` 

81 strain: Bool, optional 

82 if True, path is the path to a frame file containing gravitational 

83 wave data. All kwargs are passed to 

84 pesummary.gw.file.strain.StrainData.read 

85 cls: func, optional 

86 class to use when reading in a result file 

87 checkpoint: Bool, optional 

88 if True, treat path as the path to a checkpoint file 

89 **kwargs: dict, optional 

90 all additional kwargs are passed to the `pesummary.{}.file.read.read` 

91 function 

92 """ 

93 if not os.path.isfile(path) and "https://" in path: 

94 path = _fetch_from_url(path) 

95 elif not os.path.isfile(path) and "@" in path: 

96 path = _fetch_from_remote_server(path) 

97 extension = Path(path).suffix[1:] 

98 if cls is not None: 

99 return cls.load_file(path, **kwargs) 

100 if extension in OTHER.keys(): 

101 return OTHER[extension](path, **kwargs) 

102 elif file_format == "ini": 

103 return OTHER["ini"](path, **kwargs) 

104 elif skymap: 

105 return OTHER["fits"](path, **kwargs) 

106 elif strain: 

107 return OTHER["gwf"](path, **kwargs) 

108 elif checkpoint: 

109 return OTHER["pickle"](path, **kwargs) 

110 

111 module = importlib.import_module("pesummary.{}.file.read".format(package)) 

112 return getattr(module, "read")(path, file_format=file_format, **kwargs)