Coverage for pesummary/io/write.py: 91.5%

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

2 

3import importlib 

4 

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

6 

7 

8def write(*args, package="gw", file_format="dat", **kwargs): 

9 """Read in a results file. 

10 

11 Parameters 

12 ---------- 

13 args: tuple 

14 all args are passed to write function 

15 package: str 

16 the package you wish to use 

17 file_format: str 

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

19 function loops through all possible options 

20 kwargs: dict 

21 all kwargs passed to write function 

22 """ 

23 def _import(package, file_format): 

24 """Import format module with importlib 

25 """ 

26 return importlib.import_module( 

27 "pesummary.{}.file.formats.{}".format(package, file_format) 

28 ) 

29 

30 def _write(module, file_format, args, kwargs): 

31 """Execute the write method 

32 """ 

33 return getattr(module, "write_{}".format(file_format))(*args, **kwargs) 

34 

35 if file_format == "h5": 

36 file_format = "hdf5" 

37 

38 try: 

39 module = _import(package, file_format) 

40 return _write(module, file_format, args, kwargs) 

41 except (ImportError, AttributeError, ModuleNotFoundError): 

42 module = _import("core", file_format) 

43 return _write(module, file_format, args, kwargs) 

44 

45 

46def _multi_analysis_write( 

47 write_function, parameters, *args, labels=None, 

48 file_format="dat", outdir="./", filename=None, filenames=[], 

49 _return=False, **kwargs 

50): 

51 """Write a set of samples to file. If a 2d list of parameters are provided 

52 (i.e. samples from multiple analyses), the samples are saved to separate 

53 files 

54 

55 Parameters 

56 ---------- 

57 write_function: func 

58 the write function you wish to execute 

59 parameters: nd list 

60 list of parameters for a single analysis, or a 2d of parameters from 

61 multiple analyses 

62 *args: tuple 

63 all other arguments are passed to write_function 

64 labels: list, optional 

65 labels to use for identifying different analyses. Only used if 

66 parameters is a 2d list 

67 file_format: str, optional 

68 file format you wish to save the file to. Default 'dat' 

69 outdir: str, optional 

70 directory to save the files. Default './' 

71 filename: str, optional 

72 filename to save a single analyses to. If multiple analyses are to be 

73 saved, the filenames are '{filename}_{label}.{extension}'. Default 

74 'pesummary.{file_format}' 

75 filenames: list, optional 

76 filenames to save the files to. Only used if parameters is a 2d list 

77 **kwargs: dict, optional 

78 all other kwargs are passed to write_function 

79 """ 

80 import numpy as np 

81 from pathlib import Path 

82 from pesummary.utils.utils import logger 

83 import copy 

84 

85 _filenames = copy.deepcopy(filenames) 

86 _file_dict = {} 

87 if np.array(parameters).ndim > 1: 

88 log = True 

89 _label = kwargs.get("label", None) 

90 if labels is None and _label is not None: 

91 labels = [ 

92 "{}_{}".format(_label, idx) for idx in range(len(parameters)) 

93 ] 

94 elif labels is None: 

95 labels = np.arange(len(parameters)) 

96 if len(_filenames) and len(_filenames) != len(parameters): 

97 raise ValueError("Please provide a filename for each analysis") 

98 elif len(_filenames): 

99 log = False 

100 elif not len(_filenames): 

101 if filename is None: 

102 filename = "pesummary.{}".format(file_format) 

103 filename = Path(filename) 

104 for idx in range(len(parameters)): 

105 _filenames.append( 

106 "{}_{}{}".format(filename.stem, idx, filename.suffix) 

107 ) 

108 if log: 

109 logger.info( 

110 "Only a single set of samples can be written to a {} file. " 

111 "Saving each analysis to a separate file. Filenames will be: " 

112 "{}".format(file_format, ", ".join(_filenames)) 

113 ) 

114 for idx in range(len(parameters)): 

115 _args = [a[idx] for a in args] 

116 _file_dict[labels[idx]] = write_function( 

117 parameters[idx], *_args, file_format=file_format, outdir=outdir, 

118 filename=_filenames[idx], **kwargs 

119 ) 

120 if _return: 

121 return _file_dict 

122 else: 

123 return write_function( 

124 parameters, *args, file_format=file_format, outdir=outdir, 

125 filename=filename, **kwargs 

126 )