Coverage for bilby/core/sampler/polychord.py: 31%

1import os 

2 

3import numpy as np 

4 

5from .base_sampler import NestedSampler, signal_wrapper 

6 

7 

8class PyPolyChord(NestedSampler): 

9 

10 """ 

11 Bilby wrapper of PyPolyChord 

12 https://arxiv.org/abs/1506.00171 

13 

14 PolyChordLite is available at: 

15 https://github.com/PolyChord/PolyChordLite 

16 

17 Follow the installation instructions at their github page. 

18 

19 Keyword arguments will be passed into `pypolychord.run_polychord` into the `settings` 

20 argument. See the PolyChord documentation for what all of those mean. 

21 

22 To see what the keyword arguments are for, see the docstring of PyPolyChordSettings 

23 """ 

24 

25 sampler_name = "pypolychord" 

26 default_kwargs = dict( 

27 use_polychord_defaults=False, 

28 nlive=None, 

29 num_repeats=None, 

30 nprior=-1, 

31 do_clustering=True, 

32 feedback=1, 

33 precision_criterion=0.001, 

34 logzero=-1e30, 

35 max_ndead=-1, 

36 boost_posterior=0.0, 

37 posteriors=True, 

38 equals=True, 

39 cluster_posteriors=True, 

40 write_resume=True, 

41 write_paramnames=False, 

42 read_resume=True, 

43 write_stats=True, 

44 write_live=True, 

45 write_dead=True, 

46 write_prior=True, 

47 compression_factor=np.exp(-1), 

48 base_dir="outdir", 

49 file_root="polychord", 

50 seed=-1, 

51 grade_dims=None, 

52 grade_frac=None, 

53 nlives={}, 

54 ) 

55 hard_exit = True 

56 sampling_seed_key = "seed" 

57 

58 @signal_wrapper 

59 def run_sampler(self): 

60 import pypolychord 

61 from pypolychord.settings import PolyChordSettings 

62 

63 if self.kwargs["use_polychord_defaults"]: 

64 settings = PolyChordSettings( 

65 nDims=self.ndim, 

66 nDerived=self.ndim, 

67 base_dir=self._sample_file_directory, 

68 file_root=self.label, 

69 ) 

70 else: 

71 self._setup_dynamic_defaults() 

72 pc_kwargs = self.kwargs.copy() 

73 pc_kwargs["base_dir"] = self._sample_file_directory 

74 pc_kwargs["file_root"] = self.label 

75 pc_kwargs.pop("use_polychord_defaults") 

76 settings = PolyChordSettings( 

77 nDims=self.ndim, nDerived=self.ndim, **pc_kwargs 

78 ) 

79 self._verify_kwargs_against_default_kwargs() 

80 out = pypolychord.run_polychord( 

81 loglikelihood=self.log_likelihood, 

82 nDims=self.ndim, 

83 nDerived=self.ndim, 

84 settings=settings, 

85 prior=self.prior_transform, 

86 ) 

87 self.result.log_evidence = out.logZ 

88 self.result.log_evidence_err = out.logZerr 

89 log_likelihoods, physical_parameters = self._read_sample_file() 

90 self.result.log_likelihood_evaluations = log_likelihoods 

91 self.result.samples = physical_parameters 

92 self.calc_likelihood_count() 

93 return self.result 

94 

95 def _setup_dynamic_defaults(self): 

96 """Sets up some interdependent default argument if none are given by the user""" 

97 if not self.kwargs["grade_dims"]: 

98 self.kwargs["grade_dims"] = [self.ndim] 

99 if not self.kwargs["grade_frac"]: 

100 self.kwargs["grade_frac"] = [1.0] * len(self.kwargs["grade_dims"]) 

101 if not self.kwargs["nlive"]: 

102 self.kwargs["nlive"] = self.ndim * 25 

103 if not self.kwargs["num_repeats"]: 

104 self.kwargs["num_repeats"] = self.ndim * 5 

105 

106 def _translate_kwargs(self, kwargs): 

107 kwargs = super()._translate_kwargs(kwargs) 

108 if "nlive" not in kwargs: 

109 for equiv in self.npoints_equiv_kwargs: 

110 if equiv in kwargs: 

111 kwargs["nlive"] = kwargs.pop(equiv) 

112 

113 def log_likelihood(self, theta): 

114 """Overrides the log_likelihood so that PolyChord understands it""" 

115 return super(PyPolyChord, self).log_likelihood(theta), theta 

116 

117 def _read_sample_file(self): 

118 """ 

119 This method reads out the _equal_weights.txt file that polychord produces. 

120 The first column is omitted since it is just composed of 1s, i.e. the equal weights/ 

121 The second column are the log likelihoods, the remaining columns are the physical parameters 

122 

123 Returns 

124 ======= 

125 array_like, array_like: The log_likelihoods and the associated parameters 

126 

127 """ 

128 sample_file = ( 

129 self._sample_file_directory + "/" + self.label + "_equal_weights.txt" 

130 ) 

131 samples = np.loadtxt(sample_file) 

132 log_likelihoods = -0.5 * samples[:, 1] 

133 physical_parameters = samples[:, -self.ndim :] 

134 return log_likelihoods, physical_parameters 

135 

136 @classmethod 

137 def get_expected_outputs(cls, outdir=None, label=None): 

138 """Get lists of the expected outputs directories and files. 

139 

140 These are used by :code:`bilby_pipe` when transferring files via HTCondor. 

141 

142 Parameters 

143 ---------- 

144 outdir : str 

145 The output directory. 

146 label : str 

147 Ignored for pypolychord. 

148 

149 Returns 

150 ------- 

151 list 

152 List of file names. This will always be empty for pypolychord. 

153 list 

154 List of directory names. 

155 """ 

156 return [], [os.path.join(outdir, "chains", "")] 

157 

158 @property 

159 def _sample_file_directory(self): 

160 return self.outdir + "/chains"