Coverage for tests / test_fit_bootstrap_model.py: 21%

1# This file is part of multiprofit. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

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

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

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

12# (at your option) any later version. 

13# 

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

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

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

17# GNU General Public License for more details. 

18# 

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

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

21 

22import math 

23 

24import astropy.table 

25import lsst.gauss2d.fit as g2f 

26from lsst.multiprofit.componentconfig import ( 

27 CentroidConfig, 

28 FluxFractionParameterConfig, 

29 FluxParameterConfig, 

30 GaussianComponentConfig, 

31 ParameterConfig, 

32 SersicComponentConfig, 

33 SersicIndexParameterConfig, 

34) 

35from lsst.multiprofit.fitting.fit_bootstrap_model import ( 

36 CatalogExposurePsfBootstrap, 

37 CatalogExposureSourcesBootstrap, 

38 CatalogPsfBootstrapConfig, 

39 CatalogSourceBootstrapConfig, 

40 CatalogSourceFitterBootstrap, 

41 NoisyObservationConfig, 

42 NoisyPsfObservationConfig, 

43) 

44from lsst.multiprofit.fitting.fit_psf import ( 

45 CatalogPsfFitter, 

46 CatalogPsfFitterConfig, 

47 CatalogPsfFitterConfigData, 

48) 

49from lsst.multiprofit.fitting.fit_source import CatalogSourceFitterConfig, CatalogSourceFitterConfigData 

50from lsst.multiprofit.modelconfig import ModelConfig 

51from lsst.multiprofit.modeller import ModelFitConfig 

52from lsst.multiprofit.observationconfig import CoordinateSystemConfig 

53from lsst.multiprofit.plotting import ErrorValues, plot_catalog_bootstrap, plot_loglike 

54from lsst.multiprofit.sourceconfig import ComponentGroupConfig, SourceConfig 

55from lsst.multiprofit.utils import get_params_uniq 

56import numpy as np 

57import pytest 

58 

59shape_img = (23, 27) 

60reff_x_src, reff_y_src, rho_src, nser_src = 2.5, 3.6, -0.25, 2.0 

61 

62# TODO: These can be parameterized; should they be? 

63compute_errors_no_covar = True 

64compute_errors_from_jacobian = True 

65include_point_source = False 

66n_sources = 3 

67# Set to True for interactive debugging (but don't commit) 

68plot = False 

69 

70 

71@pytest.fixture(scope="module") 

72def channels(): 

73 return {band: g2f.Channel.get(band) for band in ("R", "G", "B")} 

74 

75 

76@pytest.fixture(scope="module") 

77def config_fitter_psfs(channels) -> dict[g2f.Channel, CatalogExposurePsfBootstrap]: 

78 config_datas = {} 

79 for idx, (band, channel) in enumerate(channels.items()): 

80 n_rows = 17 + idx * 2 

81 n_cols = 15 + idx * 2 

82 config = CatalogPsfFitterConfig( 

83 model=SourceConfig( 

84 component_groups={ 

85 "": ComponentGroupConfig( 

86 centroids={ 

87 "default": CentroidConfig( 

88 x=ParameterConfig(value_initial=n_cols / 2.0), 

89 y=ParameterConfig(value_initial=n_rows / 2.0), 

90 ), 

91 }, 

92 components_gauss={ 

93 "gauss1": GaussianComponentConfig( 

94 flux=FluxParameterConfig(value_initial=1.0, fixed=True), 

95 fluxfrac=FluxFractionParameterConfig(value_initial=0.5, fixed=False), 

96 size_x=ParameterConfig(value_initial=1.5 + 0.1 * idx), 

97 size_y=ParameterConfig(value_initial=1.7 + 0.13 * idx), 

98 rho=ParameterConfig(value_initial=-0.035 - 0.007 * idx), 

99 ), 

100 "gauss2": GaussianComponentConfig( 

101 size_x=ParameterConfig(value_initial=3.1 + 0.24 * idx), 

102 size_y=ParameterConfig(value_initial=2.7 + 0.16 * idx), 

103 rho=ParameterConfig(value_initial=0.06 + 0.012 * idx), 

104 fluxfrac=FluxFractionParameterConfig(value_initial=1.0, fixed=True), 

105 ), 

106 }, 

107 is_fractional=True, 

108 ) 

109 } 

110 ), 

111 ) 

112 config_boot = CatalogPsfBootstrapConfig( 

113 observation=NoisyPsfObservationConfig(n_rows=n_rows, n_cols=n_cols, gain=1e5), 

114 n_sources=n_sources, 

115 ) 

116 config_data = CatalogExposurePsfBootstrap(config=config, config_boot=config_boot) 

117 config_datas[channel] = config_data 

118 

119 return config_datas 

120 

121 

122@pytest.fixture(scope="module") 

123def config_fitter_source(channels) -> CatalogSourceFitterConfigData: 

124 config = CatalogSourceFitterConfig( 

125 config_fit=ModelFitConfig(fit_linear_iter=3), 

126 config_model=ModelConfig( 

127 sources={ 

128 "": SourceConfig( 

129 component_groups={ 

130 "": ComponentGroupConfig( 

131 components_gauss=( 

132 { 

133 "ps": GaussianComponentConfig( 

134 flux=FluxParameterConfig(value_initial=1000), 

135 rho=ParameterConfig(value_initial=0, fixed=True), 

136 size_x=ParameterConfig(value_initial=0, fixed=True), 

137 size_y=ParameterConfig(value_initial=0, fixed=True), 

138 ) 

139 } 

140 if include_point_source 

141 else {} 

142 ), 

143 components_sersic={ 

144 "ser": SersicComponentConfig( 

145 prior_size_mean=reff_y_src, 

146 prior_size_stddev=1.0, 

147 prior_axrat_mean=reff_x_src / reff_y_src, 

148 prior_axrat_stddev=0.2, 

149 flux=FluxParameterConfig(value_initial=5000), 

150 rho=ParameterConfig(value_initial=rho_src), 

151 size_x=ParameterConfig(value_initial=reff_x_src), 

152 size_y=ParameterConfig(value_initial=reff_y_src), 

153 sersic_index=SersicIndexParameterConfig(fixed=False, value_initial=1.0), 

154 ), 

155 }, 

156 ) 

157 } 

158 ), 

159 }, 

160 ), 

161 convert_cen_xy_to_radec=False, 

162 compute_errors_no_covar=compute_errors_no_covar, 

163 compute_errors_from_jacobian=compute_errors_from_jacobian, 

164 ) 

165 config_data = CatalogSourceFitterConfigData( 

166 channels=tuple(channels.values()), 

167 config=config, 

168 ) 

169 return config_data 

170 

171 

172@pytest.fixture(scope="module") 

173def tables_psf_fits(config_fitter_psfs) -> dict[g2f.Channel, astropy.table.Table]: 

174 fitter = CatalogPsfFitter() 

175 fits = { 

176 channel: fitter.fit( 

177 catexp=config_fitter_psf, 

178 config_data=config_fitter_psf, 

179 ) 

180 for channel, config_fitter_psf in config_fitter_psfs.items() 

181 } 

182 return fits 

183 

184 

185@pytest.fixture(scope="module") 

186def config_data_sources( 

187 config_fitter_psfs, 

188 tables_psf_fits, 

189) -> dict[g2f.Channel, CatalogExposureSourcesBootstrap]: 

190 config_datas = {} 

191 for idx, (channel, config_fitter_psf) in enumerate(config_fitter_psfs.items()): 

192 table_psf_fits = tables_psf_fits[channel] 

193 n_rows = shape_img[0] + idx * 2 

194 n_cols = shape_img[1] + idx * 2 

195 config_boot = CatalogSourceBootstrapConfig( 

196 observation=NoisyObservationConfig( 

197 n_rows=n_rows, 

198 n_cols=n_cols, 

199 band=channel.name, 

200 background=100, 

201 coordsys=CoordinateSystemConfig(x_min=-2 + 3 * idx, y_min=5 - 4 * idx), 

202 ), 

203 n_sources=n_sources, 

204 ) 

205 config_data = CatalogExposureSourcesBootstrap( 

206 config_boot=config_boot, 

207 table_psf_fits=table_psf_fits, 

208 ) 

209 config_datas[channel] = config_data 

210 

211 return config_datas 

212 

213 

214def test_fit_psf(config_fitter_psfs, tables_psf_fits): 

215 for band, results in tables_psf_fits.items(): 

216 assert len(results) == n_sources 

217 assert np.sum(results["mpf_psf_unknown_flag"]) == 0 

218 assert all(np.isfinite(list(results[0].values()))) 

219 config_data_psf = config_fitter_psfs[band] 

220 psf_model_init = config_data_psf.config.make_psf_model() 

221 psfdata = CatalogPsfFitterConfigData(config=config_data_psf.config) 

222 psf_model_fit = psfdata.psf_model 

223 psfdata.init_psf_model(results[0]) 

224 assert len(psf_model_init.components) == len(psf_model_fit.components) 

225 params_init = psf_model_init.parameters() 

226 params_fit = psf_model_fit.parameters() 

227 assert len(params_init) == len(params_fit) 

228 sigma_min_sq = config_data_psf.config.sigma_min**2 

229 for p_init, p_meas in zip(params_init, params_fit): 

230 assert p_meas.fixed == p_init.fixed 

231 if p_meas.fixed: 

232 assert p_init.value == p_meas.value 

233 else: 

234 value = p_meas.value 

235 # TODO: come up with better (noise-dependent) thresholds here 

236 if isinstance(p_init, g2f.IntegralParameterD): 

237 atol, rtol = 0, 0.02 

238 elif isinstance(p_init, g2f.ProperFractionParameterD): 

239 atol, rtol = 0.1, 0.01 

240 elif isinstance(p_init, g2f.RhoParameterD): 

241 atol, rtol = 0.05, 0.1 

242 elif isinstance(p_init, g2f.SigmaXParameterD) or isinstance(p_init, g2f.SigmaYParameterD): 

243 value = math.sqrt(value**2 + sigma_min_sq) 

244 else: 

245 atol, rtol = 0.01, 0.1 

246 assert np.isclose(p_init.value, value, atol=atol, rtol=rtol) 

247 

248 

249def test_fit_source(config_fitter_source, config_data_sources): 

250 fitter = CatalogSourceFitterBootstrap() 

251 # We don't have or need a multiband input catalog - just use the first one 

252 catalog_multi = next(iter(config_data_sources.values())).get_catalog() 

253 catexps = list(config_data_sources.values()) 

254 results = fitter.fit(catalog_multi=catalog_multi, catexps=catexps, config_data=config_fitter_source) 

255 assert len(results) == n_sources 

256 

257 model = fitter.get_model( 

258 0, 

259 catalog_multi=catalog_multi, 

260 catexps=catexps, 

261 config_data=config_fitter_source, 

262 results=results, 

263 ) 

264 

265 model_sources, priors = config_fitter_source.config.make_sources( 

266 channels=list(config_data_sources.keys()) 

267 ) 

268 model_true = g2f.ModelD(data=model.data, psfmodels=model.psfmodels, sources=model_sources) 

269 fitter.initialize_model(model_true, catalog_multi[0], catexps=catexps) 

270 params_true = tuple(param.value for param in get_params_uniq(model_true, fixed=False)) 

271 plot_catalog_bootstrap( 

272 results, histtype="step", paramvals_ref=params_true, plot_total_fluxes=True, plot_colors=True 

273 ) 

274 if plot: 

275 import matplotlib.pyplot as plt 

276 

277 plt.show() 

278 

279 assert np.sum(results["mpf_unknown_flag"]) == 0 

280 assert all(np.isfinite(list(results[0].values()))) 

281 

282 variances = [] 

283 for return_negative in (False, True): 

284 variances.append( 

285 fitter.modeller.compute_variances( 

286 model, 

287 transformed=False, 

288 options=g2f.HessianOptions(return_negative=return_negative), 

289 use_diag_only=True, 

290 ) 

291 ) 

292 assert np.all(variances[-1] > 0) 

293 if return_negative: 

294 variances = np.array(variances) 

295 variances[variances <= 0] = 0 

296 variances = list(variances) 

297 

298 # Bootstrap errors 

299 model.setup_evaluators(evaluatormode=g2f.EvaluatorMode.image) 

300 model.evaluate() 

301 img_data_old = [] 

302 for obs, output in zip(model.data, model.outputs): 

303 img_data_old.append(obs.image.data.copy()) 

304 img = obs.image.data 

305 img.flat = output.data.flat 

306 options_hessian = g2f.HessianOptions(return_negative=return_negative) 

307 variances_bootstrap = fitter.modeller.compute_variances(model, transformed=False, options=options_hessian) 

308 variances_bootstrap_diag = fitter.modeller.compute_variances( 

309 model, transformed=False, options=options_hessian, use_diag_only=True 

310 ) 

311 for obs, img_datum_old in zip(model.data, img_data_old): 

312 obs.image.data.flat = img_datum_old.flat 

313 variances_jac = fitter.modeller.compute_variances(model, transformed=False) 

314 variances_jac_diag = fitter.modeller.compute_variances(model, transformed=False, use_diag_only=True) 

315 

316 errors_plot = { 

317 "inv_hess": ErrorValues(values=np.sqrt(variances[0]), kwargs_plot={"linestyle": "-", "color": "r"}), 

318 "-inv_hess": ErrorValues(values=np.sqrt(variances[1]), kwargs_plot={"linestyle": "--", "color": "r"}), 

319 "inv_jac": ErrorValues(values=np.sqrt(variances_jac), kwargs_plot={"linestyle": "-.", "color": "r"}), 

320 "boot_hess": ErrorValues( 

321 values=np.sqrt(variances_bootstrap), kwargs_plot={"linestyle": "-", "color": "b"} 

322 ), 

323 "boot_diag": ErrorValues( 

324 values=np.sqrt(variances_bootstrap_diag), kwargs_plot={"linestyle": "--", "color": "b"} 

325 ), 

326 "boot_jac_diag": ErrorValues( 

327 values=np.sqrt(variances_jac_diag), kwargs_plot={"linestyle": "-.", "color": "m"} 

328 ), 

329 } 

330 fig, ax = plot_loglike(model, errors=errors_plot, values_reference=params_true) 

331 if plot: 

332 plt.tight_layout() 

333 plt.show()