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 astropy.table 

23import lsst.gauss2d.fit as g2f 

24from lsst.multiprofit.componentconfig import ( 

25 CentroidConfig, 

26 FluxFractionParameterConfig, 

27 FluxParameterConfig, 

28 GaussianComponentConfig, 

29 ParameterConfig, 

30 SersicComponentConfig, 

31 SersicIndexParameterConfig, 

32) 

33from lsst.multiprofit.fitting.fit_bootstrap_model import ( 

34 CatalogExposurePsfBootstrap, 

35 CatalogExposureSourcesBootstrap, 

36 CatalogPsfBootstrapConfig, 

37 CatalogSourceBootstrapConfig, 

38 CatalogSourceFitterBootstrap, 

39 NoisyObservationConfig, 

40 NoisyPsfObservationConfig, 

41) 

42from lsst.multiprofit.fitting.fit_psf import ( 

43 CatalogPsfFitter, 

44 CatalogPsfFitterConfig, 

45 CatalogPsfFitterConfigData, 

46) 

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

48from lsst.multiprofit.modelconfig import ModelConfig 

49from lsst.multiprofit.modeller import ModelFitConfig 

50from lsst.multiprofit.observationconfig import CoordinateSystemConfig 

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

52from lsst.multiprofit.sourceconfig import ComponentGroupConfig, SourceConfig 

53from lsst.multiprofit.utils import get_params_uniq 

54import numpy as np 

55import pytest 

56 

57shape_img = (23, 27) 

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

59 

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

61compute_errors_no_covar = True 

62compute_errors_from_jacobian = True 

63include_point_source = False 

64n_sources = 3 

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

66plot = False 

67 

68 

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

70def channels(): 

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

72 

73 

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

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

76 config_datas = {} 

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

78 n_rows = 17 + idx * 2 

79 n_cols = 15 + idx * 2 

80 config = CatalogPsfFitterConfig( 

81 model=SourceConfig( 

82 component_groups={ 

83 "": ComponentGroupConfig( 

84 centroids={ 

85 "default": CentroidConfig( 

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

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

88 ), 

89 }, 

90 components_gauss={ 

91 "gauss1": GaussianComponentConfig( 

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

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

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

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

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

97 ), 

98 "gauss2": GaussianComponentConfig( 

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

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

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

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

103 ), 

104 }, 

105 is_fractional=True, 

106 ) 

107 } 

108 ), 

109 ) 

110 config_boot = CatalogPsfBootstrapConfig( 

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

112 n_sources=n_sources, 

113 ) 

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

115 config_datas[channel] = config_data 

116 

117 return config_datas 

118 

119 

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

121def config_fitter_source(channels) -> CatalogSourceFitterConfigData: 

122 config = CatalogSourceFitterConfig( 

123 config_fit=ModelFitConfig(fit_linear_iter=3), 

124 config_model=ModelConfig( 

125 sources={ 

126 "": SourceConfig( 

127 component_groups={ 

128 "": ComponentGroupConfig( 

129 components_gauss=( 

130 { 

131 "ps": GaussianComponentConfig( 

132 flux=FluxParameterConfig(value_initial=1000), 

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

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

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

136 ) 

137 } 

138 if include_point_source 

139 else {} 

140 ), 

141 components_sersic={ 

142 "ser": SersicComponentConfig( 

143 prior_size_mean=reff_y_src, 

144 prior_size_stddev=1.0, 

145 prior_axrat_mean=reff_x_src / reff_y_src, 

146 prior_axrat_stddev=0.2, 

147 flux=FluxParameterConfig(value_initial=5000), 

148 rho=ParameterConfig(value_initial=rho_src), 

149 size_x=ParameterConfig(value_initial=reff_x_src), 

150 size_y=ParameterConfig(value_initial=reff_y_src), 

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

152 ), 

153 }, 

154 ) 

155 } 

156 ), 

157 }, 

158 ), 

159 convert_cen_xy_to_radec=False, 

160 compute_errors_no_covar=compute_errors_no_covar, 

161 compute_errors_from_jacobian=compute_errors_from_jacobian, 

162 ) 

163 config_data = CatalogSourceFitterConfigData( 

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

165 config=config, 

166 ) 

167 return config_data 

168 

169 

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

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

172 fitter = CatalogPsfFitter() 

173 fits = { 

174 channel: fitter.fit( 

175 catexp=config_fitter_psf, 

176 config_data=config_fitter_psf, 

177 ) 

178 for channel, config_fitter_psf in config_fitter_psfs.items() 

179 } 

180 return fits 

181 

182 

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

184def config_data_sources( 

185 config_fitter_psfs, 

186 tables_psf_fits, 

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

188 config_datas = {} 

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

190 table_psf_fits = tables_psf_fits[channel] 

191 n_rows = shape_img[0] + idx * 2 

192 n_cols = shape_img[1] + idx * 2 

193 config_boot = CatalogSourceBootstrapConfig( 

194 observation=NoisyObservationConfig( 

195 n_rows=n_rows, 

196 n_cols=n_cols, 

197 band=channel.name, 

198 background=100, 

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

200 ), 

201 n_sources=n_sources, 

202 ) 

203 config_data = CatalogExposureSourcesBootstrap( 

204 config_boot=config_boot, 

205 table_psf_fits=table_psf_fits, 

206 ) 

207 config_datas[channel] = config_data 

208 

209 return config_datas 

210 

211 

212def test_fit_psf(config_fitter_psfs, tables_psf_fits): 

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

214 assert len(results) == n_sources 

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

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

217 config_data_psf = config_fitter_psfs[band] 

218 psf_model_init = config_data_psf.config.make_psf_model() 

219 psfdata = CatalogPsfFitterConfigData(config=config_data_psf.config) 

220 psf_model_fit = psfdata.psf_model 

221 psfdata.init_psf_model(results[0]) 

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

223 params_init = psf_model_init.parameters() 

224 params_fit = psf_model_fit.parameters() 

225 assert len(params_init) == len(params_fit) 

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

227 assert p_meas.fixed == p_init.fixed 

228 if p_meas.fixed: 

229 assert p_init.value == p_meas.value 

230 else: 

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

232 if isinstance(p_init, g2f.IntegralParameterD): 

233 atol, rtol = 0, 0.02 

234 elif isinstance(p_init, g2f.ProperFractionParameterD): 

235 atol, rtol = 0.1, 0.01 

236 elif isinstance(p_init, g2f.RhoParameterD): 

237 atol, rtol = 0.05, 0.1 

238 else: 

239 atol, rtol = 0.01, 0.1 

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

241 

242 

243def test_fit_source(config_fitter_source, config_data_sources): 

244 fitter = CatalogSourceFitterBootstrap() 

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

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

247 catexps = list(config_data_sources.values()) 

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

249 assert len(results) == n_sources 

250 

251 model = fitter.get_model( 

252 0, 

253 catalog_multi=catalog_multi, 

254 catexps=catexps, 

255 config_data=config_fitter_source, 

256 results=results, 

257 ) 

258 

259 model_sources, priors = config_fitter_source.config.make_sources( 

260 channels=list(config_data_sources.keys()) 

261 ) 

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

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

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

265 plot_catalog_bootstrap( 

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

267 ) 

268 if plot: 

269 import matplotlib.pyplot as plt 

270 

271 plt.show() 

272 

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

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

275 

276 variances = [] 

277 for return_negative in (False, True): 

278 variances.append( 

279 fitter.modeller.compute_variances( 

280 model, 

281 transformed=False, 

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

283 use_diag_only=True, 

284 ) 

285 ) 

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

287 if return_negative: 

288 variances = np.array(variances) 

289 variances[variances <= 0] = 0 

290 variances = list(variances) 

291 

292 # Bootstrap errors 

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

294 model.evaluate() 

295 img_data_old = [] 

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

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

298 img = obs.image.data 

299 img.flat = output.data.flat 

300 options_hessian = g2f.HessianOptions(return_negative=return_negative) 

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

302 variances_bootstrap_diag = fitter.modeller.compute_variances( 

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

304 ) 

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

306 obs.image.data.flat = img_datum_old.flat 

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

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

309 

310 errors_plot = { 

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

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

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

314 "boot_hess": ErrorValues( 

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

316 ), 

317 "boot_diag": ErrorValues( 

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

319 ), 

320 "boot_jac_diag": ErrorValues( 

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

322 ), 

323 } 

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

325 if plot: 

326 plt.tight_layout() 

327 plt.show()