Coverage for python / lsst / meas / extensions / multiprofit / pipetasks_fit.py: 39%

243 statements  

« prev     ^ index     » next       coverage.py v7.12.0, created at 2025-12-20 01:40 -0800

1# This file is part of meas_extensions_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 "component_names_default", 

24 "model_names_default", 

25 "MultiProFitCoaddPsfFitConfig", 

26 "MultiProFitCoaddPsfFitTask", 

27 "MultiProFitCoaddObjectFitConfig", 

28 "MultiProFitCoaddPointFitConfig", 

29 "MultiProFitCoaddSersicFitConfig", 

30 "MultiProFitCoaddSersicFitTask", 

31 "MultiProFitCoaddGaussFitConfig", 

32 "MultiProFitCoaddGaussFitTask", 

33 "MultiProFitCoaddExpFitConfig", 

34 "MultiProFitCoaddExpFitTask", 

35 "MultiProFitCoaddDeVFitConfig", 

36 "MultiProFitCoaddDeVFitTask", 

37 "MultiProFitCoaddExpDeVFitConfig", 

38 "MultiProFitCoaddExpDeVFitTask", 

39) 

40 

41from abc import abstractmethod 

42import math 

43from types import SimpleNamespace 

44from typing import Any, Mapping, Sequence 

45 

46import lsst.gauss2d.fit as g2f 

47from lsst.multiprofit.componentconfig import ( 

48 GaussianComponentConfig, 

49 ParameterConfig, 

50 SersicComponentConfig, 

51 SersicIndexParameterConfig, 

52) 

53from lsst.multiprofit.fitting.fit_source import CatalogExposureSourcesABC, CatalogSourceFitterConfigData 

54from lsst.multiprofit.modelconfig import ModelConfig 

55from lsst.multiprofit.sourceconfig import ComponentGroupConfig, SourceConfig 

56from lsst.pex.config import ConfigDictField, Field 

57from lsst.pipe.tasks.fit_coadd_multiband import ( 

58 CatalogExposureInputs, 

59 CoaddMultibandFitConfig, 

60 CoaddMultibandFitConnections, 

61 CoaddMultibandFitTask, 

62) 

63from lsst.pipe.tasks.fit_coadd_psf import CoaddPsfFitConfig, CoaddPsfFitConnections, CoaddPsfFitTask 

64 

65from .fit_coadd_multiband import ( 

66 CachedBasicModelInitializer, 

67 MagnitudeDependentSizePriorConfig, 

68 MakeBasicInitializerAction, 

69 ModelInitializer, 

70 MultiProFitSourceTask, 

71 PsfComponentsActionBase, 

72 SourceTablePsfComponentsAction, 

73) 

74from .fit_coadd_psf import MultiProFitPsfTask 

75from .input_config import InputConfig 

76 

77component_names_default = SimpleNamespace( 

78 point="point", 

79 gauss="gauss", 

80 exp="exp", 

81 deV="deV", 

82 sersic="sersic", 

83) 

84 

85model_names_default = SimpleNamespace( 

86 point="Point", 

87 gauss="Gauss", 

88 exp="Exp", 

89 deV="DeV", 

90 sersic="Sersic", 

91 fixed_cen="FixedCen", 

92 shapelet_psf="ShapeletPsf", 

93) 

94 

95 

96class MultiProFitCoaddPsfFitConfig( 

97 CoaddPsfFitConfig, 

98 pipelineConnections=CoaddPsfFitConnections, 

99): 

100 """MultiProFit PSF fit task config.""" 

101 

102 def setDefaults(self): 

103 super().setDefaults() 

104 self.fit_coadd_psf.retarget(MultiProFitPsfTask) 

105 self.fit_coadd_psf.config_fit.eval_residual = False 

106 

107 

108class MultiProFitCoaddPsfFitTask(CoaddPsfFitTask): 

109 """MultiProFit PSF fit task.""" 

110 

111 ConfigClass = MultiProFitCoaddPsfFitConfig 

112 _DefaultName = "multiProFitCoaddPsfFit" 

113 

114 

115class MultiProFitCoaddObjectFitConnections(CoaddMultibandFitConnections): 

116 def __init__(self, *, config=None): 

117 super().__init__(config=config) 

118 for name, config_input in config.inputs_init.items(): 

119 if hasattr(self, name): 

120 raise ValueError( 

121 f"{config_input=} {name=} is invalid, due to being an existing attribute" f" of {self=}" 

122 ) 

123 if config_input.is_multipatch or not config_input.is_multiband: 

124 raise ValueError( 

125 f"Single-band and/or multipatch initialization config_input entries ({name})" 

126 f" are not supported yet." 

127 ) 

128 connection = config_input.get_connection(name) 

129 setattr(self, name, connection) 

130 

131 

132class MultiProFitCoaddObjectFitConfig( 

133 CoaddMultibandFitConfig, 

134 pipelineConnections=MultiProFitCoaddObjectFitConnections, 

135): 

136 """Generic MultiProFit source fit task config.""" 

137 

138 inputs_init = ConfigDictField( 

139 doc="Mapping of optional input dataset configs by name, for initialization", 

140 keytype=str, 

141 itemtype=InputConfig, 

142 default={}, 

143 ) 

144 

145 # This needs to be set, ideally in setDefaults of subclasses 

146 name_model = Field[str](doc="The name of the model", default=None) 

147 

148 def _get_source(self): 

149 return next(iter(self.fit_coadd_multiband.config_model.sources.values())) 

150 

151 def _get_component_group(self, source: SourceConfig | None = None): 

152 if source is None: 

153 source = self._get_source() 

154 return next(iter(source.component_groups.values())) 

155 

156 def add_point_source(self, name: str | None = None): 

157 """Add a point source component. 

158 

159 Parameters 

160 ---------- 

161 name 

162 The name of the component. 

163 """ 

164 if name is None: 

165 name = component_names_default.point 

166 source = self._get_source() 

167 group = self._get_component_group(source=source) 

168 if name in group.components_gauss: 

169 raise RuntimeError(f"{name=} component already exists in {source=}") 

170 group.components_gauss[name] = self.make_point_source_component() 

171 self.connections.name_table += model_names_default.point 

172 

173 def finalize( 

174 self, 

175 add_point_source: bool = False, 

176 fix_centroid: bool = False, 

177 use_shapelet_psf: bool = False, 

178 ): 

179 """Apply runtime configuration changes to this config. 

180 

181 Parameters 

182 ---------- 

183 add_point_source 

184 Whether to add a point source component. 

185 fix_centroid 

186 Whether to fix the centroid. 

187 use_shapelet_psf 

188 Whether to initialize PSF parameters from prior shapelet fits. 

189 """ 

190 if add_point_source: 

191 self.add_point_source() 

192 if fix_centroid: 

193 self.fix_centroid() 

194 if use_shapelet_psf: 

195 self.use_shapelet_psf() 

196 

197 def fix_centroid(self): 

198 """Fix (freeze) the source centroid parameters.""" 

199 group = self._get_component_group() 

200 centroids = group.centroids["default"] 

201 centroids.x.fixed = True 

202 centroids.y.fixed = True 

203 self.connections.name_table += model_names_default.fixed_cen 

204 

205 @classmethod 

206 @abstractmethod 

207 def get_model_name_default(cls) -> str: 

208 """Return the default name for this model in table columns.""" 

209 raise NotImplementedError("Subclasses must implement get_model_name_default") 

210 

211 @classmethod 

212 def get_model_name_full(cls) -> str: 

213 """Return a longer, more descriptive name for the model.""" 

214 return cls.get_model_name_default() 

215 

216 @abstractmethod 

217 def make_default_model_config(self) -> ModelConfig: 

218 """Make a default configuration object for this model.""" 

219 raise NotImplementedError("Subclasses must implement make_default_model_config") 

220 

221 @staticmethod 

222 def make_point_source_component() -> GaussianComponentConfig: 

223 """Make a point source component config (zero-size Gaussian).""" 

224 return GaussianComponentConfig( 

225 size_x=ParameterConfig(value_initial=0.0, fixed=True), 

226 size_y=ParameterConfig(value_initial=0.0, fixed=True), 

227 rho=ParameterConfig(value_initial=0.0, fixed=True), 

228 ) 

229 

230 @staticmethod 

231 def make_sersic_component(**kwargs) -> SersicComponentConfig: 

232 """Make a default Sersic component config. 

233 

234 Parameters 

235 ---------- 

236 **kwargs 

237 Keyword arguments to pass to the SersicIndexParameterConfig. 

238 

239 Returns 

240 ------- 

241 config 

242 The default-initialized config. 

243 """ 

244 return SersicComponentConfig( 

245 prior_axrat_stddev=0.8, 

246 prior_size_stddev=0.2, 

247 sersic_index=SersicIndexParameterConfig(**kwargs), 

248 ) 

249 

250 @staticmethod 

251 def make_single_model_config(group: ComponentGroupConfig) -> ModelConfig: 

252 """Make a default single-source, single component group config. 

253 

254 Parameters 

255 ---------- 

256 group 

257 The component group config for the single source. 

258 

259 Returns 

260 ------- 

261 config 

262 A model config with a single nameless source and component group. 

263 """ 

264 return ModelConfig( 

265 sources={ 

266 "": SourceConfig( 

267 component_groups={ 

268 "": group, 

269 } 

270 ) 

271 } 

272 ) 

273 

274 def setDefaults(self): 

275 super().setDefaults() 

276 self.fit_coadd_multiband.retarget(MultiProFitSourceTask) 

277 self.fit_coadd_multiband.action_psf = PsfComponentsActionBase() 

278 self.fit_coadd_multiband.bands_fit = ("u", "g", "r", "i", "z", "y") 

279 

280 self.fit_coadd_multiband.config_model = self.make_default_model_config() 

281 self.name_model = self.get_model_name_default() 

282 self.connections.name_table = self.name_model 

283 

284 def use_shapelet_psf(self): 

285 """Reconfigure self to use prior shapelet PSF fit parameters.""" 

286 self.fit_coadd_multiband.action_psf = SourceTablePsfComponentsAction() 

287 self.drop_psf_connection = True 

288 self.connections.name_table += model_names_default.shapelet_psf 

289 

290 

291class MultiProFitCoaddObjectFitTask(CoaddMultibandFitTask): 

292 """MultiProFit coadd object model fitting task.""" 

293 

294 ConfigClass = MultiProFitCoaddObjectFitConfig 

295 _DefaultName = "multiProFitCoaddObjectFit" 

296 

297 def make_kwargs(self, butlerQC, inputRefs, inputs): 

298 inputs_init = {name: (config, inputs[name][0]) for name, config in self.config.inputs_init.items()} 

299 kwargs = {} 

300 if inputs_init: 

301 kwargs["inputs_init"] = inputs_init 

302 

303 return kwargs 

304 

305 

306class MultiProFitCoaddPointFitConfig( 

307 MultiProFitCoaddObjectFitConfig, 

308 pipelineConnections=MultiProFitCoaddObjectFitConnections, 

309): 

310 """MultiProFit single Sersic model fit task config.""" 

311 

312 @classmethod 

313 def get_model_name_default(cls) -> str: 

314 return model_names_default.point 

315 

316 @classmethod 

317 def get_model_name_full(cls) -> str: 

318 return "Point Source" 

319 

320 def make_default_model_config(self) -> ModelConfig: 

321 config_group = ComponentGroupConfig() 

322 # This is a bit silly but add_point_source will look for the first 

323 # source so it must be added now. Perhaps add_point_source should 

324 # add to a config instance or only self by default 

325 self.fit_coadd_multiband.config_model = self.make_single_model_config(group=config_group) 

326 self.add_point_source() 

327 return self.fit_coadd_multiband.config_model 

328 

329 

330class MultiProFitCoaddSersicFitConfig( 

331 MultiProFitCoaddObjectFitConfig, 

332 pipelineConnections=MultiProFitCoaddObjectFitConnections, 

333): 

334 """MultiProFit single Sersic model fit task config.""" 

335 

336 def _rename_defaults( 

337 self, 

338 name_new: str, 

339 name_model: str | None = None, 

340 name_old: str | None = None, 

341 index_new: float | None = None, 

342 fix_index: bool = False, 

343 ): 

344 """Rename the default Sersic component to something more specific. 

345 

346 This is intended for fixed index models such as exponential and 

347 deVaucouleurs. 

348 

349 Parameters 

350 ---------- 

351 name_new 

352 The new name for the component. 

353 name_model 

354 The new name of the model. Default is to capitalize name_new. 

355 name_old 

356 The old name of the component. Default is to set to 

357 component_names_default.sersic. 

358 index_new 

359 The initial value for the Sersic index. 

360 fix_index 

361 Whether the fix the index to the new value. 

362 """ 

363 if name_old is None: 

364 name_old = component_names_default.sersic 

365 if name_model is None: 

366 name_model = name_new.capitalize() 

367 group = self._get_component_group() 

368 comps_sersic = group.components_sersic 

369 

370 if name_new in comps_sersic: 

371 raise RuntimeError(f"{name_new=} is already in {comps_sersic=}") 

372 

373 comp_sersic = comps_sersic[name_old] 

374 del comps_sersic[name_old] 

375 if index_new is not None: 

376 comp_sersic.sersic_index.value_initial = index_new 

377 if fix_index: 

378 comp_sersic.sersic_index.fixed = True 

379 comps_sersic[name_new] = comp_sersic 

380 

381 if prior_old := self.fit_coadd_multiband.size_priors.get(name_old): 

382 self.fit_coadd_multiband.size_priors[name_new] = prior_old 

383 del self.fit_coadd_multiband.size_priors[name_old] 

384 

385 self.name_model = name_model 

386 self.connections.name_table = name_model 

387 

388 @classmethod 

389 def get_model_name_default(cls) -> str: 

390 return model_names_default.sersic 

391 

392 @classmethod 

393 def get_model_name_full(cls) -> str: 

394 return "Sersic" 

395 

396 def make_default_model_config(self) -> ModelConfig: 

397 config_group = ComponentGroupConfig( 

398 components_sersic={ 

399 component_names_default.sersic: self.make_sersic_component(), 

400 }, 

401 ) 

402 return self.make_single_model_config(group=config_group) 

403 

404 def setDefaults(self): 

405 super().setDefaults() 

406 # This is in pixels and based on DC2. See DM-46498 for details. 

407 self.fit_coadd_multiband.size_priors[component_names_default.sersic] = ( 

408 MagnitudeDependentSizePriorConfig( 

409 intercept_mag=22.6, 

410 slope_median_per_mag=-0.15, 

411 slope_stddev_per_mag=0, 

412 ) 

413 ) 

414 

415 

416class MultiProFitCoaddSersicFitTask(MultiProFitCoaddObjectFitTask): 

417 """MultiProFit single Sersic model fit task.""" 

418 

419 ConfigClass = MultiProFitCoaddSersicFitConfig 

420 _DefaultName = "multiProFitCoaddSersicFit" 

421 

422 

423class MultiProFitCoaddGaussFitConfig( 

424 MultiProFitCoaddSersicFitConfig, 

425 pipelineConnections=MultiProFitCoaddObjectFitConnections, 

426): 

427 """MultiProFit single Gaussian model fit task config.""" 

428 

429 @classmethod 

430 def get_model_name_default(cls) -> str: 

431 return model_names_default.gauss 

432 

433 @classmethod 

434 def get_model_name_full(cls) -> str: 

435 return "Gaussian" 

436 

437 def setDefaults(self): 

438 super().setDefaults() 

439 self._rename_defaults( 

440 name_new=component_names_default.gauss, 

441 name_model=model_names_default.gauss, 

442 index_new=0.5, 

443 fix_index=True, 

444 ) 

445 

446 

447class MultiProFitCoaddGaussFitTask(MultiProFitCoaddObjectFitTask): 

448 """MultiProFit single Gaussian model fit task.""" 

449 

450 ConfigClass = MultiProFitCoaddGaussFitConfig 

451 _DefaultName = "multiProFitCoaddGaussFit" 

452 

453 

454class MultiProFitCoaddExpFitConfig( 

455 MultiProFitCoaddSersicFitConfig, 

456 pipelineConnections=MultiProFitCoaddObjectFitConnections, 

457): 

458 """MultiProFit single exponential model fit task config.""" 

459 

460 @classmethod 

461 def get_model_name_default(cls) -> str: 

462 return model_names_default.exp 

463 

464 @classmethod 

465 def get_model_name_full(cls) -> str: 

466 return "Exponential" 

467 

468 def setDefaults(self): 

469 super().setDefaults() 

470 self._rename_defaults( 

471 name_new=component_names_default.exp, 

472 name_model=model_names_default.exp, 

473 index_new=1.0, 

474 fix_index=True, 

475 ) 

476 # These are typical values from DC2 and could/should be switched to a 

477 # more data-driven prior (from HSC?) 

478 prior_size = self.fit_coadd_multiband.size_priors[component_names_default.exp] 

479 prior_size.intercept_mag = 23.4 

480 prior_size.slope_median_per_mag = -0.14 

481 

482 

483class MultiProFitCoaddExpFitTask(MultiProFitCoaddObjectFitTask): 

484 """MultiProFit single exponential model fit task.""" 

485 

486 ConfigClass = MultiProFitCoaddExpFitConfig 

487 _DefaultName = "multiProFitCoaddExpFit" 

488 

489 

490class MultiProFitCoaddDeVFitConfig( 

491 MultiProFitCoaddSersicFitConfig, 

492 pipelineConnections=MultiProFitCoaddObjectFitConnections, 

493): 

494 """MultiProFit single DeVaucouleurs model fit task config.""" 

495 

496 @classmethod 

497 def get_model_name_default(cls) -> str: 

498 return model_names_default.deV 

499 

500 @classmethod 

501 def get_model_name_full(cls) -> str: 

502 return "de Vaucouleurs" 

503 

504 def setDefaults(self): 

505 super().setDefaults() 

506 self._rename_defaults( 

507 name_new=component_names_default.deV, 

508 name_model=model_names_default.deV, 

509 index_new=4.0, 

510 fix_index=True, 

511 ) 

512 # These are typical values from DC2 and could/should be switched to a 

513 # more data-driven prior (from HSC?). See DM-46498 for details. 

514 prior_size = self.fit_coadd_multiband.size_priors[component_names_default.deV] 

515 prior_size.intercept_mag = 21.2 

516 prior_size.slope_median_per_mag = -0.14 

517 

518 

519class MultiProFitCoaddDeVFitTask(MultiProFitCoaddObjectFitTask): 

520 """MultiProFit single DeVaucouleurs model fit task.""" 

521 

522 ConfigClass = MultiProFitCoaddDeVFitConfig 

523 _DefaultName = "multiProFitCoaddDeVFit" 

524 

525 

526class CachedChainedModelInitializer(CachedBasicModelInitializer): 

527 def get_centroid_and_shape( 

528 self, 

529 source: Mapping[str, Any], 

530 catexps: list[CatalogExposureSourcesABC], 

531 config_data: CatalogSourceFitterConfigData, 

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

533 ) -> tuple[tuple[float, float], tuple[float, float, float]]: 

534 row_best = None 

535 chisq_red_min = math.inf 

536 for name, input_data in self.inputs.items(): 

537 data = input_data.data 

538 index_row = input_data.id_index.get(source["id"]) 

539 if index_row is not None: 

540 row = data[index_row] 

541 chisq_red = input_data.get_column("chisq_reduced", data=row) 

542 if chisq_red < chisq_red_min: 

543 row_best = (row, input_data) 

544 chisq_red_min = chisq_red 

545 if row_best is None: 

546 return super().get_centroid_and_shape( 

547 source=source, 

548 catexps=catexps, 

549 config_data=config_data, 

550 values_init=values_init, 

551 ) 

552 row_best, input_data = row_best 

553 cen_x, cen_y, reff_x, reff_y, rho = ( 

554 input_data.get_column(column, data=row_best) 

555 for column in ( 

556 input_data.get_column("cen_x").name, 

557 input_data.get_column("cen_y").name, 

558 input_data.get_column(f"{input_data.size_column}_x").name, 

559 input_data.get_column(f"{input_data.size_column}_y").name, 

560 input_data.get_column("rho").name, 

561 ) 

562 ) 

563 return (cen_x, cen_y), (reff_x, reff_y, rho) 

564 

565 

566class MakeCachedChainedInitializerAction(MakeBasicInitializerAction): 

567 def _make_initializer( 

568 self, 

569 catalog_multi: Sequence, 

570 catexps: list[CatalogExposureInputs], 

571 config_data: CatalogSourceFitterConfigData, 

572 ) -> ModelInitializer: 

573 sources, priors = config_data.sources_priors 

574 return CachedChainedModelInitializer(priors=priors, sources=sources) 

575 

576 

577class MultiProFitCoaddExpDeVFitConfig( 

578 MultiProFitCoaddObjectFitConfig, 

579 pipelineConnections=MultiProFitCoaddObjectFitConnections, 

580): 

581 """MultiProFit single Exponential+DeVaucouleurs model fit task config.""" 

582 

583 @classmethod 

584 def get_model_name_default(cls) -> str: 

585 return f"{model_names_default.exp}{model_names_default.deV}" 

586 

587 @classmethod 

588 def get_model_name_full(cls) -> str: 

589 return "Exponential + de Vaucouleurs" 

590 

591 def make_default_model_config(self) -> ModelConfig: 

592 config_group = ComponentGroupConfig( 

593 components_sersic={ 

594 component_names_default.exp: self.make_sersic_component(value_initial=1.0, fixed=True), 

595 component_names_default.deV: self.make_sersic_component(value_initial=4.0, fixed=True), 

596 }, 

597 ) 

598 return self.make_single_model_config(group=config_group) 

599 

600 def setDefaults(self): 

601 super().setDefaults() 

602 self.fit_coadd_multiband.action_initializer = MakeCachedChainedInitializerAction() 

603 self.fit_coadd_multiband.config_model = self.make_default_model_config() 

604 self.name_model = self.get_model_name_default() 

605 self.connections.name_table = self.name_model 

606 

607 size_priors = self.fit_coadd_multiband.size_priors 

608 size_priors[component_names_default.exp] = MagnitudeDependentSizePriorConfig( 

609 intercept_mag=23.3, 

610 slope_median_per_mag=-0.14, 

611 ) 

612 size_priors[component_names_default.deV] = MagnitudeDependentSizePriorConfig( 

613 intercept_mag=21.2, 

614 slope_median_per_mag=-0.14, 

615 ) 

616 

617 

618class MultiProFitCoaddExpDeVFitTask(MultiProFitCoaddObjectFitTask): 

619 """MultiProFit single ExpDeV model fit task.""" 

620 

621 ConfigClass = MultiProFitCoaddExpDeVFitConfig 

622 _DefaultName = "multiProFitCoaddExpDeVFit"