Coverage for frank/tests.py: 99%

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# the Free Software Foundation, either version 3 of the License, or 

9# (at your option) any later version. 

10# 

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

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

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

14# GNU General Public License for more details. 

15# 

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

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

18# 

19"""This module runs tests to confirm Frankenstein is working correctly.""" 

20 

21import numpy as np 

22import os 

23import json 

24 

25from frank.constants import rad_to_arcsec 

26from frank.hankel import DiscreteHankelTransform 

27from frank.radial_fitters import FourierBesselFitter, FrankFitter 

28from frank.debris_fitters import FrankDebrisFitter 

29from frank.geometry import ( 

30 FixedGeometry, FitGeometryGaussian, FitGeometryFourierBessel 

31) 

32from frank import utilities 

33from frank.io import load_uvtable, save_uvtable, load_sol, save_fit 

34from frank.statistical_models import VisibilityMapping 

35from frank import fit 

36 

37def test_hankel_gauss(): 

38 """Check the Hankel transform""" 

39 

40 def gauss_real_space(r): 

41 x = r 

42 return np.exp(-0.5 * x * x) 

43 

44 def gauss_vis_space(q): 

45 qs = (2 * np.pi) * q 

46 return np.exp(-0.5 * qs * qs) * (2 * np.pi) 

47 

48 DHT = DiscreteHankelTransform(5.0, 100) 

49 

50 Ir = gauss_real_space(DHT.r) 

51 Iq = gauss_vis_space(DHT.q) 

52 

53 # Test at the DHT points. 

54 # Use a large error estimate because the DHT is approximate 

55 np.testing.assert_allclose(Iq, DHT.transform(Ir, direction='forward'), 

56 atol=1e-5, rtol=0, err_msg="Forward DHT") 

57 

58 np.testing.assert_allclose(Ir, DHT.transform(Iq, direction='backward'), 

59 atol=1e-5, rtol=0, err_msg="Inverse DHT") 

60 

61 # Test at generic points. 

62 # Larger error needed 

63 q = np.linspace(0.0, 1.0, 25) 

64 np.testing.assert_allclose(gauss_vis_space(q), 

65 DHT.transform(Ir, q=q, direction='forward'), 

66 atol=1e-4, rtol=0, err_msg="Generic Forward DHT") 

67 

68 r = np.linspace(0, 5.0, 25) 

69 np.testing.assert_allclose(gauss_real_space(r), 

70 DHT.transform(Iq, q=r, direction='backward'), 

71 atol=1e-4, rtol=0, err_msg="Generic Inverse DHT") 

72 

73 # Check the coefficients matrix works 

74 Hf = DHT.coefficients(direction='forward') 

75 Hb = DHT.coefficients(direction='backward') 

76 

77 np.testing.assert_allclose(np.dot(Hf, Ir), 

78 DHT.transform(Ir, direction='forward'), 

79 rtol=1e-7, err_msg="Forward DHT Coeffs") 

80 np.testing.assert_allclose(np.dot(Hb, Iq), 

81 DHT.transform(Iq, direction='backward'), 

82 rtol=1e-7, err_msg="Inverse DHT Coeffs") 

83 

84 # Compare cached vs non-cached DHT points 

85 np.testing.assert_allclose(DHT.coefficients(q=DHT.q), 

86 DHT.coefficients(), 

87 atol=1e-12, rtol=0, 

88 err_msg="Cached forward DHT Coeffs" 

89 ) 

90 np.testing.assert_allclose(DHT.coefficients(q=DHT.r, direction='backward'), 

91 DHT.coefficients(direction='backward'), 

92 atol=1e-12, rtol=0, 

93 err_msg="Cached inverse DHT Coeffs" 

94 ) 

95 

96 

97def test_vis_mapping(): 

98 def gauss_real_space(r): 

99 x = r 

100 return np.exp(-0.5 * x * x) 

101 

102 def gauss_vis_space(q): 

103 qs = (2 * np.pi) * q / rad_to_arcsec 

104 return np.exp(-0.5 * qs * qs) * (2 * np.pi / rad_to_arcsec**2) 

105 

106 DHT = DiscreteHankelTransform(5.0/rad_to_arcsec, 100) 

107 geometry = FixedGeometry(60, 0, 0, 0) 

108 

109 VM = VisibilityMapping(DHT, geometry) 

110 

111 Ir = gauss_real_space(VM.r) 

112 Iq = gauss_vis_space(VM.q) 

113 

114 # Test at the DHT points. 

115 # Use a large error estimate because the DHT is approximate 

116 np.testing.assert_allclose(Iq, 2*VM.predict_visibilities(Ir, VM.q), 

117 atol=1e-5, rtol=0, err_msg="Forward DHT with VisibilityMapping") 

118 

119 np.testing.assert_allclose(Ir, 0.5*VM.invert_visibilities(Iq, VM.r), 

120 atol=1e-5, rtol=0, err_msg="Inverse DHT with VisibilityMapping") 

121 

122 # Test generic_dht, which uses these functions 

123 _, Iq_dht = utilities.generic_dht(VM.r, Ir, inc=60, Rmax=5, N=100) 

124 _, Ir_dht = utilities.generic_dht(VM.q, Iq, inc=60, Rmax=5, N=100, direction='backward') 

125 

126 np.testing.assert_allclose(Iq, 2*Iq_dht, 

127 atol=1e-5, rtol=0, err_msg="Forward DHT with generic_dht") 

128 

129 np.testing.assert_allclose(Ir, 0.5*Ir_dht, 

130 atol=1e-5, rtol=0, err_msg="Inverse DHT with generic_dht") 

131 

132 

133def test_import_data(): 

134 """Check the UVTable import function works for a .txt""" 

135 load_uvtable('docs/tutorials/test_datafile.txt') 

136 

137 

138def load_AS209(uv_cut=None): 

139 """Load data for subsequent tests""" 

140 uv_AS209_DSHARP = np.load('docs/tutorials/AS209_continuum.npz') 

141 geometry = FixedGeometry(dRA=-1.9e-3, dDec=2.5e-3, inc=34.97, 

142 PA=85.76) 

143 

144 if uv_cut is not None: 

145 u, v = [uv_AS209_DSHARP[x] for x in ['u', 'v']] 

146 

147 q = np.hypot(*geometry.deproject(u,v)) 

148 

149 keep = q < uv_cut 

150 

151 cut_data = {} 

152 for key in uv_AS209_DSHARP: 

153 if key not in { 'u', 'v', 'V', 'weights' }: 

154 continue 

155 cut_data[key] = uv_AS209_DSHARP[key][keep] 

156 

157 uv_AS209_DSHARP = cut_data 

158 

159 return uv_AS209_DSHARP, geometry 

160 

161def test_fit_geometry(): 

162 """Check the geometry fit on a subset of the AS209 data""" 

163 AS209, _ = load_AS209() 

164 u, v, vis, weights = [AS209[k][::100] for k in ['u', 'v', 'V', 'weights']] 

165 

166 inc_pa = [30.916257151011674, 85.46246241142246] 

167 phase_centre = [-0.6431627790617276e-3, -1.161768824369382e-3] 

168 

169 geom = FitGeometryGaussian() 

170 geom.fit(u, v, vis, weights) 

171 

172 np.testing.assert_allclose([geom.inc, geom.PA, 1e3 * geom.dRA, 

173 1e3 * geom.dDec], 

174 [30.916256948647096, 

175 85.46246845532691, 

176 -0.6434703241180601, -1.1623515516661052], 

177 err_msg="Gaussian geometry fit") 

178 

179 geom = FitGeometryGaussian(inc_pa=inc_pa) 

180 geom.fit(u, v, vis, weights) 

181 

182 np.testing.assert_allclose([geom.inc, geom.PA, 1e3 * geom.dRA, 

183 1e3 * geom.dDec], 

184 [30.916257151011674, 

185 85.46246241142246, 

186 -0.6432951224590862, -1.1619271783674576], 

187 err_msg="Gaussian geometry fit (provided inc_pa)") 

188 

189 geom = FitGeometryGaussian(phase_centre=phase_centre) 

190 geom.fit(u, v, vis, weights) 

191 

192 np.testing.assert_allclose([geom.inc, geom.PA, 1e3 * geom.dRA, 

193 1e3 * geom.dDec], 

194 [30.91601671340713, 

195 85.471787339838, 

196 -0.6431627790617276, -1.161768824369382], 

197 err_msg="Gaussian geometry fit (provided phase_centre)") 

198 

199 geom = FitGeometryGaussian(guess=[1.0, 1.0, 0.1, 0.1]) 

200 geom.fit(u, v, vis, weights) 

201 

202 np.testing.assert_allclose([geom.inc, geom.PA, 1e3 * geom.dRA, 

203 1e3 * geom.dDec], 

204 [30.91625882521282, 85.46246494153092, 

205 -0.6440453613101292, -1.1622414671266803], 

206 err_msg="Gaussian geometry fit (provided guess)") 

207 

208 geom = FitGeometryFourierBessel(1.6, 20) 

209 geom.fit(u, v, vis, weights) 

210 

211 np.testing.assert_allclose([geom.inc, geom.PA, 1e3 * geom.dRA, 

212 1e3 * geom.dDec], 

213 [33.81936473347169, 85.26142233735665, 

214 0.5611211784189547, -1.170097994325657], 

215 rtol=1e-5, 

216 err_msg="FourierBessel geometry fit") 

217 

218 

219 geom = FitGeometryFourierBessel(1.6, 20, inc_pa=inc_pa) 

220 geom.fit(u, v, vis, weights) 

221 

222 np.testing.assert_allclose([geom.inc, geom.PA, 1e3 * geom.dRA, 

223 1e3 * geom.dDec], 

224 [30.916257151011674, 85.46246241142246, 

225 1.4567005881700168, -1.658896248809076], 

226 rtol=1e-5, 

227 err_msg="FourierBessel geometry fit (provided inc_pa)") 

228 

229 

230 geom = FitGeometryFourierBessel(1.6, 20, phase_centre=phase_centre) 

231 geom.fit(u, v, vis, weights) 

232 

233 np.testing.assert_allclose([geom.inc, geom.PA, 1e3 * geom.dRA, 

234 1e3 * geom.dDec], 

235 [33.83672738960381, 85.2562498368987, 

236 -0.6431627790617276, -1.161768824369382], 

237 rtol=1e-5, 

238 err_msg="FourierBessel geometry fit (provided phase_centre)") 

239 

240 

241def test_fourier_bessel_fitter(): 

242 """Check FourierBesselFitter fitting routine with AS 209 dataset""" 

243 AS209, geometry = load_AS209() 

244 

245 u, v, vis, weights = [AS209[k] for k in ['u', 'v', 'V', 'weights']] 

246 

247 Rmax = 1.6 

248 

249 FB = FourierBesselFitter(Rmax, 20, geometry=geometry) 

250 

251 sol = FB.fit(u, v, vis, weights) 

252 expected = np.array([1.89446696e+10, 1.81772972e+10, 1.39622125e+10, 

253 1.20709653e+10, 

254 9.83716859e+09, 3.26308106e+09, 2.02453146e+08, 

255 4.73919867e+09, 

256 1.67911877e+09, 1.73161931e+08, 4.50233539e+08, 

257 3.57108466e+08, 

258 4.04216831e+09, 1.89085113e+09, 6.73819228e+08, 

259 5.50895976e+08, 

260 1.53683576e+08, 1.02413038e+08, 2.32589333e+07, 

261 3.33260713e+07 

262 ]) 

263 

264 np.testing.assert_allclose(sol.I, expected, 

265 err_msg="Testing FourierBessel Fit to AS 209") 

266 

267 

268def test_frank_fitter(): 

269 """Check FrankFitter fitting routine with AS 209 dataset""" 

270 AS209, geometry = load_AS209(uv_cut=1e6) 

271 

272 u, v, vis, weights = [AS209[k] for k in ['u', 'v', 'V', 'weights']] 

273 

274 Rmax = 1.6 

275 

276 FF = FrankFitter(Rmax, 20, geometry, alpha=1.05, weights_smooth=1e-2) 

277 

278 sol = FF.fit(u, v, vis, weights) 

279 expected = np.array([ 

280 2.007570578977623749e+10, 1.843513239499150085e+10, 

281 1.349013094584499741e+10, 1.272363099855433273e+10, 

282 1.034881472041586494e+10, 2.579145371666701317e+09, 

283 6.973651829234187603e+08, 4.127687040627769947e+09, 

284 2.502124003048851490e+09, -2.756950827897560596e+08, 

285 2.823720459944381118e+08, 8.705940396227477789e+08, 

286 3.257425109322027683e+09, 3.112003905406182289e+09, 

287 -5.145431577819123268e+08, 1.491165153547247887e+09, 

288 -5.190942564982021451e+08, 5.100334030941848755e+08, 

289 -1.922568182176418006e+08, 8.067782715878820419e+07, 

290 ]) 

291 

292 np.testing.assert_allclose(sol.I, expected, 

293 err_msg="Testing Frank Fit to AS 209") 

294 

295 

296def test_two_stage_fit(): 

297 """Check FrankFitter fitting routine with AS 209 dataset""" 

298 AS209, geometry = load_AS209(uv_cut=1e6) 

299 

300 u, v, vis, weights = [AS209[k] for k in ['u', 'v', 'V', 'weights']] 

301 

302 Rmax = 1.6 

303 

304 FF = FrankFitter(Rmax, 20, geometry, alpha=1.05, weights_smooth=1e-2) 

305 

306 # 1 step fit 

307 sol = FF.fit(u, v, vis, weights) 

308 

309 # 2 step fit 

310 preproc = FF.preprocess_visibilities(u, v, vis, weights) 

311 sol2 = FF.fit_preprocessed(preproc) 

312 

313 np.testing.assert_equal(sol.I, sol2.I, 

314 err_msg="Testing two-step fit") 

315 

316 

317def test_solve_non_negative(): 

318 """Check FrankFitter fitting routine with non-negative fit using AS 209 dataset""" 

319 AS209, geometry = load_AS209(uv_cut=1e6) 

320 

321 u, v, vis, weights = [AS209[k] for k in ['u', 'v', 'V', 'weights']] 

322 

323 Rmax = 1.6 

324 

325 FF = FrankFitter(Rmax, 20, geometry, alpha=1.05, weights_smooth=1e-2) 

326 

327 sol = FF.fit(u, v, vis, weights) 

328 I_nn = sol.solve_non_negative() 

329 

330 expected = np.array([ 

331 2.42756717e+10, 1.28541672e+10, 1.90032938e+10, 8.31444339e+09, 

332 1.30814112e+10, 1.59442160e+09, 2.93990783e+08, 5.29739902e+09, 

333 1.24011568e+09, 5.40689479e+08, 1.97475180e+08, 2.12294162e+08, 

334 4.45700329e+09, 1.67658919e+09, 8.61662448e+08, 3.81032165e+08, 

335 2.41202443e+08, 7.68452028e+07, 0.00000000e+00, 2.86208170e+07 

336 ]) 

337 

338 np.testing.assert_allclose(I_nn, expected, 

339 err_msg="Testing Frank Fit to AS 209") 

340 

341 

342def test_frank_fitter_log_normal(): 

343 """Check FrankFitter fitting routine with AS 209 dataset""" 

344 AS209, geometry = load_AS209(uv_cut=1e6) 

345 

346 u, v, vis, weights = [AS209[k] for k in ['u', 'v', 'V', 'weights']] 

347 

348 Rmax = 1.6 

349 

350 FF = FrankFitter(Rmax, 20, geometry, alpha=1.3, weights_smooth=1e-2, 

351 method='LogNormal') 

352 

353 sol = FF.fit(u, v, vis, weights) 

354 expected = np.array([ 

355 2.36087004e+10, 1.36923798e+10, 1.82805612e+10, 8.72975548e+09, 

356 1.30516037e+10, 1.28158462e+09, 8.14949172e+08, 4.74472433e+09, 

357 1.66592277e+09, 3.39438704e+08, 1.56219080e+08, 4.42345087e+08, 

358 4.13155298e+09, 2.00246824e+09, 6.07773834e+08, 5.34020982e+08, 

359 1.80820913e+08, 7.71858927e+07, 2.89354816e+07, 2.45967370e+06,]) 

360 

361 np.testing.assert_allclose(sol.MAP, expected, rtol=7e-5, 

362 err_msg="Testing Frank Log-Normal Fit to AS 209") 

363 

364 

365def test_geom_deproject(): 

366 """Check predict works properly with a different geometry""" 

367 AS209, geometry = load_AS209(uv_cut=1e6) 

368 

369 u, v, vis, weights = [AS209[k] for k in ['u', 'v', 'V', 'weights']] 

370 

371 Rmax = 1.6 

372 

373 FF = FrankFitter(Rmax, 20, geometry, alpha=1.05, weights_smooth=1e-2) 

374 sol = FF.fit(u, v, vis, weights) 

375 

376 geom2 = FixedGeometry(0,0,0,0) 

377 

378 q = np.hypot(u,v) 

379 q = np.geomspace(q.min(), q.max(), 20) 

380 

381 V = sol.predict(q, 0*q, geometry=geom2) 

382 Vexpected = [ 

383 0.3152656, 0.31260669, 0.30831366, 0.30143546, 

384 0.29055445, 0.27369896, 0.24848676, 0.2129406, 

385 0.1677038, 0.11989993, 0.08507635, 0.07491307, 

386 0.06555019, 0.01831576, -0.00173855, 0.00042803, 

387 -0.00322264, 0.00278782, -0.00978981, 0.00620259, 

388 ] 

389 

390 np.testing.assert_allclose(V, Vexpected, rtol=2e-5, atol=1e-8, 

391 err_msg="Testing predict with different geometry") 

392 

393 

394def test_fit_geometry_inside(): 

395 """Check the geometry fit embedded in a call to FrankFitter""" 

396 AS209, _ = load_AS209(uv_cut=1e6) 

397 

398 u, v, vis, weights = [AS209[k][::100] for k in ['u', 'v', 'V', 'weights']] 

399 

400 Rmax = 1.6 

401 

402 FF = FrankFitter(Rmax, 20, FitGeometryGaussian(), 

403 alpha=1.05, weights_smooth=1e-2) 

404 

405 sol = FF.fit(u, v, vis, weights) 

406 

407 geom = sol.geometry 

408 np.testing.assert_allclose([geom.inc, geom.PA, 1e3 * geom.dRA, 

409 1e3 * geom.dDec], 

410 [34.50710460482996, 86.4699107557648, 

411 0.21017246809441995, -2.109586872914908], 

412 err_msg="Gaussian geometry fit inside Frank fit") 

413 

414 

415def test_throw_error_on_bad_q_range(): 

416 """Check that frank correctly raises an error when the 

417 q range is bad.""" 

418 AS209, geometry = load_AS209() 

419 

420 u, v, vis, weights = [AS209[k][::100] for k in ['u', 'v', 'V', 'weights']] 

421 

422 Rmax = 1.6 

423 

424 FF = FrankFitter(Rmax, 20, geometry, 

425 alpha=1.05, weights_smooth=1e-2) 

426 

427 try: 

428 FF.fit(u, v, vis, weights) 

429 raise RuntimeError("Expected ValueError due to bad range") 

430 except ValueError: 

431 pass 

432 

433 

434def test_uvbin(): 

435 """Check the uv-data binning routine""" 

436 AS209, geometry = load_AS209() 

437 

438 uv = np.hypot(*geometry.deproject(AS209['u'], AS209['v'])) 

439 

440 uvbin = utilities.UVDataBinner(uv, AS209['V'], AS209['weights'], 50e3) 

441 

442 uvmin = 1e6 

443 uvmax = 1e6 + 50e3 

444 

445 idx = (uv >= uvmin) & (uv < uvmax) 

446 

447 widx = AS209['weights'][idx] 

448 

449 w = np.sum(widx) 

450 V = np.sum(widx*AS209['V'][idx]) / w 

451 q = np.sum(widx*uv[idx]) / w 

452 

453 i = (uvbin.uv >= uvmin) & (uvbin.uv < uvmax) 

454 

455 np.testing.assert_allclose(q, uvbin.uv[i]) 

456 np.testing.assert_allclose(V, uvbin.V[i]) 

457 np.testing.assert_allclose(w, uvbin.weights[i]) 

458 np.testing.assert_allclose(len(widx), uvbin.bin_counts[i]) 

459 

460def test_save_load_sol(): 

461 """Check saving/loading a frank 'sol' object""" 

462 AS209, AS209_geometry = load_AS209(uv_cut=1e6) 

463 u, v, vis, weights = [AS209[k][::100] for k in ['u', 'v', 'V', 'weights']] 

464 Rmax, N = 1.6, 20 

465 

466 # generate a sol from a standard frank fit 

467 FF = FrankFitter(Rmax, N, AS209_geometry, alpha=1.05, weights_smooth=1e-2) 

468 sol = FF.fit(u, v, vis, weights) 

469 

470 # and from a frank debris fit (has additional keys over a standard fit sol) 

471 FF_deb = FrankDebrisFitter(Rmax, N, AS209_geometry, lambda x : 0.05 * x, 

472 alpha=1.05, weights_smooth=1e-2) 

473 sol_deb = FF_deb.fit(u, v, vis, weights) 

474 

475 tmp_dir = '/tmp/frank/tests' 

476 os.makedirs(tmp_dir, exist_ok=True) 

477 

478 save_prefix = [os.path.join(tmp_dir, 'standard'), os.path.join(tmp_dir, 'debris')] 

479 sols = [sol, sol_deb] 

480 

481 for ii, jj in enumerate(save_prefix): 

482 # save the 'sol' object  

483 save_fit(u, v, vis, weights, sols[ii], prefix=jj, 

484 save_profile_fit=False, save_vis_fit=False, save_uvtables=False 

485 ) 

486 # load it 

487 load_sol(jj + '_frank_sol.obj') 

488 

489 

490def test_arcsec_baseline(): 

491 """Check utilities.arcsec_baseline""" 

492 result = utilities.arcsec_baseline(1e6) 

493 np.testing.assert_almost_equal(result, 0.2062648) 

494 

495 

496def test_radius_convert(): 

497 """Check utilities.radius_convert""" 

498 result = utilities.radius_convert(2.0, 100) 

499 assert result == 200 

500 

501 result_bwd = utilities.radius_convert(200.0, 100, conversion='au_arcsec') 

502 assert result_bwd == 2 

503 

504 

505def test_jy_convert(): 

506 """Check utilities.jy_convert""" 

507 x = 10 

508 

509 bmaj, bmin = 0.1, 0.1 

510 

511 expected = {'beam_sterad': 37547916727553.23, 

512 'beam_arcsec2': 882.5424006, 

513 'arcsec2_beam': 0.113309, 

514 'arcsec2_sterad': 425451702961.5221, 

515 'sterad_beam': 2.6632636e-12, 

516 'sterad_arcsec2': 2.3504431e-10} 

517 

518 for key in expected: 

519 result = utilities.jy_convert(x, conversion=key, bmaj=bmaj, bmin=bmin) 

520 np.testing.assert_almost_equal(result, expected[key], decimal=7) 

521 

522 

523def test_get_fit_stat_uncer(): 

524 """Check utilities.get_fit_stat_uncer""" 

525 AS209, AS209_geometry = load_AS209(uv_cut=1e6) 

526 u, v, vis, weights = [AS209[k][::100] for k in ['u', 'v', 'V', 'weights']] 

527 

528 # generate a sol from a standard frank fit 

529 FF = FrankFitter(1.6, 20, AS209_geometry, alpha=1.3, weights_smooth=1e-2) 

530 sol = FF.fit(u, v, vis, weights) 

531 

532 # call with normal model 

533 result = utilities.get_fit_stat_uncer(sol) 

534 expected = [4.29701157e+08, 4.38007127e+08, 3.81819050e+08, 2.80884179e+08, 

535 2.01486438e+08, 1.99436182e+08, 2.15565926e+08, 1.99285093e+08, 

536 1.65080363e+08, 1.49458838e+08, 1.55552558e+08, 1.55906224e+08, 

537 1.39929834e+08, 1.23399577e+08, 1.18687679e+08, 1.20329729e+08, 

538 1.19973246e+08, 1.15672282e+08, 1.05924406e+08, 7.19982652e+07] 

539 

540 np.testing.assert_allclose(result, expected, rtol=2e-5, atol=1e-8) 

541 

542 # call with lognormal model 

543 FF_logn = FrankFitter(1.6, 20, AS209_geometry, alpha=1.3, weights_smooth=1e-2, 

544 method='LogNormal') 

545 sol_logn = FF.fit(u, v, vis, weights) 

546 result_logn = utilities.get_fit_stat_uncer(sol_logn) 

547 np.testing.assert_allclose(result_logn, expected, rtol=2e-5, atol=1e-8) 

548 

549 

550def test_normalize_uv(): 

551 """Check utilities.normalize_uv""" 

552 result = utilities.normalize_uv(1e5, 1e5, 1e-3) 

553 np.testing.assert_almost_equal(result, (([1e8]), ([1e8]))) 

554 

555 

556def test_cut_data_by_baseline(): 

557 """Check utilities.cut_data_by_baseline""" 

558 AS209, AS209_geometry = load_AS209(uv_cut=1e6) 

559 u, v, vis, weights = [AS209[k][::100] for k in ['u', 'v', 'V', 'weights']] 

560 

561 # restrictive cut range to keep only a single baseline 

562 cut_range = [0, 11570] 

563 # call with no supplied geometry 

564 result = utilities.cut_data_by_baseline(u, v, vis, weights, cut_range) 

565 

566 np.testing.assert_almost_equal(result, (([9105.87121309]), 

567 ([-7126.8802574]), 

568 ([0.25705367-0.00452954j]), 

569 ([14390.94693293]) 

570 )) 

571 

572 # restrictive cut range to keep only a single baseline 

573 cut_range_geom = [0, 10370] 

574 # call with supplied geometry 

575 result_geom = utilities.cut_data_by_baseline(u, v, vis, weights, cut_range_geom, 

576 geometry=AS209_geometry) 

577 

578 np.testing.assert_almost_equal(result_geom, (([3080.37968279]), 

579 ([-12126.45120077]), 

580 ([0.01581838-0.22529848j]), 

581 ([47.91090538]) 

582 )) 

583 

584 

585def test_estimate_weights(): 

586 """Check utilities.estimate_weights""" 

587 AS209, AS209_geometry = load_AS209(uv_cut=1e6) 

588 u, v, vis, weights = [AS209[k][::100] for k in ['u', 'v', 'V', 'weights']] 

589 

590 # call with u, v, vis 

591 result = utilities.estimate_weights(u, v, vis) 

592 expected = [343.4011447938792, 1040.388154761485, 323.33779140104497, 

593 670.5799827294733, 746.6778204045879, 262.537321708902, 

594 916.0141170546902, 2478.5780126781183, 220.49922955106743, 

595 343.4011447938792] 

596 np.testing.assert_allclose(result[:10], expected, rtol=2e-5, atol=1e-8) 

597 

598 # call with only u, vis 

599 result_no_v = utilities.estimate_weights(u, vis) 

600 expected_no_v = [140.15619775524289, 140.15619775524289, 

601 136.20331899175486, 144.80828130035127, 

602 751.9714145412686, 14.69762047498323, 

603 775.4926337220135, 106.4511685363733, 

604 188.5850930080213, 299.3538060369927] 

605 np.testing.assert_allclose(result_no_v[:10], expected_no_v, rtol=2e-5, atol=1e-8) 

606 

607 # call with u, v, vis, use_median 

608 result_med = utilities.estimate_weights(u, v, vis, use_median=True) 

609 np.testing.assert_almost_equal(result_med[0], 1040.3881547614856) 

610 

611 

612def test_make_image(): 

613 """Check utilities.make_image""" 

614 AS209, AS209_geometry = load_AS209(uv_cut=1e6) 

615 u, v, vis, weights = [AS209[k][::100] for k in ['u', 'v', 'V', 'weights']] 

616 

617 # generate a sol from a standard frank fit 

618 FF = FrankFitter(1.6, 20, AS209_geometry, alpha=1.3, weights_smooth=1e-2) 

619 sol = FF.fit(u, v, vis, weights) 

620 

621 # call without projection 

622 result = utilities.make_image(sol, Npix=4, project=False) 

623 expected = (([-2.4, -0.8, 0.8, 2.4]), 

624 ([-2.4, -0.8, 0.8, 2.4]), 

625 ([[ 6.47513532e+06, -1.51846712e+07, -1.28084898e+08, -1.51846712e+07], 

626 [-1.51846712e+07, 3.53545721e+07, 3.13428201e+08, 3.53545721e+07], 

627 [-1.28084898e+08, 3.13428201e+08, 3.30602278e+09, 3.13428201e+08], 

628 [-1.51846712e+07, 3.53545721e+07, 3.13428201e+08, 3.53545721e+07]]) 

629 ) 

630 

631 # check pixel coordinates 

632 np.testing.assert_allclose(result[:2], expected[:2], rtol=2e-5, atol=1e-8) 

633 # check pixel brightness 

634 Iresult = np.asarray(result[2]) 

635 Iexpected = np.asarray(expected[2]) 

636 np.testing.assert_allclose(Iresult, Iexpected, rtol=2e-5, atol=1e-8) 

637 

638 # call with projection 

639 result_proj = utilities.make_image(sol, Npix=4, project=True) 

640 expected_proj = (([-2.4, -0.8, 0.8, 2.4]), 

641 ([-2.4, -0.8, 0.8, 2.4]), 

642 ([[ 2.40226316e+06, -8.52280178e+06, -1.37522143e+08, -8.18630605e+06], 

643 [-8.84143838e+06, 2.75149090e+07, 3.28577699e+08, 1.74953407e+07], 

644 [-1.02387759e+08, 2.33877598e+08, 3.44906494e+09, 2.25119693e+08], 

645 [-9.02378853e+06, 1.87158156e+07, 3.35326195e+08, 2.71103094e+07]]) 

646 ) 

647 

648 # check pixel coordinates 

649 np.testing.assert_allclose(result_proj[:2], expected_proj[:2], rtol=2e-5, atol=1e-8) 

650 # check pixel brightness 

651 Iresult_proj = np.asarray(result_proj[2]) 

652 Iexpected_proj = np.asarray(expected_proj[2]) 

653 np.testing.assert_allclose(Iresult_proj, Iexpected_proj, rtol=2e-5, atol=1e-8) 

654 

655 

656def test_add_vis_noise(): 

657 """Check utilities.add_vis_noise""" 

658 # dummy vis and weight 

659 vis, weights = [1.1 + 0.9j], 0.5 

660 

661 result = utilities.add_vis_noise(vis, weights, seed=47) 

662 np.testing.assert_almost_equal(result, [-0.0992665+2.74683048j]) 

663 

664 

665def test_make_mock_data(): 

666 """Check utilities.add_vis_noise""" 

667 AS209, AS209_geometry = load_AS209(uv_cut=1e6) 

668 u, v, vis, weights = [AS209[k][::100] for k in ['u', 'v', 'V', 'weights']] 

669 

670 # generate a sol from a standard frank fit 

671 FF = FrankFitter(1.6, 20, AS209_geometry, alpha=1.3, weights_smooth=1e-2) 

672 sol = FF.fit(u, v, vis, weights) 

673 

674 # call with minimal inputs 

675 result = utilities.make_mock_data(sol.r, sol.I, 3.0, u, v) 

676 expected = [ 0.06699455, 0.18302498, 0.27758414, 0.04789997, -0.00240399, 

677 0.06339718, 0.00358722, 0.28862088, 0.07058801, 0.06617371] 

678 np.testing.assert_allclose(result[1][:10], expected, rtol=2e-5, atol=1e-8) 

679 

680 # call with deprojection 

681 result_dep = utilities.make_mock_data(sol.r, sol.I, 3.0, u, v, 

682 projection='deproject', 

683 geometry=AS209_geometry) 

684 expected_dep = [0.06244746, 0.15925137, 0.2345302 , 0.0623711 , 0.00404342, 

685 0.06277, 0.00361453, 0.23649558, 0.06326574, 0.0632122 ] 

686 np.testing.assert_allclose(result_dep[1][:10], expected_dep, rtol=2e-5, atol=1e-8) 

687 

688 # call with reprojection 

689 result_rep = utilities.make_mock_data(sol.r, sol.I, 3.0, u, v, 

690 projection='reproject', 

691 geometry=AS209_geometry) 

692 expected_rep = [ 0.05219592, 0.11866411, 0.22040989, 0.03889961, -0.00133919, 

693 0.05194375, -0.00054036, 0.23372006, 0.04987515, 0.04541111] 

694 np.testing.assert_allclose(result_rep[1][:10], expected_rep, rtol=2e-5, atol=1e-8)