Coverage for pesummary/cli/summarymodify.py: 91.1%
305 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#! /usr/bin/env python
3# Licensed under an MIT style license -- see LICENSE.md
5import os
6import numpy as np
7import math
8import json
9import h5py
10from pathlib import Path
12from pesummary.utils.utils import logger, check_file_exists_and_rename
13from pesummary.utils.dict import paths_to_key
14from pesummary.utils.exceptions import InputError
15from pesummary.core.cli.parser import ArgumentParser as _ArgumentParser
16from pesummary.core.cli.actions import DelimiterSplitAction
17from pesummary.gw.cli.inputs import _GWInput
18from pesummary.gw.file.meta_file import _GWMetaFile
20__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
21__doc__ = """This executable is used to modify a PESummary metafile from the
22command line"""
25class _Input(_GWInput):
26 """Super class to handle the command line arguments
27 """
28 @property
29 def labels(self):
30 return self._labels
32 @labels.setter
33 def labels(self, labels):
34 self._labels = labels
35 if labels is not None and isinstance(labels, dict):
36 self._labels = labels
37 elif labels is not None:
38 raise InputError(
39 "Please provide an existing labels and the label you wish "
40 "to replace it with `--labels existing:new`."
41 )
43 @property
44 def config(self):
45 return self._config
47 @config.setter
48 def config(self, config):
49 self._config = config
50 if config is not None and isinstance(config, dict):
51 for key, item in config.items():
52 if not os.path.isfile(item):
53 raise FileNotFoundError(
54 "Unable to find the config file '{}'".format(item)
55 )
56 self._config[key] = _GWMetaFile._grab_config_data_from_data_file(
57 item
58 )
59 elif config is not None:
60 raise InputError(
61 "Please provide the label and config file with '--config "
62 "label:path'"
63 )
65 @property
66 def kwargs(self):
67 return self._kwargs
69 @kwargs.setter
70 def kwargs(self, kwargs):
71 self._kwargs = kwargs
72 if kwargs is not None and isinstance(kwargs, dict):
73 self._kwargs = kwargs
74 elif kwargs is not None:
75 raise InputError(
76 "Please provide the label, kwarg and value with '--kwargs "
77 "label:kwarg:value`"
78 )
80 @property
81 def replace_posterior(self):
82 return self._replace_posterior
84 @replace_posterior.setter
85 def replace_posterior(self, replace_posterior):
86 self._replace_posterior = replace_posterior
87 if replace_posterior is not None and isinstance(replace_posterior, dict):
88 self._replace_posterior = replace_posterior
89 elif replace_posterior is not None:
90 raise InputError(
91 "Please provide the label, posterior and file path with "
92 "value with '--replace_posterior "
93 "label;posterior:/path/to/posterior.dat where ';' is the chosen "
94 "delimiter and provided with '--delimiter ;`"
95 )
97 @property
98 def remove_posterior(self):
99 return self._remove_posterior
101 @remove_posterior.setter
102 def remove_posterior(self, remove_posterior):
103 self._remove_posterior = remove_posterior
104 if remove_posterior is not None and isinstance(remove_posterior, dict):
105 self._remove_posterior = remove_posterior
106 elif remove_posterior is not None:
107 raise InputError(
108 "Please provide the label and posterior with '--remove_posterior "
109 "label:posterior`"
110 )
112 @property
113 def store_skymap(self):
114 return self._store_skymap
116 @store_skymap.setter
117 def store_skymap(self, store_skymap):
118 self._store_skymap = store_skymap
119 if store_skymap is not None and isinstance(store_skymap, dict):
120 self._store_skymap = store_skymap
121 elif store_skymap is not None:
122 raise InputError(
123 "Please provide the label and path to skymap with '--store_skymap "
124 "label:path/to/skymap.fits`"
125 )
127 @property
128 def samples(self):
129 return self._samples
131 @samples.setter
132 def samples(self, samples):
133 if samples is None:
134 raise InputError(
135 "Please provide a result file that you wish to modify"
136 )
137 if len(samples) > 1:
138 raise InputError(
139 "Only a single result file can be passed"
140 )
141 samples = samples[0]
142 if not self.is_pesummary_metafile(samples):
143 raise InputError(
144 "Please provide a PESummary metafile to this executable"
145 )
146 self._samples = samples
148 @property
149 def data(self):
150 return self._data
152 @data.setter
153 def data(self, data):
154 extension = Path(self.samples).suffix
155 if extension == ".h5" or extension == ".hdf5":
156 from pesummary.core.file.formats.pesummary import PESummary
157 from pandas import DataFrame
159 with h5py.File(self.samples, "r") as f:
160 data = PESummary._convert_hdf5_to_dict(f)
161 for label in data.keys():
162 try:
163 data[label]["posterior_samples"] = DataFrame(
164 data[label]["posterior_samples"]
165 ).to_records(index=False, column_dtypes=float)
166 except KeyError:
167 pass
168 except Exception:
169 parameters = data[label]["posterior_samples"]["parameter_names"]
170 if isinstance(parameters[0], bytes):
171 parameters = [
172 parameter.decode("utf-8") for parameter in parameters
173 ]
174 samples = np.array([
175 j for j in data[label]["posterior_samples"]["samples"]
176 ].copy())
177 data[label]["posterior_samples"] = DataFrame.from_dict(
178 {
179 param: samples.T[num] for num, param in
180 enumerate(parameters)
181 }
182 ).to_records(index=False, column_dtypes=float)
183 self._data = data
184 elif extension == ".json":
185 with open(self.samples, "r") as f:
186 self._data = json.load(f)
187 else:
188 raise InputError(
189 "The extension '{}' is not recognised".format(extension)
190 )
193class Input(_Input):
194 """Class to handle the command line arguments
196 Parameters
197 ----------
198 opts: argparse.Namespace
199 Namespace object containing the command line options
201 Attributes
202 ----------
203 samples: str
204 path to a PESummary meta file that you wish to modify
205 labels: dict
206 dictionary of labels that you wish to modify. Key is the existing label
207 and item is the new label
208 """
209 def __init__(self, opts, ignore_copy=False):
210 logger.info("Command line arguments: %s" % (opts))
211 self.opts = opts
212 self.existing = None
213 self.webdir = self.opts.webdir
214 self.samples = self.opts.samples
215 self.labels = self.opts.labels
216 self.kwargs = self.opts.kwargs
217 self.config = self.opts.config
218 self.replace_posterior = self.opts.replace_posterior
219 self.remove_label = self.opts.remove_label
220 self.remove_posterior = self.opts.remove_posterior
221 self.store_skymap = self.opts.store_skymap
222 self.hdf5 = not self.opts.save_to_json
223 self.overwrite = self.opts.overwrite
224 self.force_replace = self.opts.force_replace
225 self.data = None
226 self.stored_labels = [
227 key for key in self.data.keys() if key not in
228 ["history", "strain", "version"]
229 ]
230 if self.opts.descriptions is not None:
231 import copy
232 self._labels_copy = copy.deepcopy(self._labels)
233 self._labels = self.stored_labels
234 self.descriptions = self.opts.descriptions
235 self._labels = self._labels_copy
236 else:
237 self._descriptions = None
240class ArgumentParser(_ArgumentParser):
241 def _pesummary_options(self):
242 options = super(ArgumentParser, self)._pesummary_options()
243 options.update(
244 {
245 "--labels": {
246 "nargs": "+",
247 "action": DelimiterSplitAction,
248 "help": (
249 "labels you wish to modify. Syntax: `--labels "
250 "existing:new` where ':' is the default delimiter"
251 )
252 },
253 "--config": {
254 "nargs": "+",
255 "action": DelimiterSplitAction,
256 "help": (
257 "config data you wish to modify. Syntax `--config "
258 "label:path` where label is the analysis you wish to "
259 "change and path is the path to a new configuration "
260 "file."
261 )
262 },
263 "--delimiter": {
264 "default": ":",
265 "help": (
266 "Delimiter used to seperate the existing and new "
267 "quantity"
268 )
269 },
270 "--kwargs": {
271 "nargs": "+",
272 "action": DelimiterSplitAction,
273 "help": (
274 "kwargs you wish to modify. Syntax: `--kwargs "
275 "label/kwarg:item` where '/' is a delimiter of your "
276 "choosing (it cannot be ':'), kwarg is the kwarg name "
277 "and item is the value of the kwarg"
278 )
279 },
280 "--overwrite": {
281 "action": "store_true",
282 "default": False,
283 "help": (
284 "Overwrite the supplied PESummary meta file with the "
285 "modified version"
286 )
287 },
288 "--replace_posterior": {
289 "nargs": "+",
290 "action": DelimiterSplitAction,
291 "help": (
292 "Replace the posterior for a given label. Syntax: "
293 "--replace_posterior label;a:/path/to/posterior.dat "
294 "where ';' is a delimiter of your choosing (it cannot "
295 "be '/' or ':'), a is the posterior you wish to "
296 "replace and item is a path to a one column ascii file "
297 "containing the posterior samples "
298 "(/path/to/posterior.dat)"
299 )
300 },
301 "--remove_posterior": {
302 "nargs": "+",
303 "action": DelimiterSplitAction,
304 "help": (
305 "Remove a posterior distribution for a given label. "
306 "Syntax: --remove_posterior label:a where a is the "
307 "posterior you wish to remove"
308 )
309 },
310 "--remove_label": {
311 "nargs": "+",
312 "help": "Remove an entire analysis from the input file"
313 },
314 "--store_skymap": {
315 "nargs": "+",
316 "action": DelimiterSplitAction,
317 "help": (
318 "Store the contents of a fits file in the metafile. "
319 "Syntax: --store_skymap label:path/to/skymap.fits"
320 )
321 },
322 "--force_replace": {
323 "action": "store_true",
324 "default": False,
325 "help": (
326 "Override the ValueError raised if the data is already "
327 "stored in the result file"
328 )
329 }
330 }
331 )
332 options["--descriptions"]["action"] = DelimiterSplitAction
333 return options
336def _check_label(data, label, message, logger_level="warn"):
337 """Check that a given label is stored in the data. If it is not stored
338 print a warning message
340 Parameters
341 ----------
342 data: dict
343 dictionary containing the data
344 label: str
345 name of the label you wish to check
346 message: str
347 message you wish to print in logger when the label is not stored
348 logger_level: str, optional
349 the logger level of the message
350 """
351 if label not in data.keys():
352 getattr(logger, logger_level)(message)
353 return False
354 return True
357def _modify_labels(data, labels=None):
358 """Modify the existing labels in the data
360 Parameters
361 ----------
362 data: dict
363 dictionary containing the data
364 labels: dict
365 dictionary of labels showing the existing label, key, and the new
366 label, item
367 """
368 for existing, new in labels.items():
369 if existing not in data.keys():
370 logger.warning(
371 "Unable to find label '{}' in the root of the metafile. "
372 "Checking inside the groups".format(existing)
373 )
374 for key in data.keys():
375 if existing in data[key].keys():
376 data[key][new] = data[key].pop(existing)
377 else:
378 data[new] = data.pop(existing)
379 return data
382def _modify_descriptions(data, descriptions={}):
383 """Modify the existing descriptions in the data
385 Parameters
386 ----------
387 data: dict
388 dictionary containing the data
389 descriptions: dict
390 dictionary of descriptions with label as the key and new description as
391 the item
392 """
393 message = (
394 "Unable to find label '{}' in the metafile. Unable to modify "
395 "description"
396 )
397 for label, new_desc in descriptions.items():
398 check = _check_label(data, label, message.format(label))
399 if check:
400 if "description" not in data[label].keys():
401 data[label]["description"] = []
402 data[label]["description"] = [new_desc]
403 return data
406def _modify_kwargs(data, kwargs=None):
407 """Modify kwargs that are stored in the data
409 Parameters
410 ----------
411 data: dict
412 dictionary containing the data
413 kwargs: dict
414 dictionary of kwargs showing the label as key and kwarg:value as the
415 item
416 """
417 def add_to_meta_data(data, label, string):
418 kwarg, value = string.split(":")
419 try:
420 _group, = paths_to_key(kwarg, data[label]["meta_data"])
421 group = _group[0]
422 except ValueError:
423 group = "other"
424 if group == "other" and group not in data[label]["meta_data"].keys():
425 data[label]["meta_data"]["other"] = {}
426 data[label]["meta_data"][group][kwarg] = value
427 return data
429 message = "Unable to find label '{}' in the metafile. Unable to modify kwargs"
430 for label, item in kwargs.items():
431 check = _check_label(data, label, message.format(label))
432 if check:
433 if isinstance(item, list):
434 for _item in item:
435 data = add_to_meta_data(data, label, _item)
436 else:
437 data = add_to_meta_data(data, label, item)
438 return data
441def _modify_config(data, kwargs=None):
442 """Replace the config data that is stored in the data
444 Parameters
445 ----------
446 data: dict
447 dictionary containing the data
448 kwargs: dict
449 dictionary of kwargs showing the label as key and config data as the
450 item
451 """
452 message = (
453 "Unable to find label '{}' in the metafile. Unable to modify config"
454 )
455 for label, config in kwargs.items():
456 check = _check_label(data, label, message.format(label))
457 if check:
458 data[label]["config_file"] = config
459 return data
462def _modify_posterior(data, kwargs=None):
463 """Replace a posterior distribution that is stored in the data
465 Parameters
466 ----------
467 data: dict
468 dictionary containing the data
469 kwargs: dict
470 dictionary of kwargs showing the label as key and posterior:path as the
471 item
472 """
473 def _replace_posterior(data, string):
474 posterior, path = string.split(":")
475 _data = np.genfromtxt(path, usecols=0)
476 if math.isnan(_data[0]):
477 _data = np.genfromtxt(path, names=True, usecols=0)
478 _data = _data[_data.dtype.names[0]]
479 if posterior in data[label]["posterior_samples"].dtype.names:
480 data[label]["posterior_samples"][posterior] = _data
481 else:
482 from numpy.lib.recfunctions import append_fields
484 data[label]["posterior_samples"] = append_fields(
485 data[label]["posterior_samples"], posterior, _data, usemask=False
486 )
487 return data
489 message = "Unable to find label '{}' in the metafile. Unable to modify posterior"
490 for label, item in kwargs.items():
491 check = _check_label(data, label, message.format(label))
492 if check:
493 if isinstance(item, list):
494 for _item in item:
495 data = _replace_posterior(data, _item)
496 else:
497 data = _replace_posterior(data, item)
498 return data
501def _remove_label(data, kwargs=None):
502 """Remove an analysis that is stored in the data
504 Parameters
505 ----------
506 data: dict
507 dictionary containing the data
508 kwargs: list
509 list of analysis you wish to remove
510 """
511 message = (
512 "Unable to find label '{}' in the metafile. Unable to remove"
513 )
514 for label in kwargs:
515 check = _check_label(data, label, message.format(label))
516 if check:
517 _ = data.pop(label)
518 return data
521def _remove_posterior(data, kwargs=None):
522 """Remove a posterior distribution that is stored in the data
524 Parameters
525 ----------
526 data: dict
527 dictionary containing the data
528 kwargs: dict
529 dictionary of kwargs showing the label as key and posterior as the item
530 """
531 def _rmfield(array, *fieldnames_to_remove):
532 return array[
533 [name for name in array.dtype.names if name not in fieldnames_to_remove]
534 ]
536 message = "Unable to find label '{}' in the metafile. Unable to remove posterior"
537 for label, item in kwargs.items():
538 check = _check_label(data, label, message.format(label))
539 if check:
540 group = "posterior_samples"
541 if isinstance(item, list):
542 for _item in item:
543 data[label][group] = _rmfield(data[label][group], _item)
544 else:
545 data[label][group] = _rmfield(data[label][group], item)
546 return data
549def _store_skymap(data, kwargs=None, replace=False):
550 """Store a skymap in the metafile
552 Parameters
553 ----------
554 data: dict
555 dictionary containing the data
556 kwargs: dict
557 dictionary of kwargs showing the label as key and posterior as the item
558 replace: dict
559 replace a skymap already stored in the result file
560 """
561 from pesummary.io import read
563 message = "Unable to find label '{}' in the metafile. Unable to store skymap"
564 for label, path in kwargs.items():
565 check = _check_label(data, label, message.format(label))
566 if check:
567 skymap = read(path, skymap=True)
568 if "skymap" not in data[label].keys():
569 data[label]["skymap"] = {}
570 if "meta_data" not in data[label]["skymap"].keys():
571 data[label]["skymap"]["meta_data"] = {}
572 if "data" in data[label]["skymap"].keys() and not replace:
573 raise ValueError(
574 "Skymap already found in result file for {}. If you wish to replace "
575 "the skymap, add the command line argument '--force_replace".format(
576 label
577 )
578 )
579 elif "data" in data[label]["skymap"].keys():
580 logger.warning("Replacing skymap data for {}".format(label))
581 data[label]["skymap"]["data"] = skymap
582 for key in skymap.meta_data:
583 data[label]["skymap"]["meta_data"][key] = skymap.meta_data[key]
584 return data
587def modify(data, function, **kwargs):
588 """Modify the data according to a given function
590 Parameters
591 ----------
592 data: dict
593 dictionary containing the data
594 function:
595 function you wish to use to modify the data
596 kwargs: dict
597 dictionary of kwargs for function
598 """
599 func_map = {
600 "labels": _modify_labels,
601 "descriptions": _modify_descriptions,
602 "kwargs": _modify_kwargs,
603 "add_posterior": _modify_posterior,
604 "rm_label": _remove_label,
605 "rm_posterior": _remove_posterior,
606 "skymap": _store_skymap,
607 "config": _modify_config
608 }
609 return func_map[function](data, **kwargs)
612def _main(opts):
613 """
614 """
615 args = Input(opts)
616 if not args.overwrite:
617 meta_file = os.path.join(
618 args.webdir, "modified_posterior_samples.{}".format(
619 "h5" if args.hdf5 else "json"
620 )
621 )
622 check_file_exists_and_rename(meta_file)
623 else:
624 meta_file = args.samples
625 if opts.preferred is not None and args.kwargs is None:
626 args.kwargs = {opts.preferred: "preferred:True"}
627 if opts.preferred is not None:
628 args.kwargs.update(
629 {
630 _label: "preferred:False" for _label in args.stored_labels
631 if _label != opts.preferred
632 }
633 )
634 if args.labels is not None:
635 modified_data = modify(args.data, "labels", labels=args.labels)
636 if args.descriptions is not None:
637 modified_data = modify(
638 args.data, "descriptions", descriptions=args.descriptions
639 )
640 if args.config is not None:
641 modified_data = modify(args.data, "config", kwargs=args.config)
642 if args.kwargs is not None:
643 modified_data = modify(args.data, "kwargs", kwargs=args.kwargs)
644 if args.replace_posterior is not None:
645 modified_data = modify(args.data, "add_posterior", kwargs=args.replace_posterior)
646 if args.remove_label is not None:
647 modified_data = modify(args.data, "rm_label", kwargs=args.remove_label)
648 if args.remove_posterior is not None:
649 modified_data = modify(args.data, "rm_posterior", kwargs=args.remove_posterior)
650 if args.store_skymap is not None:
651 modified_data = modify(
652 args.data, "skymap", kwargs=args.store_skymap, replace=args.force_replace
653 )
654 logger.info(
655 "Saving the modified data to '{}'".format(meta_file)
656 )
657 if args.hdf5:
658 _GWMetaFile.save_to_hdf5(
659 modified_data, list(modified_data.keys()), None, meta_file,
660 no_convert=True
661 )
662 else:
663 _GWMetaFile.save_to_json(modified_data, meta_file)
666def command_line():
667 parser = ArgumentParser(description=__doc__)
668 parser.add_known_options_to_parser(
669 [
670 "--descriptions", "--labels", "--config", "--delimiter",
671 "--samples", "--webdir", "--save_to_json", "--preferred",
672 "--kwargs", "--overwrite", "--replace_posterior",
673 "--remove_posterior", "--remove_label", "--store_skymap",
674 "--force_replace"
675 ]
676 )
677 return parser
680def main(args=None):
681 """
682 """
683 parser = command_line()
684 opts = parser.parse_args(args=args)
685 return _main(opts)
688if __name__ == "__main__":
689 main()