Coverage for pesummary/core/file/formats/numpy.py: 75.0%

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

2 

3import numpy as np 

4from pesummary.utils.utils import check_filename 

5 

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

7 

8 

9def _parameters_and_samples_from_structured_array(array): 

10 """Return the parameters and samples stored in a structured array 

11 

12 Parameters 

13 ---------- 

14 array: numpy.ndarray 

15 structured array containing the parameters and samples. Each column 

16 should correspond to the samples for a single distribution 

17 """ 

18 parameters = list(array.dtype.names) 

19 array = np.atleast_1d(array) 

20 samples = array.view((float, len(array.dtype.names))) 

21 return parameters, samples 

22 

23 

24def genfromtxt(path, **kwargs): 

25 """Return the parameters and samples stored in a `txt` file 

26 

27 Parameters 

28 ---------- 

29 path: str 

30 path to the result file you wish to read in 

31 **kwargs: dict, optional 

32 all additional kwargs are passed to the `np.genfromtxt` function 

33 """ 

34 data = np.genfromtxt(path, **kwargs) 

35 return _parameters_and_samples_from_structured_array(data) 

36 

37 

38def read_numpy(path, **kwargs): 

39 """Grab the parameters and samples in a .npy file 

40 

41 Parameters 

42 ---------- 

43 path: str 

44 path to the result file you wish to read in 

45 **kwargs: dict, optional 

46 all additional kwargs are passed to the `np.load` function 

47 """ 

48 data = load(path, **kwargs) 

49 return _parameters_and_samples_from_structured_array(data) 

50 

51 

52def load(*args, **kwargs): 

53 """Load a .npy file using the `np.load` function 

54 

55 Parameters 

56 ---------- 

57 *args: tuple 

58 all args passed to `np.load` 

59 **kwargs: dict 

60 all kwargs passed to `np.load` 

61 """ 

62 return np.load(*args, **kwargs) 

63 

64 

65def _write_numpy( 

66 parameters, samples, outdir="./", label=None, filename=None, overwrite=False, 

67 default_filename="pesummary_{}.npy", **kwargs 

68): 

69 """Write a set of samples to a csv file 

70 

71 Parameters 

72 ---------- 

73 parameters: list 

74 list of parameters 

75 samples: 2d list 

76 list of samples. Columns correspond to a given parameter 

77 outdir: str, optional 

78 directory to write the dat file 

79 label: str, optional 

80 The label of the analysis. This is used in the filename if a filename 

81 if not specified 

82 filename: str, optional 

83 The name of the file that you wish to write 

84 overwrite: Bool, optional 

85 If True, an existing file of the same name will be overwritten 

86 """ 

87 from pesummary.utils.samples_dict import SamplesDict 

88 

89 filename = check_filename( 

90 default_filename=default_filename, outdir=outdir, label=label, 

91 filename=filename, overwrite=overwrite 

92 ) 

93 _array = SamplesDict( 

94 parameters, np.array(samples).T 

95 ).to_structured_array() 

96 np.save(filename, _array) 

97 

98 

99def write_numpy( 

100 parameters, samples, outdir="./", label=None, filename=None, overwrite=False, 

101 **kwargs 

102): 

103 """Write a set of samples to a npy file 

104 

105 Parameters 

106 ---------- 

107 parameters: nd list 

108 list of parameters 

109 samples: nd list 

110 list of samples. Columns correspond to a given parameter 

111 outdir: str, optional 

112 directory to write the csv file 

113 label: str, optional 

114 The label of the analysis. This is used in the filename if a filename 

115 if not specified 

116 filename: str, optional 

117 The name of the file that you wish to write 

118 overwrite: Bool, optional 

119 If True, an existing file of the same name will be overwritten 

120 """ 

121 from pesummary.io.write import _multi_analysis_write 

122 

123 _multi_analysis_write( 

124 _write_numpy, parameters, samples, outdir=outdir, label=label, 

125 filename=filename, overwrite=overwrite, file_format="numpy", **kwargs 

126 )