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

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 

32from functools import cached_property 

33import logging 

34from typing import Any, ClassVar, Mapping, Sequence 

35 

36import astropy 

37import lsst.gauss2d.fit as g2f 

38import lsst.pex.config as pexConfig 

39import numpy as np 

40import pydantic 

41 

42from ..model_utils import make_image_gaussians 

43from ..observationconfig import ObservationConfig 

44from ..utils import frozen_arbitrary_allowed_config, get_params_uniq, set_config_from_dict 

45from .fit_psf import CatalogExposurePsfABC, CatalogPsfFitterConfig, CatalogPsfFitterConfigData 

46from .fit_source import CatalogExposureSourcesABC, CatalogSourceFitterABC, CatalogSourceFitterConfigData 

47 

48 

49class CatalogBootstrapConfig(pexConfig.Config): 

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

51 

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

53 

54 @cached_property 

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

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

57 return catalog 

58 

59 

60class ObservationNoiseConfig(pexConfig.Config): 

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

62 

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

64 by the gain to obtain counts. 

65 """ 

66 

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

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

69 

70 

71class NoisyObservationConfig(ObservationConfig, ObservationNoiseConfig): 

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

73 

74 

75class NoisyPsfObservationConfig(ObservationConfig, ObservationNoiseConfig): 

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

77 

78 

79class CatalogPsfBootstrapConfig(CatalogBootstrapConfig): 

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

81 

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

83 """ 

84 

85 observation = pexConfig.ConfigField[NoisyPsfObservationConfig]( 

86 doc="The PSF image configuration", 

87 default=NoisyPsfObservationConfig, 

88 ) 

89 

90 

91class CatalogSourceBootstrapConfig(CatalogBootstrapConfig): 

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

93 for bootstrapping. 

94 

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

96 """ 

97 

98 observation = pexConfig.ConfigField[NoisyObservationConfig]( 

99 doc="The source image configuration", 

100 default=NoisyObservationConfig, 

101 ) 

102 

103 

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

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

106 

107 model_config: ClassVar[pydantic.ConfigDict] = frozen_arbitrary_allowed_config 

108 

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

110 

111 @cached_property 

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

113 psf_model_init = self.config.make_psf_model() 

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

115 # TODO: Include this functionality in fit_psf.py 

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

117 param.value = param_init.value 

118 image = make_image_gaussians( 

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

120 n_rows=self.config_boot.observation.n_rows, 

121 n_cols=self.config_boot.observation.n_cols, 

122 ) 

123 return image.data 

124 

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

126 return self.config_boot.catalog 

127 

128 def get_psf_image( 

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

130 ) -> np.ndarray: 

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

132 image = self.image 

133 config_obs = self.config_boot.observation 

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

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

136 ) 

137 

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

139 self.config_boot.freeze() 

140 

141 

142class CatalogExposureSourcesBootstrap(CatalogExposureSourcesABC, pydantic.BaseModel): 

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

144 

145 model_config: ClassVar[pydantic.ConfigDict] = frozen_arbitrary_allowed_config 

146 

147 config_boot: CatalogSourceBootstrapConfig = pydantic.Field( 

148 title="A CatalogSourceBootstrapConfig to be frozen" 

149 ) 

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

151 

152 @cached_property 

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

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

155 return channel 

156 

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

158 return self.config_boot.catalog 

159 

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

161 psf_model = self.psf_model_data.psf_model 

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

163 return psf_model 

164 

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

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

167 return obs 

168 

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

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

171 config = CatalogPsfFitterConfig() 

172 set_config_from_dict(config, config_dict) 

173 config_data = CatalogPsfFitterConfigData(config=config) 

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

175 

176 

177class CatalogSourceFitterBootstrap(CatalogSourceFitterABC, pydantic.BaseModel): 

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

179 

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

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

182 statistics of the best-fit parameters. 

183 """ 

184 

185 model_config: ClassVar[pydantic.ConfigDict] = frozen_arbitrary_allowed_config 

186 

187 def initialize_model( 

188 self, 

189 model: g2f.ModelD, 

190 source: Mapping[str, Any], 

191 catexps: list[CatalogExposureSourcesABC], 

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

193 centroid_pixel_offset: float = 0, 

194 **kwargs: Any, 

195 ) -> None: 

196 if values_init is None: 

197 values_init = {} 

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

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

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

201 x_min = observation.image.coordsys.x_min 

202 min_x = min(min_x, x_min) 

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

204 y_min = observation.image.coordsys.y_min 

205 min_y = min(min_y, y_min) 

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

207 

208 cen_x = (min_x + max_x) / 2.0 

209 cen_y = (min_y + max_y) / 2.0 

210 

211 # One could consider setting initial values from estimated moments 

212 # here, like a real fitter would 

213 

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

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

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

217 limits_xy = (1e-6, limits_size) 

218 params_limits_init = { 

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

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

221 g2f.ReffXParameterD: (None, limits_xy), 

222 g2f.ReffYParameterD: (None, limits_xy), 

223 g2f.SigmaXParameterD: (None, limits_xy), 

224 g2f.SigmaYParameterD: (None, limits_xy), 

225 } 

226 

227 params_free = get_params_uniq(model, fixed=False) 

228 for param in params_free: 

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

230 if value_init is not None: 

231 param.value = value_init 

232 if limits_new: 

233 param.limits.min = -np.inf 

234 param.limits.max = limits_new[1] 

235 param.limits.min = limits_new[0] 

236 

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

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

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

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

241 model.evaluate() 

242 

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

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

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

246 

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

248 config_obs = catexps[idx_obs].config_boot.observation 

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

250 output_data = model.outputs[idx_obs].data 

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

252 assert image_data.size == output_data.size 

253 # This should definitely never fail 

254 assert image_data.size == sigma_inv_data.size 

255 image_data.flat = output_data.flat 

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

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

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

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

260 # (could instead initialize in get_source_observation) 

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

262 observation.mask_inv.data.flat = 1 

263 

264 def validate_fit_inputs( 

265 self, 

266 catalog_multi: Sequence, 

267 catexps: list[CatalogExposureSourcesABC], 

268 config_data: CatalogSourceFitterConfigData = None, 

269 logger: logging.Logger = None, 

270 **kwargs: Any, 

271 ) -> None: 

272 errors = [] 

273 for idx, catexp in enumerate(catexps): 

274 if not ( 

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

276 and isinstance(config_boot, CatalogSourceBootstrapConfig) 

277 ): 

278 errors.append( 

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

280 f"{CatalogSourceBootstrapConfig}" 

281 )