Coverage for pesummary/gw/file/meta_file.py: 88.0%
142 statements
« 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 numpy as np
4from pesummary.utils.utils import logger
5from pesummary.core.file.meta_file import (
6 _MetaFile, DEFAULT_HDF5_KEYS as CORE_HDF5_KEYS
7)
9__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
10DEFAULT_HDF5_KEYS = CORE_HDF5_KEYS
13class _GWMetaFile(_MetaFile):
14 """This class handles the creation of a meta file storing all information
15 from the analysis
17 Attributes
18 ----------
19 meta_file: str
20 name of the meta file storing all information
21 """
22 def __init__(
23 self, samples, labels, config, injection_data, file_versions,
24 file_kwargs, calibration=None, psd=None, approximant=None, webdir=None,
25 result_files=None, hdf5=False, existing_version=None,
26 existing_label=None, existing_samples=None, existing_psd=None,
27 existing_calibration=None, existing_approximant=None,
28 existing_config=None, existing_injection=None,
29 existing_metadata=None, priors={}, outdir=None, existing=None,
30 existing_priors={}, existing_metafile=None, package_information={},
31 mcmc_samples=False, skymap=None, existing_skymap=None,
32 filename=None, external_hdf5_links=False, hdf5_compression=None,
33 history=None, descriptions=None, gwdata=None
34 ):
35 self.calibration = calibration
36 self.psds = psd
37 self.approximant = approximant
38 self.existing_psd = existing_psd
39 self.existing_calibration = existing_calibration
40 self.existing_approximant = existing_approximant
41 self.skymap = skymap
42 self.existing_skymap = existing_skymap
43 self.gwdata = gwdata
44 super(_GWMetaFile, self).__init__(
45 samples, labels, config, injection_data, file_versions,
46 file_kwargs, webdir=webdir, result_files=result_files, hdf5=hdf5,
47 priors=priors, existing_version=existing_version,
48 existing_label=existing_label, existing_samples=existing_samples,
49 existing_injection=existing_injection,
50 existing_metadata=existing_metadata,
51 existing_config=existing_config, existing_priors=existing_priors,
52 outdir=outdir, package_information=package_information,
53 existing=existing, existing_metafile=existing_metafile,
54 mcmc_samples=mcmc_samples, filename=filename,
55 external_hdf5_links=external_hdf5_links,
56 hdf5_compression=hdf5_compression, history=history,
57 descriptions=descriptions
58 )
59 if self.calibration is None:
60 self.calibration = {label: {} for label in self.labels}
61 if self.psds is None:
62 self.psds = {label: {} for label in self.labels}
64 def _make_dictionary(self):
65 """Generate a single dictionary which stores all information
66 """
67 super(_GWMetaFile, self)._make_dictionary()
68 for num, label in enumerate(self.labels):
69 cond = all(self.calibration[label] != j for j in [{}, None])
70 if self.calibration != {} and cond:
71 self.data[label]["calibration_envelope"] = {
72 key: item for key, item in self.calibration[label].items()
73 if item is not None
74 }
75 else:
76 self.data[label]["calibration_envelope"] = {}
77 if self.psds != {} and all(self.psds[label] != j for j in [{}, None]):
78 self.data[label]["psds"] = {
79 key: item for key, item in self.psds[label].items() if item
80 is not None
81 }
82 else:
83 self.data[label]["psds"] = {}
84 if self.approximant is not None and self.approximant[label] is not None:
85 self.data[label]["approximant"] = self.approximant[label]
86 else:
87 self.data[label]["approximant"] = {}
88 if self.skymap is not None and len(self.skymap):
89 if self.skymap[label] is not None:
90 self.data[label]["skymap"] = {
91 "data": self.skymap[label],
92 "meta_data": {
93 key: item for key, item in
94 self.skymap[label].meta_data.items()
95 }
96 }
97 if self.gwdata is not None and len(self.gwdata):
98 try:
99 from gwpy.types.io.hdf5 import format_index_array_attrs
100 from pesummary.utils.dict import Dict
101 self.data["strain"] = {}
102 for key, item in self.gwdata.items():
103 if item is None:
104 continue
105 _name = item.name if item.name is not None else "unknown_name"
106 self.data["strain"][key] = Dict(
107 {_name: item.value},
108 extra_kwargs=format_index_array_attrs(item)
109 )
110 except Exception:
111 logger.warning(
112 "Failed to store the gravitational wave strain data"
113 )
115 @staticmethod
116 def save_to_hdf5(
117 data, labels, samples, meta_file, no_convert=False, mcmc_samples=False,
118 external_hdf5_links=False, compression=None, _class=None, gwdata=None
119 ):
120 """Save the metafile as a hdf5 file
121 """
122 if gwdata is not None and len(gwdata):
123 extra_keys = CORE_HDF5_KEYS + ["strain"]
124 else:
125 extra_keys = CORE_HDF5_KEYS
126 _MetaFile.save_to_hdf5(
127 data, labels, samples, meta_file, no_convert=no_convert,
128 extra_keys=extra_keys, mcmc_samples=mcmc_samples, _class=_class,
129 external_hdf5_links=external_hdf5_links, compression=compression
130 )
133class _TGRMetaFile(_GWMetaFile):
134 """Class to create a single file to contain TGR data
135 """
136 def __init__(
137 self, samples, labels, imrct_data=None, webdir=None, outdir=None,
138 file_kwargs={}
139 ):
140 super(_TGRMetaFile, self).__init__(
141 samples, labels, None, None, None, None, webdir=webdir,
142 outdir=outdir, filename="tgr_samples.h5", hdf5=True
143 )
144 self.tgr_data = {"imrct": imrct_data}
145 self.file_kwargs = file_kwargs
146 for key in self.tgr_data.keys():
147 if key not in self.file_kwargs:
148 self.file_kwargs[key] = {}
150 @staticmethod
151 def convert_posterior_samples_to_numpy(labels, samples, mcmc_samples=False):
152 """Convert a dictionary of multiple posterior samples from a
153 column-major dictionary to a row-major numpy array
155 Parameters
156 ----------
157 labels: list
158 list of unique labels for each analysis
159 samples: MultiAnalysisSamplesDict
160 dictionary of multiple posterior samples to convert to a numpy
161 array.
162 mcmc_samples: Bool, optional
163 if True, the dictionary contains seperate mcmc chains
165 Examples
166 --------
167 >>> dictionary = MultiAnalysisSamplesDict(
168 ... {"label": {"mass_1": [1,2,3], "mass_2": [1,2,3]}}
169 ... )
170 >>> dictionary = _Metafile.convert_posterior_samples_to_numpy(
171 ... dictionary.keys(), dictionary
172 ... )
173 >>> print(dictionary)
174 ... {"label": rec.array([(1., 1.), (2., 2.), (3., 3.)],
175 ... dtype=[('mass_1', '<f4'), ('mass_2', '<f4')])}
176 """
177 _convert_function = _GWMetaFile._convert_posterior_samples_to_numpy
178 converted_samples = {label: {} for label in labels}
179 for label in labels:
180 for key in ["inspiral", "postinspiral"]:
181 if "{}:{}".format(label, key) in samples.keys():
182 _samples_key = "{}:{}".format(label, key)
183 else:
184 _samples_key = key
185 converted_samples[label][key] = _convert_function(
186 samples[_samples_key], mcmc_samples=False
187 )
188 return converted_samples
190 @staticmethod
191 def save_to_hdf5(*args, **kwargs):
192 """Save the metafile as a hdf5 file
193 """
194 return _GWMetaFile.save_to_hdf5(*args, _class=_TGRMetaFile, **kwargs)
196 def _make_dictionary(self):
197 """Generate a single dictionary which stores all information
198 """
199 dictionary = self._dictionary_structure
200 dictionary.update(
201 {
202 label: {key: {} for key in self.tgr_data.keys()} for label in
203 self.labels
204 }
205 )
206 for label in self.labels:
207 dictionary[label]["posterior_samples"] = {}
208 if self.tgr_data["imrct"] is not None:
209 for analysis in ["inspiral", "postinspiral"]:
210 try:
211 _samples = self.samples["{}:{}".format(label, analysis)]
212 except KeyError:
213 _samples = self.samples[analysis]
214 parameters = _samples.keys()
215 samples = np.array([_samples[i] for i in parameters]).T
216 dictionary[label]["posterior_samples"][analysis] = {
217 "parameter_names": list(parameters),
218 "samples": samples.tolist()
219 }
220 deviations = "final_mass_final_spin_deviations"
221 _imrct_data = self.tgr_data["imrct"][label][deviations]
222 dictionary[label]["imrct"] = {
223 "final_mass_deviation": _imrct_data.x,
224 "final_spin_deviation": _imrct_data.y,
225 "pdf": _imrct_data.probs,
226 }
227 _kwargs = self.file_kwargs["imrct"]
228 if label in _kwargs.keys():
229 _kwargs = _kwargs[label]
230 dictionary[label]["imrct"]["meta_data"] = _kwargs
231 self.data = dictionary
234class GWMetaFile(object):
235 """This class handles the creation of a metafile storing all information
236 from the analysis
237 """
238 def __init__(self, inputs, history=None):
239 logger.info("Starting to generate the meta file")
240 if inputs.add_to_existing:
241 existing = inputs.existing
242 existing_metafile = inputs.existing_metafile
243 existing_samples = inputs.existing_samples
244 existing_labels = inputs.existing_labels
245 existing_psd = inputs.existing_psd
246 existing_calibration = inputs.existing_calibration
247 existing_config = inputs.existing_config
248 existing_approximant = inputs.existing_approximant
249 existing_injection = inputs.existing_injection_data
250 existing_version = inputs.existing_file_version
251 existing_metadata = inputs.existing_file_kwargs
252 existing_priors = inputs.existing_priors
253 else:
254 existing_metafile = None
255 existing_samples = None
256 existing_labels = None
257 existing_psd = None
258 existing_calibration = None
259 existing_config = None
260 existing_approximant = None
261 existing_injection = None
262 existing_version = None
263 existing_metadata = None
264 existing = None
265 existing_priors = {}
267 meta_file = _GWMetaFile(
268 inputs.samples, inputs.labels, inputs.config, inputs.injection_data,
269 inputs.file_version, inputs.file_kwargs, calibration=inputs.calibration,
270 psd=inputs.psd, hdf5=inputs.hdf5, webdir=inputs.webdir,
271 result_files=inputs.result_files, existing_version=existing_version,
272 existing_label=existing_labels, existing_samples=existing_samples,
273 existing_psd=existing_psd, existing_calibration=existing_calibration,
274 existing_approximant=existing_approximant,
275 existing_injection=existing_injection,
276 existing_metadata=existing_metadata,
277 existing_config=existing_config, priors=inputs.priors,
278 existing_priors=existing_priors, existing=existing,
279 existing_metafile=existing_metafile, approximant=inputs.approximant,
280 package_information=inputs.package_information,
281 mcmc_samples=inputs.mcmc_samples, skymap=inputs.skymap,
282 existing_skymap=inputs.existing_skymap, filename=inputs.filename,
283 external_hdf5_links=inputs.external_hdf5_links,
284 hdf5_compression=inputs.hdf5_compression, history=history,
285 gwdata=inputs.gwdata, descriptions=inputs.descriptions
286 )
287 meta_file.make_dictionary()
288 if not inputs.hdf5:
289 meta_file.save_to_json(meta_file.data, meta_file.meta_file)
290 else:
291 meta_file.save_to_hdf5(
292 meta_file.data, meta_file.labels, meta_file.samples,
293 meta_file.meta_file, mcmc_samples=meta_file.mcmc_samples,
294 external_hdf5_links=meta_file.external_hdf5_links,
295 compression=meta_file.hdf5_compression,
296 gwdata=meta_file.gwdata
297 )
298 meta_file.save_to_dat()
299 meta_file.write_marginalized_posterior_to_dat()
300 logger.info(
301 "Finishing generating the meta file. The meta file can be viewed "
302 "here: {}".format(meta_file.meta_file)
303 )
306class TGRMetaFile(_GWMetaFile):
307 """Class to create a single file to contain TGR data
308 """
309 def __init__(
310 self, samples, labels, imrct_data=None, webdir=None, outdir=None,
311 file_kwargs={}
312 ):
313 logger.info("Starting to generate the meta file")
314 meta_file = _TGRMetaFile(
315 samples, labels, imrct_data=imrct_data, webdir=webdir,
316 outdir=outdir, file_kwargs=file_kwargs
317 )
318 meta_file.make_dictionary()
319 meta_file.save_to_hdf5(
320 meta_file.data, meta_file.labels, meta_file.samples,
321 meta_file.meta_file, mcmc_samples=False
322 )
323 logger.info(
324 "Finished generating the meta file. The meta file can be viewed "
325 "here: {}".format(meta_file.meta_file)
326 )