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 get_model_radec(self, source: Mapping[str, Any], cen_x: float, cen_y: float) -> tuple[float, float]: 

188 return float(cen_x), float(cen_y) 

189 

190 def initialize_model( 

191 self, 

192 model: g2f.ModelD, 

193 source: Mapping[str, Any], 

194 catexps: list[CatalogExposureSourcesABC], 

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

196 centroid_pixel_offset: float = 0, 

197 **kwargs: Any, 

198 ) -> None: 

199 if values_init is None: 

200 values_init = {} 

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

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

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

204 x_min = observation.image.coordsys.x_min 

205 min_x = min(min_x, x_min) 

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

207 y_min = observation.image.coordsys.y_min 

208 min_y = min(min_y, y_min) 

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

210 

211 cen_x = (min_x + max_x) / 2.0 

212 cen_y = (min_y + max_y) / 2.0 

213 

214 # One could consider setting initial values from estimated moments 

215 # here, like a real fitter would 

216 

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

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

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

220 limits_xy = (1e-6, limits_size) 

221 params_limits_init = { 

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

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

224 g2f.ReffXParameterD: (None, limits_xy), 

225 g2f.ReffYParameterD: (None, limits_xy), 

226 g2f.SigmaXParameterD: (None, limits_xy), 

227 g2f.SigmaYParameterD: (None, limits_xy), 

228 } 

229 

230 params_free = get_params_uniq(model, fixed=False) 

231 for param in params_free: 

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

233 if value_init is not None: 

234 param.value = value_init 

235 if limits_new: 

236 param.limits.min = -np.inf 

237 param.limits.max = limits_new[1] 

238 param.limits.min = limits_new[0] 

239 

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

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

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

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

244 model.evaluate() 

245 

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

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

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

249 

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

251 config_obs = catexps[idx_obs].config_boot.observation 

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

253 output_data = model.outputs[idx_obs].data 

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

255 assert image_data.size == output_data.size 

256 # This should definitely never fail 

257 assert image_data.size == sigma_inv_data.size 

258 image_data.flat = output_data.flat 

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

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

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

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

263 # (could instead initialize in get_source_observation) 

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

265 observation.mask_inv.data.flat = 1 

266 

267 def validate_fit_inputs( 

268 self, 

269 catalog_multi: Sequence, 

270 catexps: list[CatalogExposureSourcesABC], 

271 config_data: CatalogSourceFitterConfigData = None, 

272 logger: logging.Logger = None, 

273 **kwargs: Any, 

274 ) -> None: 

275 errors = [] 

276 for idx, catexp in enumerate(catexps): 

277 if not ( 

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

279 and isinstance(config_boot, CatalogSourceBootstrapConfig) 

280 ): 

281 errors.append( 

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

283 f"{CatalogSourceBootstrapConfig}" 

284 )