Coverage for python/lsst/pipe/tasks/fit_coadd_multiband.py: 45%

82 statements  

« prev     ^ index     » next       coverage.py v7.2.5, created at 2023-05-13 10:16 +0000

1# This file is part of pipe_tasks. 

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 "CoaddMultibandFitConfig", "CoaddMultibandFitSubConfig", "CoaddMultibandFitSubTask", 

24 "CoaddMultibandFitTask", 

25] 

26 

27from .fit_multiband import CatalogExposure, CatalogExposureConfig 

28 

29import lsst.afw.table as afwTable 

30from lsst.meas.base import SkyMapIdGeneratorConfig 

31import lsst.pex.config as pexConfig 

32import lsst.pipe.base as pipeBase 

33import lsst.pipe.base.connectionTypes as cT 

34 

35import astropy 

36from abc import ABC, abstractmethod 

37from pydantic import Field 

38from pydantic.dataclasses import dataclass 

39from typing import Iterable 

40 

41CoaddMultibandFitBaseTemplates = { 

42 "name_coadd": "deep", 

43 "name_method": "multiprofit", 

44} 

45 

46 

47@dataclass(frozen=True, kw_only=True, config=CatalogExposureConfig) 

48class CatalogExposureInputs(CatalogExposure): 

49 table_psf_fits: astropy.table.Table = Field(title="A table of PSF fit parameters for each source") 

50 

51 def get_catalog(self): 

52 return self.catalog 

53 

54 

55class CoaddMultibandFitConnections( 

56 pipeBase.PipelineTaskConnections, 

57 dimensions=("tract", "patch", "skymap"), 

58 defaultTemplates=CoaddMultibandFitBaseTemplates, 

59): 

60 cat_ref = cT.Input( 

61 doc="Reference multiband source catalog", 

62 name="{name_coadd}Coadd_ref", 

63 storageClass="SourceCatalog", 

64 dimensions=("tract", "patch", "skymap"), 

65 ) 

66 cats_meas = cT.Input( 

67 doc="Deblended single-band source catalogs", 

68 name="{name_coadd}Coadd_meas", 

69 storageClass="SourceCatalog", 

70 dimensions=("tract", "patch", "band", "skymap"), 

71 multiple=True, 

72 ) 

73 coadds = cT.Input( 

74 doc="Exposures on which to run fits", 

75 name="{name_coadd}Coadd_calexp", 

76 storageClass="ExposureF", 

77 dimensions=("tract", "patch", "band", "skymap"), 

78 multiple=True, 

79 ) 

80 models_psf = cT.Input( 

81 doc="Input PSF model parameter catalog", 

82 # Consider allowing independent psf fit method 

83 name="{name_coadd}Coadd_psfs_{name_method}", 

84 storageClass="ArrowAstropy", 

85 dimensions=("tract", "patch", "band", "skymap"), 

86 multiple=True, 

87 ) 

88 models_scarlet = pipeBase.connectionTypes.Input( 

89 doc="Multiband scarlet models produced by the deblender", 

90 name="{name_coadd}Coadd_scarletModelData", 

91 storageClass="ScarletModelData", 

92 dimensions=("tract", "patch", "skymap"), 

93 ) 

94 cat_output = cT.Output( 

95 doc="Output source model fit parameter catalog", 

96 name="{name_coadd}Coadd_objects_{name_method}", 

97 storageClass="ArrowTable", 

98 dimensions=("tract", "patch", "skymap"), 

99 ) 

100 

101 def adjustQuantum(self, inputs, outputs, label, data_id): 

102 """Validates the `lsst.daf.butler.DatasetRef` bands against the 

103 subtask's list of bands to fit and drops unnecessary bands. 

104 

105 Parameters 

106 ---------- 

107 inputs : `dict` 

108 Dictionary whose keys are an input (regular or prerequisite) 

109 connection name and whose values are a tuple of the connection 

110 instance and a collection of associated `DatasetRef` objects. 

111 The exact type of the nested collections is unspecified; it can be 

112 assumed to be multi-pass iterable and support `len` and ``in``, but 

113 it should not be mutated in place. In contrast, the outer 

114 dictionaries are guaranteed to be temporary copies that are true 

115 `dict` instances, and hence may be modified and even returned; this 

116 is especially useful for delegating to `super` (see notes below). 

117 outputs : `Mapping` 

118 Mapping of output datasets, with the same structure as ``inputs``. 

119 label : `str` 

120 Label for this task in the pipeline (should be used in all 

121 diagnostic messages). 

122 data_id : `lsst.daf.butler.DataCoordinate` 

123 Data ID for this quantum in the pipeline (should be used in all 

124 diagnostic messages). 

125 

126 Returns 

127 ------- 

128 adjusted_inputs : `Mapping` 

129 Mapping of the same form as ``inputs`` with updated containers of 

130 input `DatasetRef` objects. All inputs involving the 'band' 

131 dimension are adjusted to put them in consistent order and remove 

132 unneeded bands. 

133 adjusted_outputs : `Mapping` 

134 Mapping of updated output datasets; always empty for this task. 

135 

136 Raises 

137 ------ 

138 lsst.pipe.base.NoWorkFound 

139 Raised if there are not enough of the right bands to run the task 

140 on this quantum. 

141 """ 

142 # Check which bands are going to be fit 

143 bands_fit, bands_read_only = self.config.get_band_sets() 

144 bands_needed = bands_fit.union(bands_read_only) 

145 

146 adjusted_inputs = {} 

147 for connection_name, (connection, dataset_refs) in inputs.items(): 

148 # Datasets without bands in their dimensions should be fine 

149 if 'band' in connection.dimensions: 

150 datasets_by_band = {dref.dataId['band']: dref for dref in dataset_refs} 

151 if not bands_needed.issubset(datasets_by_band.keys()): 

152 raise pipeBase.NoWorkFound( 

153 f'DatasetRefs={dataset_refs} have data with bands in the' 

154 f' set={set(datasets_by_band.keys())},' 

155 f' which is not a superset of the required bands={bands_needed} defined by' 

156 f' {self.config.__class__}.fit_coadd_multiband=' 

157 f'{self.config.fit_coadd_multiband._value.__class__}\'s attributes' 

158 f' bands_fit={bands_fit} and bands_read_only()={bands_read_only}.' 

159 f' Add the required bands={bands_needed.difference(datasets_by_band.keys())}.' 

160 ) 

161 # Adjust all datasets with band dimensions to include just 

162 # the needed bands, in consistent order. 

163 adjusted_inputs[connection_name] = ( 

164 connection, 

165 [datasets_by_band[band] for band in bands_needed] 

166 ) 

167 

168 # Delegate to super for more checks. 

169 inputs.update(adjusted_inputs) 

170 super().adjustQuantum(inputs, outputs, label, data_id) 

171 return adjusted_inputs, {} 

172 

173 

174class CoaddMultibandFitSubConfig(pexConfig.Config): 

175 """Configuration for implementing fitter subtasks. 

176 """ 

177 @abstractmethod 

178 def bands_read_only(self) -> set: 

179 """Return the set of bands that the Task needs to read (e.g. for 

180 defining priors) but not necessarily fit. 

181 

182 Returns 

183 ------- 

184 The set of such bands. 

185 """ 

186 

187 

188class CoaddMultibandFitSubTask(pipeBase.Task, ABC): 

189 """Subtask interface for multiband fitting of deblended sources. 

190 

191 Parameters 

192 ---------- 

193 **kwargs 

194 Additional arguments to be passed to the `lsst.pipe.base.Task` 

195 constructor. 

196 """ 

197 ConfigClass = CoaddMultibandFitSubConfig 

198 

199 def __init__(self, **kwargs): 

200 super().__init__(**kwargs) 

201 

202 @abstractmethod 

203 def run( 

204 self, catexps: Iterable[CatalogExposureInputs], cat_ref: afwTable.SourceCatalog 

205 ) -> pipeBase.Struct: 

206 """Fit models to deblended sources from multi-band inputs. 

207 

208 Parameters 

209 ---------- 

210 catexps : `typing.List [CatalogExposureInputs]` 

211 A list of catalog-exposure pairs with metadata in a given band. 

212 cat_ref : `lsst.afw.table.SourceCatalog` 

213 A reference source catalog to fit. 

214 

215 Returns 

216 ------- 

217 retStruct : `lsst.pipe.base.Struct` 

218 A struct with a cat_output attribute containing the output 

219 measurement catalog. 

220 

221 Notes 

222 ----- 

223 Subclasses may have further requirements on the input parameters, 

224 including: 

225 - Passing only one catexp per band; 

226 - Catalogs containing HeavyFootprints with deblended images; 

227 - Fitting only a subset of the sources. 

228 If any requirements are not met, the subtask should fail as soon as 

229 possible. 

230 """ 

231 

232 

233class CoaddMultibandFitConfig( 

234 pipeBase.PipelineTaskConfig, 

235 pipelineConnections=CoaddMultibandFitConnections, 

236): 

237 """Configure a CoaddMultibandFitTask, including a configurable fitting subtask. 

238 """ 

239 fit_coadd_multiband = pexConfig.ConfigurableField( 

240 target=CoaddMultibandFitSubTask, 

241 doc="Task to fit sources using multiple bands", 

242 ) 

243 idGenerator = SkyMapIdGeneratorConfig.make_field() 

244 

245 def get_band_sets(self): 

246 """Get the set of bands required by the fit_coadd_multiband subtask. 

247 

248 Returns 

249 ------- 

250 bands_fit : `set` 

251 The set of bands that the subtask will fit. 

252 bands_read_only : `set` 

253 The set of bands that the subtask will only read data 

254 (measurement catalog and exposure) for. 

255 """ 

256 try: 

257 bands_fit = self.fit_coadd_multiband.bands_fit 

258 except AttributeError: 

259 raise RuntimeError(f'{__class__}.fit_coadd_multiband must have bands_fit attribute') from None 

260 bands_read_only = self.fit_coadd_multiband.bands_read_only() 

261 return set(bands_fit), set(bands_read_only) 

262 

263 

264class CoaddMultibandFitTask(pipeBase.PipelineTask): 

265 """Fit deblended exposures in multiple bands simultaneously. 

266 

267 It is generally assumed but not enforced (except optionally by the 

268 configurable `fit_coadd_multiband` subtask) that there is only one exposure 

269 per band, presumably a coadd. 

270 """ 

271 ConfigClass = CoaddMultibandFitConfig 

272 _DefaultName = "CoaddMultibandFit" 

273 

274 def __init__(self, initInputs, **kwargs): 

275 super().__init__(initInputs=initInputs, **kwargs) 

276 self.makeSubtask("fit_coadd_multiband") 

277 

278 def runQuantum(self, butlerQC, inputRefs, outputRefs): 

279 inputs = butlerQC.get(inputRefs) 

280 id_tp = self.config.id_generator.apply(butlerQC.quantum.dataId).catalog_id 

281 # This is a roundabout way of ensuring all inputs get sorted and matched 

282 input_refs_objs = [(getattr(inputRefs, key), inputs[key]) 

283 for key in ("cats_meas", "coadds", "models_psf")] 

284 cats, exps, models_psf = [ 

285 {dRef.dataId: obj for dRef, obj in zip(refs, objs)} 

286 for refs, objs in input_refs_objs 

287 ] 

288 dataIds = set(cats).union(set(exps)) 

289 models_scarlet = inputs["models_scarlet"] 

290 catexps = [None]*len(dataIds) 

291 for idx, dataId in enumerate(dataIds): 

292 catalog = cats[dataId] 

293 exposure = exps[dataId] 

294 models_scarlet.updateCatalogFootprints( 

295 catalog=catalog, 

296 band=dataId['band'], 

297 psfModel=exposure.getPsf(), 

298 redistributeImage=exposure.image, 

299 removeScarletData=True, 

300 updateFluxColumns=False, 

301 ) 

302 catexps[idx] = CatalogExposureInputs( 

303 catalog=catalog, exposure=exposure, table_psf_fits=models_psf[dataId], 

304 dataId=dataId, id_tract_patch=id_tp, 

305 ) 

306 outputs = self.run(catexps=catexps, cat_ref=inputs['cat_ref']) 

307 butlerQC.put(outputs, outputRefs) 

308 

309 def run(self, catexps: list[CatalogExposure], cat_ref: afwTable.SourceCatalog) -> pipeBase.Struct: 

310 """Fit sources from a reference catalog using data from multiple 

311 exposures in the same region (patch). 

312 

313 Parameters 

314 ---------- 

315 catexps : `typing.List [CatalogExposure]` 

316 A list of catalog-exposure pairs in a given band. 

317 cat_ref : `lsst.afw.table.SourceCatalog` 

318 A reference source catalog to fit. 

319 

320 Returns 

321 ------- 

322 retStruct : `lsst.pipe.base.Struct` 

323 A struct with a cat_output attribute containing the output 

324 measurement catalog. 

325 

326 Notes 

327 ----- 

328 Subtasks may have further requirements; see `CoaddMultibandFitSubTask.run`. 

329 """ 

330 cat_output = self.fit_coadd_multiband.run(catalog_multi=cat_ref, catexps=catexps).output 

331 retStruct = pipeBase.Struct(cat_output=cat_output) 

332 return retStruct