Coverage for python/lsst/multiprofit/fitting/fit_bootstrap_model.py: 99%

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 

22__all__ = [ 

23 "CatalogBootstrapConfig", 

24 "CatalogExposurePsfBootstrap", 

25 "CatalogExposureSourcesBootstrap", 

26 "CatalogPsfBootstrapConfig", 

27 "CatalogSourceBootstrapConfig", 

28 "CatalogSourceFitterBootstrap", 

29 "NoisyObservationConfig", 

30] 

31 

32import logging 

33from collections.abc import Mapping, Sequence 

34from functools import cached_property 

35from typing import Any, ClassVar 

36 

37import astropy 

38import numpy as np 

39import pydantic 

40 

41import lsst.gauss2d.fit as g2f 

42import lsst.pex.config as pexConfig 

43 

44from ..model_utils import make_image_gaussians 

45from ..observationconfig import ObservationConfig 

46from ..utils import frozen_arbitrary_allowed_config, get_params_uniq, set_config_from_dict 

47from .fit_psf import CatalogExposurePsfABC, CatalogPsfFitterConfig, CatalogPsfFitterConfigData 

48from .fit_source import CatalogExposureSourcesABC, CatalogSourceFitterABC, CatalogSourceFitterConfigData 

49 

50 

51class CatalogBootstrapConfig(pexConfig.Config): 

52 """Configuration for a bootstrap source catalog fitter.""" 

53 

54 n_sources = pexConfig.Field[int](doc="Number of sources", default=1) 

55 

56 @cached_property 

57 def catalog(self) -> astropy.table.Table: 

58 catalog = astropy.table.Table({"id": np.arange(self.n_sources)}) 

59 return catalog 

60 

61 

62class ObservationNoiseConfig(pexConfig.Config): 

63 """Configuration for noise to be added to an Observation. 

64 

65 The background level is in user-defined flux units, should be multiplied 

66 by the gain to obtain counts. 

67 """ 

68 

69 background = pexConfig.Field[float](doc="Background flux per pixel", default=1e-4) 

70 gain = pexConfig.Field[float](doc="Multiplicative factor to convert flux to counts", default=1.0) 

71 

72 

73class NoisyObservationConfig(ObservationConfig, ObservationNoiseConfig): 

74 """Configuration for an observation with noise.""" 

75 

76 

77class NoisyPsfObservationConfig(ObservationConfig, ObservationNoiseConfig): 

78 """Configuration for a PSF observation with noise.""" 

79 

80 

81class CatalogPsfBootstrapConfig(CatalogBootstrapConfig): 

82 """Configuration for a catalog of noisy PSF observations for bootstrapping. 

83 

84 Each row is a stacked and normalized image of any number of point sources. 

85 """ 

86 

87 observation = pexConfig.ConfigField[NoisyPsfObservationConfig]( 

88 doc="The PSF image configuration", 

89 default=NoisyPsfObservationConfig, 

90 ) 

91 

92 

93class CatalogSourceBootstrapConfig(CatalogBootstrapConfig): 

94 """Configuration for a catalog of noisy source observations 

95 for bootstrapping. 

96 

97 Each row is a PSF-convolved observation of the sources in one band. 

98 """ 

99 

100 observation = pexConfig.ConfigField[NoisyObservationConfig]( 

101 doc="The source image configuration", 

102 default=NoisyObservationConfig, 

103 ) 

104 

105 

106class CatalogExposurePsfBootstrap(CatalogExposurePsfABC, CatalogPsfFitterConfigData, pydantic.BaseModel): 

107 """Dataclass for a PSF-convolved bootstrap fitter.""" 

108 

109 model_config: ClassVar[pydantic.ConfigDict] = frozen_arbitrary_allowed_config 

110 

111 config_boot: CatalogPsfBootstrapConfig = pydantic.Field(title="The configuration for bootstrapping") 

112 

113 @cached_property 

114 def image(self) -> np.ndarray: 

115 psf_model_init = self.config.make_psf_model() 

116 # A hacky way to initialize the psf_model property to the same values 

117 # TODO: Include this functionality in fit_psf.py 

118 for param_init, param in zip(get_params_uniq(psf_model_init), get_params_uniq(self.psf_model)): 

119 param.value = param_init.value 

120 image = make_image_gaussians( 

121 psf_model_init.gaussians(g2f.Channel.NONE), 

122 n_rows=self.config_boot.observation.n_rows, 

123 n_cols=self.config_boot.observation.n_cols, 

124 ) 

125 return image.data 

126 

127 def get_catalog(self) -> astropy.table.Table: 

128 return self.config_boot.catalog 

129 

130 def get_psf_image( 

131 self, source: astropy.table.Row | Mapping[str, Any], config: CatalogPsfFitterConfig | None = None 

132 ) -> np.ndarray: 

133 rng = np.random.default_rng(source["id"]) 

134 image = self.image 

135 config_obs = self.config_boot.observation 

136 return image + rng.standard_normal(image.shape) * np.sqrt( 

137 (image + config_obs.background) / config_obs.gain 

138 ) 

139 

140 def model_post_init(self, __context: Any) -> None: 

141 self.config_boot.freeze() 

142 

143 

144class CatalogExposureSourcesBootstrap(CatalogExposureSourcesABC, pydantic.BaseModel): 

145 """A CatalogExposure for bootstrap fitting of source catalogs.""" 

146 

147 model_config: ClassVar[pydantic.ConfigDict] = frozen_arbitrary_allowed_config 

148 

149 config_boot: CatalogSourceBootstrapConfig = pydantic.Field( 

150 title="A CatalogSourceBootstrapConfig to be frozen" 

151 ) 

152 table_psf_fits: astropy.table.Table = pydantic.Field(title="PSF fit parameters for the catalog") 

153 

154 @cached_property 

155 def channel(self) -> g2f.Channel: 

156 channel = g2f.Channel.get(self.config_boot.observation.band) 

157 return channel 

158 

159 def get_catalog(self) -> astropy.table.Table: 

160 return self.config_boot.catalog 

161 

162 def get_psf_model(self, params: Mapping[str, Any]) -> g2f.PsfModel: 

163 psf_model = self.psf_model_data.psf_model 

164 self.psf_model_data.init_psf_model(self.table_psf_fits[params["id"]]) 

165 return psf_model 

166 

167 def get_source_observation(self, source: Mapping[str, Any]) -> g2f.ObservationD: 

168 obs = self.config_boot.observation.make_observation() 

169 return obs 

170 

171 def model_post_init(self, __context: Any) -> None: 

172 config_dict = self.table_psf_fits.meta["config"] 

173 config = CatalogPsfFitterConfig() 

174 set_config_from_dict(config, config_dict) 

175 config_data = CatalogPsfFitterConfigData(config=config) 

176 object.__setattr__(self, "psf_model_data", config_data) 

177 

178 

179class CatalogSourceFitterBootstrap(CatalogSourceFitterABC, pydantic.BaseModel): 

180 """A catalog fitter that bootstraps a single model. 

181 

182 This fitter generates a different noisy image of the specified model for 

183 each row. The resulting catalog can be used to examine performance and 

184 statistics of the best-fit parameters. 

185 """ 

186 

187 model_config: ClassVar[pydantic.ConfigDict] = frozen_arbitrary_allowed_config 

188 

189 def initialize_model( 

190 self, 

191 model: g2f.ModelD, 

192 source: Mapping[str, Any], 

193 catexps: list[CatalogExposureSourcesABC], 

194 values_init: Mapping[g2f.ParameterD, float] | None = None, 

195 centroid_pixel_offset: float = 0, 

196 **kwargs: Any, 

197 ) -> None: 

198 if values_init is None: 

199 values_init = {} 

200 min_x, max_x = np.inf, -np.inf 

201 min_y, max_y = np.inf, -np.inf 

202 for idx_obs, observation in enumerate(model.data): 

203 x_min = observation.image.coordsys.x_min 

204 min_x = min(min_x, x_min) 

205 max_x = max(max_x, x_min + observation.image.n_cols * observation.image.coordsys.dx1) 

206 y_min = observation.image.coordsys.y_min 

207 min_y = min(min_y, y_min) 

208 max_y = max(max_y, y_min + observation.image.n_rows * observation.image.coordsys.dy2) 

209 

210 cen_x = (min_x + max_x) / 2.0 

211 cen_y = (min_y + max_y) / 2.0 

212 

213 # One could consider setting initial values from estimated moments 

214 # here, like a real fitter would 

215 

216 # An R_eff larger than the box size is problematic. This should also 

217 # stop unreasonable size proposals; a log10 transform isn't enough. 

218 limits_size = max(5.0, 2.0 * np.hypot(max_x - min_x, max_y - min_y)) 

219 limits_xy = (1e-6, limits_size) 

220 params_limits_init = { 

221 g2f.CentroidXParameterD: (cen_x, (min_x, max_x)), 

222 g2f.CentroidYParameterD: (cen_y, (min_y, max_y)), 

223 g2f.ReffXParameterD: (None, limits_xy), 

224 g2f.ReffYParameterD: (None, limits_xy), 

225 g2f.SigmaXParameterD: (None, limits_xy), 

226 g2f.SigmaYParameterD: (None, limits_xy), 

227 } 

228 

229 params_free = get_params_uniq(model, fixed=False) 

230 for param in params_free: 

231 value_init, limits_new = params_limits_init.get(type(param), (values_init.get(param), None)) 

232 if value_init is not None: 

233 param.value = value_init 

234 if limits_new: 

235 param.limits.min = -np.inf 

236 param.limits.max = limits_new[1] 

237 param.limits.min = limits_new[0] 

238 

239 # Should be done in get_source_observation, but it gets called first 

240 # ... and therefore does not have the initialization above 

241 # Also, this must be done per-iteration because PSF parameters vary 

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

243 model.evaluate() 

244 

245 # The offset is to keep the rng seed different from the PSF image seed 

246 # It doesn't really need to be so large but it's reasonably safe 

247 rng = np.random.default_rng(source["id"] + 10000000) 

248 

249 for idx_obs, observation in enumerate(model.data): 

250 config_obs = catexps[idx_obs].config_boot.observation 

251 image_data, sigma_inv_data = observation.image.data, observation.sigma_inv.data 

252 output_data = model.outputs[idx_obs].data 

253 # numpy does not warn if these are different lengths, so assert 

254 assert image_data.size == output_data.size 

255 # This should definitely never fail 

256 assert image_data.size == sigma_inv_data.size 

257 image_data.flat = output_data.flat 

258 sigma_inv_data.flat = np.sqrt((image_data + config_obs.background) / config_obs.gain) 

259 image_data.flat += sigma_inv_data.flat * rng.standard_normal(image_data.size) 

260 sigma_inv_data.flat = (1.0 / sigma_inv_data).flat 

261 # This is mandatory because C++ construction does no initialization 

262 # (could instead initialize in get_source_observation) 

263 # TODO: Do some timings to see which is more efficient 

264 observation.mask_inv.data.flat = 1 

265 

266 def validate_fit_inputs( 

267 self, 

268 catalog_multi: Sequence, 

269 catexps: list[CatalogExposureSourcesABC], 

270 config_data: CatalogSourceFitterConfigData = None, 

271 logger: logging.Logger = None, 

272 **kwargs: Any, 

273 ) -> None: 

274 errors = [] 

275 for idx, catexp in enumerate(catexps): 

276 if not ( 276 ↛ 280line 276 didn't jump to line 280 because the condition on line 276 was never true

277 (config_boot := getattr(catexp, "config_boot", None)) 

278 and isinstance(config_boot, CatalogSourceBootstrapConfig) 

279 ): 

280 errors.append( 

281 f"catexps[{idx=}] = {catexp} does not have a config_boot attr of type" 

282 f"{CatalogSourceBootstrapConfig}" 

283 )