Coverage for frank/filter.py: 60%

1# Frankenstein: 1D disc brightness profile reconstruction from Fourier data 

2# using non-parametric Gaussian Processes 

3# 

4# Copyright (C) 2019-2020 R. Booth, J. Jennings, M. Tazzari 

5# 

6# This program is free software: you can redistribute it and/or modify 

7# it under the terms of the GNU General Public License as published by 

8 

9# the Free Software Foundation, either version 3 of the License, or 

10# (at your option) any later version. 

11# 

12# This program is distributed in the hope that it will be useful, 

13# but WITHOUT ANY WARRANTY; without even the implied warranty of 

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

15# GNU General Public License for more details. 

16# 

17# You should have received a copy of the GNU General Public License 

18# along with this program. If not, see <https://www.gnu.org/licenses/> 

19# 

20import numpy as np 

21import scipy.sparse 

22 

23def spectral_smoothing_matrix(DHT, weights): 

24 r"""Compute the spectral smoothing prior matrix, Tij. 

25  

26 .. math:: 

27 log P(p|q, w) = -1/2 q.T . Tij . q 

28  

29 Parameters 

30 ---------- 

31 DHT : DiscreteHankelTransform object 

32 DHT to build the smoothing matrix for. 

33 weights : float > 0 

34 Weighting parameter used for the spectral smoothing matrix 

35 

36 Returns 

37 ------- 

38 Tij : Sparse Matrix 

39 Spectral smoothing p 

40 """ 

41 log_q = np.log(DHT.q) 

42 dc = (log_q[2:] - log_q[:-2]) / 2 

43 de = np.diff(log_q) 

44 

45 N = DHT.size 

46 

47 Delta = np.zeros([3, N]) 

48 Delta[0, :-2] = 1 / (dc * de[:-1]) 

49 Delta[1, 1:-1] = - (1 / de[1:] + 1 / de[:-1]) / dc 

50 Delta[2, 2:] = 1 / (dc * de[1:]) 

51 

52 Delta = scipy.sparse.dia_matrix((Delta, [-1, 0, 1]), 

53 shape=(N, N)) 

54 

55 dce = np.zeros_like(log_q) 

56 dce[1:-1] = dc 

57 dce = scipy.sparse.dia_matrix((dce.reshape(1, -1), 0), 

58 shape=(N, N)) 

59 

60 Tij = Delta.T.dot(dce.dot(Delta)) 

61 

62 return weights * Tij 

63 

64class CriticalFilter: 

65 """Optimizer for power-spectrum priors. 

66 

67 Parameters 

68 ---------- 

69 DHT : DiscreteHankelTransform 

70 A DHT object with N bins that defines H(p). The DHT is used to compute 

71 :math:`S(p)` 

72 alpha : float >= 1 

73 Order parameter of the inverse gamma prior for the power spectrum 

74 coefficients. 

75 p_0 : float >= 0, default = None, unit=Jy^2 

76 Scale parameter of the inverse gamma prior for the power spectrum 

77 coefficients. 

78 weights_smooth : float >= 0 

79 Spectral smoothness prior parameter. Zero is no smoothness prior 

80 tol : float > 0, default = 1e-3 

81 Tolerence for convergence of the power spectrum iteration. 

82 """ 

83 

84 def __init__(self, DHT, alpha, p_0, weights_smooth, 

85 tol=1e-3): 

86 

87 self._DHT = DHT 

88 

89 self._alpha = alpha 

90 self._p_0 = p_0 

91 self._rho = 1.0 

92 

93 self._tol = tol 

94 

95 self._Tij = spectral_smoothing_matrix(DHT, weights_smooth) 

96 

97 def update_power_spectrum_multiple(self, fit, binning=None): 

98 """Estimate the best fit power spectrum given the current model fit. 

99 This version works when there are multiple fields being fit simultaneously. 

100 """ 

101 

102 Nfields, Np = fit.power_spectrum.shape 

103 

104 if binning is None: 

105 binning = np.zeros(Nfields, dtype='i4') 

106 Nbins = binning.max()+1 

107 

108 Ykm = self._DHT.coefficients() 

109 Tij_pI = self._Tij + scipy.sparse.identity(self._DHT.size) 

110 

111 ds = fit.MAP 

112 

113 # Project mu to Fourier-space 

114 # Tr1 = Trace(mu mu_T . Ykm_T Ykm) = Trace( Ykm mu . (Ykm mu)^T) 

115 # = (Ykm mu)**2 

116 Tr1 = np.zeros([Nbins, Np]) 

117 Y = np.zeros([Np*Nfields, Np*Nfields]) 

118 for f in range(Nfields): 

119 s = f*Np 

120 e = s+Np 

121 

122 Tr1[binning[f]] += np.dot(Ykm, ds[f]) ** 2 

123 Y[s:e, s:e] = Ykm 

124 # Project D to Fourier-space 

125 # Drr^-1 = Ykm^T Dqq^-1 Ykm 

126 # Drr = Ykm^-1 Dqq Ykm^-T 

127 # Dqq = Ykm Drr Ykm^T 

128 # Tr2 = Trace(Dqq) 

129 tmp = np.einsum('ij,ji->i', Y, fit.Dsolve(Y.T)).reshape(Nfields, Np) 

130 

131 Tr2 = np.zeros([Nbins, Np]) 

132 ps = np.zeros([Nbins, Np]) 

133 count = np.zeros(Nbins) 

134 for f in range(Nfields): 

135 Tr2[binning[f]] += tmp[f] 

136 ps[binning[f]] += fit.power_spectrum[f] 

137 count[binning[f]] += 1 

138 ps /= count.reshape(Nbins, 1) 

139 

140 for i in range(Nbins): 

141 beta = (self._p_0 + 0.5 * (Tr1[i] + Tr2[i])) / ps[i] - \ 

142 (self._alpha - 1.0 + 0.5 * self._rho * count[i]) 

143 

144 tau = scipy.sparse.linalg.spsolve(Tij_pI, beta + np.log(ps[i])) 

145 ps[i] = np.exp(tau) 

146 

147 p = np.empty_like(fit.power_spectrum) 

148 for f in range(Nfields): 

149 p[f] = ps[binning[f]] 

150 

151 return p 

152 

153 

154 def update_power_spectrum(self, fit): 

155 """Estimate the best fit power spectrum given the current model fit""" 

156 Ykm = self._DHT.coefficients() 

157 Tij_pI = self._Tij + scipy.sparse.identity(self._DHT.size) 

158 

159 # Project mu to Fourier-space 

160 # Tr1 = Trace(mu mu_T . Ykm_T Ykm) = Trace( Ykm mu . (Ykm mu)^T) 

161 # = (Ykm mu)**2 

162 Tr1 = np.dot(Ykm, fit.MAP) ** 2 

163 # Project D to Fourier-space 

164 # Drr^-1 = Ykm^T Dqq^-1 Ykm 

165 # Drr = Ykm^-1 Dqq Ykm^-T 

166 # Dqq = Ykm Drr Ykm^T 

167 # Tr2 = Trace(Dqq) 

168 Tr2 = np.einsum('ij,ji->i', Ykm, fit.Dsolve(Ykm.T)) 

169 

170 pi = fit.power_spectrum 

171 

172 beta = (self._p_0 + 0.5 * (Tr1 + Tr2)) / pi - \ 

173 (self._alpha - 1.0 + 0.5 * self._rho) 

174 

175 tau = scipy.sparse.linalg.spsolve(Tij_pI, beta + np.log(pi)) 

176 

177 return np.exp(tau) 

178 

179 def check_convergence(self, pi_new, pi_old): 

180 """Determine whether the power-spectrum iterations have converged""" 

181 return np.all(np.abs(pi_new - pi_old) <= self._tol * pi_new) 

182 

183 

184 def covariance_MAP(self, fit, ret_inv=False): 

185 """ 

186 Covariance of the power spectrum at maximum likelihood 

187 

188 Parameters 

189 ---------- 

190 fit : _HankelRegressor 

191 Solution at maximum likelihood 

192 ret_inv : bool, default=False 

193 If True return the inverse covariance matrix instead. This  

194 avoids inverting the inverse covariace matrix 

195 

196 Returns 

197 ------- 

198 ps_cov : 2D array 

199 Covariance matrix of the power spectrum at maximum likelihood 

200 

201 Notes 

202 ----- 

203 Only valid at the location of maximum likelihood 

204 

205 """ 

206 Ykm = self._DHT.coefficients() 

207 

208 mq = np.dot(Ykm, fit.MAP) 

209 

210 mqq = np.outer(mq, mq) 

211 Dqq = np.dot(Ykm, np.dot(fit.covariance, Ykm.T)) 

212 

213 p = fit.power_spectrum 

214 

215 hess = \ 

216 + np.diag(self._p_0/p + 0.5*(mq**2 + np.diag(Dqq))/p) \ 

217 + self._Tij.todense() \ 

218 - 0.5 * np.outer(1 / p, 1 / p) * (2 * mqq + Dqq) * Dqq 

219 

220 if ret_inv: 

221 return hess 

222 else: 

223 # Invert the Hessian 

224 hess_chol = scipy.linalg.cho_factor(hess) 

225 ps_cov = scipy.linalg.cho_solve(hess_chol, np.eye(self._DHT.size)) 

226 

227 return ps_cov 

228 

229 def log_prior(self, p): 

230 r""" 

231 Compute the log prior probability, log(P(p)), 

232 

233 .. math: 

234 `log[P(p)] ~ - np.sum(p0/pi + (alpha-1)*np.log(p0/pi)) 

235 - 0.5*np.log(p) (weights_smooth*T) np.log(p)` 

236 

237 Note that frank uses log(p) in the inference hence the probability 

238 is normalized such that ..math:`\int P(p) d\log p = 1`. 

239 

240 Parameters 

241 ---------- 

242 p : array, size = N 

243 Power spectrum coefficients. 

244  

245 Returns 

246 ------- 

247 log[P(p)] : float 

248 Log prior probability 

249 

250 Notes 

251 ----- 

252 Computed up to a normalizing constant that depends on alpha and p0 

253 """ 

254 

255 # Add the power spectrum prior term 

256 xi = self._p_0 / p 

257 like = - np.sum(xi + (self._alpha-1) * np.log(xi)) 

258 

259 # Extra term due to spectral smoothness 

260 tau = np.log(p) 

261 like -= 0.5 * np.dot(tau, self._Tij.dot(tau)) 

262 

263 return like