Coverage for pesummary/cli/summarymodify.py: 91.1%

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

100

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 )

111

112 @property

113 def store_skymap(self):

114 return self._store_skymap

115

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 )

126

127 @property

128 def samples(self):

129 return self._samples

130

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

147

148 @property

149 def data(self):

150 return self._data

151

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

158

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 )

191

192

193class Input(_Input):

194 """Class to handle the command line arguments

195

196 Parameters

197 ----------

198 opts: argparse.Namespace

199 Namespace object containing the command line options

200

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

238

239

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

334

335

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

339

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

355

356

357def _modify_labels(data, labels=None):

358 """Modify the existing labels in the data

359

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

380

381

382def _modify_descriptions(data, descriptions={}):

383 """Modify the existing descriptions in the data

384

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

404

405

406def _modify_kwargs(data, kwargs=None):

407 """Modify kwargs that are stored in the data

408

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

428

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

439

440

441def _modify_config(data, kwargs=None):

442 """Replace the config data that is stored in the data

443

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

460

461

462def _modify_posterior(data, kwargs=None):

463 """Replace a posterior distribution that is stored in the data

464

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

483

484 data[label]["posterior_samples"] = append_fields(

485 data[label]["posterior_samples"], posterior, _data, usemask=False

486 )

487 return data

488

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

499

500

501def _remove_label(data, kwargs=None):

502 """Remove an analysis that is stored in the data

503

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

519

520

521def _remove_posterior(data, kwargs=None):

522 """Remove a posterior distribution that is stored in the data

523

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 ]

535

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

547

548

549def _store_skymap(data, kwargs=None, replace=False):

550 """Store a skymap in the metafile

551

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

562

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

585

586

587def modify(data, function, **kwargs):

588 """Modify the data according to a given function

589

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)

610

611

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)

664

665

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

678

679

680def main(args=None):

681 """

682 """

683 parser = command_line()

684 opts = parser.parse_args(args=args)

685 return _main(opts)

686

687

688if __name__ == "__main__":

689 main()