Coverage for frank/radial_fitters.py: 80%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

249 statements  

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# 

20"""This module contains methods for fitting a radial brightness profile to a set 

21 of deprojected visibities. 

22""" 

23import abc 

24from collections import defaultdict 

25import logging 

26import numpy as np 

27 

28from frank.constants import rad_to_arcsec 

29from frank.filter import CriticalFilter 

30from frank.hankel import DiscreteHankelTransform 

31from frank.statistical_models import ( 

32 GaussianModel, LogNormalMAPModel, VisibilityMapping 

33) 

34 

35class FrankRadialFit(metaclass=abc.ABCMeta): 

36 """ 

37 Base class for results of frank fits. 

38 

39 Parameters 

40 ---------- 

41 vis_map : VisibilityMapping object 

42 Mapping between image and visibility plane.  

43 info: dict 

44 Dictionary containing useful quantities for reproducing a fit 

45 (such as the hyperparameters used) 

46 geometry: SourceGeometry object, optional 

47 Geometry used to correct the visibilities for the source 

48 inclination. If not provided, the geometry determined during the 

49 fit will be used. 

50 """ 

51 def __init__(self, vis_map, info, geometry): 

52 self._vis_map = vis_map 

53 self._geometry = geometry 

54 self._info = info 

55 

56 def predict(self, u, v, I=None, geometry=None): 

57 r""" 

58 Predict the visibilities in the sky-plane 

59 

60 Parameters 

61 ---------- 

62 u, v : array, unit = :math:`\lambda` 

63 uv-points to predict the visibilities at 

64 I : array, optional, unit = Jy 

65 Intensity points to predict the vibilities of. If not specified, 

66 the mean will be used. The intensity should be specified at the 

67 collocation points, I[k] = :math:`I(r_k)` 

68 geometry: SourceGeometry object, optional 

69 Geometry used to correct the visibilities for the source 

70 inclination. If not provided, the geometry determined during the 

71 fit will be used 

72 

73 Returns 

74 ------- 

75 V(u,v) : array, unit = Jy 

76 Predicted visibilties of a source with a radial flux distribution 

77 given by :math:`I` and the position angle, inclination and phase 

78 centre determined by the geometry object 

79 """ 

80 if geometry is None: 

81 geometry = self._geometry 

82 

83 if I is None: 

84 I = self.I 

85 

86 if geometry is not None: 

87 u, v, wz = geometry.deproject(u, v, use3D=True) 

88 else: 

89 wz = np.zeros_like(u) 

90 

91 q = np.hypot(u, v) 

92 V = self._vis_map.predict_visibilities(I, q, wz, geometry=geometry) 

93 

94 # Undo phase centering 

95 if geometry is not None: 

96 _, _, V = geometry.undo_correction(u, v, V) 

97 

98 return V 

99 

100 def predict_deprojected(self, q=None, I=None, geometry=None, 

101 block_size=10**5, assume_optically_thick=True): 

102 r""" 

103 Predict the visibilities in the deprojected-plane 

104 

105 Parameters 

106 ---------- 

107 q : array, default = self.q, unit = :math:`\lambda` 

108 1D uv-points to predict the visibilities at 

109 I : array, optional, unit = Jy / sr 

110 Intensity points to predict the vibilities of. If not specified, 

111 the mean will be used. The intensity should be specified at the 

112 collocation points, I[k] = I(r_k) 

113 geometry: SourceGeometry object, optional 

114 Geometry used to correct the visibilities for the source 

115 inclination. If not provided, the geometry determined during the 

116 fit will be used 

117 block_size : int, default = 10**5 

118 Maximum matrix size used in the visibility calculation 

119 assume_optically_thick : bool, default = True 

120 Whether to correct the visibility amplitudes for the source 

121 inclination 

122 

123 Returns 

124 ------- 

125 V(q) : array, unit = Jy 

126 Predicted visibilties of a source with a radial flux distribution 

127 given by :math:`I`. The amplitude of the visibilities are reduced 

128 according to the inclination of the source, for consistency with 

129 `uvplot` 

130 

131 Notes 

132 ----- 

133 The visibility amplitudes are still reduced due to the projection, 

134 for consistency with `uvplot` 

135 """ 

136 if geometry is None: 

137 geometry = self._geometry 

138 

139 if I is None: 

140 I = self.I 

141 

142 V = self._vis_map.predict_visibilities(I, q, q*0, geometry=geometry) 

143 

144 return V 

145 

146 def interpolate_brightness(self, Rpts, I=None): 

147 r""" 

148 Interpolate the brightness profile to the requested set of points. 

149 

150 The interpolation is done using the Fourier-Bessel series.  

151  

152 Parameters 

153 ---------- 

154 Rpts : array, unit = arcsec 

155 The points to interpolate the brightness to. 

156 I : array, optional, unit = Jy / Sr 

157 Intensity points to interpolate. If not specified, the MAP/mean 

158 will be used. The intensity should be specified at the collocation 

159 points, I[k] = I(r_k). 

160 

161 Returns 

162 ------- 

163 I(Rpts) : array, unit = Jy / Sr 

164 Brightness at the radial points. 

165 

166 Notes 

167 ----- 

168 The resulting brightness will be consistent with higher-resolution fits 

169 as long as the original fit has sufficient resolution. By sufficient 

170 resolution we simply mean that the missing terms in the Fourier-Bessel 

171 series are negligible, which will typically be the case if the 

172 brightness was obtained from a frank fit with 100 points or more. 

173 """ 

174 Rpts = np.array(Rpts) 

175 

176 return self._vis_map.interpolate(I, Rpts, space='Real') 

177 

178 

179 @abc.abstractproperty 

180 def MAP(self): 

181 pass 

182 

183 @property 

184 def I(self): 

185 return self.MAP 

186 

187 @property 

188 def r(self): 

189 """Radius points, unit = arcsec""" 

190 return self._vis_map.r 

191 

192 @property 

193 def Rmax(self): 

194 """Maximum radius, unit = arcsec""" 

195 return self._vis_map.Rmax 

196 @property 

197 def q(self): 

198 r"""Frequency points, unit = :math:`\lambda`""" 

199 return self._vis_map.q 

200 

201 @property 

202 def Qmax(self): 

203 r"""Maximum frequency, unit = :math:`\lambda`""" 

204 return self._vis_map.Qmax 

205 

206 @property 

207 def size(self): 

208 """Number of points in reconstruction""" 

209 return self._vis_map.size 

210 

211 @property 

212 def geometry(self): 

213 """SourceGeometry object""" 

214 return self._geometry 

215 

216 @property 

217 def info(self): 

218 """Fit quantities for reference""" 

219 return self._info 

220 

221 

222class FrankGaussianFit(FrankRadialFit): 

223 """ 

224 Result of a frank fit with a Gaussian brightness model. 

225 

226 Parameters 

227 ---------- 

228 DHT : DiscreteHankelTransform 

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

230 :math:`S(p)` 

231 fit : GaussianModel object 

232 Result of fitting with MAP power spectrum. 

233 info: dict, optional 

234 Dictionary containing useful quantities for reproducing a fit 

235 (such as the hyperparameters used) 

236 geometry: SourceGeometry object, optional 

237 Geometry used to correct the visibilities for the source 

238 inclination. If not provided, the geometry determined during the 

239 fit will be used. 

240 """ 

241 

242 def __init__(self, DHT, fit, info={}, geometry=None): 

243 FrankRadialFit.__init__(self, DHT, info, geometry) 

244 self._fit = fit 

245 

246 def draw(self, N): 

247 """Compute N draws from the posterior""" 

248 return np.random.multivariate_normal(self.mean, self.covariance, N) 

249 

250 def log_likelihood(self, I=None): 

251 r""" 

252 Compute one of two types of likelihood. 

253 

254 If :math:`I` is provided, this computes 

255 

256 .. math: 

257 \log[P(I,V|S)]. 

258 

259 Otherwise the marginalized likelihood is computed, 

260 

261 .. math: 

262 \log[P(V|S)]. 

263 

264 

265 Parameters 

266 ---------- 

267 I : array, size = N, optional, unit = Jy / sr 

268 Intensity :math:`I(r)` to compute the likelihood of 

269 

270 Returns 

271 ------- 

272 log_P : float 

273 Log likelihood, :math:`\log[P(I,V|p)]` or :math:`\log[P(V|p)]` 

274 

275 Notes 

276 ----- 

277 1. The prior probability P(S) is not included. 

278 2. The likelihoods take the form: 

279 

280 .. math:: 

281 \log[P(I,V|p)] = \frac{1}{2} j^T I - \frac{1}{2} I^T D^{-1} I 

282 - \frac{1}{2} \log[\det(2 \pi S)] + H_0 

283 

284 and 

285 

286 .. math:: 

287 \log[P(V|p)] = \frac{1}{2} j^T D j 

288 + \frac{1}{2} \log[\det(D)/\det(S)] + H_0 

289 

290 where 

291 

292 .. math:: 

293 H_0 = -\frac{1}{2} V^T w V + \frac{1}{2} \sum \log(w /2 \pi) 

294 

295 is the noise likelihood. 

296 """ 

297 return self._fit.log_likelihood(I) 

298 

299 

300 def solve_non_negative(self): 

301 """Compute the best fit solution with non-negative intensities""" 

302 return self._fit.solve_non_negative() 

303 

304 @property 

305 def mean(self): 

306 """Posterior mean, unit = Jy / sr""" 

307 return self._fit.mean 

308 

309 @property 

310 def MAP(self): 

311 """Posterior maximum, unit = Jy / sr""" 

312 return self.mean 

313 

314 @property 

315 def covariance(self): 

316 """Posterior covariance, unit = (Jy / sr)**2""" 

317 return self._fit.covariance 

318 

319 @property 

320 def power_spectrum(self): 

321 """Power spectrum coefficients""" 

322 return self._fit.power_spectrum 

323 

324 

325 

326class FrankLogNormalFit(FrankRadialFit): 

327 """ 

328 Result of a frank fit with a Gaussian brightness model. 

329 

330 Parameters 

331 ---------- 

332 DHT : DiscreteHankelTransform 

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

334 :math:`S(p)` 

335 fit : LogNormalMAPModel object 

336 Result of fitting with MAP power spectrum. 

337 info: dict, optional 

338 Dictionary containing useful quantities for reproducing a fit 

339 (such as the hyperparameters used) 

340 geometry: SourceGeometry object, optional 

341 Geometry used to correct the visibilities for the source 

342 inclination. If not provided, the geometry determined during the 

343 fit will be used. 

344 """ 

345 

346 def __init__(self, DHT, fit, info={}, geometry=None): 

347 FrankRadialFit.__init__(self, DHT, info, geometry) 

348 self._fit = fit 

349 

350 def log_likelihood(self, I=None): 

351 r""" 

352 Compute the likelihood, 

353 

354 .. math: 

355 \log[P(I,V|S)]. 

356 

357 Parameters 

358 ---------- 

359 I : array, size = N, optional 

360 Intensity :math:`I(r)=exp(s0*s)` to compute the likelihood of 

361 

362 Returns 

363 ------- 

364 log_P : float 

365 Log likelihood, :math:`\log[P(I,V|p)]` 

366 

367 Notes 

368 ----- 

369 1. The prior probability P(S) is not included. 

370 2. The likelihood takes the form: 

371 

372 .. math:: 

373 \log[P(I,V|p)] = j^T I - \frac{1}{2} I^T D^{-1} I 

374 - \frac{1}{2} \log[\det(2 \pi S)] + H_0 

375 

376 where 

377 

378 .. math:: 

379 H_0 = -\frac{1}{2} V^T w V + \frac{1}{2} \sum \log(w /2 \pi) 

380 

381 is the noise likelihood. 

382 """ 

383 

384 if I is not None: 

385 return self._fit.log_likelihood(np.log(I)) 

386 else: 

387 return self._fit.log_likelihood() 

388 

389 @property 

390 def MAP(self): 

391 """Posterior maximum, unit = Jy / sr""" 

392 return np.exp((self._fit.MAP + self._fit.s_0) * self._fit.scale) 

393 

394 @property 

395 def covariance(self): 

396 """Posterior covariance, unit = log[(Jy / sr)**2]""" 

397 return self._fit.covariance 

398 

399 @property 

400 def power_spectrum(self): 

401 """Power spectrum coefficients""" 

402 return self._fit.power_spectrum 

403 

404 

405class FourierBesselFitter(object): 

406 """ 

407 Fourier-Bessel series model for fitting visibilities 

408 

409 Parameters 

410 ---------- 

411 Rmax : float, unit = arcsec 

412 Radius of support for the functions to transform, i.e., 

413 f(r) = 0 for R >= Rmax 

414 N : int 

415 Number of collocation points 

416 geometry : SourceGeometry object 

417 Geometry used to deproject the visibilities before fitting 

418 nu : int, default = 0 

419 Order of the discrete Hankel transform (DHT) 

420 block_data : bool, default = True 

421 Large temporary matrices are needed to set up the data. If block_data 

422 is True, we avoid this, limiting the memory requirement to block_size 

423 elements. 

424 assume_optically_thick : bool, default = True 

425 Whether to correct the visibility amplitudes by a factor of 

426 1 / cos(inclination); see frank.geometry.rescale_total_flux 

427 scale_height : function R --> H, optional 

428 Specifies the vertical thickness of disc as a function of radius. Both 

429 R and H should be in arcsec. Assumes a Gaussian vertical structure.  

430 Only works with assume_optically_thick=False 

431 block_size : int, default = 10**5 

432 Size of the matrices if blocking is used 

433 verbose : bool, default = False 

434 Whether to print notification messages 

435 """ 

436 

437 def __init__(self, Rmax, N, geometry, nu=0, block_data=True, 

438 assume_optically_thick=True, scale_height=None, 

439 block_size=10 ** 5, verbose=True): 

440 

441 Rmax /= rad_to_arcsec 

442 

443 self._geometry = geometry 

444 

445 self._DHT = DiscreteHankelTransform(Rmax, N, nu) 

446 

447 if assume_optically_thick: 

448 if scale_height is not None: 

449 raise ValueError("Optically thick models must have zero " 

450 "scale-height") 

451 model = 'opt_thick' 

452 elif scale_height is not None: 

453 model = 'debris' 

454 else: 

455 model = 'opt_thin' 

456 

457 self._vis_map = VisibilityMapping(self._DHT, geometry, 

458 model, scale_height=scale_height, 

459 block_data=block_data, block_size=block_size, 

460 check_qbounds=False, verbose=verbose) 

461 

462 self._info = {'Rmax' : self._DHT.Rmax * rad_to_arcsec, 

463 'N' : self._DHT.size 

464 } 

465 

466 self._verbose = verbose 

467 

468 def preprocess_visibilities(self, u, v, V, weights=1): 

469 r"""Prepare the visibilities for fitting.  

470  

471 This step will be done by the automatically fit method, but it can be  

472 expensive. This method is provided to enable the pre-processing to be 

473 done once when conducting multiple fits with different hyperprior  

474 parameters.  

475 

476 Parameters 

477 ---------- 

478 u,v : 1D array, unit = :math:`\lambda` 

479 uv-points of the visibilies 

480 V : 1D array, unit = Jy 

481 Visibility amplitudes at q 

482 weights : 1D array, optional, unit = Jy^-2 

483 Weights of the visibilities, weight = 1 / sigma^2, where sigma is 

484 the standard deviation 

485  

486 Returns 

487 ------- 

488 processed_visibilities : dict 

489 Processed visibilities needed in subsequent steps of the fit 

490 

491 Notes 

492 ----- 

493 Re-using the processed visibilities is only valid with certain parameter 

494 changes. For example N, Rmax, geometry, nu, assume_optically_thick, and 

495 scale_height cannot be changed. This will be checked before fits are  

496 conducted. 

497 """ 

498 return self._vis_map.map_visibilities(u, v, V, weights) 

499 

500 def _build_matrices(self, mapping): 

501 r""" 

502 Compute the matrices M and j from the visibility data. 

503 

504 Also compute 

505 .. math: 

506 `H0 = 0.5*\log[det(weights/(2*np.pi))] 

507 - 0.5*np.sum(V * weights * V):math:` 

508 """ 

509 self._vis_map.check_hash(mapping['hash']) 

510 

511 self._M = mapping['M'] 

512 self._j = mapping['j'] 

513 

514 self._H0 = mapping['null_likelihood'] 

515 

516 def fit_method(self): 

517 """Name of the fit method""" 

518 return type(self).__name__ 

519 

520 def fit_preprocessed(self, preproc_vis): 

521 r""" 

522 Fit the pre-processed visibilties. The last step in the fitting  

523 procedure. 

524 

525 Parameters 

526 ---------- 

527 preproc_vis : pre-processed visibilities 

528 Visibilities to fit that have been processed by 

529 `self.preprocess_visibilities`. 

530 

531 Returns 

532 ------- 

533 sol : FrankRadialFit 

534 Least-squares Fourier-Bessel series fit 

535 """ 

536 if self._verbose: 

537 logging.info(' Fitting pre-processed visibilities for brightness' 

538 ' profile using {}'.format(self.fit_method())) 

539 

540 self._build_matrices(preproc_vis) 

541 

542 return self._fit() 

543 

544 def fit(self, u, v, V, weights=1): 

545 r""" 

546 Fit the visibilties 

547 

548 Parameters 

549 ---------- 

550 u,v : 1D array, unit = :math:`\lambda` 

551 uv-points of the visibilies 

552 V : 1D array, unit = Jy 

553 Visibility amplitudes at q 

554 weights : 1D array, optional, unit = J^-2 

555 Weights of the visibilities, weight = 1 / sigma^2, where sigma is 

556 the standard deviation 

557 

558 Returns 

559 ------- 

560 sol : FrankRadialFit 

561 Least-squares Fourier-Bessel series fit 

562 """ 

563 if self._verbose: 

564 logging.info(' Fitting for brightness profile using' 

565 ' {}'.format(self.fit_method())) 

566 

567 self._geometry.fit(u, v, V, weights) 

568 

569 mapping = self.preprocess_visibilities(u, v, V, weights) 

570 self._build_matrices(mapping) 

571 

572 return self._fit() 

573 

574 def _fit(self): 

575 """Fit step. Computes the best fit given the pre-processed data""" 

576 fit = GaussianModel(self._DHT, self._M, self._j, 

577 noise_likelihood=self._H0) 

578 

579 self._sol = FrankGaussianFit(self._vis_map, fit, self._info, 

580 geometry=self._geometry.clone()) 

581 

582 return self._sol 

583 

584 

585 @property 

586 def r(self): 

587 """Radius points, unit = arcsec""" 

588 return self._DHT.r * rad_to_arcsec 

589 

590 @property 

591 def Rmax(self): 

592 """Maximum radius, unit = arcsec""" 

593 return self._DHT.Rmax * rad_to_arcsec 

594 

595 @property 

596 def q(self): 

597 r"""Frequency points, unit = :math:`\lambda`""" 

598 return self._DHT.q 

599 

600 @property 

601 def Qmax(self): 

602 r"""Maximum frequency, unit = :math:`\lambda`""" 

603 return self._DHT.Qmax 

604 

605 @property 

606 def size(self): 

607 """Number of points in reconstruction""" 

608 return self._DHT.size 

609 

610 @property 

611 def geometry(self): 

612 """Geometry object""" 

613 return self._geometry 

614 

615 

616class FrankFitter(FourierBesselFitter): 

617 """ 

618 Fit a Gaussian process model using the Discrete Hankel Transform of 

619 Baddour & Chouinard (2015). 

620 

621 The GP model is based upon Oppermann et al. (2013), which use a maximum 

622 a posteriori estimate for the power spectrum as the GP prior for the 

623 real-space coefficients 

624 

625 Parameters 

626 ---------- 

627 Rmax : float, unit = arcsec 

628 Radius of support for the functions to transform, i.e., f(r) = 0 for 

629 R >= Rmax. 

630 N : int 

631 Number of collaction points 

632 geometry : SourceGeometry object 

633 Geometry used to deproject the visibilities before fitting 

634 nu : int, default = 0 

635 Order of the discrete Hankel transform, given by J_nu(r) 

636 block_data : bool, default = True 

637 Large temporary matrices are needed to set up the data. If block_data 

638 is True, we avoid this, limiting the memory requirement to block_size 

639 elements 

640 block_size : int, default = 10**5 

641 Size of the matrices if blocking is used 

642 alpha : float >= 1, default = 1.05 

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

644 coefficients 

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

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

647 coefficients. If not provided p_0 = 1e-15 (method="Normal") or  

648 1e-35 (method="LogNormal") will be used. 

649 weights_smooth : float >= 0, default = 1e-4 

650 Spectral smoothness prior parameter. Zero is no smoothness prior 

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

652 Tolerence for convergence of the power spectrum iteration 

653 method : string, default="Normal" 

654 Model used for the brightness reconstrution. This must be one of 

655 "Normal" of "LogNormal". 

656 I_scale : float, default = 1e5, unit= Jy/Sr 

657 Brightness scale. Only used in the LogNormal model. Note the 

658 LogNormal model produces I(Rmax) = I_scale. 

659 max_iter: int, default = 2000 

660 Maximum number of fit iterations 

661 check_qbounds: bool, default = True 

662 Whether to check if the first (last) collocation point is smaller 

663 (larger) than the shortest (longest) deprojected baseline in the dataset 

664 store_iteration_diagnostics: bool, default = False 

665 Whether to store the power spectrum parameters and brightness profile 

666 for each fit iteration 

667 assume_optically_thick : bool, default = True 

668 Whether to correct the visibility amplitudes by a factor of 

669 1 / cos(inclination); see frank.geometry.rescale_total_flux 

670 scale_height : function R --> H, optional 

671 Specifies the vertical thickness of disc as a function of radius. Both 

672 R and H should be in arcsec. Assumes a Gaussian vertical structure.  

673 Only works with assume_optically_thick=False 

674 verbose: bool 

675 Whether to print notification messages 

676 convergence_failure: string, default = 'raise' 

677 Decide what to do when the frank model does not converge within max_iter. 

678 Should be one of: 

679 'raise' : raise an error 

680 'warn' : print a warning message and continue 

681 'ignore' : Ignore the error. 

682 

683 References 

684 ---------- 

685 Baddour & Chouinard (2015) 

686 DOI: https://doi.org/10.1364/JOSAA.32.000611 

687 Oppermann et al. (2013) 

688 DOI: https://doi.org/10.1103/PhysRevE.87.032136 

689 """ 

690 

691 def __init__(self, Rmax, N, geometry, nu=0, block_data=True, 

692 block_size=10 ** 5, alpha=1.05, p_0=None, weights_smooth=1e-4, 

693 tol=1e-3, method='Normal', I_scale=1e5, max_iter=2000, check_qbounds=True, 

694 store_iteration_diagnostics=False, assume_optically_thick=True, 

695 scale_height=None, verbose=True, convergence_failure='raise'): 

696 

697 if method not in {'Normal', 'LogNormal'}: 

698 raise ValueError('FrankFitter supports following mehods:\n\t' 

699 '{ "Normal", "LogNormal"}"') 

700 self._method = method 

701 

702 super(FrankFitter, self).__init__(Rmax, N, geometry, nu, block_data, 

703 assume_optically_thick, scale_height, 

704 block_size, verbose) 

705 

706 # Reinstate the bounds check: FourierBesselFitter does not check bounds 

707 self._vis_map.check_qbounds=check_qbounds 

708 

709 if p_0 is None: 

710 if method == 'Normal': 

711 p_0 = 1e-15 

712 else: 

713 p_0 = 1e-35 

714 

715 self._s_scale = np.log(I_scale) 

716 

717 self._filter = CriticalFilter( 

718 self._DHT, alpha, p_0, weights_smooth, tol 

719 ) 

720 

721 self._max_iter = max_iter 

722 

723 self._store_iteration_diagnostics = store_iteration_diagnostics 

724 

725 self._info.update({'alpha' : alpha, 'wsmooth' : weights_smooth, 

726 'p0' : p_0, 'method': method}) 

727 

728 if convergence_failure not in {'raise', 'warn', 'ignore'}: 

729 raise ValueError("convergence_failure must be one of 'raise'," 

730 f"'warn', or 'ignore', nor {convergence_failure}") 

731 self._convergence_failure = convergence_failure 

732 

733 def fit_method(self): 

734 """Name of the fit method""" 

735 return '{}: {} method'.format(type(self).__name__, self._method) 

736 

737 def _fit(self): 

738 """Fit step. Computes the best fit given the pre-processed data""" 

739 

740 if self._store_iteration_diagnostics: 

741 self._iteration_diagnostics = defaultdict(list) 

742 

743 # Inital guess for power spectrum 

744 pI = np.ones([self.size]) 

745 

746 # Do an extra iteration based on a power-law guess 

747 fit = self._perform_fit(pI, guess=np.ones_like(pI), fit_method='Normal') 

748 

749 pI = np.max(self._DHT.transform(fit.MAP)**2) 

750 pI *= (self.q / self.q[0])**-2 

751 

752 fit = self._perform_fit(pI, fit_method='Normal') 

753 

754 # Now that we've got a reasonable initial brightness, setup the 

755 # log-normal power spectrum estimate 

756 if self._method == 'LogNormal': 

757 s = np.log(np.maximum(fit.MAP, 1e-3 * fit.MAP.max())) 

758 s -= self._s_scale 

759 

760 pI = np.max(self._DHT.transform(s)**2) 

761 pI *= (self.q / self.q[0])**-4 

762 

763 fit = self._perform_fit(pI, guess=s) 

764 

765 

766 

767 count = 0 

768 pi_old = 0 

769 while (not self._filter.check_convergence(pI, pi_old) and 

770 count <= self._max_iter): 

771 

772 if self._verbose and logging.getLogger().isEnabledFor(logging.INFO): 

773 print('\r {} iteration {}'.format(type(self).__name__, count), 

774 end='', flush=True) 

775 

776 pi_old = pI.copy() 

777 pI = self._filter.update_power_spectrum(fit) 

778 

779 fit = self._perform_fit(pI, guess=fit.MAP) 

780 

781 if self._store_iteration_diagnostics: 

782 self._iteration_diagnostics['power_spectrum'].append(pI) 

783 self._iteration_diagnostics['MAP'].append(fit.MAP) 

784 

785 count += 1 

786 

787 # Check / report convergence 

788 if count < self._max_iter: 

789 if self._verbose and logging.getLogger().isEnabledFor(logging.INFO): 

790 print() 

791 logging.info(' Convergence criterion met at iteration' 

792 ' {}'.format(count-1)) 

793 else: 

794 if self._verbose and logging.getLogger().isEnabledFor(logging.INFO): 

795 print() 

796 logging.info(' Convergence criterion not met; fit stopped at' 

797 ' max_iter specified in your parameter file,' 

798 ' {}'.format(self._max_iter)) 

799 

800 msg = f'Convergence not met within {self._max_iter} ' 

801 msg += f'iterations.\nTry increasing max_iter, or ' 

802 msg += 'try increasing alpha since convergence can ' 

803 msg += 'be very slow for alpha close to 1.' 

804 if self._convergence_failure == 'raise': 

805 msg += '\nAlternatively set convergence_failure to' 

806 msg += "'warn' or 'ignore' to continue despite the" 

807 msg += 'failure.' 

808 raise RuntimeError(msg) 

809 elif self._convergence_failure == 'warn': 

810 if logging.getLogger().isEnabledFor(logging.INFO): 

811 logging.info(msg) 

812 else: 

813 print(msg) 

814 elif self._convergence_failure == 'ignore': 

815 pass 

816 

817 if self._store_iteration_diagnostics: 

818 self._iteration_diagnostics['num_iterations'] = count 

819 

820 # Save the best fit 

821 if self._method == "Normal": 

822 self._sol = FrankGaussianFit(self._vis_map, fit, self._info, 

823 geometry=self._geometry.clone()) 

824 else: 

825 self._sol = FrankLogNormalFit(self._vis_map, fit, self._info, 

826 geometry=self._geometry.clone()) 

827 

828 # Compute the power spectrum covariance at the maximum 

829 self._ps = pI 

830 self._ps_cov = None 

831 

832 return self._sol 

833 

834 def draw_powerspectrum(self, Ndraw=1): 

835 """ 

836 Draw N sets of power-spectrum parameters. 

837 

838 The draw is taken from the Laplace-approximated (Gaussian) posterior 

839 distribution for p, 

840 :math:`P(p) ~ G(p - p_MAP, p_cov)` 

841 

842 Parameters 

843 ---------- 

844 Ndraw : int, default = 1 

845 Number of draws 

846 

847 Returns 

848 ------- 

849 p : array, size = (N, Ndraw) 

850 Power spectrum draws 

851 """ 

852 

853 log_p = np.random.multivariate_normal(np.log(self._ps), 

854 self._ps_cov, Ndraw) 

855 

856 return np.exp(log_p) 

857 

858 def _perform_fit(self, p, guess=None, fit_method=None): 

859 """ 

860 Find the posterior mean and covariance given p 

861 

862 Parameters 

863 ---------- 

864 p : array, size = N 

865 Power spectrum parameters 

866 guess : array, size = N, option 

867 Initial guess for brightness used in the fit. 

868 fit_method : string, optional. 

869 One of {"Normal", "LogNormal"}. Brightness profile 

870 method used in fit. 

871 

872 Returns 

873 ------- 

874 sol : FrankRadialFit 

875 Posterior solution object for P(I|V,p) 

876 """ 

877 if fit_method is None: 

878 fit_method = self._method 

879 

880 if fit_method == 'Normal': 

881 return GaussianModel(self._DHT, self._M, self._j, p, 

882 guess=guess, 

883 noise_likelihood=self._H0) 

884 elif fit_method == 'LogNormal': 

885 return LogNormalMAPModel(self._DHT, self._M, self._j, p, 

886 guess=guess, s0=self._s_scale, 

887 noise_likelihood=self._H0) 

888 else: 

889 raise ValueError('fit_method must be one of the following:\n\t' 

890 '{"Normal", "LogNormal"}') 

891 

892 def log_prior(self, p=None): 

893 """ 

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

895 

896 .. math: 

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

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

899 

900 Parameters 

901 ---------- 

902 p : array, size = N, optional 

903 Power spectrum coefficients. If not provided, the MAP values are 

904 used 

905 

906 Returns 

907 ------- 

908 log[P(p)] : float 

909 Log prior probability 

910 

911 Notes 

912 ----- 

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

914 """ 

915 

916 if p is None: 

917 p = self._ps 

918 

919 return self._filter.log_prior(p) 

920 

921 

922 def log_likelihood(self, sol=None): 

923 r""" 

924 Compute the log likelihood :math:`log[P(p, V)]`, 

925 

926 .. math: 

927 \log[P(p, V)] ~ \log[P(V|p)] + \log[P(p)] 

928 

929 

930 Parameters 

931 ---------- 

932 sol : FrankRadialFit object, optional 

933 Posterior solution given a set power spectrum parameters, :math:`p`. 

934 If not provided, the MAP solution will be used 

935 

936 Returns 

937 ------- 

938 log[P(p, V)] : float 

939 Log prior probability 

940 

941 Notes 

942 ----- 

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

944 """ 

945 

946 if sol is None: 

947 sol = self.MAP_solution 

948 

949 return self.log_prior(sol.power_spectrum) + sol.log_likelihood() 

950 

951 def log_evidence_laplace(self): 

952 r"""Compute the evidence, p(V), for the best fit model. 

953 

954 Uses the Laplace approximation, I.e. we model the posterior p(p, V) as 

955 a Gaussian in \tau = log(p) centered on p_MAP with the covariance  

956 determined by the curvate of log P(p,V) at p_MAP: 

957 P_Laplace(\tau) ~ N(\tau-\tau_MAP, \Sigma) * p(V) 

958 where p(V) is the approximated evidence. Here  

959 \Sigma_i,j = - d^2 \log(P) / d\tau_i d\tau_j 

960 and p(V) is determined such that P_Lapace(\tau_MAP) = P(p_MAP, V). 

961 """ 

962 Sigma_inv = self._filter.covariance_MAP(self._sol, ret_inv=True) 

963 

964 sign, logdet = np.linalg.slogdet(Sigma_inv/(2*np.pi)) 

965 laplace = self.log_likelihood() - 0.5*logdet 

966 

967 return laplace 

968 

969 @property 

970 def MAP_solution(self): 

971 """Reconstruction for the maximum a posteriori power spectrum""" 

972 return self._sol 

973 

974 @property 

975 def MAP_spectrum(self): 

976 """Maximum a posteriori power spectrum""" 

977 return self._ps 

978 

979 @property 

980 def MAP_spectrum_covariance(self): 

981 """Covariance matrix of the maximum a posteriori power spectrum""" 

982 if self._ps_cov is None: 

983 self._ps_cov = self._filter.covariance_MAP(self._sol) 

984 

985 return self._ps_cov 

986 

987 @property 

988 def iteration_diagnostics(self): 

989 """Dict containing power spectrum parameters and posterior mean 

990 brightness profile at each fit iteration, and number of iterations""" 

991 return self._iteration_diagnostics