Coverage for pesummary/io/read.py: 64.1%
39 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
1# Licensed under an MIT style license -- see LICENSE.md
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
11__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
13OTHER = {
14 "fits": SkyMap.from_fits,
15 "ini": read_ini,
16 "gwf": StrainData.read,
17 "lcf": StrainData.read,
18 "pickle": read_pickle
19}
22def _fetch(ff, function, **kwargs):
23 """Copy a file from a remote location to a temporary folder ready for
24 reading
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
38def _fetch_from_url(url):
39 """Copy file from url to a temporary folder ready for reading
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})
50def _fetch_from_remote_server(ff):
51 """Copy file from remote server to a temporary folder ready for reading
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)
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.
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)
111 module = importlib.import_module("pesummary.{}.file.read".format(package))
112 return getattr(module, "read")(path, file_format=file_format, **kwargs)