Coverage for frank/io.py: 80%

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 has functions that read in a UVTable and save fit results to 

20 file. 

21""" 

22 

23import os 

24import numpy as np 

25import pickle 

26import logging 

27 

28from frank.utilities import get_fit_stat_uncer 

29 

30def load_uvtable(data_file): 

31 r""" 

32 Read in a UVTable with data to be fit 

33 

34 Parameters 

35 ---------- 

36 data_file : string 

37 UVTable with data to be fit. 

38 If in ASCII format, the table should have columns: 

39 u [lambda] v [lambda] Re(V) [Jy] Im(V) [Jy] Weight [Jy^-2] 

40 If in .npz format, the file should have arrays: 

41 "u" [lambda], "v" [lambda], "V" [Jy; complex: real + imag * 1j], 

42 "weights" [Jy^-2] 

43 

44 Returns 

45 ------- 

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

47 u and v coordinates of observations 

48 vis : array, unit = Jy 

49 Observed visibilities (complex: real + imag * 1j) 

50 weights : array, unit = Jy^-2 

51 Weights on the visibilities, of the form 

52 :math:`1 / \sigma^2` 

53 """ 

54 

55 logging.info(' Loading UVTable') 

56 

57 # Get extension removing compressed part 

58 base, extension = os.path.splitext(data_file) 

59 if extension in {'.gz', '.bz2'}: 

60 extension = os.path.splitext(base)[1] 

61 if extension not in {'.txt', '.dat'}: 

62 raise ValueError("Compressed UV tables (`.gz` or `.bz2`) must be in " 

63 "one of the formats `.txt` or `.dat`.") 

64 

65 if extension in {'.txt', '.dat'}: 

66 u, v, re, im, weights = np.genfromtxt(data_file).T 

67 vis = re + 1j*im 

68 

69 elif extension == '.npz': 

70 dat = np.load(data_file) 

71 u, v, vis, weights = [dat[i] for i in ['u', 'v', 'V', 'weights']] 

72 if not np.iscomplexobj(vis): 

73 raise ValueError("You provided a UVTable with the extension {}." 

74 " This extension requires the UVTable's variable 'V' to be" 

75 " complex (of the form Re(V) + Im(V) * 1j).".format(extension)) 

76 

77 else: 

78 raise ValueError("You provided a UVTable with the extension {}." 

79 " Please provide it as a `.txt`, `.dat`, or `.npz`." 

80 " Formats .txt and .dat may optionally be" 

81 " compressed (`.gz`, `.bz2`).".format(extension)) 

82 

83 return u, v, vis, weights 

84 

85 

86def save_uvtable(filename, u, v, vis, weights): 

87 r"""Save a uvtable to file. 

88 

89 Parameters 

90 ---------- 

91 filename : string 

92 File to save the uvtable to. 

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

94 u and v coordinates of observations 

95 vis : array, unit = Jy 

96 Complex visibilities 

97 weights : array, unit = Jy^-2 

98 Weights on the visibilities, of the form 

99 :math:`1 / \sigma^2` 

100 """ 

101 

102 extension = os.path.splitext(filename)[1] 

103 if extension not in {'.txt', '.dat', '.npz'}: 

104 raise ValueError("file extension must be 'npz', 'txt', or 'dat'.") 

105 

106 if extension in {'.txt', '.dat'}: 

107 header = 'u [lambda]\tv [lambda]\tRe(V) [Jy]\tIm(V) [Jy]\tWeight [Jy^-2]' 

108 

109 np.savetxt(filename, 

110 np.stack([u, v, vis.real, vis.imag, 

111 weights], axis=-1), 

112 header=header) 

113 

114 elif extension == '.npz': 

115 np.savez(filename, 

116 u=u, v=v, V=vis, weights=weights, 

117 units={'u': 'lambda', 'v': 'lambda', 

118 'V': 'Jy', 'weights': "Jy^-2"}) 

119 

120 

121def load_sol(sol_file): 

122 """Load a frank solution object 

123 

124 Parameters 

125 ---------- 

126 sol_file : string 

127 Filename for frank solution object, '*.obj' 

128 

129 Returns 

130 ---------- 

131 sol : _HankelRegressor object 

132 frank solution object 

133 (see frank.radial_fitters.FrankFitter) 

134 """ 

135 

136 with open(sol_file, 'rb') as f: 

137 sol = pickle.load(f) 

138 

139 return sol 

140 

141 

142def save_fit(u, v, vis, weights, sol, prefix, save_solution=True, 

143 save_profile_fit=True, save_vis_fit=True, save_uvtables=True, 

144 save_iteration_diag=False, iteration_diag=None, 

145 format='npz', 

146 ): 

147 r""" 

148 Save datafiles of fit results 

149 

150 Parameters 

151 ---------- 

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

153 u and v coordinates of observations 

154 vis : array, unit = Jy 

155 Complex visibilities 

156 weights : array, unit = Jy^-2 

157 Weights on the visibilities, of the form 

158 :math:`1 / \sigma^2` 

159 sol : _HankelRegressor object 

160 Reconstructed profile using Maximum a posteriori power spectrum 

161 (see frank.radial_fitters.FrankFitter) 

162 prefix : string 

163 Base part of the filename to which files will be saved 

164 save_solution : bool 

165 Whether to save `sol` object (see frank.radial_fitters.FrankFitter) 

166 save_profile_fit : bool 

167 Whether to save fitted brightness profile 

168 save_vis_fit : bool 

169 Whether to save fitted visibility distribution 

170 NOTE: This is deprojected 

171 save_uvtables : bool 

172 Whether to save fitted and residual UVTables 

173 NOTE: These are reprojected 

174 save_iteration_diag : bool 

175 Whether to save diagnostics of the fit iteration 

176 iteration_diag : dict 

177 Diagnostics of the fit iteration 

178 (see radial_fitters.FrankFitter.fit) 

179 format : string, default = 'npz' 

180 File format in which to save the fit's output UVTable(s) 

181 """ 

182 

183 logging.info(' Saving fit results to {}*'.format(prefix)) 

184 

185 

186 if not format in {'txt', 'dat', 'npz'}: 

187 raise ValueError("'format' must be 'npz', 'txt', or 'dat'.") 

188 

189 if save_solution: 

190 with open(prefix + '_frank_sol.obj', 'wb') as f: 

191 pickle.dump(sol, f) 

192 

193 if save_iteration_diag: 

194 with open(prefix + '_frank_iteration_diagnostics.obj', 'wb') as f: 

195 pickle.dump(iteration_diag, f) 

196 

197 if save_profile_fit: 

198 unc = get_fit_stat_uncer(sol) 

199 np.savetxt(prefix + '_frank_profile_fit.txt', 

200 np.array([sol.r, sol.I, unc]).T, 

201 header='r [arcsec]\tI [Jy/sr]\tI_uncer [Jy/sr]') 

202 

203 if save_vis_fit: 

204 np.savetxt(prefix + '_frank_vis_fit.' + format, 

205 np.array([sol.q, sol.predict_deprojected(sol.q).real]).T, 

206 header='Baseline [lambda]\tProjected Re(V) [Jy]') 

207 

208 

209 if save_uvtables: 

210 logging.info(' Saving fit and residual UVTables. N.B.: These will' 

211 ' be of comparable size to your input UVTable') 

212 

213 V_pred = sol.predict(u, v) 

214 

215 save_uvtable(prefix + '_frank_uv_fit.' + format, 

216 u, v, V_pred, weights) 

217 save_uvtable(prefix + '_frank_uv_resid.' + format, 

218 u, v, vis - V_pred, weights)