Coverage for bilby/core/fisher.py: 78%

1import numpy as np 

2import pandas as pd 

3import scipy.linalg 

4from scipy.optimize import minimize 

5 

6 

7class FisherMatrixPosteriorEstimator(object): 

8 def __init__(self, likelihood, priors, parameters=None, fd_eps=1e-6, n_prior_samples=100): 

9 """ A class to estimate posteriors using the Fisher Matrix approach 

10 

11 Parameters 

12 ---------- 

13 likelihood: bilby.core.likelihood.Likelihood 

14 A bilby likelihood object 

15 priors: bilby.core.prior.PriorDict 

16 A bilby prior object 

17 parameters: list 

18 Names of parameters to sample in 

19 fd_eps: float 

20 A parameter to control the size of perturbation used when finite 

21 differencing the likelihood 

22 n_prior_samples: int 

23 The number of prior samples to draw and use to attempt estimatation 

24 of the maximum likelihood sample. 

25 """ 

26 self.likelihood = likelihood 

27 if parameters is None: 

28 self.parameter_names = priors.non_fixed_keys 

29 else: 

30 self.parameter_names = parameters 

31 self.fd_eps = fd_eps 

32 self.n_prior_samples = n_prior_samples 

33 self.N = len(self.parameter_names) 

34 

35 self.prior_samples = [ 

36 priors.sample_subset(self.parameter_names) for _ in range(n_prior_samples) 

37 ] 

38 self.prior_bounds = [(priors[key].minimum, priors[key].maximum) for key in self.parameter_names] 

39 

40 self.prior_width_dict = {} 

41 for key in self.parameter_names: 

42 width = priors[key].width 

43 if np.isnan(width): 

44 raise ValueError(f"Prior width is ill-formed for {key}") 

45 self.prior_width_dict[key] = width 

46 

47 def log_likelihood(self, sample): 

48 self.likelihood.parameters.update(sample) 

49 return self.likelihood.log_likelihood() 

50 

51 def calculate_iFIM(self, sample): 

52 FIM = self.calculate_FIM(sample) 

53 iFIM = scipy.linalg.inv(FIM) 

54 

55 # Ensure iFIM is positive definite 

56 min_eig = np.min(np.real(np.linalg.eigvals(iFIM))) 

57 if min_eig < 0: 

58 iFIM -= 10 * min_eig * np.eye(*iFIM.shape) 

59 

60 return iFIM 

61 

62 def sample_array(self, sample, n=1): 

63 from .utils.random import rng 

64 

65 if sample == "maxL": 

66 sample = self.get_maximum_likelihood_sample() 

67 

68 self.mean = np.array(list(sample.values())) 

69 self.iFIM = self.calculate_iFIM(sample) 

70 return rng.multivariate_normal(self.mean, self.iFIM, n) 

71 

72 def sample_dataframe(self, sample, n=1): 

73 samples = self.sample_array(sample, n) 

74 return pd.DataFrame(samples, columns=self.parameter_names) 

75 

76 def calculate_FIM(self, sample): 

77 FIM = np.zeros((self.N, self.N)) 

78 for ii, ii_key in enumerate(self.parameter_names): 

79 for jj, jj_key in enumerate(self.parameter_names): 

80 FIM[ii, jj] = -self.get_second_order_derivative(sample, ii_key, jj_key) 

81 

82 return FIM 

83 

84 def get_second_order_derivative(self, sample, ii, jj): 

85 if ii == jj: 

86 return self.get_finite_difference_xx(sample, ii) 

87 else: 

88 return self.get_finite_difference_xy(sample, ii, jj) 

89 

90 def get_finite_difference_xx(self, sample, ii): 

91 # Sample grid 

92 p = self.shift_sample_x(sample, ii, 1) 

93 m = self.shift_sample_x(sample, ii, -1) 

94 

95 dx = .5 * (p[ii] - m[ii]) 

96 

97 loglp = self.log_likelihood(p) 

98 logl = self.log_likelihood(sample) 

99 loglm = self.log_likelihood(m) 

100 

101 return (loglp - 2 * logl + loglm) / dx ** 2 

102 

103 def get_finite_difference_xy(self, sample, ii, jj): 

104 # Sample grid 

105 pp = self.shift_sample_xy(sample, ii, 1, jj, 1) 

106 pm = self.shift_sample_xy(sample, ii, 1, jj, -1) 

107 mp = self.shift_sample_xy(sample, ii, -1, jj, 1) 

108 mm = self.shift_sample_xy(sample, ii, -1, jj, -1) 

109 

110 dx = .5 * (pp[ii] - mm[ii]) 

111 dy = .5 * (pp[jj] - mm[jj]) 

112 

113 loglpp = self.log_likelihood(pp) 

114 loglpm = self.log_likelihood(pm) 

115 loglmp = self.log_likelihood(mp) 

116 loglmm = self.log_likelihood(mm) 

117 

118 return (loglpp - loglpm - loglmp + loglmm) / (4 * dx * dy) 

119 

120 def shift_sample_x(self, sample, x_key, x_coef): 

121 

122 vx = sample[x_key] 

123 dvx = self.fd_eps * self.prior_width_dict[x_key] 

124 

125 shift_sample = sample.copy() 

126 shift_sample[x_key] = vx + x_coef * dvx 

127 

128 return shift_sample 

129 

130 def shift_sample_xy(self, sample, x_key, x_coef, y_key, y_coef): 

131 

132 vx = sample[x_key] 

133 vy = sample[y_key] 

134 

135 dvx = self.fd_eps * self.prior_width_dict[x_key] 

136 dvy = self.fd_eps * self.prior_width_dict[y_key] 

137 

138 shift_sample = sample.copy() 

139 shift_sample[x_key] = vx + x_coef * dvx 

140 shift_sample[y_key] = vy + y_coef * dvy 

141 return shift_sample 

142 

143 def get_maximum_likelihood_sample(self, initial_sample=None): 

144 """ A method to attempt optimization of the maximum likelihood 

145 

146 This uses a simple scipy optimization approach, starting from a number 

147 of draws from the prior to avoid problems with local optimization. 

148 

149 Note: this approach works well in small numbers of dimensions when the 

150 posterior is narrow relative to the prior. But, if the number of dimensions 

151 is large or the posterior is wide relative to the prior, the method fails 

152 to find the global maximum in high dimensional problems. 

153 """ 

154 minlogL = np.inf 

155 for i in range(self.n_prior_samples): 

156 initial_sample = self.prior_samples[i] 

157 

158 x0 = list(initial_sample.values()) 

159 

160 def neg_log_like(x, self, T=1): 

161 sample = {key: val for key, val in zip(self.parameter_names, x)} 

162 return - 1 / T * self.log_likelihood(sample) 

163 

164 out = minimize( 

165 neg_log_like, 

166 x0, 

167 args=(self, 1), 

168 bounds=self.prior_bounds, 

169 method="L-BFGS-B", 

170 ) 

171 if out.fun < minlogL: 

172 minout = out 

173 

174 return {key: val for key, val in zip(self.parameter_names, minout.x)}